From bd65d536d62c1021b9bd6ec94e1221bcacf57c85 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Sun, 12 Apr 2026 01:25:55 +0800 Subject: [PATCH] feat(storage): add PostgreSQL storage backend Add a PostgreSQL-backed storage path so tapes and sessions can move beyond JSONL files when a durable backend is needed. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- coding-agent/src/agentkit/storage/pg.py | 260 ++++++++++ .../tests/agentkit/storage/test_pg.py | 461 ++++++++++++++++++ .../plugins/test_storage_factory.py | 139 ++++++ 3 files changed, 860 insertions(+) create mode 100644 coding-agent/src/agentkit/storage/pg.py create mode 100644 coding-agent/tests/agentkit/storage/test_pg.py create mode 100644 coding-agent/tests/coding_agent/plugins/test_storage_factory.py diff --git a/coding-agent/src/agentkit/storage/pg.py b/coding-agent/src/agentkit/storage/pg.py new file mode 100644 index 000000000..373f90197 --- /dev/null +++ b/coding-agent/src/agentkit/storage/pg.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import asyncio +import importlib +import json +from collections.abc import Awaitable, Callable +from typing import Final, Protocol, cast + + +class AsyncPGPool(Protocol): + async def execute(self, query: str, *args: object) -> str: ... + + async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None: ... + + async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: ... + + async def acquire(self) -> LockConnection: ... + + async def release(self, connection: LockConnection) -> None: ... + + async def close(self) -> None: ... + + +class LockConnection(Protocol): + async def execute(self, query: str, *args: object) -> object: ... + + +class LockPool(Protocol): + async def acquire(self) -> LockConnection: ... + + async def release(self, connection: LockConnection) -> None: ... + + +PoolFactory = Callable[..., Awaitable[AsyncPGPool]] + + +class PGPool: + def __init__( + self, + *, + dsn: str, + min_size: int = 1, + max_size: int = 10, + pool_factory: PoolFactory | None = None, + ) -> None: + self._dsn: str = dsn + self._min_size: int = min_size + self._max_size: int = max_size + self._pool_factory: PoolFactory | None = pool_factory + self._pool: AsyncPGPool | None = None + self._pool_lock: asyncio.Lock = asyncio.Lock() + + async def get_pool(self) -> AsyncPGPool: + if self._pool is None: + async with self._pool_lock: + if self._pool is None: + factory = self._pool_factory or _load_asyncpg_pool_factory() + self._pool = await factory( + dsn=self._dsn, + min_size=self._min_size, + max_size=self._max_size, + ) + return self._pool + + async def close(self) -> None: + if self._pool is None: + return + await self._pool.close() + self._pool = None + + async def acquire(self) -> LockConnection: + pool = await self.get_pool() + return await pool.acquire() + + async def release(self, connection: LockConnection) -> None: + pool = await self.get_pool() + await pool.release(connection) + + +class PGSessionStore: + _CREATE_TABLE_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS agent_sessions ( + session_id TEXT PRIMARY KEY, + payload JSONB NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """ + + _UPSERT_SQL: Final[str] = """ + INSERT INTO agent_sessions (session_id, payload) + VALUES ($1, $2::jsonb) + ON CONFLICT (session_id) + DO UPDATE SET payload = EXCLUDED.payload, updated_at = NOW() + """ + + _SELECT_SQL: Final[str] = "SELECT payload FROM agent_sessions WHERE session_id = $1" + _LIST_SQL: Final[str] = "SELECT session_id FROM agent_sessions ORDER BY session_id" + _DELETE_SQL: Final[str] = "DELETE FROM agent_sessions WHERE session_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_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)) + + async def load_session(self, session_id: str) -> dict[str, object] | None: + pool = await self._ensure_schema() + row = await pool.fetchrow(self._SELECT_SQL, session_id) + if row is None: + return None + + payload = row.get("payload") + if payload is None: + return None + if not isinstance(payload, dict): + raise TypeError("postgres session payload must decode to a dict") + return cast(dict[str, object], payload) + + async def list_sessions(self) -> list[str]: + pool = await self._ensure_schema() + rows = await pool.fetch(self._LIST_SQL) + session_ids: list[str] = [] + for row in rows: + session_id = row.get("session_id") + if not isinstance(session_id, str): + raise TypeError("postgres session row must include string session_id") + session_ids.append(session_id) + return session_ids + + async def delete_session(self, session_id: str) -> None: + pool = await self._ensure_schema() + _ = await pool.execute(self._DELETE_SQL, session_id) + + +class PGTapeStore: + _CREATE_TABLE_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS agent_tapes ( + tape_id TEXT NOT NULL, + seq INTEGER NOT NULL, + entry JSONB NOT NULL, + 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] = ( + "INSERT INTO agent_tapes (tape_id, seq, entry) VALUES ($1, $2, $3::jsonb)" + ) + _LOAD_SQL: Final[str] = ( + "SELECT entry FROM agent_tapes WHERE tape_id = $1 ORDER BY seq" + ) + _LIST_SQL: Final[str] = "SELECT DISTINCT tape_id FROM agent_tapes ORDER BY tape_id" + + 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, tape_id: str, entries: list[dict[str, object]]) -> None: + if not entries: + 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") + + for offset, entry in enumerate(entries, start=1): + _ = await pool.execute( + self._INSERT_SQL, + tape_id, + max_seq + offset, + json.dumps(entry), + ) + + async def load(self, tape_id: str) -> list[dict[str, object]]: + pool = await self._ensure_schema() + rows = await pool.fetch(self._LOAD_SQL, tape_id) + loaded: list[dict[str, object]] = [] + for row in rows: + entry = row.get("entry") + if not isinstance(entry, dict): + raise TypeError("postgres tape row must include dict entry") + loaded.append(cast(dict[str, object], entry)) + return loaded + + async def list_ids(self) -> list[str]: + pool = await self._ensure_schema() + rows = await pool.fetch(self._LIST_SQL) + tape_ids: list[str] = [] + for row in rows: + tape_id = row.get("tape_id") + if not isinstance(tape_id, str): + raise TypeError("postgres tape row must include string tape_id") + tape_ids.append(tape_id) + return tape_ids + + +class PGSessionLock: + def __init__(self, *, pool: LockPool) -> None: + self._pool: LockPool = pool + self._conn: LockConnection | None = None + + async def acquire(self, session_id: str) -> None: + if self._conn is not None: + raise RuntimeError("PGSessionLock.acquire() called while already held") + + conn = await self._pool.acquire() + try: + _ = await conn.execute("SELECT pg_advisory_lock(hashtext($1))", session_id) + except BaseException: + await self._pool.release(conn) + raise + self._conn = conn + + async def release(self) -> None: + if self._conn is None: + return + try: + _ = await self._conn.execute("SELECT pg_advisory_unlock_all()") + finally: + await self._pool.release(self._conn) + self._conn = None + + +def _load_asyncpg_pool_factory() -> PoolFactory: + try: + asyncpg_module = importlib.import_module("asyncpg") + except ModuleNotFoundError as exc: + raise ImportError( + "asyncpg is required for PostgreSQL storage backends" + ) from exc + + create_pool = getattr(asyncpg_module, "create_pool", None) + if not callable(create_pool): + raise ImportError("asyncpg does not expose create_pool") + return cast(PoolFactory, create_pool) diff --git a/coding-agent/tests/agentkit/storage/test_pg.py b/coding-agent/tests/agentkit/storage/test_pg.py new file mode 100644 index 000000000..1fe626d6b --- /dev/null +++ b/coding-agent/tests/agentkit/storage/test_pg.py @@ -0,0 +1,461 @@ +import asyncio +import json +from typing import cast +from unittest.mock import AsyncMock + +import pytest +from _pytest.monkeypatch import MonkeyPatch + +from agentkit.storage.protocols import SessionStore, TapeStore +from agentkit.storage.pg import 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.closed: bool = False + self.executed: list[tuple[str, tuple[object, ...]]] = [] + + async def execute(self, query: str, *args: object) -> str: + self.executed.append((query, args)) + if "INSERT INTO agent_sessions" in query: + 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 not isinstance(payload_obj, dict): + raise TypeError("payload must decode to a dict") + self.sessions[session_id] = cast(dict[str, object], payload_obj) + return "INSERT 0 1" + if "DELETE FROM agent_sessions" in query: + (session_id,) = args + if not isinstance(session_id, str): + raise TypeError("session_id must be a string") + _ = self.sessions.pop(session_id, None) + return "DELETE 1" + if "INSERT INTO agent_tapes" in query: + tape_id, seq, payload = args + if not isinstance(tape_id, str): + raise TypeError("tape_id must be a string") + if not isinstance(seq, int): + raise TypeError("seq must be an int") + if not isinstance(payload, str): + raise TypeError("payload must be encoded json") + payload_obj = json.loads(payload) + if not isinstance(payload_obj, dict): + raise TypeError("payload must decode to a dict") + rows = self.tapes.setdefault(tape_id, []) + rows.append({"seq": seq, "entry": cast(dict[str, object], payload_obj)}) + rows.sort( + key=lambda item: item["seq"] if isinstance(item["seq"], int) else -1 + ) + return "INSERT 0 1" + if query.strip() == "SELECT 1": + return "SELECT 1" + if "CREATE TABLE IF NOT EXISTS agent_sessions" in query: + return "CREATE TABLE" + if "CREATE TABLE IF NOT EXISTS agent_tapes" 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_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 + ) + } + (session_id,) = args + if not isinstance(session_id, str): + raise TypeError("session_id must be a string") + payload = self.sessions.get(session_id) + if payload is None: + return None + return {"payload": payload} + + async def fetch(self, query: str, *args: object) -> list[dict[str, object]]: + self.executed.append((query, args)) + if "SELECT session_id FROM agent_sessions" in query: + return [ + {"session_id": session_id} + for session_id in sorted(self.sessions.keys()) + ] + if "SELECT entry 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, []) + return [ + {"entry": row["entry"]} + for row in sorted( + rows, + key=lambda r: r["seq"] if isinstance(r["seq"], int) else -1, + ) + ] + if "SELECT DISTINCT tape_id FROM agent_tapes" in query: + return [{"tape_id": tape_id} for tape_id in sorted(self.tapes.keys())] + raise AssertionError(f"unexpected fetch query: {query}") + + async def close(self) -> None: + self.closed = True + + async def acquire(self) -> "FakePool": + return self + + async def release(self, connection: object) -> None: + if connection is not self: + raise AssertionError("unexpected connection released") + + +class TestPGPool: + @pytest.mark.asyncio + async def test_creates_pool_lazily_and_reuses_instance(self): + created: list[FakePool] = [] + + async def fake_pool_factory(**kwargs: object) -> FakePool: + assert kwargs["dsn"] == "postgresql://example" + assert kwargs["min_size"] == 2 + assert kwargs["max_size"] == 5 + pool = FakePool() + created.append(pool) + return pool + + pool = PGPool( + dsn="postgresql://example", + min_size=2, + max_size=5, + pool_factory=fake_pool_factory, + ) + + first = await pool.get_pool() + second = await pool.get_pool() + + assert first is second + assert len(created) == 1 + + @pytest.mark.asyncio + async def test_close_closes_underlying_pool(self): + fake_pool = FakePool() + + async def fake_pool_factory(**_: object) -> FakePool: + return fake_pool + + pool = PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + await pool.get_pool() + + await pool.close() + + assert fake_pool.closed is True + + @pytest.mark.asyncio + async def test_concurrent_get_pool_reuses_single_factory_result(self): + created: list[FakePool] = [] + started = asyncio.Event() + release = asyncio.Event() + + async def fake_pool_factory(**_: object) -> FakePool: + pool = FakePool() + created.append(pool) + started.set() + await release.wait() + return pool + + pool = PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + + first_task = asyncio.create_task(pool.get_pool()) + await started.wait() + second_task = asyncio.create_task(pool.get_pool()) + release.set() + + first, second = await asyncio.gather(first_task, second_task) + + assert first is second + assert len(created) == 1 + + @pytest.mark.asyncio + async def test_missing_asyncpg_raises_clear_error(self, monkeypatch: MonkeyPatch): + def fake_import_module(name: str): + if name == "asyncpg": + raise ModuleNotFoundError(name) + return __import__(name) + + monkeypatch.setattr("importlib.import_module", fake_import_module) + + pool = PGPool(dsn="postgresql://example") + + with pytest.raises(ImportError, match="asyncpg is required"): + await pool.get_pool() + + @pytest.mark.asyncio + async def test_acquire_delegates_to_underlying_pool(self): + fake_pool = FakePool() + + async def fake_pool_factory(**_: object) -> FakePool: + return fake_pool + + pool = PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + + conn = await pool.acquire() + + assert conn is fake_pool + + @pytest.mark.asyncio + async def test_release_delegates_to_underlying_pool(self): + fake_pool = FakePool() + + async def fake_pool_factory(**_: object) -> FakePool: + return fake_pool + + pool = PGPool(dsn="postgresql://example", pool_factory=fake_pool_factory) + await pool.get_pool() + + await pool.release(fake_pool) + + +class TestPGSessionStore: + @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) -> PGSessionStore: + return PGSessionStore(pool=pool) + + def test_satisfies_protocol(self, store: PGSessionStore): + assert isinstance(store, SessionStore) + + @pytest.mark.asyncio + async def test_save_and_load_session(self, store: PGSessionStore): + await store.save_session("ses-1", {"model": "gpt-4.1", "turns": 3}) + + data = await store.load_session("ses-1") + + assert data == {"model": "gpt-4.1", "turns": 3} + + @pytest.mark.asyncio + async def test_load_missing_returns_none(self, store: PGSessionStore): + assert await store.load_session("missing") is None + + @pytest.mark.asyncio + async def test_save_overwrites_existing(self, store: PGSessionStore): + await store.save_session("ses-1", {"version": 1}) + await store.save_session("ses-1", {"version": 2}) + + data = await store.load_session("ses-1") + + assert data == {"version": 2} + + @pytest.mark.asyncio + async def test_list_sessions_returns_sorted_ids(self, store: PGSessionStore): + await store.save_session("b", {"x": 2}) + await store.save_session("a", {"x": 1}) + + session_ids = await store.list_sessions() + + assert session_ids == ["a", "b"] + + @pytest.mark.asyncio + async def test_delete_session(self, store: PGSessionStore): + await store.save_session("ses-1", {"x": 1}) + + await store.delete_session("ses-1") + + assert await store.load_session("ses-1") is None + + @pytest.mark.asyncio + async def test_schema_created_once( + self, store: PGSessionStore, fake_pool: FakePool + ): + await store.save_session("one", {"x": 1}) + await store.save_session("two", {"x": 2}) + + schema_calls = [ + query + for query, _ in fake_pool.executed + if "CREATE TABLE IF NOT EXISTS agent_sessions" in query + ] + assert len(schema_calls) == 1 + + +class TestPGTapeStore: + @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) -> PGTapeStore: + return PGTapeStore(pool=pool) + + def test_satisfies_protocol(self, store: PGTapeStore): + assert isinstance(store, TapeStore) + + @pytest.mark.asyncio + async def test_save_computes_seq_from_zero(self, store: PGTapeStore): + await store.save( + "tape-1", + [{"kind": "message", "payload": {"role": "user", "content": "hi"}}], + ) + + rows = await store.load("tape-1") + + assert len(rows) == 1 + assert rows[0]["kind"] == "message" + + @pytest.mark.asyncio + async def test_save_appends_after_existing(self, store: PGTapeStore): + await store.save( + "tape-1", + [{"kind": "message", "payload": {"role": "user", "content": "a"}}], + ) + await store.save( + "tape-1", + [{"kind": "message", "payload": {"role": "assistant", "content": "b"}}], + ) + + rows = await store.load("tape-1") + + assert len(rows) == 2 + first_payload = rows[0].get("payload") + second_payload = rows[1].get("payload") + assert isinstance(first_payload, dict) + assert isinstance(second_payload, dict) + assert first_payload["content"] == "a" + assert second_payload["content"] == "b" + + @pytest.mark.asyncio + async def test_save_empty_entries_is_noop(self, store: PGTapeStore): + await store.save("tape-1", []) + + assert await store.load("tape-1") == [] + + @pytest.mark.asyncio + async def test_load_empty_tape(self, store: PGTapeStore): + result = await store.load("nonexistent") + assert result == [] + + @pytest.mark.asyncio + async def test_list_ids(self, store: PGTapeStore): + await store.save( + "tape-a", [{"kind": "message", "payload": {"role": "user", "content": "a"}}] + ) + await store.save( + "tape-b", [{"kind": "message", "payload": {"role": "user", "content": "b"}}] + ) + + result = await store.list_ids() + assert result == ["tape-a", "tape-b"] + + +class MockPoolForLock: + def __init__(self) -> None: + self._conn = FakePoolConnection() + self.release = AsyncMock() + + async def acquire(self) -> FakePoolConnection: + return self._conn + + +class FakePoolConnection: + def __init__(self) -> None: + self.execute = AsyncMock() + + +class TestPGSessionLock: + @pytest.fixture + def mock_pool(self) -> MockPoolForLock: + return MockPoolForLock() + + @pytest.fixture + def lock(self, mock_pool: MockPoolForLock) -> PGSessionLock: + return PGSessionLock(pool=mock_pool) + + @pytest.mark.asyncio + async def test_acquire_calls_advisory_lock( + self, lock: PGSessionLock, mock_pool: MockPoolForLock + ): + await lock.acquire("session-abc") + + calls = mock_pool._conn.execute.call_args_list + assert len(calls) == 1 + sql = calls[0][0][0] + assert "pg_advisory_lock" in sql + assert "hashtext" in sql + + @pytest.mark.asyncio + async def test_release_unlocks_and_returns_connection( + self, lock: PGSessionLock, mock_pool: MockPoolForLock + ): + await lock.acquire("session-abc") + mock_pool._conn.execute.reset_mock() + + await lock.release() + + calls = mock_pool._conn.execute.call_args_list + assert len(calls) == 1 + sql = calls[0][0][0] + assert "pg_advisory_unlock_all" in sql + mock_pool.release.assert_awaited_once_with(mock_pool._conn) + + @pytest.mark.asyncio + async def test_release_without_acquire_is_noop(self, lock: PGSessionLock): + await lock.release() + + @pytest.mark.asyncio + async def test_release_returns_connection_even_on_error( + self, lock: PGSessionLock, mock_pool: MockPoolForLock + ): + await lock.acquire("session-abc") + mock_pool._conn.execute = AsyncMock(side_effect=Exception("unlock failed")) + + with pytest.raises(Exception, match="unlock failed"): + await lock.release() + + mock_pool.release.assert_awaited_once_with(mock_pool._conn) + + @pytest.mark.asyncio + async def test_acquire_raises_when_already_held( + self, lock: PGSessionLock, mock_pool: MockPoolForLock + ): + await lock.acquire("session-abc") + + with pytest.raises(RuntimeError, match="already held"): + await lock.acquire("session-def") + + mock_pool.release.assert_not_awaited() + + @pytest.mark.asyncio + async def test_acquire_releases_connection_when_lock_sql_fails( + self, lock: PGSessionLock, mock_pool: MockPoolForLock + ): + mock_pool._conn.execute = AsyncMock(side_effect=Exception("lock failed")) + + with pytest.raises(Exception, match="lock failed"): + await lock.acquire("session-abc") + + mock_pool.release.assert_awaited_once_with(mock_pool._conn) diff --git a/coding-agent/tests/coding_agent/plugins/test_storage_factory.py b/coding-agent/tests/coding_agent/plugins/test_storage_factory.py new file mode 100644 index 000000000..4c1ed78f9 --- /dev/null +++ b/coding-agent/tests/coding_agent/plugins/test_storage_factory.py @@ -0,0 +1,139 @@ +import pytest +from pathlib import Path +from typing import cast +from unittest.mock import MagicMock, patch + +from agentkit.storage.pg import PGPool + + +class TestStoragePluginFactory: + def test_default_backend_is_jsonl(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + plugin = StoragePlugin(data_dir=tmp_path) + + assert plugin._backend == "jsonl" + + def test_jsonl_backend_creates_jsonl_store(self, tmp_path: Path): + from coding_agent.plugins.storage import JSONLTapeStore, StoragePlugin + + plugin = StoragePlugin(data_dir=tmp_path, backend="jsonl") + store = plugin._create_tape_store() + + assert isinstance(store, JSONLTapeStore) + + def test_pg_backend_without_pool_raises(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + plugin = StoragePlugin(data_dir=tmp_path, backend="pg") + + with pytest.raises(RuntimeError, match="pg_pool"): + plugin._create_tape_store() + + def test_pg_backend_with_pool_creates_pg_store(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + mock_pg_pool = MagicMock() + mock_pg_tape_store = MagicMock() + mock_pg_session_lock = MagicMock() + + plugin = StoragePlugin(data_dir=tmp_path, backend="pg", pg_pool=mock_pg_pool) + + with patch( + "coding_agent.plugins.storage._load_pg_types", + return_value=( + MagicMock(), + mock_pg_session_lock, + MagicMock(), + mock_pg_tape_store, + ), + ): + _ = plugin._create_tape_store() + mock_pg_tape_store.assert_called_once_with(pool=mock_pg_pool) + + def test_pg_session_store_with_pool(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + mock_pg_pool = MagicMock() + mock_pg_session_store = MagicMock() + mock_pg_session_lock = MagicMock() + + plugin = StoragePlugin(data_dir=tmp_path, backend="pg", pg_pool=mock_pg_pool) + + with patch( + "coding_agent.plugins.storage._load_pg_types", + return_value=( + MagicMock(), + mock_pg_session_lock, + mock_pg_session_store, + MagicMock(), + ), + ): + _ = plugin._create_session_store() + mock_pg_session_store.assert_called_once_with(pool=mock_pg_pool) + + def test_jsonl_session_store(self, tmp_path: Path): + from agentkit.storage.session import FileSessionStore + from coding_agent.plugins.storage import StoragePlugin + + plugin = StoragePlugin(data_dir=tmp_path, backend="jsonl") + store = plugin._create_session_store() + + assert isinstance(store, FileSessionStore) + + def test_provide_storage_uses_factory(self, tmp_path: Path): + from agentkit.tape.store import ForkTapeStore + from coding_agent.plugins.storage import StoragePlugin + + plugin = StoragePlugin(data_dir=tmp_path, backend="jsonl") + storage = plugin.provide_storage() + + assert isinstance(storage, ForkTapeStore) + + def test_no_lock_when_jsonl_backend(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + plugin = StoragePlugin(data_dir=tmp_path, backend="jsonl") + + assert plugin.session_lock is None + + def test_pg_lock_stored_when_pg_backend(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + mock_pg_pool = MagicMock() + + plugin = StoragePlugin(data_dir=tmp_path, backend="pg", pg_pool=mock_pg_pool) + + assert plugin._session_lock is not None + + @pytest.mark.asyncio + async def test_shutdown_closes_pg_pool_for_mixed_backend(self, tmp_path: Path): + from coding_agent.plugins.storage import StoragePlugin + + class StubPGPool: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + async def acquire(self): + raise AssertionError("acquire should not be called in this test") + + async def release(self, connection): + _ = connection + raise AssertionError("release should not be called in this test") + + pg_pool_stub = StubPGPool() + pg_pool = cast(PGPool, cast(object, pg_pool_stub)) + plugin = StoragePlugin( + data_dir=tmp_path, + backend="pg", + pg_pool=pg_pool, + config={"session_backend": "file"}, + ) + _ = plugin.hooks()["mount"]() + + await plugin.on_shutdown() + + assert pg_pool_stub.closed is True