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

feat(core,backends,frontends,storage): conversations, inject queue, session pool, gateway mcp tools, api frontend

Browse Source
This commit is contained in:
hh
2026-08-28 03:08:30 +02:00
parent ab52fdc2b8
commit e3074c266a
28 changed files with 3543 additions and 345 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/beaver_gateway/backends/claude_sdk.py
+193 -108
View File
@@ -2,11 +2,19 @@
One :class:`ClaudeSdkBackend` per :class:`ClaudeAgent`. A live session is
one ``ClaudeSDKClient`` (one claude subprocess) and runs one turn at a
time. Sessions are keyed by ``conversation_id`` when the frontend passes
one (markdown chats) or by a text-only fingerprint of ``messages[:-1]``
time; the sessions of every agent live in one shared
:class:`~beaver_gateway.core.sessions.SessionPool` that owns TTL and
memory-pressure eviction. Sessions are keyed by ``conversation_id`` when
the caller passes one or by a text-only fingerprint of ``messages[:-1]``
for stateless callers (``/v1/messages``). Without a live session the
adapter resumes ``session_id`` from the session store, or seeds the
incoming history into the store via ``core/transcript`` and resumes that.
adapter resumes ``session_id`` from the session store (after closing any
``tool_use`` left open by a crash), or seeds the incoming history into the
store via ``core/transcript`` and resumes that.
Per-turn ``**options`` beyond the protocol's: ``kind`` (conversation kind,
picks the prompt assembly and the pool TTL), ``pinned`` (never evicted),
``tools=False`` (no MCP at all - forks and jobs), ``observer`` (callback
receiving every raw SDK message, subagent ones included), ``turn_id``.
Events on the wire are the Anthropic ``MessageStreamEvent`` family: one
``message_start``/``message_stop`` envelope per turn, block indices
@@ -36,13 +44,14 @@ import uuid
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol, Self, cast
from typing import TYPE_CHECKING, Any, Self, cast
import claude_agent_sdk
from claude_agent_sdk import (
AssistantMessage,
ClaudeAgentOptions,
ClaudeSDKClient,
MirrorErrorMessage,
ResultMessage,
StreamEvent,
TextBlock,
@@ -68,14 +77,15 @@ from beaver_gateway.core.events import (
build_thinking_delta,
build_tool_use_block_start,
)
from beaver_gateway.core.transcript import build_entries
from beaver_gateway.core.sessions import Session, SessionClient, SessionPool
from beaver_gateway.core.transcript import build_entries, close_open_tool_uses
from beaver_gateway.core.turn_capture import TurnCapture, TurnUsage
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
from anthropic.types import MessageParam
from claude_agent_sdk import SessionStore
from claude_agent_sdk import McpSdkServerConfig, SessionStore
from beaver_gateway.agents.base import BaseAgent
from beaver_gateway.agents.claude import ClaudeAgent
@@ -88,6 +98,7 @@ __all__ = [
"ClaudeSdkBackend",
"RunnerConfig",
"SessionClient",
"ToolServerFactory",
"UsageSink",
"fingerprint",
]
@@ -120,7 +131,6 @@ ENV_KEEP: tuple[str, ...] = (
)
ENV_KEEP_PREFIXES: tuple[str, ...] = ("CLAUDE_", "ANTHROPIC_", "DISABLE_")
_REAP_INTERVAL = 60.0
_STOP_REASONS: dict[str, StopReason] = {
"end_turn": "end_turn",
"tool_use": "tool_use",
@@ -131,15 +141,10 @@ _STOP_REASONS: dict[str, StopReason] = {
}
class SessionClient(Protocol):
async def connect(self) -> None: ...
async def query(self, prompt: str) -> None: ...
def receive_response(self) -> AsyncIterator[Any]: ...
async def disconnect(self) -> None: ...
ClientFactory = "Callable[[ClaudeAgentOptions], SessionClient]"
UsageSink = "Callable[[UsageEvent], Awaitable[None]]"
ToolServerFactory = "Callable[[str, str], McpSdkServerConfig | None]"
"""``(conversation_key, kind) -> in-process MCP server config`` or ``None``."""
@dataclass(frozen=True, slots=True)
@@ -158,16 +163,6 @@ class UsageEvent:
usage: TurnUsage
@dataclass
class _Live:
client: SessionClient
session_id: str | None
resumed: bool
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
last_used: float = field(default_factory=time.monotonic)
turns: int = 0
class _RunnerClient(ClaudeSDKClient):
"""``ClaudeSDKClient`` that hands the materialized resume dir to the runner uid."""
@@ -194,6 +189,8 @@ class ClaudeSdkBackend:
usage_sink: Callable[[UsageEvent], Awaitable[None]] | None = None,
client_factory: Callable[[ClaudeAgentOptions], SessionClient] | None = None,
work_dir: Path | None = None,
pool: SessionPool | None = None,
tool_server: Callable[[str, str], McpSdkServerConfig | None] | None = None,
) -> None:
self._agent = agent
self._store = session_store
@@ -203,8 +200,8 @@ class ClaudeSdkBackend:
self._work_dir = work_dir or Path(tempfile.gettempdir()) / "beaver-claude"
self._servers = _mcp_servers(agent, mcp_internal_urls)
self._mcp_disallowed = _mcp_disallowed(agent, mcp_tool_names or {})
self._sessions: dict[str, _Live] = {}
self._reaper: asyncio.Task[None] | None = None
self._pool = pool if pool is not None else SessionPool()
self._tool_server = tool_server
self._uid, self._gid = _resolve_ids(self._runner.user)
self._wrapper: Path | None = None
@@ -212,34 +209,56 @@ class ClaudeSdkBackend:
def agent(self) -> ClaudeAgent:
return self._agent
@property
def pool(self) -> SessionPool:
return self._pool
@property
def sessions(self) -> dict[str, dict[str, Any]]:
now = time.monotonic()
return {
key: {
"session_id": live.session_id,
"idle_seconds": now - live.last_used,
"turns": live.turns,
"busy": live.lock.locked(),
}
for key, live in self._sessions.items()
row["key"]: row
for row in self._pool.snapshot()
if row["agent"] == self._agent.name
}
async def __aenter__(self) -> Self:
self._reaper = asyncio.create_task(self._reap_loop())
return self
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
await self.aclose()
async def aclose(self) -> None:
if self._reaper is not None:
self._reaper.cancel()
with contextlib.suppress(BaseException):
await self._reaper
self._reaper = None
for key in list(self._sessions):
await self._close(key)
await self._pool.close_all(agent=self._agent.name)
async def close(self, key: str) -> None:
await self._pool.close(key)
async def interrupt(self, key: str) -> bool:
live = self._pool.get(key)
if live is None or not live.busy:
return False
live.interrupt_requested = True
await live.client.interrupt()
return True
def live(self, key: str) -> Session | None:
return self._pool.get(key)
async def repair_session(self, session_id: str) -> int:
"""Close ``tool_use`` blocks a crash left without a result; count added."""
key = self._store_key(session_id)
entries = await self._store.load(cast("Any", key))
if not entries:
return 0
fixes = close_open_tool_uses(cast("list[Mapping[str, Any]]", entries))
if fixes:
await self._store.append(cast("Any", key), cast("Any", fixes))
_log.warning(
"session %s: closed %d open tool_use with synthetic results",
session_id,
len(fixes),
)
return len(fixes)
async def complete(
self,
@@ -250,6 +269,11 @@ class ClaudeSdkBackend:
conversation_id: str | None = None,
session_id: str | None = None,
capture: TurnCapture | None = None,
kind: str = "deep",
pinned: bool = False,
tools: bool = True,
observer: Callable[[Any], None] | None = None,
turn_id: str | None = None,
**options: Any, # noqa: ARG002 - per-request knobs are not supported
) -> AsyncIterator[MessageStreamEvent]:
if agent.name != self._agent.name:
@@ -262,37 +286,70 @@ class ClaudeSdkBackend:
prompt = _prompt_text(history[-1].get("content"))
prior = history[:-1]
key = conversation_id or fingerprint(prior)
live = await self._acquire(key, session_id=session_id, history=prior)
spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools)
live = await self._acquire(key, session_id=session_id, history=prior, spec=spec)
message_id = f"msg_{uuid.uuid4().hex}"
yield build_message_start(message_id=message_id, model=self._agent.model)
async with live.lock:
live.running_turn = turn_id or message_id
live.last_used = time.monotonic()
try:
turn = await self._run_turn(live, prompt)
turn = await self._run_turn(live, prompt, observer=observer)
except Exception:
if live.resumed and live.turns == 0:
_log.exception(
"resume of %s failed, reseeding from history", live.session_id
)
await self._close(key)
live = await self._acquire(key, session_id=None, history=prior)
async with live.lock:
turn = await self._run_turn(live, prompt)
else:
live.running_turn = None
await self._pool.close(key)
if not (live.resumed and live.turns == 0):
raise
_log.exception(
"resume of %s failed, reseeding from history", live.session_id
)
live = await self._acquire(
key, session_id=None, history=prior, spec=spec
)
async with live.lock:
live.running_turn = turn_id or message_id
turn = await self._run_turn(live, prompt, observer=observer)
for event in turn.events:
yield event
live.turns += 1
live.last_used = time.monotonic()
live.running_turn = None
usage, interrupted = await self._after_turn(
live,
turn,
conversation_id=conversation_id,
history=history,
capture=capture,
)
if turn.result is not None and turn.result.is_error and not interrupted:
msg = f"claude: {turn.result.result or turn.result.subtype}"
raise RuntimeError(msg)
yield build_message_delta(
stop_reason=turn.stop_reason, usage=_wire_usage(usage)
)
yield build_message_stop()
async def _after_turn(
self,
live: Session,
turn: _Turn,
*,
conversation_id: str | None,
history: list[dict[str, Any]],
capture: TurnCapture | None,
) -> tuple[TurnUsage, bool]:
if turn.result is not None and turn.result.session_id:
live.session_id = turn.result.session_id
if conversation_id is None:
self._rekey(key, fingerprint([*history, *turn.synthesized]))
self._rekey(live.key, fingerprint([*history, *turn.synthesized]))
usage = _usage_of(turn.result)
interrupted = live.interrupt_requested
live.interrupt_requested = False
if capture is not None:
capture.synthesized_messages = turn.synthesized
capture.session_id = live.session_id
capture.usage = usage
capture.interrupted = interrupted
if self._usage_sink is not None:
await self._usage_sink(
UsageEvent(
@@ -304,15 +361,15 @@ class ClaudeSdkBackend:
usage=usage,
)
)
if turn.result is not None and turn.result.is_error:
msg = f"claude: {turn.result.result or turn.result.subtype}"
raise RuntimeError(msg)
yield build_message_delta(
stop_reason=turn.stop_reason, usage=_wire_usage(usage)
)
yield build_message_stop()
return usage, interrupted
async def _run_turn(self, live: _Live, prompt: str) -> _Turn:
async def _run_turn(
self,
live: Session,
prompt: str,
*,
observer: Callable[[Any], None] | None = None,
) -> _Turn:
streaming = self._agent.options.include_partial_messages
turn = _Turn()
raw: list[Any] = []
@@ -320,6 +377,16 @@ class ClaudeSdkBackend:
offset = 0
await live.client.query(prompt)
async for message in live.client.receive_response():
if observer is not None:
observer(message)
if isinstance(message, MirrorErrorMessage):
live.dirty = True
_log.error(
"session %s: mirror error, marked dirty: %s",
live.session_id,
message.error,
)
continue
if getattr(message, "parent_tool_use_id", None) is not None:
continue
if isinstance(message, StreamEvent):
@@ -357,17 +424,30 @@ class ClaudeSdkBackend:
return turn
async def _acquire(
self, key: str, *, session_id: str | None, history: list[dict[str, Any]]
) -> _Live:
live = self._sessions.get(key)
self,
key: str,
*,
session_id: str | None,
history: list[dict[str, Any]],
spec: _SessionSpec,
) -> Session:
live = self._pool.get(key)
if live is not None:
return live
resume = session_id
if resume is None and history:
if resume is not None:
await self.repair_session(resume)
elif history:
resume = await self._seed(history)
live = await self._spawn(resume)
self._sessions[key] = live
return live
await self._pool.make_room()
live = await self._spawn(resume, key=key, spec=spec)
return self._pool.add(live)
def _store_key(self, session_id: str) -> dict[str, str]:
return {
"project_key": project_key_for_directory(str(self._agent.cwd)),
"session_id": session_id,
}
async def _seed(self, history: list[dict[str, Any]]) -> str:
session_id = str(uuid.uuid4())
@@ -378,11 +458,9 @@ class ClaudeSdkBackend:
model=self._agent.model,
permission_mode=self._agent.options.permission_mode,
)
key = {
"project_key": project_key_for_directory(str(self._agent.cwd)),
"session_id": session_id,
}
await self._store.append(cast("Any", key), cast("Any", entries))
await self._store.append(
cast("Any", self._store_key(session_id)), cast("Any", entries)
)
_log.info(
"seeded session %s with %d entries from %d messages",
session_id,
@@ -391,22 +469,36 @@ class ClaudeSdkBackend:
)
return session_id
async def _spawn(self, resume: str | None) -> _Live:
options = self._build_options(resume)
async def _spawn(
self, resume: str | None, *, key: str, spec: _SessionSpec
) -> Session:
options = self._build_options(resume, key=key, spec=spec)
client = self._factory(options)
await client.connect()
_log.info(
"spawned claude: agent=%s resume=%s user=%s",
"spawned claude: agent=%s kind=%s resume=%s tools=%s user=%s",
self._agent.name,
spec.kind,
resume,
spec.tools,
self._runner.user,
)
return _Live(client=client, session_id=resume, resumed=resume is not None)
return Session(
key=key,
agent=self._agent.name,
kind=spec.kind,
client=client,
session_id=resume,
resumed=resume is not None,
pinned=spec.pinned,
)
def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient:
return _RunnerClient(options, uid=self._uid)
def _build_options(self, resume: str | None) -> ClaudeAgentOptions:
def _build_options(
self, resume: str | None, *, key: str, spec: _SessionSpec
) -> ClaudeAgentOptions:
agent = self._agent
opt = agent.options
env = dict(opt.env)
@@ -414,18 +506,25 @@ class ClaudeSdkBackend:
env["HOME"] = str(self._runner.home)
env.setdefault("CLAUDE_CONFIG_DIR", str(self._runner.home / ".claude"))
plugins = self._plugins()
sources = agent.prompt_for(spec.kind)
system_prompt = (
prompt_assembly.assemble(agent.prompt_sources)
if agent.prompt_sources
else agent.system_prompt
prompt_assembly.assemble(sources) if sources else agent.system_prompt
)
servers: dict[str, Any] = dict(self._servers) if spec.tools else {}
gateway = (
self._tool_server(key, spec.kind)
if spec.tools and self._tool_server is not None and agent.gateway_tools
else None
)
if gateway is not None:
servers[str(gateway["name"])] = gateway
return ClaudeAgentOptions(
model=agent.model or None,
effort=cast("Any", opt.effort),
system_prompt=system_prompt,
setting_sources=[],
strict_mcp_config=True,
mcp_servers=cast("Any", self._servers),
mcp_servers=cast("Any", servers),
permission_mode=cast("Any", opt.permission_mode),
tools=list(opt.tools) if opt.tools is not None else None,
disallowed_tools=[*opt.disallowed_tools, *self._mcp_disallowed],
@@ -473,8 +572,8 @@ class ClaudeSdkBackend:
target=json.dumps(target),
keep=json.dumps(keep),
prefixes=json.dumps(list(ENV_KEEP_PREFIXES)),
uid=json.dumps(self._uid),
gid=json.dumps(self._gid),
uid=repr(self._uid),
gid=repr(self._gid),
)
digest = hashlib.sha256(script.encode("utf-8")).hexdigest()[:12]
path = self._work_dir / f"claude-exec-{digest}.py"
@@ -487,30 +586,16 @@ class ClaudeSdkBackend:
return path
def _rekey(self, old: str, new: str) -> None:
live = self._sessions.pop(old, None)
if live is None:
return
stale = self._sessions.pop(new, None)
self._sessions[new] = live
if stale is not None and stale is not live:
stale = self._pool.rekey(old, new)
if stale is not None:
asyncio.get_running_loop().create_task(_disconnect(stale))
async def _close(self, key: str) -> None:
live = self._sessions.pop(key, None)
if live is not None:
await _disconnect(live)
async def _reap_loop(self) -> None:
ttl = self._agent.options.idle_session_ttl
while True:
await asyncio.sleep(_REAP_INTERVAL)
if ttl <= 0:
continue
now = time.monotonic()
for key, live in list(self._sessions.items()):
if not live.lock.locked() and now - live.last_used > ttl:
_log.info("closing idle session %s (%s)", live.session_id, key)
await self._close(key)
@dataclass(frozen=True, slots=True)
class _SessionSpec:
kind: str
pinned: bool
tools: bool
@dataclass
@@ -541,7 +626,7 @@ os.execve(TARGET, [TARGET, *sys.argv[1:]], env)
"""
async def _disconnect(live: _Live) -> None:
async def _disconnect(live: Session) -> None:
try:
await live.client.disconnect()
except Exception: # noqa: BLE001