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

This commit is contained in:
hh
2026-08-28 03:08:30 +02:00
parent 50b7057fa4
commit 33ccc78fec
28 changed files with 3543 additions and 345 deletions
+16 -4
View File
@@ -42,15 +42,27 @@ class ClaudeOptions(BaseModel):
"""Extra inherited variable names to let through the env whitelist."""
max_turns: int | None = None
idle_session_ttl: float = 1800.0
"""Seconds a live session may sit unused before it is closed."""
session_store_flush: str = "batched"
session_store_flush: str = "eager"
"""``eager`` mirrors every transcript frame as it lands, so a gateway
killed mid-turn loses at most the frame in flight."""
class ClaudeAgent(BaseAgent):
cwd: Path
system_prompt: str = ""
prompt_sources: tuple[PromptSource, ...] = ()
prompt_sources_by_kind: Mapping[str, tuple[PromptSource, ...]] = Field(
default_factory=dict
)
"""Per conversation kind (``master``/``branch``/``deep``/``job``/``fork``)
assembly; falls back to ``prompt_sources``. Constant per kind (§3.12)."""
skill_sets: tuple[Path, ...] = ()
gateway_tools: tuple[str, ...] = ()
"""Gateway tools exposed in-process (``read_conversation``, ``spawn``,
``say``, ``schedule``, ``inject``); empty = no gateway MCP server."""
options: ClaudeOptions = Field(default_factory=ClaudeOptions)
def prompt_for(self, kind: str) -> tuple[PromptSource, ...]:
return self.prompt_sources_by_kind.get(kind, self.prompt_sources)
+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
+47 -1
View File
@@ -22,6 +22,7 @@ from __future__ import annotations
import asyncio
import contextlib
import functools
import logging
import signal
from contextlib import AsyncExitStack
@@ -43,13 +44,18 @@ from beaver_gateway.backends.claude_sdk import (
)
from beaver_gateway.backends.raycast import RaycastBackend
from beaver_gateway.core.auth import TokenStore
from beaver_gateway.core.bus import EventBus
from beaver_gateway.core.conversations import Conversations
from beaver_gateway.core.gateway_tools import build_tool_server
from beaver_gateway.core.registry import AgentRegistry, McpRegistry
from beaver_gateway.core.sessions import SessionPool
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.mcp.internal_app import build_internal_app
from beaver_gateway.settings import Settings
from beaver_gateway.storage import Database, PostgresSessionStore, Usage, append_usage
if TYPE_CHECKING:
from claude_agent_sdk import McpSdkServerConfig
from fastmcp import FastMCP
from fastmcp.tools.base import Tool as FastMCPTool
from starlette.applications import Starlette
@@ -140,15 +146,32 @@ async def _async_main() -> None:
# surfaces at startup instead of mid-conversation.
mcp_tools = await _prefetch_mcp_tools(mcp_servers)
pool = SessionPool()
bus = EventBus()
late = _LateConversations()
session_store = PostgresSessionStore(db)
backends: dict[str, Backend] = await _build_backends(
settings=settings,
agents=agents,
stack=stack,
db=db,
session_store=session_store,
mcp_internal_urls=internal_urls,
mcp_servers=mcp_servers,
mcp_tools=mcp_tools,
pool=pool,
late=late,
)
conversations = Conversations(
db=db,
agents=agents,
backends=backends,
bus=bus,
pool=pool,
store=session_store,
texts=gateway.texts,
)
late.conversations = conversations
runtime = GatewayRuntime(
agents=agents,
@@ -161,6 +184,9 @@ async def _async_main() -> None:
admin_pass=settings.admin_pass,
session_secret=settings.session_secret,
frontends=tuple(gateway.frontends),
conversations=conversations,
bus=bus,
pool=pool,
)
for fe in gateway.frontends:
@@ -186,7 +212,10 @@ async def _async_main() -> None:
# it in this path and exit cleanly (Phase 0 DoD).
return
await conversations.start()
stack.push_async_callback(conversations.stop)
async with asyncio.TaskGroup() as tg:
tg.create_task(pool.reap_loop())
if internal_app is not None:
tg.create_task(_serve_internal_mcp(internal_app, settings=settings))
for fe in gateway.frontends:
@@ -252,15 +281,31 @@ async def _serve_internal_mcp(app: Starlette, *, settings: Settings) -> None:
await server.serve()
class _LateConversations:
"""Backends need a tool-server factory before the service that backs it exists."""
conversations: Conversations | None = None
def server(
self, key: str, _kind: str, names: tuple[str, ...]
) -> McpSdkServerConfig | None:
if self.conversations is None or not names:
return None
return build_tool_server(self.conversations, conversation_key=key, names=names)
async def _build_backends(
*,
settings: Settings,
agents: AgentRegistry,
stack: AsyncExitStack,
db: Database,
session_store: PostgresSessionStore,
mcp_internal_urls: dict[str, str],
mcp_servers: dict[str, FastMCP],
mcp_tools: dict[str, list[FastMCPTool]],
pool: SessionPool,
late: _LateConversations,
) -> dict[str, Backend]:
"""Construct one backend per agent name.
@@ -285,7 +330,6 @@ async def _build_backends(
for a in raycast_agents:
backends[a.name] = raycast_backend
session_store = PostgresSessionStore(db)
runner = RunnerConfig(user=settings.claude_runner_user, home=settings.claude_home)
mcp_tool_names = {
name: [t.name for t in tools] for name, tools in mcp_tools.items()
@@ -321,6 +365,8 @@ async def _build_backends(
mcp_tool_names=mcp_tool_names,
runner=runner,
usage_sink=record_usage,
pool=pool,
tool_server=functools.partial(late.server, names=a.gateway_tools),
)
await stack.enter_async_context(adapter)
backends[a.name] = adapter
+2
View File
@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any
from beaver_gateway.agents.base import BaseAgent, ExposedMcp
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.agents.raycast import RaycastAgent
from beaver_gateway.core.conversations import ConversationTexts
from beaver_gateway.core.registry import Gateway
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.mcp.types import HttpMcp, McpServer, PythonToolMcp, StdioMcp
@@ -37,6 +38,7 @@ _PUBLIC_NAMES: dict[str, Any] = {
"McpServer": McpServer,
"ExposedMcp": ExposedMcp,
"Gateway": Gateway,
"ConversationTexts": ConversationTexts,
}
_McpInstance = StdioMcp | HttpMcp | PythonToolMcp
+2 -1
View File
@@ -60,12 +60,13 @@ class TokenStoreError(ValueError):
"""Malformed ``BOOTSTRAP_TOKENS`` value or duplicate token."""
VALID_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin"})
VALID_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin", "api"})
"""The scopes a ``Token.scope`` may hold (Phase 4.3 admin UI enforces).
* ``*`` — wildcard, may use any frontend
* ``messages`` — Anthropic Messages frontend only
* ``mcp`` — MCP server frontend only
* ``api`` — conversations API (``/api``) and its SSE
* ``admin`` — reserved for programmatic admin access; the AdminFrontend
itself authenticates via session cookies, not bearer tokens, so this
scope is unused today and kept for forward compatibility.
+69
View File
@@ -0,0 +1,69 @@
"""In-process event bus: what the gateway does, as a stream frontends can tap.
Events are plain dicts ``{"type", "seq", "ts", ...}``. Publishers never
block; a subscriber that falls behind loses its oldest events rather than
stalling the turn that produced them. ``/api/events`` serialises the
stream as SSE, the panel builds its subagent tree from ``stream`` events
carrying ``parent_tool_use_id``.
"""
from __future__ import annotations
import asyncio
import contextlib
import itertools
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import AsyncIterator
__all__ = ["Event", "EventBus"]
Event = dict[str, Any]
class EventBus:
def __init__(self, *, maxsize: int = 2000) -> None:
self._maxsize = maxsize
self._subscribers: set[asyncio.Queue[Event]] = set()
self._seq = itertools.count(1)
def publish(self, type_: str, **data: Any) -> Event:
event: Event = {
"type": type_,
"seq": next(self._seq),
"ts": datetime.now(UTC).isoformat(timespec="milliseconds"),
**data,
}
for queue in list(self._subscribers):
if queue.full():
with contextlib.suppress(asyncio.QueueEmpty):
queue.get_nowait()
queue.put_nowait(event)
return event
@contextlib.asynccontextmanager
async def subscribe(self) -> AsyncIterator[asyncio.Queue[Event]]:
queue: asyncio.Queue[Event] = asyncio.Queue(maxsize=self._maxsize)
self._subscribers.add(queue)
try:
yield queue
finally:
self._subscribers.discard(queue)
async def stream(
self, *, conversation_id: str | None = None
) -> AsyncIterator[Event]:
async with self.subscribe() as queue:
while True:
event = await queue.get()
if (
conversation_id is None
or event.get("conversation_id") == conversation_id
):
yield event
@property
def subscribers(self) -> int:
return len(self._subscribers)
@@ -496,12 +496,14 @@ def _walk_prefix(
s.name for s in inc.structure if isinstance(s, ToolSegment)
)
inc_text_count = sum(1 for s in inc.structure if isinstance(s, TextSegment))
if inc_skeleton != st.skeleton:
# Files rendered without tool callouts (§3.10) carry no skeleton:
# prose alone decides whether the turn matched.
if inc_skeleton and inc_skeleton != st.skeleton:
return spliced_groups, i
if inc.text == st.spoken_text:
spliced_groups.append(list(st.messages))
continue
if inc_text_count != st.text_segment_count:
if inc_skeleton and inc_text_count != st.text_segment_count:
return spliced_groups, i
spliced = _splice_assistant_group(stored_group=st, incoming=inc)
if spliced is None:
File diff suppressed because it is too large Load Diff
+169
View File
@@ -0,0 +1,169 @@
"""In-process MCP server with the gateway's own tools (§3.1, §3.2).
One server per live session so every tool knows which conversation is
calling; ``alwaysLoad`` keeps the tools out of tool search. Which names a
session gets comes from ``ClaudeAgent.gateway_tools``.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, cast
from claude_agent_sdk import create_sdk_mcp_server, tool
if TYPE_CHECKING:
from collections.abc import Iterable
from claude_agent_sdk import McpSdkServerConfig, SdkMcpTool
from beaver_gateway.core.conversations import Conversations
__all__ = ["SERVER_NAME", "TOOL_NAMES", "build_tool_server"]
_log = logging.getLogger("beaver_gateway.core.gateway_tools")
SERVER_NAME = "gateway"
TOOL_NAMES = ("read_conversation", "spawn", "say", "schedule", "inject")
def build_tool_server(
conversations: Conversations, *, conversation_key: str, names: Iterable[str]
) -> McpSdkServerConfig | None:
wanted = set(names)
unknown = wanted - set(TOOL_NAMES)
if unknown:
msg = f"unknown gateway tools: {sorted(unknown)}"
raise ValueError(msg)
tools = [t for t in _tools(conversations, conversation_key) if t.name in wanted]
if not tools:
return None
server = create_sdk_mcp_server(SERVER_NAME, tools=tools)
return cast("McpSdkServerConfig", {**server, "alwaysLoad": True})
def _tools(conversations: Conversations, key: str) -> list[SdkMcpTool[Any]]:
async def current() -> Any:
conv = await conversations.get(key)
if conv is None:
msg = f"conversation {key} not found"
raise LookupError(msg)
return conv
@tool(
"read_conversation",
"Read another conversation (a branch, the master, a deep chat) as plain "
"text. `window` limits it to the last N user turns.",
{
"type": "object",
"properties": {
"id": {"type": "string", "description": "conversation id"},
"window": {"type": "integer", "minimum": 1},
},
"required": ["id"],
},
)
async def read_conversation(args: dict[str, Any]) -> dict[str, Any]:
conv = await conversations.get(str(args["id"]))
if conv is None:
return _error(f"conversation {args['id']} not found")
text = await conversations.read(conv, window=args.get("window"))
return _text(text or "(empty)")
@tool(
"spawn",
"Open a new conversation of the given kind (branch = your own thread, "
"deep = a long research chat, job = a headless task). `seed` is how it "
"starts: clean (nothing), morning (handout), copy (copy of this "
"conversation, last `window` turns), brief (your `text`). Returns the id.",
{
"type": "object",
"properties": {
"kind": {"type": "string", "enum": ["branch", "deep", "job"]},
"seed": {
"type": "string",
"enum": ["clean", "morning", "copy", "brief"],
"default": "clean",
},
"text": {"type": "string", "description": "brief for seed=brief"},
"title": {"type": "string"},
"window": {"type": "integer", "minimum": 1},
},
"required": ["kind"],
},
)
async def spawn(args: dict[str, Any]) -> dict[str, Any]:
parent = await current()
child = await conversations.spawn(
kind=str(args["kind"]),
agent=parent.agent_name,
seed=str(args.get("seed") or "clean"),
parent=parent,
text=args.get("text"),
title=args.get("title"),
window=args.get("window"),
origin="mcp",
)
return _text(f"spawned {child.kind} {child.external_id}")
@tool(
"say",
"Say something to the human in the frontend this conversation is bound "
"to. The only way an inject-started turn can speak; silence is simply "
"not calling it.",
{"text": str},
)
async def say(args: dict[str, Any]) -> dict[str, Any]:
conv = await current()
await conversations.say(conv, str(args["text"]))
return _text("ok")
@tool(
"schedule",
"Promise yourself an inject later: `at` is `+15m`, `+2h`, `+1d` or an "
"ISO datetime; `text` arrives in this conversation at that time.",
{"at": str, "text": str},
)
async def schedule(args: dict[str, Any]) -> dict[str, Any]:
conv = await current()
row = await conversations.schedule(conv, str(args["at"]), str(args["text"]))
return _text(f"scheduled #{row.id} at {row.execute_at.isoformat()}")
@tool(
"inject",
"Put a system-origin message into another conversation's queue.",
{
"type": "object",
"properties": {
"conversation": {"type": "string"},
"text": {"type": "string"},
"urgency": {
"type": "string",
"enum": ["normal", "urgent"],
"default": "normal",
},
},
"required": ["conversation", "text"],
},
)
async def inject(args: dict[str, Any]) -> dict[str, Any]:
target = await conversations.get(str(args["conversation"]))
if target is None:
return _error(f"conversation {args['conversation']} not found")
item = await conversations.inject(
target,
str(args["text"]),
urgency=cast("Any", args.get("urgency") or "normal"),
origin="агент",
)
return _text(f"queued #{item.id}")
return [read_conversation, spawn, say, schedule, inject]
def _text(text: str) -> dict[str, Any]:
return {"content": [{"type": "text", "text": text}]}
def _error(text: str) -> dict[str, Any]:
return {"content": [{"type": "text", "text": text}], "is_error": True}
+130
View File
@@ -0,0 +1,130 @@
"""Persisted per-conversation queue with priorities ``urgent > user > normal`` (§3.4).
One ``ClaudeSDKClient`` runs one turn at a time, so ordering has to happen
before the client: the rows here are the queue, ``core/conversations``
runs one worker per conversation over them. A row that is still
``running`` when the gateway starts was cut by a restart; it is flagged
``interrupted`` and never re-run.
"""
from __future__ import annotations
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Literal
from sqlmodel import col, select
from beaver_gateway.storage.models import InjectQueueItem
if TYPE_CHECKING:
from collections.abc import Iterable, Sequence
from beaver_gateway.storage.db import Database
__all__ = ["PRIORITY_RANK", "InjectQueue", "Priority", "inject_header"]
Priority = Literal["urgent", "user", "normal"]
PRIORITY_RANK: dict[str, int] = {"urgent": 0, "user": 1, "normal": 2}
def inject_header(origin: str) -> str:
return f"[инжект: {origin} - это не Бобёр, отвечать не нужно, голос не обязателен]"
class InjectQueue:
def __init__(self, db: Database) -> None:
self._db = db
async def push(
self, *, conversation_id: int, priority: Priority, origin: str, text: str
) -> InjectQueueItem:
if priority not in PRIORITY_RANK:
msg = f"unknown priority {priority!r}"
raise ValueError(msg)
row = InjectQueueItem(
conversation_id=conversation_id, priority=priority, origin=origin, text=text
)
async with self._db.session() as session:
session.add(row)
await session.commit()
await session.refresh(row)
return row
async def pending(self, conversation_id: int) -> list[InjectQueueItem]:
async with self._db.session() as session:
result = await session.exec(
select(InjectQueueItem)
.where(
InjectQueueItem.conversation_id == conversation_id,
InjectQueueItem.status == "queued",
)
.order_by(col(InjectQueueItem.created_at), col(InjectQueueItem.id))
)
rows = list(result.all())
rows.sort(key=lambda r: (PRIORITY_RANK.get(r.priority, 9), r.created_at))
return rows
async def conversations_with_pending(self) -> list[int]:
async with self._db.session() as session:
result = await session.exec(
select(InjectQueueItem.conversation_id)
.where(InjectQueueItem.status == "queued")
.distinct()
)
return list(result.all())
async def start(self, items: Iterable[InjectQueueItem], turn_id: str) -> None:
await self._mark(items, status="running", turn_id=turn_id, delivered=True)
async def finish(
self, items: Iterable[InjectQueueItem], status: str = "done"
) -> None:
await self._mark(items, status=status)
async def interrupted(self) -> Sequence[InjectQueueItem]:
async with self._db.session() as session:
result = await session.exec(
select(InjectQueueItem).where(InjectQueueItem.status == "running")
)
rows = list(result.all())
for row in rows:
row.status = "interrupted"
session.add(row)
await session.commit()
return rows
async def recent(
self, conversation_id: int, *, limit: int = 50
) -> list[InjectQueueItem]:
async with self._db.session() as session:
result = await session.exec(
select(InjectQueueItem)
.where(InjectQueueItem.conversation_id == conversation_id)
.order_by(col(InjectQueueItem.id).desc())
.limit(limit)
)
return list(result.all())
async def _mark(
self,
items: Iterable[InjectQueueItem],
*,
status: str,
turn_id: str | None = None,
delivered: bool = False,
) -> None:
ids = [i.id for i in items if i.id is not None]
if not ids:
return
async with self._db.session() as session:
result = await session.exec(
select(InjectQueueItem).where(col(InjectQueueItem.id).in_(ids))
)
for row in result.all():
row.status = status
if turn_id is not None:
row.turn_id = turn_id
if delivered:
row.delivered_at = datetime.now(UTC)
session.add(row)
await session.commit()
+3
View File
@@ -16,6 +16,7 @@ if TYPE_CHECKING:
from collections.abc import Iterable, Iterator
from beaver_gateway.agents.base import BaseAgent
from beaver_gateway.core.conversations import ConversationTexts
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.mcp.types import McpServerT
@@ -81,3 +82,5 @@ class Gateway:
agents: list[BaseAgent] = field(default_factory=list)
mcps: list[McpServerT] = field(default_factory=list)
frontends: list[Frontend] = field(default_factory=list)
texts: ConversationTexts | None = None
"""Merge prompt and seed bodies for ``core/conversations`` (§8.2-8.3)."""
+257
View File
@@ -0,0 +1,257 @@
"""Pool of live Agent SDK sessions across every Claude agent (§3.2).
One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess).
The pool owns the two decisions the adapters used to make on their own:
when a session is closed for idleness (TTL by conversation kind) and
which one goes when memory runs out (measured RSS of the subprocess tree
against the cgroup limit, ``max_live`` where there is no limit). Eviction
only ever picks ``idle && !running_turn && !pending_question`` sessions
that are neither pinned (the master) nor ``dirty`` (mirror gap not yet
repaired); forks and jobs go first.
"""
from __future__ import annotations
import asyncio
import contextlib
import logging
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol
import psutil
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterator, Mapping
__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]
_log = logging.getLogger("beaver_gateway.core.sessions")
DEFAULT_TTL: Mapping[str, float | None] = {
"master": None,
"branch": 7200.0,
"deep": 1800.0,
"job": 0.0,
"fork": 0.0,
}
"""Idle seconds before a session is closed; ``None`` = never (pinned kinds)."""
_EVICT_ORDER = {"fork": 0, "job": 0, "deep": 1, "branch": 2, "master": 3}
_RSS_HEADROOM = 0.8
class SessionClient(Protocol):
async def connect(self) -> None: ...
async def query(self, prompt: str) -> None: ...
def receive_response(self) -> AsyncIterator[Any]: ...
async def interrupt(self) -> None: ...
async def disconnect(self) -> None: ...
@dataclass
class Session:
key: str
agent: str
kind: str
client: SessionClient
session_id: str | None
resumed: bool
pinned: bool = False
dirty: bool = False
running_turn: str | None = None
pending_question: bool = False
interrupt_requested: bool = False
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
last_used: float = field(default_factory=time.monotonic)
created_at: float = field(default_factory=time.monotonic)
turns: int = 0
@property
def busy(self) -> bool:
return self.lock.locked() or self.running_turn is not None
@property
def evictable(self) -> bool:
return not (self.pinned or self.dirty or self.busy or self.pending_question)
@property
def pid(self) -> int | None:
transport = getattr(self.client, "_transport", None)
process = getattr(transport, "_process", None)
pid = getattr(process, "pid", None)
return pid if isinstance(pid, int) else None
class SessionPool:
def __init__(
self,
*,
rss_limit: int | None = None,
max_live: int = 8,
ttl: Mapping[str, float | None] = DEFAULT_TTL,
reap_interval: float = 60.0,
) -> None:
self._sessions: dict[str, Session] = {}
self._rss_limit = rss_limit if rss_limit is not None else cgroup_limit()
self._max_live = max_live
self._ttl = dict(ttl)
self._reap_interval = reap_interval
def __len__(self) -> int:
return len(self._sessions)
def __iter__(self) -> Iterator[Session]:
return iter(list(self._sessions.values()))
def __contains__(self, key: object) -> bool:
return key in self._sessions
@property
def rss_limit(self) -> int | None:
return self._rss_limit
def get(self, key: str) -> Session | None:
return self._sessions.get(key)
def add(self, session: Session) -> Session:
self._sessions[session.key] = session
return session
def pop(self, key: str) -> Session | None:
return self._sessions.pop(key, None)
def rekey(self, old: str, new: str) -> Session | None:
session = self._sessions.pop(old, None)
if session is None:
return None
stale = self._sessions.pop(new, None)
session.key = new
self._sessions[new] = session
return stale if stale is not session else None
def ttl_for(self, kind: str) -> float | None:
return self._ttl.get(kind, self._ttl.get("deep"))
def rss(self) -> int:
try:
children = psutil.Process(os.getpid()).children(recursive=True)
except psutil.Error:
return 0
total = 0
for child in children:
with contextlib.suppress(psutil.Error):
total += child.memory_info().rss
return total
@staticmethod
def rss_of(session: Session) -> int | None:
pid = session.pid
if pid is None:
return None
try:
process = psutil.Process(pid)
return process.memory_info().rss + sum(
c.memory_info().rss for c in process.children(recursive=True)
)
except psutil.Error:
return None
def over_limit(self) -> bool:
if self._rss_limit is not None:
return self.rss() > self._rss_limit * _RSS_HEADROOM
return len(self._sessions) >= self._max_live
def victims(self) -> list[Session]:
candidates = [s for s in self._sessions.values() if s.evictable]
candidates.sort(key=lambda s: (_EVICT_ORDER.get(s.kind, 1), s.last_used))
return candidates
async def make_room(self) -> int:
closed = 0
while self.over_limit():
victims = self.victims()
if not victims:
_log.warning(
"session pool over limit (%d live, rss=%d) but nothing evictable",
len(self._sessions),
self.rss(),
)
break
await self.close(victims[0].key)
closed += 1
return closed
async def close(self, key: str) -> None:
session = self._sessions.pop(key, None)
if session is None:
return
_log.info("closing session %s (%s, %s)", session.session_id, session.kind, key)
try:
await session.client.disconnect()
except Exception: # noqa: BLE001
_log.exception("disconnect failed for session %s", session.session_id)
async def close_all(self, *, agent: str | None = None) -> None:
for session in list(self._sessions.values()):
if agent is None or session.agent == agent:
await self.close(session.key)
async def reap_once(self) -> int:
now = time.monotonic()
closed = 0
for session in list(self._sessions.values()):
ttl = self.ttl_for(session.kind)
if ttl is None or not session.evictable:
continue
if now - session.last_used > ttl:
await self.close(session.key)
closed += 1
return closed
async def reap_loop(self) -> None:
while True:
await asyncio.sleep(self._reap_interval)
try:
await self.reap_once()
await self.make_room()
except Exception: # noqa: BLE001
_log.exception("session reaper failed")
def snapshot(self) -> list[dict[str, Any]]:
now = time.monotonic()
return [
{
"key": s.key,
"agent": s.agent,
"kind": s.kind,
"session_id": s.session_id,
"pid": s.pid,
"rss": self.rss_of(s),
"idle_seconds": round(now - s.last_used, 1),
"age_seconds": round(now - s.created_at, 1),
"turns": s.turns,
"busy": s.busy,
"running_turn": s.running_turn,
"pending_question": s.pending_question,
"pinned": s.pinned,
"dirty": s.dirty,
}
for s in self._sessions.values()
]
def cgroup_limit() -> int | None:
for path in (
"/sys/fs/cgroup/memory.max",
"/sys/fs/cgroup/memory/memory.limit_in_bytes",
):
try:
raw = Path(path).read_text(encoding="ascii").strip()
except OSError:
continue
if raw.isdigit() and int(raw) < 1 << 60:
return int(raw)
return None
+229 -1
View File
@@ -27,7 +27,17 @@ except ImportError: # pragma: no cover
if TYPE_CHECKING:
from collections.abc import Iterable, Mapping
__all__ = ["CLI_VERSION", "build_entries", "messages_from_entries"]
__all__ = [
"CLI_VERSION",
"build_entries",
"close_open_tool_uses",
"messages_from_entries",
"open_tool_uses",
"prompt_count",
"render_messages",
"strip_tool_entries",
"window_entries",
]
CLI_VERSION = _cli_version
_ENTRYPOINT = "sdk-py"
@@ -275,3 +285,221 @@ def _zero_usage() -> dict[str, Any]:
"iterations": [],
"speed": "standard",
}
# ---- repair, windows, projections ---------------------------------------
_PROMPT_TYPES = ("user", "assistant")
_INTERRUPTED = "прервано"
def open_tool_uses(
entries: Iterable[Mapping[str, Any]],
) -> list[tuple[Mapping[str, Any], dict[str, Any]]]:
"""``(assistant entry, tool_use block)`` pairs that never got a ``tool_result``."""
closed: set[str] = set()
uses: list[tuple[Mapping[str, Any], dict[str, Any]]] = []
for entry in entries:
content = _entry_content(entry)
if entry.get("type") == "user":
closed.update(
str(b.get("tool_use_id", ""))
for b in content
if b.get("type") == "tool_result"
)
elif entry.get("type") == "assistant":
uses.extend((entry, b) for b in content if b.get("type") == "tool_use")
return [
(owner, block)
for owner, block in uses
if str(block.get("id", "")) not in closed
]
def close_open_tool_uses(
entries: list[Mapping[str, Any]], *, text: str = _INTERRUPTED
) -> list[dict[str, Any]]:
"""Synthetic ``tool_result`` entries on the leaf, one per open ``tool_use``.
Appending the result to the session store gives the next ``resume`` a
transcript the CLI accepts: an assistant message ending in ``tool_use``
without its result is rejected by the API on the next call.
"""
pending = open_tool_uses(entries)
if not pending:
return []
leaf = next(
(
e
for e in reversed(entries)
if e.get("type") in _PROMPT_TYPES and e.get("uuid")
),
None,
)
parent = str(leaf["uuid"]) if leaf is not None else None
stamp = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S.000Z")
out: list[dict[str, Any]] = []
for owner, block in pending:
uid = _new_uuid()
result = {
"type": "tool_result",
"tool_use_id": block.get("id", ""),
"content": text,
"is_error": True,
}
out.append(
{
"parentUuid": parent,
"promptId": owner.get("promptId"),
"type": "user",
"message": {"role": "user", "content": [result]},
"uuid": uid,
"timestamp": stamp,
"toolUseResult": text,
"sourceToolAssistantUUID": owner.get("uuid"),
**{k: owner[k] for k in _COMMON_KEYS if k in owner},
}
)
parent = uid
return out
_COMMON_KEYS = (
"isSidechain",
"userType",
"entrypoint",
"cwd",
"sessionId",
"version",
"gitBranch",
)
def window_entries(
entries: Iterable[Mapping[str, Any]], *, window: int | None = None
) -> list[dict[str, Any]]:
"""``user``/``assistant`` entries only, last ``window`` prompts, relinked."""
kept = [
dict(e)
for e in entries
if e.get("type") in _PROMPT_TYPES and isinstance(e.get("uuid"), str)
]
if window is not None and window > 0:
starts = [i for i, e in enumerate(kept) if _is_prompt_entry(e)]
if len(starts) > window:
kept = kept[starts[-window] :]
return _relink(kept)
def strip_tool_entries(entries: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
"""Drop ``tool_result`` entries and ``tool_use``/``thinking`` blocks; relink."""
out: list[dict[str, Any]] = []
for raw in entries:
entry = dict(raw)
message = entry.get("message")
if not isinstance(message, dict):
out.append(entry)
continue
content = message.get("content")
if isinstance(content, list):
blocks = [
b
for b in content
if isinstance(b, dict)
and b.get("type") not in ("tool_use", "tool_result", "thinking")
]
if not blocks:
continue
entry["message"] = {**message, "content": blocks}
out.append(entry)
return _relink(out)
def prompt_count(entries: Iterable[Mapping[str, Any]]) -> int:
return sum(1 for e in entries if _is_prompt_entry(e))
def render_messages(
messages: Iterable[Mapping[str, Any]], *, window: int | None = None
) -> str:
"""Plain-text projection for ``read_conversation``: ``user:``/``assistant:`` turns.
Tool calls collapse to a one-line summary per assistant turn; tool
results and thinking are dropped.
"""
turns: list[str] = []
tools: list[str] = []
current: list[str] = []
def flush() -> None:
if not current and not tools:
return
body = "\n\n".join(current).strip()
if tools:
body = (body + "\n" if body else "") + "(tools: " + ", ".join(tools) + ")"
turns.append("assistant:\n" + body)
current.clear()
tools.clear()
for message in messages:
role = message.get("role")
content = message.get("content")
if role == "user":
if _tool_results(content):
continue
flush()
turns.append("user:\n" + _text_of_content(content))
continue
for block in _assistant_blocks(content):
if block.get("type") == "text" and block.get("text"):
current.append(str(block["text"]))
elif block.get("type") == "tool_use":
tools.append(str(block.get("name", "")))
flush()
if window is not None and window > 0:
starts = [i for i, t in enumerate(turns) if t.startswith("user:")]
if len(starts) > window:
turns = turns[starts[-window] :]
return "\n\n".join(turns)
def _entry_content(entry: Mapping[str, Any]) -> list[dict[str, Any]]:
message = entry.get("message")
if not isinstance(message, dict):
return []
content = message.get("content")
if not isinstance(content, list):
return []
return [b for b in content if isinstance(b, dict)]
def _is_prompt_entry(entry: Mapping[str, Any]) -> bool:
if entry.get("type") != "user":
return False
message = entry.get("message")
if not isinstance(message, dict):
return False
content = message.get("content")
if isinstance(content, str):
return True
return not _tool_results(content)
def _relink(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
parent: str | None = None
for entry in entries:
entry["parentUuid"] = parent
parent = entry.get("uuid")
return entries
def _text_of_content(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
return "\n\n".join(
str(b.get("text", ""))
for b in content
if isinstance(b, dict) and b.get("type") == "text" and b.get("text")
)
return ""
+2
View File
@@ -35,3 +35,5 @@ class TurnCapture:
"""Backend session that ran the turn; persist and pass back as ``session_id``."""
usage: TurnUsage | None = None
interrupted: bool = False
"""The turn was cut by ``interrupt()`` (urgent inject), not by an error."""
+63
View File
@@ -0,0 +1,63 @@
"""Server-sent events helpers shared by the markdown and API frontends."""
from __future__ import annotations
import asyncio
import contextlib
import json
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import AsyncIterator
__all__ = ["HEARTBEAT_INTERVAL", "SSE_HEADERS", "events_with_heartbeat", "sse_pack"]
HEARTBEAT_INTERVAL = 15.0
"""Seconds of backend silence before a comment frame keeps the socket warm."""
SSE_HEADERS = {
"Cache-Control": "no-cache, no-transform",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
}
KEEPALIVE = b": keepalive\n\n"
async def events_with_heartbeat(
events: AsyncIterator[Any], interval: float = HEARTBEAT_INTERVAL
) -> AsyncIterator[Any]:
"""Pass ``events`` through, yielding ``None`` after ``interval`` seconds of silence.
One in-flight ``__anext__`` task is reused across timeouts: a second
consumer on the same async generator raises ``RuntimeError``.
Cancellation of the outer scope cancels that task instead of leaving
it dangling.
"""
src = events.__aiter__()
next_task: asyncio.Task[Any] | None = None
try:
while True:
if next_task is None:
next_task = asyncio.ensure_future(src.__anext__())
done, _pending = await asyncio.wait({next_task}, timeout=interval)
if not done:
yield None
continue
task = next_task
next_task = None
try:
result = task.result()
except StopAsyncIteration:
return
yield result
finally:
if next_task is not None and not next_task.done():
next_task.cancel()
with contextlib.suppress(BaseException):
await next_task
def sse_pack(event: str, data: dict[str, Any]) -> bytes:
body = json.dumps(data, ensure_ascii=False)
return f"event: {event}\ndata: {body}\n\n".encode()
@@ -0,0 +1,5 @@
"""``ApiFrontend`` - the conversations API and event stream (§3.9)."""
from beaver_gateway.frontends.api.frontend import ApiFrontend
__all__ = ["ApiFrontend"]
@@ -0,0 +1,491 @@
"""``ApiFrontend`` - ``/api/conversations``, SSE events, sessions, usage (§3.9).
Bearer scope ``api``. Every write goes through ``core/conversations``; the
frontend only shapes JSON. ``/api/events`` and
``/api/conversations/{id}/events`` replay the gateway bus as SSE with the
same keepalive the markdown frontend uses, so a proxy never sees a
silent socket.
"""
from __future__ import annotations
import json
import logging
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any, cast
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from sqlalchemy import func
from sqlalchemy import select as sa_select
from sqlmodel import col
from beaver_gateway.core import audit
from beaver_gateway.core.conversations import KINDS, SEEDS
from beaver_gateway.frontends._auth import require_token
from beaver_gateway.frontends._sse import (
KEEPALIVE,
SSE_HEADERS,
events_with_heartbeat,
sse_pack,
)
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.storage.models import Usage
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from beaver_gateway.core.conversations import Conversations
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.storage.models import Conversation
_log = logging.getLogger("beaver_gateway.frontends.api")
__all__ = ["ApiFrontend"]
SCOPE = "api"
class ApiFrontend(Frontend):
def __init__(
self,
*,
host: str = "0.0.0.0", # noqa: S104
port: int = 8004,
public_base_url: str | None = None,
) -> None:
self.host = host
self.port = port
self.public_base_url = public_base_url.rstrip("/") if public_base_url else None
self._app: FastAPI | None = None
def configure(self, runtime: GatewayRuntime) -> None:
if runtime.conversations is None or runtime.bus is None:
msg = "ApiFrontend needs runtime.conversations and runtime.bus"
raise RuntimeError(msg)
self._app = _build_app(runtime)
async def serve(self) -> None:
import uvicorn
if self._app is None:
msg = "configure() must be called before serve()"
raise RuntimeError(msg)
server = uvicorn.Server(
uvicorn.Config(self._app, host=self.host, port=self.port, log_level="info")
)
await server.serve()
def _build_app(runtime: GatewayRuntime) -> FastAPI: # noqa: PLR0915
app = FastAPI(title="beaver-gateway / API")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=False,
allow_methods=["*"],
allow_headers=["*"],
)
conversations = cast("Conversations", runtime.conversations)
async def body_of(request: Request) -> dict[str, Any]:
if not await request.body():
return {}
try:
data = await request.json()
except json.JSONDecodeError as exc:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"invalid JSON: {exc}"
) from exc
if not isinstance(data, dict):
raise HTTPException(status.HTTP_400_BAD_REQUEST, "body must be an object")
return data
async def conv_of(public_id: str) -> Conversation:
conv = await conversations.get(public_id)
if conv is None:
raise HTTPException(
status.HTTP_404_NOT_FOUND, f"unknown conversation {public_id}"
)
return conv
def text_of(data: dict[str, Any], key: str = "text") -> str:
text = data.get(key)
if not isinstance(text, str) or not text.strip():
raise HTTPException(status.HTTP_400_BAD_REQUEST, f"missing `{key}`")
return text
def int_or_none(data: dict[str, Any], key: str) -> int | None:
value = data.get(key)
if value is None:
return None
if not isinstance(value, int) or value < 1:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"`{key}` must be a positive int"
)
return value
@app.get("/healthz")
async def healthz() -> dict[str, str]:
return {"status": "ok"}
@app.get("/api/conversations")
async def list_conversations(request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
q = request.query_params
rows = await conversations.find(
status=q.get("status"), kind=q.get("kind"), limit=int(q.get("limit", "200"))
)
return {"conversations": [conversations.public(r) for r in rows]}
@app.post("/api/conversations", status_code=status.HTTP_201_CREATED)
async def create_conversation(request: Request) -> dict[str, Any]:
token = await require_token(request, runtime, scope=SCOPE)
data = await body_of(request)
kind = str(data.get("kind") or "deep")
agent = data.get("agent")
if kind not in KINDS or kind == "fork":
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"kind must be one of {KINDS[:-1]}"
)
if not isinstance(agent, str) or agent not in runtime.agents:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "unknown or missing `agent`"
)
seed = str(data.get("seed") or "clean")
if seed not in SEEDS:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"seed must be one of {SEEDS}"
)
parent = await conv_of(str(data["parent"])) if data.get("parent") else None
try:
conv = await conversations.spawn(
kind=kind,
agent=agent,
seed=seed,
parent=parent,
text=data.get("text"),
title=data.get("title"),
window=int_or_none(data, "window"),
origin="api",
)
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
await audit.log(
runtime,
actor=f"token:{token}",
kind="api_spawn",
agent_name=agent,
conversation=conv.external_id,
seed=seed,
)
return await conversations.describe(conv)
@app.get("/api/conversations/{public_id}")
async def get_conversation(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
out = await conversations.describe(conv)
out["queue"] = [
{
"id": i.id,
"priority": i.priority,
"origin": i.origin,
"status": i.status,
"created_at": i.created_at.isoformat(),
"text": i.text[:200],
}
for i in await conversations.queue.recent(cast("int", conv.id), limit=20)
]
return out
@app.get("/api/conversations/{public_id}/messages")
async def get_messages(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
raw = request.query_params.get("window")
window = int(raw) if raw and raw.isdigit() else None
return {
"id": conv.external_id,
"text": await conversations.read(conv, window=window),
}
@app.post(
"/api/conversations/{public_id}/messages", status_code=status.HTTP_202_ACCEPTED
)
async def post_message(public_id: str, request: Request) -> dict[str, Any]:
token = await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
data = await body_of(request)
item = await conversations.post(
conv, text_of(data), origin=str(data.get("origin") or "user")
)
await audit.log(
runtime,
actor=f"token:{token}",
kind="api_message",
agent_name=conv.agent_name,
conversation=conv.external_id,
)
return {"id": conv.external_id, "item": item.id, "status": item.status}
@app.post(
"/api/conversations/{public_id}/inject", status_code=status.HTTP_202_ACCEPTED
)
async def post_inject(public_id: str, request: Request) -> dict[str, Any]:
token = await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
data = await body_of(request)
urgency = str(data.get("urgency") or "normal")
if urgency not in ("normal", "urgent"):
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "urgency must be normal|urgent"
)
item = await conversations.inject(
conv,
text_of(data),
urgency=cast("Any", urgency),
origin=str(data.get("origin") or "api"),
)
await audit.log(
runtime,
actor=f"token:{token}",
kind="api_inject",
agent_name=conv.agent_name,
conversation=conv.external_id,
urgency=urgency,
)
return {"id": conv.external_id, "item": item.id, "priority": item.priority}
@app.post("/api/conversations/{public_id}/say")
async def post_say(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
return await conversations.say(conv, text_of(await body_of(request)))
@app.post(
"/api/conversations/{public_id}/branch", status_code=status.HTTP_201_CREATED
)
async def post_branch(public_id: str, request: Request) -> dict[str, Any]:
token = await require_token(request, runtime, scope=SCOPE)
parent = await conv_of(public_id)
data = await body_of(request)
seed = str(data.get("seed_mode") or data.get("seed") or "morning")
if seed not in SEEDS:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"seed must be one of {SEEDS}"
)
agent = str(data.get("agent") or parent.agent_name)
try:
child = await conversations.spawn(
kind="branch",
agent=agent,
seed=seed,
parent=parent,
text=data.get("text"),
title=data.get("title"),
window=int_or_none(data, "window"),
origin="api",
)
except (ValueError, LookupError) as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
await audit.log(
runtime,
actor=f"token:{token}",
kind="api_branch",
agent_name=agent,
conversation=child.external_id,
parent=parent.external_id,
seed=seed,
)
return await conversations.describe(child)
@app.post("/api/conversations/{public_id}/merge")
async def post_merge(public_id: str, request: Request) -> dict[str, Any]:
token = await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
try:
result = await conversations.merge(conv)
except (ValueError, LookupError) as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
await audit.log(
runtime,
actor=f"token:{token}",
kind="api_merge",
agent_name=conv.agent_name,
conversation=conv.external_id,
)
return {
"id": conv.external_id,
"status": "merged",
"fork": result.conversation.external_id,
"text": result.text,
}
@app.post("/api/conversations/{public_id}/fork")
async def post_fork(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
data = await body_of(request)
try:
result = await conversations.fork(
conv,
text_of(data, "prompt"),
window=int_or_none(data, "window"),
strip_tools=bool(data.get("strip_tools", False)),
)
except (ValueError, RuntimeError) as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
return {
"id": conv.external_id,
"fork": result.conversation.external_id,
"text": result.text,
}
@app.post("/api/conversations/{public_id}/bind")
async def post_bind(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
data = await body_of(request)
frontend = text_of(data, "frontend")
external_id = text_of(data, "external_id")
await conversations.bind(
conv,
frontend=frontend,
external_id=external_id,
visible=bool(data.get("visible", True)),
)
return await conversations.describe(conv)
@app.patch("/api/conversations/{public_id}/flags")
async def patch_flags(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
data = await body_of(request)
return conversations.public(await conversations.set_flags(conv, data))
@app.patch("/api/conversations/{public_id}")
async def patch_conversation(public_id: str, request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
data = await body_of(request)
try:
if isinstance(data.get("status"), str):
conv = await conversations.set_status(conv, data["status"])
if isinstance(data.get("title"), str):
conv = await conversations.set_title(conv, data["title"])
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
return conversations.public(conv)
@app.get("/api/conversations/{public_id}/events")
async def conversation_events(public_id: str, request: Request) -> Any:
await require_token(request, runtime, scope=SCOPE)
conv = await conv_of(public_id)
return _sse(runtime, conversation_id=conv.external_id)
@app.get("/api/events")
async def all_events(request: Request) -> Any:
await require_token(request, runtime, scope=SCOPE)
return _sse(runtime, conversation_id=None)
@app.get("/api/sessions")
async def sessions(request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
pool = runtime.pool
return {
"rss": pool.rss() if pool is not None else None,
"rss_limit": pool.rss_limit if pool is not None else None,
"sessions": pool.snapshot() if pool is not None else [],
}
@app.get("/api/schedules")
async def schedules(request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
raw = request.query_params.get("conversation")
conv = await conv_of(raw) if raw else None
return {
"schedules": [
{
"id": s.id,
"conversation_row": s.conversation_id,
"execute_at": s.execute_at.isoformat(),
"text": s.text,
"delivered_at": s.delivered_at.isoformat()
if s.delivered_at
else None,
}
for s in await conversations.schedules(conv)
]
}
@app.get("/api/usage")
async def usage(request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
hours = float(request.query_params.get("hours", "24"))
since = datetime.now(UTC) - timedelta(hours=hours)
columns = (
func.count(),
func.coalesce(func.sum(Usage.input_tokens), 0),
func.coalesce(func.sum(Usage.output_tokens), 0),
func.coalesce(func.sum(Usage.cache_read_tokens), 0),
func.coalesce(func.sum(Usage.cache_creation_tokens), 0),
func.coalesce(func.sum(Usage.cost_usd), 0.0),
)
async with runtime.db.session() as session:
by_agent = (
await session.execute( # ty: ignore[deprecated]
sa_select(col(Usage.agent_name), *columns)
.where(col(Usage.ts) >= since.replace(tzinfo=None))
.group_by(col(Usage.agent_name))
)
).all()
by_conversation = (
await session.execute( # ty: ignore[deprecated]
sa_select(col(Usage.conversation_id), *columns)
.where(col(Usage.ts) >= since.replace(tzinfo=None))
.group_by(col(Usage.conversation_id))
)
).all()
return {
"since": since.isoformat(timespec="seconds"),
"by_agent": [_usage_row("agent", r) for r in by_agent],
"by_conversation": [_usage_row("conversation", r) for r in by_conversation],
}
@app.exception_handler(HTTPException)
async def http_error(_request: Request, exc: HTTPException) -> JSONResponse:
return JSONResponse(
status_code=exc.status_code,
content={"error": exc.detail},
headers=exc.headers,
)
return app
def _usage_row(label: str, row: Any) -> dict[str, Any]:
key, turns, inp, out, cache_read, cache_creation, cost = row
return {
label: key,
"turns": turns,
"input": inp,
"output": out,
"cache_read": cache_read,
"cache_creation": cache_creation,
"cost_usd": round(float(cost or 0.0), 4),
}
def _sse(runtime: GatewayRuntime, *, conversation_id: str | None) -> StreamingResponse:
async def gen() -> AsyncIterator[bytes]:
stream = runtime.bus.stream(conversation_id=conversation_id)
yield sse_pack("hello", {"type": "hello", "conversation_id": conversation_id})
async for event in events_with_heartbeat(stream):
if event is None:
yield KEEPALIVE
continue
yield sse_pack(str(event["type"]), event)
return StreamingResponse(gen(), media_type="text/event-stream", headers=SSE_HEADERS)
+7 -1
View File
@@ -12,7 +12,7 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable, Mapping, Sequence
@@ -77,6 +77,12 @@ class GatewayRuntime:
# to keep ``anthropic.types`` out of the runtime import graph for
# this base module.
turn_log_handlers: list[TurnLogHandler] = field(default_factory=list)
# M1b: conversations service, event bus and the shared session pool.
# ``Any`` for the same import-graph reason as above; ``None`` only in
# tests that build a runtime without them.
conversations: Any = None
bus: Any = None
pool: Any = None
class Frontend(ABC):
+84 -194
View File
@@ -13,7 +13,13 @@ Concurrency model: an in-memory ``set[Path]`` of files currently in
flight. Two concurrent requests for the same file → the second gets
409. The set is single-process (one gateway instance) — that's by
design; the markdown frontend is the only writer in its vault from
the gateway side.
the gateway side. The turn itself runs through ``core/conversations``
(one turn per conversation, ``running_turn`` in the DB), so a message
posted to the same conversation via ``/api`` waits its turn.
A chat file is a ``deep`` conversation bound as
``(markdown, <vault-relative path>)``; frontmatter carries only ``agent``
and ``conversation_id`` (§3.10), tool calls are never rendered.
Cross-frontend logging: when ``log_all_chats=True``, ``configure()``
registers a handler on ``runtime.turn_log_handlers`` so every other
@@ -24,7 +30,6 @@ shape.
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
@@ -33,7 +38,7 @@ import tempfile
import time
from collections.abc import AsyncIterator
from pathlib import Path
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
import aiofile
from anthropic.types import RawContentBlockStopEvent
@@ -44,29 +49,28 @@ from fastapi.responses import JSONResponse, StreamingResponse
from beaver_gateway.core import audit
from beaver_gateway.core.conversation_store import (
diff_and_fork,
load_conversation,
load_messages,
mint_conversation,
rewrite_messages,
set_session_id,
)
from beaver_gateway.core.turn_capture import TurnCapture
from beaver_gateway.core.turn_record import TurnRecord
from beaver_gateway.frontends._accumulate import StreamAccumulator
from beaver_gateway.frontends._auth import require_token
from beaver_gateway.frontends._sse import (
KEEPALIVE,
SSE_HEADERS,
events_with_heartbeat,
sse_pack,
)
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.frontends.markdown import parser, renderer
from beaver_gateway.frontends.markdown.crossfront import (
CrossFrontendLogger,
fingerprint_messages,
)
from beaver_gateway.frontends.markdown.crossfront import CrossFrontendLogger
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Callable
from anthropic.types import MessageParam
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.storage.models import Conversation
_log = logging.getLogger("beaver_gateway.frontends.markdown")
@@ -89,12 +93,7 @@ _STREAM_FLUSH_DEBOUNCE = 0.4
# disk round-trip).
_SSE_FLUSH_DEBOUNCE = 0.1
# Interval between SSE comment-frames sent when the backend is silent
# (e.g. claude is mid-thinking on a large context). The Obsidian plugin
# and any intermediate proxies will hold the connection open as long as
# bytes keep flowing; a comment-frame is the cheapest legal SSE keepalive.
# Set well under typical proxy/client idle timeouts (60s).
_SSE_HEARTBEAT_INTERVAL = 15.0
FRONTEND = "markdown"
class MarkdownFrontend(Frontend):
@@ -134,6 +133,9 @@ class MarkdownFrontend(Frontend):
self._crossfront: CrossFrontendLogger | None = None
def configure(self, runtime: GatewayRuntime) -> None:
if runtime.conversations is None:
msg = "MarkdownFrontend needs runtime.conversations"
raise RuntimeError(msg)
self._runtime = runtime
self.vault_path.mkdir(parents=True, exist_ok=True)
if self.log_all_chats:
@@ -286,17 +288,7 @@ class MarkdownFrontend(Frontend):
self._busy.discard(file_path)
return StreamingResponse(
gen(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache, no-transform",
"Connection": "keep-alive",
# nginx default-buffers SSE bodies; this header tells
# both nginx and uvicorn-behind-proxy to flush as we
# write. Harmless if the deployment has no reverse
# proxy in front.
"X-Accel-Buffering": "no",
},
gen(), media_type="text/event-stream", headers=SSE_HEADERS
)
return app
@@ -384,16 +376,19 @@ class MarkdownFrontend(Frontend):
# stored history, and feed the aligned messages to the backend
# - see ``core/conversation_store.py`` for the full rationale.
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
runtime=runtime, metadata=parsed.metadata, agent_name=agent.name
runtime=runtime,
metadata=parsed.metadata,
agent_name=agent.name,
file_path=file_path,
)
outcome = diff_and_fork(stored=stored_msgs, incoming=parsed.turns)
capture = TurnCapture()
events = backend.complete(
agent=agent,
events = runtime.conversations.turn(
conv,
messages=outcome.messages,
system=None,
origin="user",
capture=capture,
**_session_options(conv, outcome.divergence_index),
use_session=outcome.divergence_index is None,
)
try:
message = await self._stream_to_file(
@@ -422,7 +417,7 @@ class MarkdownFrontend(Frontend):
await self._persist_canonical_history(
runtime=runtime,
conversation_id=conv.id,
conversation_id=cast("int", conv.id),
persist_messages=outcome.persist_messages,
new_user_text=parsed.turns[-1].text,
capture=capture,
@@ -487,7 +482,7 @@ class MarkdownFrontend(Frontend):
elif content_override is None:
file_text = await _read_or_empty(file_path)
else:
yield _sse_pack(
yield sse_pack(
"error",
{
"status_code": status.HTTP_400_BAD_REQUEST,
@@ -503,7 +498,7 @@ class MarkdownFrontend(Frontend):
default=self.default_agent,
)
if not agent_name:
yield _sse_pack(
yield sse_pack(
"error",
{
"status_code": status.HTTP_400_BAD_REQUEST,
@@ -516,7 +511,7 @@ class MarkdownFrontend(Frontend):
return
if not parsed.messages:
yield _sse_pack(
yield sse_pack(
"done",
{
"status": "nothing_to_do",
@@ -527,7 +522,7 @@ class MarkdownFrontend(Frontend):
return
if parser.last_role(parsed.messages) == "assistant":
yield _sse_pack(
yield sse_pack(
"done",
{
"status": "nothing_to_do",
@@ -539,7 +534,7 @@ class MarkdownFrontend(Frontend):
agent = runtime.agents.get(agent_name)
if agent is None:
yield _sse_pack(
yield sse_pack(
"error",
{
"status_code": status.HTTP_404_NOT_FOUND,
@@ -549,7 +544,7 @@ class MarkdownFrontend(Frontend):
return
backend = runtime.backends.get(agent.name)
if backend is None:
yield _sse_pack(
yield sse_pack(
"error",
{
"status_code": status.HTTP_503_SERVICE_UNAVAILABLE,
@@ -575,7 +570,10 @@ class MarkdownFrontend(Frontend):
)
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
runtime=runtime, metadata=parsed.metadata, agent_name=agent.name
runtime=runtime,
metadata=parsed.metadata,
agent_name=agent.name,
file_path=file_path,
)
_log.info(
"chat/stream: file=%s conv_external_id=%s conv_id=%d "
@@ -602,12 +600,12 @@ class MarkdownFrontend(Frontend):
agent.name,
conv.session_id,
)
events = backend.complete(
agent=agent,
events = runtime.conversations.turn(
conv,
messages=outcome.messages,
system=None,
origin="user",
capture=capture,
**_session_options(conv, outcome.divergence_index),
use_session=outcome.divergence_index is None,
)
acc = StreamAccumulator()
@@ -624,13 +622,9 @@ class MarkdownFrontend(Frontend):
return _reattach_frontmatter(parsed.metadata, new_body)
try:
async for ev in _events_with_heartbeat(events):
async for ev in events_with_heartbeat(events):
if ev is None:
# Backend is quiet (claude mid-thinking, MCP slow,
# whatever). SSE comment-frame keeps the TCP socket
# warm so the plugin / uvicorn / any reverse proxy
# doesn't time the request out before we finish.
yield b": keepalive\n\n"
yield KEEPALIVE
continue
acc.feed(ev)
now = time.monotonic()
@@ -643,7 +637,7 @@ class MarkdownFrontend(Frontend):
# render to the same prefix as before they closed
# (we don't surface the tool-call args in markdown).
if payload is not None and payload != last_payload:
yield _sse_pack("delta", {"new_content": payload})
yield sse_pack("delta", {"new_content": payload})
last_payload = payload
last_flush = now
except Exception as exc: # noqa: BLE001 — wire any backend failure as an SSE error frame
@@ -662,7 +656,7 @@ class MarkdownFrontend(Frontend):
await _write_atomic(
file_path, _reattach_frontmatter(parsed.metadata, new_body)
)
yield _sse_pack(
yield sse_pack(
"error",
{
"status_code": status.HTTP_500_INTERNAL_SERVER_ERROR,
@@ -684,7 +678,7 @@ class MarkdownFrontend(Frontend):
await self._persist_canonical_history(
runtime=runtime,
conversation_id=conv.id,
conversation_id=cast("int", conv.id),
persist_messages=outcome.persist_messages,
new_user_text=parsed.turns[-1].text,
capture=capture,
@@ -704,7 +698,7 @@ class MarkdownFrontend(Frontend):
except Exception: # noqa: BLE001
_log.exception("turn_log_handler raised; continuing")
yield _sse_pack(
yield sse_pack(
"done",
{
"status": "ok",
@@ -795,71 +789,53 @@ class MarkdownFrontend(Frontend):
rendered = renderer.render_assistant_message(message)
new_body = renderer.append_to_body(parsed.body, rendered)
new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD)
# Recompute fingerprint so a future cross-frontend hit on this
# same conversation can find it. Stored as hex string in
# frontmatter — only the markdown frontend reads it.
assistant_param: MessageParam = {
"role": "assistant",
"content": _flatten_assistant_text(message),
}
updated_messages: list[MessageParam] = [*parsed.messages, assistant_param]
updated_metadata = dict(parsed.metadata)
updated_metadata.pop("fingerprint", None)
updated_metadata["agent"] = agent_name
updated_metadata["conversation_id"] = conv_external_id
updated_metadata["fingerprint"] = fingerprint_messages(updated_messages)
new_content = _reattach_frontmatter(updated_metadata, new_body)
if write_disk:
await _write_atomic(file_path, new_content)
return new_content
async def _resolve_conversation(
self, *, runtime: GatewayRuntime, metadata: dict[str, Any], agent_name: str
) -> tuple[Any, str, list[dict[str, Any]]]:
"""Resolve the conversation row + stored messages for this request.
self,
*,
runtime: GatewayRuntime,
metadata: dict[str, Any],
agent_name: str,
file_path: Path,
) -> tuple[Conversation, str, list[dict[str, Any]]]:
"""Resolve the ``deep`` conversation for this file + its stored messages.
Looks up by frontmatter ``conversation_id``, mints a new row if
missing, and returns ``(conv, external_id, stored_messages)``.
``conv.id`` is guaranteed non-None because both
``load_conversation`` (after refresh on a committed row) and
``mint_conversation`` (post-commit refresh) populate it. We
coerce with a runtime check so the rest of the handler can
treat it as ``int``.
Frontmatter ``conversation_id`` wins; a file that lost it is found
by its visible ``(markdown, path)`` binding; otherwise a new
conversation is created. The binding follows the file: a moved
chat re-binds to its new path on the next turn.
"""
conversations = runtime.conversations
rel = file_path.relative_to(self.vault_path).as_posix()
raw = metadata.get("conversation_id")
lookup_id = raw if isinstance(raw, str) and raw else None
conv = await conversations.get(raw) if isinstance(raw, str) and raw else None
if conv is None:
conv = await conversations.find_bound(frontend=FRONTEND, external_id=rel)
if conv is None:
conv = await conversations.create(
kind="deep", agent=agent_name, origin=FRONTEND, title=file_path.stem
)
_log.info("minted conversation %s for %s", conv.external_id, rel)
bound = [
b
for b in await conversations.bindings(conv)
if b.frontend == FRONTEND and b.visible and b.external_id == rel
]
if not bound:
await conversations.bind(conv, frontend=FRONTEND, external_id=rel)
await conversations.touch_user(conv)
if conv.id is None:
msg = "conversation row missing primary key after commit"
raise RuntimeError(msg)
async with runtime.db.session() as session:
conv = None
if lookup_id is not None:
conv = await load_conversation(
session, frontend="markdown", external_id=lookup_id
)
if conv is None:
_log.info(
"_resolve_conversation: frontmatter conv_id=%s "
"not found in DB, will mint new",
lookup_id,
)
else:
_log.info(
"_resolve_conversation: LOADED existing conv "
"id=%d external_id=%s",
conv.id or -1,
conv.external_id,
)
if conv is None:
conv = await mint_conversation(
session, frontend="markdown", agent_name=agent_name
)
_log.info(
"_resolve_conversation: MINTED new conv "
"id=%d external_id=%s agent=%s",
conv.id or -1,
conv.external_id,
agent_name,
)
if conv.id is None:
msg = "conversation row missing primary key after commit"
raise RuntimeError(msg)
stored = await load_messages(session, conversation_id=conv.id)
return conv, conv.external_id, stored
@@ -894,12 +870,6 @@ class MarkdownFrontend(Frontend):
await rewrite_messages(
session, conversation_id=conversation_id, messages=canonical
)
if capture.session_id is not None:
await set_session_id(
session,
conversation_id=conversation_id,
session_id=capture.session_id,
)
_log.info(
"_persist_canonical_history: conv_id=%d DB committed", conversation_id
)
@@ -927,71 +897,6 @@ class MarkdownFrontend(Frontend):
# ---- module-level utilities ----------------------------------------------
def _session_options(conv: Any, divergence_index: int | None) -> dict[str, Any]:
"""Backend options that pin the turn to the conversation's live session.
A divergence means the file's history no longer matches what the
session saw, so the stored ``session_id`` is withheld and the backend
seeds a fresh one from the aligned messages.
"""
return {
"conversation_id": conv.external_id,
"session_id": conv.session_id if divergence_index is None else None,
}
async def _events_with_heartbeat(
events: AsyncIterator[Any], interval: float = _SSE_HEARTBEAT_INTERVAL
) -> AsyncIterator[Any]:
"""Wrap an async event stream with idle-time heartbeat markers.
Yields ``None`` every ``interval`` seconds during silence; real
events pass through unchanged. When the wrapped iterator is
exhausted, this generator returns. Cancellation propagates: if the
outer scope is cancelled we cancel the pending ``__anext__`` task
instead of leaving it dangling.
"""
src = events.__aiter__()
next_task: asyncio.Task[Any] | None = None
try:
while True:
# Reuse the in-flight task across timeouts. Spawning a fresh
# ``__anext__()`` while the previous one is still pending
# puts two consumers on the same async generator — that
# raises ``RuntimeError: anext(): asynchronous generator is
# already running``.
if next_task is None:
next_task = asyncio.ensure_future(src.__anext__())
done, _pending = await asyncio.wait({next_task}, timeout=interval)
if not done:
yield None
continue
task = next_task
next_task = None
try:
result = task.result()
except StopAsyncIteration:
return
yield result
finally:
if next_task is not None and not next_task.done():
next_task.cancel()
with contextlib.suppress(BaseException):
await next_task
def _sse_pack(event: str, data: dict[str, Any]) -> bytes:
r"""Format one Server-Sent Event frame.
Uses named events (``event: <name>``) so the plugin can dispatch on
type without parsing JSON discriminators. ``ensure_ascii=False`` so
multibyte content rides through verbatim instead of becoming
``\uXXXX`` blobs that bloat the wire.
"""
body = json.dumps(data, ensure_ascii=False)
return f"event: {event}\ndata: {body}\n\n".encode()
async def _read_or_empty(path: Path) -> str:
"""Return file contents, or empty string if the file doesn't exist."""
# ``path.exists()`` here is a metadata stat — microseconds — and
@@ -1083,21 +988,6 @@ def _fallback_synthesized(message: Any) -> list[dict[str, Any]]:
return [{"role": "assistant", "content": content}]
def _flatten_assistant_text(message: Any) -> str:
"""Pull all text blocks from an assistant ``Message`` and join them.
Used when we need the assistant content as a plain string for
fingerprinting / equality with a parser-shaped history (parser
already drops thinking + tool_use from assistant turns).
"""
chunks = [
getattr(block, "text", "") or ""
for block in getattr(message, "content", ())
if getattr(block, "type", None) == "text"
]
return "\n\n".join(c for c in chunks if c)
def _render_error_block(exc: BaseException) -> str:
"""Render a backend failure as an Assistant turn with a ``[!error]-`` callout."""
msg = str(exc) or exc.__class__.__name__
@@ -9,11 +9,10 @@ from other frontends).
from __future__ import annotations
import json
import re
from typing import TYPE_CHECKING, Any, cast
from anthropic.types import Message, TextBlock, ThinkingBlock, ToolUseBlock
from anthropic.types import Message, TextBlock, ThinkingBlock
if TYPE_CHECKING:
from collections.abc import Iterable
@@ -61,9 +60,8 @@ def render_assistant_message(message: Message) -> str:
* ``ThinkingBlock`` → ``> [!thinking]-`` collapsed callout
* ``TextBlock`` → plain text (the spoken answer)
* ``ToolUseBlock`` → ``> [!tool]- <name>`` callout with the ``input``
JSON quoted inside. Tool *results* are not persisted — see
module docstring on ``parser.py`` for why.
* ``ToolUseBlock`` → nothing (§3.10: tool calls never reach the file;
"what the agent is doing" is the activity panel fed by SSE)
Blank lines separate adjacent blocks; trailing newline guarantees
the next ``---`` / ``### User:`` marker lands on its own line.
@@ -147,10 +145,7 @@ def _render_block(block: object) -> Iterable[str]:
if isinstance(block, ThinkingBlock):
yield from _render_thinking(block.thinking or "")
return
if isinstance(block, ToolUseBlock):
yield from _render_tool_use(block)
return
# Unknown block type — skip silently rather than corrupting the file.
# Tool-use blocks and unknown block types never reach the file.
def _render_thinking(text: str) -> Iterable[str]:
@@ -159,17 +154,6 @@ def _render_thinking(text: str) -> Iterable[str]:
yield f"> {line}" if line else ">"
def _render_tool_use(block: ToolUseBlock) -> Iterable[str]:
title = summarize_tool_input(block.name, block.input)
yield f"> [!tool]- {title}"
yield "> **input:**"
yield "> ```json"
pretty = json.dumps(block.input, indent=2, ensure_ascii=False, sort_keys=True)
for line in pretty.splitlines():
yield f"> {line}" if line else ">"
yield "> ```"
def adaptive_fence(content: str) -> str:
"""Return a backtick fence at least one longer than the longest run in ``content``.
+14 -1
View File
@@ -16,13 +16,26 @@ from beaver_gateway.storage.db import (
revoke_token,
touch_token,
)
from beaver_gateway.storage.models import AuditLog, Token, TranscriptEntry, Usage
from beaver_gateway.storage.models import (
AuditLog,
Conversation,
ConversationBinding,
InjectQueueItem,
Schedule,
Token,
TranscriptEntry,
Usage,
)
from beaver_gateway.storage.session_store import PostgresSessionStore
__all__ = [
"AuditLog",
"Conversation",
"ConversationBinding",
"Database",
"InjectQueueItem",
"PostgresSessionStore",
"Schedule",
"Token",
"TranscriptEntry",
"Usage",
+23 -3
View File
@@ -105,9 +105,29 @@ def _add_missing_columns(conn: Connection) -> None:
if column.name in existing:
continue
kind = column.type.compile(conn.dialect)
conn.execute(
text(f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}")
)
ddl = f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}"
default = _column_default(column)
if default is not None:
ddl += f" DEFAULT {default}"
conn.execute(text(ddl))
def _column_default(column: Any) -> str | None:
"""Literal for ``ADD COLUMN ... DEFAULT`` so old rows get the model default."""
server_default = getattr(column, "server_default", None)
if server_default is not None:
return str(server_default.arg.text)
default = getattr(column, "default", None)
if default is None or default.is_callable:
return None
value = default.arg
if isinstance(value, bool):
return "TRUE" if value else "FALSE"
if isinstance(value, int | float):
return repr(value)
if isinstance(value, str):
return "'" + value.replace("'", "''") + "'"
return None
# ---- Token CRUD ---------------------------------------------------------
+97 -9
View File
@@ -1,6 +1,6 @@
"""SQLModel tables.
Four tables, all flat, no FK relationships modelled (``actor`` and
Flat tables, no FK relationships modelled (``actor`` and
``agent_name`` are stored as strings — joining audit→token by name is
fine at this volume; we'll introduce FKs when the admin UI actually
demands them).
@@ -74,15 +74,16 @@ class AuditLog(SQLModel, table=True):
class Conversation(SQLModel, table=True):
"""One chat thread, scoped to a frontend.
"""One conversation (§3.1): a master thread, a branch, a deep chat or a job.
``external_id`` is the identifier the frontend uses to find this
thread again on the next request — for the markdown frontend it's a
uuid we mint and persist into the file's frontmatter, for the
anthropic frontend it'd be the same metadata.conversation_id the
client passes. Unique per ``(frontend, external_id)`` because two
frontends sharing a uuid is fine; the same frontend reusing one is
a bug.
``external_id`` is the public id (uuid) every frontend, the API and the
usage table refer to; ``frontend`` names the frontend that created the
row (``markdown``, ``api``, ``mcp``, ``system``). Where a conversation
is *visible* lives in :class:`ConversationBinding`.
``running_turn`` survives a restart: a non-null value at startup means
the gateway died mid-turn and the transcript needs its open ``tool_use``
closed before the session is resumed (``core/conversations``).
"""
__tablename__ = "conversations"
@@ -95,8 +96,92 @@ class Conversation(SQLModel, table=True):
external_id: str = Field(index=True)
agent_name: str = Field(index=True)
session_id: str | None = Field(default=None, index=True)
kind: str = Field(default="deep", index=True)
parent_id: int | None = Field(default=None, index=True)
title: str | None = Field(default=None)
status: str = Field(default="open", index=True)
running_turn: str | None = Field(default=None)
pending_question: bool = Field(default=False)
flags: dict[str, Any] = Field(
default_factory=dict,
sa_column=Column(
JSON().with_variant(JSONB(), "postgresql"),
nullable=False,
server_default=text("'{}'"),
),
)
created_at: datetime = Field(default_factory=_utcnow)
updated_at: datetime = Field(default_factory=_utcnow)
last_user_activity_at: datetime | None = Field(default=None)
last_activity_at: datetime | None = Field(default=None)
class ConversationBinding(SQLModel, table=True):
"""Where a conversation shows up: a Telegram topic id, a vault-relative path.
Invariant (§3.1): at most one *visible* binding per frontend per
conversation - enforced by the partial unique index. The same external
id may point at several conversations over time (a renamed topic gets a
new branch), only one of them visible.
"""
__tablename__ = "conversation_bindings"
__table_args__ = (
Index(
"uq_binding_visible",
"conversation_id",
"frontend",
unique=True,
postgresql_where=text("visible"),
sqlite_where=text("visible"),
),
Index("ix_binding_lookup", "frontend", "external_id"),
)
id: int | None = Field(default=None, primary_key=True)
conversation_id: int = Field(index=True)
frontend: str
external_id: str
visible: bool = Field(default=True)
created_at: datetime = Field(default_factory=_utcnow)
class InjectQueueItem(SQLModel, table=True):
"""Persisted per-conversation queue (§3.4), priorities ``urgent > user > normal``.
``status`` walks ``queued -> running -> done``; a row still ``running``
at startup was cut by a restart and becomes ``interrupted`` - it is
never re-run, the agent gets an inject saying so instead.
"""
__tablename__ = "inject_queue"
id: int | None = Field(default=None, primary_key=True)
conversation_id: int = Field(index=True)
priority: str = Field(index=True)
origin: str = Field(default="system")
text: str
status: str = Field(default="queued", index=True)
turn_id: str | None = Field(default=None)
created_at: datetime = Field(default_factory=_utcnow)
delivered_at: datetime | None = Field(default=None)
class Schedule(SQLModel, table=True):
"""Deferred inject written by the ``schedule`` tool (§3.6).
M1b only records the promise; the executor (pgqueuer, M3) will move
these into its own job table and this one goes away.
"""
__tablename__ = "schedules"
id: int | None = Field(default=None, primary_key=True)
conversation_id: int = Field(index=True)
execute_at: datetime = Field(index=True)
text: str
created_at: datetime = Field(default_factory=_utcnow)
delivered_at: datetime | None = Field(default=None)
class ConversationMessage(SQLModel, table=True):
@@ -193,7 +278,10 @@ class Usage(SQLModel, table=True):
__all__ = [
"AuditLog",
"Conversation",
"ConversationBinding",
"ConversationMessage",
"InjectQueueItem",
"Schedule",
"Token",
"TranscriptEntry",
"Usage",