beaver/beaver-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(backends,storage,core,markdown,infra): claude agent sdk backend, session store, transcript seeding, runner isolation

Browse Source
This commit is contained in:
hh
2026-08-28 01:56:39 +02:00
parent b3a584a362
commit 7424d52f88
28 changed files with 2154 additions and 875 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/beaver_gateway/storage/__init__.py
+10 -5
View File
@@ -1,14 +1,14 @@
"""SQLModel-backed persistence.
Two tables — :class:`Token`, :class:`AuditLog` — plus a thin
:class:`Database` wrapper around an async SQLAlchemy engine. The
``GatewayRuntime`` carries the handle so auth (token verify / touch)
and admin (token CRUD + audit listing) can reach it.
Tokens, audit, conversations, Agent SDK transcript entries and per-turn
usage, plus a thin :class:`Database` wrapper around an async SQLAlchemy
engine. ``GatewayRuntime`` carries the handle.
"""
from beaver_gateway.storage.db import (
Database,
append_audit,
append_usage,
create_token,
list_active_tokens,
list_audit_records,
@@ -16,13 +16,18 @@ from beaver_gateway.storage.db import (
revoke_token,
touch_token,
)
from beaver_gateway.storage.models import AuditLog, Token
from beaver_gateway.storage.models import AuditLog, Token, TranscriptEntry, Usage
from beaver_gateway.storage.session_store import PostgresSessionStore
__all__ = [
"AuditLog",
"Database",
"PostgresSessionStore",
"Token",
"TranscriptEntry",
"Usage",
"append_audit",
"append_usage",
"create_token",
"list_active_tokens",
"list_audit_records",
src/beaver_gateway/storage/db.py
+10 -1
View File
@@ -22,7 +22,7 @@ from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession
from beaver_gateway.storage.models import AuditLog, Token
from beaver_gateway.storage.models import AuditLog, Token, Usage
if TYPE_CHECKING:
from collections.abc import Sequence
@@ -183,9 +183,18 @@ async def list_audit_records(
return result.all()
# ---- Usage --------------------------------------------------------------
async def append_usage(session: AsyncSession, row: Usage) -> None:
session.add(row)
await session.commit()
__all__ = [
"Database",
"append_audit",
"append_usage",
"create_token",
"list_active_tokens",
"list_audit_records",
src/beaver_gateway/storage/models.py
+74 -2
View File
@@ -23,8 +23,10 @@ needs an id uses ``Optional[int]`` so SQLAlchemy can autoincrement.
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import UniqueConstraint
from sqlalchemy import JSON, Column, Index, UniqueConstraint, text
from sqlalchemy.dialects.postgresql import JSONB
from sqlmodel import Field, SQLModel
@@ -92,6 +94,7 @@ class Conversation(SQLModel, table=True):
frontend: str = Field(index=True)
external_id: str = Field(index=True)
agent_name: str = Field(index=True)
session_id: str | None = Field(default=None, index=True)
created_at: datetime = Field(default_factory=_utcnow)
updated_at: datetime = Field(default_factory=_utcnow)
@@ -125,4 +128,73 @@ class ConversationMessage(SQLModel, table=True):
created_at: datetime = Field(default_factory=_utcnow)
__all__ = ["AuditLog", "Conversation", "ConversationMessage", "Token"]
class TranscriptEntry(SQLModel, table=True):
"""One Agent SDK transcript line, mirrored from the session store protocol.
``(project_key, session_id, subpath)`` is the ``SessionKey``; ``subpath``
is ``""`` for the main transcript so the unique constraints stay simple.
``uuid`` is the SDK's idempotency key - the partial unique index rejects
a replayed batch, entries without a uuid are appended as-is.
"""
__tablename__ = "transcript_entries"
__table_args__ = (
UniqueConstraint(
"project_key", "session_id", "subpath", "seq", name="uq_transcript_seq"
),
Index(
"uq_transcript_uuid",
"project_key",
"session_id",
"subpath",
"uuid",
unique=True,
postgresql_where=text("uuid IS NOT NULL"),
sqlite_where=text("uuid IS NOT NULL"),
),
)
id: int | None = Field(default=None, primary_key=True)
project_key: str = Field(index=True)
session_id: str = Field(index=True)
subpath: str = Field(default="")
seq: int
uuid: str | None = Field(default=None)
entry: dict[str, Any] = Field(
sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=False)
)
origin: str | None = Field(default=None)
source: str | None = Field(default=None)
turn_id: str | None = Field(default=None)
created_at: datetime = Field(default_factory=_utcnow)
class Usage(SQLModel, table=True):
"""Per-turn token accounting from ``ResultMessage.usage``."""
__tablename__ = "usage"
id: int | None = Field(default=None, primary_key=True)
ts: datetime = Field(default_factory=_utcnow, index=True)
agent_name: str = Field(index=True)
conversation_id: str | None = Field(default=None, index=True)
session_id: str | None = Field(default=None, index=True)
model: str
effort: str | None = Field(default=None)
input_tokens: int = 0
output_tokens: int = 0
cache_read_tokens: int = 0
cache_creation_tokens: int = 0
cost_usd: float | None = Field(default=None)
duration_ms: int | None = Field(default=None)
num_turns: int | None = Field(default=None)
__all__ = [
"AuditLog",
"Conversation",
"ConversationMessage",
"Token",
"TranscriptEntry",
"Usage",
]
src/beaver_gateway/storage/session_store.py
+126
View File
@@ -0,0 +1,126 @@
"""``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 ""