From 2e1a88ee52482ed80cb833d8d1af3458811c97f9 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 5 Jun 2026 17:05:30 +0800 Subject: [PATCH] feat(storage): add sqlite durable fencing Implement ADR-0068 PR3 local SQLite transactional fencing for session, tape, checkpoint, and runtime writes. Verification: - uv run pytest tests/coding_agent/test_sqlite_local_durable_fencing.py tests/ui/test_session_owner_store.py tests/ui/test_session_manager_owner_checks.py tests/ui/test_session_manager_public_api.py tests/cli/test_entrypoint_contract.py tests/cli/test_local_runtime.py tests/cli/test_repl.py tests/ui/test_http_server.py::TestSessionCreation::test_build_session_manager_enables_owner_store_for_local_sqlite_bundle tests/ui/test_http_server.py::TestSessionCreation::test_build_session_manager_enables_owner_store_for_equivalent_local_sqlite_bundle_paths tests/ui/test_http_server.py::TestSessionCreation::test_build_session_manager_does_not_enable_owner_store_for_non_bundle_storage -q - uv run ruff check ; uv run ruff format --check ; git diff --check --- .../src/coding_agent/cli/local_runtime.py | 29 +- coding-agent/src/coding_agent/cli/main.py | 3 + coding-agent/src/coding_agent/cli/repl.py | 4 +- .../src/coding_agent/local_durable_store.py | 1258 +++++++++++++++++ .../src/coding_agent/local_storage.py | 18 + .../coding_agent/runs/checkpoint_restore.py | 14 +- .../runs/runtime_checkpoint_restore.py | 3 + .../src/coding_agent/server/http_server.py | 47 + .../coding_agent/server/session_manager.py | 378 ++++- .../server/stores/session_owner_store.py | 112 ++ .../tests/cli/test_entrypoint_contract.py | 12 + coding-agent/tests/cli/test_local_runtime.py | 46 +- coding-agent/tests/cli/test_repl.py | 21 +- .../test_sqlite_local_durable_fencing.py | 737 ++++++++++ coding-agent/tests/ui/test_http_server.py | 72 +- .../ui/test_session_manager_public_api.py | 33 + 16 files changed, 2714 insertions(+), 73 deletions(-) create mode 100644 coding-agent/src/coding_agent/local_durable_store.py create mode 100644 coding-agent/tests/coding_agent/test_sqlite_local_durable_fencing.py diff --git a/coding-agent/src/coding_agent/cli/local_runtime.py b/coding-agent/src/coding_agent/cli/local_runtime.py index eda2bc05d..31ccc20bd 100644 --- a/coding-agent/src/coding_agent/cli/local_runtime.py +++ b/coding-agent/src/coding_agent/cli/local_runtime.py @@ -7,11 +7,14 @@ import importlib import logging from typing import Any, Protocol, runtime_checkable +from uuid import uuid4 from coding_agent.adapter import PipelineAdapter from coding_agent.core.config import Config from coding_agent.local_storage import local_sqlite_storage_config +from coding_agent.local_storage import local_sqlite_path_from_storage_config from coding_agent.server.session_manager import SessionManager +from coding_agent.server.stores.session_owner_store import SQLiteSessionOwnerStore logger = logging.getLogger(__name__) @@ -114,7 +117,7 @@ async def submit_approval_response( async def resume_session(self, session_id: str, **kwargs: Any) -> Any: ... - def attach_runtime( + async def attach_runtime( self, managed_session: LocalCliManagedSession, *, @@ -146,11 +149,25 @@ class ServerBackedLocalCliSessionManager: def __post_init__(self) -> None: if self.owner_renew_interval_seconds <= 0: raise ValueError("owner_renew_interval_seconds must be positive") + owner_store = self.owner_store + owner_id = self.owner_id + fencing_token = self.fencing_token + if owner_store is None: + owner_store = SQLiteSessionOwnerStore( + local_sqlite_path_from_storage_config( + dict(self.storage_config), + ) + ) + owner_id = owner_id or f"local-cli:{uuid4().hex}" + fencing_token = fencing_token or 1 + self.owner_store = owner_store + self.owner_id = owner_id + self.fencing_token = fencing_token self._delegate = SessionManager( storage_config=dict(self.storage_config), - owner_store=self.owner_store, - owner_id=self.owner_id, - fencing_token=self.fencing_token, + owner_store=owner_store, + owner_id=owner_id, + fencing_token=fencing_token, ) async def create_session(self, **kwargs: Any) -> str: @@ -269,7 +286,7 @@ async def submit_approval_response( async def resume_session(self, session_id: str, **kwargs: Any) -> Any: return await self._delegate.resume_session(session_id, **kwargs) - def attach_runtime( + async def attach_runtime( self, managed_session: LocalCliManagedSession, *, @@ -283,7 +300,7 @@ def attach_runtime( adapter=pipeline_adapter, ) managed_session.tape_id = pipeline_ctx.tape.tape_id - self._delegate._persist_session(managed_session) + await self._delegate._persist_session_async(managed_session) async def close(self) -> None: if self._owner_renew_task is not None: diff --git a/coding-agent/src/coding_agent/cli/main.py b/coding-agent/src/coding_agent/cli/main.py index fa79e4e24..8e8db71b4 100644 --- a/coding-agent/src/coding_agent/cli/main.py +++ b/coding-agent/src/coding_agent/cli/main.py @@ -476,6 +476,8 @@ async def _resume_managed_session( raise click.UsageError("Pass --session or --last.") resolved_session_id = await _latest_managed_session_id(session_manager) session = await session_manager.get_session_async(resolved_session_id) + await session_manager.acquire_session_owner(resolved_session_id) + await session_manager.start_owner_lease_renewal() consumer = HeadlessConsumer(auto_approve=config.approval_mode == "yolo") task = asyncio.create_task( session_manager.resume_session( @@ -612,6 +614,7 @@ async def _run_managed_one_shot(config: Config, goal: str, consumer) -> None: base_url=config.base_url, max_steps=config.max_steps, ) + await session_manager.start_owner_lease_renewal() session = await session_manager.get_session_async(session_id) task = asyncio.create_task(session_manager.run_agent(session_id, goal)) await _stream_managed_session_wire(session.wire, task, consumer) diff --git a/coding-agent/src/coding_agent/cli/repl.py b/coding-agent/src/coding_agent/cli/repl.py index 2db4b00a9..ed3621da8 100644 --- a/coding-agent/src/coding_agent/cli/repl.py +++ b/coding-agent/src/coding_agent/cli/repl.py @@ -175,12 +175,13 @@ async def _initialize_managed_session(self) -> None: ) self.context["session_id"] = session_id managed_session = self._session_manager.get_session(session_id) - self._session_manager.attach_runtime( + await self._session_manager.attach_runtime( managed_session, pipeline=self._pipeline, pipeline_ctx=self._pipeline_ctx, pipeline_adapter=self._pipeline_adapter, ) + await self._session_manager.start_owner_lease_renewal() self._managed_session_initialized = True async def _create_managed_session(self) -> str: @@ -386,6 +387,7 @@ async def run(self): self._footer.disable() if self._pipeline_adapter is not None: await self._pipeline_adapter.close() + await self._session_manager.close() print_pt("\nSession ended.\n", output=prompt_output) diff --git a/coding-agent/src/coding_agent/local_durable_store.py b/coding-agent/src/coding_agent/local_durable_store.py new file mode 100644 index 000000000..9c6bc07de --- /dev/null +++ b/coding-agent/src/coding_agent/local_durable_store.py @@ -0,0 +1,1258 @@ +from __future__ import annotations + +import json +import sqlite3 +import threading +from collections.abc import Callable, Mapping +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Final, cast + +from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot +from agentkit.storage.protocols import TapeInfo, TapeSearchResult +from agentkit.storage.sqlite import ( + SQLiteCheckpointStore, + SQLiteTapeStore, + _checkpoint_snapshot_from_payload, + _checkpoint_snapshot_to_payload, + _entry_anchor_type, + _entry_nested_str, + _json_object_from_text, + _optional_entry_str, +) + +from coding_agent.runtime_store import ( + AgentRunRecord, + AgentInteractionRecord, + RuntimeEventRecord, + RunMessageSnapshotRecord, + SQLiteRuntimeStore, + _agent_run_from_sqlite_row, + _agent_run_sqlite_values, + _datetime_to_json, + _interaction_from_sqlite_row, + _interaction_sqlite_values, + _json_to_sql, + _runtime_event_from_sqlite_row, +) +from coding_agent.server.stores.session_owner_store import ( + OwnerAuthority, + SQLiteSessionOwnerStore, + SessionOwnershipConflictError, + SessionOwnershipConflictReason, + _datetime_from_sqlite_text, + _datetime_to_sqlite_text, +) +from coding_agent.server.stores.session_store import SQLiteSessionStore, SessionPayload + + +class SQLiteLocalDurableStore: + """SQLite-local protected mutation facade. + + This facade is intentionally local to the SQLite bundle path. It does not replace + the generic store protocols; it provides the transaction shape required by local + durable fencing: owner epoch check, target ownership check, and mutation in the + same SQLite transaction. + """ + + _SESSION_TAPES_SCHEMA_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS session_tapes ( + session_id TEXT PRIMARY KEY, + tape_id TEXT UNIQUE NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS session_tapes_tape_id_idx + ON session_tapes (tape_id); + """ + + def __init__(self, path: Path) -> None: + self._path = path + self._lock = threading.Lock() + self._path.parent.mkdir(parents=True, exist_ok=True) + self._owner_store = SQLiteSessionOwnerStore(path) + with self._connect() as connection: + connection.executescript(SQLiteSessionStore._CREATE_TABLE_SQL) + connection.executescript(SQLiteTapeStore._CREATE_SCHEMA_SQL) + connection.executescript(SQLiteCheckpointStore._CREATE_SCHEMA_SQL) + connection.executescript(SQLiteRuntimeStore._CREATE_SCHEMA_SQL) + connection.executescript(self._SESSION_TAPES_SCHEMA_SQL) + + def _connect(self) -> sqlite3.Connection: + connection = sqlite3.connect(self._path, timeout=5.0) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout = 5000") + return connection + + async def acquire_owner( + self, + session_id: str, + owner_id: str, + lease_seconds: float = 30.0, + ) -> OwnerAuthority: + return await self._owner_store.acquire_authority( + session_id, + owner_id, + lease_seconds=lease_seconds, + ) + + async def renew_owner( + self, + authority: OwnerAuthority, + lease_seconds: float = 30.0, + ) -> OwnerAuthority: + return await self._owner_store.renew_authority( + authority, + lease_seconds=lease_seconds, + ) + + async def save_session( + self, + authority: OwnerAuthority, + payload: SessionPayload, + ) -> None: + _require_json_object("session payload", payload) + payload_id = payload.get("id") + if payload_id != authority.session_id: + raise SessionOwnershipConflictError( + "session payload belongs to another owner" + ) + payload_session_id = payload.get("session_id") + if ( + payload_session_id is not None + and payload_session_id != authority.session_id + ): + raise SessionOwnershipConflictError( + "session payload belongs to another owner" + ) + tape_id = payload.get("tape_id") + if tape_id is not None and not isinstance(tape_id, str): + raise TypeError("session payload tape_id must be a string") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + if tape_id: + self._bind_tape(connection, authority.session_id, tape_id) + connection.execute( + """ + INSERT INTO agent_http_sessions (session_id, payload, updated_at) + VALUES (?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(session_id) + DO UPDATE SET payload = excluded.payload, + updated_at = CURRENT_TIMESTAMP + """, + (authority.session_id, json.dumps(payload, sort_keys=True)), + ) + + async def delete_session(self, authority: OwnerAuthority) -> None: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + connection.execute( + "DELETE FROM agent_http_sessions WHERE session_id = ?", + (authority.session_id,), + ) + connection.execute( + "DELETE FROM session_tapes WHERE session_id = ?", + (authority.session_id,), + ) + + def load_session(self, session_id: str) -> SessionPayload | None: + _require_non_empty("session_id", session_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT payload FROM agent_http_sessions WHERE session_id = ?", + (session_id,), + ).fetchone() + if row is None: + return None + payload = json.loads(str(row["payload"])) + if not isinstance(payload, dict): + raise TypeError("sqlite session payload must be a JSON object") + return cast(SessionPayload, payload) + + async def append_tape_entries( + self, + authority: OwnerAuthority, + tape_id: str, + entries: list[dict[str, Any]], + ) -> None: + _require_non_empty("tape_id", tape_id) + if not entries: + return + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=tape_id, + session_id=authority.session_id, + ) + row = connection.execute( + "SELECT COALESCE(MAX(seq), -1) AS max_seq FROM tape_entries WHERE tape_id = ?", + (tape_id,), + ).fetchone() + max_seq = _row_required_int(row, "max_seq", context="tape max seq") + values: list[tuple[object, ...]] = [] + now = datetime.now(UTC).isoformat() + for offset, entry in enumerate(entries, start=1): + _require_json_object("tape entry", entry) + values.append( + ( + tape_id, + max_seq + offset, + json.dumps(entry, sort_keys=True), + _optional_entry_str(entry, "kind"), + _entry_nested_str(entry, "run_id"), + _entry_nested_str(entry, "tool_call_id"), + _entry_anchor_type(entry), + now, + ) + ) + connection.executemany( + """ + INSERT INTO tape_entries ( + tape_id, seq, entry_json, kind, run_id, tool_call_id, + anchor_type, created_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + values, + ) + + async def load_tape(self, tape_id: str) -> list[dict[str, Any]]: + _require_non_empty("tape_id", tape_id) + with self._lock, self._connect() as connection: + rows = connection.execute( + """ + SELECT entry_json + FROM tape_entries + WHERE tape_id = ? + ORDER BY seq + """, + (tape_id,), + ).fetchall() + return [ + _json_object_from_text(row["entry_json"], context="tape entry") + for row in rows + ] + + async def truncate_tape( + self, + authority: OwnerAuthority, + tape_id: str, + keep: int, + ) -> None: + _require_non_empty("tape_id", tape_id) + if keep < 0: + raise ValueError("keep must be >= 0") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=tape_id, + session_id=authority.session_id, + ) + connection.execute( + "DELETE FROM tape_entries WHERE tape_id = ? AND seq >= ?", + (tape_id, keep), + ) + + def session_id_for_tape(self, tape_id: str) -> str | None: + _require_non_empty("tape_id", tape_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT session_id FROM session_tapes WHERE tape_id = ?", + (tape_id,), + ).fetchone() + if row is None: + return None + session_id = row["session_id"] + if not isinstance(session_id, str): + raise TypeError("session_tapes session_id must be text") + return session_id + + async def create_agent_run( + self, + authority: OwnerAuthority, + record: AgentRunRecord, + ) -> AgentRunRecord: + if record.session_id != authority.session_id: + raise SessionOwnershipConflictError("agent run belongs to another session") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + existing = connection.execute( + "SELECT session_id FROM agent_runs WHERE run_id = ?", + (record.run_id,), + ).fetchone() + if existing is not None and existing["session_id"] != authority.session_id: + raise SessionOwnershipConflictError( + "agent run target belongs to another session" + ) + connection.execute( + """ + INSERT INTO agent_runs ( + run_id, session_id, tape_id, parent_run_id, agent_id, status, + started_at, ended_at, metadata, result, error + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(run_id) + DO UPDATE SET + tape_id = excluded.tape_id, + parent_run_id = excluded.parent_run_id, + agent_id = excluded.agent_id, + status = excluded.status, + started_at = excluded.started_at, + ended_at = excluded.ended_at, + metadata = excluded.metadata, + result = excluded.result, + error = excluded.error + """, + _agent_run_sqlite_values(record), + ) + return record + + async def update_agent_run( + self, + authority: OwnerAuthority, + run_id: str, + *, + status: str, + ended_at: datetime | None, + metadata: dict[str, Any], + result: dict[str, Any], + error: str | None, + ) -> AgentRunRecord: + _require_non_empty("run_id", run_id) + _require_non_empty("status", status) + _require_json_object("metadata", metadata) + _require_json_object("result", result) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=run_id, + session_id=authority.session_id, + ) + connection.execute( + """ + UPDATE agent_runs + SET status = ?, + ended_at = ?, + metadata = ?, + result = ?, + error = ? + WHERE run_id = ? + """, + ( + status, + None if ended_at is None else _datetime_to_json(ended_at), + _json_to_sql(cast(dict[str, Any], metadata)), + _json_to_sql(cast(dict[str, Any], result)), + error, + run_id, + ), + ) + row = connection.execute( + "SELECT * FROM agent_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + raise KeyError(f"agent run not found after update: {run_id}") + return _agent_run_from_sqlite_row(row) + + async def append_runtime_event( + self, + authority: OwnerAuthority, + record: RuntimeEventRecord, + ) -> RuntimeEventRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=record.run_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT runtime_events.* + FROM runtime_events + JOIN agent_runs ON agent_runs.run_id = runtime_events.run_id + WHERE runtime_events.event_id = ? + """, + (record.event_id,), + ).fetchone() + if existing is not None: + if existing["run_id"] != record.run_id: + raise SessionOwnershipConflictError( + "runtime event target belongs to another run" + ) + return _runtime_event_from_sqlite_row(existing) + connection.execute( + """ + INSERT INTO runtime_events ( + event_id, run_id, event_kind, payload, created_at + ) + VALUES (?, ?, ?, ?, ?) + """, + ( + record.event_id, + record.run_id, + record.event_kind, + _json_to_sql(record.payload), + _datetime_to_json(record.created_at), + ), + ) + row = connection.execute( + "SELECT * FROM runtime_events WHERE event_id = ?", + (record.event_id,), + ).fetchone() + if row is None: + raise RuntimeError("sqlite runtime event insert returned no row") + return _runtime_event_from_sqlite_row(row) + + async def replay_runtime_events( + self, + run_id: str, + *, + after_sequence: int = 0, + limit: int = 1000, + ) -> list[RuntimeEventRecord]: + _require_non_empty("run_id", run_id) + if after_sequence < 0: + raise ValueError("after_sequence must be >= 0") + if limit <= 0: + raise ValueError("limit must be positive") + with self._lock, self._connect() as connection: + rows = connection.execute( + """ + SELECT * FROM runtime_events + WHERE run_id = ? AND sequence > ? + ORDER BY sequence + LIMIT ? + """, + (run_id, after_sequence, limit), + ).fetchall() + return [_runtime_event_from_sqlite_row(row) for row in rows] + + async def save_message_snapshot( + self, + authority: OwnerAuthority, + record: RunMessageSnapshotRecord, + ) -> RunMessageSnapshotRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=record.run_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT run_message_snapshots.run_id, agent_runs.session_id + FROM run_message_snapshots + JOIN agent_runs ON agent_runs.run_id = run_message_snapshots.run_id + WHERE run_message_snapshots.snapshot_id = ? + """, + (record.snapshot_id,), + ).fetchone() + if existing is not None and ( + existing["run_id"] != record.run_id + or existing["session_id"] != authority.session_id + ): + raise SessionOwnershipConflictError( + "message snapshot target belongs to another session" + ) + connection.execute( + """ + INSERT INTO run_message_snapshots ( + snapshot_id, run_id, messages, metadata, created_at + ) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(snapshot_id) + DO UPDATE SET + messages = excluded.messages, + metadata = excluded.metadata, + created_at = excluded.created_at + """, + ( + record.snapshot_id, + record.run_id, + _json_to_sql(cast(list[Any], record.messages)), + _json_to_sql(record.metadata), + _datetime_to_json(record.created_at), + ), + ) + return record + + async def create_agent_interaction( + self, + authority: OwnerAuthority, + record: AgentInteractionRecord, + ) -> AgentInteractionRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=record.run_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT agent_interactions.run_id, agent_runs.session_id + FROM agent_interactions + JOIN agent_runs ON agent_runs.run_id = agent_interactions.run_id + WHERE agent_interactions.interaction_id = ? + """, + (record.interaction_id,), + ).fetchone() + if existing is not None: + if ( + existing["run_id"] != record.run_id + or existing["session_id"] != authority.session_id + ): + raise SessionOwnershipConflictError( + "interaction target belongs to another session" + ) + row = connection.execute( + "SELECT * FROM agent_interactions WHERE interaction_id = ?", + (record.interaction_id,), + ).fetchone() + if row is None: + raise RuntimeError( + "sqlite agent interaction lookup returned no row" + ) + return _interaction_from_sqlite_row(row) + connection.execute( + """ + INSERT INTO agent_interactions ( + interaction_id, run_id, interaction_kind, status, + request_payload, response_payload, metadata, + created_at, resolved_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + _interaction_sqlite_values(record), + ) + row = connection.execute( + "SELECT * FROM agent_interactions WHERE interaction_id = ?", + (record.interaction_id,), + ).fetchone() + if row is None: + raise RuntimeError("sqlite agent interaction insert returned no row") + return _interaction_from_sqlite_row(row) + + async def resolve_agent_interaction( + self, + authority: OwnerAuthority, + interaction_id: str, + *, + status: str, + response_payload: dict[str, Any], + resolved_at: datetime, + ) -> AgentInteractionRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + existing = connection.execute( + """ + SELECT agent_interactions.* + FROM agent_interactions + JOIN agent_runs ON agent_runs.run_id = agent_interactions.run_id + WHERE agent_interactions.interaction_id = ? + AND agent_runs.session_id = ? + """, + (interaction_id, authority.session_id), + ).fetchone() + if existing is None: + raise SessionOwnershipConflictError( + "interaction target belongs to another session" + ) + if existing["resolved_at"] is None: + connection.execute( + """ + UPDATE agent_interactions + SET status = ?, + response_payload = ?, + resolved_at = ? + WHERE interaction_id = ? + """, + ( + status, + _json_to_sql(response_payload), + _datetime_to_json(resolved_at), + interaction_id, + ), + ) + row = connection.execute( + "SELECT * FROM agent_interactions WHERE interaction_id = ?", + (interaction_id,), + ).fetchone() + if row is None: + raise RuntimeError("sqlite agent interaction resolve returned no row") + return _interaction_from_sqlite_row(row) + + async def claim_attached_executor_run( + self, + authorities: Mapping[str, OwnerAuthority], + *, + session_id: str | None, + executor_kind: str, + claim_metadata: dict[str, Any], + ) -> AgentRunRecord | None: + _require_non_empty("executor_kind", executor_kind) + _require_json_object("claim_metadata", claim_metadata) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + filter_sql = "" + values: list[object] = [] + if session_id is not None: + filter_sql = "AND session_id = ?" + values.append(session_id) + rows = connection.execute( + f""" + SELECT * FROM agent_runs + WHERE status IN ('requested', 'expired') + {filter_sql} + ORDER BY started_at, run_id + """, + values, + ).fetchall() + candidates = [_agent_run_from_sqlite_row(row) for row in rows] + candidates = [ + run + for run in candidates + if run.metadata.get("executor_ref_kind") + in {"external_worker", "local_attached"} + and run.metadata.get("executor_kind") == executor_kind + ] + selected: AgentRunRecord | None = None + authority: OwnerAuthority | None = None + for candidate in candidates: + candidate_authority = authorities.get(candidate.session_id) + if candidate_authority is None: + continue + try: + self._assert_authority(connection, candidate_authority) + except SessionOwnershipConflictError: + continue + selected = candidate + authority = candidate_authority + break + if selected is None or authority is None: + return None + metadata = cast(dict[str, Any], {**selected.metadata, **claim_metadata}) + connection.execute( + """ + UPDATE agent_runs + SET status = 'claimed', + metadata = ? + WHERE run_id = ? + AND session_id = ? + """, + ( + _json_to_sql(metadata), + selected.run_id, + authority.session_id, + ), + ) + row = connection.execute( + "SELECT * FROM agent_runs WHERE run_id = ?", + (selected.run_id,), + ).fetchone() + if row is None: + raise KeyError(f"agent run not found after claim: {selected.run_id}") + return _agent_run_from_sqlite_row(row) + + def session_id_for_run(self, run_id: str) -> str | None: + _require_non_empty("run_id", run_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT session_id FROM agent_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + return None + session_id = row["session_id"] + if not isinstance(session_id, str): + raise TypeError("agent run session_id must be text") + return session_id + + async def save_checkpoint( + self, + authority: OwnerAuthority, + snapshot: CheckpointSnapshot, + ) -> None: + meta = snapshot.meta + if meta.session_id != authority.session_id: + raise SessionOwnershipConflictError("checkpoint belongs to another session") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=meta.tape_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT session_id, tape_id + FROM checkpoints + WHERE checkpoint_id = ? + """, + (meta.checkpoint_id,), + ).fetchone() + if existing is not None and ( + existing["session_id"] != authority.session_id + or existing["tape_id"] != meta.tape_id + ): + raise SessionOwnershipConflictError( + "checkpoint target belongs to another session" + ) + payload = _checkpoint_snapshot_to_payload(snapshot) + connection.execute( + """ + INSERT INTO checkpoints ( + checkpoint_id, tape_id, session_id, entry_count, window_start, + label, snapshot_json, created_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(checkpoint_id) + DO UPDATE SET + entry_count = excluded.entry_count, + window_start = excluded.window_start, + label = excluded.label, + snapshot_json = excluded.snapshot_json, + created_at = excluded.created_at + """, + ( + meta.checkpoint_id, + meta.tape_id, + meta.session_id, + meta.entry_count, + meta.window_start, + meta.label, + json.dumps(payload, sort_keys=True), + meta.created_at.isoformat(), + ), + ) + + async def load_checkpoint(self, checkpoint_id: str) -> CheckpointSnapshot | None: + _require_non_empty("checkpoint_id", checkpoint_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT snapshot_json FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + if row is None: + return None + return _checkpoint_snapshot_from_payload( + _json_object_from_text(row["snapshot_json"], context="checkpoint snapshot") + ) + + async def delete_checkpoint( + self, + authority: OwnerAuthority, + checkpoint_id: str, + ) -> None: + _require_non_empty("checkpoint_id", checkpoint_id) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + row = connection.execute( + "SELECT session_id FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + if row is None: + return + if row["session_id"] != authority.session_id: + raise SessionOwnershipConflictError( + "checkpoint target belongs to another session" + ) + connection.execute( + "DELETE FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ) + + async def restore_checkpoint_state( + self, + authority: OwnerAuthority, + snapshot: CheckpointSnapshot, + session_payload: SessionPayload, + ) -> None: + meta = snapshot.meta + if meta.session_id != authority.session_id: + raise SessionOwnershipConflictError("checkpoint belongs to another session") + if session_payload.get("id") != authority.session_id: + raise SessionOwnershipConflictError( + "session payload belongs to another owner" + ) + if session_payload.get("tape_id") != meta.tape_id: + raise SessionOwnershipConflictError( + "checkpoint restore session payload has mismatched tape id" + ) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=meta.tape_id, + session_id=authority.session_id, + ) + checkpoint_row = connection.execute( + """ + SELECT session_id, tape_id + FROM checkpoints + WHERE checkpoint_id = ? + """, + (meta.checkpoint_id,), + ).fetchone() + if checkpoint_row is None: + raise KeyError(f"checkpoint not found: {meta.checkpoint_id}") + if ( + checkpoint_row["session_id"] != authority.session_id + or checkpoint_row["tape_id"] != meta.tape_id + ): + raise SessionOwnershipConflictError( + "checkpoint target belongs to another session" + ) + connection.execute( + "DELETE FROM tape_entries WHERE tape_id = ? AND seq >= ?", + (meta.tape_id, meta.entry_count), + ) + connection.execute( + """ + INSERT INTO agent_http_sessions (session_id, payload, updated_at) + VALUES (?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(session_id) + DO UPDATE SET payload = excluded.payload, + updated_at = CURRENT_TIMESTAMP + """, + (authority.session_id, json.dumps(session_payload, sort_keys=True)), + ) + connection.execute( + """ + DELETE FROM checkpoints + WHERE tape_id = ? + AND session_id = ? + AND entry_count > ? + """, + (meta.tape_id, authority.session_id, meta.entry_count), + ) + + def session_id_for_checkpoint(self, checkpoint_id: str) -> str | None: + _require_non_empty("checkpoint_id", checkpoint_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT session_id FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + if row is None: + return None + session_id = row["session_id"] + if not isinstance(session_id, str): + raise TypeError("checkpoint session_id must be text") + return session_id + + def _assert_authority( + self, + connection: sqlite3.Connection, + authority: OwnerAuthority, + ) -> None: + row = connection.execute( + """ + SELECT owner_id, lease_expires_at, fencing_token + FROM session_owners + WHERE session_id = ? + """, + (authority.session_id,), + ).fetchone() + if row is None: + raise SessionOwnershipConflictError( + "session owner lease is missing", + reason=SessionOwnershipConflictReason.MISSING_OWNER, + ) + lease_expires_at = _datetime_from_sqlite_text(row["lease_expires_at"]) + if lease_expires_at <= datetime.now(UTC): + raise SessionOwnershipConflictError( + "session owner lease has expired", + reason=SessionOwnershipConflictReason.EXPIRED_LEASE, + ) + if ( + row["owner_id"] != authority.owner_id + or row["fencing_token"] != authority.epoch + ): + raise SessionOwnershipConflictError("stale owner or fencing token rejected") + + def _bind_tape( + self, + connection: sqlite3.Connection, + session_id: str, + tape_id: str, + ) -> None: + now = _datetime_to_sqlite_text(datetime.now(UTC)) + session_row = connection.execute( + "SELECT tape_id FROM session_tapes WHERE session_id = ?", + (session_id,), + ).fetchone() + if session_row is not None and session_row["tape_id"] != tape_id: + raise SessionOwnershipConflictError("session tape target cannot be rebound") + tape_row = connection.execute( + "SELECT session_id FROM session_tapes WHERE tape_id = ?", + (tape_id,), + ).fetchone() + if tape_row is not None and tape_row["session_id"] != session_id: + raise SessionOwnershipConflictError( + "tape target belongs to another session" + ) + connection.execute( + """ + INSERT INTO session_tapes (session_id, tape_id, created_at, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(session_id) + DO UPDATE SET + tape_id = excluded.tape_id, + updated_at = excluded.updated_at + """, + (session_id, tape_id, now, now), + ) + + def _assert_tape_belongs_to_session( + self, + connection: sqlite3.Connection, + *, + tape_id: str, + session_id: str, + ) -> None: + row = connection.execute( + """ + SELECT session_id + FROM session_tapes + WHERE tape_id = ? + """, + (tape_id,), + ).fetchone() + if row is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + if row["session_id"] != session_id: + raise SessionOwnershipConflictError( + "tape target belongs to another session" + ) + + def _assert_run_belongs_to_session( + self, + connection: sqlite3.Connection, + *, + run_id: str, + session_id: str, + ) -> None: + row = connection.execute( + "SELECT session_id FROM agent_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + raise KeyError(f"agent run not found: {run_id}") + if row["session_id"] != session_id: + raise SessionOwnershipConflictError("agent run belongs to another session") + + +def _require_non_empty(field_name: str, value: str) -> None: + if not isinstance(value, str) or not value: + raise ValueError(f"{field_name} must be non-empty") + + +def _require_json_object(field_name: str, value: object) -> None: + if not isinstance(value, dict): + raise TypeError(f"{field_name} must be a JSON object") + + +def _row_required_int(row: sqlite3.Row | None, key: str, *, context: str) -> int: + if row is None: + raise TypeError(f"{context} row is missing") + value = row[key] + if not isinstance(value, int): + raise TypeError(f"{context} {key} must be an int") + return value + + +class FencedSQLiteTapeStore: + def __init__( + self, + *, + durable_store: SQLiteLocalDurableStore, + path: Path, + authority_for_session: Callable[[str], OwnerAuthority], + ) -> None: + self._durable_store = durable_store + self._delegate = SQLiteTapeStore(path) + self._authority_for_session = authority_for_session + + async def save(self, tape_id: str, entries: list[dict[str, Any]]) -> None: + session_id = self._require_session_id_for_tape(tape_id) + await self._durable_store.append_tape_entries( + self._authority_for_session(session_id), + tape_id, + entries, + ) + + async def load(self, tape_id: str) -> list[dict[str, Any]]: + return await self._delegate.load(tape_id) + + async def list_ids(self) -> list[str]: + return await self._delegate.list_ids() + + async def truncate(self, tape_id: str, keep: int) -> None: + session_id = self._require_session_id_for_tape(tape_id) + await self._durable_store.truncate_tape( + self._authority_for_session(session_id), + tape_id, + keep, + ) + + async def info(self, tape_id: str) -> TapeInfo | None: + return await self._delegate.info(tape_id) + + async def search( + self, + *, + tape_id: str | None = None, + kind: str | None = None, + run_id: str | None = None, + tool_call_id: str | None = None, + anchor_type: str | None = None, + limit: int = 100, + ) -> list[TapeSearchResult]: + return await self._delegate.search( + tape_id=tape_id, + kind=kind, + run_id=run_id, + tool_call_id=tool_call_id, + anchor_type=anchor_type, + limit=limit, + ) + + def _require_session_id_for_tape(self, tape_id: str) -> str: + session_id = self._durable_store.session_id_for_tape(tape_id) + if session_id is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + return session_id + + +class FencedSQLiteCheckpointStore: + def __init__( + self, + *, + durable_store: SQLiteLocalDurableStore, + path: Path, + authority_for_session: Callable[[str], OwnerAuthority], + ) -> None: + self._durable_store = durable_store + self._delegate = SQLiteCheckpointStore(path) + self._authority_for_session = authority_for_session + + async def save(self, snapshot: CheckpointSnapshot) -> None: + session_id = snapshot.meta.session_id + if session_id is None: + raise SessionOwnershipConflictError( + "checkpoint target is not bound to a session" + ) + await self._durable_store.save_checkpoint( + self._authority_for_session(session_id), + snapshot, + ) + + async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: + return await self._delegate.load(checkpoint_id) + + async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: + return await self._delegate.list_by_tape(tape_id) + + async def delete(self, checkpoint_id: str) -> None: + session_id = self._durable_store.session_id_for_checkpoint(checkpoint_id) + if session_id is None: + return + await self._durable_store.delete_checkpoint( + self._authority_for_session(session_id), + checkpoint_id, + ) + + +class FencedSQLiteRuntimeStore: + def __init__( + self, + *, + durable_store: SQLiteLocalDurableStore, + path: Path, + authority_for_session: Callable[[str], OwnerAuthority], + authorities: Callable[[], Mapping[str, OwnerAuthority]], + ) -> None: + self._durable_store = durable_store + self._delegate = SQLiteRuntimeStore(path) + self._authority_for_session = authority_for_session + self._authorities = authorities + + async def create_agent_run(self, record: AgentRunRecord) -> AgentRunRecord: + return await self._durable_store.create_agent_run( + self._authority_for_session(record.session_id), + record, + ) + + async def update_agent_run( + self, + run_id: str, + *, + status: str, + ended_at: datetime | None, + metadata: dict[str, Any], + result: dict[str, Any], + error: str | None, + ) -> AgentRunRecord: + session_id = self._require_session_id_for_run(run_id) + return await self._durable_store.update_agent_run( + self._authority_for_session(session_id), + run_id, + status=status, + ended_at=ended_at, + metadata=metadata, + result=result, + error=error, + ) + + async def load_agent_run(self, run_id: str) -> AgentRunRecord | None: + return await self._delegate.load_agent_run(run_id) + + async def list_agent_runs(self, session_id: str) -> list[AgentRunRecord]: + return await self._delegate.list_agent_runs(session_id) + + async def claim_attached_executor_run( + self, + *, + session_id: str | None, + executor_kind: str, + claim_metadata: dict[str, Any], + ) -> AgentRunRecord | None: + return await self._durable_store.claim_attached_executor_run( + self._authorities(), + session_id=session_id, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + + async def claim_external_worker_run( + self, + *, + session_id: str | None, + executor_kind: str, + claim_metadata: dict[str, Any], + ) -> AgentRunRecord | None: + return await self._durable_store.claim_attached_executor_run( + self._authorities(), + session_id=session_id, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + + async def append_runtime_event( + self, + record: RuntimeEventRecord, + ) -> RuntimeEventRecord: + session_id = self._require_session_id_for_run(record.run_id) + return await self._durable_store.append_runtime_event( + self._authority_for_session(session_id), + record, + ) + + async def load_runtime_event( + self, + event_id: str, + ) -> RuntimeEventRecord | None: + return await self._delegate.load_runtime_event(event_id) + + async def replay_runtime_events( + self, + run_id: str, + *, + after_sequence: int = 0, + limit: int = 1000, + ) -> list[RuntimeEventRecord]: + return await self._delegate.replay_runtime_events( + run_id, + after_sequence=after_sequence, + limit=limit, + ) + + async def save_message_snapshot( + self, + record: RunMessageSnapshotRecord, + ) -> RunMessageSnapshotRecord: + session_id = self._require_session_id_for_run(record.run_id) + return await self._durable_store.save_message_snapshot( + self._authority_for_session(session_id), + record, + ) + + async def load_message_snapshot( + self, + snapshot_id: str, + ) -> RunMessageSnapshotRecord | None: + return await self._delegate.load_message_snapshot(snapshot_id) + + async def list_message_snapshots( + self, + run_id: str, + ) -> list[RunMessageSnapshotRecord]: + return await self._delegate.list_message_snapshots(run_id) + + async def create_agent_interaction( + self, + record: AgentInteractionRecord, + ) -> AgentInteractionRecord: + session_id = self._require_session_id_for_run(record.run_id) + return await self._durable_store.create_agent_interaction( + self._authority_for_session(session_id), + record, + ) + + async def load_agent_interaction( + self, + interaction_id: str, + ) -> AgentInteractionRecord | None: + return await self._delegate.load_agent_interaction(interaction_id) + + async def list_agent_interactions( + self, + run_id: str, + ) -> list[AgentInteractionRecord]: + return await self._delegate.list_agent_interactions(run_id) + + async def resolve_agent_interaction( + self, + interaction_id: str, + *, + status: str, + response_payload: dict[str, Any], + resolved_at: datetime, + ) -> AgentInteractionRecord: + interaction = await self._delegate.load_agent_interaction(interaction_id) + if interaction is None: + raise KeyError(f"agent interaction not found: {interaction_id}") + session_id = self._require_session_id_for_run(interaction.run_id) + return await self._durable_store.resolve_agent_interaction( + self._authority_for_session(session_id), + interaction_id, + status=status, + response_payload=response_payload, + resolved_at=resolved_at, + ) + + def _require_session_id_for_run(self, run_id: str) -> str: + session_id = self._durable_store.session_id_for_run(run_id) + if session_id is None: + raise KeyError(f"agent run not found: {run_id}") + return session_id diff --git a/coding-agent/src/coding_agent/local_storage.py b/coding-agent/src/coding_agent/local_storage.py index 969ff5527..b41bd6cbc 100644 --- a/coding-agent/src/coding_agent/local_storage.py +++ b/coding-agent/src/coding_agent/local_storage.py @@ -8,6 +8,10 @@ LOCAL_SQLITE_FILENAME = "local.sqlite3" +def normalize_storage_path(path: str) -> Path: + return Path(os.path.abspath(os.path.expanduser(path))) + + def local_data_dir() -> Path: return Path(os.environ.get("AGENT_DATA_DIR", "./data")) @@ -42,6 +46,20 @@ def local_sqlite_path_from_storage_config( local_path = paths.get("local") if isinstance(local_path, str) and local_path.strip(): return Path(local_path) + explicit_paths = [ + storage_config.get(path_key) + for path_key in ( + "http_session_path", + "tape_path", + "checkpoint_path", + "runtime_path", + ) + ] + if all(isinstance(path, str) and path.strip() for path in explicit_paths): + resolved_paths = [normalize_storage_path(str(path)) for path in explicit_paths] + first = resolved_paths[0] + if all(path == first for path in resolved_paths): + return first return local_sqlite_path(data_dir) diff --git a/coding-agent/src/coding_agent/runs/checkpoint_restore.py b/coding-agent/src/coding_agent/runs/checkpoint_restore.py index 4c713360c..bbc055eba 100644 --- a/coding-agent/src/coding_agent/runs/checkpoint_restore.py +++ b/coding-agent/src/coding_agent/runs/checkpoint_restore.py @@ -5,6 +5,7 @@ from typing import Any, Protocol from agentkit.checkpoint import CheckpointService +from agentkit.checkpoint.models import CheckpointSnapshot from agentkit.storage.protocols import TapeStore from agentkit.tape.models import Entry from agentkit.tape.tape import Tape @@ -55,6 +56,10 @@ def attach_runtime_binding( ] CloseCheckpointRuntime = Callable[[CheckpointRestoreSession], Awaitable[None]] PersistCheckpointSession = Callable[[CheckpointRestoreSession], Awaitable[None]] +RestoreDurableCheckpointState = Callable[ + [CheckpointRestoreSession, CheckpointSnapshot], + Awaitable[None], +] def serialize_checkpoint_session_config( @@ -133,6 +138,7 @@ class CheckpointRestoreService: prepare_runtime: PrepareCheckpointRuntime close_runtime: CloseCheckpointRuntime persist_session: PersistCheckpointSession + restore_durable_state: RestoreDurableCheckpointState | None = None async def restore( self, @@ -173,7 +179,6 @@ async def restore( previous_base_url = session.base_url await self.close_runtime(session) - await self.tape_store.truncate(session.tape_id, meta.entry_count) session.tape_id = runtime.ctx.tape.tape_id session.provider_name = restored_config.provider_name session.model_name = restored_config.model_name @@ -191,8 +196,12 @@ async def restore( ctx=runtime.ctx, adapter=runtime.adapter, ) - await self.persist_session(session) + if self.restore_durable_state is not None: + await self.restore_durable_state(session, snapshot) + return + await self.tape_store.truncate(session.tape_id, meta.entry_count) + await self.persist_session(session) checkpoints = await self.checkpoint_service.list(runtime.ctx.tape.tape_id) for checkpoint_meta in checkpoints: if checkpoint_meta.entry_count > meta.entry_count: @@ -205,6 +214,7 @@ async def restore( "CheckpointRestoreSession", "CheckpointRestoredRuntime", "CheckpointSessionConfig", + "RestoreDurableCheckpointState", "checkpoint_session_config_from_extra", "serialize_checkpoint_session_config", ] diff --git a/coding-agent/src/coding_agent/runs/runtime_checkpoint_restore.py b/coding-agent/src/coding_agent/runs/runtime_checkpoint_restore.py index 080356514..97b62a661 100644 --- a/coding-agent/src/coding_agent/runs/runtime_checkpoint_restore.py +++ b/coding-agent/src/coding_agent/runs/runtime_checkpoint_restore.py @@ -16,6 +16,7 @@ CheckpointSessionConfig, CloseCheckpointRuntime, PersistCheckpointSession, + RestoreDurableCheckpointState, ) from coding_agent.runs.checkpoint_runtime import ( CheckpointRuntimeAdapterFactory, @@ -72,6 +73,7 @@ class RuntimeCheckpointRestoreService: runtime_preparation_request: RuntimePreparationRequestFactory close_runtime: CloseCheckpointRuntime persist_session: PersistCheckpointSession + restore_durable_state: RestoreDurableCheckpointState | None = None async def restore( self, @@ -84,6 +86,7 @@ async def restore( prepare_runtime=self._prepare_runtime, close_runtime=self.close_runtime, persist_session=self.persist_session, + restore_durable_state=self.restore_durable_state, ).restore(session, checkpoint_id) async def _prepare_runtime( diff --git a/coding-agent/src/coding_agent/server/http_server.py b/coding-agent/src/coding_agent/server/http_server.py index c017450b7..c385547ed 100644 --- a/coding-agent/src/coding_agent/server/http_server.py +++ b/coding-agent/src/coding_agent/server/http_server.py @@ -87,6 +87,8 @@ RunMessageSnapshotRecord, RuntimeEventRecord, ) +from coding_agent.local_storage import local_sqlite_path_from_storage_config +from coding_agent.local_storage import normalize_storage_path from coding_agent.events import DisplayEvent, DisplayEventStreamProjector from coding_agent.scheduled_runs import ( PGScheduledRunStore, @@ -242,6 +244,7 @@ ) from coding_agent.server.session_manager import Session, SessionManager from coding_agent.server.stores.session_owner_store import ( + SQLiteSessionOwnerStore, SessionOwnershipConflictError, SessionOwnershipConflictReason, SessionOwnerStore, @@ -722,6 +725,33 @@ def _storage_uses_pg_http_sessions(storage_config: dict[str, Any]) -> bool: return tape_backend == "pg" +def _storage_uses_local_sqlite_bundle(storage_config: dict[str, Any]) -> bool: + backend_keys = ( + "http_session_backend", + "tape_backend", + "checkpoint_backend", + "runtime_backend", + ) + if any( + str(storage_config.get(key, "")).strip().lower() != "sqlite" + for key in backend_keys + ): + return False + local_path = normalize_storage_path( + str(local_sqlite_path_from_storage_config(storage_config)) + ) + path_keys = ( + "http_session_path", + "tape_path", + "checkpoint_path", + "runtime_path", + ) + return all( + normalize_storage_path(str(storage_config.get(key, ""))) == local_path + for key in path_keys + ) + + def _configured_owner_id(storage_config: dict[str, Any]) -> str: owner_id = storage_config.get("owner_id") if isinstance(owner_id, str) and owner_id.strip(): @@ -760,8 +790,25 @@ def _build_session_manager() -> SessionManager: except Exception: logger.exception("Production config validation failed") raise + uses_local_sqlite_bundle = _storage_uses_local_sqlite_bundle(storage_config) manager = SessionManager( storage_config=storage_config, + owner_store=( + SQLiteSessionOwnerStore( + local_sqlite_path_from_storage_config(storage_config) + ) + if uses_local_sqlite_bundle + else None + ), + owner_id=( + _configured_owner_id(storage_config) if uses_local_sqlite_bundle else None + ), + fencing_token=1 if uses_local_sqlite_bundle else None, + owner_lease_seconds=( + _configured_owner_lease_seconds(storage_config) + if uses_local_sqlite_bundle + else 30.0 + ), cloud_workspace_client_factory=( cloud_client_factory_from_config(_load_cloud_workspace_config()) if _load_cloud_workspace_config().get("enabled") is True diff --git a/coding-agent/src/coding_agent/server/session_manager.py b/coding-agent/src/coding_agent/server/session_manager.py index 924ba5c17..ea9e8b8bc 100644 --- a/coding-agent/src/coding_agent/server/session_manager.py +++ b/coding-agent/src/coding_agent/server/session_manager.py @@ -50,8 +50,16 @@ from coding_agent.core import config as core_config from coding_agent.local_storage import ( local_sqlite_storage_config, + local_sqlite_path_from_storage_config, + normalize_storage_path, with_local_sqlite_bundle_paths, ) +from coding_agent.local_durable_store import ( + FencedSQLiteCheckpointStore, + FencedSQLiteRuntimeStore, + FencedSQLiteTapeStore, + SQLiteLocalDurableStore, +) from coding_agent.plugins.storage import JSONLTapeStore from coding_agent.providers.base import ToolSchema from coding_agent.runtime_store import ( @@ -127,6 +135,7 @@ SessionStore, create_session_store, ) +from coding_agent.server.stores.session_owner_store import OwnerAuthority from coding_agent.server.stores.session_owner_store import SessionOwnerStoreProtocol from coding_agent.server.stores.session_owner_store import SessionOwnershipConflictError from coding_agent.server.stores.session_owner_store import ( @@ -804,6 +813,15 @@ def __init__( ) self._pg_pool = pg_pool self._owns_pg_pool = False + self._local_durable_store = self._create_local_durable_store( + data_dir=data_dir, + store=store, + tape_store=tape_store, + checkpoint_store=checkpoint_store, + checkpoint_service=checkpoint_service, + runtime_store=runtime_store, + owner_store=owner_store, + ) self._store = store or self._create_http_session_store() self._session_cache: dict[str, Session] = {} self._approval_stores: dict[str, ApprovalStore] = {} @@ -812,6 +830,15 @@ def __init__( self._session_turn_locks: dict[str, asyncio.Lock] = {} self._session_workspace_export_counts: dict[str, int] = {} self._tape_store = tape_store or self._create_tape_store(data_dir) + if self._local_durable_store is not None and tape_store is None: + self._tape_store = FencedSQLiteTapeStore( + durable_store=self._local_durable_store, + path=local_sqlite_path_from_storage_config( + self._storage_config, + data_dir, + ), + authority_for_session=self._owner_authority_for_session, + ) self._agent_observation_store = observation_store or JsonlAgentObservationStore( data_dir / "observability" ) @@ -822,6 +849,19 @@ def __init__( resolved_checkpoint_store = checkpoint_store or self._create_checkpoint_store( data_dir ) + if ( + self._local_durable_store is not None + and checkpoint_store is None + and checkpoint_service is None + ): + resolved_checkpoint_store = FencedSQLiteCheckpointStore( + durable_store=self._local_durable_store, + path=local_sqlite_path_from_storage_config( + self._storage_config, + data_dir, + ), + authority_for_session=self._owner_authority_for_session, + ) self._checkpoint_service = checkpoint_service or CheckpointService( resolved_checkpoint_store ) @@ -841,6 +881,16 @@ def __init__( self._runtime_store = ( runtime_store if runtime_store is not None else self._create_runtime_store() ) + if self._local_durable_store is not None and runtime_store is None: + self._runtime_store = FencedSQLiteRuntimeStore( + durable_store=self._local_durable_store, + path=local_sqlite_path_from_storage_config( + self._storage_config, + data_dir, + ), + authority_for_session=self._owner_authority_for_session, + authorities=lambda: dict(self._owner_authorities), + ) async def runtime_session_is_recoverable(session_id: str) -> bool: if self._owner_store is None: @@ -967,6 +1017,11 @@ async def runtime_session_is_recoverable(session_id: str) -> bool: ), close_runtime=self._runtime_closer.close, persist_session=self._persist_session_async, + restore_durable_state=( + self._restore_checkpoint_durable_state + if self._local_durable_store is not None + else None + ), ) self._runtime_checkpoint_restore_orchestration = ( RuntimeCheckpointRestoreOrchestrationService( @@ -1118,6 +1173,60 @@ def configure_owner_leases( self._owner_id = owner_id self._fencing_token = fencing_token self._owner_lease_seconds = owner_lease_seconds + self._owner_authorities: dict[str, OwnerAuthority] = {} + + def _create_local_durable_store( + self, + *, + data_dir: Path, + store: SessionStore | None, + tape_store: TapeStore | None, + checkpoint_store: CheckpointStore | None, + checkpoint_service: CheckpointService | None, + runtime_store: RuntimeStore | None, + owner_store: SessionOwnerStoreProtocol | None, + ) -> SQLiteLocalDurableStore | None: + if owner_store is None or not callable( + getattr(owner_store, "acquire_authority", None) + ): + return None + if any( + value is not None + for value in ( + store, + tape_store, + checkpoint_store, + checkpoint_service, + runtime_store, + ) + ): + return None + config = self._storage_config + sqlite_backend_keys = ( + "http_session_backend", + "tape_backend", + "checkpoint_backend", + "runtime_backend", + ) + if any( + str(config.get(key, "")).lower() != "sqlite" for key in sqlite_backend_keys + ): + return None + local_path = normalize_storage_path( + str(local_sqlite_path_from_storage_config(config, data_dir)) + ) + path_keys = ( + "http_session_path", + "tape_path", + "checkpoint_path", + "runtime_path", + ) + configured_paths = [ + normalize_storage_path(str(config.get(key, ""))) for key in path_keys + ] + if any(path != local_path for path in configured_paths): + return None + return SQLiteLocalDurableStore(local_path) def configure_workspace_metadata_store( self, @@ -1824,7 +1933,7 @@ async def prepare_session_turn(self, session_id: str) -> Session: async def _assert_owner(self, session_id: str) -> None: if self._owner_store is None: return - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") owner = await self._owner_store.get_owner(session_id) @@ -1841,6 +1950,19 @@ async def _assert_owner(self, session_id: str) -> None: current_owner_id = owner.owner_id current_fencing_token = owner.fencing_token + authority = self._owner_authorities.get(session_id) + if authority is not None: + if ( + current_owner_id != authority.owner_id + or current_fencing_token != authority.epoch + ): + raise SessionOwnershipConflictError( + "stale owner or fencing token rejected" + ) + return + + if self._fencing_token is None: + raise SessionOwnershipConflictError("stale owner or fencing token rejected") if ( current_owner_id != self._owner_id @@ -1848,6 +1970,14 @@ async def _assert_owner(self, session_id: str) -> None: ): raise SessionOwnershipConflictError("stale owner or fencing token rejected") + def _owner_authority_for_session(self, session_id: str) -> OwnerAuthority: + authority = self._owner_authorities.get(session_id) + if authority is None: + raise SessionOwnershipConflictError( + "durable mutation requires owner authority" + ) + return authority + async def authorize_event_stream(self, session_id: str) -> None: await self._assert_owner(session_id) @@ -1864,10 +1994,19 @@ def run_guarded() -> T: async def _persist_session_async(self, session: Session) -> None: self._session_cache[session.id] = session + payload = cast(dict[str, Any], session.to_store_data()) + if self._local_durable_store is not None: + authority = self._owner_authorities.get(session.id) + if authority is None: + raise SessionOwnershipConflictError( + "session metadata mutation requires owner authority" + ) + await self._local_durable_store.save_session(authority, payload) + return await self._run_store_io( self._store.save, session.id, - cast(dict[str, Any], session.to_store_data()), + payload, ) async def _persist_workspace_record_for_session(self, session: Session) -> None: @@ -1925,7 +2064,20 @@ async def _persist_workspace_record_for_session(self, session: Session) -> None: async def _acquire_owner_for_session(self, session_id: str) -> None: if self._owner_store is None: return - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: + raise SessionOwnershipConflictError("stale owner or fencing token rejected") + acquire_authority = getattr(self._owner_store, "acquire_authority", None) + if callable(acquire_authority): + authority = await acquire_authority( + session_id, + self._owner_id, + lease_seconds=self._owner_lease_seconds, + ) + if not isinstance(authority, OwnerAuthority): + raise TypeError("acquire_authority must return OwnerAuthority") + self._owner_authorities[session_id] = authority + return + if self._fencing_token is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") acquired = await self._owner_store.acquire( session_id, @@ -1939,11 +2091,19 @@ async def _acquire_owner_for_session(self, session_id: str) -> None: async def _holds_owner_lease(self, session_id: str) -> bool: if self._owner_store is None: return False - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") owner = await self._owner_store.get_owner(session_id) if owner is None: return False + authority = self._owner_authorities.get(session_id) + if authority is not None: + return ( + owner.owner_id == authority.owner_id + and owner.fencing_token == authority.epoch + ) + if self._fencing_token is None: + raise SessionOwnershipConflictError("stale owner or fencing token rejected") return ( owner.owner_id == self._owner_id and owner.fencing_token == self._fencing_token @@ -1952,11 +2112,20 @@ async def _holds_owner_lease(self, session_id: str) -> bool: async def _holds_active_owner_lease(self, session_id: str) -> bool: if self._owner_store is None: return False - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") owner = await self._owner_store.get_owner(session_id) if owner is None: return False + authority = self._owner_authorities.get(session_id) + if authority is not None: + return ( + owner.owner_id == authority.owner_id + and owner.fencing_token == authority.epoch + and owner.lease_expires_at > datetime.now(UTC) + ) + if self._fencing_token is None: + raise SessionOwnershipConflictError("stale owner or fencing token rejected") return ( owner.owner_id == self._owner_id and owner.fencing_token == self._fencing_token @@ -1966,7 +2135,7 @@ async def _holds_active_owner_lease(self, session_id: str) -> bool: async def release_owned_sessions(self) -> None: if self._owner_store is None: return - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") for session_id in await self.list_sessions_async(): await self._release_owner_lease_for_session(session_id) @@ -1974,22 +2143,28 @@ async def release_owned_sessions(self) -> None: async def _release_owner_lease_for_session(self, session_id: str) -> None: if self._owner_store is None: return - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") if not await self._holds_owner_lease(session_id): return + authority = self._owner_authorities.get(session_id) + fencing_token = ( + authority.epoch if authority is not None else self._fencing_token + ) + if fencing_token is None: + raise SessionOwnershipConflictError("stale owner or fencing token rejected") try: released = await self._owner_store.release( session_id, self._owner_id, - self._fencing_token, + fencing_token, ) except Exception: logger.warning( "Failed to release owner lease for session %s owned by %s with fencing token %s", session_id, self._owner_id, - self._fencing_token, + fencing_token, exc_info=True, ) return @@ -1998,14 +2173,16 @@ async def _release_owner_lease_for_session(self, session_id: str) -> None: "Failed to release owner lease for session %s owned by %s with fencing token %s", session_id, self._owner_id, - self._fencing_token, + fencing_token, ) + return + self._owner_authorities.pop(session_id, None) async def renew_owner_leases(self) -> list[str]: lost_active_sessions: list[str] = [] if self._owner_store is None: return lost_active_sessions - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") now = datetime.now(UTC) for session_id in await self.list_sessions_async(): @@ -2014,28 +2191,58 @@ async def renew_owner_leases(self) -> list[str]: if await self._cancel_active_turn_after_owner_loss(session_id): lost_active_sessions.append(session_id) continue - if ( - owner.owner_id != self._owner_id - or owner.fencing_token != self._fencing_token - or owner.lease_expires_at <= now - ): + authority = self._owner_authorities.get(session_id) + if authority is not None: + owns_session = ( + owner.owner_id == authority.owner_id + and owner.fencing_token == authority.epoch + and owner.lease_expires_at > now + ) + log_token = authority.epoch + else: + if self._fencing_token is None: + raise SessionOwnershipConflictError( + "stale owner or fencing token rejected" + ) + owns_session = ( + owner.owner_id == self._owner_id + and owner.fencing_token == self._fencing_token + and owner.lease_expires_at > now + ) + log_token = self._fencing_token + if not owns_session: if await self._cancel_active_turn_after_owner_loss(session_id): lost_active_sessions.append(session_id) continue try: - renewed = await self._owner_store.renew( - session_id, - self._owner_id, - lease_seconds=self._owner_lease_seconds, - new_fencing_token=self._fencing_token, - current_fencing_token=self._fencing_token, - ) + renew_authority = getattr(self._owner_store, "renew_authority", None) + if authority is not None and callable(renew_authority): + renewed_authority = await renew_authority( + authority, + lease_seconds=self._owner_lease_seconds, + ) + if not isinstance(renewed_authority, OwnerAuthority): + raise TypeError("renew_authority must return OwnerAuthority") + self._owner_authorities[session_id] = renewed_authority + renewed = True + else: + if self._fencing_token is None: + raise SessionOwnershipConflictError( + "stale owner or fencing token rejected" + ) + renewed = await self._owner_store.renew( + session_id, + self._owner_id, + lease_seconds=self._owner_lease_seconds, + new_fencing_token=self._fencing_token, + current_fencing_token=self._fencing_token, + ) except Exception: logger.warning( "Failed to renew owner lease for session %s owned by %s with fencing token %s", session_id, self._owner_id, - self._fencing_token, + log_token, exc_info=True, ) continue @@ -2044,7 +2251,7 @@ async def renew_owner_leases(self) -> list[str]: "Failed to renew owner lease for session %s owned by %s with fencing token %s", session_id, self._owner_id, - self._fencing_token, + log_token, ) if await self._cancel_active_turn_after_owner_loss(session_id): lost_active_sessions.append(session_id) @@ -2070,7 +2277,7 @@ async def _cancel_active_turn_after_owner_loss(self, session_id: str) -> bool: async def backfill_owner_leases(self) -> None: if self._owner_store is None: return - if self._owner_id is None or self._fencing_token is None: + if self._owner_id is None: raise SessionOwnershipConflictError("stale owner or fencing token rejected") now = datetime.now(UTC) @@ -2079,12 +2286,32 @@ async def backfill_owner_leases(self) -> None: owner = await self._owner_store.get_owner(session_id) if owner is not None and owner.lease_expires_at > now: continue - acquired = await self._owner_store.acquire( - session_id, - self._owner_id, - lease_seconds=self._owner_lease_seconds, - fencing_token=self._fencing_token, + acquire_authority = getattr( + self._owner_store, "acquire_authority", None ) + if callable(acquire_authority): + authority = await acquire_authority( + session_id, + self._owner_id, + lease_seconds=self._owner_lease_seconds, + ) + if not isinstance(authority, OwnerAuthority): + raise TypeError("acquire_authority must return OwnerAuthority") + self._owner_authorities[session_id] = authority + acquired = True + log_token = authority.epoch + else: + if self._fencing_token is None: + raise SessionOwnershipConflictError( + "stale owner or fencing token rejected" + ) + acquired = await self._owner_store.acquire( + session_id, + self._owner_id, + lease_seconds=self._owner_lease_seconds, + fencing_token=self._fencing_token, + ) + log_token = self._fencing_token except Exception: logger.warning( "Failed to backfill owner lease for session %s owned by %s with fencing token %s", @@ -2099,7 +2326,7 @@ async def backfill_owner_leases(self) -> None: "Failed to backfill owner lease for session %s owned by %s with fencing token %s", session_id, self._owner_id, - self._fencing_token, + log_token, ) async def get_session_async(self, session_id: str) -> Session: @@ -2183,11 +2410,41 @@ async def _remove_session_async_no_lock(self, session_id: str) -> None: await self._runtime_closer.close(session) await self._finalize_provisioned_cloud_workspace_on_close(session) self._session_cache.pop(session_id, None) - await self._run_store_io(self._store.delete, session_id) + if self._local_durable_store is not None: + authority = self._owner_authority_for_session(session_id) + await self._local_durable_store.delete_session(authority) + else: + await self._run_store_io(self._store.delete, session_id) await self._release_owner_lease_for_session(session_id) self._approval_stores.pop(session_id, None) self._session_turn_locks.pop(session_id, None) + async def _rollback_partially_created_session(self, session_id: str) -> None: + self._session_cache.pop(session_id, None) + try: + if self._local_durable_store is not None: + authority = self._owner_authority_for_session(session_id) + await self._local_durable_store.delete_session(authority) + else: + await self._run_store_io(self._store.delete, session_id) + except asyncio.CancelledError: + pass + except BaseException: + logger.exception( + "Failed to delete partially created session during rollback: %s", + session_id, + ) + try: + await self._release_owner_lease_for_session(session_id) + except asyncio.CancelledError: + pass + except BaseException: + logger.exception( + "Failed to release partially created owner lease during rollback: %s", + session_id, + ) + self._approval_stores.pop(session_id, None) + async def remove_session_async(self, session_id: str) -> None: async with self._lock: await self._remove_session_async_no_lock(session_id) @@ -2255,7 +2512,27 @@ async def _request_approval(req: ApprovalRequest) -> ApprovalResponse: async def _restore_checkpoint(self, session: Session, checkpoint_id: str) -> None: await self._runtime_checkpoint_restore_service.restore(session, checkpoint_id) + async def _restore_checkpoint_durable_state( + self, + session: Any, + snapshot: Any, + ) -> None: + typed_session = cast(Session, session) + if self._local_durable_store is None: + await self._persist_session_async(typed_session) + return + authority = self._owner_authority_for_session(typed_session.id) + await self._local_durable_store.restore_checkpoint_state( + authority, + snapshot, + cast(dict[str, Any], typed_session.to_store_data()), + ) + def _persist_session(self, session: Session) -> None: + if self._local_durable_store is not None: + raise RuntimeError( + "synchronous session persistence is unavailable for fenced local SQLite" + ) self._session_cache[session.id] = session self._store.save(session.id, cast(dict[str, Any], session.to_store_data())) @@ -2353,30 +2630,11 @@ async def create_session( async with self._lock: try: + await self._acquire_owner_for_session(session_id) await self._persist_session_async(session) await self._persist_workspace_record_for_session(session) - await self._acquire_owner_for_session(session_id) except BaseException: - self._session_cache.pop(session_id, None) - try: - await self._release_owner_lease_for_session(session_id) - except asyncio.CancelledError: - pass - except BaseException: - logger.exception( - "Failed to release partially created owner lease during rollback: %s", - session_id, - ) - try: - await self._run_store_io(self._store.delete, session_id) - except asyncio.CancelledError: - pass - except BaseException: - logger.exception( - "Failed to delete partially created session during rollback: %s", - session_id, - ) - self._approval_stores.pop(session_id, None) + await self._rollback_partially_created_session(session_id) raise logger.info(f"Created session: {session_id}") @@ -2453,11 +2711,19 @@ def has_session(self, session_id: str) -> bool: return self._store.load(session_id) is not None def register_session(self, session: Session) -> None: + if self._local_durable_store is not None: + raise RuntimeError( + "synchronous session registration is unavailable for fenced local SQLite" + ) self._runtime_closer.close_sync_safe(session) self._approval_stores[session.id] = session.approval_store self._persist_session(session) def remove_session(self, session_id: str) -> None: + if self._local_durable_store is not None: + raise RuntimeError( + "synchronous session removal is unavailable for fenced local SQLite" + ) if not self.has_session(session_id): raise KeyError(f"Session not found: {session_id}") session = self.get_session(session_id) @@ -2469,6 +2735,10 @@ def remove_session(self, session_id: str) -> None: self._session_turn_locks.pop(session_id, None) def clear_sessions(self) -> None: + if self._local_durable_store is not None: + raise RuntimeError( + "synchronous session clearing is unavailable for fenced local SQLite" + ) cleared_session_ids = set(self._session_cache) for session in list(self._session_cache.values()): self._runtime_closer.close_sync_safe(session) diff --git a/coding-agent/src/coding_agent/server/stores/session_owner_store.py b/coding-agent/src/coding_agent/server/stores/session_owner_store.py index f7c42fdfe..5908b233b 100644 --- a/coding-agent/src/coding_agent/server/stores/session_owner_store.py +++ b/coding-agent/src/coding_agent/server/stores/session_owner_store.py @@ -38,6 +38,17 @@ class SessionOwnerRecord: fencing_token: int +@dataclass(frozen=True) +class OwnerAuthority: + session_id: str + owner_id: str + epoch: int + + @property + def fencing_token(self) -> int: + return self.epoch + + class SessionOwnerBackend(Protocol): async def acquire( self, @@ -346,6 +357,107 @@ async def get_owner(self, session_id: str) -> SessionOwnerRecord | None: fencing_token=_required_sqlite_int(row, "fencing_token"), ) + async def acquire_authority( + self, + session_id: str, + owner_id: str, + lease_seconds: float = 30.0, + ) -> OwnerAuthority: + _require_owner_input( + session_id=session_id, + owner_id=owner_id, + lease_seconds=lease_seconds, + fencing_token=1, + ) + now = datetime.now(UTC) + expires_at = now + timedelta(seconds=lease_seconds) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + row = connection.execute( + """ + SELECT owner_id, lease_expires_at, fencing_token + FROM session_owners + WHERE session_id = ? + """, + (session_id,), + ).fetchone() + if row is None: + epoch = 1 + else: + current_owner_id = _required_sqlite_str(row, "owner_id") + current_expiry = _datetime_from_sqlite_text(row["lease_expires_at"]) + current_epoch = _required_sqlite_int(row, "fencing_token") + if current_expiry > now and current_owner_id != owner_id: + raise SessionOwnershipConflictError( + "stale owner or fencing token rejected" + ) + if current_expiry > now: + epoch = current_epoch + else: + epoch = current_epoch + 1 + connection.execute( + """ + INSERT INTO session_owners ( + session_id, owner_id, lease_expires_at, fencing_token, + created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(session_id) + DO UPDATE SET + owner_id = excluded.owner_id, + lease_expires_at = excluded.lease_expires_at, + fencing_token = excluded.fencing_token, + updated_at = excluded.updated_at + """, + ( + session_id, + owner_id, + _datetime_to_sqlite_text(expires_at), + epoch, + _datetime_to_sqlite_text(now), + _datetime_to_sqlite_text(now), + ), + ) + return OwnerAuthority(session_id=session_id, owner_id=owner_id, epoch=epoch) + + async def renew_authority( + self, + authority: OwnerAuthority, + lease_seconds: float = 30.0, + ) -> OwnerAuthority: + _require_owner_input( + session_id=authority.session_id, + owner_id=authority.owner_id, + lease_seconds=lease_seconds, + fencing_token=authority.epoch, + ) + now = datetime.now(UTC) + expires_at = now + timedelta(seconds=lease_seconds) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + cursor = connection.execute( + """ + UPDATE session_owners + SET lease_expires_at = ?, + updated_at = ? + WHERE session_id = ? + AND owner_id = ? + AND lease_expires_at > ? + AND fencing_token = ? + """, + ( + _datetime_to_sqlite_text(expires_at), + _datetime_to_sqlite_text(now), + authority.session_id, + authority.owner_id, + _datetime_to_sqlite_text(now), + authority.epoch, + ), + ) + if cursor.rowcount != 1: + raise SessionOwnershipConflictError("stale owner or fencing token rejected") + return authority + def _require_owner_input( *, diff --git a/coding-agent/tests/cli/test_entrypoint_contract.py b/coding-agent/tests/cli/test_entrypoint_contract.py index 14ac19104..5db5f8f79 100644 --- a/coding-agent/tests/cli/test_entrypoint_contract.py +++ b/coding-agent/tests/cli/test_entrypoint_contract.py @@ -337,6 +337,9 @@ async def get_session_async(self, session_id: str): async def run_agent(self, session_id: str, prompt: str) -> None: calls.append(("run_agent", {"session_id": session_id, "prompt": prompt})) + async def start_owner_lease_renewal(self) -> None: + calls.append(("start_owner_lease_renewal", None)) + async def close(self) -> None: calls.append(("close", None)) @@ -388,6 +391,7 @@ async def close(self) -> None: "max_steps": 3, }, ), + ("start_owner_lease_renewal", None), ("get_session_async", "session-managed"), ( "run_agent", @@ -518,10 +522,16 @@ async def get_session_async(self, session_id: str): calls.append(("get_session_async", session_id)) return SimpleNamespace(wire=wire) + async def acquire_session_owner(self, session_id: str) -> None: + calls.append(("acquire_session_owner", session_id)) + async def resume_session(self, session_id: str, **kwargs): calls.append(("resume_session", {"session_id": session_id, **kwargs})) return SimpleNamespace(run_id="run-resumed") + async def start_owner_lease_renewal(self) -> None: + calls.append(("start_owner_lease_renewal", None)) + async def close(self) -> None: calls.append(("close", None)) @@ -553,6 +563,8 @@ async def close(self) -> None: assert calls == [ calls[0], ("get_session_async", "session-managed"), + ("acquire_session_owner", "session-managed"), + ("start_owner_lease_renewal", None), ( "resume_session", { diff --git a/coding-agent/tests/cli/test_local_runtime.py b/coding-agent/tests/cli/test_local_runtime.py index e8c42b1df..1ce3f2ec8 100644 --- a/coding-agent/tests/cli/test_local_runtime.py +++ b/coding-agent/tests/cli/test_local_runtime.py @@ -4,6 +4,7 @@ from types import SimpleNamespace from coding_agent.cli.local_runtime import ServerBackedLocalCliSessionManager +from coding_agent.local_storage import local_sqlite_storage_config class FakeManagedSession: @@ -27,7 +28,7 @@ def __init__(self) -> None: self.release_owned_sessions_called = False self.close_called = False - def _persist_session(self, session) -> None: + async def _persist_session_async(self, session) -> None: self.persisted.append(session) async def acquire_session_owner(self, session_id: str) -> None: @@ -51,7 +52,9 @@ async def close(self) -> None: self.close_called = True -def test_server_backed_local_cli_attach_runtime_uses_session_binding_delegate() -> None: +async def test_server_backed_local_cli_attach_runtime_uses_session_binding_delegate() -> ( + None +): manager = object.__new__(ServerBackedLocalCliSessionManager) delegate = FakeDelegate() manager._delegate = delegate @@ -60,7 +63,7 @@ def test_server_backed_local_cli_attach_runtime_uses_session_binding_delegate() ctx = SimpleNamespace(tape=SimpleNamespace(tape_id="repl-tape")) adapter = object() - manager.attach_runtime( + await manager.attach_runtime( session, pipeline=pipeline, pipeline_ctx=ctx, @@ -125,3 +128,40 @@ async def test_server_backed_local_cli_renews_owner_leases_until_close() -> None assert manager._owner_renew_task is None assert delegate.release_owned_sessions_called is True assert delegate.close_called is True + + +async def test_server_backed_local_cli_default_sqlite_owner_starts_renewal( + tmp_path, +) -> None: + manager = ServerBackedLocalCliSessionManager( + storage_config=local_sqlite_storage_config(tmp_path), + owner_renew_interval_seconds=0.001, + ) + + await manager.start_owner_lease_renewal() + try: + assert manager.owner_store is not None + assert manager._owner_renew_task is not None + finally: + await manager.close() + + +async def test_server_backed_local_cli_default_owner_id_is_process_unique( + tmp_path, +) -> None: + first = ServerBackedLocalCliSessionManager( + storage_config=local_sqlite_storage_config(tmp_path), + ) + second = ServerBackedLocalCliSessionManager( + storage_config=local_sqlite_storage_config(tmp_path), + ) + + try: + assert isinstance(first.owner_id, str) + assert isinstance(second.owner_id, str) + assert first.owner_id.startswith("local-cli:") + assert second.owner_id.startswith("local-cli:") + assert first.owner_id != second.owner_id + finally: + await first.close() + await second.close() diff --git a/coding-agent/tests/cli/test_repl.py b/coding-agent/tests/cli/test_repl.py index 1a7ff184a..a23ac2043 100644 --- a/coding-agent/tests/cli/test_repl.py +++ b/coding-agent/tests/cli/test_repl.py @@ -1334,6 +1334,13 @@ def __init__(self) -> None: self.build_calls: list[str] = [] self.persisted_models: list[str] = [] self.closed_adapters: list[object] = [] + self._runtime_maintenance_admission = ( + session_manager_module.RuntimeMaintenanceAdmissionService( + turn_lock_for=self._turn_lock_for, + assert_owner=self._assert_owner, + load_session=self.get_session_async, + ) + ) self._runtime_closer = SimpleNamespace( close_adapter=self.close_runtime_adapter ) @@ -1418,6 +1425,9 @@ async def fake_create_session(**kwargs): return "session-init" class FakeSessionManager: + def __init__(self) -> None: + self.renewal_started = False + async def create_session(self, **kwargs): return await fake_create_session(**kwargs) @@ -1430,7 +1440,7 @@ def get_session(self, session_id: str): tape_id=None, ) - def attach_runtime( + async def attach_runtime( self, managed_session, *, @@ -1444,6 +1454,9 @@ def attach_runtime( managed_session.tape_id = pipeline_ctx.tape.tape_id assert managed_session.tape_id == "repl-init-tape" + async def start_owner_lease_renewal(self) -> None: + self.renewal_started = True + fake_pipeline = object() fake_ctx = SimpleNamespace( config={"tool_registry": "registry-a"}, @@ -1488,6 +1501,7 @@ async def close(self) -> None: "mode": "interactive", } assert session._pipeline_ctx is fake_ctx + assert fake_session_manager.renewal_started is True @pytest.mark.asyncio async def test_session_new_hook_creates_repl_origin_session(self, monkeypatch): @@ -1558,7 +1572,7 @@ def get_session(self, session_id: str): def register_session(self, managed_session): raise AssertionError("initialize should not re-register live runtime") - def attach_runtime( + async def attach_runtime( self, managed_session, *, @@ -1572,6 +1586,9 @@ def attach_runtime( managed_session.tape_id = pipeline_ctx.tape.tape_id return None + async def start_owner_lease_renewal(self) -> None: + return None + fake_pipeline = object() fake_ctx = SimpleNamespace( config={"tool_registry": "registry-a"}, diff --git a/coding-agent/tests/coding_agent/test_sqlite_local_durable_fencing.py b/coding-agent/tests/coding_agent/test_sqlite_local_durable_fencing.py new file mode 100644 index 000000000..e07123499 --- /dev/null +++ b/coding-agent/tests/coding_agent/test_sqlite_local_durable_fencing.py @@ -0,0 +1,737 @@ +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta +from pathlib import Path + +import pytest + +from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot +from agentkit.checkpoint import CheckpointService +from coding_agent.local_durable_store import SQLiteLocalDurableStore +from coding_agent.local_storage import local_sqlite_path +from coding_agent.runtime_store import AgentRunRecord, RuntimeEventRecord +from coding_agent.server.session_manager import SessionManager +from coding_agent.server.stores.session_owner_store import ( + SQLiteSessionOwnerStore, + SessionOwnershipConflictError, +) +from coding_agent.server.stores.session_store import InMemorySessionStore +from coding_agent.server.stores.session_store import SQLiteSessionStore + + +class FakeCheckpointStore: + async def save(self, snapshot: CheckpointSnapshot) -> None: + del snapshot + + async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: + del checkpoint_id + return None + + async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: + del tape_id + return [] + + async def delete(self, checkpoint_id: str) -> None: + del checkpoint_id + + +@pytest.mark.asyncio +async def test_sqlite_owner_authority_renew_does_not_advance_epoch( + tmp_path: Path, +) -> None: + store = SQLiteSessionOwnerStore(tmp_path / "local.sqlite3") + + authority = await store.acquire_authority( + "session-a", + "owner-a", + lease_seconds=30.0, + ) + renewed = await store.renew_authority(authority, lease_seconds=30.0) + + owner = await store.get_owner("session-a") + assert renewed == authority + assert owner is not None + assert owner.owner_id == "owner-a" + assert owner.fencing_token == authority.epoch == 1 + + +@pytest.mark.asyncio +async def test_sqlite_owner_authority_takeover_advances_epoch_in_database( + tmp_path: Path, +) -> None: + store = SQLiteSessionOwnerStore(tmp_path / "local.sqlite3") + + dead = await store.acquire_authority( + "session-a", + "owner-dead", + lease_seconds=0.001, + ) + await asyncio.sleep(0.01) + + taken = await store.acquire_authority( + "session-a", + "owner-b", + lease_seconds=30.0, + ) + + assert dead.epoch == 1 + assert taken.owner_id == "owner-b" + assert taken.epoch == 2 + + +@pytest.mark.asyncio +async def test_sqlite_owner_authority_same_owner_expired_reacquire_advances_epoch( + tmp_path: Path, +) -> None: + store = SQLiteSessionOwnerStore(tmp_path / "local.sqlite3") + + stale = await store.acquire_authority( + "session-a", + "owner-a", + lease_seconds=0.001, + ) + await asyncio.sleep(0.01) + reacquired = await store.acquire_authority( + "session-a", + "owner-a", + lease_seconds=30.0, + ) + + assert stale.epoch == 1 + assert reacquired.epoch == 2 + + +@pytest.mark.asyncio +async def test_session_manager_uses_sqlite_owner_authority_epochs( + tmp_path: Path, +) -> None: + owner_store = SQLiteSessionOwnerStore(tmp_path / "local.sqlite3") + manager = SessionManager( + store=InMemorySessionStore(), + checkpoint_service=CheckpointService(FakeCheckpointStore()), + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + + session_id = await manager.create_session() + owner = await owner_store.get_owner(session_id) + + assert owner is not None + assert owner.owner_id == "owner-a" + assert owner.fencing_token == 1 + + await manager.renew_owner_leases() + renewed = await owner_store.get_owner(session_id) + + assert renewed is not None + assert renewed.owner_id == "owner-a" + assert renewed.fencing_token == 1 + + await manager.release_owned_sessions() + assert await owner_store.get_owner(session_id) is None + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_uses_protected_session_save( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + + session_id = await manager.create_session() + + owner = await owner_store.get_owner(session_id) + payload = SQLiteSessionStore(sqlite_path).load(session_id) + + assert owner is not None + assert owner.fencing_token == 1 + assert payload is not None + assert payload["id"] == session_id + assert payload["tape_id"] is None + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_fences_session_delete( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + session_id = await manager.create_session() + assert await owner_store.release(session_id, "owner-a", 1) is True + await owner_store.acquire_authority(session_id, "owner-b", lease_seconds=30.0) + + with pytest.raises(SessionOwnershipConflictError): + await manager.remove_session_async(session_id) + + assert SQLiteSessionStore(sqlite_path).load(session_id) is not None + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_create_rollback_deletes_before_owner_release( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + session_store = SQLiteSessionStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + release_observations: list[tuple[str, bool]] = [] + original_release = manager._release_owner_lease_for_session + + async def fail_workspace_persist(session) -> None: + raise RuntimeError(f"workspace persist failed: {session.id}") + + async def observe_release(session_id: str) -> None: + release_observations.append( + (session_id, session_store.load(session_id) is None) + ) + await original_release(session_id) + + monkeypatch.setattr( + manager, + "_persist_workspace_record_for_session", + fail_workspace_persist, + ) + monkeypatch.setattr( + manager, + "_release_owner_lease_for_session", + observe_release, + ) + + with pytest.raises(RuntimeError, match="workspace persist failed"): + await manager.create_session() + + assert release_observations + session_id, session_deleted_before_release = release_observations[0] + assert session_deleted_before_release is True + assert session_store.load(session_id) is None + assert await owner_store.get_owner(session_id) is None + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_rejects_sync_mutations( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + owner_store = SQLiteSessionOwnerStore(local_sqlite_path(tmp_path)) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + session_id = await manager.create_session() + session = await manager.get_session_async(session_id) + + with pytest.raises(RuntimeError, match="synchronous session persistence"): + manager._persist_session(session) + with pytest.raises(RuntimeError, match="synchronous session removal"): + manager.remove_session(session_id) + with pytest.raises(RuntimeError, match="synchronous session clearing"): + manager.clear_sessions() + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_fences_tape_save( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + session_id = await manager.create_session() + session = await manager.get_session_async(session_id) + session.tape_id = "tape-a" + await manager._persist_session_async(session) + + await manager._tape_store.save( + "tape-a", + [{"kind": "message", "payload": {"text": "before takeover"}}], + ) + + assert await owner_store.release("session-missing", "owner-a", 1) is False + assert await owner_store.release(session_id, "owner-a", 1) is True + await owner_store.acquire_authority(session_id, "owner-b", lease_seconds=30.0) + + with pytest.raises(SessionOwnershipConflictError): + await manager._tape_store.save( + "tape-a", + [{"kind": "message", "payload": {"text": "stale owner"}}], + ) + + assert await manager._tape_store.load("tape-a") == [ + {"kind": "message", "payload": {"text": "before takeover"}} + ] + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_fences_runtime_event( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + session_id = await manager.create_session() + session = await manager.get_session_async(session_id) + session.tape_id = "tape-a" + await manager._persist_session_async(session) + + started_at = datetime.now(UTC) + await manager._runtime_store.create_agent_run( + AgentRunRecord( + run_id="run-a", + session_id=session_id, + tape_id="tape-a", + parent_run_id=None, + agent_id="agent-a", + status="running", + started_at=started_at, + ) + ) + await manager._runtime_store.append_runtime_event( + RuntimeEventRecord( + event_id="event-before", + run_id="run-a", + event_kind="progress", + payload={"step": "before"}, + created_at=started_at, + ) + ) + + assert await owner_store.release(session_id, "owner-a", 1) is True + await owner_store.acquire_authority(session_id, "owner-b", lease_seconds=30.0) + + with pytest.raises(SessionOwnershipConflictError): + await manager._runtime_store.append_runtime_event( + RuntimeEventRecord( + event_id="event-after", + run_id="run-a", + event_kind="progress", + payload={"step": "after"}, + created_at=started_at + timedelta(seconds=1), + ) + ) + + assert [ + event.event_id + for event in await manager._runtime_store.replay_runtime_events("run-a") + ] == ["event-before"] + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_fences_runtime_claim( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + session_id = await manager.create_session() + session = await manager.get_session_async(session_id) + session.tape_id = "tape-a" + await manager._persist_session_async(session) + await manager._runtime_store.create_agent_run( + AgentRunRecord( + run_id="run-a", + session_id=session_id, + tape_id="tape-a", + parent_run_id=None, + agent_id=None, + status="requested", + started_at=datetime.now(UTC), + metadata={ + "executor_ref_kind": "external_worker", + "executor_kind": "test-worker", + }, + ) + ) + + assert await owner_store.release(session_id, "owner-a", 1) is True + await owner_store.acquire_authority(session_id, "owner-b", lease_seconds=30.0) + + assert ( + await manager._runtime_store.claim_external_worker_run( + session_id=None, + executor_kind="test-worker", + claim_metadata={"worker_id": "worker-a"}, + ) + is None + ) + assert ( + await manager._runtime_store.claim_external_worker_run( + session_id=session_id, + executor_kind="test-worker", + claim_metadata={"worker_id": "worker-a"}, + ) + is None + ) + + run = await manager._runtime_store.load_agent_run("run-a") + assert run is not None + assert run.status == "requested" + + +@pytest.mark.asyncio +async def test_session_manager_default_sqlite_bundle_fences_checkpoint_save( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AGENT_DATA_DIR", str(tmp_path)) + sqlite_path = local_sqlite_path(tmp_path) + owner_store = SQLiteSessionOwnerStore(sqlite_path) + manager = SessionManager( + owner_store=owner_store, + owner_id="owner-a", + fencing_token=999, + ) + session_id = await manager.create_session() + session = await manager.get_session_async(session_id) + session.tape_id = "tape-a" + await manager._persist_session_async(session) + created_at = datetime.now(UTC) + + checkpoint = CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-a", + tape_id="tape-a", + session_id=session_id, + entry_count=0, + window_start=0, + created_at=created_at, + label="before", + ), + tape_entries=(), + plugin_states={}, + ) + await manager._checkpoint_service._store.save(checkpoint) + + assert await owner_store.release(session_id, "owner-a", 1) is True + await owner_store.acquire_authority(session_id, "owner-b", lease_seconds=30.0) + + with pytest.raises(SessionOwnershipConflictError): + await manager._checkpoint_service._store.save( + CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-a", + tape_id="tape-a", + session_id=session_id, + entry_count=1, + window_start=0, + created_at=created_at + timedelta(seconds=1), + label="after", + ), + tape_entries=({"kind": "message"},), + plugin_states={}, + ) + ) + + loaded = await manager._checkpoint_service._store.load("checkpoint-a") + assert loaded is not None + assert loaded.meta.entry_count == 0 + assert loaded.meta.label == "before" + + +@pytest.mark.asyncio +async def test_sqlite_local_durable_store_rejects_stale_epoch_for_session_metadata( + tmp_path: Path, +) -> None: + store = SQLiteLocalDurableStore(tmp_path / "local.sqlite3") + stale = await store.acquire_owner("session-a", "owner-a", lease_seconds=0.001) + await asyncio.sleep(0.01) + current = await store.acquire_owner("session-a", "owner-b", lease_seconds=30.0) + + with pytest.raises(SessionOwnershipConflictError): + await store.save_session( + stale, {"id": "session-a", "session_id": "session-a", "tape_id": "tape-a"} + ) + + await store.save_session( + current, + {"id": "session-a", "session_id": "session-a", "tape_id": "tape-a"}, + ) + assert store.load_session("session-a") == { + "id": "session-a", + "session_id": "session-a", + "tape_id": "tape-a", + } + + +@pytest.mark.asyncio +async def test_sqlite_local_durable_store_rejects_session_payload_id_mismatch( + tmp_path: Path, +) -> None: + store = SQLiteLocalDurableStore(tmp_path / "local.sqlite3") + owner_a = await store.acquire_owner("session-a", "owner-a") + + with pytest.raises(SessionOwnershipConflictError): + await store.save_session( + owner_a, + {"id": "session-b", "session_id": "session-a", "tape_id": "tape-a"}, + ) + + +@pytest.mark.asyncio +async def test_sqlite_local_durable_store_rejects_cross_session_tape_write( + tmp_path: Path, +) -> None: + store = SQLiteLocalDurableStore(tmp_path / "local.sqlite3") + owner_a = await store.acquire_owner("session-a", "owner-a") + owner_b = await store.acquire_owner("session-b", "owner-b") + + await store.save_session( + owner_a, + {"id": "session-a", "session_id": "session-a", "tape_id": "tape-a"}, + ) + await store.save_session( + owner_b, + {"id": "session-b", "session_id": "session-b", "tape_id": "tape-b"}, + ) + await store.append_tape_entries( + owner_a, + "tape-a", + [{"kind": "message", "payload": {"text": "owned by session-a"}}], + ) + + with pytest.raises(SessionOwnershipConflictError): + await store.append_tape_entries( + owner_b, + "tape-a", + [{"kind": "message", "payload": {"text": "wrong owner"}}], + ) + + assert await store.load_tape("tape-a") == [ + {"kind": "message", "payload": {"text": "owned by session-a"}} + ] + + +@pytest.mark.asyncio +async def test_sqlite_local_durable_store_rejects_cross_session_runtime_write( + tmp_path: Path, +) -> None: + store = SQLiteLocalDurableStore(tmp_path / "local.sqlite3") + owner_a = await store.acquire_owner("session-a", "owner-a") + owner_b = await store.acquire_owner("session-b", "owner-b") + started_at = datetime.now(UTC) + + run = AgentRunRecord( + run_id="run-a", + session_id="session-a", + tape_id="tape-a", + parent_run_id=None, + agent_id="agent-a", + status="running", + started_at=started_at, + ) + await store.create_agent_run(owner_a, run) + + with pytest.raises(SessionOwnershipConflictError): + await store.append_runtime_event( + owner_b, + RuntimeEventRecord( + event_id="event-wrong-owner", + run_id="run-a", + event_kind="progress", + payload={"step": "wrong-owner"}, + created_at=started_at, + ), + ) + + await store.append_runtime_event( + owner_a, + RuntimeEventRecord( + event_id="event-owner-a", + run_id="run-a", + event_kind="progress", + payload={"step": "owner-a"}, + created_at=started_at + timedelta(seconds=1), + ), + ) + assert [event.event_id for event in await store.replay_runtime_events("run-a")] == [ + "event-owner-a" + ] + + +@pytest.mark.asyncio +async def test_sqlite_local_durable_store_rejects_checkpoint_rebind( + tmp_path: Path, +) -> None: + store = SQLiteLocalDurableStore(tmp_path / "local.sqlite3") + owner_a = await store.acquire_owner("session-a", "owner-a") + owner_b = await store.acquire_owner("session-b", "owner-b") + created_at = datetime.now(UTC) + + await store.save_session( + owner_a, + {"id": "session-a", "session_id": "session-a", "tape_id": "tape-a"}, + ) + await store.save_session( + owner_b, + {"id": "session-b", "session_id": "session-b", "tape_id": "tape-b"}, + ) + await store.save_checkpoint( + owner_a, + CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-a", + tape_id="tape-a", + session_id="session-a", + entry_count=0, + window_start=0, + created_at=created_at, + label="owned", + ), + tape_entries=(), + plugin_states={}, + ), + ) + + with pytest.raises(SessionOwnershipConflictError): + await store.save_checkpoint( + owner_b, + CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-a", + tape_id="tape-b", + session_id="session-b", + entry_count=0, + window_start=0, + created_at=created_at, + label="rebind", + ), + tape_entries=(), + plugin_states={}, + ), + ) + + checkpoint = await store.load_checkpoint("checkpoint-a") + assert checkpoint is not None + assert checkpoint.meta.session_id == "session-a" + assert checkpoint.meta.tape_id == "tape-a" + + +@pytest.mark.asyncio +async def test_sqlite_local_durable_store_restores_checkpoint_state_atomically( + tmp_path: Path, +) -> None: + store = SQLiteLocalDurableStore(tmp_path / "local.sqlite3") + owner = await store.acquire_owner("session-a", "owner-a") + created_at = datetime.now(UTC) + await store.save_session( + owner, + {"id": "session-a", "session_id": "session-a", "tape_id": "tape-a"}, + ) + await store.append_tape_entries( + owner, + "tape-a", + [ + {"kind": "message", "payload": {"text": "keep"}}, + {"kind": "message", "payload": {"text": "drop"}}, + ], + ) + await store.save_checkpoint( + owner, + CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-keep", + tape_id="tape-a", + session_id="session-a", + entry_count=1, + window_start=0, + created_at=created_at, + label="keep", + ), + tape_entries=({"kind": "message", "payload": {"text": "keep"}},), + plugin_states={}, + ), + ) + await store.save_checkpoint( + owner, + CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-drop", + tape_id="tape-a", + session_id="session-a", + entry_count=2, + window_start=0, + created_at=created_at + timedelta(seconds=1), + label="drop", + ), + tape_entries=( + {"kind": "message", "payload": {"text": "keep"}}, + {"kind": "message", "payload": {"text": "drop"}}, + ), + plugin_states={}, + ), + ) + + await store.restore_checkpoint_state( + owner, + CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="checkpoint-keep", + tape_id="tape-a", + session_id="session-a", + entry_count=1, + window_start=0, + created_at=created_at, + label="keep", + ), + tape_entries=({"kind": "message", "payload": {"text": "keep"}},), + plugin_states={}, + ), + { + "id": "session-a", + "session_id": "session-a", + "tape_id": "tape-a", + "provider_name": "restored-provider", + }, + ) + + assert await store.load_tape("tape-a") == [ + {"kind": "message", "payload": {"text": "keep"}} + ] + assert store.load_session("session-a") == { + "id": "session-a", + "session_id": "session-a", + "tape_id": "tape-a", + "provider_name": "restored-provider", + } + assert await store.load_checkpoint("checkpoint-keep") is not None + assert await store.load_checkpoint("checkpoint-drop") is None diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 2aaed6a2e..e27f9c701 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -2002,22 +2002,84 @@ def test_build_session_manager_requires_explicit_positive_fencing_token_for_pg_h with pytest.raises(ValueError, match="storage.fencing_token"): _build_session_manager() - def test_build_session_manager_does_not_enable_owner_store_for_non_pg_http_sessions( + def test_build_session_manager_enables_owner_store_for_local_sqlite_bundle( + self, monkeypatch, tmp_path + ): + local_path = tmp_path / "local.sqlite3" + monkeypatch.setattr( + "coding_agent.server.http_server._load_storage_config", + lambda: { + "http_session_backend": "sqlite", + "http_session_path": str(local_path), + "tape_backend": "sqlite", + "tape_path": str(local_path), + "checkpoint_backend": "sqlite", + "checkpoint_path": str(local_path), + "runtime_backend": "sqlite", + "runtime_path": str(local_path), + "owner_id": "server-a", + "owner_lease_seconds": 40.0, + }, + ) + + manager = _build_session_manager() + try: + assert manager._owner_store is not None + assert type(manager._owner_store).__name__ == "SQLiteSessionOwnerStore" + assert manager._owner_id == "server-a" + assert manager.owner_lease_seconds == 40.0 + assert manager._local_durable_store is not None + finally: + asyncio.run(manager.close()) + + def test_build_session_manager_enables_owner_store_for_equivalent_local_sqlite_bundle_paths( + self, monkeypatch, tmp_path + ): + local_path = tmp_path / "local.sqlite3" + equivalent_parent = tmp_path / "nested" / ".." / "local.sqlite3" + monkeypatch.setattr( + "coding_agent.server.http_server._load_storage_config", + lambda: { + "http_session_backend": "sqlite", + "http_session_path": str(local_path), + "tape_backend": "sqlite", + "tape_path": str(equivalent_parent), + "checkpoint_backend": "sqlite", + "checkpoint_path": str(local_path.resolve()), + "runtime_backend": "sqlite", + "runtime_path": str(equivalent_parent), + "owner_id": "server-a", + "owner_lease_seconds": 40.0, + }, + ) + + manager = _build_session_manager() + try: + assert manager._owner_store is not None + assert type(manager._owner_store).__name__ == "SQLiteSessionOwnerStore" + assert manager._local_durable_store is not None + finally: + asyncio.run(manager.close()) + + def test_build_session_manager_does_not_enable_owner_store_for_non_bundle_storage( self, monkeypatch ): monkeypatch.setattr( "coding_agent.server.http_server._load_storage_config", lambda: { "http_session_backend": "redis", - "tape_backend": "pg", - "checkpoint_backend": "pg", - "dsn": "postgresql://example", + "tape_backend": "jsonl", + "checkpoint_backend": "fs", + "runtime_backend": "jsonl", }, ) manager = _build_session_manager() - assert manager._owner_store is None + try: + assert manager._owner_store is None + finally: + asyncio.run(manager.close()) async def test_renew_owner_leases_exits_when_owner_leases_are_not_configured( self, monkeypatch diff --git a/coding-agent/tests/ui/test_session_manager_public_api.py b/coding-agent/tests/ui/test_session_manager_public_api.py index 708d728cd..2202e5300 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -36,6 +36,7 @@ TurnEnd, ) from coding_agent.server.session_manager import MockProvider, Session, SessionManager +from coding_agent.server.stores.session_owner_store import SQLiteSessionOwnerStore from coding_agent.server.stores.session_store import ( InMemorySessionStore, PGSessionMetadataStore, @@ -623,6 +624,38 @@ def test_session_manager_uses_storage_paths_local_for_sqlite_bundle( assert manager._runtime_store._path == custom_path +def test_session_manager_uses_equivalent_explicit_sqlite_paths_as_bundle( + tmp_path, +) -> None: + custom_path = tmp_path / "custom" / "bundle.sqlite3" + equivalent_path = tmp_path / "custom" / "nested" / ".." / "bundle.sqlite3" + + storage_config = { + "http_session_backend": "sqlite", + "http_session_path": str(custom_path), + "tape_backend": "sqlite", + "tape_path": str(equivalent_path), + "checkpoint_backend": "sqlite", + "checkpoint_path": str(custom_path.resolve()), + "runtime_backend": "sqlite", + "runtime_path": str(equivalent_path), + } + + manager = SessionManager( + storage_config=storage_config, + owner_store=SQLiteSessionOwnerStore(custom_path), + owner_id="test-owner", + fencing_token=1, + ) + + resolved_path = custom_path.resolve() + assert manager._store._path.resolve() == resolved_path + assert manager._local_durable_store is not None + assert manager._tape_store._delegate._path.resolve() == resolved_path + assert manager._checkpoint_service._store._delegate._path.resolve() == resolved_path + assert manager._runtime_store._delegate._path.resolve() == resolved_path + + def test_session_manager_explicit_sqlite_paths_override_storage_paths_local( tmp_path, ) -> None: