From 3810a647d6287fee47ef9203fc2f9f0c59641d61 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 16:32:58 +0800 Subject: [PATCH 01/42] docs: add postgres phase1 completion ADR Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../postgres-phase1-completion-task-packet.md | 45 ++++++++++++++++ ...se1-postgresql-http-session-persistence.md | 54 +++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 coding-agent/.opencode/prompts/postgres-phase1-completion-task-packet.md create mode 100644 coding-agent/docs/adr/0012-complete-phase1-postgresql-http-session-persistence.md 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..cd6b207df --- /dev/null +++ b/coding-agent/docs/adr/0012-complete-phase1-postgresql-http-session-persistence.md @@ -0,0 +1,54 @@ +# 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 + +- [ ] `test_pg_checkpoint_store_round_trip_snapshot` +- [ ] `test_pg_checkpoint_store_list_by_tape_returns_sorted_meta` +- [ ] `test_create_persistence_bundle_builds_pg_backends_from_storage_config` +- [ ] `test_session_manager_defaults_use_configured_pg_persistence_bundle` +- [ ] `test_http_server_uses_configured_persistence_bundle_for_global_session_manager` +- [ ] `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 -k "pg or checkpoint or persistence_bundle or session_manager_defaults" -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` From 9e7dbf624cf1391d5a6402035fa79429bd2a2204 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 16:33:06 +0800 Subject: [PATCH 02/42] feat(storage): add postgres checkpoint store Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/agentkit/storage/pg.py | 175 +++++++++++ .../tests/agentkit/storage/test_pg.py | 288 +++++++++++++++++- 2 files changed, 448 insertions(+), 15 deletions(-) diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py index 002c457b4..9388bc5a8 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 @@ -231,6 +246,88 @@ 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) + """ + _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'" + ) + _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, + json.dumps(meta_payload), + json.dumps(list(snapshot.tape_entries)), + json.dumps(snapshot.plugin_states), + json.dumps(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 @@ -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/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py index cb8cde892..9994494ca 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, ...]]] = [] @@ -71,6 +81,28 @@ 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) + for item in (meta_json, entries_json, state_json, extra_json) + ): + raise TypeError("checkpoint payloads must be encoded json") + self.checkpoints[checkpoint_id] = { + "checkpoint_id": checkpoint_id, + "tape_id": tape_id, + "meta": json.loads(meta_json), + "entries": json.loads(entries_json), + "plugin_states": json.loads(state_json), + "extra": json.loads(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 +120,49 @@ 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 MAX(seq) AS max_seq FROM agent_tapes" in 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 +192,23 @@ 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 + ] + rows.sort( + key=lambda row: ( + cast(dict[str, object], row["meta"])["created_at"] + if isinstance(row["meta"], dict) + else "" + ) + ) + return rows raise AssertionError(f"unexpected fetch query: {query}") async def close(self) -> None: @@ -148,6 +222,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 +323,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() @@ -451,6 +573,142 @@ 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_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 + + class MockPoolForLock: def __init__(self) -> None: self._conn = FakePoolConnection() From 827c4a1187a6bfecc65f53bb48c796ae3ab86bfb Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 16:33:14 +0800 Subject: [PATCH 03/42] feat(ui): back http session metadata with asyncpg Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 207 +++++++++++++++++- .../tests/ui/test_session_persistence.py | 137 ++++++++++++ 2 files changed, 343 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index c3031cb9d..a502fcf43 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -1,11 +1,15 @@ from __future__ import annotations +import asyncio import json import logging import os -from collections.abc import Callable +import threading +from collections.abc import Callable, Coroutine from typing import Protocol, cast +from agentkit.storage.pg import PGPool + from ..redaction import redact_sensitive_text, redact_url_credentials logger = logging.getLogger(__name__) @@ -46,6 +50,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] = {} @@ -124,12 +134,207 @@ def check_health(self) -> bool: return False +class PGSessionMetadataStore: + _CREATE_TABLE_SQL = """ + 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 = """ + 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 = "SELECT payload FROM agent_http_sessions WHERE session_id = $1" + _LIST_SQL = "SELECT session_id FROM agent_http_sessions ORDER BY session_id" + _DELETE_SQL = "DELETE FROM agent_http_sessions WHERE session_id = $1" + + def __init__( + self, + *, + pool: AsyncPGSessionPool, + ) -> None: + self._pool = pool + self._schema_ready = False + self._loop = asyncio.new_event_loop() + self._loop_thread = threading.Thread( + target=self._run_loop, + name="pg-session-metadata-store", + daemon=True, + ) + self._loop_thread.start() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def _run_sync(self, operation: Coroutine[object, object, object]) -> object: + future = asyncio.run_coroutine_threadsafe(operation, self._loop) + return future.result() + + 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, + json.dumps(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 + if not isinstance(row_obj, dict): + raise TypeError("postgres session metadata row must decode to a dict") + payload = row_obj.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 rows_obj: + if not isinstance(row, dict): + raise TypeError("postgres session metadata list rows must be dicts") + session_id = row.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 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: + await self._pool.close() + + try: + _ = self._run_sync(_close_pool()) + finally: + self._loop.call_soon_threadsafe(self._loop.stop) + self._loop_thread.join(timeout=5) + self._loop.close() + + 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") + ) resolved_redis_url = redis_url or os.environ.get("AGENT_SESSION_REDIS_URL") + if resolved_backend == "pg": + if not resolved_dsn: + raise ValueError("PG session store requires dsn") + return PGSessionMetadataStore( + pool=pg_pool or PGPool(dsn=resolved_dsn), + ) + + 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() diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 7d0f38148..41427c6bb 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json from datetime import datetime import pytest @@ -8,6 +9,7 @@ from coding_agent.ui.session_manager import Session, SessionManager from coding_agent.ui.session_store import ( InMemorySessionStore, + PGSessionMetadataStore, RedisSessionStore, create_session_store, ) @@ -128,6 +130,141 @@ def smembers(self, key: str) -> set[str]: assert "persisted-session" in store.list_sessions() +def test_pg_session_metadata_store_round_trips_session_metadata() -> None: + 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, str) + self.sessions[session_id] = json.loads(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) -> dict[str, 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 {"payload": payload} + if query.strip() == "SELECT 1": + return {"?column?": 1} + raise AssertionError(f"unexpected fetchrow query: {query}") + + async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: + self.executed.append((query, args)) + if "SELECT session_id FROM agent_http_sessions" in query: + return [ + {"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() + + store = PGSessionMetadataStore( + pool=FakePGPool(), + ) + 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 + + +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) + + store.close() + + assert pool.pool.closed is True + + +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) + + def test_rehydrate_clears_non_restart_safe_runtime_state() -> None: store = InMemorySessionStore() manager = SessionManager(store=store) From 7baf6e7a5bd46ec90df52df3379e6c3bb65fbe5e Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 16:33:35 +0800 Subject: [PATCH 04/42] feat(ui): wire postgres session persistence defaults Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/coding_agent/agent.toml | 2 + .../src/coding_agent/ui/http_server.py | 13 ++- .../src/coding_agent/ui/session_manager.py | 98 +++++++++++++++++-- .../ui/test_session_manager_public_api.py | 39 ++++++++ 4 files changed, 143 insertions(+), 9 deletions(-) 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..a97d06eb3 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -15,6 +15,7 @@ from fastapi.middleware.cors import CORSMiddleware from sse_starlette.sse import EventSourceResponse +from agentkit.config.loader import load_config as load_agent_toml from coding_agent.approval import ApprovalPolicy from coding_agent.ui.session_manager import Session, SessionManager from coding_agent.ui.schemas import ( @@ -58,8 +59,17 @@ SESSION_IDLE_TIMEOUT_MINUTES = 30 +def _load_storage_config() -> dict[str, Any]: + config_path = Path(__file__).resolve().parent.parent / "agent.toml" + return cast(dict[str, Any], load_agent_toml(config_path).extra.get("storage", {})) + + +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: @@ -87,6 +97,7 @@ async def lifespan(app: FastAPI): # Close all sessions for session_id in list(session_manager.list_sessions()): await session_manager.close_session(session_id) + await session_manager.close() logger.info("HTTP server shut down") diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 791eaac25..ec1759ea8 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -15,15 +15,17 @@ from pathlib import Path from typing import Any, Literal, 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 @@ -40,8 +42,6 @@ SessionStore, create_session_store, ) -from agentkit.storage.protocols import CheckpointStore, TapeStore -from agentkit.checkpoint.models import CheckpointMeta logger = logging.getLogger(__name__) @@ -273,6 +273,20 @@ 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; add agentkit.storage.pg before using tape_backend='pg'" + ) from exc + return ( + getattr(pg_module, "PGPool"), + getattr(pg_module, "PGTapeStore"), + getattr(pg_module, "PGCheckpointStore"), + ) + + class _WireConsumer: def __init__( self, @@ -296,26 +310,81 @@ 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._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._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) -> object: + PGPool, _, _ = _load_pg_storage_types() + if self._pg_pool is not None: + return self._pg_pool + + dsn_obj = self._storage_config.get("dsn") + if not isinstance(dsn_obj, str) or not dsn_obj.strip(): + raise RuntimeError("PG storage requires storage.dsn") + self._pg_pool = PGPool(dsn=dsn_obj) + return self._pg_pool + + def _create_http_session_store(self) -> SessionStore: + configured_backend = self._storage_config.get("http_session_backend") + 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 self._storage_config.get("tape_backend") == "pg": + configured_backend = "pg" + + backend = configured_backend 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, + ) + + 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: + default_backend = ( + "pg" if self._storage_config.get("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) @@ -720,6 +789,19 @@ async def close_session(self, session_id: str) -> None: logger.info(f"Closed session: {session_id}") + async def close(self) -> None: + store_close = getattr(self._store, "close", None) + if callable(store_close): + store_close_result = store_close() + if isawaitable(store_close_result): + await store_close_result + + close = getattr(self._pg_pool, "close", None) + if callable(close): + close_result = close() + if isawaitable(close_result): + await close_result + async def run_agent( self, session_id: str, 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..e1f22000e 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -20,6 +20,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 +153,44 @@ async def test_create_session_persists_configured_restart_metadata_by_default() assert payload["base_url"] == "http://llm.default" +def test_session_manager_uses_pg_backends_when_storage_config_requests_pg() -> None: + class FakePGPool: + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + + 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 + + def test_create_session_store_warns_and_falls_back_when_redis_unreachable( caplog: pytest.LogCaptureFixture, ) -> None: From 039aff50af4a8cd32cba8d06734fcfd305b87ba7 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 16:55:58 +0800 Subject: [PATCH 05/42] fix(ui): handle asyncpg record session rows Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 31 +++++++++--- .../tests/ui/test_session_persistence.py | 48 +++++++++++-------- 2 files changed, 52 insertions(+), 27 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index a502fcf43..8a47aa589 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -5,7 +5,7 @@ import logging import os import threading -from collections.abc import Callable, Coroutine +from collections.abc import Callable, Coroutine, Iterable from typing import Protocol, cast from agentkit.storage.pg import PGPool @@ -215,9 +215,10 @@ async def _load() -> SessionPayload | None: ) if row_obj is None: return None - if not isinstance(row_obj, dict): - raise TypeError("postgres session metadata row must decode to a dict") - payload = row_obj.get("payload") + 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): @@ -246,9 +247,10 @@ async def _list_sessions() -> list[str]: raise TypeError("postgres session metadata list result must be a list") session_ids: list[str] = [] for row in rows_obj: - if not isinstance(row, dict): - raise TypeError("postgres session metadata list rows must be dicts") - session_id = row.get("session_id") + 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" @@ -354,6 +356,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/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 41427c6bb..a49b30af6 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -131,6 +131,13 @@ def smembers(self, key: str) -> set[str]: 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]] = {} @@ -154,22 +161,22 @@ async def execute(self, query: str, *args: object) -> str: return "CREATE TABLE" raise AssertionError(f"unexpected execute query: {query}") - async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None: + 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 {"payload": payload} + return None if payload is None else FakeRecord({"payload": payload}) if query.strip() == "SELECT 1": - return {"?column?": 1} + return FakeRecord({"?column?": 1}) raise AssertionError(f"unexpected fetchrow query: {query}") - async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: + 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 [ - {"session_id": session_id} + FakeRecord({"session_id": session_id}) for session_id in sorted(self.sessions.keys()) ] raise AssertionError(f"unexpected fetch query: {query}") @@ -194,24 +201,25 @@ async def get_pool(self) -> FakeAsyncPGPool: async def close(self) -> None: await self.pool.close() - store = PGSessionMetadataStore( - pool=FakePGPool(), - ) - session = Session( - id="pg-session", - created_at=datetime.now(), - last_activity=datetime.now(), - approval_store=ApprovalStore(), - ) + store = PGSessionMetadataStore(pool=FakePGPool()) + 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()) + store.save(session.id, session.to_store_data()) - payload = store.load("pg-session") + 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 + assert payload is not None + assert payload["id"] == "pg-session" + assert store.list_sessions() == ["pg-session"] + assert store.check_health() is True + finally: + store.close() def test_pg_session_metadata_store_closes_background_loop_and_pool() -> None: From d7ff944b9900344c3079a973e377841369f6b712 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 16:56:23 +0800 Subject: [PATCH 06/42] fix(ui): reuse pg pool for session defaults Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 25 +++-- .../ui/test_session_manager_public_api.py | 95 +++++++++++++++++++ 2 files changed, 113 insertions(+), 7 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index ec1759ea8..22695d150 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -39,6 +39,7 @@ WireMessage, ) from coding_agent.ui.session_store import ( + AsyncPGSessionPool, SessionStore, create_session_store, ) @@ -334,19 +335,22 @@ def __init__( ) self._create_agent = create_agent_fn - def _get_pg_pool(self) -> object: + def _get_pg_pool(self) -> AsyncPGSessionPool: PGPool, _, _ = _load_pg_storage_types() if self._pg_pool is not None: - return self._pg_pool + return cast(AsyncPGSessionPool, self._pg_pool) dsn_obj = self._storage_config.get("dsn") if not isinstance(dsn_obj, str) or not dsn_obj.strip(): raise RuntimeError("PG storage requires storage.dsn") self._pg_pool = PGPool(dsn=dsn_obj) - return self._pg_pool + 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 ( @@ -354,14 +358,20 @@ def _create_http_session_store(self) -> SessionStore: and legacy_backend.strip().lower() == "pg" ): configured_backend = "pg" - elif self._storage_config.get("tape_backend") == "pg": + elif tape_backend == "pg": configured_backend = "pg" - backend = configured_backend if isinstance(configured_backend, str) else None + backend = ( + configured_backend.strip().lower() + if isinstance(configured_backend, str) + else None + ) dsn = self._storage_config.get("dsn") + pg_pool = self._get_pg_pool() if backend == "pg" else None return create_session_store( backend=backend, dsn=dsn if isinstance(dsn, str) else None, + pg_pool=pg_pool, ) def _create_tape_store(self, data_dir: Path) -> TapeStore: @@ -372,9 +382,10 @@ def _create_tape_store(self, data_dir: Path) -> TapeStore: return JSONLTapeStore(data_dir / "tapes") def _create_checkpoint_store(self, data_dir: Path) -> CheckpointStore: - default_backend = ( - "pg" if self._storage_config.get("tape_backend") == "pg" else "fs" + 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() 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 e1f22000e..f3821e955 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,7 @@ from __future__ import annotations import types from datetime import datetime +from pathlib import Path from unittest.mock import patch import pytest @@ -158,6 +159,12 @@ class FakePGPool: def __init__(self, *, dsn: str) -> None: self.dsn = dsn + 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 @@ -191,6 +198,94 @@ def __init__(self, *, pool: FakePGPool) -> None: assert manager._tape_store.pool is manager._checkpoint_service._store.pool +def test_session_manager_reuses_shared_pg_pool_for_http_session_store() -> None: + class FakePGPool: + def __init__(self, *, dsn: str) -> None: + self.dsn = dsn + + 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 isinstance(create_store.call_args.kwargs["pg_pool"], FakePGPool) + + +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_create_session_store_warns_and_falls_back_when_redis_unreachable( caplog: pytest.LogCaptureFixture, ) -> None: From 492caa9abf1b2b357d5105e048baf905c1223394 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 17:17:38 +0800 Subject: [PATCH 07/42] fix(ui): isolate pg pools across event loops Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 3 +- .../src/coding_agent/ui/session_store.py | 7 ++- .../ui/test_session_manager_public_api.py | 13 +++- .../tests/ui/test_session_persistence.py | 59 +++++++++++++++++++ 4 files changed, 77 insertions(+), 5 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 22695d150..0d5b33250 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -367,11 +367,10 @@ def _create_http_session_store(self) -> SessionStore: else None ) dsn = self._storage_config.get("dsn") - pg_pool = self._get_pg_pool() if backend == "pg" else None return create_session_store( backend=backend, dsn=dsn if isinstance(dsn, str) else None, - pg_pool=pg_pool, + pg_pool=None, ) def _create_tape_store(self, data_dir: Path) -> TapeStore: diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 8a47aa589..95fe13db0 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -304,7 +304,12 @@ async def _close_pool() -> None: finally: self._loop.call_soon_threadsafe(self._loop.stop) self._loop_thread.join(timeout=5) - self._loop.close() + if self._loop_thread.is_alive(): + logger.warning( + "Timed out waiting for postgres session metadata loop thread to stop; skipping event loop close" + ) + else: + self._loop.close() def create_session_store( 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 f3821e955..8a0a34b34 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -156,8 +156,11 @@ async def test_create_session_persists_configured_restart_metadata_by_default() 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 @@ -196,12 +199,17 @@ def __init__(self, *, pool: FakePGPool) -> None: 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_reuses_shared_pg_pool_for_http_session_store() -> None: +def test_session_manager_uses_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: @@ -231,7 +239,8 @@ def __init__(self, *, pool: FakePGPool) -> None: assert create_store.call_count == 1 assert create_store.call_args.kwargs["backend"] == "pg" assert create_store.call_args.kwargs["dsn"] == "postgresql://example" - assert isinstance(create_store.call_args.kwargs["pg_pool"], FakePGPool) + 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: diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index a49b30af6..af67b629f 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -2,6 +2,7 @@ import json from datetime import datetime +from unittest.mock import patch import pytest @@ -267,6 +268,64 @@ async def close(self) -> None: assert pool.pool.closed is True +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_create_session_store_builds_pg_store_from_explicit_backend() -> None: store = create_session_store(backend="pg", dsn="postgresql://example") From c109c4a2f856df761252754207823f41d4d38017 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 17 Apr 2026 09:35:50 +0000 Subject: [PATCH 08/42] fix(ui): allow injected pg session pools without dsn Agent-Logs-Url: https://github.com/delete-cloud/Coding-Agent-Loop/sessions/19385579-068c-4ded-9dce-a9591e6b6020 Co-authored-by: delete-cloud <76783546+delete-cloud@users.noreply.github.com> --- .../src/coding_agent/ui/session_store.py | 4 +-- .../tests/ui/test_session_persistence.py | 26 ++++++++++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 95fe13db0..6375a23ea 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -330,8 +330,8 @@ def create_session_store( ) resolved_redis_url = redis_url or os.environ.get("AGENT_SESSION_REDIS_URL") if resolved_backend == "pg": - if not resolved_dsn: - raise ValueError("PG session store requires dsn") + if pg_pool is None and not resolved_dsn: + raise ValueError("PG session store requires dsn or pg_pool") return PGSessionMetadataStore( pool=pg_pool or PGPool(dsn=resolved_dsn), ) diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index af67b629f..8fa18d388 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -328,8 +328,32 @@ def fake_run_sync(operation): def test_create_session_store_builds_pg_store_from_explicit_backend() -> None: store = create_session_store(backend="pg", dsn="postgresql://example") + try: + assert isinstance(store, PGSessionMetadataStore) + finally: + store.close() + + +def test_create_session_store_accepts_injected_pg_pool_without_dsn() -> None: + class FakePGPool: + 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 = FakePGPool() + store = create_session_store(backend="pg", pg_pool=pool) + try: + assert isinstance(store, PGSessionMetadataStore) + assert store._pool is pool + finally: + store.close() - assert isinstance(store, PGSessionMetadataStore) + assert pool.closed is True def test_rehydrate_clears_non_restart_safe_runtime_state() -> None: From 48f6699f53695b09ba0554f2fc765b8a84fdbf2f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 17 Apr 2026 09:38:06 +0000 Subject: [PATCH 09/42] test(ui): tidy injected pg pool regression coverage Agent-Logs-Url: https://github.com/delete-cloud/Coding-Agent-Loop/sessions/19385579-068c-4ded-9dce-a9591e6b6020 Co-authored-by: delete-cloud <76783546+delete-cloud@users.noreply.github.com> --- coding-agent/tests/ui/test_session_persistence.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 8fa18d388..6a9d242f9 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -335,7 +335,7 @@ def test_create_session_store_builds_pg_store_from_explicit_backend() -> None: def test_create_session_store_accepts_injected_pg_pool_without_dsn() -> None: - class FakePGPool: + class MockPGPool: def __init__(self) -> None: self.closed = False @@ -345,7 +345,7 @@ async def get_pool(self) -> object: async def close(self) -> None: self.closed = True - pool = FakePGPool() + pool = MockPGPool() store = create_session_store(backend="pg", pg_pool=pool) try: assert isinstance(store, PGSessionMetadataStore) From e0ae8e5fdce9d4ea95fd7c8b8102375cfb69a4ef Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 17:59:16 +0800 Subject: [PATCH 10/42] fix(ui): normalize pg session dsns Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 3 +- .../src/coding_agent/ui/session_store.py | 8 +- .../ui/test_session_manager_public_api.py | 19 ++++ .../tests/ui/test_session_persistence.py | 90 ++++++++++++++++++- 4 files changed, 116 insertions(+), 4 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 0d5b33250..cebf52962 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -343,7 +343,8 @@ def _get_pg_pool(self) -> AsyncPGSessionPool: dsn_obj = self._storage_config.get("dsn") if not isinstance(dsn_obj, str) or not dsn_obj.strip(): raise RuntimeError("PG storage requires storage.dsn") - self._pg_pool = PGPool(dsn=dsn_obj) + dsn = dsn_obj.strip() + self._pg_pool = PGPool(dsn=dsn) return cast(AsyncPGSessionPool, self._pg_pool) def _create_http_session_store(self) -> SessionStore: diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 6375a23ea..53a14d092 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -328,12 +328,16 @@ def create_session_store( 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": - if pg_pool is None and not resolved_dsn: + 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 = PGPool(dsn=normalized_dsn) return PGSessionMetadataStore( - pool=pg_pool or PGPool(dsn=resolved_dsn), + pool=pg_pool, ) if resolved_backend == "memory": 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 8a0a34b34..d9dd949e0 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -295,6 +295,25 @@ def __init__(self, *, pool: FakePGPool) -> None: 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_create_session_store_warns_and_falls_back_when_redis_unreachable( caplog: pytest.LogCaptureFixture, ) -> None: diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 6a9d242f9..6f389c170 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -326,8 +326,96 @@ def fake_run_sync(operation): assert fake_loop.closed is False +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 = 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): + operation.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), + ): + store = PGSessionMetadataStore(pool=FakePool()) + with patch.object(store, "_run_sync", fake_run_sync): + store.close() + + assert len(created_loops) == 1 + assert len(created_threads) == 1 + assert created_threads[0].started is True + assert created_threads[0].join_calls == [5] + assert created_loops[0].stop_calls == 1 + assert created_loops[0].closed is False + + +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: @@ -347,8 +435,8 @@ async def close(self) -> None: pool = MockPGPool() store = create_session_store(backend="pg", pg_pool=pool) + assert isinstance(store, PGSessionMetadataStore) try: - assert isinstance(store, PGSessionMetadataStore) assert store._pool is pool finally: store.close() From 2142e9159beaa198adbe947781e7a48a383c28c0 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 20:37:38 +0800 Subject: [PATCH 11/42] fix(storage): pass python objects to pg codecs Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/agentkit/storage/pg.py | 10 +-- .../tests/agentkit/storage/test_pg.py | 76 +++++++++++++++---- 2 files changed, 68 insertions(+), 18 deletions(-) diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py index 9388bc5a8..240b61e41 100644 --- a/coding-agent/src/agentkit/storage/pg.py +++ b/coding-agent/src/agentkit/storage/pg.py @@ -125,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() @@ -305,10 +305,10 @@ async def save(self, snapshot: CheckpointSnapshot) -> None: self._INSERT_SQL, meta.checkpoint_id, meta.tape_id, - json.dumps(meta_payload), - json.dumps(list(snapshot.tape_entries)), - json.dumps(snapshot.plugin_states), - json.dumps(snapshot.extra), + meta_payload, + list(snapshot.tape_entries), + snapshot.plugin_states, + snapshot.extra, ) async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: diff --git a/coding-agent/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py index 9994494ca..c18f319f1 100644 --- a/coding-agent/tests/agentkit/storage/test_pg.py +++ b/coding-agent/tests/agentkit/storage/test_pg.py @@ -33,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) @@ -60,20 +63,22 @@ async def execute(self, query: str, *args: object) -> str: else: raise TypeError("seq must be an int or list[int]") - if isinstance(third_arg, str): + if isinstance(third_arg, (str, dict)): payload_values = [third_arg] elif isinstance(third_arg, list) and all( - isinstance(item, str) for item in third_arg + isinstance(item, (str, dict)) for item in third_arg ): payload_values = third_arg else: - raise TypeError("payload must be encoded json or list[str]") + raise TypeError("payload must be encoded json/dict or list thereof") 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) + 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)}) @@ -90,17 +95,25 @@ async def execute(self, query: str, *args: object) -> str: if not isinstance(tape_id, str): raise TypeError("tape_id must be a string") if not all( - isinstance(item, str) + isinstance(item, (str, dict, list)) for item in (meta_json, entries_json, state_json, extra_json) ): - raise TypeError("checkpoint payloads must be encoded 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), - "entries": json.loads(entries_json), - "plugin_states": json.loads(state_json), - "extra": json.loads(extra_json), + "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": @@ -446,6 +459,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 @@ -708,6 +735,29 @@ async def test_pg_checkpoint_store_schema_created_once( ] 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: From 44520455a83af6c4dc281de5398d584875c3bf85 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 20:37:48 +0800 Subject: [PATCH 12/42] fix(ui): fail fast on pg session startup Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 46 +++++--- .../tests/ui/test_session_persistence.py | 100 ++++++++++++++++-- 2 files changed, 122 insertions(+), 24 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 53a14d092..af367a251 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -6,7 +6,7 @@ import os import threading from collections.abc import Callable, Coroutine, Iterable -from typing import Protocol, cast +from typing import Final, Protocol, cast from agentkit.storage.pg import PGPool @@ -135,40 +135,60 @@ def check_health(self) -> bool: class PGSessionMetadataStore: - _CREATE_TABLE_SQL = """ + _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 = """ + _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 = "SELECT payload FROM agent_http_sessions WHERE session_id = $1" - _LIST_SQL = "SELECT session_id FROM agent_http_sessions ORDER BY session_id" - _DELETE_SQL = "DELETE FROM agent_http_sessions WHERE session_id = $1" + _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" + ) + _DELETE_SQL: Final[str] = "DELETE FROM agent_http_sessions WHERE session_id = $1" + _LOOP_READY_TIMEOUT_SECONDS: Final[float] = 5.0 def __init__( self, *, pool: AsyncPGSessionPool, ) -> None: - self._pool = pool - self._schema_ready = False - self._loop = asyncio.new_event_loop() - self._loop_thread = threading.Thread( + self._pool: AsyncPGSessionPool = 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_ready.set() self._loop.run_forever() def _run_sync(self, operation: Coroutine[object, object, object]) -> object: @@ -196,7 +216,7 @@ async def _save() -> None: _ = await cast(Callable[..., Coroutine[object, object, object]], execute)( self._UPSERT_SQL, session_id, - json.dumps(data), + data, ) _ = self._run_sync(_save()) @@ -246,7 +266,7 @@ async def _list_sessions() -> list[str]: if not isinstance(rows_obj, list): raise TypeError("postgres session metadata list result must be a list") session_ids: list[str] = [] - for row in rows_obj: + for row in cast(list[object], rows_obj): row_dict = _coerce_row_dict( row=row, context="postgres session metadata list row" ) @@ -302,7 +322,7 @@ async def _close_pool() -> None: try: _ = self._run_sync(_close_pool()) finally: - self._loop.call_soon_threadsafe(self._loop.stop) + _ = self._loop.call_soon_threadsafe(self._loop.stop) self._loop_thread.join(timeout=5) if self._loop_thread.is_alive(): logger.warning( diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 6f389c170..670bc315e 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -1,7 +1,8 @@ from __future__ import annotations -import json from datetime import datetime +from types import TracebackType +from typing import cast from unittest.mock import patch import pytest @@ -150,8 +151,8 @@ async def execute(self, query: str, *args: object) -> str: if "INSERT INTO agent_http_sessions" in query: session_id, payload = args assert isinstance(session_id, str) - assert isinstance(payload, str) - self.sessions[session_id] = json.loads(payload) + assert isinstance(payload, dict) + self.sessions[session_id] = payload return "INSERT 0 1" if "DELETE FROM agent_http_sessions" in query: (session_id,) = args @@ -202,7 +203,8 @@ async def get_pool(self) -> FakeAsyncPGPool: async def close(self) -> None: await self.pool.close() - store = PGSessionMetadataStore(pool=FakePGPool()) + pg_pool = FakePGPool() + store = PGSessionMetadataStore(pool=pg_pool) try: session = Session( id="pg-session", @@ -219,10 +221,82 @@ async def close(self) -> 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_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 + created_loops.append(self) + + def run_forever(self) -> None: + 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 + 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: @@ -339,7 +413,7 @@ async def close(self) -> None: class FakeThread: def __init__(self, *, target, name: str, daemon: bool) -> None: - self.target = target + self.target = cast(object, target) self.name = name self.daemon = daemon self.started = False @@ -373,8 +447,11 @@ def call_soon_threadsafe(self, callback) -> None: def close(self) -> None: self.closed = True - def fake_run_sync(operation): - operation.close() + def fake_run_sync(operation: object) -> None: + cast(object, operation) + close = getattr(operation, "close") + cast(object, close) + close() return None with ( @@ -383,14 +460,15 @@ def fake_run_sync(operation): ), patch("coding_agent.ui.session_store.threading.Thread", side_effect=FakeThread), ): - store = PGSessionMetadataStore(pool=FakePool()) - with patch.object(store, "_run_sync", fake_run_sync): - store.close() + 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 == [5] + assert created_threads[0].join_calls == [0.1] assert created_loops[0].stop_calls == 1 assert created_loops[0].closed is False From 2c362918612075ed42ab7a70bde985132e8f595a Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 23:00:30 +0800 Subject: [PATCH 13/42] fix(ui): offload pg session store access Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 40 ++++--- .../src/coding_agent/ui/session_manager.py | 109 ++++++++++++++---- .../tests/ui/test_session_persistence.py | 1 - 3 files changed, 110 insertions(+), 40 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index a97d06eb3..4271a95ed 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -375,7 +375,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=len(await session_manager.list_sessions_async()), + version="2.0.0", ) @@ -383,7 +385,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 @@ -457,10 +459,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") @@ -543,10 +545,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") @@ -578,12 +580,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.""" @@ -597,7 +599,7 @@ async def event_generator() -> AsyncIterator[dict[str, str]]: yield {"event": "ping", "data": ""} except asyncio.CancelledError: # Client disconnected - session_manager.remove_event_queue(session_id, queue) + await session_manager.remove_event_queue_async(session_id, queue) raise return EventSourceResponse(event_generator()) @@ -611,8 +613,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") @@ -628,7 +630,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: @@ -662,7 +664,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) @@ -693,7 +695,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: @@ -720,10 +722,10 @@ 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): + 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) await _broadcast_event( session, {"event": "SessionClosed", "data": json.dumps({"session_id": session_id})}, @@ -748,7 +750,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, @@ -756,7 +758,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 cebf52962..859d57780 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -9,11 +9,12 @@ 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 @@ -47,6 +48,7 @@ logger = logging.getLogger(__name__) _CHECKPOINT_SESSION_CONFIG_KEY = "session_restart_config" +T = TypeVar("T") class MockProvider: @@ -428,6 +430,73 @@ 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: + 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 list_sessions_async(self) -> list[str]: + return await self._run_store_io(self._store.list_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) + await self._persist_session_async(session) + + 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) + await self._persist_session_async(session) + + async def check_health_async(self) -> bool: + return bool(await self._run_store_io(self._store.check_health)) + + async def remove_session_async(self, session_id: str) -> None: + if not await self.has_session_async(session_id): + raise KeyError(f"Session not found: {session_id}") + session = await self.get_session_async(session_id) + self._close_runtime_sync_safe(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 @@ -660,7 +729,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 @@ -786,7 +855,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(): @@ -796,7 +865,7 @@ async def close_session(self, session_id: str) -> None: except (asyncio.CancelledError, asyncio.TimeoutError): pass - self.remove_session(session_id) + await self.remove_session_async(session_id) logger.info(f"Closed session: {session_id}") @@ -823,10 +892,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 = { @@ -855,7 +924,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"] = "" @@ -876,7 +945,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") @@ -900,7 +969,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, @@ -926,7 +995,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( @@ -946,7 +1015,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( @@ -961,7 +1030,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, @@ -969,7 +1038,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, @@ -990,7 +1059,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( @@ -1003,7 +1072,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, @@ -1067,7 +1136,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 @@ -1102,7 +1171,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( @@ -1117,7 +1186,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") @@ -1138,7 +1207,7 @@ async def capture_checkpoint( 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) @@ -1149,7 +1218,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/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 670bc315e..037af1269 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -1,7 +1,6 @@ from __future__ import annotations from datetime import datetime -from types import TracebackType from typing import cast from unittest.mock import patch From 167f4d5ab63b0470222cedf73e2240f5984b54dd Mon Sep 17 00:00:00 2001 From: kina0630 Date: Fri, 17 Apr 2026 23:23:38 +0800 Subject: [PATCH 14/42] fix(ui): harden session store review fixes Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 15 +++++++-- .../src/coding_agent/ui/session_manager.py | 4 +-- .../src/coding_agent/ui/session_store.py | 11 ++++++- coding-agent/tests/ui/test_http_server.py | 16 ++++++++++ .../ui/test_session_manager_public_api.py | 28 ++++++++++++++++ .../tests/ui/test_session_persistence.py | 32 +++++++++++++++++++ 6 files changed, 100 insertions(+), 6 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 4271a95ed..42ef00cb1 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -16,6 +16,7 @@ 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 ( @@ -61,7 +62,17 @@ def _load_storage_config() -> dict[str, Any]: config_path = Path(__file__).resolve().parent.parent / "agent.toml" - return cast(dict[str, Any], load_agent_toml(config_path).extra.get("storage", {})) + try: + return cast( + dict[str, Any], load_agent_toml(config_path).extra.get("storage", {}) + ) + except ConfigError: + logger.warning( + "Unable to load storage config from %s; using defaults", + config_path, + exc_info=True, + ) + return {} def _build_session_manager() -> SessionManager: @@ -95,7 +106,7 @@ async def lifespan(app: FastAPI): pass # Close all sessions - for session_id in list(session_manager.list_sessions()): + for session_id in await session_manager.list_sessions_async(): await session_manager.close_session(session_id) await session_manager.close() diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 859d57780..8dbd9b492 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -344,7 +344,7 @@ def _get_pg_pool(self) -> AsyncPGSessionPool: dsn_obj = self._storage_config.get("dsn") if not isinstance(dsn_obj, str) or not dsn_obj.strip(): - raise RuntimeError("PG storage requires storage.dsn") + raise RuntimeError("PG storage requires storage_config['dsn']") dsn = dsn_obj.strip() self._pg_pool = PGPool(dsn=dsn) return cast(AsyncPGSessionPool, self._pg_pool) @@ -488,8 +488,6 @@ async def check_health_async(self) -> bool: return bool(await self._run_store_io(self._store.check_health)) async def remove_session_async(self, session_id: str) -> None: - if not await self.has_session_async(session_id): - raise KeyError(f"Session not found: {session_id}") session = await self.get_session_async(session_id) self._close_runtime_sync_safe(session) self._session_cache.pop(session_id, None) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index af367a251..33beca9ac 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from concurrent.futures import TimeoutError as FutureTimeoutError import json import logging import os @@ -156,6 +157,7 @@ class PGSessionMetadataStore: ) _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, @@ -193,7 +195,14 @@ def _run_loop(self) -> None: def _run_sync(self, operation: Coroutine[object, object, object]) -> object: future = asyncio.run_coroutine_threadsafe(operation, self._loop) - return future.result() + try: + return future.result(timeout=self._SYNC_OPERATION_TIMEOUT_SECONDS) + except FutureTimeoutError as exc: + future.cancel() + 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() diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 1aa5f6ecd..505a052f8 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -3,7 +3,9 @@ from __future__ import annotations import asyncio +import importlib import json +import sys import types from datetime import datetime, timedelta from unittest.mock import patch @@ -11,6 +13,7 @@ import pytest from httpx import AsyncClient, ASGITransport from httpx_sse import aconnect_sse +from agentkit.errors import ConfigError from agentkit.checkpoint.models import CheckpointMeta from agentkit.providers.models import DoneEvent, TextEvent, ToolCallEvent @@ -1215,6 +1218,19 @@ 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: + monkeypatch.delitem(sys.modules, "coding_agent.ui.http_server", raising=False) + + with patch("agentkit.config.loader.load_config") as load_config: + load_config.side_effect = ConfigError("config file not found") + http_server = importlib.import_module("coding_agent.ui.http_server") + + assert http_server._load_storage_config() == {} + assert http_server.session_manager._storage_config == {} + + 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 d9dd949e0..81f4d46d0 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -6,6 +6,7 @@ import pytest +from agentkit.errors import ConfigError from agentkit.checkpoint import CheckpointService from agentkit.checkpoint.models import CheckpointMeta from agentkit.tape.tape import Tape @@ -314,6 +315,33 @@ def __init__(self, *, dsn: str) -> None: assert pool.dsn == "postgresql://example" +@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 + + def test_create_session_store_warns_and_falls_back_when_redis_unreachable( caplog: pytest.LogCaptureFixture, ) -> None: diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 037af1269..3d0ff6230 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -1,7 +1,10 @@ from __future__ import annotations +import asyncio from datetime import datetime +from concurrent.futures import TimeoutError as FutureTimeoutError from typing import cast +from unittest.mock import MagicMock from unittest.mock import patch import pytest @@ -472,6 +475,35 @@ def fake_run_sync(operation: object) -> None: 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_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 ") From 1c58e90233abc836c0c61c524fe437287c92c24b Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 09:27:05 +0800 Subject: [PATCH 15/42] fix(ui): finish async session cleanup Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 13 +++- .../src/coding_agent/ui/session_manager.py | 34 +++++----- .../src/coding_agent/ui/session_store.py | 37 ++++++++++ coding-agent/tests/ui/test_http_server.py | 53 +++++++++++++++ .../ui/test_session_manager_public_api.py | 68 +++++++++++++++++++ 5 files changed, 186 insertions(+), 19 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 42ef00cb1..d7aaca25b 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -66,7 +66,7 @@ def _load_storage_config() -> dict[str, Any]: return cast( dict[str, Any], load_agent_toml(config_path).extra.get("storage", {}) ) - except ConfigError: + except (ConfigError, OSError): logger.warning( "Unable to load storage config from %s; using defaults", config_path, @@ -346,6 +346,13 @@ async def _broadcast_event(session: Session, event: dict[str, str]) -> None: await session_manager.broadcast_event(session.id, event) +async def _cleanup_event_queue_on_disconnect( + session_id: str, + queue: asyncio.Queue[dict[str, str]], +) -> None: + await asyncio.shield(session_manager.remove_event_queue_async(session_id, queue)) + + async def _cleanup_idle_sessions() -> None: """Background task to clean up idle sessions.""" while True: @@ -387,7 +394,7 @@ async def stream_wire_messages(wire: LocalWire) -> AsyncIterator[dict[str, str]] async def liveness_check(request: Request) -> HealthResponse: return HealthResponse( status="healthy", - sessions=len(await session_manager.list_sessions_async()), + sessions=await session_manager.count_sessions_async(), version="2.0.0", ) @@ -610,7 +617,7 @@ async def event_generator() -> AsyncIterator[dict[str, str]]: yield {"event": "ping", "data": ""} except asyncio.CancelledError: # Client disconnected - await session_manager.remove_event_queue_async(session_id, queue) + await _cleanup_event_queue_on_disconnect(session_id, queue) raise return EventSourceResponse(event_generator()) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 8dbd9b492..f942fc5ef 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -461,6 +461,9 @@ async def has_session_async(self, session_id: str) -> bool: 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() @@ -487,6 +490,14 @@ async def remove_event_queue_async( 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) self._close_runtime_sync_safe(session) @@ -527,7 +538,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( @@ -548,12 +559,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) @@ -642,7 +653,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: @@ -868,17 +879,8 @@ async def close_session(self, session_id: str) -> None: logger.info(f"Closed session: {session_id}") async def close(self) -> None: - store_close = getattr(self._store, "close", None) - if callable(store_close): - store_close_result = store_close() - if isawaitable(store_close_result): - await store_close_result - - close = getattr(self._pg_pool, "close", None) - if callable(close): - close_result = close() - if isawaitable(close_result): - await close_result + await self._close_resource_async(self._store) + await self._close_resource_async(self._pg_pool) async def run_agent( self, @@ -1201,7 +1203,7 @@ 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]: diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 33beca9ac..673e59e6c 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -44,6 +44,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: ... @@ -67,6 +69,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) @@ -124,6 +129,9 @@ def list_sessions(self) -> list[str]: session_ids.append(member) return sorted(session_ids) + def count_sessions(self) -> int: + 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) @@ -155,6 +163,7 @@ class PGSessionMetadataStore: _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 @@ -292,6 +301,34 @@ async def _list_sessions() -> list[str]: 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() diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 505a052f8..e9fb48b62 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -24,6 +24,7 @@ from coding_agent.ui.http_server import ( APPROVAL_TIMEOUT_SECONDS, SESSION_IDLE_TIMEOUT_MINUTES, + _cleanup_event_queue_on_disconnect, _broadcast_event, _session_to_dict, stream_wire_messages, @@ -139,6 +140,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") @@ -735,6 +759,35 @@ 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)] + class TestGetSession: """Tests for get session endpoint.""" 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 81f4d46d0..268ce54ec 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -1,4 +1,6 @@ from __future__ import annotations +import asyncio +import threading import types from datetime import datetime from pathlib import Path @@ -342,6 +344,72 @@ def load(self, session_id: str): 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_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] + + def test_create_session_store_warns_and_falls_back_when_redis_unreachable( caplog: pytest.LogCaptureFixture, ) -> None: From b112863c3d20d35057e80e43d2fb28245dd8880b Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 09:47:31 +0800 Subject: [PATCH 16/42] fix(ui): preserve sessions on shutdown Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 2 +- .../src/coding_agent/ui/session_manager.py | 18 ++++++ .../ui/test_session_manager_public_api.py | 33 ++++++++++ .../tests/ui/test_session_persistence.py | 63 +++++++++++++++++++ 4 files changed, 115 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index d7aaca25b..7466e25d8 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -107,7 +107,7 @@ async def lifespan(app: FastAPI): # Close all sessions for session_id in await session_manager.list_sessions_async(): - await session_manager.close_session(session_id) + await session_manager.shutdown_session_runtime(session_id) await session_manager.close() logger.info("HTTP server shut down") diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index f942fc5ef..411822ea8 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -878,6 +878,24 @@ async def close_session(self, session_id: str) -> None: 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 + + await self._close_runtime(session) + session.task = None + session.turn_in_progress = False + session.last_activity = datetime.now() + await self._persist_session_async(session) + async def close(self) -> None: await self._close_resource_async(self._store) await self._close_resource_async(self._pg_pool) 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 268ce54ec..89edfb3e1 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -364,6 +364,39 @@ def close(self) -> None: assert close_thread_ids[0] != loop_thread_id +@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() + + 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.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()) diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 3d0ff6230..a7336120d 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -234,6 +234,69 @@ async def close(self) -> None: 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] = [] From 6b9ee98e13034f1337952377b24fa5715470e861 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 10:04:45 +0800 Subject: [PATCH 17/42] fix(ui): harden async session store access Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 8 ++- .../src/coding_agent/ui/session_store.py | 9 +++- .../ui/test_session_manager_public_api.py | 50 +++++++++++++++++++ .../tests/ui/test_session_persistence.py | 28 +++++++++++ 4 files changed, 89 insertions(+), 6 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 411822ea8..cc07d373c 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -1132,14 +1132,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) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 673e59e6c..e8cd0d93d 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -203,7 +203,14 @@ def _run_loop(self) -> None: self._loop.run_forever() def _run_sync(self, operation: Coroutine[object, object, object]) -> object: - future = asyncio.run_coroutine_threadsafe(operation, self._loop) + 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: 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 89edfb3e1..b8c8fe6c5 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -443,6 +443,56 @@ async def fake_persist_session_async(current_session: Session) -> None: 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: diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index a7336120d..0eed7e6a3 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -567,6 +567,34 @@ async def close(self) -> None: 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 ") From 581903d57ab48d3b6e0812e96359e40d709a3815 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 10:23:57 +0800 Subject: [PATCH 18/42] fix(ui): harden pg session loop lifecycle Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 17 ++++- .../tests/ui/test_session_persistence.py | 66 ++++++++++++++++++- 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index e8cd0d93d..ffa633957 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -199,7 +199,7 @@ def _handle_loop_start_failure(self) -> None: def _run_loop(self) -> None: asyncio.set_event_loop(self._loop) - self._loop_ready.set() + self._loop.call_soon(self._loop_ready.set) self._loop.run_forever() def _run_sync(self, operation: Coroutine[object, object, object]) -> object: @@ -375,14 +375,25 @@ async def _close_pool() -> None: try: _ = self._run_sync(_close_pool()) finally: - _ = self._loop.call_soon_threadsafe(self._loop.stop) + 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: - self._loop.close() + try: + self._loop.close() + except RuntimeError: + logger.warning( + "Postgres session metadata loop was already closed during shutdown" + ) def create_session_store( diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 0eed7e6a3..f9f621903 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -3,7 +3,7 @@ import asyncio from datetime import datetime from concurrent.futures import TimeoutError as FutureTimeoutError -from typing import cast +from typing import Callable, cast from unittest.mock import MagicMock from unittest.mock import patch @@ -311,9 +311,15 @@ async def close(self) -> 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: @@ -357,6 +363,7 @@ def is_alive(self) -> bool: 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() @@ -465,6 +472,63 @@ def fake_run_sync(operation): 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_test_does_not_start_real_thread() -> None: created_loops: list[FakeLoop] = [] created_threads: list[FakeThread] = [] From 6819b0d3c9a0ce2620e1c72f11069b6f7cc3db48 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 10:38:25 +0800 Subject: [PATCH 19/42] fix(storage): make pg checkpoints retry-safe Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/agentkit/storage/pg.py | 11 ++++- .../tests/agentkit/storage/test_pg.py | 49 +++++++++++++++++-- .../ui/test_session_manager_public_api.py | 1 - 3 files changed, 54 insertions(+), 7 deletions(-) diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py index 240b61e41..569e87948 100644 --- a/coding-agent/src/agentkit/storage/pg.py +++ b/coding-agent/src/agentkit/storage/pg.py @@ -267,6 +267,12 @@ class PGCheckpointStore: 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 " @@ -351,7 +357,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 diff --git a/coding-agent/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py index c18f319f1..77802536c 100644 --- a/coding-agent/tests/agentkit/storage/test_pg.py +++ b/coding-agent/tests/agentkit/storage/test_pg.py @@ -215,11 +215,7 @@ async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: if payload["tape_id"] == tape_id ] rows.sort( - key=lambda row: ( - cast(dict[str, object], row["meta"])["created_at"] - if isinstance(row["meta"], dict) - else "" - ) + key=lambda row: str(cast(dict[str, object], row["meta"])["created_at"]) ) return rows raise AssertionError(f"unexpected fetch query: {query}") @@ -666,6 +662,49 @@ async def test_pg_checkpoint_store_round_trip_snapshot( 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 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 b8c8fe6c5..373086a0e 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -8,7 +8,6 @@ import pytest -from agentkit.errors import ConfigError from agentkit.checkpoint import CheckpointService from agentkit.checkpoint.models import CheckpointMeta from agentkit.tape.tape import Tape From 8fc6d23c0cd8cc26701cd37f1dfc0d319149d414 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 10:51:07 +0800 Subject: [PATCH 20/42] fix(ui): preserve idle semantics on shutdown Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/coding_agent/ui/session_manager.py | 1 - coding-agent/tests/ui/test_session_manager_public_api.py | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index cc07d373c..bcc6f65a8 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -893,7 +893,6 @@ async def shutdown_session_runtime(self, session_id: str) -> None: await self._close_runtime(session) session.task = None session.turn_in_progress = False - session.last_activity = datetime.now() await self._persist_session_async(session) async def close(self) -> None: 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 373086a0e..a38994890 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -372,6 +372,7 @@ async def test_shutdown_session_runtime_preserves_persisted_session_metadata() - 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: @@ -390,6 +391,7 @@ async def close(self) -> 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 From 627ac94ec845000ca73ec7ecd1df100754c4ab94 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 15:18:02 +0800 Subject: [PATCH 21/42] fix(ui): stop persisting event queues Queue fan-out state is runtime-only and should stay in memory to avoid redundant store writes. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 4 -- .../ui/test_session_manager_public_api.py | 37 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index bcc6f65a8..0345c6a8f 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -475,7 +475,6 @@ async def add_event_queue_async( ) -> None: session = await self.get_session_async(session_id) session.event_queues.append(queue) - await self._persist_session_async(session) async def remove_event_queue_async( self, @@ -485,7 +484,6 @@ async def remove_event_queue_async( session = await self.get_session_async(session_id) if queue in session.event_queues: session.event_queues.remove(queue) - await self._persist_session_async(session) async def check_health_async(self) -> bool: return bool(await self._run_store_io(self._store.check_health)) @@ -809,7 +807,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, @@ -819,7 +816,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, 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 a38994890..19df290b8 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -156,6 +156,43 @@ 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 == [] + + def test_session_manager_uses_pg_backends_when_storage_config_requests_pg() -> None: class FakePGPool: instances: list[FakePGPool] = [] From f33080fde87774cd3645d16247c516359ca4daa8 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 15:18:06 +0800 Subject: [PATCH 22/42] fix(storage): ignore pg close errors Shutdown should keep draining the loop and thread even if the metadata pool close fails. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 8 ++- .../tests/ui/test_session_persistence.py | 57 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index ffa633957..64be08c80 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -373,7 +373,13 @@ async def _close_pool() -> None: await self._pool.close() try: - _ = self._run_sync(_close_pool()) + 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) diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index f9f621903..b0e13bfae 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -529,6 +529,63 @@ def fake_run_sync(operation): 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] = [] From 97bbf423e26b99fc4b80c6efe25ef6c4bb6b2cd5 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 15:18:10 +0800 Subject: [PATCH 23/42] fix(ui): broadcast from loaded session Use the provided session object for SSE fan-out so the hot path stays local to the connection state. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/coding_agent/ui/http_server.py | 15 ++++++++++++++- coding-agent/tests/ui/test_http_server.py | 15 +++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 7466e25d8..3823ee3f9 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -343,7 +343,20 @@ 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) + before_count = len(session.event_queues) + session.event_queues = [queue for queue in session.event_queues if not queue.full()] + pruned_count = before_count - len(session.event_queues) + if pruned_count: + logger.info( + "Pruned %d full event queue(s) for session %s", + pruned_count, + session.id, + ) + for queue in session.event_queues: + try: + await queue.put(event) + except Exception: + logger.debug("Dropping closed event queue", exc_info=True) async def _cleanup_event_queue_on_disconnect( diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index e9fb48b62..6ef429b1f 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -1173,6 +1173,21 @@ 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 + class TestWaitForApproval: """Tests for the approval wait function.""" From ec9846b8f56d29722eb5f2e2a826e736cc87a2ad Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 15:40:59 +0800 Subject: [PATCH 24/42] fix(ui): serialize async store io Session store calls now stay ordered per manager instance so executor offloading cannot race in-memory or sync-backed stores. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 6 ++-- .../ui/test_session_manager_public_api.py | 36 +++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 0345c6a8f..a61e97550 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -326,6 +326,7 @@ def __init__( 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 self._create_tape_store(data_dir) @@ -431,8 +432,9 @@ def _turn_lock_for(self, session_id: str) -> asyncio.Lock: return lock async def _run_store_io(self, func: Callable[..., T], /, *args: object) -> T: - loop = asyncio.get_running_loop() - return await loop.run_in_executor(None, partial(func, *args)) + 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 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 19df290b8..2c8acf179 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -193,6 +193,42 @@ async def test_event_queue_mutations_do_not_persist_async_runtime_only_state() - 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] = [] From 4156432e7b75cdf2f12e1ccdccc26fdbecb4bb3f Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 15:41:05 +0800 Subject: [PATCH 25/42] fix(store): use scard for redis counts Health checks no longer need to pull the full session index when Redis can return the set cardinality directly. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 8 ++ .../tests/ui/test_session_persistence.py | 89 +++++++++++++++++++ 2 files changed, 97 insertions(+) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 64be08c80..c47946014 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -38,6 +38,8 @@ def srem(self, key: str, value: str) -> object: ... def smembers(self, key: str) -> set[str] | set[bytes]: ... + def scard(self, key: str) -> int: ... + class SessionStore(Protocol): def save(self, session_id: str, data: SessionPayload) -> None: ... @@ -130,6 +132,12 @@ def list_sessions(self) -> list[str]: 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: diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index b0e13bfae..559e5cdf8 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -89,6 +89,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 @@ -112,6 +113,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() @@ -134,6 +136,93 @@ 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_pg_session_metadata_store_round_trips_session_metadata() -> None: class FakeRecord: def __init__(self, payload: dict[str, object]) -> None: From 20262b2273e4900f32582b692a217cdc77544d5f Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 20:56:01 +0800 Subject: [PATCH 26/42] fix(storage): sort pg checkpoints by timestamp Checkpoint ordering now matches datetime semantics even when ISO strings use different timezone offsets. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/agentkit/storage/pg.py | 2 +- .../tests/agentkit/storage/test_pg.py | 69 ++++++++++++++++++- 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py index 569e87948..79333575f 100644 --- a/coding-agent/src/agentkit/storage/pg.py +++ b/coding-agent/src/agentkit/storage/pg.py @@ -280,7 +280,7 @@ class PGCheckpointStore: ) _LIST_SQL: Final[str] = ( "SELECT meta FROM agent_checkpoints WHERE tape_id = $1 " - "ORDER BY meta->>'created_at'" + "ORDER BY (meta->>'created_at')::timestamptz" ) _DELETE_SQL: Final[str] = "DELETE FROM agent_checkpoints WHERE checkpoint_id = $1" diff --git a/coding-agent/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py index 77802536c..f0200ba70 100644 --- a/coding-agent/tests/agentkit/storage/test_pg.py +++ b/coding-agent/tests/agentkit/storage/test_pg.py @@ -214,9 +214,18 @@ async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: for payload in self.checkpoints.values() if payload["tape_id"] == tape_id ] - rows.sort( - key=lambda row: str(cast(dict[str, object], row["meta"])["created_at"]) - ) + 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}") @@ -384,6 +393,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: From 6225db91c27f1f910bdd3cadafa96d6f6c65719b Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 20:56:10 +0800 Subject: [PATCH 27/42] fix(ui): prune blocked event queues SSE fan-out now fails fast on slow or broken consumers so request handlers never wait on per-client queue backpressure. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 34 +++++++++++++------ coding-agent/tests/ui/test_http_server.py | 33 ++++++++++++++++++ 2 files changed, 57 insertions(+), 10 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 3823ee3f9..caef2e43c 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -343,21 +343,35 @@ 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.""" - before_count = len(session.event_queues) - session.event_queues = [queue for queue in session.event_queues if not queue.full()] - pruned_count = before_count - len(session.event_queues) - if pruned_count: - logger.info( - "Pruned %d full event queue(s) for session %s", - pruned_count, - session.id, - ) + active_queues: list[asyncio.Queue[dict[str, str]]] = [] + full_pruned_count = 0 + failed_pruned_count = 0 + for queue in session.event_queues: try: - await queue.put(event) + 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, diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 6ef429b1f..6b0d1ba4e 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -8,6 +8,7 @@ import sys import types from datetime import datetime, timedelta +from typing import cast from unittest.mock import patch import pytest @@ -1188,6 +1189,38 @@ async def test_broadcast_uses_provided_session_without_manager_lookup(self): 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.""" From ef19d73ab874aad10494adff22cf1c9795de9f8a Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 20:56:17 +0800 Subject: [PATCH 28/42] fix(store): relax optional redis scard typing Clients that only expose set membership operations should still satisfy the Redis session store contract because count_sessions already falls back cleanly. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 2 - .../tests/ui/test_session_persistence.py | 39 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index c47946014..021fed13e 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -38,8 +38,6 @@ def srem(self, key: str, value: str) -> object: ... def smembers(self, key: str) -> set[str] | set[bytes]: ... - def scard(self, key: str) -> int: ... - class SessionStore(Protocol): def save(self, session_id: str, data: SessionPayload) -> None: ... diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 559e5cdf8..04c3d7d8e 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -223,6 +223,45 @@ def smembers(self, key: str) -> set[str]: 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: From deb44f72b73f7d0d772ef3981267df6705bac219 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 22:21:13 +0800 Subject: [PATCH 29/42] fix(ui): harden session shutdown teardown Teardown now waits for async runtime cleanup and refuses to remove a session while its turn task survives cancellation. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 19 +++- .../tests/ui/test_session_manager_runtime.py | 91 +++++++++++++++++++ 2 files changed, 107 insertions(+), 3 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index a61e97550..b7b54b18c 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -402,6 +402,9 @@ def _create_checkpoint_store(self, data_dir: Path) -> CheckpointStore: 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) @@ -411,12 +414,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 @@ -500,7 +505,7 @@ async def _close_resource_async(self, resource: object) -> None: async def remove_session_async(self, session_id: str) -> None: session = await self.get_session_async(session_id) - self._close_runtime_sync_safe(session) + 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) @@ -871,6 +876,10 @@ 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" + ) await self.remove_session_async(session_id) @@ -887,6 +896,10 @@ async def shutdown_session_runtime(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" + ) await self._close_runtime(session) session.task = 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..56308691e 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 @@ -271,6 +273,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()) From 855b080dbceffc504cc00ca412196877cbb53734 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 22:21:54 +0800 Subject: [PATCH 30/42] fix(ui): make event cleanup best effort Disconnect and close flows now treat missing sessions as expected teardown races and stop reporting a closed session when manager shutdown actually fails. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 35 +++++++++++++++---- coding-agent/tests/ui/test_http_server.py | 32 +++++++++++++++++ 2 files changed, 60 insertions(+), 7 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index caef2e43c..e0bd134a1 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -377,7 +377,16 @@ async def _cleanup_event_queue_on_disconnect( session_id: str, queue: asyncio.Queue[dict[str, str]], ) -> None: - await asyncio.shield(session_manager.remove_event_queue_async(session_id, queue)) + 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: @@ -639,13 +648,24 @@ 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 session_id not in session_manager._session_cache: + break + current_session = session_manager._session_cache.get(session_id) + if ( + current_session is not None + and queue not in current_session.event_queues + ): + break # Send keepalive yield {"event": "ping", "data": ""} except asyncio.CancelledError: # Client disconnected - await _cleanup_event_queue_on_disconnect(session_id, queue) raise + finally: + await _cleanup_event_queue_on_disconnect(session_id, queue) return EventSourceResponse(event_generator()) @@ -771,16 +791,17 @@ async def close_session( raise HTTPException(status_code=404, detail="Session not found") session = await session_manager.get_session_async(session_id) + + try: + await session_manager.close_session(session_id) + except Exception as exc: + raise HTTPException(status_code=500, detail=str(exc)) from exc + 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) diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 6b0d1ba4e..67f77e3a3 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -789,6 +789,21 @@ async def fake_remove_event_queue_async(current_session_id: str, queue) -> None: 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) + class TestGetSession: """Tests for get session endpoint.""" @@ -852,6 +867,23 @@ 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"] == "close exploded" + class TestSessionTimeout: """Tests for session idle timeout.""" From d95224d8e30649748fdda1f1b3d5b453b728ce24 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sat, 18 Apr 2026 22:28:38 +0800 Subject: [PATCH 31/42] fix(store): wait for timed-out cancels Timed-out PG metadata operations now request cancellation and perform a brief completion handshake so callers do not assume the coroutine stopped instantly. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 7 ++++ .../tests/ui/test_session_persistence.py | 40 ++++++++++++++++++- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 021fed13e..5e5706e24 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -220,7 +220,14 @@ def _run_sync(self, operation: Coroutine[object, object, object]) -> object: 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" diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index 04c3d7d8e..de3fddfa4 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -4,7 +4,7 @@ from datetime import datetime from concurrent.futures import TimeoutError as FutureTimeoutError from typing import Callable, cast -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call from unittest.mock import patch import pytest @@ -816,6 +816,44 @@ async def close(self) -> None: 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 ): From 94312775d875a9fc591f36de87f640eb03e1f642 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 00:02:07 +0800 Subject: [PATCH 32/42] fix(ui): harden shutdown and SSE checks Server shutdown now continues past per-session teardown failures and SSE keepalive termination uses public session APIs instead of cache internals. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 23 ++-- coding-agent/tests/ui/test_http_server.py | 115 ++++++++++++++++++ 2 files changed, 130 insertions(+), 8 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index e0bd134a1..c64eab2e2 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -106,9 +106,18 @@ async def lifespan(app: FastAPI): pass # Close all sessions - for session_id in await session_manager.list_sessions_async(): - await session_manager.shutdown_session_runtime(session_id) - await session_manager.close() + 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") @@ -651,12 +660,10 @@ async def event_generator() -> AsyncIterator[dict[str, str]]: if event.get("event") == "SessionClosed": break except asyncio.TimeoutError: - if session_id not in session_manager._session_cache: + if not await session_manager.has_session_async(session_id): break - current_session = session_manager._session_cache.get(session_id) - if ( - current_session is not None - and queue not in current_session.event_queues + if not await session_manager.has_event_queue_async( + session_id, queue ): break # Send keepalive diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 67f77e3a3..f4bcba390 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -14,6 +14,7 @@ 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 @@ -27,6 +28,7 @@ SESSION_IDLE_TIMEOUT_MINUTES, _cleanup_event_queue_on_disconnect, _broadcast_event, + get_events, _session_to_dict, stream_wire_messages, _wire_message_to_event, @@ -804,6 +806,119 @@ async def fake_remove_event_queue_async(current_session_id: str, queue) -> None: 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 = 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 + ) + + +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.""" From 9038e939b45aa3740ed1f207cdd22dccbfcef887 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 00:02:16 +0800 Subject: [PATCH 33/42] fix(ui): clarify pg backend dependency errors PG storage setup now points operators at missing optional dependencies and fails clearly when the backend module is present but missing required exports. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 23 +++++++++++++- .../tests/ui/test_session_manager_runtime.py | 31 ++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index b7b54b18c..09b3b6a93 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -281,8 +281,21 @@ def _load_pg_storage_types() -> tuple[Any, Any, Any]: pg_module = importlib.import_module("agentkit.storage.pg") except ImportError as exc: raise RuntimeError( - "PG backend is not available; add agentkit.storage.pg before using tape_backend='pg'" + "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"), @@ -465,6 +478,14 @@ async def has_session_async(self, session_id: str) -> bool: 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) diff --git a/coding-agent/tests/ui/test_session_manager_runtime.py b/coding-agent/tests/ui/test_session_manager_runtime.py index 56308691e..d358e6f57 100644 --- a/coding-agent/tests/ui/test_session_manager_runtime.py +++ b/coding-agent/tests/ui/test_session_manager_runtime.py @@ -17,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 @@ -59,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() From 75a5c823962e158a686e942abcb7fadaab543265 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 00:33:20 +0800 Subject: [PATCH 34/42] docs(adr): align phase1 acceptance criteria Replace placeholder test names with the actual verification coverage used in this PR so the ADR remains auditable. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- ...lete-phase1-postgresql-http-session-persistence.md | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) 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 index cd6b207df..3dfd2e61a 100644 --- 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 @@ -32,12 +32,11 @@ Specifically: ## Acceptance Criteria -- [ ] `test_pg_checkpoint_store_round_trip_snapshot` -- [ ] `test_pg_checkpoint_store_list_by_tape_returns_sorted_meta` -- [ ] `test_create_persistence_bundle_builds_pg_backends_from_storage_config` -- [ ] `test_session_manager_defaults_use_configured_pg_persistence_bundle` -- [ ] `test_http_server_uses_configured_persistence_bundle_for_global_session_manager` -- [ ] `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 -k "pg or checkpoint or persistence_bundle or session_manager_defaults" -v` +- [ ] 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 From d995dd320f453a540eca98816bd8fe026d6b0dbe Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 01:17:56 +0800 Subject: [PATCH 35/42] fix(ui): return 404 when sessions disappear during close Map the close-session race to not found instead of surfacing a 500. Keep the import-fallback test from polluting module state when it reloads the HTTP server. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 2 + coding-agent/tests/ui/test_http_server.py | 54 ++++++++++++++++--- 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index c64eab2e2..5eae3f163 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -801,6 +801,8 @@ async def close_session( try: 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: raise HTTPException(status_code=500, detail=str(exc)) from exc diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index f4bcba390..da0a73bc1 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -999,6 +999,39 @@ async def failing_close_session(current_session_id: str) -> None: assert response.status_code == 500 assert response.json()["detail"] == "close exploded" + 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}" + class TestSessionTimeout: """Tests for session idle timeout.""" @@ -1469,14 +1502,23 @@ async def test_wait_for_approval_request_can_be_approved_via_http_endpoint( 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) - with patch("agentkit.config.loader.load_config") as load_config: - load_config.side_effect = ConfigError("config file not found") - http_server = importlib.import_module("coding_agent.ui.http_server") - - assert http_server._load_storage_config() == {} - assert http_server.session_manager._storage_config == {} + try: + with patch("agentkit.config.loader.load_config") as load_config: + load_config.side_effect = ConfigError("config file not found") + 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 class TestIntegration: From ef6bf9a11420abff58401035f60fe31f0b659da9 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 01:18:00 +0800 Subject: [PATCH 36/42] fix(ui): reuse shared pg pool for session storage Keep the injected pool path from loading PG types again and pass the shared pool into the HTTP session store. Refresh the regression coverage so it asserts the real shared-pool contract. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 6 ++- .../ui/test_session_manager_public_api.py | 37 ++++++++++++++++++- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 09b3b6a93..aef88dccf 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -352,10 +352,11 @@ def __init__( self._create_agent = create_agent_fn def _get_pg_pool(self) -> AsyncPGSessionPool: - PGPool, _, _ = _load_pg_storage_types() 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']") @@ -384,10 +385,11 @@ def _create_http_session_store(self) -> SessionStore: else None ) dsn = self._storage_config.get("dsn") + pg_pool = self._get_pg_pool() if backend == "pg" else None return create_session_store( backend=backend, dsn=dsn if isinstance(dsn, str) else None, - pg_pool=None, + pg_pool=pg_pool, ) def _create_tape_store(self, data_dir: Path) -> TapeStore: 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 2c8acf179..1727be2cd 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -278,7 +278,7 @@ def __init__(self, *, pool: FakePGPool) -> None: assert len(FakePGPool.instances) == 2 -def test_session_manager_uses_dedicated_pg_pool_for_http_session_store() -> None: +def test_session_manager_reuses_pg_pool_for_http_session_store() -> None: class FakePGPool: instances: list[FakePGPool] = [] @@ -314,7 +314,7 @@ def __init__(self, *, pool: FakePGPool) -> None: 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 create_store.call_args.kwargs["pg_pool"] is FakePGPool.instances[0] assert len(FakePGPool.instances) == 1 @@ -389,6 +389,39 @@ def __init__(self, *, dsn: str) -> None: 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_reuses_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 injected_pool + + @pytest.mark.asyncio async def test_remove_session_async_loads_store_once_when_not_cached() -> None: class CountingStore(InMemorySessionStore): From 6b3266480db3570500b1286d96c76a1d80573218 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 01:50:15 +0800 Subject: [PATCH 37/42] fix(ui): close session without stale existence check Collapse session lookup and close into one KeyError-handled path so sessions deleted between checks still return 404. Add a regression covering the retrieval-time race before broadcast. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 6 +---- coding-agent/tests/ui/test_http_server.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 5eae3f163..412c23ec5 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -794,12 +794,8 @@ async def close_session( api_key: str | None = Depends(verify_api_key), ) -> CloseSessionResponse: """Close session and release resources.""" - if not await session_manager.has_session_async(session_id): - raise HTTPException(status_code=404, detail="Session not found") - - session = await session_manager.get_session_async(session_id) - 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 diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index da0a73bc1..d8844b948 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -1032,6 +1032,32 @@ async def disappearing_close_session(current_session_id: str) -> None: 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}" + class TestSessionTimeout: """Tests for session idle timeout.""" From 54e4ee614b26813996a58188e14455dd6ea102a3 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 11:20:40 +0800 Subject: [PATCH 38/42] fix(ui): fail fast on invalid redis config Stop swallowing ConfigError during Redis session store setup so malformed configuration no longer silently downgrades persistence. Keep the connection-failure fallback and add a regression for the config-error path. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 2 +- .../ui/test_session_manager_public_api.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 5e5706e24..ee4db5a58 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -458,7 +458,7 @@ def create_session_store( redis_client_factory=redis_client_factory, ) return RedisSessionStore(client=client, redis_url=resolved_redis_url) - except Exception as exc: + except OSError as exc: safe_exc = redact_sensitive_text(str(exc)) logger.warning( "Redis session store unavailable at %s; falling back to in-memory store: %s", 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 1727be2cd..3c962836c 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -13,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, @@ -618,6 +619,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: From 26ab735a6f02d1661887069a16f979e67abb4522 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 11:47:57 +0800 Subject: [PATCH 39/42] fix(ui): preserve ownership of shared pg pools Keep PG session metadata stores from closing injected pools that are shared with the tape and checkpoint backends. Add regressions for both owned and injected pool shutdown paths so the lifecycle stays explicit. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_store.py | 13 +++-- .../tests/ui/test_session_persistence.py | 48 ++++++++++++++++++- 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index ee4db5a58..988d6ba4b 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -178,8 +178,10 @@ 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() @@ -383,7 +385,8 @@ async def _check_health() -> bool: def close(self) -> None: async def _close_pool() -> None: - await self._pool.close() + if self._owns_pool: + await self._pool.close() try: try: @@ -434,14 +437,14 @@ def create_session_store( 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 = PGPool(dsn=normalized_dsn) - return PGSessionMetadataStore( - pool=pg_pool, - ) + 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() diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index de3fddfa4..b7e56efcf 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -535,13 +535,57 @@ async def close(self) -> None: await self.pool.close() pool = FakePGPool() - store = PGSessionMetadataStore(pool=pool) + 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: @@ -928,7 +972,7 @@ async def close(self) -> None: finally: store.close() - assert pool.closed is True + assert pool.closed is False def test_rehydrate_clears_non_restart_safe_runtime_state() -> None: From d03eedd10b1e62b09c789b49997238367ac43279 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 12:10:41 +0800 Subject: [PATCH 40/42] fix(ui): harden event and storage backends Handle disappearing SSE sessions cleanly, keep Redis fallback working for redis-py connection errors while still surfacing ConfigError, and give the HTTP PG metadata store a dedicated pool to avoid cross-loop reuse with tape and checkpoint storage. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 9 ++- .../src/coding_agent/ui/session_manager.py | 3 +- .../src/coding_agent/ui/session_store.py | 5 +- coding-agent/tests/ui/test_http_server.py | 61 ++++++++++++++++++- .../ui/test_session_manager_public_api.py | 8 +-- .../tests/ui/test_session_persistence.py | 19 ++++++ 6 files changed, 94 insertions(+), 11 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 412c23ec5..ccec12336 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -662,9 +662,12 @@ async def event_generator() -> AsyncIterator[dict[str, str]]: except asyncio.TimeoutError: if not await session_manager.has_session_async(session_id): break - if not await session_manager.has_event_queue_async( - session_id, queue - ): + try: + if not await session_manager.has_event_queue_async( + session_id, queue + ): + break + except KeyError: break # Send keepalive yield {"event": "ping", "data": ""} diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index aef88dccf..d8329b9ed 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -385,11 +385,10 @@ def _create_http_session_store(self) -> SessionStore: else None ) dsn = self._storage_config.get("dsn") - pg_pool = self._get_pg_pool() if backend == "pg" else None return create_session_store( backend=backend, dsn=dsn if isinstance(dsn, str) else None, - pg_pool=pg_pool, + pg_pool=None, ) def _create_tape_store(self, data_dir: Path) -> TapeStore: diff --git a/coding-agent/src/coding_agent/ui/session_store.py b/coding-agent/src/coding_agent/ui/session_store.py index 988d6ba4b..267cf42c3 100644 --- a/coding-agent/src/coding_agent/ui/session_store.py +++ b/coding-agent/src/coding_agent/ui/session_store.py @@ -9,6 +9,7 @@ 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 @@ -461,7 +462,9 @@ def create_session_store( redis_client_factory=redis_client_factory, ) return RedisSessionStore(client=client, redis_url=resolved_redis_url) - except OSError as exc: + 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", diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index d8844b948..b357886d3 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -7,6 +7,7 @@ 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 @@ -860,7 +861,7 @@ async def fake_wait_for(awaitable, timeout): } ) response = await get_events(request, session_id, None) - event_generator = response.body_iterator + event_generator = cast(AsyncIterator[dict[str, str]], response.body_iterator) event = await anext(event_generator) assert event == {"event": "ping", "data": ""} @@ -874,6 +875,64 @@ async def fake_wait_for(awaitable, timeout): 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): 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 3c962836c..e979761b2 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -279,7 +279,7 @@ def __init__(self, *, pool: FakePGPool) -> None: assert len(FakePGPool.instances) == 2 -def test_session_manager_reuses_pg_pool_for_http_session_store() -> None: +def test_session_manager_creates_dedicated_pg_pool_for_http_session_store() -> None: class FakePGPool: instances: list[FakePGPool] = [] @@ -315,7 +315,7 @@ def __init__(self, *, pool: FakePGPool) -> None: 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 FakePGPool.instances[0] + assert create_store.call_args.kwargs["pg_pool"] is None assert len(FakePGPool.instances) == 1 @@ -405,7 +405,7 @@ def test_session_manager_reuses_injected_pg_pool_without_importing_pg_types() -> assert pool is injected_pool -def test_session_manager_http_session_store_reuses_injected_pg_pool() -> None: +def test_session_manager_http_session_store_does_not_reuse_injected_pg_pool() -> None: injected_pool = object() manager = SessionManager.__new__(SessionManager) manager._storage_config = { @@ -420,7 +420,7 @@ def test_session_manager_http_session_store_reuses_injected_pg_pool() -> None: _ = SessionManager._create_http_session_store(manager) assert create_store.call_args.kwargs["backend"] == "pg" - assert create_store.call_args.kwargs["pg_pool"] is injected_pool + assert create_store.call_args.kwargs["pg_pool"] is None @pytest.mark.asyncio diff --git a/coding-agent/tests/ui/test_session_persistence.py b/coding-agent/tests/ui/test_session_persistence.py index b7e56efcf..a086fea8d 100644 --- a/coding-agent/tests/ui/test_session_persistence.py +++ b/coding-agent/tests/ui/test_session_persistence.py @@ -80,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 From 3455afcbe367fbaad28da45d626cb8e95ce11271 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 12:47:45 +0800 Subject: [PATCH 41/42] fix(ui): hide close-session internals Return a generic 500 from close-session while logging the underlying exception server-side so unexpected failures do not leak implementation details or secrets to clients. Update the close-session regression to assert the sanitized response. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/http_server.py | 3 ++- coding-agent/tests/ui/test_http_server.py | 22 ++++++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index ccec12336..30ae768bd 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -803,7 +803,8 @@ async def close_session( except KeyError as exc: raise HTTPException(status_code=404, detail=_key_error_detail(exc)) from exc except Exception as exc: - raise HTTPException(status_code=500, detail=str(exc)) from exc + logger.exception("Unexpected error while closing session %s", session_id) + raise HTTPException(status_code=500, detail="Internal server error") from exc await _broadcast_event( session, diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index b357886d3..a42c324c8 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -1056,7 +1056,7 @@ async def failing_close_session(current_session_id: str) -> None: response = await client.delete(f"/sessions/{session_id}") assert response.status_code == 500 - assert response.json()["detail"] == "close exploded" + assert response.json()["detail"] == "Internal server error" async def test_close_session_returns_404_when_session_disappears_during_close( self, client, monkeypatch @@ -1117,6 +1117,26 @@ async def disappearing_get_session_async(current_session_id: str): 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.""" From 0a492312b0fdcaea79322eeedf74611e375d16e6 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 19 Apr 2026 14:30:09 +0800 Subject: [PATCH 42/42] fix(storage): harden pg review fixes - fail fast on invalid agent.toml storage config - keep SessionManager from closing injected pg pools - make pg tape appends use one advisory-locked insert Amp-Thread-ID: https://ampcode.com/threads/T-019da160-819b-73af-9baf-f39e0dbcfbcb Co-authored-by: Amp --- coding-agent/src/agentkit/storage/pg.py | 37 +++++------- .../src/coding_agent/ui/http_server.py | 6 +- .../src/coding_agent/ui/session_manager.py | 5 +- .../tests/agentkit/storage/test_pg.py | 59 ++++++++----------- coding-agent/tests/ui/test_http_server.py | 22 ++++++- .../ui/test_session_manager_public_api.py | 17 ++++++ 6 files changed, 86 insertions(+), 60 deletions(-) diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py index 79333575f..98d48dc53 100644 --- a/coding-agent/src/agentkit/storage/pg.py +++ b/coding-agent/src/agentkit/storage/pg.py @@ -165,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" @@ -197,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() diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index 30ae768bd..0375994c3 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -66,7 +66,11 @@ def _load_storage_config() -> dict[str, Any]: return cast( dict[str, Any], load_agent_toml(config_path).extra.get("storage", {}) ) - except (ConfigError, OSError): + 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, diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index d8329b9ed..02d5969dd 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -335,6 +335,7 @@ def __init__( ): 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] = {} @@ -362,6 +363,7 @@ def _get_pg_pool(self) -> AsyncPGSessionPool: 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: @@ -930,7 +932,8 @@ async def shutdown_session_runtime(self, session_id: str) -> None: async def close(self) -> None: await self._close_resource_async(self._store) - await self._close_resource_async(self._pg_pool) + if self._owns_pg_pool: + await self._close_resource_async(self._pg_pool) async def run_agent( self, diff --git a/coding-agent/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py index f0200ba70..e17f79a30 100644 --- a/coding-agent/tests/agentkit/storage/test_pg.py +++ b/coding-agent/tests/agentkit/storage/test_pg.py @@ -50,32 +50,20 @@ 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 + if not ( + isinstance(payload_values, list) + and all(isinstance(item, (str, dict)) for item in payload_values) ): - seq_values = second_arg - else: - raise TypeError("seq must be an int or list[int]") - - if isinstance(third_arg, (str, dict)): - payload_values = [third_arg] - elif isinstance(third_arg, list) and all( - isinstance(item, (str, dict)) for item in third_arg - ): - payload_values = third_arg - else: raise TypeError("payload must be encoded json/dict or list thereof") - - 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): + 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 ) @@ -146,19 +134,6 @@ async def execute(self, query: str, *args: object) -> str: 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" in 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 @@ -615,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") diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index a42c324c8..e03e4c22d 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -1612,7 +1612,9 @@ def test_http_server_import_falls_back_when_agent_toml_is_unreadable( try: with patch("agentkit.config.loader.load_config") as load_config: - load_config.side_effect = ConfigError("config file not found") + 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() == {} @@ -1626,6 +1628,24 @@ def test_http_server_import_falls_back_when_agent_toml_is_unreadable( 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 e979761b2..935412617 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -470,6 +470,23 @@ def close(self) -> None: 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()