diff --git a/coding-agent/.opencode/prompts/postgres-phase1-completion-task-packet.md b/coding-agent/.opencode/prompts/postgres-phase1-completion-task-packet.md new file mode 100644 index 000000000..68bfca5cb --- /dev/null +++ b/coding-agent/.opencode/prompts/postgres-phase1-completion-task-packet.md @@ -0,0 +1,45 @@ +Goal: +Complete Phase 1 PostgreSQL persistence for the HTTP/UI session stack by adding PG checkpoint storage and config-driven default persistence wiring for session metadata, tapes, and checkpoints. + +Scope: +- Add `PGCheckpointStore` in `agentkit.storage.pg` with matching protocol tests. +- Add a synchronous UI PostgreSQL session metadata store and one app-layer persistence factory for default SessionManager/HTTP wiring. +- Update config defaults/tests so HTTP startup and SessionManager defaults can use file or PostgreSQL storage through one shared path. + +Out of scope: +- Phase 2 ownership, fencing, approval routing, or cross-instance event routing. +- New migration tooling or schema-management framework beyond the repo's existing inline schema creation pattern. + +Context: +- ADRs: + - `docs/adr/0012-complete-phase1-postgresql-http-session-persistence.md` + - `docs/adr/0001-checkpoint-captures-serialized-tape-and-plugin-state.md` + - `docs/adr/0003-http-sessions-use-one-stable-tape-timeline.md` + - `docs/adr/0005-checkpoint-restore-uses-truncate-rollback.md` +- Relevant files: + - `src/agentkit/storage/pg.py` + - `src/coding_agent/ui/session_store.py` + - `src/coding_agent/ui/session_manager.py` + - `src/coding_agent/ui/http_server.py` + - `src/coding_agent/agent.toml` + - `tests/agentkit/storage/test_pg.py` + - `tests/ui/test_session_persistence.py` + - `tests/ui/test_session_manager_public_api.py` + - `tests/ui/test_http_server.py` + +Target tests: +- `uv run pytest tests/agentkit/storage/test_pg.py -k "checkpoint or session or tape" -v` +- `uv run pytest tests/ui/test_session_persistence.py tests/ui/test_session_manager_public_api.py tests/ui/test_http_server.py -k "pg or checkpoint or persistence" -v` +- `uv run pytest tests/ -v` + +Loop policy: +- Engineer implements the smallest correct change and runs the target tests. +- Reviewer reviews only the resulting diff and affected tests. +- Reviewer reports only P1/P2 findings. +- Engineer fixes only accepted P1/P2 findings and reruns the same target tests. +- Verifier reruns the exact target tests and reports pass/fail only. + +Stop conditions: +- At most one review/fix/retest cycle. +- Escalate architectural redirection or scope expansion to the human. +- Ignore non-blocking optimization suggestions. diff --git a/coding-agent/docs/adr/0012-complete-phase1-postgresql-http-session-persistence.md b/coding-agent/docs/adr/0012-complete-phase1-postgresql-http-session-persistence.md new file mode 100644 index 000000000..3dfd2e61a --- /dev/null +++ b/coding-agent/docs/adr/0012-complete-phase1-postgresql-http-session-persistence.md @@ -0,0 +1,53 @@ +# ADR-0012: Complete Phase 1 PostgreSQL HTTP session persistence + +**Status**: Proposed +**Date**: 2026-04-17 + +## Context + +The repository already includes reusable PostgreSQL persistence primitives for session metadata and tapes in `agentkit.storage.pg`, and the HTTP/UI session stack already has stable tape and checkpoint semantics backed by tests. But the Phase 1 goal for HTTP session persistence is still incomplete on `main`. + +`SessionManager` still defaults to the local JSONL tape store and filesystem checkpoint store, `http_server.py` still constructs `SessionManager()` directly with those defaults, and the UI session metadata layer still only provides in-memory and Redis-backed implementations. There is also no PostgreSQL `CheckpointStore` implementation in `agentkit`, which means the current HTTP checkpoint path cannot persist checkpoint snapshots/catalogs in PostgreSQL even when tape persistence uses PostgreSQL. + +Phase 1 is still explicitly single-instance and keeps active runtime state (`task`, `approval_event`, `event_queues`, `runtime_pipeline`, `runtime_ctx`, `runtime_adapter`) local-authoritative. This ADR only closes the durable storage gap for restart-safe metadata, stable tape history, and checkpoint snapshots/catalogs. + +## Decision + +Complete Phase 1 by adding PostgreSQL-backed checkpoint persistence in `agentkit` and one config-driven app-layer persistence path for the HTTP/UI session stack. + +Specifically: + +- Add `PGCheckpointStore` to `agentkit.storage.pg` implementing the existing `CheckpointStore` protocol. +- Add a synchronous PostgreSQL-backed UI session metadata store in `coding_agent.ui.session_store` rather than widening `SessionManager` to the async `agentkit` session protocol. +- Add a small app-layer persistence factory in `coding_agent.ui` that builds the default session metadata store, tape store, checkpoint store, and shutdown hook from the existing flat `[storage]` config in `src/coding_agent/agent.toml`. +- Make both `SessionManager()` defaults and HTTP server startup use that same persistence factory instead of hard-coded JSONL/filesystem/in-memory defaults. +- Keep runtime execution ownership, approval routing, and event routing out of scope for this ADR and Phase 1. + +## Alternatives Rejected + +- Reuse the async `agentkit.PGSessionStore` directly inside `SessionManager` — rejected because it would widen a synchronous UI/session-manager boundary into a larger refactor. +- Keep checkpoints on the filesystem while only moving tapes to PostgreSQL — rejected because checkpoint restore belongs to the same stable tape timeline and Phase 1 needs durable checkpoint snapshots/catalogs. +- Add PostgreSQL support by sprinkling separate wiring logic across `SessionManager`, `http_server.py`, and REPL startup — rejected because Phase 1 needs one config-driven default path, not divergent per-entrypoint behavior. +- Expand into Phase 2 owner/lease/fencing work while touching persistence — rejected because Phase 1 remains single-instance and local-authoritative for active runtime state. + +## Acceptance Criteria + +- [ ] PostgreSQL checkpoint persistence is covered by tests in `tests/agentkit/storage/test_pg.py`, including round-trip snapshot persistence and ordered listing by tape. +- [ ] Config-driven persistence bundle construction from `[storage]` configuration is covered by tests in `tests/coding_agent/plugins/test_storage_factory.py`. +- [ ] `SessionManager` default persistence wiring is covered by tests in `tests/ui/test_session_manager_public_api.py` and `tests/ui/test_session_persistence.py`. +- [ ] HTTP server startup uses the configured persistence bundle for the global session manager, covered by tests in `tests/ui/test_http_server.py`. +- [ ] `uv run pytest tests/agentkit/storage/test_pg.py tests/coding_agent/plugins/test_storage_factory.py tests/ui/test_session_persistence.py tests/ui/test_session_manager_public_api.py tests/ui/test_http_server.py -v` + +## References + +- `src/agentkit/storage/pg.py` +- `src/agentkit/storage/protocols.py` +- `src/agentkit/checkpoint/service.py` +- `src/coding_agent/ui/session_store.py` +- `src/coding_agent/ui/session_manager.py` +- `src/coding_agent/ui/http_server.py` +- `src/coding_agent/agent.toml` +- `tests/agentkit/storage/test_pg.py` +- `tests/ui/test_session_persistence.py` +- `tests/ui/test_session_manager_public_api.py` +- `tests/ui/test_http_server.py` diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py index 002c457b4..98d48dc53 100644 --- a/coding-agent/src/agentkit/storage/pg.py +++ b/coding-agent/src/agentkit/storage/pg.py @@ -4,8 +4,11 @@ import importlib import json from collections.abc import Awaitable, Callable +from datetime import datetime from typing import Final, Protocol, cast +from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot + class AsyncPGPool(Protocol): async def execute(self, query: str, *args: object) -> str: ... @@ -31,6 +34,17 @@ async def acquire(self) -> LockConnection: ... async def release(self, connection: LockConnection) -> None: ... +class CodecConnection(Protocol): + async def set_type_codec( + self, + typename: str, + *, + encoder: Callable[[object], str], + decoder: Callable[[str], object], + schema: str, + ) -> None: ... + + PoolFactory = Callable[..., Awaitable[AsyncPGPool]] @@ -59,6 +73,7 @@ async def get_pool(self) -> AsyncPGPool: dsn=self._dsn, min_size=self._min_size, max_size=self._max_size, + init=_init_connection_codecs, ) return self._pool @@ -110,7 +125,7 @@ async def _ensure_schema(self) -> AsyncPGPool: async def save_session(self, session_id: str, data: dict[str, object]) -> None: pool = await self._ensure_schema() - _ = await pool.execute(self._UPSERT_SQL, session_id, json.dumps(data)) + _ = await pool.execute(self._UPSERT_SQL, session_id, data) async def load_session(self, session_id: str) -> dict[str, object] | None: pool = await self._ensure_schema() @@ -150,13 +165,21 @@ class PGTapeStore: PRIMARY KEY (tape_id, seq) ) """ - _MAX_SEQ_SQL: Final[str] = ( - "SELECT MAX(seq) AS max_seq FROM agent_tapes WHERE tape_id = $1" - ) _INSERT_SQL: Final[str] = """ + WITH advisory_lock AS ( + SELECT pg_advisory_xact_lock(hashtext($1)) + ), + max_seq AS ( + SELECT COALESCE(MAX(seq), -1) AS value + FROM agent_tapes + WHERE tape_id = $1 + ) INSERT INTO agent_tapes (tape_id, seq, entry) - SELECT $1, seq, entry::jsonb - FROM unnest($2::int[], $3::text[]) AS batch(seq, entry) + SELECT + $1, + (SELECT value FROM max_seq) + batch.ordinality, + batch.entry::jsonb + FROM advisory_lock, unnest($2::text[]) WITH ORDINALITY AS batch(entry, ordinality) """ _LOAD_SQL: Final[str] = ( "SELECT entry FROM agent_tapes WHERE tape_id = $1 ORDER BY seq" @@ -182,25 +205,8 @@ async def save(self, tape_id: str, entries: list[dict[str, object]]) -> None: return pool = await self._ensure_schema() - row = await pool.fetchrow(self._MAX_SEQ_SQL, tape_id) - max_seq = -1 - if row is not None: - raw_max_seq = row.get("max_seq") - if raw_max_seq is None: - max_seq = -1 - elif isinstance(raw_max_seq, int): - max_seq = raw_max_seq - else: - raise TypeError("postgres tape row must include int max_seq or None") - - seq_values = [max_seq + offset for offset, _ in enumerate(entries, start=1)] payload_values = [json.dumps(entry) for entry in entries] - _ = await pool.execute( - self._INSERT_SQL, - tape_id, - seq_values, - payload_values, - ) + _ = await pool.execute(self._INSERT_SQL, tape_id, payload_values) async def load(self, tape_id: str) -> list[dict[str, object]]: pool = await self._ensure_schema() @@ -231,6 +237,94 @@ async def truncate(self, tape_id: str, keep: int) -> None: _ = await pool.execute(self._TRUNCATE_SQL, tape_id, keep) +class PGCheckpointStore: + _CREATE_TABLE_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS agent_checkpoints ( + checkpoint_id TEXT PRIMARY KEY, + tape_id TEXT NOT NULL, + meta JSONB NOT NULL, + entries JSONB NOT NULL, + plugin_states JSONB NOT NULL, + extra JSONB NOT NULL + ) + """ + _INSERT_SQL: Final[str] = """ + INSERT INTO agent_checkpoints ( + checkpoint_id, + tape_id, + meta, + entries, + plugin_states, + extra + ) + VALUES ($1, $2, $3::jsonb, $4::jsonb, $5::jsonb, $6::jsonb) + ON CONFLICT (checkpoint_id) DO UPDATE SET + tape_id = EXCLUDED.tape_id, + meta = EXCLUDED.meta, + entries = EXCLUDED.entries, + plugin_states = EXCLUDED.plugin_states, + extra = EXCLUDED.extra + """ + _LOAD_SQL: Final[str] = ( + "SELECT meta, entries, plugin_states, extra FROM agent_checkpoints " + "WHERE checkpoint_id = $1" + ) + _LIST_SQL: Final[str] = ( + "SELECT meta FROM agent_checkpoints WHERE tape_id = $1 " + "ORDER BY (meta->>'created_at')::timestamptz" + ) + _DELETE_SQL: Final[str] = "DELETE FROM agent_checkpoints WHERE checkpoint_id = $1" + + def __init__(self, *, pool: PGPool) -> None: + self._pool: PGPool = pool + self._schema_ready: bool = False + + async def _ensure_schema(self) -> AsyncPGPool: + pool = await self._pool.get_pool() + if not self._schema_ready: + _ = await pool.execute(self._CREATE_TABLE_SQL) + self._schema_ready = True + return pool + + async def save(self, snapshot: CheckpointSnapshot) -> None: + pool = await self._ensure_schema() + meta = snapshot.meta + meta_payload = { + "checkpoint_id": meta.checkpoint_id, + "tape_id": meta.tape_id, + "session_id": meta.session_id, + "entry_count": meta.entry_count, + "window_start": meta.window_start, + "created_at": meta.created_at.isoformat(), + "label": meta.label, + } + _ = await pool.execute( + self._INSERT_SQL, + meta.checkpoint_id, + meta.tape_id, + meta_payload, + list(snapshot.tape_entries), + snapshot.plugin_states, + snapshot.extra, + ) + + async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: + pool = await self._ensure_schema() + row = await pool.fetchrow(self._LOAD_SQL, checkpoint_id) + if row is None: + return None + return _snapshot_from_row(row) + + async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: + pool = await self._ensure_schema() + rows = await pool.fetch(self._LIST_SQL, tape_id) + return [_meta_from_raw(_required_dict(row, "meta")) for row in rows] + + async def delete(self, checkpoint_id: str) -> None: + pool = await self._ensure_schema() + _ = await pool.execute(self._DELETE_SQL, checkpoint_id) + + class PGSessionLock: def __init__(self, *, pool: LockPool) -> None: self._pool: LockPool = pool @@ -254,7 +348,10 @@ async def release(self) -> None: try: _ = await self._conn.execute("SELECT pg_advisory_unlock_all()") finally: - await self._pool.release(self._conn) + pool_release: Awaitable[None] = cast( + Awaitable[None], self._pool.release(self._conn) + ) + await pool_release self._conn = None @@ -270,3 +367,81 @@ def _load_asyncpg_pool_factory() -> PoolFactory: if not callable(create_pool): raise ImportError("asyncpg does not expose create_pool") return cast(PoolFactory, create_pool) + + +async def _init_connection_codecs(connection: CodecConnection) -> None: + for typename in ("json", "jsonb"): + await connection.set_type_codec( + typename, + encoder=json.dumps, + decoder=json.loads, + schema="pg_catalog", + ) + + +def _required_dict(row: dict[str, object], key: str) -> dict[str, object]: + value = row.get(key) + if not isinstance(value, dict): + raise TypeError(f"postgres checkpoint row must include dict {key}") + return cast(dict[str, object], value) + + +def _required_list(row: dict[str, object], key: str) -> list[object]: + value = row.get(key) + if not isinstance(value, list): + raise TypeError(f"postgres checkpoint row must include list {key}") + return value + + +def _meta_from_raw(raw: dict[str, object]) -> CheckpointMeta: + checkpoint_id = raw.get("checkpoint_id") + tape_id = raw.get("tape_id") + session_id = raw.get("session_id") + entry_count = raw.get("entry_count") + window_start = raw.get("window_start") + created_at = raw.get("created_at") + label = raw.get("label") + + if not isinstance(checkpoint_id, str): + raise TypeError("postgres checkpoint meta must include string checkpoint_id") + if not isinstance(tape_id, str): + raise TypeError("postgres checkpoint meta must include string tape_id") + if session_id is not None and not isinstance(session_id, str): + raise TypeError( + "postgres checkpoint meta must include string session_id or None" + ) + if not isinstance(entry_count, int): + raise TypeError("postgres checkpoint meta must include int entry_count") + if not isinstance(window_start, int): + raise TypeError("postgres checkpoint meta must include int window_start") + if not isinstance(created_at, str): + raise TypeError("postgres checkpoint meta must include string created_at") + if label is not None and not isinstance(label, str): + raise TypeError("postgres checkpoint meta must include string label or None") + + return CheckpointMeta( + checkpoint_id=checkpoint_id, + tape_id=tape_id, + session_id=session_id, + entry_count=entry_count, + window_start=window_start, + created_at=datetime.fromisoformat(created_at), + label=label, + ) + + +def _snapshot_from_row(row: dict[str, object]) -> CheckpointSnapshot: + entries_raw = _required_list(row, "entries") + plugin_states = _required_dict(row, "plugin_states") + extra = _required_dict(row, "extra") + tape_entries: list[dict[str, object]] = [] + for entry in entries_raw: + if not isinstance(entry, dict): + raise TypeError("postgres checkpoint entries must contain dict items") + tape_entries.append(cast(dict[str, object], entry)) + return CheckpointSnapshot( + meta=_meta_from_raw(_required_dict(row, "meta")), + tape_entries=tuple(tape_entries), + plugin_states=plugin_states, + extra=extra, + ) diff --git a/coding-agent/src/coding_agent/agent.toml b/coding-agent/src/coding_agent/agent.toml index a3b0a10fb..30bce483a 100644 --- a/coding-agent/src/coding_agent/agent.toml +++ b/coding-agent/src/coding_agent/agent.toml @@ -23,6 +23,8 @@ enabled = [ [storage] tape_backend = "jsonl" +http_session_backend = "memory" +checkpoint_backend = "fs" doc_backend = "lancedb" [storage.paths] diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 0033ce22f..0375994c3 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -15,6 +15,8 @@ from fastapi.middleware.cors import CORSMiddleware from sse_starlette.sse import EventSourceResponse +from agentkit.config.loader import load_config as load_agent_toml +from agentkit.errors import ConfigError from coding_agent.approval import ApprovalPolicy from coding_agent.ui.session_manager import Session, SessionManager from coding_agent.ui.schemas import ( @@ -58,8 +60,31 @@ SESSION_IDLE_TIMEOUT_MINUTES = 30 +def _load_storage_config() -> dict[str, Any]: + config_path = Path(__file__).resolve().parent.parent / "agent.toml" + try: + return cast( + dict[str, Any], load_agent_toml(config_path).extra.get("storage", {}) + ) + except (ConfigError, OSError) as exc: + if isinstance(exc, ConfigError): + detail = exc.args[0] if exc.args and isinstance(exc.args[0], str) else "" + if not detail.startswith("config file not found:"): + raise + logger.warning( + "Unable to load storage config from %s; using defaults", + config_path, + exc_info=True, + ) + return {} + + +def _build_session_manager() -> SessionManager: + return SessionManager(storage_config=_load_storage_config()) + + # Global session manager -session_manager = SessionManager() +session_manager = _build_session_manager() def _key_error_detail(exc: KeyError) -> str: @@ -85,8 +110,18 @@ async def lifespan(app: FastAPI): pass # Close all sessions - for session_id in list(session_manager.list_sessions()): - await session_manager.close_session(session_id) + try: + for session_id in await session_manager.list_sessions_async(): + try: + await session_manager.shutdown_session_runtime(session_id) + except Exception: + logger.warning( + "Failed to shut down runtime for session %s during server shutdown", + session_id, + exc_info=True, + ) + finally: + await session_manager.close() logger.info("HTTP server shut down") @@ -321,7 +356,50 @@ def _wire_message_to_event(msg: WireMessage) -> dict[str, str]: async def _broadcast_event(session: Session, event: dict[str, str]) -> None: """Broadcast event to all connected clients.""" - await session_manager.broadcast_event(session.id, event) + active_queues: list[asyncio.Queue[dict[str, str]]] = [] + full_pruned_count = 0 + failed_pruned_count = 0 + + for queue in session.event_queues: + try: + queue.put_nowait(event) + active_queues.append(queue) + except asyncio.QueueFull: + full_pruned_count += 1 + except Exception: + failed_pruned_count += 1 + logger.debug("Dropping closed event queue", exc_info=True) + + session.event_queues = active_queues + + if full_pruned_count: + logger.info( + "Pruned %d full event queue(s) for session %s", + full_pruned_count, + session.id, + ) + if failed_pruned_count: + logger.info( + "Pruned %d failed event queue(s) for session %s", + failed_pruned_count, + session.id, + ) + + +async def _cleanup_event_queue_on_disconnect( + session_id: str, + queue: asyncio.Queue[dict[str, str]], +) -> None: + try: + await asyncio.shield( + session_manager.remove_event_queue_async(session_id, queue) + ) + except KeyError: + logger.debug( + "Event queue cleanup skipped for already-removed session %s", + session_id, + exc_info=True, + ) async def _cleanup_idle_sessions() -> None: @@ -364,7 +442,9 @@ async def stream_wire_messages(wire: LocalWire) -> AsyncIterator[dict[str, str]] @limiter.limit(RateLimits.HEALTH) async def liveness_check(request: Request) -> HealthResponse: return HealthResponse( - status="healthy", sessions=len(session_manager.list_sessions()), version="2.0.0" + status="healthy", + sessions=await session_manager.count_sessions_async(), + version="2.0.0", ) @@ -372,7 +452,7 @@ async def liveness_check(request: Request) -> HealthResponse: @limiter.limit(RateLimits.HEALTH) async def readiness_check(request: Request, response: Response) -> ReadinessResponse: try: - session_store_ok = bool(session_manager._store.check_health()) + session_store_ok = bool(await session_manager.check_health_async()) except Exception: logger.exception("Session store readiness check failed") session_store_ok = False @@ -446,10 +526,10 @@ async def send_prompt( if not prompt_text: raise HTTPException(status_code=422, detail="Prompt is required") - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): raise HTTPException(status_code=404, detail="Session not found") - session = session_manager.get_session(session_id) + session = await session_manager.get_session_async(session_id) if session.turn_in_progress or (session.task and not session.task.done()): raise HTTPException(status_code=409, detail="Turn already in progress") @@ -532,10 +612,10 @@ async def approve_request( if is_approved is None: raise HTTPException(status_code=422, detail="approved is required") - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): raise HTTPException(status_code=404, detail="Session not found") - session = session_manager.get_session(session_id) + session = await session_manager.get_session_async(session_id) if session.approval_store.get_request(req_id) is None: raise HTTPException(status_code=400, detail="No pending approval request") @@ -567,12 +647,12 @@ async def get_events( api_key: str | None = Depends(verify_api_key), ) -> EventSourceResponse: """Persistent SSE event stream (fan-out supported).""" - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): raise HTTPException(status_code=404, detail="Session not found") - session = session_manager.get_session(session_id) + session = await session_manager.get_session_async(session_id) queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=100) - session_manager.add_event_queue(session_id, queue) + await session_manager.add_event_queue_async(session_id, queue) async def event_generator() -> AsyncIterator[dict[str, str]]: """Generate events from queue.""" @@ -581,13 +661,25 @@ async def event_generator() -> AsyncIterator[dict[str, str]]: try: event = await asyncio.wait_for(queue.get(), timeout=30.0) yield event + if event.get("event") == "SessionClosed": + break except asyncio.TimeoutError: + if not await session_manager.has_session_async(session_id): + break + try: + if not await session_manager.has_event_queue_async( + session_id, queue + ): + break + except KeyError: + break # Send keepalive yield {"event": "ping", "data": ""} except asyncio.CancelledError: # Client disconnected - session_manager.remove_event_queue(session_id, queue) raise + finally: + await _cleanup_event_queue_on_disconnect(session_id, queue) return EventSourceResponse(event_generator()) @@ -600,8 +692,8 @@ async def get_session( api_key: str | None = Depends(verify_api_key), ) -> dict[str, Any]: """Get session state.""" - if session_manager.has_session(session_id): - return session_manager.get_session_info(session_id) + if await session_manager.has_session_async(session_id): + return await session_manager.get_session_info_async(session_id) raise HTTPException(status_code=404, detail="Session not found") @@ -617,7 +709,7 @@ async def capture_checkpoint( body: CheckpointCaptureRequest | None = None, api_key: str | None = Depends(verify_api_key), ) -> CheckpointMetadataResponse: - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): raise HTTPException(status_code=404, detail="Session not found") try: @@ -651,7 +743,7 @@ async def list_checkpoints( session_id: str, api_key: str | None = Depends(verify_api_key), ) -> CheckpointListResponse: - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): raise HTTPException(status_code=404, detail="Session not found") checkpoints = await session_manager.list_checkpoints(session_id) @@ -682,7 +774,7 @@ async def restore_checkpoint( checkpoint_id: str, api_key: str | None = Depends(verify_api_key), ) -> CheckpointRestoreResponse: - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): raise HTTPException(status_code=404, detail="Session not found") try: @@ -709,20 +801,20 @@ async def close_session( api_key: str | None = Depends(verify_api_key), ) -> CloseSessionResponse: """Close session and release resources.""" - if not session_manager.has_session(session_id): - raise HTTPException(status_code=404, detail="Session not found") + try: + session = await session_manager.get_session_async(session_id) + await session_manager.close_session(session_id) + except KeyError as exc: + raise HTTPException(status_code=404, detail=_key_error_detail(exc)) from exc + except Exception as exc: + logger.exception("Unexpected error while closing session %s", session_id) + raise HTTPException(status_code=500, detail="Internal server error") from exc - session = session_manager.get_session(session_id) await _broadcast_event( session, {"event": "SessionClosed", "data": json.dumps({"session_id": session_id})}, ) - try: - await session_manager.close_session(session_id) - except Exception as e: - logger.exception(f"Error closing session in manager: {e}") - logger.info(f"Closed session: {session_id}") return CloseSessionResponse(status="closed", session_id=session_id) @@ -737,7 +829,7 @@ async def wait_for_approval( It will block until the user responds via the /approve endpoint or the timeout expires. """ - if not session_manager.has_session(session_id): + if not await session_manager.has_session_async(session_id): return ApprovalResponse( session_id=session_id, request_id=approval_req.request_id, @@ -745,7 +837,7 @@ async def wait_for_approval( feedback="Session not found", ) - session = session_manager.get_session(session_id) + session = await session_manager.get_session_async(session_id) event = _wire_message_to_event(approval_req) await _broadcast_event(session, event) response = await session_manager.wait_for_http_approval( diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 791eaac25..02d5969dd 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -9,21 +9,24 @@ import uuid from collections.abc import AsyncIterator from collections.abc import Callable +from functools import partial from dataclasses import dataclass, field from datetime import datetime from inspect import isawaitable from pathlib import Path -from typing import Any, Literal, cast +from typing import Any, Literal, TypeVar, cast +from agentkit.storage.checkpoint_fs import FSCheckpointStore +from agentkit.checkpoint.models import CheckpointMeta +from agentkit.checkpoint import CheckpointService +from agentkit.storage.protocols import CheckpointStore, TapeStore from coding_agent.adapter import PipelineAdapter from coding_agent.approval import ApprovalCoordinator, ApprovalPolicy from coding_agent.approval.store import ApprovalStore from coding_agent.core import config as core_config from coding_agent.plugins.storage import JSONLTapeStore -from agentkit.checkpoint import CheckpointService from coding_agent.providers.base import ChatProvider, ToolSchema from agentkit.providers.models import DoneEvent, TextEvent -from agentkit.storage.checkpoint_fs import FSCheckpointStore from agentkit.tape.tape import Tape from agentkit.tape.models import Entry from coding_agent.wire.local import LocalWire @@ -37,15 +40,15 @@ WireMessage, ) from coding_agent.ui.session_store import ( + AsyncPGSessionPool, SessionStore, create_session_store, ) -from agentkit.storage.protocols import CheckpointStore, TapeStore -from agentkit.checkpoint.models import CheckpointMeta logger = logging.getLogger(__name__) _CHECKPOINT_SESSION_CONFIG_KEY = "session_restart_config" +T = TypeVar("T") class MockProvider: @@ -273,6 +276,33 @@ def _checkpoint_session_config_from_extra( ) +def _load_pg_storage_types() -> tuple[Any, Any, Any]: + try: + pg_module = importlib.import_module("agentkit.storage.pg") + except ImportError as exc: + raise RuntimeError( + "PG backend is not available; ensure agentkit.storage.pg and its PostgreSQL " + "optional dependencies are installed before using tape_backend='pg' " + "(for example, install/include the PG extra or `asyncpg`)." + ) from exc + required_symbols = ("PGPool", "PGTapeStore", "PGCheckpointStore") + missing_symbols = [ + symbol for symbol in required_symbols if not hasattr(pg_module, symbol) + ] + if missing_symbols: + raise RuntimeError( + "PG backend is missing required exports from agentkit.storage.pg: " + f"{', '.join(missing_symbols)}. Ensure the installed PG backend package " + "version includes the PostgreSQL storage implementation and its optional " + "dependencies." + ) + return ( + getattr(pg_module, "PGPool"), + getattr(pg_module, "PGTapeStore"), + getattr(pg_module, "PGCheckpointStore"), + ) + + class _WireConsumer: def __init__( self, @@ -296,29 +326,101 @@ def __init__( self, store: SessionStore | None = None, *, + storage_config: dict[str, Any] | None = None, + pg_pool: object | None = None, tape_store: TapeStore | None = None, checkpoint_store: CheckpointStore | None = None, checkpoint_service: CheckpointService | None = None, create_agent_fn: Callable[..., tuple[Any, Any]] | None = None, ): - self._store = store or create_session_store() + self._storage_config = storage_config or {} + self._pg_pool = pg_pool + self._owns_pg_pool = False + self._store = store or self._create_http_session_store() self._session_cache: dict[str, Session] = {} self._approval_stores: dict[str, ApprovalStore] = {} self._lock = asyncio.Lock() + self._store_io_guard = asyncio.Lock() self._session_turn_locks: dict[str, asyncio.Lock] = {} data_dir = Path(os.environ.get("AGENT_DATA_DIR", "./data")) - self._tape_store = tape_store or JSONLTapeStore(data_dir / "tapes") - resolved_checkpoint_store = checkpoint_store or FSCheckpointStore( - data_dir / "checkpoints" + self._tape_store = tape_store or self._create_tape_store(data_dir) + resolved_checkpoint_store = checkpoint_store or self._create_checkpoint_store( + data_dir ) self._checkpoint_service = checkpoint_service or CheckpointService( resolved_checkpoint_store ) self._create_agent = create_agent_fn + def _get_pg_pool(self) -> AsyncPGSessionPool: + if self._pg_pool is not None: + return cast(AsyncPGSessionPool, self._pg_pool) + + PGPool, _, _ = _load_pg_storage_types() + + dsn_obj = self._storage_config.get("dsn") + if not isinstance(dsn_obj, str) or not dsn_obj.strip(): + raise RuntimeError("PG storage requires storage_config['dsn']") + dsn = dsn_obj.strip() + self._pg_pool = PGPool(dsn=dsn) + self._owns_pg_pool = True + return cast(AsyncPGSessionPool, self._pg_pool) + + def _create_http_session_store(self) -> SessionStore: + configured_backend = self._storage_config.get("http_session_backend") + tape_backend = ( + str(self._storage_config.get("tape_backend", "jsonl")).strip().lower() + ) + if configured_backend is None: + legacy_backend = self._storage_config.get("session_backend") + if ( + isinstance(legacy_backend, str) + and legacy_backend.strip().lower() == "pg" + ): + configured_backend = "pg" + elif tape_backend == "pg": + configured_backend = "pg" + + backend = ( + configured_backend.strip().lower() + if isinstance(configured_backend, str) + else None + ) + dsn = self._storage_config.get("dsn") + return create_session_store( + backend=backend, + dsn=dsn if isinstance(dsn, str) else None, + pg_pool=None, + ) + + def _create_tape_store(self, data_dir: Path) -> TapeStore: + backend = str(self._storage_config.get("tape_backend", "jsonl")).strip().lower() + if backend == "pg": + _, PGTapeStore, _ = _load_pg_storage_types() + return cast(TapeStore, PGTapeStore(pool=self._get_pg_pool())) + return JSONLTapeStore(data_dir / "tapes") + + def _create_checkpoint_store(self, data_dir: Path) -> CheckpointStore: + tape_backend = ( + str(self._storage_config.get("tape_backend", "jsonl")).strip().lower() + ) + default_backend = "pg" if tape_backend == "pg" else "fs" + backend = ( + str(self._storage_config.get("checkpoint_backend", default_backend)) + .strip() + .lower() + ) + if backend == "pg": + _, _, PGCheckpointStore = _load_pg_storage_types() + return cast(CheckpointStore, PGCheckpointStore(pool=self._get_pg_pool())) + return FSCheckpointStore(data_dir / "checkpoints") + async def _close_runtime(self, session: Session) -> None: adapter = session.runtime_adapter self._invalidate_runtime(session) + await self._close_runtime_adapter(adapter) + + async def _close_runtime_adapter(self, adapter: object | None) -> None: if adapter is None: return close = getattr(adapter, "close", None) @@ -328,12 +430,14 @@ async def _close_runtime(self, session: Session) -> None: await close_result def _close_runtime_sync_safe(self, session: Session) -> None: + adapter = session.runtime_adapter + self._invalidate_runtime(session) try: loop = asyncio.get_running_loop() except RuntimeError: - asyncio.run(self._close_runtime(session)) + asyncio.run(self._close_runtime_adapter(adapter)) return - loop.create_task(self._close_runtime(session)) + _ = loop.create_task(self._close_runtime_adapter(adapter)) def _create_agent_for_session(self, **kwargs: Any) -> tuple[Any, Any]: factory = self._create_agent @@ -348,6 +452,89 @@ def _turn_lock_for(self, session_id: str) -> asyncio.Lock: self._session_turn_locks[session_id] = lock return lock + async def _run_store_io(self, func: Callable[..., T], /, *args: object) -> T: + async with self._store_io_guard: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, partial(func, *args)) + + async def _persist_session_async(self, session: Session) -> None: + self._session_cache[session.id] = session + await self._run_store_io( + self._store.save, + session.id, + cast(dict[str, Any], session.to_store_data()), + ) + + async def get_session_async(self, session_id: str) -> Session: + session = self._session_cache.get(session_id) + if session is not None: + return session + loaded = await self._run_store_io(self._store.load, session_id) + if loaded is None: + raise KeyError(f"Session not found: {session_id}") + return self._hydrate_session( + Session.from_store_data(cast(dict[str, Any], loaded)) + ) + + async def has_session_async(self, session_id: str) -> bool: + if session_id in self._session_cache: + return True + return await self._run_store_io(self._store.load, session_id) is not None + + async def has_event_queue_async( + self, + session_id: str, + queue: asyncio.Queue[dict[str, Any]], + ) -> bool: + session = await self.get_session_async(session_id) + return queue in session.event_queues + + async def list_sessions_async(self) -> list[str]: + return await self._run_store_io(self._store.list_sessions) + + async def count_sessions_async(self) -> int: + return await self._run_store_io(self._store.count_sessions) + + async def get_session_info_async(self, session_id: str) -> dict[str, Any]: + session = await self.get_session_async(session_id) + return session.as_dict() + + async def add_event_queue_async( + self, + session_id: str, + queue: asyncio.Queue[dict[str, Any]], + ) -> None: + session = await self.get_session_async(session_id) + session.event_queues.append(queue) + + async def remove_event_queue_async( + self, + session_id: str, + queue: asyncio.Queue[dict[str, Any]], + ) -> None: + session = await self.get_session_async(session_id) + if queue in session.event_queues: + session.event_queues.remove(queue) + + async def check_health_async(self) -> bool: + return bool(await self._run_store_io(self._store.check_health)) + + async def _close_resource_async(self, resource: object) -> None: + close = getattr(resource, "close", None) + if not callable(close): + return + close_result = await self._run_store_io(close) + if isawaitable(close_result): + await close_result + + async def remove_session_async(self, session_id: str) -> None: + session = await self.get_session_async(session_id) + await self._close_runtime(session) + self._session_cache.pop(session_id, None) + await self._run_store_io(self._store.delete, session_id) + self._approval_stores.pop(session_id, None) + self._session_turn_locks.pop(session_id, None) + async def _restore_tape(self, tape_id: str | None) -> Tape | None: if tape_id is None: return None @@ -380,7 +567,7 @@ async def _request_approval(req: ApprovalRequest) -> ApprovalResponse: session.pending_approval = session.approval_coordinator.projection() session.approval_event.clear() session.approval_response = None - self._persist_session(session) + await self._persist_session_async(session) await session.wire.send(req) try: response = await session.approval_coordinator.wait_for_response( @@ -401,12 +588,12 @@ async def _request_approval(req: ApprovalRequest) -> ApprovalResponse: } session.approval_event.set() session.pending_approval = session.approval_coordinator.projection() - self._persist_session(session) + await self._persist_session_async(session) return response finally: session.pending_approval = session.approval_coordinator.projection() session.approval_response = None - self._persist_session(session) + await self._persist_session_async(session) return _WireConsumer(session.wire, _request_approval) @@ -495,7 +682,7 @@ async def _restore_checkpoint(self, session: Session, checkpoint_id: str) -> Non session.runtime_pipeline = pipeline session.runtime_ctx = ctx session.runtime_adapter = adapter - self._persist_session(session) + await self._persist_session_async(session) checkpoints = await self._checkpoint_service.list(ctx.tape.tape_id) for checkpoint_meta in checkpoints: @@ -580,7 +767,7 @@ async def create_session( ) async with self._lock: - self._persist_session(session) + await self._persist_session_async(session) logger.info(f"Created session: {session_id}") return session_id @@ -651,7 +838,6 @@ def add_event_queue( ) -> None: session = self.get_session(session_id) session.event_queues.append(queue) - self._persist_session(session) def remove_event_queue( self, @@ -661,7 +847,6 @@ def remove_event_queue( session = self.get_session(session_id) if queue in session.event_queues: session.event_queues.remove(queue) - self._persist_session(session) async def broadcast_event( self, @@ -706,7 +891,7 @@ async def close_session(self, session_id: str) -> None: KeyError: If session not found """ async with self._lock: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) # Cancel any running task if session.task and not session.task.done(): @@ -715,11 +900,41 @@ async def close_session(self, session_id: str) -> None: await asyncio.wait_for(session.task, timeout=5.0) except (asyncio.CancelledError, asyncio.TimeoutError): pass + if not session.task.done(): + raise RuntimeError( + f"Session task for {session_id} did not stop after cancellation" + ) - self.remove_session(session_id) + await self.remove_session_async(session_id) logger.info(f"Closed session: {session_id}") + async def shutdown_session_runtime(self, session_id: str) -> None: + """Release runtime resources without deleting persisted session metadata.""" + async with self._lock: + session = await self.get_session_async(session_id) + + if session.task and not session.task.done(): + session.task.cancel() + try: + await asyncio.wait_for(session.task, timeout=5.0) + except (asyncio.CancelledError, asyncio.TimeoutError): + pass + if not session.task.done(): + raise RuntimeError( + f"Session task for {session_id} did not stop after cancellation" + ) + + await self._close_runtime(session) + session.task = None + session.turn_in_progress = False + await self._persist_session_async(session) + + async def close(self) -> None: + await self._close_resource_async(self._store) + if self._owns_pg_pool: + await self._close_resource_async(self._pg_pool) + async def run_agent( self, session_id: str, @@ -730,10 +945,10 @@ async def run_agent( raise RuntimeError("turn already in progress") async with turn_lock: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) session.last_activity = datetime.now() session.turn_in_progress = True - self._persist_session(session) + await self._persist_session_async(session) try: approval_mode_map = { @@ -762,7 +977,7 @@ async def run_agent( tape=await self._restore_tape(session.tape_id), ) session.tape_id = ctx.tape.tape_id - self._persist_session(session) + await self._persist_session_async(session) ctx.config["wire_consumer"] = None ctx.config["agent_id"] = "" @@ -783,7 +998,7 @@ async def run_agent( ctx.config["wire_consumer"] = consumer await adapter.run_turn(prompt) session.tape_id = ctx.tape.tape_id - self._persist_session(session) + await self._persist_session_async(session) except Exception as exc: await self._close_runtime(session) logger.exception("HTTP session turn failed") @@ -807,7 +1022,7 @@ async def run_agent( if session.task is None or session.task is not current_task: session.turn_in_progress = False session.last_activity = datetime.now() - self._persist_session(session) + await self._persist_session_async(session) async def submit_approval( self, @@ -833,7 +1048,7 @@ async def submit_approval( Raises: KeyError: If session not found """ - session = self.get_session(session_id) + session = await self.get_session_async(session_id) # Create approval response and submit to ApprovalStore response = ApprovalResponse( @@ -853,7 +1068,7 @@ async def submit_approval( "feedback": feedback, } session.approval_event.set() - self._persist_session(session) + await self._persist_session_async(session) logger.info(f"Approval submitted for session {session_id}: {approved}") else: logger.warning( @@ -868,7 +1083,7 @@ async def wait_for_http_approval( approval_req: ApprovalRequest, timeout_seconds: float, ) -> ApprovalResponse: - if not self.has_session(session_id): + if not await self.has_session_async(session_id): return ApprovalResponse( session_id=session_id, request_id=approval_req.request_id, @@ -876,7 +1091,7 @@ async def wait_for_http_approval( feedback="Session not found", ) - session = self.get_session(session_id) + session = await self.get_session_async(session_id) if not session.turn_in_progress: return ApprovalResponse( session_id=session_id, @@ -897,7 +1112,7 @@ async def wait_for_http_approval( session.pending_approval = session.approval_coordinator.projection() session.approval_event.clear() session.approval_response = None - self._persist_session(session) + await self._persist_session_async(session) try: response = await session.approval_coordinator.wait_for_response( @@ -910,7 +1125,7 @@ async def wait_for_http_approval( session.pending_approval = session.approval_coordinator.projection() session.approval_response = None _ = session.approval_event.set() - self._persist_session(session) + await self._persist_session_async(session) return ApprovalResponse( session_id=session_id, @@ -952,14 +1167,12 @@ async def cleanup_idle_sessions(self, max_idle_minutes: int = 30) -> list[str]: List of closed session IDs """ now = datetime.now() - closed = [] - - async with self._lock: - session_ids = list(self._store.list_sessions()) + closed: list[str] = [] + session_ids = await self.list_sessions_async() for session_id in session_ids: try: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) idle_time = now - session.last_activity if idle_time.total_seconds() > max_idle_minutes * 60: await self.close_session(session_id) @@ -974,7 +1187,7 @@ async def cleanup_idle_sessions(self, max_idle_minutes: int = 30) -> list[str]: return closed async def ensure_session_runtime(self, session_id: str) -> Any: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) if session.runtime_ctx is not None and session.runtime_adapter is not None: return session.runtime_ctx @@ -1009,7 +1222,7 @@ async def ensure_session_runtime(self, session_id: str) -> Any: session.runtime_ctx = ctx session.runtime_adapter = adapter session.tape_id = ctx.tape.tape_id - self._persist_session(session) + await self._persist_session_async(session) return ctx async def capture_checkpoint( @@ -1024,7 +1237,7 @@ async def capture_checkpoint( raise RuntimeError("turn already in progress") async with turn_lock: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) if session.turn_in_progress or (session.task and not session.task.done()): raise RuntimeError("turn already in progress") @@ -1041,11 +1254,11 @@ async def capture_checkpoint( ctx, label=label, extra=payload ) session.tape_id = ctx.tape.tape_id - self._persist_session(session) + await self._persist_session_async(session) return checkpoint async def list_checkpoints(self, session_id: str) -> list[CheckpointMeta]: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) if session.tape_id is None: return [] return await self._checkpoint_service.list(session.tape_id) @@ -1056,7 +1269,7 @@ async def restore_checkpoint(self, session_id: str, checkpoint_id: str) -> None: raise RuntimeError("turn already in progress") async with turn_lock: - session = self.get_session(session_id) + session = await self.get_session_async(session_id) if session.turn_in_progress or (session.task and not session.task.done()): raise RuntimeError("turn already in progress") await self._restore_checkpoint(session, checkpoint_id) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index c3031cb9d..267cf42c3 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -1,10 +1,16 @@ from __future__ import annotations +import asyncio +from concurrent.futures import TimeoutError as FutureTimeoutError import json import logging import os -from collections.abc import Callable -from typing import Protocol, cast +import threading +from collections.abc import Callable, Coroutine, Iterable +from typing import Final, Protocol, cast + +from agentkit.config.loader import ConfigError +from agentkit.storage.pg import PGPool from ..redaction import redact_sensitive_text, redact_url_credentials @@ -39,6 +45,8 @@ def save(self, session_id: str, data: SessionPayload) -> None: ... def load(self, session_id: str) -> SessionPayload | None: ... + def count_sessions(self) -> int: ... + def list_sessions(self) -> list[str]: ... def delete(self, session_id: str) -> None: ... @@ -46,6 +54,12 @@ def delete(self, session_id: str) -> None: ... def check_health(self) -> bool: ... +class AsyncPGSessionPool(Protocol): + async def get_pool(self) -> object: ... + + async def close(self) -> None: ... + + class InMemorySessionStore: def __init__(self) -> None: self._sessions: dict[str, SessionPayload] = {} @@ -56,6 +70,9 @@ def save(self, session_id: str, data: SessionPayload) -> None: def load(self, session_id: str) -> SessionPayload | None: return self._sessions.get(session_id) + def count_sessions(self) -> int: + return len(self._sessions) + def get(self, session_id: str) -> SessionPayload | None: return self.load(session_id) @@ -113,6 +130,15 @@ def list_sessions(self) -> list[str]: session_ids.append(member) return sorted(session_ids) + def count_sessions(self) -> int: + scard = getattr(self._client, "scard", None) + if callable(scard): + count = scard(self._index_key) + if not isinstance(count, int): + raise TypeError("redis scard must return an int") + return count + return len(cast(set[str | bytes], self._client.smembers(self._index_key))) + def delete(self, session_id: str) -> None: _ = self._client.delete(self._key_for(session_id)) _ = self._client.srem(self._index_key, session_id) @@ -124,12 +150,309 @@ def check_health(self) -> bool: return False +class PGSessionMetadataStore: + _CREATE_TABLE_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS agent_http_sessions ( + session_id TEXT PRIMARY KEY, + payload JSONB NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """ + _UPSERT_SQL: Final[str] = """ + INSERT INTO agent_http_sessions (session_id, payload) + VALUES ($1, $2::jsonb) + ON CONFLICT (session_id) + DO UPDATE SET payload = EXCLUDED.payload, updated_at = NOW() + """ + _SELECT_SQL: Final[str] = ( + "SELECT payload FROM agent_http_sessions WHERE session_id = $1" + ) + _LIST_SQL: Final[str] = ( + "SELECT session_id FROM agent_http_sessions ORDER BY session_id" + ) + _COUNT_SQL: Final[str] = "SELECT COUNT(*) AS session_count FROM agent_http_sessions" + _DELETE_SQL: Final[str] = "DELETE FROM agent_http_sessions WHERE session_id = $1" + _LOOP_READY_TIMEOUT_SECONDS: Final[float] = 5.0 + _SYNC_OPERATION_TIMEOUT_SECONDS: Final[float] = 30.0 + + def __init__( + self, + *, + pool: AsyncPGSessionPool, + owns_pool: bool = False, + ) -> None: + self._pool: AsyncPGSessionPool = pool + self._owns_pool = owns_pool + self._schema_ready: bool = False + self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop() + self._loop_ready: threading.Event = threading.Event() + self._loop_thread: threading.Thread = threading.Thread( + target=self._run_loop, + name="pg-session-metadata-store", + daemon=True, + ) + self._loop_thread.start() + ready = self._loop_ready.wait(self._LOOP_READY_TIMEOUT_SECONDS) + if not ready: + self._handle_loop_start_failure() + raise RuntimeError("postgres session metadata loop thread failed to start") + + def _handle_loop_start_failure(self) -> None: + try: + _ = self._loop.call_soon_threadsafe(self._loop.stop) + except RuntimeError: + pass + self._loop_thread.join(timeout=0.1) + if not self._loop_thread.is_alive(): + self._loop.close() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.call_soon(self._loop_ready.set) + self._loop.run_forever() + + def _run_sync(self, operation: Coroutine[object, object, object]) -> object: + try: + future = asyncio.run_coroutine_threadsafe(operation, self._loop) + except RuntimeError as exc: + operation.close() + raise RuntimeError( + "postgres session metadata loop is not available; " + "the store may have been closed or failed to start" + ) from exc + try: + return future.result(timeout=self._SYNC_OPERATION_TIMEOUT_SECONDS) + except FutureTimeoutError as exc: + cancellation_complete = threading.Event() + future.add_done_callback(lambda _future: cancellation_complete.set()) + future.cancel() + try: + _ = future.result(timeout=0.1) + except (FutureTimeoutError, asyncio.CancelledError, Exception): + pass + cancellation_complete.wait(timeout=0.1) + raise TimeoutError( + "postgres session metadata operation timed out " + f"after {self._SYNC_OPERATION_TIMEOUT_SECONDS} seconds" + ) from exc + + async def _ensure_schema(self) -> object: + asyncpg_pool = await self._pool.get_pool() + if not self._schema_ready: + execute = getattr(asyncpg_pool, "execute", None) + if not callable(execute): + raise TypeError("postgres session metadata pool must expose execute") + _ = await cast(Callable[..., Coroutine[object, object, object]], execute)( + self._CREATE_TABLE_SQL + ) + self._schema_ready = True + return asyncpg_pool + + def save(self, session_id: str, data: SessionPayload) -> None: + async def _save() -> None: + pool = await self._ensure_schema() + execute = getattr(pool, "execute", None) + if not callable(execute): + raise TypeError("postgres session metadata pool must expose execute") + _ = await cast(Callable[..., Coroutine[object, object, object]], execute)( + self._UPSERT_SQL, + session_id, + data, + ) + + _ = self._run_sync(_save()) + + def load(self, session_id: str) -> SessionPayload | None: + async def _load() -> SessionPayload | None: + pool = await self._ensure_schema() + fetchrow = getattr(pool, "fetchrow", None) + if not callable(fetchrow): + raise TypeError("postgres session metadata pool must expose fetchrow") + row_obj = await cast( + Callable[..., Coroutine[object, object, object]], fetchrow + )( + self._SELECT_SQL, + session_id, + ) + if row_obj is None: + return None + row_dict = _coerce_row_dict( + row=row_obj, context="postgres session metadata row" + ) + payload = row_dict.get("payload") + if payload is None: + return None + if not isinstance(payload, dict): + raise TypeError( + "postgres session metadata payload must decode to a dict" + ) + return cast(SessionPayload, payload) + + result = self._run_sync(_load()) + if result is None: + return None + if not isinstance(result, dict): + raise TypeError("postgres session metadata payload must be a JSON object") + return cast(SessionPayload, result) + + def list_sessions(self) -> list[str]: + async def _list_sessions() -> list[str]: + pool = await self._ensure_schema() + fetch = getattr(pool, "fetch", None) + if not callable(fetch): + raise TypeError("postgres session metadata pool must expose fetch") + rows_obj = await cast( + Callable[..., Coroutine[object, object, object]], fetch + )(self._LIST_SQL) + if not isinstance(rows_obj, list): + raise TypeError("postgres session metadata list result must be a list") + session_ids: list[str] = [] + for row in cast(list[object], rows_obj): + row_dict = _coerce_row_dict( + row=row, context="postgres session metadata list row" + ) + session_id = row_dict.get("session_id") + if not isinstance(session_id, str): + raise TypeError( + "postgres session metadata row must include string session_id" + ) + session_ids.append(session_id) + return session_ids + + result = self._run_sync(_list_sessions()) + if not isinstance(result, list): + raise TypeError("postgres session metadata list result must be a list") + return cast(list[str], result) + + def count_sessions(self) -> int: + async def _count_sessions() -> int: + pool = await self._ensure_schema() + fetchrow = getattr(pool, "fetchrow", None) + if not callable(fetchrow): + raise TypeError("postgres session metadata pool must expose fetchrow") + row_obj = await cast( + Callable[..., Coroutine[object, object, object]], fetchrow + )(self._COUNT_SQL) + if row_obj is None: + raise TypeError( + "postgres session metadata count result must not be None" + ) + row_dict = _coerce_row_dict( + row=row_obj, context="postgres session metadata count row" + ) + session_count = row_dict.get("session_count") + if not isinstance(session_count, int): + raise TypeError( + "postgres session metadata count row must include int session_count" + ) + return session_count + + result = self._run_sync(_count_sessions()) + if not isinstance(result, int): + raise TypeError("postgres session metadata count result must be an int") + return result + + def delete(self, session_id: str) -> None: + async def _delete() -> None: + pool = await self._ensure_schema() + execute = getattr(pool, "execute", None) + if not callable(execute): + raise TypeError("postgres session metadata pool must expose execute") + _ = await cast(Callable[..., Coroutine[object, object, object]], execute)( + self._DELETE_SQL, + session_id, + ) + + _ = self._run_sync(_delete()) + + def check_health(self) -> bool: + try: + + async def _check_health() -> bool: + pool = await self._ensure_schema() + fetchrow = getattr(pool, "fetchrow", None) + if not callable(fetchrow): + raise TypeError( + "postgres session metadata pool must expose fetchrow" + ) + row_obj = await cast( + Callable[..., Coroutine[object, object, object]], fetchrow + )("SELECT 1") + return row_obj is not None + + return bool(self._run_sync(_check_health())) + except Exception: + return False + + def close(self) -> None: + async def _close_pool() -> None: + if self._owns_pool: + await self._pool.close() + + try: + try: + _ = self._run_sync(_close_pool()) + except Exception: + logger.warning( + "Failed to close postgres session metadata pool during shutdown", + exc_info=True, + ) + finally: + try: + _ = self._loop.call_soon_threadsafe(self._loop.stop) + except RuntimeError: + logger.warning( + "Postgres session metadata loop was already closed while " + "scheduling stop during shutdown" + ) + self._loop_thread.join(timeout=5) + if self._loop_thread.is_alive(): + logger.warning( + "Timed out waiting for postgres session metadata loop thread to stop; skipping event loop close" + ) + else: + try: + self._loop.close() + except RuntimeError: + logger.warning( + "Postgres session metadata loop was already closed during shutdown" + ) + + def create_session_store( *, + backend: str | None = None, + dsn: str | None = None, redis_url: str | None = None, redis_client_factory: Callable[[str], RedisClient] | None = None, + pg_pool: AsyncPGSessionPool | None = None, ) -> SessionStore: + resolved_backend = ( + (backend or os.environ.get("AGENT_SESSION_BACKEND") or "").strip().lower() + ) + resolved_dsn = ( + dsn + or os.environ.get("AGENT_SESSION_PG_DSN") + or os.environ.get("AGENT_STORAGE_DSN") + ) + normalized_dsn = resolved_dsn.strip() if isinstance(resolved_dsn, str) else None resolved_redis_url = redis_url or os.environ.get("AGENT_SESSION_REDIS_URL") + if resolved_backend == "pg": + owns_pool = False + if pg_pool is None and not normalized_dsn: + raise ValueError("PG session store requires dsn or pg_pool") + if pg_pool is None: + assert normalized_dsn is not None + pg_pool = cast(AsyncPGSessionPool, PGPool(dsn=normalized_dsn)) + owns_pool = True + return PGSessionMetadataStore(pool=pg_pool, owns_pool=owns_pool) + + if resolved_backend == "memory": + return InMemorySessionStore() + + if resolved_backend not in {"", "redis"}: + raise ValueError(f"unsupported session store backend: {resolved_backend}") + if not resolved_redis_url: return InMemorySessionStore() @@ -140,6 +463,8 @@ def create_session_store( ) return RedisSessionStore(client=client, redis_url=resolved_redis_url) except Exception as exc: + if isinstance(exc, ConfigError): + raise safe_exc = redact_sensitive_text(str(exc)) logger.warning( "Redis session store unavailable at %s; falling back to in-memory store: %s", @@ -149,6 +474,21 @@ def create_session_store( return InMemorySessionStore() +def _coerce_row_dict(*, row: object, context: str) -> dict[str, object]: + if isinstance(row, dict): + return row + if not isinstance(row, Iterable): + raise TypeError(f"{context} must be convertible to a dict") + try: + row_items = cast(Iterable[tuple[object, object]], row) + row_dict_obj = dict(row_items) + except Exception as exc: + raise TypeError(f"{context} must be convertible to a dict") from exc + if not isinstance(row_dict_obj, dict): + raise TypeError(f"{context} must decode to a dict") + return cast(dict[str, object], row_dict_obj) + + def _create_redis_client( *, redis_url: str, diff --git a/coding-agent/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py index cb8cde892..e17f79a30 100644 --- a/coding-agent/tests/agentkit/storage/test_pg.py +++ b/coding-agent/tests/agentkit/storage/test_pg.py @@ -6,14 +6,24 @@ import pytest from _pytest.monkeypatch import MonkeyPatch -from agentkit.storage.protocols import SessionStore, TapeStore -from agentkit.storage.pg import PGPool, PGSessionLock, PGSessionStore, PGTapeStore +from datetime import UTC, datetime + +from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot +from agentkit.storage.protocols import CheckpointStore, SessionStore, TapeStore +from agentkit.storage.pg import ( + PGCheckpointStore, + PGPool, + PGSessionLock, + PGSessionStore, + PGTapeStore, +) class FakePool: def __init__(self) -> None: self.sessions: dict[str, dict[str, object]] = {} self.tapes: dict[str, list[dict[str, int | dict[str, object]]]] = {} + self.checkpoints: dict[str, dict[str, object]] = {} self.closed: bool = False self.executed: list[tuple[str, tuple[object, ...]]] = [] @@ -23,9 +33,12 @@ async def execute(self, query: str, *args: object) -> str: session_id, payload = args if not isinstance(session_id, str): raise TypeError("session_id must be a string") - if not isinstance(payload, str): - raise TypeError("payload must be encoded json") - payload_obj = json.loads(payload) + if isinstance(payload, str): + payload_obj = json.loads(payload) + elif isinstance(payload, dict): + payload_obj = payload + else: + raise TypeError("payload must be encoded json or dict") if not isinstance(payload_obj, dict): raise TypeError("payload must decode to a dict") self.sessions[session_id] = cast(dict[str, object], payload_obj) @@ -37,33 +50,23 @@ async def execute(self, query: str, *args: object) -> str: _ = self.sessions.pop(session_id, None) return "DELETE 1" if "INSERT INTO agent_tapes" in query: - tape_id, second_arg, third_arg = args + tape_id, payload_values = args if not isinstance(tape_id, str): raise TypeError("tape_id must be a string") rows = self.tapes.setdefault(tape_id, []) - if isinstance(second_arg, int): - seq_values = [second_arg] - elif isinstance(second_arg, list) and all( - isinstance(item, int) for item in second_arg - ): - seq_values = second_arg - else: - raise TypeError("seq must be an int or list[int]") - - if isinstance(third_arg, str): - payload_values = [third_arg] - elif isinstance(third_arg, list) and all( - isinstance(item, str) for item in third_arg + if not ( + isinstance(payload_values, list) + and all(isinstance(item, (str, dict)) for item in payload_values) ): - payload_values = third_arg - else: - raise TypeError("payload must be encoded json or list[str]") - - if len(seq_values) != len(payload_values): - raise TypeError("seq and payload batch lengths must match") - - for seq, payload in zip(seq_values, payload_values, strict=True): - payload_obj = json.loads(payload) + raise TypeError("payload must be encoded json/dict or list thereof") + max_seq = max( + (row["seq"] for row in rows if isinstance(row["seq"], int)), default=-1 + ) + for offset, payload in enumerate(payload_values, start=1): + seq = max_seq + offset + payload_obj = ( + json.loads(payload) if isinstance(payload, str) else payload + ) if not isinstance(payload_obj, dict): raise TypeError("payload must decode to a dict") rows.append({"seq": seq, "entry": cast(dict[str, object], payload_obj)}) @@ -71,6 +74,36 @@ async def execute(self, query: str, *args: object) -> str: key=lambda item: item["seq"] if isinstance(item["seq"], int) else -1 ) return "INSERT 0 1" + if "INSERT INTO agent_checkpoints" in query: + checkpoint_id, tape_id, meta_json, entries_json, state_json, extra_json = ( + args + ) + if not isinstance(checkpoint_id, str): + raise TypeError("checkpoint_id must be a string") + if not isinstance(tape_id, str): + raise TypeError("tape_id must be a string") + if not all( + isinstance(item, (str, dict, list)) + for item in (meta_json, entries_json, state_json, extra_json) + ): + raise TypeError("checkpoint payloads must be json-compatible values") + self.checkpoints[checkpoint_id] = { + "checkpoint_id": checkpoint_id, + "tape_id": tape_id, + "meta": json.loads(meta_json) + if isinstance(meta_json, str) + else meta_json, + "entries": json.loads(entries_json) + if isinstance(entries_json, str) + else entries_json, + "plugin_states": json.loads(state_json) + if isinstance(state_json, str) + else state_json, + "extra": json.loads(extra_json) + if isinstance(extra_json, str) + else extra_json, + } + return "INSERT 0 1" if query.strip() == "SELECT 1": return "SELECT 1" if "DELETE FROM agent_tapes WHERE tape_id = $1 AND seq >= $2" in query: @@ -88,24 +121,36 @@ async def execute(self, query: str, *args: object) -> str: return "CREATE TABLE" if "CREATE TABLE IF NOT EXISTS agent_tapes" in query: return "CREATE TABLE" + if "CREATE TABLE IF NOT EXISTS agent_checkpoints" in query: + return "CREATE TABLE" + if "DELETE FROM agent_checkpoints" in query: + (checkpoint_id,) = args + if not isinstance(checkpoint_id, str): + raise TypeError("checkpoint_id must be a string") + _ = self.checkpoints.pop(checkpoint_id, None) + return "DELETE 1" raise AssertionError(f"unexpected execute query: {query}") async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None: self.executed.append((query, args)) if "SELECT payload FROM agent_sessions" not in query: - if "SELECT MAX(seq) AS max_seq FROM agent_tapes" not in query: - raise AssertionError(f"unexpected fetchrow query: {query}") - (tape_id,) = args - if not isinstance(tape_id, str): - raise TypeError("tape_id must be a string") - rows = self.tapes.get(tape_id, []) - if not rows: - return {"max_seq": None} - return { - "max_seq": max( - row["seq"] if isinstance(row["seq"], int) else -1 for row in rows - ) - } + if ( + "SELECT meta, entries, plugin_states, extra FROM agent_checkpoints" + in query + ): + (checkpoint_id,) = args + if not isinstance(checkpoint_id, str): + raise TypeError("checkpoint_id must be a string") + payload = self.checkpoints.get(checkpoint_id) + if payload is None: + return None + return { + "meta": payload["meta"], + "entries": payload["entries"], + "plugin_states": payload["plugin_states"], + "extra": payload["extra"], + } + raise AssertionError(f"unexpected fetchrow query: {query}") (session_id,) = args if not isinstance(session_id, str): raise TypeError("session_id must be a string") @@ -135,6 +180,28 @@ async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: ] if "SELECT DISTINCT tape_id FROM agent_tapes" in query: return [{"tape_id": tape_id} for tape_id in sorted(self.tapes.keys())] + if "SELECT meta FROM agent_checkpoints" in query: + (tape_id,) = args + if not isinstance(tape_id, str): + raise TypeError("tape_id must be a string") + rows = [ + {"meta": payload["meta"]} + for payload in self.checkpoints.values() + if payload["tape_id"] == tape_id + ] + if "::timestamptz" in query: + rows.sort( + key=lambda row: datetime.fromisoformat( + str(cast(dict[str, object], row["meta"])["created_at"]) + ) + ) + else: + rows.sort( + key=lambda row: str( + cast(dict[str, object], row["meta"])["created_at"] + ) + ) + return rows raise AssertionError(f"unexpected fetch query: {query}") async def close(self) -> None: @@ -148,6 +215,28 @@ async def release(self, connection: object) -> None: raise AssertionError("unexpected connection released") +class FakeInitConnection: + def __init__(self) -> None: + self.codec_calls: list[dict[str, object]] = [] + + async def set_type_codec( + self, + typename: str, + *, + encoder: object, + decoder: object, + schema: str, + ) -> None: + self.codec_calls.append( + { + "typename": typename, + "encoder": encoder, + "decoder": decoder, + "schema": schema, + } + ) + + class TestPGPool: @pytest.mark.asyncio async def test_creates_pool_lazily_and_reuses_instance(self): @@ -227,6 +316,32 @@ def fake_import_module(name: str): with pytest.raises(ImportError, match="asyncpg is required"): await pool.get_pool() + @pytest.mark.asyncio + async def test_registers_json_and_jsonb_codecs_on_pool_connections(self): + seen_connections: list[FakeInitConnection] = [] + + async def fake_pool_factory(**kwargs: object) -> FakePool: + init = kwargs.get("init") + if not callable(init): + raise AssertionError("expected asyncpg init callback") + conn = FakeInitConnection() + seen_connections.append(conn) + await init(conn) + return FakePool() + + pool = PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + + _ = await pool.get_pool() + + assert len(seen_connections) == 1 + codec_names = [ + cast(str, call["typename"]) for call in seen_connections[0].codec_calls + ] + assert codec_names == ["json", "jsonb"] + assert all( + call["schema"] == "pg_catalog" for call in seen_connections[0].codec_calls + ) + @pytest.mark.asyncio async def test_acquire_delegates_to_underlying_pool(self): fake_pool = FakePool() @@ -253,6 +368,60 @@ async def fake_pool_factory(**_: object) -> FakePool: await pool.release(fake_pool) +@pytest.mark.asyncio +async def test_pg_checkpoint_store_orders_list_by_created_at_timestamp() -> None: + async def fake_pool_factory(**_: object) -> FakePool: + return fake_pool + + fake_pool = FakePool() + store = PGCheckpointStore( + pool=PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + ) + + older = CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="cp-older", + tape_id="tape-1", + session_id="session-1", + entry_count=1, + window_start=0, + created_at=datetime.fromisoformat("2026-04-18T12:30:00+02:00"), + label="older", + ), + tape_entries=tuple(), + plugin_states={}, + extra={}, + ) + newer = CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id="cp-newer", + tape_id="tape-1", + session_id="session-1", + entry_count=2, + window_start=0, + created_at=datetime.fromisoformat("2026-04-18T11:45:00+00:00"), + label="newer", + ), + tape_entries=tuple(), + plugin_states={}, + extra={}, + ) + + await store.save(newer) + await store.save(older) + + metas = await store.list_by_tape("tape-1") + + assert [meta.checkpoint_id for meta in metas] == ["cp-older", "cp-newer"] + list_queries = [ + query + for query, _args in fake_pool.executed + if "SELECT meta FROM agent_checkpoints" in query + ] + assert list_queries + assert "::timestamptz" in list_queries[-1] + + class TestPGSessionStore: @pytest.fixture def fake_pool(self) -> FakePool: @@ -324,6 +493,20 @@ async def test_schema_created_once( ] assert len(schema_calls) == 1 + @pytest.mark.asyncio + async def test_save_session_passes_python_object_to_codec_enabled_pool( + self, store: PGSessionStore, fake_pool: FakePool + ): + await store.save_session("ses-1", {"model": "gpt-4.1"}) + + insert_calls = [ + args + for query, args in fake_pool.executed + if "INSERT INTO agent_sessions" in query + ] + assert len(insert_calls) == 1 + assert isinstance(insert_calls[0][1], dict) + class TestPGTapeStore: @pytest.fixture @@ -407,6 +590,22 @@ async def test_save_batches_insert_round_trip( assert len(insert_calls) == 1 + @pytest.mark.asyncio + async def test_save_uses_single_atomic_append_statement( + self, store: PGTapeStore, fake_pool: FakePool + ): + await store.save( + "tape-atomic", + [{"kind": "message", "payload": {"role": "user", "content": "x"}}], + ) + + tape_queries = [ + query for query, _args in fake_pool.executed if "agent_tapes" in query + ] + + assert any("pg_advisory_xact_lock" in query for query in tape_queries) + assert all("SELECT MAX(seq) AS max_seq" not in query for query in tape_queries) + @pytest.mark.asyncio async def test_load_empty_tape(self, store: PGTapeStore): result = await store.load("nonexistent") @@ -451,6 +650,208 @@ async def test_truncate_rejects_negative_keep(self, store: PGTapeStore): await store.truncate("tape-truncate", -1) +class TestPGCheckpointStore: + @pytest.fixture + def fake_pool(self) -> FakePool: + return FakePool() + + @pytest.fixture + def pool(self, fake_pool: FakePool) -> PGPool: + async def fake_pool_factory(**_: object) -> FakePool: + return fake_pool + + return PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + + @pytest.fixture + def store(self, pool: PGPool) -> PGCheckpointStore: + return PGCheckpointStore(pool=pool) + + def _snapshot( + self, checkpoint_id: str, tape_id: str, *, created_at: datetime + ) -> CheckpointSnapshot: + return CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id=checkpoint_id, + tape_id=tape_id, + session_id="session-1", + entry_count=2, + window_start=1, + created_at=created_at, + label=checkpoint_id, + ), + tape_entries=( + { + "id": "e-1", + "kind": "message", + "payload": {"content": "a"}, + "timestamp": 1.0, + }, + { + "id": "e-2", + "kind": "message", + "payload": {"content": "b"}, + "timestamp": 2.0, + }, + ), + plugin_states={"topic": {"current": checkpoint_id}}, + extra={"source": checkpoint_id}, + ) + + def test_satisfies_protocol(self, store: PGCheckpointStore): + assert isinstance(store, CheckpointStore) + + @pytest.mark.asyncio + async def test_pg_checkpoint_store_round_trip_snapshot( + self, store: PGCheckpointStore + ): + snapshot = self._snapshot( + "cp-roundtrip", + "tape-roundtrip", + created_at=datetime(2026, 4, 17, tzinfo=UTC), + ) + + await store.save(snapshot) + + loaded = await store.load("cp-roundtrip") + + assert loaded == snapshot + + @pytest.mark.asyncio + async def test_pg_checkpoint_store_overwrites_existing_checkpoint_id( + self, store: PGCheckpointStore + ): + first_snapshot = self._snapshot( + "cp-retry", + "tape-first", + created_at=datetime(2026, 4, 17, tzinfo=UTC), + ) + second_snapshot = self._snapshot( + "cp-retry", + "tape-second", + created_at=datetime(2026, 4, 17, 1, tzinfo=UTC), + ) + second_snapshot = CheckpointSnapshot( + meta=CheckpointMeta( + checkpoint_id=second_snapshot.meta.checkpoint_id, + tape_id=second_snapshot.meta.tape_id, + session_id="session-2", + entry_count=4, + window_start=2, + created_at=second_snapshot.meta.created_at, + label="retry-save", + ), + tape_entries=( + { + "id": "e-3", + "kind": "message", + "payload": {"content": "retry"}, + "timestamp": 3.0, + }, + ), + plugin_states={"topic": {"current": "retry"}}, + extra={"source": "retry"}, + ) + + await store.save(first_snapshot) + await store.save(second_snapshot) + + loaded = await store.load("cp-retry") + + assert loaded == second_snapshot + + @pytest.mark.asyncio + async def test_pg_checkpoint_store_list_by_tape_returns_sorted_meta( + self, store: PGCheckpointStore + ): + later = self._snapshot( + "cp-later", + "tape-a", + created_at=datetime(2026, 4, 17, 1, tzinfo=UTC), + ) + earlier = self._snapshot( + "cp-earlier", + "tape-a", + created_at=datetime(2026, 4, 17, 0, tzinfo=UTC), + ) + other = self._snapshot( + "cp-other", + "tape-b", + created_at=datetime(2026, 4, 17, 2, tzinfo=UTC), + ) + + await store.save(later) + await store.save(earlier) + await store.save(other) + + listed = await store.list_by_tape("tape-a") + + assert listed == [earlier.meta, later.meta] + + @pytest.mark.asyncio + async def test_pg_checkpoint_store_delete_removes_snapshot( + self, store: PGCheckpointStore + ): + snapshot = self._snapshot( + "cp-delete", + "tape-delete", + created_at=datetime(2026, 4, 17, tzinfo=UTC), + ) + + await store.save(snapshot) + await store.delete("cp-delete") + + assert await store.load("cp-delete") is None + + @pytest.mark.asyncio + async def test_pg_checkpoint_store_schema_created_once( + self, store: PGCheckpointStore, fake_pool: FakePool + ): + await store.save( + self._snapshot( + "cp-one", + "tape-one", + created_at=datetime(2026, 4, 17, tzinfo=UTC), + ) + ) + await store.save( + self._snapshot( + "cp-two", + "tape-two", + created_at=datetime(2026, 4, 17, 1, tzinfo=UTC), + ) + ) + + schema_calls = [ + query + for query, _ in fake_pool.executed + if "CREATE TABLE IF NOT EXISTS agent_checkpoints" in query + ] + assert len(schema_calls) == 1 + + @pytest.mark.asyncio + async def test_pg_checkpoint_store_save_passes_python_objects_to_codec_enabled_pool( + self, store: PGCheckpointStore, fake_pool: FakePool + ): + snapshot = self._snapshot( + "checkpoint-codec", + "tape-codec", + created_at=datetime(2025, 1, 1, tzinfo=UTC), + ) + + await store.save(snapshot) + + insert_calls = [ + args + for query, args in fake_pool.executed + if "INSERT INTO agent_checkpoints" in query + ] + assert len(insert_calls) == 1 + assert isinstance(insert_calls[0][2], dict) + assert isinstance(insert_calls[0][3], list) + assert isinstance(insert_calls[0][4], dict) + assert isinstance(insert_calls[0][5], dict) + + class MockPoolForLock: def __init__(self) -> None: self._conn = FakePoolConnection() diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 1aa5f6ecd..e03e4c22d 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -3,14 +3,20 @@ from __future__ import annotations import asyncio +import importlib import json +import sys import types +from collections.abc import AsyncIterator from datetime import datetime, timedelta +from typing import cast from unittest.mock import patch import pytest from httpx import AsyncClient, ASGITransport from httpx_sse import aconnect_sse +from starlette.requests import Request +from agentkit.errors import ConfigError from agentkit.checkpoint.models import CheckpointMeta from agentkit.providers.models import DoneEvent, TextEvent, ToolCallEvent @@ -21,7 +27,9 @@ from coding_agent.ui.http_server import ( APPROVAL_TIMEOUT_SECONDS, SESSION_IDLE_TIMEOUT_MINUTES, + _cleanup_event_queue_on_disconnect, _broadcast_event, + get_events, _session_to_dict, stream_wire_messages, _wire_message_to_event, @@ -136,6 +144,29 @@ async def test_healthz_reports_store_backed_session_count(self, client): assert health.json()["sessions"] == 1 assert session_manager.has_session(session_id) + async def test_healthz_uses_count_sessions_async(self, client, monkeypatch): + async def fake_count_sessions_async() -> int: + return 7 + + def fail_list_sessions_async(): + raise AssertionError("healthz should not call list_sessions_async") + + monkeypatch.setattr( + session_manager, + "count_sessions_async", + fake_count_sessions_async, + ) + monkeypatch.setattr( + session_manager, + "list_sessions_async", + fail_list_sessions_async, + ) + + health = await client.get("/healthz") + + assert health.status_code == 200 + assert health.json()["sessions"] == 7 + async def test_readyz_reports_dependencies_ready(self, client): ready = await client.get("/readyz") @@ -732,6 +763,221 @@ async def test_multiple_queues_in_session(self, client): assert hasattr(session, "event_queues") assert isinstance(session.event_queues, list) + async def test_event_queue_cleanup_is_shielded_on_disconnect(self, monkeypatch): + session_id = "disconnect-session" + queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=1) + + cleanup_started = asyncio.Event() + cleanup_released = asyncio.Event() + cleaned: list[tuple[str, object]] = [] + + async def fake_remove_event_queue_async(current_session_id: str, queue) -> None: + cleaned.append((current_session_id, queue)) + cleanup_started.set() + await cleanup_released.wait() + + monkeypatch.setattr( + session_manager, + "remove_event_queue_async", + fake_remove_event_queue_async, + ) + + task = asyncio.create_task( + _cleanup_event_queue_on_disconnect(session_id, queue) + ) + await asyncio.wait_for(cleanup_started.wait(), timeout=1) + cleanup_released.set() + await task + + assert len(cleaned) == 1 + assert cleaned == [(session_id, queue)] + + async def test_event_queue_cleanup_ignores_missing_session(self, monkeypatch): + queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=1) + + async def fake_remove_event_queue_async(current_session_id: str, queue) -> None: + _ = (current_session_id, queue) + raise KeyError("Session not found: removed") + + monkeypatch.setattr( + session_manager, + "remove_event_queue_async", + fake_remove_event_queue_async, + ) + + await _cleanup_event_queue_on_disconnect("removed", queue) + + async def test_event_generator_uses_public_session_apis_for_keepalive_exit( + self, client, monkeypatch + ): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + + has_session_calls: list[str] = [] + get_session_calls: list[str] = [] + has_session_results = iter([True, True, False]) + + class FakeEventSourceResponse: + def __init__(self, body_iterator): + self.body_iterator = body_iterator + + async def fake_has_session_async(current_session_id: str) -> bool: + has_session_calls.append(current_session_id) + return next(has_session_results) + + async def fake_get_session_async(current_session_id: str): + get_session_calls.append(current_session_id) + return session_manager.get_session(current_session_id) + + monkeypatch.setattr( + session_manager, "has_session_async", fake_has_session_async + ) + monkeypatch.setattr( + session_manager, "get_session_async", fake_get_session_async + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.EventSourceResponse", + FakeEventSourceResponse, + ) + + real_wait_for = asyncio.wait_for + + async def fake_wait_for(awaitable, timeout): + if timeout == 30.0: + awaitable.close() + raise asyncio.TimeoutError + return await real_wait_for(awaitable, timeout) + + monkeypatch.setattr( + "coding_agent.ui.http_server.asyncio.wait_for", fake_wait_for + ) + + request = Request( + { + "type": "http", + "method": "GET", + "path": f"/sessions/{session_id}/events", + "headers": [], + } + ) + response = await get_events(request, session_id, None) + event_generator = cast(AsyncIterator[dict[str, str]], response.body_iterator) + event = await anext(event_generator) + + assert event == {"event": "ping", "data": ""} + + with pytest.raises(StopAsyncIteration): + await anext(event_generator) + + assert has_session_calls == [session_id, session_id, session_id] + assert len(get_session_calls) >= 2 + assert all( + call_session_id == session_id for call_session_id in get_session_calls + ) + + async def test_event_generator_exits_cleanly_when_session_disappears_during_keepalive( + self, client, monkeypatch + ): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + + has_session_results = iter([True, True]) + + class FakeEventSourceResponse: + def __init__(self, body_iterator): + self.body_iterator = body_iterator + + async def fake_has_session_async(current_session_id: str) -> bool: + assert current_session_id == session_id + return next(has_session_results) + + async def fake_has_event_queue_async(current_session_id: str, queue) -> bool: + _ = queue + assert current_session_id == session_id + raise KeyError(f"Session not found: {session_id}") + + monkeypatch.setattr( + session_manager, "has_session_async", fake_has_session_async + ) + monkeypatch.setattr( + session_manager, "has_event_queue_async", fake_has_event_queue_async + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.EventSourceResponse", + FakeEventSourceResponse, + ) + + real_wait_for = asyncio.wait_for + + async def fake_wait_for(awaitable, timeout): + if timeout == 30.0: + awaitable.close() + raise asyncio.TimeoutError + return await real_wait_for(awaitable, timeout) + + monkeypatch.setattr( + "coding_agent.ui.http_server.asyncio.wait_for", fake_wait_for + ) + + request = Request( + { + "type": "http", + "method": "GET", + "path": f"/sessions/{session_id}/events", + "headers": [], + } + ) + response = await get_events(request, session_id, None) + event_generator = cast(AsyncIterator[dict[str, str]], response.body_iterator) + + with pytest.raises(StopAsyncIteration): + await anext(event_generator) + + +class TestLifespanShutdown: + async def test_lifespan_shutdown_continues_after_session_failure(self, monkeypatch): + observed_shutdowns: list[str] = [] + close_calls: list[str] = [] + + async def fake_cleanup_idle_sessions() -> None: + try: + while True: + await asyncio.sleep(3600) + except asyncio.CancelledError: + raise + + async def fake_list_sessions_async() -> list[str]: + return ["session-a", "session-b"] + + async def fake_shutdown_session_runtime(session_id: str) -> None: + observed_shutdowns.append(session_id) + if session_id == "session-a": + raise RuntimeError("boom") + + async def fake_close() -> None: + close_calls.append("closed") + + monkeypatch.setattr( + "coding_agent.ui.http_server._cleanup_idle_sessions", + fake_cleanup_idle_sessions, + ) + monkeypatch.setattr( + session_manager, "list_sessions_async", fake_list_sessions_async + ) + monkeypatch.setattr( + session_manager, + "shutdown_session_runtime", + fake_shutdown_session_runtime, + ) + monkeypatch.setattr(session_manager, "close", fake_close) + + cm = app.router.lifespan_context(app) + await cm.__aenter__() + await cm.__aexit__(None, None, None) + + assert observed_shutdowns == ["session-a", "session-b"] + assert close_calls == ["closed"] + class TestGetSession: """Tests for get session endpoint.""" @@ -795,6 +1041,102 @@ async def test_close_session_broadcasts_event(self, client): assert any(e["event"] == "SessionClosed" for e in received_events) + async def test_close_session_returns_error_when_manager_close_fails( + self, client, monkeypatch + ): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + + async def failing_close_session(current_session_id: str) -> None: + assert current_session_id == session_id + raise RuntimeError("close exploded") + + monkeypatch.setattr(session_manager, "close_session", failing_close_session) + + response = await client.delete(f"/sessions/{session_id}") + + assert response.status_code == 500 + assert response.json()["detail"] == "Internal server error" + + async def test_close_session_returns_404_when_session_disappears_during_close( + self, client, monkeypatch + ): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + + async def fake_has_session_async(current_session_id: str) -> bool: + assert current_session_id == session_id + return True + + async def fake_get_session_async(current_session_id: str): + assert current_session_id == session_id + return session_manager.get_session(current_session_id) + + async def disappearing_close_session(current_session_id: str) -> None: + assert current_session_id == session_id + raise KeyError(f"Session not found: {session_id}") + + monkeypatch.setattr( + session_manager, "has_session_async", fake_has_session_async + ) + monkeypatch.setattr( + session_manager, "get_session_async", fake_get_session_async + ) + monkeypatch.setattr( + session_manager, "close_session", disappearing_close_session + ) + + response = await client.delete(f"/sessions/{session_id}") + + assert response.status_code == 404 + assert response.json()["detail"] == f"Session not found: {session_id}" + + async def test_close_session_returns_404_when_session_disappears_before_load( + self, client, monkeypatch + ): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + + async def fake_has_session_async(current_session_id: str) -> bool: + assert current_session_id == session_id + return True + + async def disappearing_get_session_async(current_session_id: str): + assert current_session_id == session_id + raise KeyError(f"Session not found: {session_id}") + + monkeypatch.setattr( + session_manager, "has_session_async", fake_has_session_async + ) + monkeypatch.setattr( + session_manager, "get_session_async", disappearing_get_session_async + ) + + response = await client.delete(f"/sessions/{session_id}") + + assert response.status_code == 404 + assert response.json()["detail"] == f"Session not found: {session_id}" + + async def test_close_session_hides_unexpected_internal_error_detail( + self, client, monkeypatch, caplog + ): + create_resp = await client.post("/sessions", json={}) + session_id = create_resp.json()["session_id"] + + async def failing_close_session(current_session_id: str) -> None: + assert current_session_id == session_id + raise RuntimeError("dsn=postgresql://user:secret@example/db") + + monkeypatch.setattr(session_manager, "close_session", failing_close_session) + + with caplog.at_level("ERROR"): + response = await client.delete(f"/sessions/{session_id}") + + assert response.status_code == 500 + assert response.json()["detail"] == "Internal server error" + assert "dsn=postgresql://user:secret@example/db" not in response.text + assert "Unexpected error while closing session" in caplog.text + class TestSessionTimeout: """Tests for session idle timeout.""" @@ -1117,6 +1459,53 @@ async def test_broadcast_to_multiple_queues(self): assert await queue1.get() == event assert await queue2.get() == event + async def test_broadcast_uses_provided_session_without_manager_lookup(self): + session = register_session("broadcast-without-lookup") + queue = asyncio.Queue() + session.event_queues = [queue] + event = {"event": "Test", "data": "{}"} + + with patch.object( + session_manager, + "broadcast_event", + side_effect=AssertionError("manager lookup should be skipped"), + ): + await _broadcast_event(session, event) + + assert await queue.get() == event + + async def test_broadcast_prunes_full_queue_without_blocking(self): + session = register_session("broadcast-full-queue") + full_queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=1) + await full_queue.put({"event": "Old", "data": "{}"}) + healthy_queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=1) + session.event_queues = [full_queue, healthy_queue] + event = {"event": "Test", "data": "{}"} + + await _broadcast_event(session, event) + + assert session.event_queues == [healthy_queue] + assert full_queue.qsize() == 1 + assert await healthy_queue.get() == event + + async def test_broadcast_prunes_failed_queue(self): + session = register_session("broadcast-failed-queue") + + class BrokenQueue: + def put_nowait(self, item: object) -> None: + _ = item + raise RuntimeError("queue closed") + + healthy_queue: asyncio.Queue[dict[str, str]] = asyncio.Queue(maxsize=1) + broken_queue = cast(asyncio.Queue[dict[str, str]], BrokenQueue()) + session.event_queues = [broken_queue, healthy_queue] + event = {"event": "Test", "data": "{}"} + + await _broadcast_event(session, event) + + assert session.event_queues == [healthy_queue] + assert await healthy_queue.get() == event + class TestWaitForApproval: """Tests for the approval wait function.""" @@ -1215,6 +1604,48 @@ async def test_wait_for_approval_request_can_be_approved_via_http_endpoint( assert approval_response.feedback == "approved over http" +def test_http_server_import_falls_back_when_agent_toml_is_unreadable( + monkeypatch, +) -> None: + original_module = sys.modules.get("coding_agent.ui.http_server") + monkeypatch.delitem(sys.modules, "coding_agent.ui.http_server", raising=False) + + try: + with patch("agentkit.config.loader.load_config") as load_config: + load_config.side_effect = ConfigError( + "config file not found: /tmp/missing-agent.toml" + ) + http_server = importlib.import_module("coding_agent.ui.http_server") + + assert http_server._load_storage_config() == {} + assert http_server.session_manager._storage_config == {} + finally: + if original_module is None: + monkeypatch.delitem( + sys.modules, "coding_agent.ui.http_server", raising=False + ) + else: + sys.modules["coding_agent.ui.http_server"] = original_module + + +def test_http_server_import_raises_on_invalid_agent_toml(monkeypatch) -> None: + original_module = sys.modules.get("coding_agent.ui.http_server") + monkeypatch.delitem(sys.modules, "coding_agent.ui.http_server", raising=False) + + try: + with patch("agentkit.config.loader.load_config") as load_config: + load_config.side_effect = ConfigError("missing [agent] section") + with pytest.raises(ConfigError, match=r"missing \[agent\] section"): + importlib.import_module("coding_agent.ui.http_server") + finally: + if original_module is None: + monkeypatch.delitem( + sys.modules, "coding_agent.ui.http_server", raising=False + ) + else: + sys.modules["coding_agent.ui.http_server"] = original_module + + class TestIntegration: """Integration tests for the full flow.""" 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 dbb1e4230..935412617 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -1,6 +1,9 @@ from __future__ import annotations +import asyncio +import threading import types from datetime import datetime +from pathlib import Path from unittest.mock import patch import pytest @@ -10,6 +13,7 @@ from agentkit.tape.tape import Tape from coding_agent.approval.store import ApprovalStore from coding_agent.approval import ApprovalPolicy +from agentkit.config.loader import ConfigError from coding_agent.wire.protocol import ( ApprovalRequest, CompletionStatus, @@ -20,6 +24,7 @@ from coding_agent.ui.session_manager import MockProvider, Session, SessionManager from coding_agent.ui.session_store import ( InMemorySessionStore, + PGSessionMetadataStore, RedisSessionStore, create_session_store, ) @@ -152,6 +157,467 @@ async def test_create_session_persists_configured_restart_metadata_by_default() assert payload["base_url"] == "http://llm.default" +def test_event_queue_mutations_do_not_persist_sync_runtime_only_state() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session = Session( + id="runtime-only-sync-queues", + created_at=datetime.now(), + last_activity=datetime.now(), + approval_store=ApprovalStore(), + ) + manager.register_session(session) + queue: asyncio.Queue[dict[str, object]] = asyncio.Queue() + + with patch.object(manager, "_persist_session", side_effect=AssertionError): + manager.add_event_queue(session.id, queue) + manager.remove_event_queue(session.id, queue) + + assert session.event_queues == [] + + +@pytest.mark.asyncio +async def test_event_queue_mutations_do_not_persist_async_runtime_only_state() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session = Session( + id="runtime-only-async-queues", + created_at=datetime.now(), + last_activity=datetime.now(), + approval_store=ApprovalStore(), + ) + manager.register_session(session) + queue: asyncio.Queue[dict[str, object]] = asyncio.Queue() + + with patch.object(manager, "_persist_session_async", side_effect=AssertionError): + await manager.add_event_queue_async(session.id, queue) + await manager.remove_event_queue_async(session.id, queue) + + assert session.event_queues == [] + + +@pytest.mark.asyncio +async def test_run_store_io_serializes_executor_calls() -> None: + manager = SessionManager(store=InMemorySessionStore()) + gate = threading.Event() + order: list[str] = [] + concurrent_runs = 0 + max_concurrent_runs = 0 + + def blocking_store_call(name: str) -> str: + nonlocal concurrent_runs, max_concurrent_runs + concurrent_runs += 1 + max_concurrent_runs = max(max_concurrent_runs, concurrent_runs) + order.append(f"start:{name}") + try: + gate.wait() + finally: + order.append(f"end:{name}") + concurrent_runs -= 1 + return name + + async def call_store(name: str) -> str: + return await manager._run_store_io(blocking_store_call, name) + + first = asyncio.create_task(call_store("first")) + await asyncio.sleep(0) + second = asyncio.create_task(call_store("second")) + await asyncio.sleep(0.05) + + gate.set() + results = await asyncio.gather(first, second) + + assert results == ["first", "second"] + assert max_concurrent_runs == 1 + assert order == ["start:first", "end:first", "start:second", "end:second"] + + +def test_session_manager_uses_pg_backends_when_storage_config_requests_pg() -> None: + class FakePGPool: + instances: list[FakePGPool] = [] + + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + self.__class__.instances.append(self) + + async def get_pool(self) -> FakePGPool: + return self + + async def close(self) -> None: + return None + + class FakePGTapeStore: + def __init__(self, *, pool: FakePGPool) -> None: + self.pool = pool + + class FakePGCheckpointStore: + def __init__(self, *, pool: FakePGPool) -> None: + self.pool = pool + + with ( + patch("coding_agent.ui.session_manager.create_session_store") as create_store, + patch( + "coding_agent.ui.session_manager._load_pg_storage_types", + return_value=(FakePGPool, FakePGTapeStore, FakePGCheckpointStore), + ), + ): + create_store.return_value = PGSessionMetadataStore( + pool=FakePGPool(dsn="postgresql://example") + ) + manager = SessionManager( + storage_config={ + "tape_backend": "pg", + "checkpoint_backend": "pg", + "http_session_backend": "pg", + "dsn": "postgresql://example", + } + ) + + assert isinstance(manager._store, PGSessionMetadataStore) + assert isinstance(manager._tape_store, FakePGTapeStore) + assert isinstance(manager._checkpoint_service._store, FakePGCheckpointStore) + assert manager._tape_store.pool is manager._checkpoint_service._store.pool + assert manager._store._pool is not manager._tape_store.pool + assert len(FakePGPool.instances) == 2 + + +def test_session_manager_creates_dedicated_pg_pool_for_http_session_store() -> None: + class FakePGPool: + instances: list[FakePGPool] = [] + + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + self.__class__.instances.append(self) + + class FakePGTapeStore: + def __init__(self, *, pool: FakePGPool) -> None: + self.pool = pool + + class FakePGCheckpointStore: + def __init__(self, *, pool: FakePGPool) -> None: + self.pool = pool + + with ( + patch("coding_agent.ui.session_manager.create_session_store") as create_store, + patch( + "coding_agent.ui.session_manager._load_pg_storage_types", + return_value=(FakePGPool, FakePGTapeStore, FakePGCheckpointStore), + ), + ): + create_store.return_value = InMemorySessionStore() + _ = SessionManager( + storage_config={ + "tape_backend": "pg", + "checkpoint_backend": "pg", + "http_session_backend": "pg", + "dsn": "postgresql://example", + } + ) + + assert create_store.call_count == 1 + assert create_store.call_args.kwargs["backend"] == "pg" + assert create_store.call_args.kwargs["dsn"] == "postgresql://example" + assert create_store.call_args.kwargs["pg_pool"] is None + assert len(FakePGPool.instances) == 1 + + +def test_session_manager_normalizes_tape_backend_for_http_session_pg_default() -> None: + manager = SessionManager.__new__(SessionManager) + manager._storage_config = { + "tape_backend": " PG ", + "dsn": "postgresql://example", + } + manager._pg_pool = None + + class FakePGPool: + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + + with ( + patch("coding_agent.ui.session_manager.create_session_store") as create_store, + patch( + "coding_agent.ui.session_manager._load_pg_storage_types", + return_value=(FakePGPool, object(), object()), + ), + ): + create_store.return_value = InMemorySessionStore() + + _ = SessionManager._create_http_session_store(manager) + + assert create_store.call_args.kwargs["backend"] == "pg" + + +def test_session_manager_normalizes_tape_backend_for_checkpoint_default() -> None: + manager = SessionManager.__new__(SessionManager) + manager._storage_config = { + "tape_backend": " PG ", + "dsn": "postgresql://example", + } + + class FakePGPool: + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + + class FakePGCheckpointStore: + def __init__(self, *, pool: FakePGPool) -> None: + self.pool = pool + + manager._pg_pool = None + + with patch( + "coding_agent.ui.session_manager._load_pg_storage_types", + return_value=(FakePGPool, object(), FakePGCheckpointStore), + ): + store = SessionManager._create_checkpoint_store(manager, Path("/tmp/data")) + + assert isinstance(store, FakePGCheckpointStore) + + +def test_session_manager_strips_dsn_before_creating_pg_pool() -> None: + class FakePGPool: + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + + manager = SessionManager.__new__(SessionManager) + manager._storage_config = {"dsn": " postgresql://example "} + manager._pg_pool = None + + with patch( + "coding_agent.ui.session_manager._load_pg_storage_types", + return_value=(FakePGPool, object(), object()), + ): + pool = SessionManager._get_pg_pool(manager) + + assert isinstance(pool, FakePGPool) + assert pool.dsn == "postgresql://example" + + +def test_session_manager_reuses_injected_pg_pool_without_importing_pg_types() -> None: + injected_pool = object() + manager = SessionManager.__new__(SessionManager) + manager._storage_config = {} + manager._pg_pool = injected_pool + + with patch( + "coding_agent.ui.session_manager._load_pg_storage_types", + side_effect=AssertionError("should not import pg types when pool injected"), + ): + pool = SessionManager._get_pg_pool(manager) + + assert pool is injected_pool + + +def test_session_manager_http_session_store_does_not_reuse_injected_pg_pool() -> None: + injected_pool = object() + manager = SessionManager.__new__(SessionManager) + manager._storage_config = { + "http_session_backend": "pg", + "dsn": "postgresql://example", + } + manager._pg_pool = injected_pool + + with patch("coding_agent.ui.session_manager.create_session_store") as create_store: + create_store.return_value = InMemorySessionStore() + + _ = SessionManager._create_http_session_store(manager) + + assert create_store.call_args.kwargs["backend"] == "pg" + assert create_store.call_args.kwargs["pg_pool"] is None + + +@pytest.mark.asyncio +async def test_remove_session_async_loads_store_once_when_not_cached() -> None: + class CountingStore(InMemorySessionStore): + def __init__(self) -> None: + super().__init__() + self.load_calls = 0 + + def load(self, session_id: str): + self.load_calls += 1 + return super().load(session_id) + + store = CountingStore() + manager = SessionManager(store=store) + session = Session( + id="remove-me", + created_at=datetime.now(), + last_activity=datetime.now(), + approval_store=ApprovalStore(), + ) + store.save(session.id, session.to_store_data()) + + await manager.remove_session_async("remove-me") + + assert store.load_calls == 1 + assert store.load("remove-me") is None + + +@pytest.mark.asyncio +async def test_close_offloads_sync_store_close_to_executor() -> None: + close_called = threading.Event() + loop_thread_id = threading.get_ident() + close_thread_ids: list[int] = [] + + class ClosableStore(InMemorySessionStore): + def close(self) -> None: + close_thread_ids.append(threading.get_ident()) + close_called.set() + + manager = SessionManager(store=ClosableStore()) + + await manager.close() + + assert close_called.is_set() is True + assert close_thread_ids == [close_thread_ids[0]] + assert close_thread_ids[0] != loop_thread_id + + +@pytest.mark.asyncio +async def test_close_does_not_close_injected_pg_pool() -> None: + class FakePool: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + pool = FakePool() + manager = SessionManager(store=InMemorySessionStore(), pg_pool=pool) + + await manager.close() + + assert pool.closed is False + + +@pytest.mark.asyncio +async def test_shutdown_session_runtime_preserves_persisted_session_metadata() -> None: + store = InMemorySessionStore() + manager = SessionManager(store=store) + session_id = await manager.create_session() + session = manager.get_session(session_id) + session.turn_in_progress = True + session.runtime_pipeline = object() + session.runtime_ctx = object() + previous_last_activity = session.last_activity + + class FakeAdapter: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + adapter = FakeAdapter() + session.runtime_adapter = adapter + + await manager.shutdown_session_runtime(session_id) + + payload = store.get(session_id) + assert payload is not None + assert payload["id"] == session_id + assert adapter.closed is True + reloaded = manager.get_session(session_id) + assert reloaded.last_activity == previous_last_activity + assert reloaded.turn_in_progress is False + assert reloaded.runtime_pipeline is None + assert reloaded.runtime_ctx is None + assert reloaded.runtime_adapter is None + + +@pytest.mark.asyncio +async def test_capture_checkpoint_persists_tape_id_via_async_store_path() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + runtime_ctx = types.SimpleNamespace(tape=Tape(tape_id="stable-tape")) + session.runtime_pipeline = object() + session.runtime_ctx = runtime_ctx + session.runtime_adapter = object() + + expected = CheckpointMeta( + checkpoint_id="cp-save", + tape_id="stable-tape", + session_id=session_id, + entry_count=0, + window_start=0, + created_at=datetime.now(), + label="manual save", + ) + persist_calls: list[str] = [] + + class FakeCheckpointService: + async def capture(self, ctx, label: str | None = None, extra=None): + del label, extra + assert ctx is runtime_ctx + return expected + + async def fake_persist_session_async(current_session: Session) -> None: + persist_calls.append(current_session.id) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + manager, "_checkpoint_service", FakeCheckpointService(), raising=False + ) + mp.setattr( + manager, + "_persist_session_async", + fake_persist_session_async, + raising=False, + ) + checkpoint = await manager.capture_checkpoint(session_id, label="manual save") + + assert checkpoint == expected + assert persist_calls == [session_id] + + +@pytest.mark.asyncio +async def test_cleanup_idle_sessions_uses_async_store_helpers() -> None: + store = InMemorySessionStore() + manager = SessionManager(store=store) + first_id = await manager.create_session() + second_id = await manager.create_session() + + first = manager.get_session(first_id) + second = manager.get_session(second_id) + first.last_activity = datetime.now() + second.last_activity = datetime.now() + + sync_calls: list[str] = [] + async_calls: list[str] = [] + closed_sessions: list[str] = [] + + def fail_sync_list_sessions() -> list[str]: + sync_calls.append("list") + raise AssertionError("cleanup_idle_sessions should not call sync list_sessions") + + def fail_sync_get_session(_session_id: str) -> Session: + sync_calls.append("get") + raise AssertionError("cleanup_idle_sessions should not call sync get_session") + + async def fake_list_sessions_async() -> list[str]: + async_calls.append("list") + return [first_id, second_id] + + async def fake_get_session_async(session_id: str) -> Session: + async_calls.append(f"get:{session_id}") + return first if session_id == first_id else second + + async def fake_close_session(session_id: str) -> None: + closed_sessions.append(session_id) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr(manager._store, "list_sessions", fail_sync_list_sessions) + mp.setattr(manager, "get_session", fail_sync_get_session) + mp.setattr(manager, "list_sessions_async", fake_list_sessions_async) + mp.setattr(manager, "get_session_async", fake_get_session_async) + mp.setattr(manager, "close_session", fake_close_session) + + closed = await manager.cleanup_idle_sessions(max_idle_minutes=10_000) + + assert closed == [] + assert closed_sessions == [] + assert sync_calls == [] + assert async_calls == ["list", f"get:{first_id}", f"get:{second_id}"] + + def test_create_session_store_warns_and_falls_back_when_redis_unreachable( caplog: pytest.LogCaptureFixture, ) -> None: @@ -170,6 +636,24 @@ def failing_factory(url: str): assert "redis://example:6379/0" in caplog.text +def test_create_session_store_raises_config_error_when_redis_factory_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + def failing_factory(url: str): + raise ConfigError("invalid redis config") + + with ( + caplog.at_level("WARNING"), + pytest.raises(ConfigError, match="invalid redis config"), + ): + _ = create_session_store( + redis_url="redis://example:6379/0", + redis_client_factory=failing_factory, + ) + + assert "falling back to in-memory" not in caplog.text + + def test_create_session_store_redacts_reformatted_redis_exception_text( caplog: pytest.LogCaptureFixture, ) -> None: diff --git a/coding-agent/tests/ui/test_session_manager_runtime.py b/coding-agent/tests/ui/test_session_manager_runtime.py index b104f68b9..d358e6f57 100644 --- a/coding-agent/tests/ui/test_session_manager_runtime.py +++ b/coding-agent/tests/ui/test_session_manager_runtime.py @@ -2,6 +2,8 @@ import asyncio import types +from typing import cast +from unittest.mock import patch import pytest @@ -15,7 +17,11 @@ ToolCallDelta, TurnEnd, ) -from coding_agent.ui.session_manager import MockProvider, SessionManager +from coding_agent.ui.session_manager import ( + MockProvider, + SessionManager, + _load_pg_storage_types, +) from coding_agent.ui.session_store import InMemorySessionStore @@ -57,6 +63,31 @@ def fake_create_agent(**kwargs): assert captured_kwargs["api_key"] is None +def test_load_pg_storage_types_reports_missing_optional_dependencies() -> None: + with pytest.MonkeyPatch.context() as mp: + fake_import_error = ModuleNotFoundError("No module named 'asyncpg'") + mp.setattr( + "coding_agent.ui.session_manager.importlib.import_module", + lambda name: (_ for _ in ()).throw(fake_import_error), + ) + + with pytest.raises(RuntimeError, match="optional dependencies"): + _load_pg_storage_types() + + +def test_load_pg_storage_types_reports_missing_exports() -> None: + fake_module = types.SimpleNamespace(PGPool=object(), PGTapeStore=object()) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "coding_agent.ui.session_manager.importlib.import_module", + lambda name: fake_module, + ) + + with pytest.raises(RuntimeError, match="PGCheckpointStore"): + _load_pg_storage_types() + + @pytest.mark.asyncio async def test_run_agent_emits_error_turn_end_when_bootstrap_fails() -> None: manager = SessionManager() @@ -271,6 +302,95 @@ def fake_create_agent(**kwargs): assert session.runtime_adapter is None +@pytest.mark.asyncio +async def test_remove_session_async_awaits_runtime_close() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + close_started = asyncio.Event() + close_released = asyncio.Event() + + class FakeAdapter: + async def close(self) -> None: + close_started.set() + await close_released.wait() + + session.runtime_pipeline = object() + session.runtime_ctx = object() + session.runtime_adapter = FakeAdapter() + + task = asyncio.create_task(manager.remove_session_async(session_id)) + await asyncio.wait_for(close_started.wait(), timeout=1) + + assert task.done() is False + assert manager.has_session(session_id) is True + + close_released.set() + await task + + assert manager.has_session(session_id) is False + + +@pytest.mark.asyncio +async def test_close_session_raises_if_task_survives_cancellation() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + + class FakeTask: + def __init__(self) -> None: + self.cancel_calls = 0 + + def done(self) -> bool: + return False + + def cancel(self) -> None: + self.cancel_calls += 1 + + task = cast(asyncio.Task[None], FakeTask()) + session.task = task + + with patch( + "coding_agent.ui.session_manager.asyncio.wait_for", + side_effect=asyncio.TimeoutError, + ): + with pytest.raises(RuntimeError, match="did not stop after cancellation"): + await manager.close_session(session_id) + + assert manager.has_session(session_id) is True + assert cast(FakeTask, task).cancel_calls == 1 + + +@pytest.mark.asyncio +async def test_shutdown_session_runtime_raises_if_task_survives_cancellation() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + + class FakeTask: + def __init__(self) -> None: + self.cancel_calls = 0 + + def done(self) -> bool: + return False + + def cancel(self) -> None: + self.cancel_calls += 1 + + task = cast(asyncio.Task[None], FakeTask()) + session.task = task + + with patch( + "coding_agent.ui.session_manager.asyncio.wait_for", + side_effect=asyncio.TimeoutError, + ): + with pytest.raises(RuntimeError, match="did not stop after cancellation"): + await manager.shutdown_session_runtime(session_id) + + assert manager.has_session(session_id) is True + assert cast(FakeTask, task).cancel_calls == 1 + + @pytest.mark.asyncio async def test_register_session_closes_cached_runtime() -> None: manager = SessionManager(store=InMemorySessionStore()) diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 7d0f38148..a086fea8d 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -1,6 +1,11 @@ from __future__ import annotations +import asyncio from datetime import datetime +from concurrent.futures import TimeoutError as FutureTimeoutError +from typing import Callable, cast +from unittest.mock import MagicMock, call +from unittest.mock import patch import pytest @@ -8,6 +13,7 @@ from coding_agent.ui.session_manager import Session, SessionManager from coding_agent.ui.session_store import ( InMemorySessionStore, + PGSessionMetadataStore, RedisSessionStore, create_session_store, ) @@ -74,6 +80,25 @@ def failing_factory(url: str): assert "falling back to in-memory" in caplog.text +def test_create_session_store_falls_back_when_redis_factory_raises_connection_error( + caplog: pytest.LogCaptureFixture, +) -> None: + class FakeRedisConnectionError(Exception): + pass + + def failing_factory(url: str): + raise FakeRedisConnectionError(f"cannot connect to {url}") + + with caplog.at_level("WARNING"): + store = create_session_store( + redis_url="redis://example:6379/0", + redis_client_factory=failing_factory, + ) + + assert isinstance(store, InMemorySessionStore) + assert "falling back to in-memory" in caplog.text + + def test_in_memory_session_store_reports_healthy() -> None: assert InMemorySessionStore().check_health() is True @@ -83,6 +108,7 @@ class FakeRedisClient: def __init__(self) -> None: self._data: dict[str, str] = {} self._index: set[str] = set() + self.smembers_calls = 0 def ping(self) -> bool: return True @@ -106,6 +132,7 @@ def srem(self, key: str, value: str) -> None: def smembers(self, key: str) -> set[str]: assert key == "coding-agent:sessions:index" + self.smembers_calls += 1 return set(self._index) client = FakeRedisClient() @@ -128,6 +155,845 @@ def smembers(self, key: str) -> set[str]: assert "persisted-session" in store.list_sessions() +def test_redis_session_store_count_uses_scard_when_available() -> None: + class FakeRedisClient: + def __init__(self) -> None: + self.smembers_calls = 0 + self.scard_calls = 0 + + def ping(self) -> bool: + return True + + def set(self, key: str, value: str) -> None: + _ = (key, value) + raise AssertionError("unused") + + def get(self, key: str) -> str | None: + _ = key + raise AssertionError("unused") + + def delete(self, key: str) -> None: + _ = key + raise AssertionError("unused") + + def sadd(self, key: str, value: str) -> None: + _ = (key, value) + raise AssertionError("unused") + + def srem(self, key: str, value: str) -> None: + _ = (key, value) + raise AssertionError("unused") + + def smembers(self, key: str) -> set[str]: + _ = key + self.smembers_calls += 1 + return {"unexpected"} + + def scard(self, key: str) -> int: + assert key == "coding-agent:sessions:index" + self.scard_calls += 1 + return 7 + + client = FakeRedisClient() + store = RedisSessionStore(client=client, redis_url="redis://test") + + assert store.count_sessions() == 7 + assert client.scard_calls == 1 + assert client.smembers_calls == 0 + + +def test_redis_session_store_count_falls_back_to_smembers_without_scard() -> None: + class FakeRedisClient: + def __init__(self) -> None: + self.smembers_calls = 0 + + def ping(self) -> bool: + return True + + def set(self, key: str, value: str) -> None: + _ = (key, value) + raise AssertionError("unused") + + def get(self, key: str) -> str | None: + _ = key + raise AssertionError("unused") + + def delete(self, key: str) -> None: + _ = key + raise AssertionError("unused") + + def sadd(self, key: str, value: str) -> None: + _ = (key, value) + raise AssertionError("unused") + + def srem(self, key: str, value: str) -> None: + _ = (key, value) + raise AssertionError("unused") + + def smembers(self, key: str) -> set[str]: + assert key == "coding-agent:sessions:index" + self.smembers_calls += 1 + return {"one", "two", "three"} + + client = FakeRedisClient() + store = RedisSessionStore(client=client, redis_url="redis://test") + + assert store.count_sessions() == 3 + assert client.smembers_calls == 1 + + +def test_redis_session_store_accepts_client_without_scard_method() -> None: + class FakeRedisClient: + def __init__(self) -> None: + self._data: dict[str, str] = {} + self._index: set[str] = set() + + def ping(self) -> bool: + return True + + def set(self, key: str, value: str) -> None: + self._data[key] = value + + def get(self, key: str) -> str | None: + return self._data.get(key) + + def delete(self, key: str) -> None: + self._data.pop(key, None) + + def sadd(self, key: str, value: str) -> None: + assert key == "coding-agent:sessions:index" + self._index.add(value) + + def srem(self, key: str, value: str) -> None: + assert key == "coding-agent:sessions:index" + self._index.discard(value) + + def smembers(self, key: str) -> set[str]: + assert key == "coding-agent:sessions:index" + return set(self._index) + + client = FakeRedisClient() + store = RedisSessionStore(client=cast(object, client), redis_url="redis://test") + + store.save("session-a", {"id": "session-a"}) + store.save("session-b", {"id": "session-b"}) + + assert store.count_sessions() == 2 + + +def test_pg_session_metadata_store_round_trips_session_metadata() -> None: + class FakeRecord: + def __init__(self, payload: dict[str, object]) -> None: + self._payload = payload + + def __iter__(self): + return iter(self._payload.items()) + + class FakeAsyncPGPool: + def __init__(self) -> None: + self.sessions: dict[str, dict[str, object]] = {} + self.executed: list[tuple[str, tuple[object, ...]]] = [] + self.closed = False + + async def execute(self, query: str, *args: object) -> str: + self.executed.append((query, args)) + if "INSERT INTO agent_http_sessions" in query: + session_id, payload = args + assert isinstance(session_id, str) + assert isinstance(payload, dict) + self.sessions[session_id] = payload + return "INSERT 0 1" + if "DELETE FROM agent_http_sessions" in query: + (session_id,) = args + assert isinstance(session_id, str) + self.sessions.pop(session_id, None) + return "DELETE 1" + if "CREATE TABLE IF NOT EXISTS agent_http_sessions" in query: + return "CREATE TABLE" + raise AssertionError(f"unexpected execute query: {query}") + + async def fetchrow(self, query: str, *args: object) -> object | None: + self.executed.append((query, args)) + if "SELECT payload FROM agent_http_sessions" in query: + (session_id,) = args + assert isinstance(session_id, str) + payload = self.sessions.get(session_id) + return None if payload is None else FakeRecord({"payload": payload}) + if query.strip() == "SELECT 1": + return FakeRecord({"?column?": 1}) + raise AssertionError(f"unexpected fetchrow query: {query}") + + async def fetch(self, query: str, *args: object) -> list[object]: + self.executed.append((query, args)) + if "SELECT session_id FROM agent_http_sessions" in query: + return [ + FakeRecord({"session_id": session_id}) + for session_id in sorted(self.sessions.keys()) + ] + raise AssertionError(f"unexpected fetch query: {query}") + + async def acquire(self): + raise AssertionError("unused") + + async def release(self, connection): + _ = connection + raise AssertionError("unused") + + async def close(self) -> None: + self.closed = True + + class FakePGPool: + def __init__(self) -> None: + self.pool = FakeAsyncPGPool() + + async def get_pool(self) -> FakeAsyncPGPool: + return self.pool + + async def close(self) -> None: + await self.pool.close() + + pg_pool = FakePGPool() + store = PGSessionMetadataStore(pool=pg_pool) + try: + session = Session( + id="pg-session", + created_at=datetime.now(), + last_activity=datetime.now(), + approval_store=ApprovalStore(), + ) + + store.save(session.id, session.to_store_data()) + + payload = store.load("pg-session") + + assert payload is not None + assert payload["id"] == "pg-session" + assert store.list_sessions() == ["pg-session"] + assert store.check_health() is True + insert_calls = [ + args + for query, args in pg_pool.pool.executed + if "INSERT INTO agent_http_sessions" in query + ] + assert len(insert_calls) == 1 + assert isinstance(insert_calls[0][1], dict) + finally: + store.close() + + +def test_pg_session_metadata_store_count_sessions_uses_count_query() -> None: + class FakeRecord: + def __init__(self, payload: dict[str, object]) -> None: + self._payload = payload + + def __iter__(self): + return iter(self._payload.items()) + + class FakeAsyncPGPool: + def __init__(self) -> None: + self.executed: list[tuple[str, tuple[object, ...]]] = [] + self.closed = False + + async def execute(self, query: str, *args: object) -> str: + self.executed.append((query, args)) + if "CREATE TABLE IF NOT EXISTS agent_http_sessions" in query: + return "CREATE TABLE" + raise AssertionError(f"unexpected execute query: {query}") + + async def fetchrow(self, query: str, *args: object) -> object | None: + self.executed.append((query, args)) + if "SELECT COUNT(*) AS session_count FROM agent_http_sessions" in query: + return FakeRecord({"session_count": 3}) + raise AssertionError(f"unexpected fetchrow query: {query}") + + async def fetch(self, query: str, *args: object) -> list[object]: + self.executed.append((query, args)) + raise AssertionError(f"unexpected fetch query: {query}") + + async def acquire(self): + raise AssertionError("unused") + + async def release(self, connection): + _ = connection + raise AssertionError("unused") + + async def close(self) -> None: + self.closed = True + + class FakePGPool: + def __init__(self) -> None: + self.pool = FakeAsyncPGPool() + + async def get_pool(self) -> FakeAsyncPGPool: + return self.pool + + async def close(self) -> None: + await self.pool.close() + + pg_pool = FakePGPool() + store = PGSessionMetadataStore(pool=pg_pool) + try: + assert store.count_sessions() == 3 + count_calls = [ + args + for query, args in pg_pool.pool.executed + if "SELECT COUNT(*) AS session_count FROM agent_http_sessions" in query + ] + assert count_calls == [()] + finally: + store.close() + + +def test_pg_session_metadata_store_waits_for_loop_start_before_returning() -> None: + created_loops: list[FakeLoop] = [] + created_threads: list[FakeThread] = [] + + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + class FakeLoop: + def __init__(self) -> None: + self.closed = False + self.scheduled_callbacks: list[Callable[[], object]] = [] + created_loops.append(self) + + def call_soon(self, callback) -> None: + self.scheduled_callbacks.append(callback) + + def run_forever(self) -> None: + for callback in list(self.scheduled_callbacks): + callback() + return None + + def stop(self) -> None: + return None + + def call_soon_threadsafe(self, callback) -> None: + callback() + + def close(self) -> None: + self.closed = True + + class FakeThread: + def __init__(self, *, target, name: str, daemon: bool) -> None: + self.target = target + self.name = name + self.daemon = daemon + self.started = False + created_threads.append(self) + + def start(self) -> None: + self.started = True + self.target() + + def join(self, timeout: float | None = None) -> None: + _ = timeout + + def is_alive(self) -> bool: + return False + + with ( + patch( + "coding_agent.ui.session_store.asyncio.new_event_loop", side_effect=FakeLoop + ), + patch("coding_agent.ui.session_store.asyncio.set_event_loop"), + patch("coding_agent.ui.session_store.threading.Thread", side_effect=FakeThread), + ): + store = PGSessionMetadataStore(pool=FakePool()) + + try: + assert len(created_loops) == 1 + assert len(created_threads) == 1 + assert created_threads[0].started is True + assert store._loop_ready.is_set() is True + assert len(created_loops[0].scheduled_callbacks) == 1 + finally: + store._run_sync = lambda operation: (operation.close(), None)[1] + store.close() + + +def test_pg_session_metadata_store_closes_background_loop_and_pool() -> None: + class FakeAsyncPGPool: + def __init__(self) -> None: + self.closed = False + + async def execute(self, query: str, *args: object) -> str: + _ = (query, args) + raise AssertionError("unused") + + async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None: + _ = (query, args) + raise AssertionError("unused") + + async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: + _ = (query, args) + raise AssertionError("unused") + + async def acquire(self): + raise AssertionError("unused") + + async def release(self, connection): + _ = connection + raise AssertionError("unused") + + async def close(self) -> None: + self.closed = True + + class FakePGPool: + def __init__(self) -> None: + self.pool = FakeAsyncPGPool() + + async def get_pool(self) -> FakeAsyncPGPool: + return self.pool + + async def close(self) -> None: + await self.pool.close() + + pool = FakePGPool() + store = PGSessionMetadataStore(pool=pool, owns_pool=True) + + store.close() + + assert pool.pool.closed is True + + +def test_pg_session_metadata_store_leaves_injected_pool_open() -> None: + class FakeAsyncPGPool: + def __init__(self) -> None: + self.closed = False + + async def execute(self, query: str, *args: object) -> str: + _ = (query, args) + raise AssertionError("unused") + + async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None: + _ = (query, args) + raise AssertionError("unused") + + async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: + _ = (query, args) + raise AssertionError("unused") + + async def acquire(self): + raise AssertionError("unused") + + async def release(self, connection): + _ = connection + raise AssertionError("unused") + + async def close(self) -> None: + self.closed = True + + class FakePGPool: + def __init__(self) -> None: + self.pool = FakeAsyncPGPool() + + async def get_pool(self) -> FakeAsyncPGPool: + return self.pool + + async def close(self) -> None: + await self.pool.close() + + pool = FakePGPool() + store = PGSessionMetadataStore(pool=pool) + + assert store.close() is None + assert pool.pool.closed is False + + +def test_pg_session_metadata_store_skips_loop_close_when_thread_does_not_stop() -> None: + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + class FakeThread: + def __init__(self) -> None: + self.join_calls: list[float | None] = [] + + def start(self) -> None: + return None + + def join(self, timeout: float | None = None) -> None: + self.join_calls.append(timeout) + + def is_alive(self) -> bool: + return True + + class FakeLoop: + def __init__(self) -> None: + self.closed = False + self.stop_calls = 0 + + def run_forever(self) -> None: + return None + + def stop(self) -> None: + self.stop_calls += 1 + + def call_soon_threadsafe(self, callback) -> None: + callback() + + def close(self) -> None: + self.closed = True + + store = PGSessionMetadataStore(pool=FakePool()) + fake_thread = FakeThread() + fake_loop = FakeLoop() + + def fake_run_sync(operation): + operation.close() + return None + + with ( + patch.object(store, "_loop_thread", fake_thread), + patch.object(store, "_loop", fake_loop), + patch.object(store, "_run_sync", fake_run_sync), + ): + store.close() + + assert fake_thread.join_calls == [5] + assert fake_loop.stop_calls == 1 + assert fake_loop.closed is False + + +def test_pg_session_metadata_store_close_handles_already_closed_loop() -> None: + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + class FakeThread: + def __init__(self) -> None: + self.join_calls: list[float | None] = [] + + def start(self) -> None: + return None + + def join(self, timeout: float | None = None) -> None: + self.join_calls.append(timeout) + + def is_alive(self) -> bool: + return False + + class FakeLoop: + def __init__(self) -> None: + self.close_calls = 0 + + def stop(self) -> None: + return None + + def call_soon_threadsafe(self, callback) -> None: + _ = callback + raise RuntimeError("loop already closed") + + def close(self) -> None: + self.close_calls += 1 + raise RuntimeError("loop already closed") + + store = PGSessionMetadataStore(pool=FakePool()) + fake_thread = FakeThread() + fake_loop = FakeLoop() + + def fake_run_sync(operation): + operation.close() + return None + + with ( + patch.object(store, "_loop_thread", fake_thread), + patch.object(store, "_loop", fake_loop), + patch.object(store, "_run_sync", fake_run_sync), + patch("coding_agent.ui.session_store.logger.warning") as warning, + ): + store.close() + + assert fake_thread.join_calls == [5] + assert fake_loop.close_calls == 1 + assert warning.call_count == 2 + + +def test_pg_session_metadata_store_close_swallows_pool_close_errors() -> None: + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + class FakeThread: + def __init__(self) -> None: + self.join_calls: list[float | None] = [] + + def start(self) -> None: + return None + + def join(self, timeout: float | None = None) -> None: + self.join_calls.append(timeout) + + def is_alive(self) -> bool: + return False + + class FakeLoop: + def __init__(self) -> None: + self.stop_calls = 0 + self.close_calls = 0 + + def stop(self) -> None: + self.stop_calls += 1 + + def call_soon_threadsafe(self, callback) -> None: + callback() + + def close(self) -> None: + self.close_calls += 1 + + store = PGSessionMetadataStore(pool=FakePool()) + fake_thread = FakeThread() + fake_loop = FakeLoop() + + def fake_run_sync(operation): + operation.close() + raise OSError("pool close failed") + + with ( + patch.object(store, "_loop_thread", fake_thread), + patch.object(store, "_loop", fake_loop), + patch.object(store, "_run_sync", fake_run_sync), + patch("coding_agent.ui.session_store.logger.warning") as warning, + ): + store.close() + + assert fake_thread.join_calls == [5] + assert fake_loop.stop_calls == 1 + assert fake_loop.close_calls == 1 + warning.assert_called_once() + + +def test_pg_session_metadata_store_close_test_does_not_start_real_thread() -> None: + created_loops: list[FakeLoop] = [] + created_threads: list[FakeThread] = [] + + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + class FakeThread: + def __init__(self, *, target, name: str, daemon: bool) -> None: + self.target = cast(object, target) + self.name = name + self.daemon = daemon + self.started = False + self.join_calls: list[float | None] = [] + created_threads.append(self) + + def start(self) -> None: + self.started = True + + def join(self, timeout: float | None = None) -> None: + self.join_calls.append(timeout) + + def is_alive(self) -> bool: + return True + + class FakeLoop: + def __init__(self) -> None: + self.closed = False + self.stop_calls = 0 + created_loops.append(self) + + def run_forever(self) -> None: + return None + + def stop(self) -> None: + self.stop_calls += 1 + + def call_soon_threadsafe(self, callback) -> None: + callback() + + def close(self) -> None: + self.closed = True + + def fake_run_sync(operation: object) -> None: + cast(object, operation) + close = getattr(operation, "close") + cast(object, close) + close() + return None + + with ( + patch( + "coding_agent.ui.session_store.asyncio.new_event_loop", side_effect=FakeLoop + ), + patch("coding_agent.ui.session_store.threading.Thread", side_effect=FakeThread), + ): + with pytest.raises( + RuntimeError, match="postgres session metadata loop thread failed to start" + ): + PGSessionMetadataStore(pool=FakePool()) + + assert len(created_loops) == 1 + assert len(created_threads) == 1 + assert created_threads[0].started is True + assert created_threads[0].join_calls == [0.1] + assert created_loops[0].stop_calls == 1 + assert created_loops[0].closed is False + + +def test_pg_session_metadata_store_run_sync_times_out_and_cancels_future() -> None: + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + future = MagicMock() + future.result.side_effect = FutureTimeoutError() + + store = PGSessionMetadataStore(pool=FakePool()) + try: + operation = asyncio.sleep(0) + with patch( + "coding_agent.ui.session_store.asyncio.run_coroutine_threadsafe", + return_value=future, + ): + with pytest.raises( + TimeoutError, match="postgres session metadata operation timed out" + ): + store._run_sync(operation) + + future.cancel.assert_called_once_with() + operation.close() + finally: + store.close() + + +def test_pg_session_metadata_store_run_sync_waits_for_cancellation_callback() -> None: + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + future = MagicMock() + future.result.side_effect = [FutureTimeoutError(), FutureTimeoutError()] + callbacks: list[Callable[[object], object]] = [] + + def add_done_callback(callback): + callbacks.append(callback) + + future.add_done_callback.side_effect = add_done_callback + + store = PGSessionMetadataStore(pool=FakePool()) + try: + operation = asyncio.sleep(0) + with patch( + "coding_agent.ui.session_store.asyncio.run_coroutine_threadsafe", + return_value=future, + ): + with pytest.raises( + TimeoutError, match="postgres session metadata operation timed out" + ): + store._run_sync(operation) + + future.cancel.assert_called_once_with() + assert len(callbacks) == 1 + future.result.assert_has_calls([call(timeout=30.0), call(timeout=0.1)]) + callbacks[0](future) + operation.close() + finally: + store.close() + + +def test_pg_session_metadata_store_run_sync_raises_clear_error_when_loop_unavailable() -> ( + None +): + class FakePool: + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + return None + + store = PGSessionMetadataStore(pool=FakePool()) + try: + operation = asyncio.sleep(0) + with patch( + "coding_agent.ui.session_store.asyncio.run_coroutine_threadsafe", + side_effect=RuntimeError("loop closed"), + ): + with pytest.raises( + RuntimeError, + match="postgres session metadata loop is not available", + ): + store._run_sync(operation) + + assert operation.cr_frame is None + finally: + store.close() + + +def test_create_session_store_strips_dsn_before_building_pg_pool() -> None: + with patch("coding_agent.ui.session_store.PGPool") as pg_pool_cls: + store = create_session_store(backend="pg", dsn=" postgresql://example ") + assert isinstance(store, PGSessionMetadataStore) + try: + assert pg_pool_cls.call_args.kwargs["dsn"] == "postgresql://example" + finally: + store._loop.call_soon_threadsafe(store._loop.stop) + store._loop_thread.join(timeout=5) + if not store._loop_thread.is_alive(): + store._loop.close() + + +def test_create_session_store_rejects_all_whitespace_dsn() -> None: + with pytest.raises(ValueError, match="PG session store requires dsn or pg_pool"): + create_session_store(backend="pg", dsn=" ") + + +def test_create_session_store_builds_pg_store_from_explicit_backend() -> None: + store = create_session_store(backend="pg", dsn="postgresql://example") + assert isinstance(store, PGSessionMetadataStore) + try: + assert isinstance(store, PGSessionMetadataStore) + finally: + store.close() + + +def test_create_session_store_accepts_injected_pg_pool_without_dsn() -> None: + class MockPGPool: + def __init__(self) -> None: + self.closed = False + + async def get_pool(self) -> object: + raise AssertionError("unused") + + async def close(self) -> None: + self.closed = True + + pool = MockPGPool() + store = create_session_store(backend="pg", pg_pool=pool) + assert isinstance(store, PGSessionMetadataStore) + try: + assert store._pool is pool + finally: + store.close() + + assert pool.closed is False + + def test_rehydrate_clears_non_restart_safe_runtime_state() -> None: store = InMemorySessionStore() manager = SessionManager(store=store)