127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
"""``SessionStore`` adapter for the Claude Agent SDK on top of :class:`Database`.
|
|
|
|
Entries are stored verbatim as JSON (``jsonb`` on Postgres), ordered by a
|
|
per-key ``seq``. ``append`` is idempotent on ``entry["uuid"]`` because the
|
|
SDK re-delivers a batch on retry; entries without a uuid are appended as
|
|
they come. Runs on SQLite too - the conformance suite uses that in tests.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
from claude_agent_sdk import SessionStore
|
|
from sqlalchemy import delete, func
|
|
from sqlmodel import col, select
|
|
|
|
from beaver_gateway.storage.models import TranscriptEntry
|
|
|
|
if TYPE_CHECKING:
|
|
from claude_agent_sdk.types import (
|
|
SessionKey,
|
|
SessionListSubkeysKey,
|
|
SessionStoreEntry,
|
|
)
|
|
|
|
from beaver_gateway.storage.db import Database
|
|
|
|
__all__ = ["PostgresSessionStore"]
|
|
|
|
|
|
class PostgresSessionStore(SessionStore):
|
|
def __init__(self, db: Database) -> None:
|
|
self._db = db
|
|
|
|
async def append(self, key: SessionKey, entries: list[SessionStoreEntry]) -> None:
|
|
if not entries:
|
|
return
|
|
project_key, session_id, subpath = _parts(key)
|
|
wanted = [
|
|
u for u in (e.get("uuid") for e in entries) if isinstance(u, str) and u
|
|
]
|
|
async with self._db.session() as session:
|
|
known: set[str] = set()
|
|
if wanted:
|
|
result = await session.exec(
|
|
select(TranscriptEntry.uuid).where(
|
|
TranscriptEntry.project_key == project_key,
|
|
TranscriptEntry.session_id == session_id,
|
|
TranscriptEntry.subpath == subpath,
|
|
col(TranscriptEntry.uuid).in_(wanted),
|
|
)
|
|
)
|
|
known = {u for u in result.all() if u is not None}
|
|
seq_result = await session.exec(
|
|
select(func.max(TranscriptEntry.seq)).where(
|
|
TranscriptEntry.project_key == project_key,
|
|
TranscriptEntry.session_id == session_id,
|
|
TranscriptEntry.subpath == subpath,
|
|
)
|
|
)
|
|
seq = seq_result.one() or 0
|
|
for entry in entries:
|
|
uuid = entry.get("uuid")
|
|
uuid_str = uuid if isinstance(uuid, str) and uuid else None
|
|
if uuid_str is not None:
|
|
if uuid_str in known:
|
|
continue
|
|
known.add(uuid_str)
|
|
seq += 1
|
|
session.add(
|
|
TranscriptEntry(
|
|
project_key=project_key,
|
|
session_id=session_id,
|
|
subpath=subpath,
|
|
seq=seq,
|
|
uuid=uuid_str,
|
|
entry=cast("dict[str, Any]", dict(entry)),
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
async def load(self, key: SessionKey) -> list[SessionStoreEntry] | None:
|
|
project_key, session_id, subpath = _parts(key)
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(TranscriptEntry.entry)
|
|
.where(
|
|
TranscriptEntry.project_key == project_key,
|
|
TranscriptEntry.session_id == session_id,
|
|
TranscriptEntry.subpath == subpath,
|
|
)
|
|
.order_by(col(TranscriptEntry.seq))
|
|
)
|
|
rows = result.all()
|
|
if not rows:
|
|
return None
|
|
return [cast("SessionStoreEntry", row) for row in rows]
|
|
|
|
async def list_subkeys(self, key: SessionListSubkeysKey) -> list[str]:
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(TranscriptEntry.subpath)
|
|
.where(
|
|
TranscriptEntry.project_key == key["project_key"],
|
|
TranscriptEntry.session_id == key["session_id"],
|
|
TranscriptEntry.subpath != "",
|
|
)
|
|
.distinct()
|
|
)
|
|
return sorted(result.all())
|
|
|
|
async def delete(self, key: SessionKey) -> None:
|
|
project_key, session_id, subpath = _parts(key)
|
|
stmt = delete(TranscriptEntry).where(
|
|
col(TranscriptEntry.project_key) == project_key,
|
|
col(TranscriptEntry.session_id) == session_id,
|
|
)
|
|
if key.get("subpath"):
|
|
stmt = stmt.where(col(TranscriptEntry.subpath) == subpath)
|
|
async with self._db.session() as session:
|
|
await session.execute(stmt) # ty: ignore[deprecated]
|
|
await session.commit()
|
|
|
|
|
|
def _parts(key: SessionKey) -> tuple[str, str, str]:
|
|
return key["project_key"], key["session_id"], key.get("subpath") or ""
|