diff --git a/docs/index.md b/docs/index.md index 2299240..a96260d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -33,6 +33,7 @@ | [`mcp.md`](mcp.md) | MCP servers — configs, tool exposure, per-user keys (BYOK) | | [`sessions.md`](sessions.md) | Session model, dual-buffer design, context compression | | [`store.md`](store.md) | KV store — per-session/per-user key-value persistence | +| [`todo-store.md`](todo-store.md) | Todo plans — dedicated table, migration, crash recovery, protocol step insertion | | [`recall.md`](recall.md) | Scheduled callbacks — headless recall system | | [`tasks.md`](tasks.md) | Background tasks — detach tool calls, TaskManager, completion notes, message queue | | [`websocket.md`](websocket.md) | WebSocket protocol — all events, stop mechanism | diff --git a/docs/store.md b/docs/store.md index 0e24f12..61c999b 100644 --- a/docs/store.md +++ b/docs/store.md @@ -7,9 +7,12 @@ Provides a simple API for tools to persist per-session state without inventing ad-hoc storage mechanisms. Current consumers: -- **todo** — task plans (`scope='todo'`) - **scratchpad** — working notes (`scope='scratchpad'`) +Todo plans used to live here (`scope='todo'`) and now have their own table — +see [todo-store.md](todo-store.md). The rows were migrated on first start and +removed from `session_store`. + Future consumers: recall context, wizard state, user preferences per session. ## Schema @@ -45,9 +48,9 @@ - `user_id` — nullable for legacy single-user mode. - `session_id` — the chat session UUID. -- `scope` — namespace: `todo`, `scratchpad`, or any custom string. +- `scope` — namespace: `scratchpad`, or any custom string. - `key` — arbitrary string identifier within the scope. ## Wiring -`KvStore` is created as a singleton in `navi/api/deps.py` (`get_kv_store()`) and injected into `build_default_registries()`, which passes it to `TodoTool` and `ScratchpadTool` on startup. +`KvStore` is created as a singleton in `navi/core/container.py` (`get_kv_store()` exposes it to routes) and injected into `build_default_registries()`, which passes it to `ScratchpadTool` on startup. `TodoTool` gets the dedicated `PgTodoStore` from the same container. diff --git a/docs/todo-store.md b/docs/todo-store.md new file mode 100644 index 0000000..5837e49 --- /dev/null +++ b/docs/todo-store.md @@ -0,0 +1,102 @@ +# Todo store — plans that outlive the process + +The `todo` tool's plan lives in its own PostgreSQL table, one row per +`(user_id, session_id)`. It used to be an opaque JSON string inside +`session_store` (`scope='todo'`); this page covers why it moved and what the +row now carries. + +## Why a separate table + +The old row already survived restarts — the problem was that it could not tell +a live plan from one a crashed run abandoned. Reads were untyped JSON parsed on +every call, and there was no timestamp the recovery path could trust. The +dedicated row adds a `revision` counter and an `updated_at` that the agent can +reason about, plus room for per-step provenance that protocols will use. + +## Schema + +```sql +CREATE TABLE session_todos ( + user_id TEXT NOT NULL DEFAULT '', + session_id TEXT NOT NULL, + tasks JSONB NOT NULL DEFAULT '[]'::jsonb, + revision INTEGER NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (user_id, session_id) +); +``` + +`tasks` is an array of objects, one per step: + +| Field | Meaning | +|-------|---------| +| `text` | The step as shown to the model and the client. | +| `status` | `pending` \| `in_progress` \| `done` \| `failed` \| `skipped`. | +| `validation` | How the result was verified. Required for `done`. | +| `milestone` | Grouping label from the plan (`"A"`, `"B"`, …). | +| `source` | `"plan"` for the agent's own plan, `"protocol"` for injected protocol steps. | +| `step_id` | Stable id of the step inside the protocol definition. | +| `protocol_run_id` | Which protocol run injected the step. | + +Rows written before these fields existed load with their defaults, so no +backfill is needed. + +`revision` starts at `0` and increments on every save. It is not a concurrency +control yet — it exists so the recovery note can fire once per plan version. + +## Migration + +`PgTodoStore` creates the table and runs the migration lazily on first use +(same pattern as `KvStore`). The migration copies every `session_store` row with +`scope='todo'` into `session_todos` and then **deletes** the source rows: + +- `ON CONFLICT DO NOTHING` makes the copy idempotent, so a plan written to the + new table is never overwritten by a stale KV copy; +- deleting the source is what stops a plan the user cleared from being + resurrected on the next process start. + +If `session_store` does not exist yet, the migration is logged and skipped — +the table simply starts empty. Running the old code after this migration will +not see existing plans; that is the cost of the cut-over and the reason it is +done in one release. + +## Crash recovery + +When the plan reminder is built (`get_progress_message`), a plan with a step +still `in_progress` whose row was last updated *before this process started* is +flagged once: + +> Recovery: this plan was last changed at … and step N is still +> 'in_progress'. If no run is working on it right now, a previous run stopped +> mid-step — check the actual state … before continuing. + +The note is keyed by `(session_id, revision)`, so it appears once per plan +version and not on every iteration. A plan updated by the current process gets +no note. The wording is deliberately conditional ("if no run is working on it +right now") because a multi-worker deployment can legitimately have another +process writing the row. + +## Wiring + +`PgTodoStore(pool)` is created in `navi/core/container.py`, stored on +`AppContainer.todo_store`, passed into `build_default_registries()` and onward +to `TodoTool`. `todo.set_store()` makes it the module-global backend. If no +store is wired, `TodoTool` falls back to `KvTodoStore` over the shared +`KvStore` — that path is what unit tests use and what an instance that has not +migrated yet would still read. + +`load_tasks_for(session_id, user_id, kv=None, store=None)` resolves in order: +explicit `store`, explicit `kv`, module-global store. The REST endpoint +`GET /sessions/{id}/todos` passes both (the container's store, and the KV store +as a fallback), so its payload shape is unchanged. + +## Helpers for protocols + +`navi/tools/todo.py` exposes two internals that the protocol work builds on: + +- `insert_steps(session_id, steps, after=None, source="protocol", protocol_run_id=…, step_ids=…)` + splices a block into the existing plan after the last step the agent reached + (or after an explicit index) and returns the new 1-based indexes. Existing + steps and statuses are preserved — no re-issue of the whole plan via `set`. +- `_current_position(tasks)` finds that insertion point (last `done` / + `in_progress` / `failed` step). diff --git a/manuals/todo.md b/manuals/todo.md index eea4253..735d710 100644 --- a/manuals/todo.md +++ b/manuals/todo.md @@ -5,7 +5,7 @@ The tool is the record of *what you actually verified*. Marking a step `done` without a `validation` field is rejected outright, which is the point: it stops "done" from meaning "I moved on". -Tasks are per-session and do not survive a server restart. +Tasks are per-session and live in PostgreSQL (`session_todos`), so they survive a server restart. When the previous run died with a step still `in_progress`, the plan reminder says so once — re-check that step's real result before marking it done. ## Parameters diff --git a/navi/api/deps.py b/navi/api/deps.py index d3e6950..59a040f 100644 --- a/navi/api/deps.py +++ b/navi/api/deps.py @@ -18,6 +18,7 @@ from navi.memory import MemoryStore from navi.mcp import McpManager from navi.store import KvStore +from navi.todo_store import PgTodoStore from navi.core.scheduler import RecallScheduler from navi.workers.base import Worker from navi.auth.deps import ( @@ -62,6 +63,10 @@ return _resolve_container().kv_store +def get_todo_store() -> PgTodoStore: + return _resolve_container().todo_store + + def get_tool_registry() -> ToolRegistry: return _resolve_container().tool_registry diff --git a/navi/api/routes/sessions.py b/navi/api/routes/sessions.py index dde07f8..8a5954c 100644 --- a/navi/api/routes/sessions.py +++ b/navi/api/routes/sessions.py @@ -19,6 +19,7 @@ get_profile_registry, get_scheduler, get_session_store, + get_todo_store, require_admin, require_user, ) @@ -35,6 +36,7 @@ from navi.profiles.base import admin_only_blocked from navi.session_files import delete_session_dir, ensure_session_dir, is_forbidden, list_session_files, safe_filename, session_dir from navi.store import KvStore +from navi.todo_store import PgTodoStore log = structlog.get_logger() @@ -484,11 +486,12 @@ session_id: str, store: Annotated[SessionStore, Depends(get_session_store)], kv: Annotated[KvStore, Depends(get_kv_store)], + todos: Annotated[PgTodoStore, Depends(get_todo_store)], user: Annotated[User, Depends(require_user)], ) -> dict: """Return the session's current todo list (the parent agent's plan). - Reads the parent session's KV row via explicit (user_id, session_id) — no + Reads the parent session's todo row via explicit (user_id, session_id) — no ContextVars. Sub-agent todos live in isolated rows and are not exposed here (Phase 3). Used by the TUI to seed the side panel on attach/switch. """ @@ -498,7 +501,7 @@ if session is None: raise HTTPException(status_code=404, detail="Session not found") check_session_access(session, user, permission="navi.sessions.read_all") - tasks = await load_tasks_for(session_id, session.user_id, kv=kv) + tasks = await load_tasks_for(session_id, session.user_id, kv=kv, store=todos) return {"session_id": session_id, "tasks": tasks} diff --git a/navi/core/container.py b/navi/core/container.py index 0e082ba..324e495 100644 --- a/navi/core/container.py +++ b/navi/core/container.py @@ -20,6 +20,7 @@ from navi.memory import MemoryStore from navi.mcp import McpManager from navi.store import KvStore + from navi.todo_store import PgTodoStore from navi.workers.base import Worker from navi.push.service import PushService @@ -32,6 +33,7 @@ memory_store: "MemoryStore" = None # type: ignore[assignment] session_store: "SessionStore" = None # type: ignore[assignment] kv_store: "KvStore" = None # type: ignore[assignment] + todo_store: "PgTodoStore" = None # type: ignore[assignment] scheduler: "RecallScheduler" = None # type: ignore[assignment] tool_registry: "ToolRegistry" = None # type: ignore[assignment] profile_registry: "ProfileRegistry" = None # type: ignore[assignment] @@ -106,6 +108,7 @@ from navi.mcp import McpManager, load_mcp_servers from navi.mcp.tools import McpTool from navi.store import KvStore + from navi.todo_store import PgTodoStore from navi.workers import build_default_workers if not settings.database_url: @@ -117,6 +120,7 @@ session_store = PgSessionStore(pool) memory_store = MemoryStore(pool) kv_store = KvStore(pool) + todo_store = PgTodoStore(pool) scheduler = RecallScheduler(pool) # Create MCP manager early so build_default_registries can wire it in @@ -139,6 +143,7 @@ kv_store=kv_store, mcp_manager=mcp_manager, terminal_manager=terminal_manager, + todo_store=todo_store, ) # Wire embedding backend into memory store @@ -226,6 +231,7 @@ memory_store=memory_store, session_store=session_store, kv_store=kv_store, + todo_store=todo_store, scheduler=scheduler, tool_registry=tool_registry, profile_registry=profile_registry, diff --git a/navi/core/registry.py b/navi/core/registry.py index ab2180e..a8d1df6 100644 --- a/navi/core/registry.py +++ b/navi/core/registry.py @@ -207,6 +207,7 @@ kv_store=None, mcp_manager=None, terminal_manager=None, + todo_store=None, ) -> tuple[ToolRegistry, ProfileRegistry, BackendRegistry, ContextProviderRegistry]: """Build and populate registries with all built-in components.""" from navi.core.ai_helper import AIHelper @@ -276,7 +277,7 @@ builtins = [FilesystemTool(ai_helper=ai_helper), CodeExecTool(), TerminalTool(terminal_manager=terminal_manager), SshExecTool(), ImageViewTool(), ShareFileTool(), ContentPublishTool(), - TodoTool(kv_store=kv_store), ScratchpadTool(kv_store=kv_store), + TodoTool(kv_store=kv_store, todo_store=todo_store), ScratchpadTool(kv_store=kv_store), ReflectTool(ai_helper=ai_helper), PlanTool(), PeerTool(), diff --git a/navi/todo_store.py b/navi/todo_store.py new file mode 100644 index 0000000..4017a2d --- /dev/null +++ b/navi/todo_store.py @@ -0,0 +1,202 @@ +"""Persistent todo plans — one row per (user, session) in PostgreSQL. + +Todo used to live as an opaque JSON string inside the generic KV store +(``session_store``, ``scope='todo'``). It survived restarts there too, but the +row was untyped, every read parsed JSON, and nothing could tell a live plan +from one abandoned by a crashed run. This module gives the plan its own table +with a revision counter and an ``updated_at`` the recovery path can reason +about. + +The KV row is *migrated*, then removed: the first time the table is created, +existing plans are copied over so a session that was mid-task before this +change resumes with its plan and statuses intact, and the old rows are deleted +afterwards (a stale copy would otherwise resurrect a plan the user cleared). + +The tool API is unchanged — :class:`navi.tools.todo.TodoTool` still speaks +strings; structured per-step fields (``source``, ``step_id``, +``protocol_run_id``) ride along in the JSON and default to empty. +""" + +from __future__ import annotations + +import asyncio +import json +from datetime import datetime +from typing import Any, Protocol + +import asyncpg +import structlog + +log = structlog.get_logger() + +_DDL = """ +CREATE TABLE IF NOT EXISTS session_todos ( + user_id TEXT NOT NULL DEFAULT '', + session_id TEXT NOT NULL, + tasks JSONB NOT NULL DEFAULT '[]'::jsonb, + revision INTEGER NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (user_id, session_id) +); +""" + +# One-time copy of the plans that predate this table. ``ON CONFLICT DO NOTHING`` +# makes it idempotent: a plan that already exists here is newer than (or equal +# to) its KV copy, and must not be overwritten by a stale row. +_MIGRATE_INSERT = """ +INSERT INTO session_todos (user_id, session_id, tasks, revision, updated_at) +SELECT COALESCE(user_id, ''), session_id, value::jsonb, 0, updated_at + FROM session_store + WHERE scope = 'todo' AND key = 'tasks' AND value <> '' +ON CONFLICT (user_id, session_id) DO NOTHING +""" + +# The KV row must not survive the migration. Otherwise a plan the user clears +# here would be resurrected from the stale copy on the next process start. +_MIGRATE_DROP_SOURCE = """ +DELETE FROM session_store WHERE scope = 'todo' AND key = 'tasks' +""" + + +def _norm_uid(user_id: str | None) -> str: + return user_id or "" + + +class TodoStore(Protocol): + """Storage contract the todo tool depends on.""" + + async def load(self, user_id: str | None, session_id: str) -> list[dict[str, Any]] | None: + ... + + async def load_meta(self, user_id: str | None, session_id: str) -> dict | None: + """The row's revision/updated_at, or None when no plan exists.""" + ... + + async def save(self, user_id: str | None, session_id: str, tasks: list[dict[str, Any]]) -> None: + ... + + async def clear(self, user_id: str | None, session_id: str) -> None: + ... + + +class PgTodoStore: + """Todo plans in their own table; DDL and migration run once, lazily.""" + + def __init__(self, pool: asyncpg.Pool) -> None: + self._pool = pool + self._ready = False + self._lock = asyncio.Lock() + + async def _ensure_ready(self) -> None: + if self._ready: + return + async with self._lock: + if self._ready: + return + async with self._pool.acquire() as conn: + await conn.execute(_DDL) + try: + await conn.execute(_MIGRATE_INSERT) + await conn.execute(_MIGRATE_DROP_SOURCE) + except asyncpg.PostgresError as e: + # A missing session_store (fresh install, table created by + # another subsystem later) is not a reason to lose todos for + # good — log it and carry on with an empty table. + log.warning("todo.migration_skipped", error=str(e)) + self._ready = True + + async def load(self, user_id: str | None, session_id: str) -> list[dict[str, Any]] | None: + await self._ensure_ready() + async with self._pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT tasks FROM session_todos WHERE user_id = $1 AND session_id = $2", + _norm_uid(user_id), session_id, + ) + if row is None: + return None + raw = row["tasks"] + if isinstance(raw, str): + # asyncpg returns JSONB as str unless a codec is registered. + try: + raw = json.loads(raw) + except ValueError: + return None + return list(raw) if isinstance(raw, list) else None + + async def load_meta(self, user_id: str | None, session_id: str) -> dict | None: + await self._ensure_ready() + async with self._pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT revision, updated_at FROM session_todos " + "WHERE user_id = $1 AND session_id = $2", + _norm_uid(user_id), session_id, + ) + return ( + {"revision": row["revision"], "updated_at": row["updated_at"]} + if row is not None else None + ) + + async def save(self, user_id: str | None, session_id: str, tasks: list[dict[str, Any]]) -> None: + await self._ensure_ready() + async with self._pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO session_todos (user_id, session_id, tasks, revision, updated_at) + VALUES ($1, $2, $3::jsonb, 1, now()) + ON CONFLICT (user_id, session_id) DO UPDATE SET + tasks = EXCLUDED.tasks, + revision = session_todos.revision + 1, + updated_at = now() + """, + _norm_uid(user_id), session_id, json.dumps(tasks, ensure_ascii=False), + ) + + async def clear(self, user_id: str | None, session_id: str) -> None: + await self._ensure_ready() + async with self._pool.acquire() as conn: + await conn.execute( + "DELETE FROM session_todos WHERE user_id = $1 AND session_id = $2", + _norm_uid(user_id), session_id, + ) + + +class KvTodoStore: + """Adapter over the legacy KV row. + + Kept for two reasons: tests inject a fake KvStore, and an instance that has + not run the migration yet still answers from the old location. + """ + + def __init__(self, kv) -> None: + self._kv = kv + + async def load(self, user_id: str | None, session_id: str) -> list[dict[str, Any]] | None: + raw = await self._kv.get(_norm_uid(user_id), session_id, "todo", "tasks") + if not raw: + return None + try: + data = json.loads(raw) + except ValueError: + return None + return list(data) if isinstance(data, list) else None + + async def load_meta(self, user_id: str | None, session_id: str) -> dict | None: + # The KV row carries no revision and no timestamp the store exposes. + return None + + async def save(self, user_id: str | None, session_id: str, tasks: list[dict[str, Any]]) -> None: + await self._kv.set( + _norm_uid(user_id), session_id, "todo", "tasks", + json.dumps(tasks, ensure_ascii=False), + ) + + async def clear(self, user_id: str | None, session_id: str) -> None: + await self._kv.clear_scope(_norm_uid(user_id), session_id, "todo") + + +def plan_timestamp(meta: dict | None) -> datetime | None: + """``updated_at`` from a :meth:`TodoStore.load_meta` result, if it has one.""" + if not meta: + return None + value = meta.get("updated_at") + return value if isinstance(value, datetime) else None diff --git a/navi/tools/todo.py b/navi/tools/todo.py index 09c3a8d..474c825 100644 --- a/navi/tools/todo.py +++ b/navi/tools/todo.py @@ -1,12 +1,13 @@ -"""Session-scoped task plan manager — backed by PostgreSQL KV store.""" +"""Session-scoped task plan manager — backed by PostgreSQL (session_todos).""" from __future__ import annotations -import json -from dataclasses import dataclass +from dataclasses import dataclass, fields +from datetime import UTC, datetime import structlog from navi.llm.base import Message +from navi.todo_store import KvTodoStore, plan_timestamp from navi.tools._internal.base import ( Tool, ToolContext, @@ -60,18 +61,46 @@ status: str = "pending" validation: str = "" # how the result was verified (required for done, encouraged for failed) milestone: str = "" # grouping label (e.g. "A"/"B") from the plan's Milestones; "" = ungrouped + # Provenance, used once protocols inject their own steps into a plan: + source: str = "plan" # "plan" | "protocol" + step_id: str = "" # stable id inside the protocol definition + protocol_run_id: str = "" # which run injected this step -# Global KV store reference — injected at startup by registry.py +_TASK_FIELDS = tuple(f.name for f in fields(_Task)) + + +def _task_dict(t: _Task) -> dict: + return {name: getattr(t, name) for name in _TASK_FIELDS} + + +# Storage: a dedicated PgTodoStore when the container wires one, otherwise the +# legacy KV row (tests inject a fake KvStore, and an instance that has not run +# the migration still answers from the old location). +_store = None _kv_store = None +def set_store(store) -> None: + """Inject the todo store (PgTodoStore) — called once at startup.""" + global _store + _store = store + + def set_kv_store(kv) -> None: - """Inject the shared KvStore instance (called once at startup).""" + """Inject the shared KvStore instance (legacy path; called once at startup).""" global _kv_store _kv_store = kv +def _active_store(): + if _store is not None: + return _store + if _kv_store is not None: + return KvTodoStore(_kv_store) + return None + + def _sid(explicit: str | None = None) -> str: # Sub-agents set current_todo_session_id to their ephemeral run id so their # plan/todo stay isolated from the parent session's todo. It must win over @@ -91,23 +120,24 @@ async def _load_tasks(sid: str) -> list[_Task]: - if _kv_store is None: + store = _active_store() + if store is None: return [] - raw = await _kv_store.get(_uid(), sid, "todo", "tasks") - if not raw: + data = await store.load(_uid(), sid) + if not data: return [] try: - data = json.loads(raw) - return [_Task(**item) for item in data] + # Rows written before the provenance fields existed simply default. + return [_Task(**{k: v for k, v in item.items() if k in _TASK_FIELDS}) for item in data] except Exception: return [] async def _save_tasks(sid: str, tasks: list[_Task]) -> None: - if _kv_store is None: + store = _active_store() + if store is None: return - data = [{"text": t.text, "status": t.status, "validation": t.validation, "milestone": t.milestone} for t in tasks] - await _kv_store.set(_uid(), sid, "todo", "tasks", json.dumps(data)) + await store.save(_uid(), sid, [_task_dict(t) for t in tasks]) class TodoTool(Tool): @@ -166,7 +196,9 @@ "required": ["op"], } - def __init__(self, kv_store=None) -> None: + def __init__(self, kv_store=None, todo_store=None) -> None: + if todo_store is not None: + set_store(todo_store) if kv_store is not None: set_kv_store(kv_store) @@ -257,8 +289,12 @@ return ToolResult(success=True, output=self._render(sid, tasks)) if op == "clear": - if _kv_store is not None: - await _kv_store.clear_scope(_uid(ctx.user_id if ctx else None), sid, "todo") + store = _active_store() + if store is not None: + # Same user scope as load/save: the ContextVar wins, and the + # explicit ctx.user_id is only a fallback. Using ctx alone here + # would delete a different row than the one the tool writes. + await store.clear(_uid(), sid) return ToolResult(success=True, output="Plan cleared.") return ToolResult( @@ -341,6 +377,14 @@ icon_failed = _STATUS_ICON["failed"] lines = [f"TODO progress: {done}/{n} {icon_done} {failed} {icon_failed} ({in_progress} in progress, {pending} pending)"] + # A plan that was in flight when the previous run died reads exactly + # like a live one: a step sits at in_progress and the statuses look + # current. Say so once, so the model re-checks reality instead of + # trusting a step that may have been half-done (or not done at all). + restart_note = await _restart_note(session_id, tasks) + if restart_note: + lines.append(restart_note) + # List remaining tasks remaining = [t for t in tasks if t.status not in ("done", "skipped")] if remaining: @@ -441,26 +485,21 @@ async def load_tasks_for( - session_id: str, user_id: str | None, kv=None + session_id: str, user_id: str | None, kv=None, store=None ) -> list[dict]: """Load the todo list for an explicit (user_id, session_id) pair as JSON-serialisable dicts — no ContextVars. Used by the REST endpoint (GET /sessions/{id}/todos) and the TodoUpdated emitter, both of which run outside tool execution where the per-session - ContextVars are not reliably set. Pass ``kv`` to bypass the module-global - store (the REST route resolves its own KvStore via deps). + ContextVars are not reliably set. Precedence: an explicitly passed + ``store``, then an explicitly passed ``kv``, then the module-global store + (the dedicated table when wired, else the legacy KV row). """ - store = kv if kv is not None else _kv_store - if store is None: + backend = store if store is not None else (KvTodoStore(kv) if kv is not None else _active_store()) + if backend is None: return [] - raw = await store.get(user_id, session_id, "todo", "tasks") - if not raw: - return [] - try: - data = json.loads(raw) - except Exception: - return [] + data = await backend.load(user_id, session_id) or [] return [ { "index": i + 1, @@ -470,6 +509,7 @@ "milestone": item.get("milestone", ""), } for i, item in enumerate(data) + if isinstance(item, dict) ] @@ -483,3 +523,98 @@ icon = _STATUS_ICON.get(t.status, "?") lines.append(f" {icon} [{i}] {t.text} ({t.status})") return lines + + +# ── Crash recovery ────────────────────────────────────────────────────────── + +_PROCESS_STARTED = datetime.now(UTC) +_restart_notified: set[tuple[str, int]] = set() + + +async def _restart_note(sid: str, tasks: list[_Task]) -> str | None: + """One-shot note when an in-progress step predates this process. + + ``updated_at`` older than the process start means no live run owns the + plan: either the server restarted mid-task or the plan was abandoned. + Either way the in_progress marker says nothing about the real world, so + the model is told to verify before continuing. + + Only the dedicated store carries a timestamp; the KV fallback returns no + metadata and no note is produced. + """ + if not any(t.status == "in_progress" for t in tasks): + return None + store = _active_store() + if store is None: + return None + try: + meta = await store.load_meta(_uid(), sid) + except Exception: + return None + updated = plan_timestamp(meta) + if updated is None: + return None + if updated.tzinfo is None: + updated = updated.replace(tzinfo=UTC) + if updated >= _PROCESS_STARTED: + return None + key = (sid, int((meta or {}).get("revision") or 0)) + if key in _restart_notified: + return None + _restart_notified.add(key) + idx = next(i for i, t in enumerate(tasks, 1) if t.status == "in_progress") + return ( + f"Recovery: this plan was last changed at {updated:%Y-%m-%d %H:%M} UTC and step {idx} " + "is still 'in_progress'. If no run is working on it right now, a previous run stopped " + "mid-step — check the actual state of that step's work (files, services, command " + "output) before continuing: mark it done only if its result really exists, otherwise " + "redo it or mark it failed with what you found. If the plan belongs to an older task, " + "clear it and plan again." + ) + + +# ── Plan surgery (used by protocol_start) ─────────────────────────────────── + +def _current_position(tasks: list[_Task]) -> int: + """0-based index of the last step the agent has reached, or -1.""" + for i in range(len(tasks) - 1, -1, -1): + if tasks[i].status in ("done", "in_progress", "failed"): + return i + return -1 + + +async def insert_steps( + session_id: str, + steps: list[str], + *, + after: int | None = None, + source: str = "protocol", + protocol_run_id: str = "", + step_ids: list[str] | None = None, +) -> list[int]: + """Splice steps into the existing plan and return their 1-based indexes. + + ``after`` is the 0-based index the new block follows; ``None`` places it + after the last step the agent has already reached (done / in_progress / + failed), so the agent's own progress is never reordered — the plan simply + grows underneath its current position. + + No progress is lost: existing steps and statuses are preserved, which is + why this exists instead of asking the model to re-issue the whole plan + with 'set'. + """ + tasks = await _load_tasks(session_id) + pos = (after if after is not None else _current_position(tasks)) + 1 + pos = max(0, min(pos, len(tasks))) + new = [ + _Task( + text=str(text), + source=source, + protocol_run_id=protocol_run_id, + step_id=(step_ids[i] if step_ids and i < len(step_ids) else ""), + ) + for i, text in enumerate(steps) + ] + tasks[pos:pos] = new + await _save_tasks(session_id, tasks) + return list(range(pos + 1, pos + 1 + len(new))) diff --git a/tests/unit/store/test_todo_store.py b/tests/unit/store/test_todo_store.py new file mode 100644 index 0000000..a46e543 --- /dev/null +++ b/tests/unit/store/test_todo_store.py @@ -0,0 +1,136 @@ +"""Tests for navi.todo_store — the dedicated session_todos table.""" + +import json +from datetime import datetime, timezone + +from navi.todo_store import KvTodoStore, PgTodoStore, plan_timestamp +from tests.conftest_factory import FakeConnection, FakePool, FakeRecord + + +def _pg(conn: FakeConnection | None = None) -> tuple[PgTodoStore, FakeConnection]: + conn = conn or FakeConnection() + store = PgTodoStore(FakePool(conn)) + # Skip the DDL/migration round trip: those two execute() calls would eat + # the results a test queued for the query under test. + store._ready = True + return store, conn + + +class TestPgTodoStore: + async def test_load_returns_rows(self): + store, conn = _pg() + conn.enqueue(FakeRecord(tasks='[{"text": "a", "status": "done"}]')) + assert await store.load("u1", "s1") == [{"text": "a", "status": "done"}] + op, sql, args = conn.calls[-1] + assert op == "fetchrow" + assert "FROM session_todos" in sql + assert args == ("u1", "s1") + + async def test_load_accepts_already_decoded_json(self): + store, conn = _pg() + conn.enqueue(FakeRecord(tasks=[{"text": "a", "status": "pending"}])) + assert await store.load(None, "s1") == [{"text": "a", "status": "pending"}] + assert conn.calls[-1][2] == ("", "s1") + + async def test_load_missing_row_is_none(self): + store, conn = _pg() + assert await store.load("u1", "s1") is None + + async def test_load_garbage_row_is_none(self): + store, conn = _pg() + conn.enqueue(FakeRecord(tasks="{not json")) + assert await store.load("u1", "s1") is None + + async def test_load_meta(self): + store, conn = _pg() + now = datetime.now(timezone.utc) + conn.enqueue(FakeRecord(revision=3, updated_at=now)) + assert await store.load_meta("u1", "s1") == {"revision": 3, "updated_at": now} + assert "SELECT revision, updated_at" in conn.calls[-1][1] + + async def test_save_writes_json_and_bumps_revision(self): + store, conn = _pg() + await store.save("u1", "s1", [{"text": "a", "status": "done"}]) + op, sql, args = conn.calls[-1] + assert op == "execute" + assert "INSERT INTO session_todos" in sql + assert "revision = session_todos.revision + 1" in sql + assert args[0] == "u1" and args[1] == "s1" + assert json.loads(args[2]) == [{"text": "a", "status": "done"}] + + async def test_clear_deletes_row(self): + store, conn = _pg() + await store.clear("u1", "s1") + op, sql, args = conn.calls[-1] + assert op == "execute" + assert sql.startswith("DELETE FROM session_todos") + assert args == ("u1", "s1") + + async def test_ddl_and_migration_run_once(self): + conn = FakeConnection() + store = PgTodoStore(FakePool(conn)) + await store.load("u1", "s1") + first = [c[1] for c in conn.calls if c[0] == "execute"] + assert any("CREATE TABLE IF NOT EXISTS session_todos" in q for q in first) + assert any("FROM session_store" in q for q in first) + # The legacy row is removed after the copy, so clearing a plan later is + # not undone by the next process start. + assert any("DELETE FROM session_store" in q for q in first) + n = len(conn.calls) + await store.load("u1", "s1") + # Only the SELECT was issued the second time. + assert len(conn.calls) == n + 1 + assert conn.calls[-1][1].startswith("SELECT tasks") + + async def test_migration_failure_does_not_disable_store(self): + import asyncpg + + conn = FakeConnection() + # execute() for DDL returns the queued exception first, which is + # consumed by _DDL; queue the failure for the migration call instead + # by making the first call succeed implicitly. + conn.enqueue("OK") + conn.enqueue(asyncpg.UndefinedTableError("session_store")) + store = PgTodoStore(FakePool(conn)) + await store._ensure_ready() + assert store._ready is True + + +class TestKvTodoStore: + class _Kv: + def __init__(self): + self.data = {} + + async def get(self, uid, sid, scope, key): + return self.data.get((uid, sid, scope, key)) + + async def set(self, uid, sid, scope, key, value): + self.data[(uid, sid, scope, key)] = value + + async def clear_scope(self, uid, sid, scope): + for k in [k for k in self.data if k[:3] == (uid, sid, scope)]: + del self.data[k] + + async def test_round_trip(self): + kv = self._Kv() + store = KvTodoStore(kv) + await store.save("u1", "s1", [{"text": "a"}]) + assert await store.load("u1", "s1") == [{"text": "a"}] + await store.clear("u1", "s1") + assert await store.load("u1", "s1") is None + + async def test_no_meta(self): + assert await KvTodoStore(self._Kv()).load_meta("u1", "s1") is None + + async def test_invalid_json_is_none(self): + kv = self._Kv() + await kv.set("u1", "s1", "todo", "tasks", "{broken") + assert await KvTodoStore(kv).load("u1", "s1") is None + + +def test_plan_timestamp(): + now = datetime.now(timezone.utc) + assert plan_timestamp({"updated_at": now}) == now + assert plan_timestamp({"updated_at": "nope"}) is None + assert plan_timestamp({}) is None + assert plan_timestamp(None) is None diff --git a/tests/unit/tools/test_todo.py b/tests/unit/tools/test_todo.py index 4e7f3bb..dbad00e 100644 --- a/tests/unit/tools/test_todo.py +++ b/tests/unit/tools/test_todo.py @@ -1,5 +1,7 @@ """Unit tests for navi.tools.todo.""" +from datetime import UTC, datetime + import pytest from navi.llm.base import Message @@ -52,8 +54,11 @@ store = FakeKvStore() from navi.tools import todo as _mod _mod._kv_store = store + _mod._store = None + _mod._restart_notified.clear() yield store _mod._kv_store = None + _mod._store = None # ── TodoTool execute tests ──────────────────────────────────────────────────── @@ -615,3 +620,119 @@ assert "'op'" in result.error assert "update" in result.error assert "add" in result.error + + +# ── Dedicated store: provenance, plan surgery, crash recovery ──────────────── + +class _FakeTodoStore: + """In-memory stand-in for PgTodoStore (which needs a live pool).""" + + def __init__(self, tasks=None, meta=None): + self.tasks = list(tasks or []) + self.meta = meta + self.saved = None + + async def load(self, user_id, session_id): + return [dict(t) for t in self.tasks] + + async def load_meta(self, user_id, session_id): + return self.meta + + async def save(self, user_id, session_id, tasks): + self.tasks = [dict(t) for t in tasks] + self.saved = tasks + + async def clear(self, user_id, session_id): + self.tasks = [] + + +@pytest.mark.asyncio +async def test_task_provenance_survives_a_round_trip(_fake_kv): + from navi.tools import todo as _mod + + store = _FakeTodoStore() + _mod._store = store + ctx = ToolContext(session_id="sess1", user_id="user1") + tool = TodoTool() + await tool.execute({"op": "set", "tasks": ["step one"]}, ctx=ctx) + + assert store.saved[0]["source"] == "plan" + assert store.saved[0]["step_id"] == "" + assert store.saved[0]["protocol_run_id"] == "" + + +@pytest.mark.asyncio +async def test_insert_steps_grafts_after_the_current_position(_fake_kv): + from navi.tools import todo as _mod + + store = _FakeTodoStore(tasks=[ + {"text": "one", "status": "done", "validation": "ok"}, + {"text": "two", "status": "in_progress"}, + {"text": "three", "status": "pending"}, + ]) + _mod._store = store + + added = await _mod.insert_steps( + "sess1", ["p1", "p2"], protocol_run_id="run-1", step_ids=["s1", "s2"], + ) + + assert added == [3, 4] + assert [t["text"] for t in store.tasks] == ["one", "two", "p1", "p2", "three"] + assert store.tasks[1]["status"] == "in_progress" # nothing reordered + assert store.tasks[2]["source"] == "protocol" + assert store.tasks[2]["step_id"] == "s1" + assert store.tasks[2]["protocol_run_id"] == "run-1" + + +@pytest.mark.asyncio +async def test_insert_steps_into_an_empty_plan(_fake_kv): + from navi.tools import todo as _mod + + _mod._store = _FakeTodoStore() + assert await _mod.insert_steps("sess1", ["p1"]) == [1] + + +@pytest.mark.asyncio +async def test_restart_note_fires_once_for_a_stale_in_progress_step(_fake_kv): + from datetime import timedelta + + from navi.tools import todo as _mod + + stale = datetime.now(UTC) - timedelta(hours=2) + _mod._store = _FakeTodoStore( + tasks=[{"text": "one", "status": "in_progress"}], + meta={"revision": 4, "updated_at": stale}, + ) + + first = await get_progress_message("sess1") + second = await get_progress_message("sess1") + + assert first is not None and "Recovery:" in first.content + assert "step 1" in first.content + assert second is not None and "Recovery:" not in second.content + + +@pytest.mark.asyncio +async def test_no_restart_note_for_a_fresh_plan(_fake_kv): + from navi.tools import todo as _mod + + _mod._store = _FakeTodoStore( + tasks=[{"text": "one", "status": "in_progress"}], + meta={"revision": 1, "updated_at": datetime.now(UTC)}, + ) + msg = await get_progress_message("sess1") + assert msg is not None and "Recovery:" not in msg.content + + +@pytest.mark.asyncio +async def test_no_restart_note_without_an_in_progress_step(_fake_kv): + from datetime import timedelta + + from navi.tools import todo as _mod + + _mod._store = _FakeTodoStore( + tasks=[{"text": "one", "status": "done", "validation": "ok"}], + meta={"revision": 1, "updated_at": datetime.now(UTC) - timedelta(days=1)}, + ) + msg = await get_progress_message("sess1") + assert msg is not None and "Recovery:" not in msg.content