shikiru-cli 0.2.0__py3-none-any.whl

shikiru/store/retention.py

@@ -0,0 +1,92 @@
+"""Row retention: the finished history that can be pruned from the state DB."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shikiru.domain.commands import CommandStatus
+from shikiru.domain.escalation import EscalationStatus
+from shikiru.domain.events import EventKind
+from shikiru.domain.tasks import Status
+from shikiru.store.database import format_timestamp
+
+if TYPE_CHECKING:
+    import sqlite3  # noqa: TID251
+    from collections.abc import Sequence
+    from datetime import datetime
+
+_FINISHED = (Status.DONE.value, Status.CANCELLED.value)
+_FAULT_KINDS = (EventKind.ENGINE_FAULT.value, EventKind.ENGINE_RECOVERED.value)
+_NEWEST_EVENT = "id < (SELECT MAX(id) FROM events)"
+
+
+class StateRetention:
+    """Selects and deletes rows past the retention horizon, never the ones still needed.
+
+    The newest event and command row always stay, so ids are never reused: a reused command id
+    would match the messages and events filed under the old one. Events not yet covered by the
+    notification cursor stay, so an urgent event is never pruned before it is delivered.
+    """
+
+    def __init__(self, connection: sqlite3.Connection) -> None:
+        """Use `connection`, the engine's writer."""
+        self._connection = connection
+
+    def settled_task_ids(self, finished_before: datetime) -> list[str]:
+        """Return finished tasks older than `finished_before` that still have events.
+
+        A task with an open escalation is not settled: its packet and evidence are still read.
+        """
+        rows = self._connection.execute(
+            f"""
+            SELECT id FROM tasks
+            WHERE status IN ({_marks(_FINISHED)})
+              AND COALESCE(finished_at, created_at) < ?
+              AND EXISTS (SELECT 1 FROM events WHERE events.task_id = tasks.id)
+              AND NOT EXISTS (
+                SELECT 1 FROM escalations
+                WHERE escalations.task_id = tasks.id AND escalations.status = ?
+              )
+            ORDER BY seq
+            """,
+            (*_FINISHED, format_timestamp(finished_before), EscalationStatus.OPEN.value),
+        ).fetchall()
+        return [row["id"] for row in rows]
+
+    def delete_task_events(self, task_ids: Sequence[str], *, delivered_up_to: int) -> int:
+        """Delete the events of `task_ids` that were delivered; return how many went."""
+        deleted = 0
+        for task_id in task_ids:
+            deleted += self._connection.execute(
+                f"DELETE FROM events WHERE task_id = ? AND id <= ? AND {_NEWEST_EVENT}",
+                (task_id, delivered_up_to),
+            ).rowcount
+        return deleted
+
+    def delete_engine_events(self, before: datetime, *, delivered_up_to: int) -> int:
+        """Delete delivered engine-level events older than `before`, but not fault records.
+
+        A fault or its recovery stays: `status` derives the open fault streaks from them.
+        """
+        return self._connection.execute(
+            f"""
+            DELETE FROM events
+            WHERE task_id IS NULL AND ts < ? AND id <= ? AND {_NEWEST_EVENT}
+              AND kind NOT IN ({_marks(_FAULT_KINDS)})
+            """,
+            (format_timestamp(before), delivered_up_to, *_FAULT_KINDS),
+        ).rowcount
+
+    def delete_finished_commands(self, before: datetime) -> int:
+        """Delete answered commands older than `before`; a pending command is never deleted."""
+        return self._connection.execute(
+            """
+            DELETE FROM commands
+            WHERE status != ? AND ts < ? AND id < (SELECT MAX(id) FROM commands)
+            """,
+            (CommandStatus.PENDING.value, format_timestamp(before)),
+        ).rowcount
+
+
+def _marks(values: Sequence[str]) -> str:
+    return ", ".join("?" for _ in values)

shikiru/store/runs.py

@@ -0,0 +1,227 @@
+"""Agent session runs: one per task role's session lineage."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shikiru.agents.contracts import (
+    UNKNOWN_BILLING_DECLARATION,
+    BillingDeclaration,
+    Usage,
+)
+from shikiru.domain.runs import Run, RunStatus
+from shikiru.domain.tasks import Role
+from shikiru.store.database import (
+    format_timestamp,
+    parse_required_timestamp,
+    parse_timestamp,
+    require_after_insert,
+)
+
+if TYPE_CHECKING:
+    import sqlite3  # noqa: TID251
+
+    from shikiru.agents.contracts import Rung
+    from shikiru.clock import Clock
+
+_RUN_COLUMNS = (
+    "id, task_id, role, agent, model, effort, session_id, resume_command, "
+    "last_message, status, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, "
+    "cost_usd, cost_basis, context_used_pct, usage_reported, billing_mode, billing_source, "
+    "started_at, ended_at"
+)
+MAX_TAKEN_OVER_ROLE_COUNT = len(Role)
+
+
+def _row_to_run(row: sqlite3.Row) -> Run:
+    return Run(
+        id=row["id"],
+        task_id=row["task_id"],
+        role=Role(row["role"]),
+        agent=row["agent"],
+        model=row["model"],
+        effort=row["effort"],
+        session_id=row["session_id"],
+        resume_command=row["resume_command"],
+        last_message=row["last_message"],
+        status=RunStatus(row["status"]),
+        usage=Usage(
+            input_tokens=row["input_tokens"],
+            output_tokens=row["output_tokens"],
+            cache_read_tokens=row["cache_read_tokens"],
+            cache_write_tokens=row["cache_write_tokens"],
+            cost_usd=row["cost_usd"],
+            cost_basis=row["cost_basis"],
+            context_used_pct=row["context_used_pct"],
+        ),
+        usage_reported=(bool(row["usage_reported"]) if row["usage_reported"] is not None else None),
+        billing_mode=row["billing_mode"],
+        billing_source=row["billing_source"],
+        started_at=parse_required_timestamp(row["started_at"]),
+        ended_at=parse_timestamp(row["ended_at"]),
+    )
+
+
+def fetch_run(connection: sqlite3.Connection, run_id: int) -> Run | None:
+    """Read one run by id, or `None` if it doesn't exist."""
+    row = connection.execute(f"SELECT {_RUN_COLUMNS} FROM runs WHERE id = ?", (run_id,)).fetchone()
+    return _row_to_run(row) if row is not None else None
+
+
+def fetch_runs_for_task(connection: sqlite3.Connection, task_id: str) -> list[Run]:
+    """Read every run recorded for `task_id`, most recently started first."""
+    rows = connection.execute(
+        f"SELECT {_RUN_COLUMNS} FROM runs WHERE task_id = ? ORDER BY id DESC", (task_id,)
+    ).fetchall()
+    return [_row_to_run(row) for row in rows]
+
+
+def fetch_taken_over_roles_for_task(
+    connection: sqlite3.Connection,
+    task_id: str,
+) -> tuple[Role, ...]:
+    rows = connection.execute(
+        "SELECT role FROM runs WHERE task_id = ? AND status = ? ORDER BY id DESC LIMIT ?",
+        (task_id, RunStatus.TAKEN_OVER.value, MAX_TAKEN_OVER_ROLE_COUNT),
+    ).fetchall()
+    return tuple(Role(row["role"]) for row in rows)
+
+
+def fetch_all_runs(connection: sqlite3.Connection) -> list[Run]:
+    """Read every run, most recently started first."""
+    rows = connection.execute(f"SELECT {_RUN_COLUMNS} FROM runs ORDER BY id DESC").fetchall()
+    return [_row_to_run(row) for row in rows]
+
+
+class RunStore:
+    """Starts, updates and reads agent session runs."""
+
+    def __init__(self, connection: sqlite3.Connection, clock: Clock) -> None:
+        """Use `connection` (the engine's writer) and `clock` for every write below."""
+        self._connection = connection
+        self._clock = clock
+
+    def start(
+        self,
+        task_id: str,
+        role: Role,
+        rung: Rung,
+        *,
+        billing_declaration: BillingDeclaration = UNKNOWN_BILLING_DECLARATION,
+    ) -> Run:
+        """Start a new live run for `task_id`'s `role` on `rung`."""
+        cursor = self._connection.execute(
+            """
+            INSERT INTO runs (
+              task_id, role, agent, model, effort, status, usage_reported,
+              billing_mode, billing_source, started_at
+            ) VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?, ?)
+            """,
+            (
+                task_id,
+                role.value,
+                rung.agent,
+                rung.model,
+                rung.effort,
+                RunStatus.LIVE.value,
+                billing_declaration.mode,
+                billing_declaration.source,
+                format_timestamp(self._clock.now()),
+            ),
+        )
+        run_id = require_after_insert(
+            cursor.lastrowid, "sqlite did not allocate a rowid for the new run"
+        )
+        run = fetch_run(self._connection, run_id)
+        return require_after_insert(run, f"run {run_id} vanished immediately after insert")
+
+    def get_live(self, task_id: str, role: Role) -> Run | None:
+        """Return `task_id`'s live run for `role`, if any."""
+        return self._get_with_status(task_id, role, RunStatus.LIVE)
+
+    def get_taken_over(self, task_id: str, role: Role) -> Run | None:
+        """Return `task_id`'s run for `role` that the Operator has taken over, if any."""
+        return self._get_with_status(task_id, role, RunStatus.TAKEN_OVER)
+
+    def mark_taken_over(self, run_id: int) -> None:
+        """Mark `run_id` taken over: its session runs interactively in the Operator's pane."""
+        self._set_status_keeping_open(run_id, RunStatus.TAKEN_OVER)
+
+    def mark_live(self, run_id: int) -> None:
+        """Mark a taken-over `run_id` live again: the engine drives its session."""
+        self._set_status_keeping_open(run_id, RunStatus.LIVE)
+
+    def has_taken_over(self, task_id: str) -> bool:
+        """Report whether the Operator has taken over any role of `task_id`."""
+        row = self._connection.execute(
+            "SELECT 1 FROM runs WHERE task_id = ? AND status = ? LIMIT 1",
+            (task_id, RunStatus.TAKEN_OVER.value),
+        ).fetchone()
+        return row is not None
+
+    def _get_with_status(self, task_id: str, role: Role, status: RunStatus) -> Run | None:
+        row = self._connection.execute(
+            f"SELECT {_RUN_COLUMNS} FROM runs WHERE task_id = ? AND role = ? AND status = ?",
+            (task_id, role.value, status.value),
+        ).fetchone()
+        return _row_to_run(row) if row is not None else None
+
+    def _set_status_keeping_open(self, run_id: int, status: RunStatus) -> None:
+        self._connection.execute("UPDATE runs SET status = ? WHERE id = ?", (status.value, run_id))
+
+    def for_task(self, task_id: str) -> list[Run]:
+        """Return every run recorded for `task_id`."""
+        return fetch_runs_for_task(self._connection, task_id)
+
+    def set_session(self, run_id: int, session_id: str | None, resume_command: str | None) -> None:
+        """Record `run_id`'s launched session id and resume command.
+
+        `session_id` is `None` when `SessionStarted` says the session can't be resumed,
+        and `resume_command` is `None` when the agent's `resume_command()` returns none.
+        """
+        self._connection.execute(
+            "UPDATE runs SET session_id = ?, resume_command = ? WHERE id = ?",
+            (session_id, resume_command, run_id),
+        )
+
+    def set_rung(self, run_id: int, rung: Rung) -> None:
+        """Record a same-agent model or effort change on an existing session."""
+        self._connection.execute(
+            "UPDATE runs SET model = ?, effort = ? WHERE id = ?",
+            (rung.model, rung.effort, run_id),
+        )
+
+    def set_last_message(self, run_id: int, message: str) -> None:
+        """Persist `run_id`'s latest known message from the agent."""
+        self._connection.execute("UPDATE runs SET last_message = ? WHERE id = ?", (message, run_id))
+
+    def record_usage(self, run_id: int, usage: Usage) -> None:
+        """Overwrite `run_id`'s usage with `usage`'s latest cumulative snapshot."""
+        self._connection.execute(
+            """
+            UPDATE runs SET input_tokens = ?, output_tokens = ?, cache_read_tokens = ?,
+              cache_write_tokens = ?, cost_usd = ?, cost_basis = ?, context_used_pct = ?,
+              usage_reported = 1
+            WHERE id = ?
+            """,
+            (
+                usage.input_tokens,
+                usage.output_tokens,
+                usage.cache_read_tokens,
+                usage.cache_write_tokens,
+                usage.cost_usd,
+                usage.cost_basis,
+                usage.context_used_pct,
+                run_id,
+            ),
+        )
+
+    def close(self, run_id: int) -> None:
+        """Mark `run_id` closed: the task ended, the session is not resumable."""
+        self._set_status(run_id, RunStatus.CLOSED)
+
+    def _set_status(self, run_id: int, status: RunStatus) -> None:
+        self._connection.execute(
+            "UPDATE runs SET status = ?, ended_at = ? WHERE id = ?",
+            (status.value, format_timestamp(self._clock.now()), run_id),
+        )