Compare commits

...
6 Commits
49 changed files with 6722 additions and 1271 deletions

No files matched your search

+1
View File
@@ -11,6 +11,7 @@ wheels/
# Sloppy
docs/
t/
# IDE
.idea
+13 -17
View File
@@ -44,26 +44,22 @@ RUN apt-get update \
&& apt-get install -y --no-install-recommends ca-certificates git \
&& rm -rf /var/lib/apt/lists/*
# Bun native binary (glibc) — fast `npm install` replacement, only used
# to drop the claude CLI into the image. Not invoked at runtime.
# Bun native binary (glibc): `bunx` runs the stdio MCP servers declared in
# the user's config. The claude CLI itself ships inside the
# claude-agent-sdk wheel (`_bundled/claude`), nothing to install here.
COPY --from=oven/bun:1-slim /usr/local/bin/bun /usr/local/bin/bun
RUN ln -s /usr/local/bin/bun /usr/local/bin/bunx
# `--trust` is required: without it bun skips the postinstall step that
# fetches claude's native binary (anthropics/claude-code#50203). The
# postinstall itself is bun's smoke check — if it fails the layer
# fails. We deliberately DO NOT run `claude --version` here: claude
# touches `$HOME` on every invocation (creates `/root/.claude/`,
# `/root/.claude.json`, sometimes `/root/.config/claude/`), and those
# build-time artifacts seed the runtime named-volume `claude-home`
# with stale "haven't onboarded" state, so the user gets re-prompted
# for trust/bypass dialogs on every rebuild and the subscription auth
# can land on a tainted credential file.
ENV BUN_INSTALL=/usr/local/bun-global \
PATH=/usr/local/bun-global/bin:/app/.venv/bin:$PATH
RUN bun install -g --trust @anthropic-ai/claude-code \
&& test -x "$(command -v claude)" \
&& rm -rf /root/.claude /root/.claude.json /root/.config/claude
# The model process runs as `beaver-runner`, not as the gateway: the
# adapter spawns claude under this uid with a whitelisted environment.
# `acl` lets the entrypoint grant it write access to the vault sub-mounts
# without chowning files that Obsidian Sync keeps rewriting as root.
RUN apt-get update \
&& apt-get install -y --no-install-recommends acl \
&& rm -rf /var/lib/apt/lists/* \
&& useradd --system --uid 1001 --create-home --shell /usr/sbin/nologin beaver-runner
ENV PATH=/app/.venv/bin:$PATH
COPY --from=builder /app/.venv /app/.venv
COPY --from=builder /app /app
+16 -21
View File
@@ -12,7 +12,7 @@ from beaver_gateway.agents.base import ExposedMcp
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.agents.raycast import RaycastAgent, RemoteTool, UserPreferences
from beaver_gateway.core.registry import Gateway
from beaver_gateway.core.turn_record import TurnRecord, slugify
from beaver_gateway.core.turn_record import slugify
from beaver_gateway.frontends.admin import AdminFrontend
from beaver_gateway.frontends.anthropic import AnthropicMessagesFrontend
from beaver_gateway.frontends.markdown import MarkdownFrontend
@@ -20,20 +20,19 @@ from beaver_gateway.frontends.mcp_server import McpServerFrontend
from beaver_gateway.mcp.types import McpServer
def chat_log_path(record: TurnRecord, vault: Path) -> Path:
"""Decide where a logged chat from another frontend lands in the vault.
def chat_path(title: str, agent: str, vault: Path) -> Path: # noqa: ARG001
"""Where a new chat file lands in the vault.
Called by ``MarkdownFrontend`` (with ``log_all_chats=True``) the first
time a conversation needs a file — continuation turns are matched by
fingerprint and stick to the file picked here. Return value can be
Called by ``MarkdownFrontend`` for every conversation that needs a
file: a ``deep`` chat spawned by the dispatcher, a ``/v1/messages``
chat (``title`` = first user message), or with ``log_all_chats=True``
the archive of a stateless agent's turns. Return value can be
absolute or relative; relative paths are anchored under ``vault``.
Layout below: ``<vault>/<YYYY-MM>/<YYYY-MM-DD>_<topic>.md`` where
``topic`` is a slug of the very first user message in the chat.
Layout below: ``<vault>/<YYYY-MM>/<YYYY-MM-DD>_<topic>.md``.
"""
today = date.today()
topic = slugify(record.first_user_text, maxlen=40)
return vault / f"{today:%Y-%m}" / f"{today:%Y-%m-%d}_{topic}.md"
return vault / f"{today:%Y-%m}" / f"{today:%Y-%m-%d}_{slugify(title, maxlen=40)}.md"
def current_time() -> str:
@@ -208,17 +207,13 @@ gateway = Gateway(
vault_path=Path(tempfile.mkdtemp(prefix="beaver-vault-")).resolve(),
default_agent="research",
log_all_chats=True,
# ``log_path`` (optional) overrides the default
# ``{vault}/_logs/<agent>/<date>_<hex>.md`` layout for chats
# logged from OTHER frontends (Anthropic Messages, admin
# in-browser chat). Defined as a top-of-file function so the
# types are explicit and the IDE can hover them; a lambda
# works too, but a real ``def`` keeps the signature visible
# and lets you docstring it. Heads up: any custom path
# forces ``warm_index`` to scan the entire vault on startup
# so the fingerprint→file map survives a restart no matter
# where you put files.
log_path=chat_log_path,
# ``chat_path`` (optional) overrides the default
# ``{vault}/_logs/<agent>/<date>_<slug>.md`` layout. Heads up:
# any custom path forces ``warm_index`` to scan the entire
# vault on startup so the fingerprint→file map of archived
# stateless chats survives a restart no matter where you put
# files.
chat_path=chat_path,
),
],
)
+6 -6
View File
@@ -13,11 +13,13 @@ dependencies = [
"anthropic>=0.103.0",
"anyio>=4.13.0",
"argon2-cffi>=25.1.0",
"claude-agent-sdk>=0.2.146",
"fastapi>=0.136.1",
"fastmcp>=3.3.1",
"greenlet>=3.5.0",
"itsdangerous>=2.2.0",
"jinja2>=3.1.6",
"psutil>=7.2.2",
"psycopg[binary]>=3.3.4",
"pydantic>=2.13.4",
"pydantic-settings>=2.14.1",
@@ -30,11 +32,9 @@ dependencies = [
[project.optional-dependencies]
local = [
"raycast-api",
"claude-code-api",
]
prod = [
"raycast-api",
"claude-code-api",
]
[tool.uv]
@@ -47,10 +47,6 @@ raycast-api = [
{ path = "../raycast-api", editable = true, extra = "local" },
{ git = "https://git.kotikot.com/beaver/raycast-api.git", extra = "prod" },
]
claude-code-api = [
{ path = "../claude-code-api", editable = true, extra = "local" },
{ git = "https://git.kotikot.com/beaver/claude-code-api.git", extra = "prod" },
]
[project.scripts]
beaver-gateway = "beaver_gateway:main"
@@ -67,3 +63,7 @@ dev = [
"ruff>=0.15.13",
"ty>=0.0.37",
]
[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]
+8 -2
View File
@@ -2,7 +2,7 @@
``BaseAgent`` is the structural contract every backend-specific agent
honours. Subclasses add **capabilities as fields** (``ClaudeAgent.cwd``,
``RaycastAgent.streaming``) — backends dispatch on type, not on a runtime
``RaycastAgent.streaming``) - backends dispatch on type, not on a runtime
matrix.
"""
@@ -15,10 +15,16 @@ from pydantic import BaseModel, ConfigDict
@dataclass(frozen=True, slots=True)
class ExposedMcp:
"""Reference to an ``McpServer`` (by name) exposed to a single agent."""
"""Reference to an ``McpServer`` (by name) exposed to a single agent.
``tools`` is an allowlist of tool names (``None`` = every tool);
``deny`` is a tuple of ``fnmatch`` patterns removed on top of that,
e.g. ``("delete_*",)``.
"""
name: str
tools: tuple[str, ...] | None = None
deny: tuple[str, ...] = ()
class BaseAgent(BaseModel):
+77 -229
View File
@@ -1,258 +1,106 @@
"""Claude Code agent definition.
"""Claude agent definition, backed by the Claude Agent SDK.
There is no ``streaming`` field, but there is
:attr:`ClaudeCodeOptions.include_partial_messages`, and the difference
matters. Token-level deltas are not a property of the agent, they are a
property of the transport: the PTY transport reads a JSONL that only
ever contains finished blocks, so it cannot stream no matter what the
agent asks for, while ``transport="stream_json"`` can. The flag lives
next to the transport that grants it rather than pretending to be a
free-standing capability.
``BaseAgent.system_prompt`` maps onto the claude CLI's
``--system-prompt`` — i.e. it really *is* the agent's system prompt,
not "added on top of claude-code's giant built-in". The additive slot
``--append-system-prompt`` is exposed via
:attr:`ClaudeCodeOptions.append_system_prompt` for the rare case
when the user wants claude-code's planning conventions / dynamic
sections *and* a delta on top. Difference, measured empirically:
* tools (names, JSON schemas, embedded guidance like Bash's "prefer
Read over cat/head/tail") survive both flags — they ride the
Anthropic API ``tools=[]`` field, not the system prompt text;
* ``--system-prompt`` drops ~8.6k tokens of claude-code's *textual*
baseline — agent persona, multi-step-work conventions ("use
TaskCreate proactively", etc.), and the dynamic per-machine sections
(cwd, env info, git status, memory paths — see the
``--exclude-dynamic-system-prompt-sections`` flag note in
``claude --help``: "Only applies with the default system prompt
(ignored with --system-prompt)").
We pick override as the default because it preserves the principle of
least surprise: the ``system_prompt=...`` you wrote on the agent is
what claude actually receives. If you need claude-code's full
planning/dynamic-context behaviour on top of your prompt, opt in via
``options=ClaudeCodeOptions(append_system_prompt=...)`` and leave
``system_prompt`` for your agent's identity (or vice versa — set
``system_prompt=""``-ish and put everything in ``append_*``).
Per-agent passthrough of ``claude_code_api.BackendOptions`` lives in
:class:`ClaudeCodeOptions`, attached to :class:`ClaudeAgent` as
``options``. Every tunable knob ``BackendOptions`` supports — except
the ones derived from the agent's own primary surface
(``cwd`` / ``model`` / ``system_prompt`` / ``available_native_tools`` →
``allowed_tools`` / ``expose_mcps`` → ``mcp_servers``) — is exposed
there. Defaults match upstream except where correctness demands
otherwise (``wait_for_turn_duration=True``,
``dangerously_skip_permissions=True``).
The system prompt is ``system_prompt`` verbatim or, per conversation kind,
the granules named in ``prompts`` assembled at every session spawn (see
``core/prompt.py``). ``skill_sets`` are directories of ``<skill>/SKILL.md``
folders; each becomes a local SDK plugin. Nothing from disk is loaded
otherwise: the adapter runs with ``setting_sources=[]``.
"""
from __future__ import annotations
from collections.abc import Mapping # noqa: TC003 — pydantic runtime
from pathlib import Path # noqa: TC003 — pydantic runtime
from typing import Literal
from collections.abc import Mapping # noqa: TC003 - pydantic runtime
from pathlib import Path # noqa: TC003 - pydantic runtime
from pydantic import BaseModel, ConfigDict, Field
from pydantic import BaseModel, ConfigDict, Field, model_validator
from beaver_gateway.agents.base import BaseAgent
from beaver_gateway.core.kinds import KINDS, Kind
from beaver_gateway.core.prompt import PromptSource # noqa: TC001 - pydantic runtime
HistoryInjectionMode = Literal["native_jsonl", "concat_message"]
"""Mirrors ``claude_code_api.HistoryInjectionMode`` to avoid an import
cycle on the user-facing path (configs may be loaded without
``claude-code-api`` installed, e.g. ``--extra prod`` minus claude)."""
Transport = Literal["pty", "stream_json"]
"""Mirrors ``claude_code_api.Transport``, for the same reason."""
__all__ = ["ClaudeAgent", "ClaudeOptions", "Prompts"]
class ClaudeCodeOptions(BaseModel):
"""Per-agent passthrough for ``claude_code_api.BackendOptions``.
Mirrors every field of ``BackendOptions`` except those that come
from the agent's primary surface
(``cwd`` / ``model`` / ``system_prompt`` /
``available_native_tools`` / ``expose_mcps``). Defaults match
upstream, with two exceptions for correctness:
* ``wait_for_turn_duration=True`` — without it, extended-thinking
turns drop the text response because ``TurnManager`` returns on
the first terminal assistant record (the thinking snapshot) and
never reads the second one (the actual text). Phase 2.2 PROGRESS
describes the incident in detail.
* ``dangerously_skip_permissions=True`` — Beaver Gateway is a
single-user trusted-config product; the user already wrote the
``config.py`` that spawns claude. Permission prompts on a
headless backend mean stuck turns.
Any field on :class:`claude_code_api.BackendOptions` we forward
here lives on this model with the same name and type, so a future
``BackendOptions`` field addition is a one-liner here plus a
one-liner in :mod:`backends.claude_code`.
"""
class ClaudeOptions(BaseModel):
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
transport: Transport = "pty"
"""Which claude process backs this agent.
``stream_json`` runs ``claude -p`` with stream-json on stdin and
stdout: prompts go into a pipe, events come back as they happen, and
a native ``result`` record closes each turn. ``pty`` (the default,
kept for backwards compatibility) drives the interactive TUI and
tails the session JSONL, which costs a multi-second readiness wait
per spawn — the thing ``startup_delay`` bounds — and brings the
swallowed-paste failure mode with it.
Switching also makes the PTY-only knobs below inert
(``startup_delay`` / ``file_wait_timeout`` /
``turn_duration_timeout`` / ``wait_for_turn_duration``) and drops
interactive-only ``extra_args`` such as ``--remote-control``, which
would take claude out of print mode. Everything else — model,
prompts, tools, MCPs, history seeding — behaves identically."""
warmup_turn: bool | None = None
"""Spend a throwaway turn when a session spawns, so MCP tools are
live on the agent's first reply.
``None`` (the default) decides per agent: on when it exposes any
MCP, off when it doesn't. Set it explicitly to override.
It exists because the CLI begins connecting MCP servers when a turn
starts and doesn't wait for them — measured on the pi, turn one saw
29 tools with both servers ``pending``, turn two saw 90 with both
``connected``. For an agent whose whole job is those tools, a first
reply without them is worse than a couple of extra seconds. Nothing
cheaper works: sleeping longer changes nothing, and neither a
``/status`` command nor the SDK's ``initialize`` handshake touches
the MCP client.
``stream_json`` only — the PTY transport's readiness wait covers
this incidentally."""
include_partial_messages: bool = False
"""Stream token-level deltas to the client instead of one delta per
finished block. Requires ``transport="stream_json"``; inert on the
PTY transport, whose JSONL contains only finished blocks. Off by
default because it multiplies the SSE event count — worth it for
interactive chat, pointless for a client that waits for the whole
message anyway."""
append_system_prompt: str | None = None
"""Maps to claude CLI's ``--append-system-prompt``. Opting in
re-attaches claude-code's full built-in prompt (persona, planning
conventions like "use TaskCreate proactively", and the dynamic
per-machine sections — cwd, env info, git status, memory paths)
plus this string on top. By default we ship only
:attr:`BaseAgent.system_prompt` via ``--system-prompt`` (~8.6k
tokens lighter); use this when the agent should behave like a
real claude-code coding session and your text is a delta on top
of those built-ins. Tool schemas survive either way — they ride
the Anthropic ``tools=[]`` channel, not the prompt text."""
disallowed_tools: tuple[str, ...] = ()
"""Tool names that claude must refuse to call. Combines with the
agent's ``available_native_tools`` allowlist — disallow wins."""
effort: str | None = None
include_partial_messages: bool = True
"""Token-level deltas on the wire; off means one delta per finished block."""
permission_mode: str = "bypassPermissions"
"""``claude --permission-mode`` value. Only meaningful when
``dangerously_skip_permissions=False``; otherwise the CLI bypasses
the prompt path entirely."""
tools: tuple[str, ...] | None = None
"""Base set of built-in tools; ``None`` keeps the CLI default set."""
dangerously_skip_permissions: bool = True
"""Pass ``--dangerously-skip-permissions`` to claude. ``True`` by
default — see the class docstring for why."""
disallowed_tools: tuple[str, ...] = ()
add_dirs: tuple[str, ...] = ()
env: Mapping[str, str] = Field(default_factory=dict)
"""Extra environment for the claude subprocess (always passed through)."""
effort: str | None = None
"""``claude --effort`` value (typically ``low`` / ``medium`` /
``high``). Tunes reasoning effort budget for newer models."""
env_keep: tuple[str, ...] = ()
"""Extra inherited variable names to let through the env whitelist."""
add_dir: tuple[str, ...] = ()
"""Extra directories claude is allowed to read/edit beyond
``cwd``. Maps to repeated ``--add-dir`` flags."""
max_turns: int | None = None
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."""
settings: str | None = None
"""Path to a ``--settings`` JSON file claude should load (hooks,
MCP servers, etc. that don't belong in our internal aggregator)."""
extra_args: tuple[str, ...] = ()
"""Raw extra argv to append to the claude command. Escape hatch
for flags we haven't surfaced as first-class options."""
class Prompts(BaseModel):
"""Prompt assembly per conversation kind (§3.12): the granules, in order.
extra_env: Mapping[str, str] = Field(default_factory=dict)
"""Additional environment variables for the spawned claude
process. Layered on top of the gateway's own env after
``preserve_provider_env`` is applied."""
A kind left ``None`` is not served by the agent; ``ClaudeAgent.kinds``
follows from the kinds set here.
"""
preserve_provider_env: bool = False
"""When ``False`` (default), the spawn env strips
``ANTHROPIC_API_KEY`` / ``ANTHROPIC_AUTH_TOKEN`` /
``ANTHROPIC_BASE_URL`` so claude uses subscription auth instead of
leaking through whatever the gateway process inherited."""
model_config = ConfigDict(frozen=True)
history_injection_mode: HistoryInjectionMode = "native_jsonl"
"""How prior turns are seeded into a fresh session when no live
session matches: ``native_jsonl`` writes a hand-crafted transcript
and ``--resume``s; ``concat_message`` instead folds history into
the first user prompt. ``native_jsonl`` is more faithful."""
master: tuple[PromptSource, ...] | None = None
branch: tuple[PromptSource, ...] | None = None
deep: tuple[PromptSource, ...] | None = None
job: tuple[PromptSource, ...] | None = None
fork: tuple[PromptSource, ...] | None = None
wait_for_turn_duration: bool = True
"""Keep reading JSONL until the ``turn_duration`` heartbeat
arrives, instead of returning on the first terminal assistant
record. ``True`` by default — see class docstring."""
def for_kind(self, kind: Kind) -> tuple[PromptSource, ...] | None:
return getattr(self, kind)
include_meta_user: bool = False
"""Surface claude's ``isMeta=True`` user records (local-command
caveats) as ``UserMessage`` events. Off by default — they're not
part of the real conversation."""
startup_delay: float = 60.0
"""Upper bound (seconds) on waiting for claude's TUI to enable
bracketed-paste mode after spawn. claude-code-api polls PTY output
for the DECSET 2004 marker and returns as soon as it arrives, so a
generous cap costs nothing on a warm host — it only bounds the case
where the marker never shows.
It has to be generous, because timing out here means writing the
prompt into a terminal that hasn't enabled bracketed paste, where our
``ESC[200~`` framing is parsed as junk and the message is silently
dropped. Observed on a cold Raspberry Pi: 10s elapsed with *zero*
bytes of PTY output, the prompt went into the void, and the turn hung
until a human typed into the session by hand."""
file_wait_timeout: float = 30.0
"""How long to wait for the session JSONL to appear after spawn.
Failure here usually means a CLI auth / config problem."""
turn_duration_timeout: float = 5.0
"""How long to wait for the ``turn_duration`` heartbeat once a
terminal assistant has been seen. Bound on extra latency when
``wait_for_turn_duration=True``."""
idle_session_ttl: float = 1800.0
"""Seconds a pooled claude session may sit unused before it is
terminated. The backend pools live PTYs by conversation-history
fingerprint; when a conversation forks, the session behind the old
fingerprint becomes unreachable but stays resident — one live
``claude`` (hundreds of MB) each. 30 minutes keeps an active chat
warm across a coffee break while bounding what a fork storm can
strand. There is no count cap: concurrent conversations should all
stay warm. ``0`` disables reaping."""
@property
def kinds(self) -> tuple[Kind, ...]:
return tuple(k for k in KINDS if self.for_kind(k) is not None)
class ClaudeAgent(BaseAgent):
"""Agent backed by ``claude-code-api``.
``cwd`` is the working directory the claude CLI is spawned in;
typed as ``Path`` (Pydantic coerces strings) since claude-code-api
accepts ``str | os.PathLike[str]``. ``available_native_tools`` stays
``tuple[str, ...]`` because Claude Code's tool set is a moving target
(Bash/Read/Edit/Write/WebSearch/etc. — versions add and rename), so
pinning it as a ``Literal`` would rot the type each release.
"""
cwd: Path
available_native_tools: tuple[str, ...] = ()
options: ClaudeCodeOptions = Field(default_factory=ClaudeCodeOptions)
"""Per-agent passthrough for the underlying claude-code-api
``BackendOptions``. See :class:`ClaudeCodeOptions`."""
system_prompt: str = ""
"""Verbatim prompt for agents without ``prompts`` (tests, one-offs)."""
prompts: Prompts = Field(default_factory=Prompts)
kinds: tuple[Kind, ...] = ()
"""Conversation kinds this agent serves; ``create``/``spawn`` reject the
rest. Defaults to the kinds ``prompts`` covers, or ``("deep",)`` for a
verbatim ``system_prompt``."""
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)
@model_validator(mode="after")
def _kinds_follow_prompts(self) -> ClaudeAgent:
covered = self.prompts.kinds
if not self.kinds:
object.__setattr__(self, "kinds", covered or ("deep",))
elif covered:
missing = [k for k in self.kinds if k not in covered]
if missing:
msg = f"agent {self.name!r} serves {missing} without a prompt"
raise ValueError(msg)
return self
def prompt_for(self, kind: Kind) -> tuple[PromptSource, ...] | None:
return self.prompts.for_kind(kind)
def serves(self, kind: str) -> bool:
return kind in self.kinds
+1 -1
View File
@@ -1,6 +1,6 @@
"""Backend adapters.
Each backend wraps a provider-specific SDK (``raycast-api``, ``claude-code-api``)
Each backend wraps a provider SDK (``raycast-api``, ``claude-agent-sdk``)
and yields the unified :class:`~beaver_gateway.core.events.MessageStreamEvent`
family. The Anthropic-style frontend serialises events straight to SSE.
"""
+9 -2
View File
@@ -5,8 +5,15 @@ into a stream of :class:`~beaver_gateway.core.events.MessageStreamEvent`
records. The frontend serializes whatever comes out straight to SSE, so
backends are the only place where provider quirks are translated.
Implementations are plain :class:`typing.Protocol` conformers — no ABC
subclassing — to keep them swappable in tests with bare async generators.
Implementations are plain :class:`typing.Protocol` conformers - no ABC
subclassing - to keep them swappable in tests with bare async generators.
``**options`` is the one extension point. Known keys, all optional and
ignored by backends that don't keep state: ``conversation_id`` (stable id
the backend may pin a live session to), ``session_id`` (backend session to
resume when nothing is live), ``capture`` (a
:class:`~beaver_gateway.core.turn_capture.TurnCapture` the backend fills
after the stream closes).
"""
from __future__ import annotations
-502
View File
@@ -1,502 +0,0 @@
"""Claude Code backend adapter.
One :class:`ClaudeCodeBackendAdapter` per :class:`ClaudeAgent`. The
underlying :class:`claude_code_api.ClaudeCodeBackend` bakes ``cwd`` /
``model`` / ``system_prompt`` / MCP wiring into a single
:class:`~claude_code_api.BackendOptions` at construction time, so a
single backend instance is conceptually bound to one agent (different
agents would mean different cwds / system prompts / exposed MCPs and
thus different live-session pools).
Per :meth:`complete` we:
* hand the full Anthropic-style ``messages`` list to
``ClaudeCodeBackend.complete`` — it does its own fingerprint-based
session lookup, so we never need to track sessions ourselves;
* turn its events into one ``message_start`` … ``message_stop``
envelope per ``complete`` call, with content-block indices increasing
monotonically across the whole turn;
* close the envelope on the ``ResultMessage``.
How the blocks inside that envelope are produced depends on the agent's
transport. By default each ``AssistantMessage`` is re-emitted as
``content_block_start`` + one delta + ``content_block_stop`` per block —
a "stream" of finished blocks, because the PTY transport has nothing
finer to offer. With ``transport="stream_json"`` *and*
``include_partial_messages``, the backend instead feeds us the real
token-level ``StreamEvent``s and we emit from those; the whole-block
records still arrive and are still kept for :class:`TurnCapture`, but
emitting from both would duplicate every block on the wire.
The per-request ``system`` parameter is intentionally **ignored** —
``BackendOptions.system_prompt`` is fixed at session-spawn time, and the
agent's ``system_prompt`` is the canonical identity of the agent.
"""
from __future__ import annotations
import json
import logging
import uuid
from collections.abc import Mapping # runtime import: isinstance in _emit_stream_event
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Self
from claude_code_api import (
AssistantMessage,
BackendOptions,
ClaudeCodeBackend,
ResultMessage,
StreamEvent,
TextBlock,
ThinkingBlock,
ToolUseBlock,
synthesize_turn_messages,
)
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.core.events import (
StopReason,
build_content_block_stop,
build_input_json_delta,
build_message_delta,
build_message_start,
build_message_stop,
build_signature_delta,
build_text_block_start,
build_text_delta,
build_thinking_block_start,
build_thinking_delta,
build_tool_use_block_start,
)
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterable
from anthropic.types import MessageParam
from beaver_gateway.agents.base import BaseAgent
from beaver_gateway.core.events import MessageStreamEvent
_log = logging.getLogger("beaver_gateway.backends.claude_code")
__all__ = ["ClaudeCodeBackendAdapter", "TurnCapture"]
@dataclass
class TurnCapture:
"""Side-channel sink for per-turn metadata.
Pass an instance via ``ClaudeCodeBackendAdapter.complete(capture=...)``.
After the stream finishes, :attr:`synthesized_messages` holds the
full assistant↔tool-result cycle (from
:func:`claude_code_api.synthesize_turn_messages`) — i.e. the exact
list of canonical Anthropic-shape messages claude-code-api stashed
the live session under. The markdown frontend uses this to write the
conversation history to its DB so a subsequent turn's prefix
fingerprint hits the same session.
Other backends (anthropic, raycast) ignore the kwarg — it lands in
their ``**options`` and is silently dropped.
"""
synthesized_messages: list[dict[str, Any]] = field(default_factory=list)
_CLAUDE_TO_ANTHROPIC_STOP: dict[str, StopReason] = {
"end_turn": "end_turn",
"tool_use": "tool_use",
"max_tokens": "max_tokens",
"stop_sequence": "stop_sequence",
"refusal": "refusal",
}
def _map_stop_reason(raw: str | None) -> StopReason:
"""Map claude-code's stop reason into Anthropic's vocabulary.
Unknown / missing values collapse to ``end_turn`` so the client sees
a clean finish rather than a wire-format error.
"""
if raw is None:
return "end_turn"
return _CLAUDE_TO_ANTHROPIC_STOP.get(raw, "end_turn")
def _build_mcp_servers(
agent: ClaudeAgent, mcp_internal_urls: Mapping[str, str]
) -> dict[str, dict[str, Any]] | None:
"""Render ``agent.expose_mcps`` into ``BackendOptions.mcp_servers``.
Each exposed MCP is a streamable-HTTP pointer at the gateway's
internal aggregator (built by :mod:`beaver_gateway.mcp.internal_app`).
``None`` keeps claude-code from materializing an ``--mcp-config``
file when the agent exposes nothing.
"""
if not agent.expose_mcps:
return None
servers: dict[str, dict[str, Any]] = {}
for em in agent.expose_mcps:
url = mcp_internal_urls.get(em.name)
if url is None:
msg = (
f"agent {agent.name!r} exposes MCP {em.name!r} "
"but no internal URL is registered for it"
)
raise ValueError(msg)
servers[em.name] = {"type": "http", "url": url}
return servers
def _build_backend_options(
agent: ClaudeAgent, mcp_internal_urls: Mapping[str, str]
) -> BackendOptions:
"""Compose the per-agent :class:`BackendOptions`.
Agent-primary fields:
* ``cwd`` / ``model`` come from the agent directly;
* ``system_prompt`` carries :attr:`BaseAgent.system_prompt`
verbatim — i.e. wire-level ``--system-prompt`` (~8.6k tokens
lighter than ``--append-system-prompt`` because claude-code's
persona/planning conventions and dynamic sections drop out;
tool schemas survive via the API ``tools=[]`` channel);
* ``append_system_prompt`` carries
:attr:`ClaudeCodeOptions.append_system_prompt`, normally
``None``. Setting it re-attaches claude-code's built-in prompt
*and* this delta — opt-in for "claude as a real coding session";
* ``allowed_tools`` follows the PLAN: when the user lists native
tools we restrict to those *plus* a per-MCP wildcard so MCP tools
stay reachable; when no native list is declared we leave
``allowed_tools`` empty (= all tools allowed by claude-code's
default);
* ``mcp_servers`` comes from :func:`_build_mcp_servers`.
Every other tunable knob is passed through from
:attr:`ClaudeAgent.options`. Our default overrides
(``wait_for_turn_duration=True``,
``dangerously_skip_permissions=True``) live on
:class:`ClaudeCodeOptions`, not here, so a user who builds
``ClaudeCodeOptions(...)`` explicitly inherits the same defaults
instead of getting whatever claude-code-api ships.
"""
allowed_tools: tuple[str, ...] = ()
if agent.available_native_tools:
mcp_wildcards = tuple(f"mcp__{em.name}" for em in agent.expose_mcps)
allowed_tools = tuple(agent.available_native_tools) + mcp_wildcards
opt = agent.options
return BackendOptions(
cwd=agent.cwd,
model=agent.model or None,
system_prompt=agent.system_prompt,
append_system_prompt=opt.append_system_prompt,
allowed_tools=allowed_tools,
mcp_servers=_build_mcp_servers(agent, mcp_internal_urls),
transport=opt.transport,
include_partial_messages=opt.include_partial_messages,
# ``None`` means "decide from the agent": an agent that exposes
# MCPs needs them on its first reply, one that doesn't shouldn't
# pay for a warm-up it gains nothing from.
warmup_turn=(
bool(agent.expose_mcps) if opt.warmup_turn is None else opt.warmup_turn
),
disallowed_tools=opt.disallowed_tools,
permission_mode=opt.permission_mode,
dangerously_skip_permissions=opt.dangerously_skip_permissions,
effort=opt.effort,
add_dir=opt.add_dir,
settings=opt.settings,
extra_args=opt.extra_args,
extra_env=opt.extra_env,
preserve_provider_env=opt.preserve_provider_env,
history_injection_mode=opt.history_injection_mode,
wait_for_turn_duration=opt.wait_for_turn_duration,
include_meta_user=opt.include_meta_user,
startup_delay=opt.startup_delay,
file_wait_timeout=opt.file_wait_timeout,
turn_duration_timeout=opt.turn_duration_timeout,
idle_session_ttl=opt.idle_session_ttl,
)
class ClaudeCodeBackendAdapter:
"""One ``claude-code-api`` backend bound to a single :class:`ClaudeAgent`.
Owns the underlying :class:`ClaudeCodeBackend`'s lifecycle through
the async-context-manager protocol so :mod:`beaver_gateway.cli` can
park it in its ``AsyncExitStack``.
"""
def __init__(
self, *, agent: ClaudeAgent, mcp_internal_urls: Mapping[str, str]
) -> None:
self._agent = agent
options = _build_backend_options(agent, mcp_internal_urls)
self._backend = ClaudeCodeBackend(options)
self._streaming = (
options.transport == "stream_json" and options.include_partial_messages
)
"""Whether the backend feeds us token-level ``StreamEvent``s.
When it does, blocks are emitted from those and the whole-block
``AssistantMessage`` records are used only for bookkeeping —
emitting from both would duplicate every block on the wire."""
@property
def agent(self) -> ClaudeAgent:
return self._agent
@property
def live_session_count(self) -> int:
return self._backend.live_session_count
@property
def live_sessions(self) -> dict[str, Any]:
"""Live claude processes keyed by claude session_id.
Pass-through to the underlying ``ClaudeCodeBackend``. The value
is a ``PtyClaudeProcess`` or a ``StreamClaudeProcess`` depending
on the agent's transport; both implement the
``captured_output()`` / ``add_output_listener`` / ``write``
surface the admin terminal consumes, so it works either way —
with the caveat that the stream transport's "terminal" is a
read-only view of the JSON event stream rather than a TUI you
can type into. Typed as ``Any`` to avoid leaking the lower
layer's type into the gateway's public surface.
"""
return self._backend.live_sessions
async def __aenter__(self) -> Self:
await self._backend.__aenter__()
return self
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
await self._backend.__aexit__(exc_type, exc, tb)
async def aclose(self) -> None:
await self._backend.aclose()
async def complete(
self,
*,
agent: BaseAgent,
messages: Iterable[MessageParam],
system: str | None = None, # noqa: ARG002 — see module docstring
capture: TurnCapture | None = None,
**options: Any, # noqa: ARG002 — no per-request knobs for claude-code yet
) -> AsyncIterator[MessageStreamEvent]:
if not isinstance(agent, ClaudeAgent):
msg = (
"ClaudeCodeBackendAdapter requires ClaudeAgent, "
f"got {type(agent).__name__}"
)
raise TypeError(msg)
if agent.name != self._agent.name:
# Adapter is per-agent; routing a different agent through it
# would mean a different cwd / system_prompt / MCP set than
# the live-session pool was spawned with.
msg = (
f"ClaudeCodeBackendAdapter bound to {self._agent.name!r} "
f"got request for {agent.name!r}"
)
raise ValueError(msg)
msgs_list: list[Mapping[str, Any]] = list(messages)
_log.info(
"complete: agent=%s n_messages=%d capture=%s live_sessions=%d",
agent.name,
len(msgs_list),
capture is not None,
self._backend.live_session_count,
)
message_id = f"msg_{uuid.uuid4().hex}"
yield build_message_start(message_id=message_id, model=agent.model)
next_index = 0
stop_reason: str | None = None
usage: Mapping[str, Any] | None = None
n_text = 0
n_thinking = 0
n_tool_use = 0
# We keep raw events so we can hand them to
# ``synthesize_turn_messages`` after the stream closes — the
# markdown frontend stores the result in its conversation
# history so the next turn's prefix matches the backend's
# session-pool fingerprint. UserMessage (tool_result) events
# are silently discarded from the wire but kept here.
raw_events: list[Any] = []
async for event in self._backend.complete(msgs_list):
raw_events.append(event)
if isinstance(event, StreamEvent):
# Only ever arrives when ``self._streaming`` — the
# backend doesn't emit these otherwise. Indices are
# already rebased onto the whole turn by claude-code-api.
for ev in _emit_stream_event(event.event):
yield ev
elif isinstance(event, AssistantMessage):
for block in event.content:
if isinstance(block, TextBlock):
n_text += 1
elif isinstance(block, ThinkingBlock):
n_thinking += 1
elif isinstance(block, ToolUseBlock):
n_tool_use += 1
if not self._streaming:
for ev in _emit_block(block, next_index):
yield ev
next_index += 1
elif isinstance(event, ResultMessage):
# ResultMessage is the terminal event from TurnManager
# — we capture its stop_reason / usage for the envelope
# below. We DO NOT break here: an early break would
# raise GeneratorExit inside claude-code-api's
# ``complete`` coroutine before it gets a chance to
# stash the live session under the post-turn
# fingerprint, so every continuation would miss the
# cache and reseed. Let the inner generator exit
# naturally instead.
stop_reason = event.stop_reason
usage = event.usage
# UserMessage (tool_result records) and SystemMessage
# (turn_duration heartbeats) carry no content for the
# /v1/messages caller — skip silently on the wire, but they
# ARE retained in ``raw_events`` for synthesis below.
if capture is not None:
capture.synthesized_messages = synthesize_turn_messages(raw_events)
_log.info(
"complete: agent=%s DONE text=%d thinking=%d tool_use=%d stop=%s synth=%d",
agent.name,
n_text,
n_thinking,
n_tool_use,
stop_reason,
len(capture.synthesized_messages) if capture is not None else 0,
)
yield build_message_delta(
stop_reason=_map_stop_reason(stop_reason), usage=_normalize_usage(usage)
)
yield build_message_stop()
def _emit_stream_event(event: Mapping[str, Any]) -> Iterable[MessageStreamEvent]:
"""Render one raw Anthropic streaming event through our own builders.
The payload is already Anthropic-shaped, so forwarding it verbatim
is tempting — but it arrives as a plain dict and the frontend
serializes pydantic models, so it would have to be validated against
the SDK's ``RawMessageStreamEvent`` union anyway. Rebuilding it from
the fields we recognize is the same amount of work and fails the
same way :func:`_emit_block` already does: a block or delta type
this gateway doesn't know is dropped rather than raising mid-stream
on a claude release that added one.
Only ``content_block_*`` events reach here — claude-code-api strips
the inner per-request message envelopes, because the caller owns the
single envelope spanning the whole turn.
"""
etype = event.get("type")
index = event.get("index")
if not isinstance(index, int):
return ()
if etype == "content_block_start":
block = event.get("content_block")
if not isinstance(block, Mapping):
return ()
btype = block.get("type")
if btype == "text":
return (build_text_block_start(index),)
if btype == "thinking":
return (build_thinking_block_start(index),)
if btype == "tool_use":
return (
build_tool_use_block_start(
index,
tool_use_id=str(block.get("id", "")),
name=str(block.get("name", "")),
),
)
return ()
if etype == "content_block_delta":
delta = event.get("delta")
if not isinstance(delta, Mapping):
return ()
dtype = delta.get("type")
if dtype == "text_delta":
return (build_text_delta(index, str(delta.get("text", ""))),)
if dtype == "thinking_delta":
return (build_thinking_delta(index, str(delta.get("thinking", ""))),)
if dtype == "signature_delta":
return (build_signature_delta(index, str(delta.get("signature", ""))),)
if dtype == "input_json_delta":
return (build_input_json_delta(index, str(delta.get("partial_json", ""))),)
return ()
if etype == "content_block_stop":
return (build_content_block_stop(index),)
return ()
def _emit_block(
block: TextBlock | ThinkingBlock | ToolUseBlock | Any, index: int
) -> Iterable[MessageStreamEvent]:
"""Render one ``claude-code`` content block as Anthropic stream events.
``ToolResultBlock`` would arrive only on user-role records — we
don't emit it here because :meth:`complete` skips ``UserMessage``.
"""
if isinstance(block, TextBlock):
return (
build_text_block_start(index),
build_text_delta(index, block.text),
build_content_block_stop(index),
)
if isinstance(block, ThinkingBlock):
return (
build_thinking_block_start(index),
build_thinking_delta(index, block.thinking),
build_signature_delta(index, block.signature),
build_content_block_stop(index),
)
if isinstance(block, ToolUseBlock):
partial = json.dumps(block.input, separators=(",", ":"), ensure_ascii=False)
return (
build_tool_use_block_start(index, tool_use_id=block.id, name=block.name),
build_input_json_delta(index, partial),
build_content_block_stop(index),
)
return ()
def _normalize_usage(usage: Mapping[str, Any] | None) -> dict[str, int] | None:
"""Coerce claude-code's ``usage`` dict to Anthropic ``MessageDeltaUsage`` shape.
claude-code copies whatever the JSONL ``usage`` record carried —
fields can be missing, strings, or ints. We pass through only the
fields ``MessageDeltaUsage`` knows about and discard the rest so an
odd ``cache_creation`` object structure doesn't fail pydantic
validation downstream.
"""
if not usage:
return None
out: dict[str, int] = {}
for key in (
"input_tokens",
"output_tokens",
"cache_creation_input_tokens",
"cache_read_input_tokens",
):
value = usage.get(key)
if isinstance(value, int):
out[key] = value
return out or None
+855
View File
@@ -0,0 +1,855 @@
"""Claude Agent SDK backend adapter.
One :class:`ClaudeSdkBackend` per :class:`ClaudeAgent`. A live session is
one ``ClaudeSDKClient`` (one claude subprocess) and runs one turn at a
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 (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
rebased across the API calls claude makes inside the turn.
Process isolation: claude is spawned through a small exec wrapper that
drops every inherited environment variable outside a whitelist and, when
``RunnerConfig.user`` is set, switches to that uid before exec (done in the
wrapper rather than via ``subprocess(user=...)``, which uvloop rejects).
"""
from __future__ import annotations
import asyncio
import contextlib
import fnmatch
import hashlib
import json
import logging
import os
import pwd
import shutil
import sys
import tempfile
import time
import uuid
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Self, cast
import claude_agent_sdk
from claude_agent_sdk import (
AssistantMessage,
ClaudeAgentOptions,
ClaudeSDKClient,
MirrorErrorMessage,
ResultMessage,
StreamEvent,
TextBlock,
ThinkingBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
project_key_for_directory,
)
from beaver_gateway.core import prompt as prompt_assembly
from beaver_gateway.core.events import (
StopReason,
build_content_block_stop,
build_input_json_delta,
build_message_delta,
build_message_start,
build_message_stop,
build_signature_delta,
build_text_block_start,
build_text_delta,
build_thinking_block_start,
build_thinking_delta,
build_tool_use_block_start,
)
from beaver_gateway.core.kinds import as_kind
from beaver_gateway.core.sessions import Session, SessionClient, SessionPool
from beaver_gateway.core.transcript import (
build_entries,
close_open_tool_uses,
fingerprint,
text_of,
)
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 McpSdkServerConfig, SessionStore
from beaver_gateway.agents.base import BaseAgent
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.core.events import MessageStreamEvent
_log = logging.getLogger("beaver_gateway.backends.claude_sdk")
__all__ = [
"ClaudeSdkBackend",
"RunnerConfig",
"SessionClient",
"ToolServerFactory",
"UsageSink",
"fingerprint",
]
ENV_KEEP: tuple[str, ...] = (
"PATH",
"HOME",
"LANG",
"LANGUAGE",
"LC_ALL",
"LC_CTYPE",
"TZ",
"TERM",
"USER",
"LOGNAME",
"SHELL",
"TMPDIR",
"PWD",
"SSL_CERT_FILE",
"SSL_CERT_DIR",
"NODE_EXTRA_CA_CERTS",
"NODE_OPTIONS",
"IS_SANDBOX",
"HTTP_PROXY",
"HTTPS_PROXY",
"NO_PROXY",
"http_proxy",
"https_proxy",
"no_proxy",
)
ENV_KEEP_PREFIXES: tuple[str, ...] = ("CLAUDE_", "ANTHROPIC_", "DISABLE_")
_STOP_REASONS: dict[str, StopReason] = {
"end_turn": "end_turn",
"tool_use": "tool_use",
"max_tokens": "max_tokens",
"stop_sequence": "stop_sequence",
"refusal": "refusal",
"pause_turn": "pause_turn",
}
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)
class RunnerConfig:
user: str | None = None
home: Path | None = None
@dataclass(frozen=True, slots=True)
class UsageEvent:
agent_name: str
model: str
effort: str | None
conversation_id: str | None
session_id: str | None
usage: TurnUsage
class _RunnerClient(ClaudeSDKClient):
"""``ClaudeSDKClient`` that hands the materialized resume dir to the runner uid."""
def __init__(self, options: ClaudeAgentOptions, *, uid: int | None) -> None:
super().__init__(options=options)
self._runner_uid = uid
async def _connect_inner(self, prompt: Any, actual_prompt: Any) -> None:
materialized = self._materialized
if materialized is not None and self._runner_uid is not None:
_chown_tree(materialized.config_dir, self._runner_uid)
await super()._connect_inner(prompt, actual_prompt)
class ClaudeSdkBackend:
def __init__(
self,
*,
agent: ClaudeAgent,
mcp_internal_urls: Mapping[str, str],
session_store: SessionStore,
mcp_tool_names: Mapping[str, Sequence[str]] | None = None,
runner: RunnerConfig | None = None,
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
self._runner = runner or RunnerConfig()
self._usage_sink = usage_sink
self._factory = client_factory or self._default_factory
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._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
@property
def agent(self) -> ClaudeAgent:
return self._agent
@property
def pool(self) -> SessionPool:
return self._pool
@property
def sessions(self) -> dict[str, dict[str, Any]]:
return {
row["key"]: row
for row in self._pool.snapshot()
if row["agent"] == self._agent.name
}
async def __aenter__(self) -> Self:
return self
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
await self.aclose()
async def aclose(self) -> None:
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,
*,
agent: BaseAgent,
messages: Iterable[MessageParam],
system: str | None = None, # noqa: ARG002 - the agent owns its prompt
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:
msg = f"backend bound to {self._agent.name!r}, got {agent.name!r}"
raise ValueError(msg)
history = [dict(m) for m in messages]
if not history or history[-1].get("role") != "user":
msg = "the last message must be a user turn"
raise ValueError(msg)
prompt = _prompt_text(history[-1].get("content"))
prior = history[:-1]
key = conversation_id or fingerprint(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, observer=observer)
except Exception:
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(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(
agent_name=self._agent.name,
model=self._agent.model,
effort=self._agent.options.effort,
conversation_id=conversation_id,
session_id=live.session_id,
usage=usage,
)
)
return usage, interrupted
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] = []
next_index = 0
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):
event = message.event
if event.get("type") == "message_start":
offset = next_index
continue
index = event.get("index")
if isinstance(index, int):
next_index = max(next_index, offset + index + 1)
if streaming:
turn.events.extend(_emit_stream_event(event, offset + index))
elif isinstance(message, AssistantMessage):
raw.append(message)
if not streaming:
for block in message.content:
turn.events.extend(_emit_block(block, next_index))
next_index += 1
elif isinstance(message, UserMessage):
raw.append(message)
elif isinstance(message, ResultMessage):
turn.result = message
turn.stop_reason = _STOP_REASONS.get(
message.stop_reason or "", "end_turn"
)
turn.synthesized = synthesize_turn_messages(raw)
_log.info(
"turn: agent=%s session=%s events=%d synthesized=%d stop=%s",
self._agent.name,
live.session_id,
len(turn.events),
len(turn.synthesized),
turn.stop_reason,
)
return turn
async def _acquire(
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 not None:
await self.repair_session(resume)
elif history:
resume = await self._seed(history)
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())
entries = build_entries(
history,
session_id=session_id,
cwd=str(self._agent.cwd),
model=self._agent.model,
permission_mode=self._agent.options.permission_mode,
)
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,
len(entries),
len(history),
)
return session_id
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 kind=%s resume=%s tools=%s user=%s",
self._agent.name,
spec.kind,
resume,
spec.tools,
self._runner.user,
)
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, *, key: str, spec: _SessionSpec
) -> ClaudeAgentOptions:
agent = self._agent
opt = agent.options
env = dict(opt.env)
if self._runner.home is not None:
env["HOME"] = str(self._runner.home)
env.setdefault("CLAUDE_CONFIG_DIR", str(self._runner.home / ".claude"))
plugins = self._plugins()
sources = agent.prompt_for(as_kind(spec.kind))
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", 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],
cwd=str(agent.cwd),
add_dirs=list(opt.add_dirs),
env=env,
cli_path=str(self._exec_wrapper(extra_keep=tuple(env))),
include_partial_messages=opt.include_partial_messages,
session_store=self._store,
session_store_flush=cast("Any", opt.session_store_flush),
resume=resume,
plugins=cast("Any", plugins),
skills="all" if plugins else None,
max_turns=opt.max_turns,
stderr=lambda line: _log.warning(
"claude[%s]: %s", agent.name, line.rstrip()
),
)
def _plugins(self) -> list[dict[str, str]]:
plugins: list[dict[str, str]] = []
root = self._work_dir / "plugins" / self._agent.name
for raw in sorted(self._agent.skill_sets, key=str):
source = Path(str(raw))
name = source.name
target = root / name
if target.exists():
shutil.rmtree(target)
shutil.copytree(source, target / "skills")
(target / ".claude-plugin").mkdir(parents=True, exist_ok=True)
(target / ".claude-plugin" / "plugin.json").write_text(
json.dumps({"name": name, "version": "0.0.0"}), encoding="utf-8"
)
_chmod_tree(target)
plugins.append({"type": "local", "path": str(target)})
return plugins
def _exec_wrapper(self, *, extra_keep: tuple[str, ...]) -> Path:
if self._wrapper is not None:
return self._wrapper
keep = sorted({*ENV_KEEP, *self._agent.options.env_keep, *extra_keep})
target = _claude_binary()
script = _WRAPPER.format(
python=sys.executable,
target=json.dumps(target),
keep=json.dumps(keep),
prefixes=json.dumps(list(ENV_KEEP_PREFIXES)),
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"
path.parent.mkdir(parents=True, exist_ok=True)
if not path.exists() or path.read_text(encoding="utf-8") != script:
path.write_text(script, encoding="utf-8")
path.chmod(0o755)
path.parent.chmod(0o755)
self._wrapper = path
return path
def _rekey(self, old: str, new: str) -> None:
stale = self._pool.rekey(old, new)
if stale is not None:
asyncio.get_running_loop().create_task(_disconnect(stale))
@dataclass(frozen=True, slots=True)
class _SessionSpec:
kind: str
pinned: bool
tools: bool
@dataclass
class _Turn:
events: list[Any] = field(default_factory=list)
synthesized: list[dict[str, Any]] = field(default_factory=list)
result: ResultMessage | None = None
stop_reason: StopReason = "end_turn"
_WRAPPER = """#!{python}
import os
import sys
TARGET = {target}
KEEP = set({keep})
PREFIXES = tuple({prefixes})
UID = {uid}
GID = {gid}
env = {{
k: v for k, v in os.environ.items() if k in KEEP or k.startswith(PREFIXES)
}}
if UID is not None and os.getuid() != UID:
os.setgroups([])
os.setgid(GID)
os.setuid(UID)
os.execve(TARGET, [TARGET, *sys.argv[1:]], env)
"""
async def _disconnect(live: Session) -> None:
try:
await live.client.disconnect()
except Exception: # noqa: BLE001
_log.exception("disconnect failed for session %s", live.session_id)
def _claude_binary() -> str:
bundled = Path(claude_agent_sdk.__file__).parent / "_bundled" / "claude"
if bundled.is_file():
return str(bundled)
found = shutil.which("claude")
if found is None:
msg = "claude CLI not found: neither bundled in claude_agent_sdk nor on PATH"
raise FileNotFoundError(msg)
return found
def _resolve_ids(user: str | None) -> tuple[int | None, int | None]:
if user is None:
return None, None
record = pwd.getpwuid(int(user)) if user.isdigit() else pwd.getpwnam(user)
return record.pw_uid, record.pw_gid
def _chown_tree(root: Path, uid: int) -> None:
for path in [root, *root.rglob("*")]:
with contextlib.suppress(OSError):
os.chown(path, uid, -1)
def _chmod_tree(root: Path) -> None:
for path in [root, *root.rglob("*")]:
with contextlib.suppress(OSError):
path.chmod(0o755 if path.is_dir() else 0o644)
def _mcp_servers(
agent: ClaudeAgent, urls: Mapping[str, str]
) -> dict[str, dict[str, Any]]:
servers: dict[str, dict[str, Any]] = {}
for exposed in agent.expose_mcps:
url = urls.get(exposed.name)
if url is None:
msg = f"agent {agent.name!r} exposes MCP {exposed.name!r} without a URL"
raise ValueError(msg)
servers[exposed.name] = {"type": "http", "url": url}
return servers
def _mcp_disallowed(
agent: ClaudeAgent, catalog: Mapping[str, Sequence[str]]
) -> list[str]:
out: list[str] = []
for exposed in agent.expose_mcps:
if exposed.tools is None and not exposed.deny:
continue
names = list(catalog.get(exposed.name, ()))
if not names:
_log.warning(
"MCP %r has no tool catalog; disallowing the whole server for %s",
exposed.name,
agent.name,
)
out.append(f"mcp__{exposed.name}")
continue
for name in names:
allowed = exposed.tools is None or name in exposed.tools
denied = any(fnmatch.fnmatchcase(name, pat) for pat in exposed.deny)
if not allowed or denied:
out.append(f"mcp__{exposed.name}__{name}")
return out
def _prompt_text(content: Any) -> str:
text = text_of(content)
if not text:
msg = "user message has no text content"
raise ValueError(msg)
return text
def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
out: list[dict[str, Any]] = []
for message in raw:
if isinstance(message, AssistantMessage):
out.append(
{
"role": "assistant",
"content": [_block_to_dict(b) for b in message.content],
}
)
elif isinstance(message, UserMessage):
content = message.content
if isinstance(content, list) and content:
out.append(
{"role": "user", "content": [_block_to_dict(b) for b in content]}
)
return out
def _block_to_dict(block: Any) -> dict[str, Any]:
if isinstance(block, TextBlock):
return {"type": "text", "text": block.text}
if isinstance(block, ToolUseBlock):
return {
"type": "tool_use",
"id": block.id,
"name": block.name,
"input": block.input,
}
if isinstance(block, ToolResultBlock):
result: dict[str, Any] = {
"type": "tool_result",
"tool_use_id": block.tool_use_id,
"content": block.content,
}
if block.is_error is not None:
result["is_error"] = block.is_error
return result
if isinstance(block, ThinkingBlock):
return {
"type": "thinking",
"thinking": block.thinking,
"signature": block.signature,
}
msg = f"unknown content block type: {type(block).__name__}"
raise TypeError(msg)
def _usage_of(result: ResultMessage | None) -> TurnUsage:
if result is None:
return TurnUsage()
usage = result.usage or {}
return TurnUsage(
input_tokens=_int(usage.get("input_tokens")),
output_tokens=_int(usage.get("output_tokens")),
cache_read_tokens=_int(usage.get("cache_read_input_tokens")),
cache_creation_tokens=_int(usage.get("cache_creation_input_tokens")),
cost_usd=result.total_cost_usd,
duration_ms=result.duration_ms,
num_turns=result.num_turns,
)
def _wire_usage(usage: TurnUsage) -> dict[str, int]:
return {
"input_tokens": usage.input_tokens,
"output_tokens": usage.output_tokens,
"cache_read_input_tokens": usage.cache_read_tokens,
"cache_creation_input_tokens": usage.cache_creation_tokens,
}
def _int(value: Any) -> int:
return value if isinstance(value, int) else 0
def _emit_stream_event(
event: Mapping[str, Any], index: int
) -> Iterable[MessageStreamEvent]:
etype = event.get("type")
if etype == "content_block_start":
block = event.get("content_block")
if not isinstance(block, Mapping):
return ()
btype = block.get("type")
if btype == "text":
return (build_text_block_start(index),)
if btype == "thinking":
return (build_thinking_block_start(index),)
if btype == "tool_use":
return (
build_tool_use_block_start(
index,
tool_use_id=str(block.get("id", "")),
name=str(block.get("name", "")),
),
)
return ()
if etype == "content_block_delta":
delta = event.get("delta")
if not isinstance(delta, Mapping):
return ()
dtype = delta.get("type")
if dtype == "text_delta":
return (build_text_delta(index, str(delta.get("text", ""))),)
if dtype == "thinking_delta":
return (build_thinking_delta(index, str(delta.get("thinking", ""))),)
if dtype == "signature_delta":
return (build_signature_delta(index, str(delta.get("signature", ""))),)
if dtype == "input_json_delta":
return (build_input_json_delta(index, str(delta.get("partial_json", ""))),)
return ()
if etype == "content_block_stop":
return (build_content_block_stop(index),)
return ()
def _emit_block(block: Any, index: int) -> Iterable[MessageStreamEvent]:
if isinstance(block, TextBlock):
return (
build_text_block_start(index),
build_text_delta(index, block.text),
build_content_block_stop(index),
)
if isinstance(block, ThinkingBlock):
return (
build_thinking_block_start(index),
build_thinking_delta(index, block.thinking),
build_signature_delta(index, block.signature),
build_content_block_stop(index),
)
if isinstance(block, ToolUseBlock):
partial = json.dumps(block.input, separators=(",", ":"), ensure_ascii=False)
return (
build_tool_use_block_start(index, tool_use_id=block.id, name=block.name),
build_input_json_delta(index, partial),
build_content_block_stop(index),
)
return ()
+4 -1
View File
@@ -26,11 +26,12 @@ tools we route it back to the underlying MCP in-process, append the
result as a ``tool`` message, and re-issue the stream — all inside one
Anthropic envelope (one ``message_start`` … one ``message_stop``). The
tool_use blocks DO surface to the caller (mirrors what
``ClaudeCodeBackendAdapter`` does), but tool_results stay internal.
``ClaudeSdkBackend`` does), but tool_results stay internal.
"""
from __future__ import annotations
import fnmatch
import json
import logging
import uuid
@@ -297,6 +298,8 @@ def _build_agent_catalog(
for mt in mcp_tools.get(em.name, []):
if em.tools is not None and mt.name not in em.tools:
continue
if any(fnmatch.fnmatchcase(mt.name, pat) for pat in em.deny):
continue
wire_name = f"{em.name}__{mt.name}"
routing[wire_name] = (em.name, mt.name)
local_tools.append(
+96 -13
View File
@@ -22,6 +22,7 @@ from __future__ import annotations
import asyncio
import contextlib
import functools
import logging
import signal
from contextlib import AsyncExitStack
@@ -36,16 +37,25 @@ from raycast_api.config import Config as RaycastConfig
from beaver_gateway import config_loader
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.agents.raycast import RaycastAgent
from beaver_gateway.backends.claude_code import ClaudeCodeBackendAdapter
from beaver_gateway.backends.claude_sdk import (
ClaudeSdkBackend,
RunnerConfig,
UsageEvent,
)
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
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
@@ -122,8 +132,8 @@ async def _async_main() -> None:
await token_store.start()
stack.push_async_callback(token_store.stop)
# Internal MCP URLs must exist before we construct any
# ClaudeCodeBackendAdapter — adapters bake the URLs into their
# ``BackendOptions.mcp_servers`` at construction time. The
# ClaudeSdkBackend - adapters bake the URLs into their
# ``mcp_servers`` at construction time. The
# ``mcp_servers`` map is used by the Raycast backend, which
# needs in-process ``list_tools`` / ``call_tool`` access (the
# Raycast wire has no native MCP concept).
@@ -136,14 +146,33 @@ 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,
frontends=gateway.frontends,
)
late.conversations = conversations
runtime = GatewayRuntime(
agents=agents,
@@ -156,6 +185,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:
@@ -181,7 +213,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:
@@ -247,27 +282,42 @@ 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.
The Raycast ``Client`` is shared across every ``RaycastAgent``
(bearer + device-id are process-wide), so we open it lazily — only
when at least one ``RaycastAgent`` is present — and close it via
(bearer + device-id are process-wide), so we open it lazily - only
when at least one ``RaycastAgent`` is present - and close it via
the caller's exit stack.
Each :class:`ClaudeAgent` gets its own
:class:`ClaudeCodeBackendAdapter`: ``BackendOptions`` pins
``cwd`` / ``model`` / ``system_prompt`` / ``mcp_servers`` for the
lifetime of the underlying ``ClaudeCodeBackend``, so different
agents can't share one.
Each :class:`ClaudeAgent` gets its own :class:`ClaudeSdkBackend`
(own live-session pool, own prompt and MCP set); all of them share
the session store and the usage sink.
"""
backends: dict[str, Backend] = {}
@@ -281,10 +331,43 @@ async def _build_backends(
for a in raycast_agents:
backends[a.name] = raycast_backend
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()
}
async def record_usage(event: UsageEvent) -> None:
row = Usage(
agent_name=event.agent_name,
conversation_id=event.conversation_id,
session_id=event.session_id,
model=event.model,
effort=event.effort,
input_tokens=event.usage.input_tokens,
output_tokens=event.usage.output_tokens,
cache_read_tokens=event.usage.cache_read_tokens,
cache_creation_tokens=event.usage.cache_creation_tokens,
cost_usd=event.usage.cost_usd,
duration_ms=event.usage.duration_ms,
num_turns=event.usage.num_turns,
)
try:
async with db.session() as session:
await append_usage(session, row)
except Exception: # noqa: BLE001
_log.exception("usage write failed for %s", event.agent_name)
for a in agents:
if isinstance(a, ClaudeAgent):
adapter = ClaudeCodeBackendAdapter(
agent=a, mcp_internal_urls=mcp_internal_urls
adapter = ClaudeSdkBackend(
agent=a,
mcp_internal_urls=mcp_internal_urls,
session_store=session_store,
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)
+16 -2
View File
@@ -63,6 +63,7 @@ __all__ = [
"load_messages",
"mint_conversation",
"rewrite_messages",
"set_session_id",
]
@@ -121,6 +122,17 @@ async def mint_conversation(
return row
async def set_session_id(
session: AsyncSession, *, conversation_id: int, session_id: str | None
) -> None:
conv = await session.get(Conversation, conversation_id)
if conv is None or conv.session_id == session_id:
return
conv.session_id = session_id
session.add(conv)
await session.commit()
async def load_messages(
session: AsyncSession, *, conversation_id: int
) -> list[dict[str, Any]]:
@@ -484,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
+177
View File
@@ -0,0 +1,177 @@
"""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
from beaver_gateway.core.kinds import as_kind
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`). A branch "
"keeps your agent, other kinds get their frontend's default unless "
"`agent` says otherwise. Returns the id.",
{
"type": "object",
"properties": {
"kind": {"type": "string", "enum": ["branch", "deep", "job"]},
"agent": {"type": "string", "description": "agent name"},
"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()
try:
child = await conversations.spawn(
kind=as_kind(str(args["kind"])),
agent=args.get("agent"),
seed=str(args.get("seed") or "clean"),
parent=parent,
text=args.get("text"),
title=args.get("title"),
window=args.get("window"),
origin="mcp",
)
except (ValueError, LookupError) as exc:
return _error(str(exc))
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()
+18
View File
@@ -0,0 +1,18 @@
"""Conversation kinds (§3.1) as one closed type for agents, frontends, service."""
from __future__ import annotations
from typing import Literal, get_args
__all__ = ["KINDS", "Kind", "as_kind"]
Kind = Literal["master", "branch", "deep", "job", "fork"]
KINDS: tuple[Kind, ...] = get_args(Kind)
def as_kind(value: str) -> Kind:
for kind in KINDS:
if value == kind:
return kind
msg = f"unknown conversation kind {value!r}"
raise ValueError(msg)
+42
View File
@@ -0,0 +1,42 @@
"""System prompt assembly from a list of source files.
The gateway holds no prompt text: an agent names its granules (paths from
``config.py``) and :func:`assemble` concatenates them in that order, so the
result is byte-for-byte identical for every session of the same agent as
long as the files are. A source is a path, or a ``(tag, path)`` pair whose
content is wrapped in ``<tag>...</tag>`` - the markup lives here, the vault
keeps plain markdown. Each granule's hash is logged at assembly so a
drifted prompt can be traced to the file that changed.
"""
from __future__ import annotations
import hashlib
import logging
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Iterable
__all__ = ["PromptSource", "assemble"]
_log = logging.getLogger("beaver_gateway.core.prompt")
PromptSource = str | Path | tuple[str, str | Path]
def assemble(sources: Iterable[PromptSource]) -> str:
parts: list[str] = []
for source in sources:
tag, raw = source if isinstance(source, tuple) else (None, source)
path = Path(raw)
text = path.read_text(encoding="utf-8").strip()
digest = hashlib.sha256(text.encode("utf-8")).hexdigest()[:12]
_log.info(
"prompt granule %s tag=%s sha=%s bytes=%d", path, tag, digest, len(text)
)
if not text:
continue
parts.append(f"<{tag}>\n{text}\n</{tag}>" if tag else text)
return "\n\n".join(parts) + "\n"
+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
+543
View File
@@ -0,0 +1,543 @@
"""Anthropic messages <-> Agent SDK transcript entries.
:func:`build_entries` renders a message list as the entries the Claude CLI
itself writes (reference: ``t/spike_sdk/entries_reference.json``, CLI
2.1.248 via ``import_session_to_store``): one ``user`` entry per prompt,
one ``assistant`` entry per content block sharing a message id, one
``user`` entry per ``tool_result`` parented on the matching ``tool_use``
entry. Only ``user``/``assistant`` entries are produced - no attachments,
titles or queue markers. Appending the result to a session store and
resuming that session id seeds an external history into the SDK.
:func:`messages_from_entries` is the projection back, used by tests and by
anything that needs Anthropic-shape history out of a mirrored transcript.
"""
from __future__ import annotations
import hashlib
import uuid as _uuid
from collections.abc import Mapping
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
try:
from claude_agent_sdk._cli_version import __cli_version__ as _cli_version
except ImportError: # pragma: no cover
_cli_version = "2.1.248"
if TYPE_CHECKING:
from collections.abc import Iterable
__all__ = [
"CLI_VERSION",
"build_entries",
"close_open_tool_uses",
"fingerprint",
"messages_from_entries",
"open_tool_uses",
"prompt_count",
"render_messages",
"strip_tool_entries",
"text_of",
"window_entries",
]
CLI_VERSION = _cli_version
_ENTRYPOINT = "sdk-py"
_USER_TYPE = "external"
_PROMPT_SOURCE = "sdk"
_GIT_BRANCH = "HEAD"
def build_entries(
messages: Iterable[Mapping[str, Any]],
*,
session_id: str,
cwd: str,
model: str,
permission_mode: str = "bypassPermissions",
now: datetime | None = None,
) -> list[dict[str, Any]]:
stamp = (now or datetime.now(UTC)).strftime("%Y-%m-%dT%H:%M:%S.") + (
f"{(now or datetime.now(UTC)).microsecond // 1000:03d}Z"
)
common = {
"isSidechain": False,
"userType": _USER_TYPE,
"entrypoint": _ENTRYPOINT,
"cwd": cwd,
"sessionId": session_id,
"version": CLI_VERSION,
"gitBranch": _GIT_BRANCH,
}
entries: list[dict[str, Any]] = []
parent: str | None = None
prompt_id: str | None = None
tool_use_owner: dict[str, str] = {}
for message in messages:
role = message.get("role")
content = message.get("content")
if role == "user":
results = _tool_results(content)
if results:
for block in results:
owner = tool_use_owner.get(
str(block.get("tool_use_id", "")), parent
)
uid = _new_uuid()
entries.append(
{
"parentUuid": owner,
"promptId": prompt_id,
"type": "user",
"message": {"role": "user", "content": [block]},
"uuid": uid,
"timestamp": stamp,
"toolUseResult": _result_text(block),
"sourceToolAssistantUUID": owner,
**common,
}
)
parent = uid
continue
prompt_id = _new_uuid()
uid = _new_uuid()
entries.append(
{
"parentUuid": parent,
"promptId": prompt_id,
"type": "user",
"message": {"role": "user", "content": _user_content(content)},
"uuid": uid,
"timestamp": stamp,
"permissionMode": permission_mode,
"promptSource": _PROMPT_SOURCE,
**common,
}
)
parent = uid
elif role == "assistant":
blocks = _assistant_blocks(content)
message_id = f"msg_{_uuid.uuid4().hex[:24]}"
request_id = f"req_{_uuid.uuid4().hex[:24]}"
stop_reason = (
"tool_use"
if any(b.get("type") == "tool_use" for b in blocks)
else "end_turn"
)
for block in blocks:
uid = _new_uuid()
entries.append(
{
"parentUuid": parent,
"message": {
"model": model,
"id": message_id,
"type": "message",
"role": "assistant",
"content": [block],
"stop_reason": stop_reason,
"stop_sequence": None,
"stop_details": None,
"usage": _zero_usage(),
"diagnostics": None,
},
"requestId": request_id,
"type": "assistant",
"uuid": uid,
"timestamp": stamp,
**common,
}
)
if block.get("type") == "tool_use" and block.get("id"):
tool_use_owner[str(block["id"])] = uid
parent = uid
else:
msg = f"message role must be user or assistant, got {role!r}"
raise ValueError(msg)
return entries
def messages_from_entries(entries: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
out: list[dict[str, Any]] = []
last_message_id: str | None = None
for entry in entries:
kind = entry.get("type")
message = entry.get("message")
if kind not in ("user", "assistant") or not isinstance(message, dict):
continue
content = message.get("content")
if kind == "user":
results = _tool_results(content)
if (
results
and out
and out[-1]["role"] == "user"
and _tool_results(out[-1]["content"])
):
out[-1]["content"].extend(results)
else:
out.append({"role": "user", "content": _user_content(content)})
last_message_id = None
continue
blocks = _assistant_blocks(content)
message_id = message.get("id")
if (
out
and out[-1]["role"] == "assistant"
and message_id is not None
and message_id == last_message_id
):
out[-1]["content"].extend(blocks)
else:
out.append({"role": "assistant", "content": blocks})
last_message_id = message_id
return out
def _new_uuid() -> str:
return str(_uuid.uuid4())
def _tool_results(content: Any) -> list[dict[str, Any]]:
if not isinstance(content, list):
return []
return [
_clean_block(b)
for b in content
if isinstance(b, dict) and b.get("type") == "tool_result"
]
def _user_content(content: Any) -> Any:
if isinstance(content, str):
return content
if isinstance(content, list):
return [_clean_block(b) for b in content if isinstance(b, dict)]
msg = f"user content must be str or list, got {type(content).__name__}"
raise TypeError(msg)
def _assistant_blocks(content: Any) -> list[dict[str, Any]]:
if isinstance(content, str):
return [{"type": "text", "text": content}]
if isinstance(content, list):
blocks = [_clean_block(b) for b in content if isinstance(b, dict)]
if blocks:
return blocks
msg = "assistant content must be a non-empty list of blocks"
raise ValueError(msg)
def _clean_block(block: dict[str, Any]) -> dict[str, Any]:
kind = block.get("type")
if kind == "tool_use":
return {
"type": "tool_use",
"id": block.get("id", ""),
"name": block.get("name", ""),
"input": block.get("input", {}),
}
if kind == "tool_result":
out: dict[str, Any] = {
"type": "tool_result",
"tool_use_id": block.get("tool_use_id", ""),
"content": block.get("content", ""),
}
if block.get("is_error") is not None:
out["is_error"] = bool(block["is_error"])
return out
if kind == "thinking":
return {
"type": "thinking",
"thinking": block.get("thinking", ""),
"signature": block.get("signature", ""),
}
if kind == "text":
return {"type": "text", "text": str(block.get("text", ""))}
return dict(block)
def _result_text(block: dict[str, Any]) -> str:
content = block.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
return "\n".join(
str(b.get("text", ""))
for b in content
if isinstance(b, dict) and b.get("type") == "text"
)
return str(content)
def _zero_usage() -> dict[str, Any]:
return {
"input_tokens": 0,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"output_tokens": 0,
"service_tier": "standard",
"server_tool_use": {"web_search_requests": 0, "web_fetch_requests": 0},
"cache_creation": {
"ephemeral_1h_input_tokens": 0,
"ephemeral_5m_input_tokens": 0,
},
"inference_geo": "not_available",
"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 ""
def fingerprint(messages: Iterable[Mapping[str, Any]]) -> str:
"""Text-only hash of a history; stateless callers are keyed by it."""
turns: list[tuple[str, str]] = []
for message in messages:
text = text_of(message.get("content"))
if not text:
continue
role = str(message.get("role", ""))
if turns and turns[-1][0] == role:
turns[-1] = (role, turns[-1][1] + "\n" + text)
else:
turns.append((role, text))
digest = hashlib.sha1(usedforsecurity=False)
for role, text in turns:
digest.update(role.encode("utf-8"))
digest.update(b"\x00")
digest.update(text.strip().encode("utf-8"))
digest.update(b"\x01")
return digest.hexdigest()
def text_of(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts = [
str(b.get("text", ""))
for b in content
if isinstance(b, Mapping) and b.get("type") == "text"
]
return "\n".join(p for p in parts if p)
return ""
+39
View File
@@ -0,0 +1,39 @@
"""Side channel for what a backend learned during one turn.
Frontends pass ``capture=TurnCapture()`` through ``Backend.complete``'s
``**options``. Backends that keep state per conversation (the Claude SDK
adapter) fill it in after the stream closes; backends that don't
(anthropic, raycast) drop the kwarg and the frontend falls back to a
text-only history.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
__all__ = ["TurnCapture", "TurnUsage"]
@dataclass(frozen=True, slots=True)
class TurnUsage:
input_tokens: int = 0
output_tokens: int = 0
cache_read_tokens: int = 0
cache_creation_tokens: int = 0
cost_usd: float | None = None
duration_ms: int | None = None
num_turns: int | None = None
@dataclass
class TurnCapture:
synthesized_messages: list[dict[str, Any]] = field(default_factory=list)
"""The assistant/tool_result cycle of the turn as Anthropic messages."""
session_id: str | None = None
"""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()
@@ -564,7 +564,7 @@ def _collect_pty_sessions(runtime: GatewayRuntime) -> list[dict[str, Any]]:
"""Enumerate live PTY sessions across all backends.
A backend qualifies if it exposes a ``live_sessions`` mapping
(currently only ``ClaudeCodeBackendAdapter``). Other backend types
(none since the SDK backend; kept for M1b). Other backend types
are quietly skipped — the admin terminal viewer only makes sense for
PTY-backed agents.
+145 -91
View File
@@ -4,39 +4,44 @@ Exposes the gateway as an Anthropic-compatible Messages endpoint, so any
client that already speaks Anthropic (Cursor, Cline, the official SDK,
``curl``) can hit a configured agent by passing its name as ``model``.
The frontend is intentionally thin: it authenticates the bearer token,
resolves ``body.model`` to an agent + its backend, and then either
streams the backend's events straight to SSE or accumulates them into a
single ``Message`` for ``stream=false`` callers. All provider quirks
already live in the backend adapters; we don't translate here.
Phase 1.4 wires only ``RaycastAgent`` through ``RaycastBackend``;
``ClaudeAgent`` lands in Phase 2 and will plug into the same dispatch
table without changes to this module.
A Claude agent behind this endpoint is a ``deep`` conversation: the client
knows nothing about our ids, so the text fingerprint of the history it
sends is the ``(anthropic, fingerprint)`` binding of the conversation,
rebound after every turn to the fingerprint the next request will carry.
A history nobody has seen becomes a new conversation, materialized by the
home frontend of ``deep`` (the vault file), and every reply is published
on the bus so that file follows the chat. Other agents (Raycast) stay
stateless and are only archived through ``turn_log_handlers``.
"""
from __future__ import annotations
import json
import logging
import uuid
from typing import TYPE_CHECKING, Any
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.responses import JSONResponse, StreamingResponse
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.core import audit
from beaver_gateway.core.transcript import fingerprint, text_of
from beaver_gateway.core.turn_capture import TurnCapture
from beaver_gateway.core.turn_record import TurnRecord
from beaver_gateway.frontends._accumulate import StreamAccumulator, accumulate
from beaver_gateway.frontends._accumulate import StreamAccumulator
from beaver_gateway.frontends._auth import require_token
from beaver_gateway.frontends.base import Frontend
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Awaitable, Callable
from anthropic.types import Message, MessageParam
from beaver_gateway.core.conversations import Conversations
from beaver_gateway.core.events import MessageStreamEvent
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.storage.models import Conversation
_log = logging.getLogger("beaver_gateway.frontends.anthropic")
@@ -44,10 +49,16 @@ _log = logging.getLogger("beaver_gateway.frontends.anthropic")
__all__ = ["AnthropicMessagesFrontend"]
FRONTEND = "anthropic"
_TITLE_LEN = 80
class AnthropicMessagesFrontend(Frontend):
"""FastAPI app behind ``POST /v1/messages`` + ``GET /v1/models``."""
name = FRONTEND
kinds = ("deep",)
def __init__(
self,
*,
@@ -57,14 +68,6 @@ class AnthropicMessagesFrontend(Frontend):
) -> None:
self.host = host
self.port = port
# External URL prefix the reverse proxy (Caddy/nginx/Cloudflare/…)
# uses to reach this frontend, e.g. ``https://api.example.com/ai``.
# The frontend's internal paths (``/v1/messages``, ``/v1/models``)
# are appended to it when the admin dashboard renders copy-pastable
# URLs. Trailing slash is stripped so the concatenation is
# idempotent. ``None`` means "advertise raw ``host:port``" (dev /
# no proxy) — the dashboard then derives the base from the
# browser's own request hostname.
self.public_base_url = public_base_url.rstrip("/") if public_base_url else None
self._runtime: GatewayRuntime | None = None
self._app: FastAPI | None = None
@@ -74,9 +77,6 @@ class AnthropicMessagesFrontend(Frontend):
self._app = self._build_app(runtime)
async def serve(self) -> None:
# Local import: uvicorn pulls in a lot, no reason to load it when
# something else (a test, a script) imports this module just for
# the FastAPI factory.
import uvicorn
if self._app is None:
@@ -134,9 +134,6 @@ class AnthropicMessagesFrontend(Frontend):
backend = runtime.backends.get(agent.name)
if backend is None:
# Agent exists in config but its backend isn't wired in
# this phase (e.g. ClaudeAgent before Phase 2, or a
# RaycastAgent without RAYCAST_BEARER set at startup).
raise HTTPException(
status.HTTP_503_SERVICE_UNAVAILABLE,
f"no backend configured for agent {agent.name!r}",
@@ -144,7 +141,12 @@ class AnthropicMessagesFrontend(Frontend):
messages = body.get("messages") or []
system = body.get("system")
system_str = system if isinstance(system, str) else None
stream_flag = bool(body.get("stream", False))
if not messages or messages[-1].get("role") != "user":
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "the last message must be a user turn"
)
_log.info(
"messages: actor=%s agent=%s stream=%s msgs=%d",
@@ -153,10 +155,6 @@ class AnthropicMessagesFrontend(Frontend):
stream_flag,
len(messages),
)
# Record at request acceptance, not at stream completion:
# a long streaming response can be aborted mid-flight by
# the client, and we still want the row in the audit trail.
# Detail stays small — no message bodies, no system prompt.
await audit.log(
runtime,
actor=f"token:{token_name}",
@@ -166,74 +164,139 @@ class AnthropicMessagesFrontend(Frontend):
msgs=len(messages),
)
# Forward per-request knobs the Anthropic body may carry —
# backend adapters layer these over per-agent defaults. Only
# values explicitly present (not Anthropic-defaulted ones we
# never received) are forwarded, so the agent's default still
# wins when the caller omits the field.
options: dict[str, Any] = {}
if isinstance(body.get("temperature"), int | float):
options["temperature"] = body["temperature"]
if isinstance(agent, ClaudeAgent) and runtime.conversations is not None:
conv, history = await _resolve(runtime.conversations, agent, messages)
turn_id = f"turn_{uuid.uuid4().hex[:12]}"
capture = TurnCapture()
events = runtime.conversations.turn(
conv,
messages=messages,
origin="user",
capture=capture,
turn_id=turn_id,
)
events = backend.complete(
agent=agent,
messages=messages,
system=system if isinstance(system, str) else None,
**options,
)
async def after(message: Message) -> None:
await _finish(
runtime,
conv,
messages=messages,
message=message,
capture=capture,
turn_id=turn_id,
history=history,
)
else:
options: dict[str, Any] = {}
if isinstance(body.get("temperature"), int | float):
options["temperature"] = body["temperature"]
events = backend.complete(
agent=agent, messages=messages, system=system_str, **options
)
system_str = system if isinstance(system, str) else None
if stream_flag:
# Side-accumulate while streaming so we can still emit a
# ``TurnRecord`` after the response closes. The buffered
# ``Message`` lives only in this coroutine's frame; SSE
# bytes still flow to the client unchanged.
return StreamingResponse(
_sse_and_broadcast(
events,
runtime=runtime,
async def after(message: Message) -> None:
await _broadcast_turn(
runtime,
agent_name=agent.name,
input_messages=messages,
output_message=message,
system=system_str,
model=model,
),
)
if stream_flag:
return StreamingResponse(
_sse(events, model=model, after=after),
media_type="text/event-stream",
)
message = await accumulate(events, model=model)
await _broadcast_turn(
runtime,
agent_name=agent.name,
input_messages=messages,
output_message=message,
system=system_str,
)
acc = StreamAccumulator()
async for ev in events:
acc.feed(ev)
message = acc.finalize(model=model)
await after(message)
return JSONResponse(content=message.model_dump(mode="json"))
return app
async def _sse_and_broadcast(
async def _resolve(
conversations: Conversations, agent: ClaudeAgent, messages: list[MessageParam]
) -> tuple[Conversation, list[dict[str, str]] | None]:
prior = [dict(m) for m in messages[:-1]]
if prior:
conv = await conversations.find_bound(
frontend=FRONTEND, external_id=fingerprint(prior)
)
if conv is not None:
if conv.agent_name != agent.name:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
f"this chat runs on {conv.agent_name!r}, not {agent.name!r}",
)
return conv, None
first = next(
(text_of(m.get("content")) for m in messages if m.get("role") == "user"), ""
)
try:
conv = await conversations.create(
kind="deep",
agent=agent.name,
origin=FRONTEND,
title=" ".join(first.split())[:_TITLE_LEN] or None,
)
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
await conversations.materialize(conv)
history = [
{"role": str(m.get("role", "")), "text": text_of(m.get("content"))}
for m in prior
]
return conv, history
async def _finish(
runtime: GatewayRuntime,
conv: Conversation,
*,
messages: list[MessageParam],
message: Message,
capture: TurnCapture,
turn_id: str,
history: list[dict[str, str]] | None,
) -> None:
reply = "\n\n".join(
getattr(b, "text", "")
for b in message.content
if getattr(b, "type", "") == "text"
).strip()
tail = capture.synthesized_messages or [
{"role": "assistant", "content": [{"type": "text", "text": reply}]}
]
await runtime.conversations.bind(
conv, frontend=FRONTEND, external_id=fingerprint([*messages, *tail])
)
runtime.bus.publish(
"reply",
conversation_id=conv.external_id,
turn_id=turn_id,
source=FRONTEND,
prompt=text_of(messages[-1].get("content")),
text=reply,
history=history,
)
async def _sse(
events: AsyncIterator[MessageStreamEvent],
*,
runtime: GatewayRuntime,
agent_name: str,
input_messages: list[MessageParam],
system: str | None,
model: str,
after: Callable[[Message], Awaitable[None]],
) -> AsyncIterator[bytes]:
r"""Serialize an event stream to SSE; broadcast a :class:`TurnRecord` after.
r"""Serialize an event stream to SSE, then hand the assembled ``Message`` on.
Each event becomes ``event: <type>\ndata: <json>\n\n`` — the shape
Each event becomes ``event: <type>\ndata: <json>\n\n`` - the shape
the Anthropic SDK's SSE decoder expects. Errors mid-stream are
swallowed into a synthetic ``error`` event so the client sees the
failure rather than a hung connection.
The same events feed a :class:`StreamAccumulator` on the side so that
once the SSE response closes we can hand a fully-assembled
``Message`` to every ``runtime.turn_log_handlers`` entry (the
markdown frontend's archive logger lives in there). Broadcast
failures are caught — they must never bubble up to the client.
"""
acc = StreamAccumulator()
try:
@@ -248,14 +311,10 @@ async def _sse_and_broadcast(
)
yield f"event: error\ndata: {err}\n\n".encode()
return
message = acc.finalize(model=model)
await _broadcast_turn(
runtime,
agent_name=agent_name,
input_messages=input_messages,
output_message=message,
system=system,
)
try:
await after(acc.finalize(model=model))
except Exception: # noqa: BLE001
_log.exception("post-turn hook failed")
async def _broadcast_turn(
@@ -266,11 +325,6 @@ async def _broadcast_turn(
output_message: Message,
system: str | None,
) -> None:
"""Fire each ``turn_log_handlers`` entry with a fresh :class:`TurnRecord`.
Handler exceptions are caught and logged — they're observability
plumbing, not part of the user-visible request path.
"""
if not runtime.turn_log_handlers:
return
record = TurnRecord(
@@ -278,7 +332,7 @@ async def _broadcast_turn(
input_messages=list(input_messages),
output_message=output_message,
system=system,
source="anthropic",
source=FRONTEND,
)
for handler in runtime.turn_log_handlers:
try:
@@ -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,541 @@
"""``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 SEEDS
from beaver_gateway.core.kinds import Kind, as_kind
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):
name = "api"
kinds = ("master", "branch", "deep", "job")
def __init__(
self,
*,
host: str = "0.0.0.0", # noqa: S104
port: int = 8004,
public_base_url: str | None = None,
master_agent: str | None = None,
branch_agent: str | None = None,
deep_agent: str | None = None,
job_agent: 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.master_agent = master_agent
self.branch_agent = branch_agent
self.deep_agent = deep_agent
self.job_agent = job_agent
self._app: FastAPI | None = None
def agent_for(self, kind: Kind) -> str | None:
return {
"master": self.master_agent,
"branch": self.branch_agent,
"deep": self.deep_agent,
"job": self.job_agent,
}.get(kind)
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/agents")
async def list_agents(request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope=SCOPE)
return {
"agents": [
{
"name": a.name,
"model": a.model,
"kinds": list(getattr(a, "kinds", ())),
}
for a in runtime.agents
],
"frontends": [
{
"name": fe.name,
"kinds": list(fe.kinds),
"default_agents": {
k: fe.agent_for(k) for k in fe.kinds if fe.agent_for(k)
},
}
for fe in conversations.frontends
],
}
@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)
agent = data.get("agent")
try:
kind = as_kind(str(data.get("kind") or "deep"))
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
if kind == "fork":
raise HTTPException(status.HTTP_400_BAD_REQUEST, "fork is internal")
if agent is not None and (
not isinstance(agent, str) or agent not in runtime.agents
):
raise HTTPException(status.HTTP_400_BAD_REQUEST, "unknown `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, LookupError) 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=conv.agent_name,
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 = data.get("agent")
if agent is not None and not isinstance(agent, str):
raise HTTPException(status.HTTP_400_BAD_REQUEST, "`agent` must be a string")
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=child.agent_name,
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")
try:
await conversations.bind(
conv,
frontend=frontend,
external_id=external_id,
visible=bool(data.get("visible", True)),
)
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
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)
+31 -3
View File
@@ -12,16 +12,18 @@ 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
from beaver_gateway.backends.base import Backend
from beaver_gateway.core.auth import TokenStore
from beaver_gateway.core.kinds import Kind
from beaver_gateway.core.registry import AgentRegistry, McpRegistry
from beaver_gateway.core.turn_record import TurnRecord
from beaver_gateway.storage import Database
from beaver_gateway.storage.models import Conversation, ConversationBinding
TurnLogHandler = Callable[[TurnRecord], Awaitable[None]]
@@ -36,7 +38,7 @@ class GatewayRuntime:
the name in hand, so the indirection lives one step earlier.
``mcp_internal_urls`` is filled in Phase 2.1: one loopback URL per
declared ``McpServer`` so ``ClaudeCodeBackendAdapter`` (Phase 2.2)
declared ``McpServer`` so ``ClaudeSdkBackend``
can pass them to ``BackendOptions.mcp_servers`` without re-running
discovery.
@@ -77,13 +79,39 @@ 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):
"""Listens on a port, dispatches into the gateway."""
"""Listens on a port, dispatches into the gateway.
A frontend that shows conversations declares ``name`` (the binding
key) and ``kinds`` (which conversation kinds it shows);
``core/conversations`` refuses to bind a conversation to a frontend
outside its declaration. The first frontend in declaration order
whose ``materialize`` returns a binding is the *home* of that kind:
``spawn`` calls it so a new conversation gets a window (a vault file,
a Telegram topic). ``agent_for`` names the default agent for a kind
so callers may omit ``agent``. Stateless frontends (MCP, admin) keep
the defaults and stay outside the routing.
"""
name: str = ""
kinds: tuple[Kind, ...] = ()
@abstractmethod
def configure(self, runtime: GatewayRuntime) -> None: ...
@abstractmethod
async def serve(self) -> None: ...
def agent_for(self, kind: Kind) -> str | None: # noqa: ARG002
return None
async def materialize(self, conv: Conversation) -> ConversationBinding | None: # noqa: ARG002
return None
@@ -38,14 +38,16 @@ if TYPE_CHECKING:
_log = logging.getLogger("beaver_gateway.frontends.markdown.crossfront")
# User hook: take a turn + vault root, return where the new file should
# live. Returning a relative ``Path`` is treated as relative to the
# vault. ``None`` (the default) keeps the built-in
# ``{vault}/{logged_subdir}/{agent}/{YYYY-MM-DD}_{hex8}.md`` layout.
LogPathFn = "Callable[[TurnRecord, Path], Path]"
__all__ = [
"ChatPathFn",
"CrossFrontendLogger",
"fingerprint_messages",
"strip_trailing_user_scaffold",
]
__all__ = ["CrossFrontendLogger", "LogPathFn", "fingerprint_messages"]
ChatPathFn = "Callable[[str, str, Path], Path]"
"""``(title, agent, vault) -> path`` of a new chat file; relative = under the
vault. ``None`` keeps ``{vault}/{logged_subdir}/{agent}/{YYYY-MM-DD}_{slug}.md``."""
def fingerprint_messages(messages: Iterable[MessageParam]) -> str:
@@ -96,18 +98,18 @@ class CrossFrontendLogger:
*,
vault_path: Path,
logged_subdir: str,
log_path: Callable[[TurnRecord, Path], Path] | None = None,
chat_path: Callable[[str, str, Path], Path] | None = None,
) -> None:
self._vault = vault_path
self._root = vault_path / logged_subdir
self._index: dict[str, Path] = {}
self._lock = asyncio.Lock()
self._log_path_fn = log_path
self._chat_path_fn = chat_path
# When the user supplies a custom path function, files can land
# anywhere in the vault — so we have to scan the whole vault on
# startup to rebuild the fingerprint→path map. With the default
# layout we can bound the scan to ``_logs/``.
self._scan_root = vault_path if log_path is not None else self._root
self._scan_root = vault_path if chat_path is not None else self._root
def warm_index(self) -> None:
"""Scan logged files synchronously, populating the fingerprint map.
@@ -171,7 +173,7 @@ class CrossFrontendLogger:
if target.exists():
existing = target.read_text(encoding="utf-8")
parsed = frontmatter.loads(existing)
body = _strip_trailing_user_scaffold(parsed.content)
body = strip_trailing_user_scaffold(parsed.content)
# We append only the *new* user turn (the last one in
# input_messages, since prior turns are already on disk)
# plus the assistant reply.
@@ -206,13 +208,15 @@ class CrossFrontendLogger:
def _new_file_path(self, record: TurnRecord) -> Path:
"""Pick a fresh filename for a brand-new conversation.
With a user-supplied ``log_path`` we delegate to it (joining a
With a user-supplied ``chat_path`` we delegate to it (joining a
relative result with the vault root). Without one, we fall back
to ``{logged_subdir}/{agent}/{date}_{hex8}.md`` and ensure the
``.md`` suffix in case the user picks a non-md extension by hand.
"""
if self._log_path_fn is not None:
result = self._log_path_fn(record, self._vault)
if self._chat_path_fn is not None:
result = self._chat_path_fn(
record.first_user_text, record.agent_name, self._vault
)
if not result.is_absolute():
result = self._vault / result
if result.suffix != ".md":
@@ -270,7 +274,7 @@ def _render_full_history(messages: list[MessageParam], assistant: Any) -> str:
return body
def _strip_trailing_user_scaffold(body: str) -> str:
def strip_trailing_user_scaffold(body: str) -> str:
"""Drop a trailing empty ``### User:`` block if present.
Cross-frontend turns aren't typed into the file by the human — they
@@ -0,0 +1,50 @@
"""Vault file IO shared by the markdown frontend and its mirror."""
from __future__ import annotations
import contextlib
import os
import tempfile
from typing import TYPE_CHECKING, Any
import aiofile
import frontmatter
if TYPE_CHECKING:
from pathlib import Path
__all__ = ["read_or_empty", "reattach_frontmatter", "write_atomic"]
async def read_or_empty(path: Path) -> str:
if not path.exists(): # noqa: ASYNC240
return ""
async with aiofile.async_open(path, "r", encoding="utf-8") as f:
return await f.read()
async def write_atomic(path: Path, content: str) -> None:
"""Write via tmp + ``os.replace`` in the same directory: no torn reads."""
path.parent.mkdir(parents=True, exist_ok=True)
fd, tmp_path = tempfile.mkstemp(
prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent)
)
try:
async with aiofile.async_open(tmp_path, "w", encoding="utf-8") as f:
await f.write(content)
os.close(fd)
os.replace(tmp_path, path) # noqa: PTH105
except BaseException:
with contextlib.suppress(OSError):
os.close(fd)
with contextlib.suppress(OSError):
os.unlink(tmp_path) # noqa: PTH108
raise
def reattach_frontmatter(metadata: dict[str, Any], body: str) -> str:
"""Re-emit a ``.md`` with YAML frontmatter; no block at all for empty metadata."""
if not metadata:
return body if body.endswith("\n") else body + "\n"
post = frontmatter.Post(content=body, **metadata)
return frontmatter.dumps(post) + "\n"
+178 -283
View File
@@ -13,7 +13,16 @@ 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. The
frontend is the home of ``deep``: ``materialize`` gives a conversation
spawned elsewhere its file, and :class:`.mirror.ChatMirror` keeps that
file in step with replies produced outside ``/chat``.
Cross-frontend logging: when ``log_all_chats=True``, ``configure()``
registers a handler on ``runtime.turn_log_handlers`` so every other
@@ -28,44 +37,48 @@ import asyncio
import contextlib
import json
import logging
import os
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
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from beaver_gateway.backends.claude_code import ClaudeCodeBackendAdapter, TurnCapture
from beaver_gateway.core import audit
from beaver_gateway.core.conversation_store import (
diff_and_fork,
load_conversation,
load_messages,
mint_conversation,
rewrite_messages,
)
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
from beaver_gateway.frontends.markdown.files import (
read_or_empty,
reattach_frontmatter,
write_atomic,
)
from beaver_gateway.frontends.markdown.mirror import FRONTEND, ChatMirror
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Callable
from anthropic.types import MessageParam
from beaver_gateway.core.kinds import Kind
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.storage.models import Conversation, ConversationBinding
_log = logging.getLogger("beaver_gateway.frontends.markdown")
@@ -88,17 +101,13 @@ _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
class MarkdownFrontend(Frontend):
"""FastAPI app behind ``POST /chat`` driven by Obsidian-vault files."""
name = FRONTEND
kinds = ("deep",)
def __init__(
self,
*,
@@ -108,7 +117,7 @@ class MarkdownFrontend(Frontend):
default_agent: str | None = None,
log_all_chats: bool = False,
logged_subdir: str = "_logs",
log_path: Callable[[TurnRecord, Path], Path] | None = None,
chat_path: Callable[[str, str, Path], Path] | None = None,
public_base_url: str | None = None,
) -> None:
self.vault_path = Path(vault_path).expanduser().resolve()
@@ -117,7 +126,7 @@ class MarkdownFrontend(Frontend):
self.default_agent = default_agent
self.log_all_chats = log_all_chats
self.logged_subdir = logged_subdir
self.log_path = log_path
self.chat_path = chat_path
# External URL prefix when behind a reverse proxy — same role as
# on the other bearer frontends. Trailing slash trimmed for
# idempotent concatenation; ``None`` means "no proxy / advertise
@@ -131,15 +140,25 @@ class MarkdownFrontend(Frontend):
# request reliably loses the race to 409.
self._busy: set[Path] = set()
self._crossfront: CrossFrontendLogger | None = None
self._mirror: ChatMirror | 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)
self._mirror = ChatMirror(
vault_path=self.vault_path,
runtime=runtime,
logged_subdir=self.logged_subdir,
chat_path=self.chat_path,
)
if self.log_all_chats:
self._crossfront = CrossFrontendLogger(
vault_path=self.vault_path,
logged_subdir=self.logged_subdir,
log_path=self.log_path,
chat_path=self.chat_path,
)
# Scan the existing logged files synchronously here so the
# fingerprint→path map is populated before the first
@@ -148,6 +167,19 @@ class MarkdownFrontend(Frontend):
runtime.turn_log_handlers.append(self._crossfront.handle)
self._app = self._build_app(runtime)
@property
def mirror(self) -> ChatMirror:
if self._mirror is None:
msg = "configure() must be called first"
raise RuntimeError(msg)
return self._mirror
def agent_for(self, kind: Kind) -> str | None:
return self.default_agent if kind == "deep" else None
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
return await self.mirror.materialize(conv)
async def serve(self) -> None:
import uvicorn
@@ -158,7 +190,13 @@ class MarkdownFrontend(Frontend):
self._app, host=self.host, port=self.port, log_level="info"
)
server = uvicorn.Server(config)
await server.serve()
mirror = asyncio.create_task(self.mirror.run())
try:
await server.serve()
finally:
mirror.cancel()
with contextlib.suppress(asyncio.CancelledError):
await mirror
# ---- app builder ---------------------------------------------------
@@ -285,17 +323,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
@@ -312,11 +340,11 @@ class MarkdownFrontend(Frontend):
content_override: Any,
agent_override: str | None,
) -> Any:
write_disk = content_override is None
if isinstance(content_override, str):
await _write_atomic(file_path, content_override)
file_text = content_override
elif content_override is None:
file_text = await _read_or_empty(file_path)
file_text = await read_or_empty(file_path)
else:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "`content` must be a string when present"
@@ -381,22 +409,21 @@ class MarkdownFrontend(Frontend):
# Resolve / mint the conversation row, align incoming against
# stored history, and feed the aligned messages to the backend
# — see ``core/conversation_store.py`` for the full rationale.
# If the backend isn't claude-code (no ``TurnCapture`` support)
# we fall through to the legacy parser-only path.
# - 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 | None = (
TurnCapture() if isinstance(backend, ClaudeCodeBackendAdapter) else None
)
kwargs: dict[str, Any] = {}
if capture is not None:
kwargs["capture"] = capture
events = backend.complete(
agent=agent, messages=outcome.messages, system=None, **kwargs
capture = TurnCapture()
events = runtime.conversations.turn(
conv,
messages=outcome.messages,
origin="user",
capture=capture,
use_session=outcome.divergence_index is None,
)
try:
message = await self._stream_to_file(
@@ -405,6 +432,7 @@ class MarkdownFrontend(Frontend):
parsed=parsed,
model=agent.model or agent.name,
filename=filename,
write_disk=write_disk,
)
except HTTPException:
raise
@@ -419,11 +447,12 @@ class MarkdownFrontend(Frontend):
message=message,
agent_name=agent.name,
conv_external_id=conv_external_id,
write_disk=write_disk,
)
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,
@@ -480,17 +509,15 @@ class MarkdownFrontend(Frontend):
writers in their respective halves of Obsidian Sync. Final
content is identical on both sides, so Sync no-ops.
"""
# File-text resolution + early bailouts. ``content_override`` is
# still written to disk on the gateway side because that's the
# state the rest of the request consumes; it just doesn't keep
# ticking after that single write.
# With ``content`` the plugin is the only writer of the file
# (§3.10): the gateway never touches disk in that case.
write_disk = content_override is None
if isinstance(content_override, str):
await _write_atomic(file_path, content_override)
file_text = content_override
elif content_override is None:
file_text = await _read_or_empty(file_path)
file_text = await read_or_empty(file_path)
else:
yield _sse_pack(
yield sse_pack(
"error",
{
"status_code": status.HTTP_400_BAD_REQUEST,
@@ -506,7 +533,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,
@@ -519,7 +546,7 @@ class MarkdownFrontend(Frontend):
return
if not parsed.messages:
yield _sse_pack(
yield sse_pack(
"done",
{
"status": "nothing_to_do",
@@ -530,7 +557,7 @@ class MarkdownFrontend(Frontend):
return
if parser.last_role(parsed.messages) == "assistant":
yield _sse_pack(
yield sse_pack(
"done",
{
"status": "nothing_to_do",
@@ -542,7 +569,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,
@@ -552,7 +579,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,
@@ -577,9 +604,18 @@ class MarkdownFrontend(Frontend):
msgs=len(parsed.messages),
)
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
runtime=runtime, metadata=parsed.metadata, agent_name=agent.name
)
try:
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
runtime=runtime,
metadata=parsed.metadata,
agent_name=agent.name,
file_path=file_path,
)
except HTTPException as exc:
yield sse_pack(
"error", {"status_code": exc.status_code, "detail": exc.detail}
)
return
_log.info(
"chat/stream: file=%s conv_external_id=%s conv_id=%d "
"stored_msgs=%d incoming_turns=%d",
@@ -598,21 +634,19 @@ class MarkdownFrontend(Frontend):
len(outcome.messages),
len(outcome.persist_messages),
)
capture: TurnCapture | None = (
TurnCapture() if isinstance(backend, ClaudeCodeBackendAdapter) else None
)
kwargs: dict[str, Any] = {}
if capture is not None:
kwargs["capture"] = capture
capture = TurnCapture()
_log.info(
"chat/stream: file=%s calling backend.complete agent=%s capture=%s",
"chat/stream: file=%s calling backend.complete agent=%s session=%s",
filename,
agent.name,
capture is not None,
conv.session_id,
)
events = backend.complete(
agent=agent, messages=outcome.messages, system=None, **kwargs
events = runtime.conversations.turn(
conv,
messages=outcome.messages,
origin="user",
capture=capture,
use_session=outcome.divergence_index is None,
)
acc = StreamAccumulator()
@@ -626,16 +660,12 @@ class MarkdownFrontend(Frontend):
return None
rendered = renderer.render_assistant_message(partial)
new_body = renderer.append_to_body(parsed.body, rendered)
return _reattach_frontmatter(parsed.metadata, new_body)
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()
@@ -648,7 +678,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
@@ -663,10 +693,11 @@ class MarkdownFrontend(Frontend):
new_body, renderer.render_assistant_message(partial)
)
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
await _write_atomic(
file_path, _reattach_frontmatter(parsed.metadata, new_body)
)
yield _sse_pack(
if write_disk:
await write_atomic(
file_path, reattach_frontmatter(parsed.metadata, new_body)
)
yield sse_pack(
"error",
{
"status_code": status.HTTP_500_INTERNAL_SERVER_ERROR,
@@ -683,11 +714,12 @@ class MarkdownFrontend(Frontend):
message=message,
agent_name=agent.name,
conv_external_id=conv_external_id,
write_disk=write_disk,
)
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,
@@ -707,7 +739,7 @@ class MarkdownFrontend(Frontend):
except Exception: # noqa: BLE001
_log.exception("turn_log_handler raised; continuing")
yield _sse_pack(
yield sse_pack(
"done",
{
"status": "ok",
@@ -727,6 +759,7 @@ class MarkdownFrontend(Frontend):
parsed: parser.ParsedFile,
model: str,
filename: str,
write_disk: bool = True,
) -> Any:
"""Drain ``events`` into a ``Message``, flushing partials to disk.
@@ -745,13 +778,15 @@ class MarkdownFrontend(Frontend):
acc = StreamAccumulator()
async def flush_partial() -> None:
if not write_disk:
return
partial = acc.finalize(model=model)
if not partial.content:
return
rendered = renderer.render_assistant_message(partial)
new_body = renderer.append_to_body(parsed.body, rendered)
await _write_atomic(
file_path, _reattach_frontmatter(parsed.metadata, new_body)
await write_atomic(
file_path, reattach_frontmatter(parsed.metadata, new_body)
)
try:
@@ -774,9 +809,10 @@ class MarkdownFrontend(Frontend):
new_body, renderer.render_assistant_message(partial)
)
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
await _write_atomic(
file_path, _reattach_frontmatter(parsed.metadata, new_body)
)
if write_disk:
await write_atomic(
file_path, reattach_frontmatter(parsed.metadata, new_body)
)
raise
return acc.finalize(model=model)
@@ -788,75 +824,62 @@ class MarkdownFrontend(Frontend):
message: Any,
agent_name: str,
conv_external_id: str,
write_disk: bool = True,
) -> str:
"""Render the assistant turn, append to the file, refresh frontmatter."""
"""Render the assistant turn, refresh frontmatter, write if we own the file."""
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)
await _write_atomic(file_path, new_content)
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:
try:
conv = await conversations.create(
kind="deep", agent=agent_name, origin=FRONTEND, title=file_path.stem
)
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
_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
@@ -867,21 +890,17 @@ class MarkdownFrontend(Frontend):
conversation_id: int,
persist_messages: list[dict[str, Any]],
new_user_text: str,
capture: TurnCapture | None,
capture: TurnCapture,
message: Any,
) -> None:
"""Stamp the DB with the post-turn canonical Anthropic-shape history.
Combines the matched/spliced prior state, the new user prompt,
and the synthesized assistant↔tool cycle from the backend (or
a text-only fallback for backends without ``TurnCapture``).
and the synthesized assistant/tool cycle from the backend (or a
text-only fallback for backends that left ``capture`` empty).
"""
new_user_msg = {"role": "user", "content": new_user_text}
synthesized = (
capture.synthesized_messages
if capture is not None
else _fallback_synthesized(message)
)
synthesized = capture.synthesized_messages or _fallback_synthesized(message)
canonical = [*persist_messages, new_user_msg, *synthesized]
_log.info(
"_persist_canonical_history: conv_id=%d writing %d msgs "
@@ -919,115 +938,6 @@ class MarkdownFrontend(Frontend):
return candidate
# ---- module-level utilities ----------------------------------------------
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
# gating an async read on whether the file is there is exactly the
# check we want. Switching to anyio.Path / aiofiles.os just to
# silence the async-pathlib lint would cost a dep edge for no
# practical win.
if not path.exists(): # noqa: ASYNC240
return ""
async with aiofile.async_open(path, "r", encoding="utf-8") as f:
return await f.read()
async def _write_atomic(path: Path, content: str) -> None:
"""Write ``content`` to ``path`` via tmp + ``os.replace`` (atomic)."""
path.parent.mkdir(parents=True, exist_ok=True)
# ``NamedTemporaryFile`` keeps the file open which complicates
# ``os.replace`` on some platforms. Build the tmp name manually,
# write+fsync, then rename. Same-directory so the rename is atomic.
tmp_name = tempfile.mkstemp(
prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent)
)
fd, tmp_path = tmp_name
try:
async with aiofile.async_open(tmp_path, "w", encoding="utf-8") as f:
await f.write(content)
os.close(fd)
# ``os.replace`` is the atomic primitive — ``Path.replace`` is a
# thin wrapper around the same syscall; either works, ``os.`` is
# the one Linux/POSIX docs reach for.
os.replace(tmp_path, path) # noqa: PTH105
except BaseException:
# Cleanup on failure: close fd, remove tmp.
with contextlib.suppress(OSError):
os.close(fd)
with contextlib.suppress(OSError):
os.unlink(tmp_path) # noqa: PTH108
raise
def _reattach_frontmatter(metadata: dict[str, Any], body: str) -> str:
r"""Re-emit a ``.md`` file with YAML frontmatter at the top.
Empty metadata → no frontmatter block (avoid littering every file
with a hollow ``---\n---``).
"""
if not metadata:
return body if body.endswith("\n") else body + "\n"
import frontmatter as _fm
post = _fm.Post(content=body, **metadata)
return _fm.dumps(post) + "\n"
def _fallback_synthesized(message: Any) -> list[dict[str, Any]]:
"""Build a single-assistant ``synthesized_messages`` list from a raw ``Message``.
@@ -1065,21 +975,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__
@@ -0,0 +1,169 @@
"""Vault files for ``deep`` conversations that were not typed into a file (§3.10).
``materialize`` is the markdown frontend's answer to ``spawn(kind=deep)``:
a new file in the vault with ``agent`` + ``conversation_id`` frontmatter
and the ``(markdown, path)`` binding. ``run`` tails the gateway bus and
appends every ``reply`` of a markdown-bound conversation to its file -
the seed turn of a spawn, a message posted through ``/api``, a turn
that came in over ``/v1/messages`` - and stamps the same exchange into
the canonical history, so a continuation typed in Obsidian aligns
against the store and resumes the same SDK session instead of reseeding.
"""
from __future__ import annotations
import asyncio
import logging
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
import frontmatter
from beaver_gateway.core.conversation_store import load_messages, rewrite_messages
from beaver_gateway.core.turn_record import slugify
from beaver_gateway.frontends.markdown import renderer
from beaver_gateway.frontends.markdown.crossfront import strip_trailing_user_scaffold
from beaver_gateway.frontends.markdown.files import (
read_or_empty,
reattach_frontmatter,
write_atomic,
)
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
from pathlib import Path
from beaver_gateway.core.bus import Event
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.storage.models import Conversation, ConversationBinding
__all__ = ["FRONTEND", "ChatMirror"]
_log = logging.getLogger("beaver_gateway.frontends.markdown.mirror")
FRONTEND = "markdown"
class ChatMirror:
def __init__(
self,
*,
vault_path: Path,
runtime: GatewayRuntime,
logged_subdir: str,
chat_path: Callable[[str, str, Path], Path] | None = None,
) -> None:
self._vault = vault_path
self._runtime = runtime
self._root = vault_path / logged_subdir
self._chat_path_fn = chat_path
self._lock = asyncio.Lock()
def path_for(self, title: str, agent: str) -> Path:
if self._chat_path_fn is not None:
path = self._chat_path_fn(title, agent, self._vault)
if not path.is_absolute():
path = self._vault / path
else:
day = datetime.now(UTC).strftime("%Y-%m-%d")
path = self._root / agent / f"{day}_{slugify(title, maxlen=60)}.md"
if path.suffix != ".md":
path = path.with_suffix(".md")
path = path.resolve()
path.relative_to(self._vault)
candidate, n = path, 1
while candidate.exists():
n += 1
candidate = path.with_name(f"{path.stem} ({n}){path.suffix}")
return candidate
async def materialize(self, conv: Conversation) -> ConversationBinding:
path = self.path_for(conv.title or conv.external_id, conv.agent_name)
rel = path.relative_to(self._vault).as_posix()
async with self._lock:
await write_atomic(path, reattach_frontmatter(_frontmatter(conv), ""))
binding = await self._runtime.conversations.bind(
conv, frontend=FRONTEND, external_id=rel
)
_log.info("materialized %s as %s", conv.external_id, rel)
return binding
async def bound_path(self, conv: Conversation) -> Path | None:
for b in await self._runtime.conversations.bindings(conv):
if b.frontend == FRONTEND and b.visible:
return self._vault / b.external_id
return None
async def run(self) -> None:
async for event in self._runtime.bus.stream():
if event.get("type") != "reply":
continue
try:
await self._on_reply(event)
except Exception: # noqa: BLE001
_log.exception("mirror of reply %s failed", event.get("turn_id"))
async def _on_reply(self, event: Event) -> None:
conv = await self._runtime.conversations.get(str(event["conversation_id"]))
if conv is None:
return
await self.append(
conv,
prompt=str(event.get("prompt") or ""),
text=str(event.get("text") or ""),
history=event.get("history"),
)
async def append(
self,
conv: Conversation,
*,
prompt: str,
text: str,
history: Sequence[dict[str, Any]] | None = None,
) -> None:
path = await self.bound_path(conv)
if path is None:
return
prior = [
{"role": str(h["role"]), "content": str(h["text"])}
for h in history or ()
if h.get("text")
]
async with self._lock:
parsed = frontmatter.loads(await read_or_empty(path))
body = strip_trailing_user_scaffold(parsed.content)
if prior and "### " not in body:
for m in prior:
body = renderer.append_to_body(body, _render(m))
body = renderer.append_to_body(body, renderer.render_user_text(prompt))
body = renderer.append_to_body(body, renderer.render_assistant_text(text))
body = renderer.append_to_body(body, renderer.USER_SCAFFOLD)
metadata = {**parsed.metadata, **_frontmatter(conv)}
metadata.pop("fingerprint", None)
await write_atomic(path, reattach_frontmatter(metadata, body))
await self._persist(conv, prior=prior, prompt=prompt, text=text)
async def _persist(
self, conv: Conversation, *, prior: list[dict[str, Any]], prompt: str, text: str
) -> None:
if conv.id is None:
return
async with self._runtime.db.session() as session:
stored = await load_messages(session, conversation_id=conv.id)
canonical = [
*(stored or prior),
{"role": "user", "content": prompt},
{"role": "assistant", "content": [{"type": "text", "text": text}]},
]
await rewrite_messages(session, conversation_id=conv.id, messages=canonical)
def _frontmatter(conv: Conversation) -> dict[str, Any]:
return {"agent": conv.agent_name, "conversation_id": conv.external_id}
def _render(message: dict[str, Any]) -> str:
if message["role"] == "user":
return renderer.render_user_text(message["content"])
return renderer.render_assistant_text(message["content"])
@@ -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
@@ -54,6 +53,11 @@ def render_user_text(content: str) -> str:
return f"### User:\n\n{content.strip()}\n"
def render_assistant_text(text: str) -> str:
r"""Render a text-only assistant turn as ``### Assistant:\n\n<text>``."""
return f"### Assistant:\n\n{text.strip()}\n"
def render_assistant_message(message: Message) -> str:
"""Render an assistant ``Message`` (with content blocks) into a turn block.
@@ -61,9 +65,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 +150,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 +159,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``.
+1 -1
View File
@@ -17,7 +17,7 @@ heavy setups (PRD §6 cites the ~95%→~71% tool-selection drop on flat
namespaces) but real and reachable.
The aggregator returns both the app and a ``{name: url}`` map; Phase
2.2's ``ClaudeCodeBackendAdapter`` plugs the map directly into
``ClaudeSdkBackend`` plugs the map directly into
``BackendOptions.mcp_servers``. ``/mcp/all/`` is NOT included in that
map — claude-code-agents always get per-domain framing; only the
external MCP frontend (Phase 3.1) reverse-proxies the flat endpoint.
+6
View File
@@ -27,6 +27,12 @@ class Settings(BaseSettings):
internal_mcp_port: int = 8765
config_path: Path = Path("/config/config.py")
claude_runner_user: str | None = None
"""Unix user the claude subprocess runs as. Needs the gateway to be root."""
claude_home: Path | None = None
"""``HOME`` for the claude subprocess (its ``~/.claude`` lives there)."""
raycast_bearer: str | None = None
raycast_config_path: Path = Path("/config/raycast.json")
raycast_device_id: str | None = None
+23 -5
View File
@@ -1,14 +1,14 @@
"""SQLModel-backed persistence.
Two tables — :class:`Token`, :class:`AuditLog` — plus a thin
:class:`Database` wrapper around an async SQLAlchemy engine. The
``GatewayRuntime`` carries the handle so auth (token verify / touch)
and admin (token CRUD + audit listing) can reach it.
Tokens, audit, conversations, Agent SDK transcript entries and per-turn
usage, plus a thin :class:`Database` wrapper around an async SQLAlchemy
engine. ``GatewayRuntime`` carries the handle.
"""
from beaver_gateway.storage.db import (
Database,
append_audit,
append_usage,
create_token,
list_active_tokens,
list_audit_records,
@@ -16,13 +16,31 @@ from beaver_gateway.storage.db import (
revoke_token,
touch_token,
)
from beaver_gateway.storage.models import AuditLog, Token
from beaver_gateway.storage.models import (
AuditLog,
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",
"append_audit",
"append_usage",
"create_token",
"list_active_tokens",
"list_audit_records",
+51 -2
View File
@@ -18,15 +18,17 @@ import json
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
from sqlalchemy import inspect, text
from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession
from beaver_gateway.storage.models import AuditLog, Token
from beaver_gateway.storage.models import AuditLog, Token, Usage
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy import Connection
from sqlalchemy.ext.asyncio import AsyncEngine
@@ -72,9 +74,14 @@ class Database:
self._engine: AsyncEngine = create_async_engine(_to_async_url(url), echo=False)
async def create_all(self) -> None:
"""Idempotent ``CREATE TABLE IF NOT EXISTS`` for every model."""
"""Create missing tables, then add columns models gained since.
Nullable columns only - that is the whole migration story until
Alembic is worth it.
"""
async with self._engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
await conn.run_sync(_add_missing_columns)
def session(self) -> AsyncSession:
"""Open a fresh :class:`AsyncSession` (use as ``async with``).
@@ -90,6 +97,39 @@ class Database:
await self._engine.dispose()
def _add_missing_columns(conn: Connection) -> None:
inspector = inspect(conn)
for table in SQLModel.metadata.sorted_tables:
existing = {c["name"] for c in inspector.get_columns(table.name)}
for column in table.columns:
if column.name in existing:
continue
kind = column.type.compile(conn.dialect)
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 ---------------------------------------------------------
@@ -183,9 +223,18 @@ async def list_audit_records(
return result.all()
# ---- Usage --------------------------------------------------------------
async def append_usage(session: AsyncSession, row: Usage) -> None:
session.add(row)
await session.commit()
__all__ = [
"Database",
"append_audit",
"append_usage",
"create_token",
"list_active_tokens",
"list_audit_records",
+171 -11
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).
@@ -23,8 +23,10 @@ needs an id uses ``Optional[int]`` so SQLAlchemy can autoincrement.
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import UniqueConstraint
from sqlalchemy import JSON, Column, Index, UniqueConstraint, text
from sqlalchemy.dialects.postgresql import JSONB
from sqlmodel import Field, SQLModel
@@ -72,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"
@@ -92,8 +95,93 @@ class Conversation(SQLModel, table=True):
frontend: str = Field(index=True)
external_id: str = Field(index=True)
agent_name: str = Field(index=True)
session_id: str | None = Field(default=None, index=True)
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):
@@ -125,4 +213,76 @@ class ConversationMessage(SQLModel, table=True):
created_at: datetime = Field(default_factory=_utcnow)
__all__ = ["AuditLog", "Conversation", "ConversationMessage", "Token"]
class TranscriptEntry(SQLModel, table=True):
"""One Agent SDK transcript line, mirrored from the session store protocol.
``(project_key, session_id, subpath)`` is the ``SessionKey``; ``subpath``
is ``""`` for the main transcript so the unique constraints stay simple.
``uuid`` is the SDK's idempotency key - the partial unique index rejects
a replayed batch, entries without a uuid are appended as-is.
"""
__tablename__ = "transcript_entries"
__table_args__ = (
UniqueConstraint(
"project_key", "session_id", "subpath", "seq", name="uq_transcript_seq"
),
Index(
"uq_transcript_uuid",
"project_key",
"session_id",
"subpath",
"uuid",
unique=True,
postgresql_where=text("uuid IS NOT NULL"),
sqlite_where=text("uuid IS NOT NULL"),
),
)
id: int | None = Field(default=None, primary_key=True)
project_key: str = Field(index=True)
session_id: str = Field(index=True)
subpath: str = Field(default="")
seq: int
uuid: str | None = Field(default=None)
entry: dict[str, Any] = Field(
sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=False)
)
origin: str | None = Field(default=None)
source: str | None = Field(default=None)
turn_id: str | None = Field(default=None)
created_at: datetime = Field(default_factory=_utcnow)
class Usage(SQLModel, table=True):
"""Per-turn token accounting from ``ResultMessage.usage``."""
__tablename__ = "usage"
id: int | None = Field(default=None, primary_key=True)
ts: datetime = Field(default_factory=_utcnow, index=True)
agent_name: str = Field(index=True)
conversation_id: str | None = Field(default=None, index=True)
session_id: str | None = Field(default=None, index=True)
model: str
effort: str | None = Field(default=None)
input_tokens: int = 0
output_tokens: int = 0
cache_read_tokens: int = 0
cache_creation_tokens: int = 0
cost_usd: float | None = Field(default=None)
duration_ms: int | None = Field(default=None)
num_turns: int | None = Field(default=None)
__all__ = [
"AuditLog",
"Conversation",
"ConversationBinding",
"ConversationMessage",
"InjectQueueItem",
"Schedule",
"Token",
"TranscriptEntry",
"Usage",
]
+126
View File
@@ -0,0 +1,126 @@
"""``SessionStore`` adapter for the Claude Agent SDK on top of :class:`Database`.
Entries are stored verbatim as JSON (``jsonb`` on Postgres), ordered by a
per-key ``seq``. ``append`` is idempotent on ``entry["uuid"]`` because the
SDK re-delivers a batch on retry; entries without a uuid are appended as
they come. Runs on SQLite too - the conformance suite uses that in tests.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, cast
from claude_agent_sdk import SessionStore
from sqlalchemy import delete, func
from sqlmodel import col, select
from beaver_gateway.storage.models import TranscriptEntry
if TYPE_CHECKING:
from claude_agent_sdk.types import (
SessionKey,
SessionListSubkeysKey,
SessionStoreEntry,
)
from beaver_gateway.storage.db import Database
__all__ = ["PostgresSessionStore"]
class PostgresSessionStore(SessionStore):
def __init__(self, db: Database) -> None:
self._db = db
async def append(self, key: SessionKey, entries: list[SessionStoreEntry]) -> None:
if not entries:
return
project_key, session_id, subpath = _parts(key)
wanted = [
u for u in (e.get("uuid") for e in entries) if isinstance(u, str) and u
]
async with self._db.session() as session:
known: set[str] = set()
if wanted:
result = await session.exec(
select(TranscriptEntry.uuid).where(
TranscriptEntry.project_key == project_key,
TranscriptEntry.session_id == session_id,
TranscriptEntry.subpath == subpath,
col(TranscriptEntry.uuid).in_(wanted),
)
)
known = {u for u in result.all() if u is not None}
seq_result = await session.exec(
select(func.max(TranscriptEntry.seq)).where(
TranscriptEntry.project_key == project_key,
TranscriptEntry.session_id == session_id,
TranscriptEntry.subpath == subpath,
)
)
seq = seq_result.one() or 0
for entry in entries:
uuid = entry.get("uuid")
uuid_str = uuid if isinstance(uuid, str) and uuid else None
if uuid_str is not None:
if uuid_str in known:
continue
known.add(uuid_str)
seq += 1
session.add(
TranscriptEntry(
project_key=project_key,
session_id=session_id,
subpath=subpath,
seq=seq,
uuid=uuid_str,
entry=cast("dict[str, Any]", dict(entry)),
)
)
await session.commit()
async def load(self, key: SessionKey) -> list[SessionStoreEntry] | None:
project_key, session_id, subpath = _parts(key)
async with self._db.session() as session:
result = await session.exec(
select(TranscriptEntry.entry)
.where(
TranscriptEntry.project_key == project_key,
TranscriptEntry.session_id == session_id,
TranscriptEntry.subpath == subpath,
)
.order_by(col(TranscriptEntry.seq))
)
rows = result.all()
if not rows:
return None
return [cast("SessionStoreEntry", row) for row in rows]
async def list_subkeys(self, key: SessionListSubkeysKey) -> list[str]:
async with self._db.session() as session:
result = await session.exec(
select(TranscriptEntry.subpath)
.where(
TranscriptEntry.project_key == key["project_key"],
TranscriptEntry.session_id == key["session_id"],
TranscriptEntry.subpath != "",
)
.distinct()
)
return sorted(result.all())
async def delete(self, key: SessionKey) -> None:
project_key, session_id, subpath = _parts(key)
stmt = delete(TranscriptEntry).where(
col(TranscriptEntry.project_key) == project_key,
col(TranscriptEntry.session_id) == session_id,
)
if key.get("subpath"):
stmt = stmt.where(col(TranscriptEntry.subpath) == subpath)
async with self._db.session() as session:
await session.execute(stmt) # ty: ignore[deprecated]
await session.commit()
def _parts(key: SessionKey) -> tuple[str, str, str]:
return key["project_key"], key["session_id"], key.get("subpath") or ""
+421
View File
@@ -0,0 +1,421 @@
import tempfile
from pathlib import Path
from typing import Any
import pytest
from anthropic.types import (
RawContentBlockDeltaEvent,
RawContentBlockStartEvent,
RawContentBlockStopEvent,
RawMessageDeltaEvent,
RawMessageStartEvent,
RawMessageStopEvent,
)
from claude_agent_sdk import (
AssistantMessage,
InMemorySessionStore,
ResultMessage,
StreamEvent,
TextBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
project_key_for_directory,
)
from beaver_gateway.agents.base import ExposedMcp
from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions, Prompts
from beaver_gateway.backends.claude_sdk import (
ClaudeSdkBackend,
RunnerConfig,
UsageEvent,
fingerprint,
)
from beaver_gateway.core.transcript import messages_from_entries
from beaver_gateway.core.turn_capture import TurnCapture
def _stream(index: int, text: str) -> list[StreamEvent]:
def ev(event: dict[str, Any]) -> StreamEvent:
return StreamEvent(uuid="u", session_id="s", event=event)
return [
ev(
{
"type": "content_block_start",
"index": index,
"content_block": {"type": "text", "text": ""},
}
),
ev(
{
"type": "content_block_delta",
"index": index,
"delta": {"type": "text_delta", "text": text},
}
),
ev({"type": "content_block_stop", "index": index}),
]
def _result(session_id: str) -> ResultMessage:
return ResultMessage(
subtype="success",
duration_ms=10,
duration_api_ms=5,
is_error=False,
num_turns=1,
session_id=session_id,
stop_reason="end_turn",
total_cost_usd=0.01,
usage={
"input_tokens": 3,
"output_tokens": 7,
"cache_read_input_tokens": 100,
"cache_creation_input_tokens": 20,
},
)
class FakeClient:
instances: list["FakeClient"] = []
def __init__(self, options: Any) -> None:
self.options = options
self.prompts: list[str] = []
self.connected = False
self.session_id = options.resume or "fresh-session"
FakeClient.instances.append(self)
async def connect(self) -> None:
self.connected = True
async def query(self, prompt: str) -> None:
self.prompts.append(prompt)
async def interrupt(self) -> None:
self.interrupted = True
async def receive_response(self):
start = {"type": "message_start", "message": {}}
yield StreamEvent(uuid="u", session_id="s", event=start)
for e in _stream(0, "calling "):
yield e
tool = ToolUseBlock(id="toolu_1", name="Read", input={"path": "x"})
yield StreamEvent(
uuid="u",
session_id="s",
event={
"type": "content_block_start",
"index": 1,
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "Read"},
},
)
yield StreamEvent(
uuid="u", session_id="s", event={"type": "content_block_stop", "index": 1}
)
yield AssistantMessage(content=[TextBlock(text="calling "), tool], model="m")
yield UserMessage(
content=[ToolResultBlock(tool_use_id="toolu_1", content="42")]
)
yield StreamEvent(uuid="u", session_id="s", event=start)
for e in _stream(0, "done"):
yield e
yield AssistantMessage(content=[TextBlock(text="done")], model="m")
yield StreamEvent(
uuid="sub",
session_id="s",
event={"type": "content_block_stop", "index": 5},
parent_tool_use_id="toolu_9",
)
yield _result(self.session_id)
async def disconnect(self) -> None:
self.connected = False
@pytest.fixture
def cwd() -> Path:
return Path(tempfile.mkdtemp(prefix="beaver-sdk-"))
def _backend(
cwd: Path, store: InMemorySessionStore, sink=None, **agent_kwargs
) -> ClaudeSdkBackend:
FakeClient.instances.clear()
agent = ClaudeAgent(
name="a",
model="claude-x",
system_prompt="hi",
cwd=cwd,
options=ClaudeOptions(effort="low"),
**agent_kwargs,
)
return ClaudeSdkBackend(
agent=agent,
mcp_internal_urls={"firefly": "http://127.0.0.1:1/mcp/firefly/"},
session_store=store,
mcp_tool_names={"firefly": ["list_account", "delete_account", "store_account"]},
usage_sink=sink,
client_factory=FakeClient,
work_dir=cwd / "work",
)
async def _drain(events) -> list[Any]:
return [e async for e in events]
async def test_stream_envelope_and_index_rebase(cwd: Path) -> None:
backend = _backend(
cwd, InMemorySessionStore(), expose_mcps=(ExposedMcp(name="firefly"),)
)
capture = TurnCapture()
events = await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": "hi"}],
conversation_id="conv-1",
capture=capture,
)
)
assert isinstance(events[0], RawMessageStartEvent)
assert isinstance(events[-1], RawMessageStopEvent)
assert isinstance(events[-2], RawMessageDeltaEvent)
assert events[-2].usage.output_tokens == 7
assert events[-2].usage.cache_read_input_tokens == 100
starts = [e.index for e in events if isinstance(e, RawContentBlockStartEvent)]
assert starts == [0, 1, 2]
stops = [e.index for e in events if isinstance(e, RawContentBlockStopEvent)]
assert stops == [0, 1, 2]
deltas = [e.delta.text for e in events if isinstance(e, RawContentBlockDeltaEvent)]
assert deltas == ["calling ", "done"]
assert capture.session_id == "fresh-session"
assert capture.usage is not None and capture.usage.cost_usd == 0.01
assert capture.synthesized_messages == [
{
"role": "assistant",
"content": [
{"type": "text", "text": "calling "},
{
"type": "tool_use",
"id": "toolu_1",
"name": "Read",
"input": {"path": "x"},
},
],
},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "42"}
],
},
{"role": "assistant", "content": [{"type": "text", "text": "done"}]},
]
client = FakeClient.instances[0]
assert client.prompts == ["hi"]
opts = client.options
assert opts.setting_sources == []
assert opts.strict_mcp_config is True
assert opts.permission_mode == "bypassPermissions"
assert opts.mcp_servers == {
"firefly": {"type": "http", "url": "http://127.0.0.1:1/mcp/firefly/"}
}
assert opts.session_store is not None
assert opts.resume is None
assert Path(opts.cli_path).exists()
wrapper = Path(opts.cli_path).read_text()
assert "CLAUDE_" in wrapper and "DATABASE_URL" not in wrapper
async def test_conversation_reuses_live_session(cwd: Path) -> None:
backend = _backend(cwd, InMemorySessionStore())
for i in range(2):
await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": f"turn {i}"}],
conversation_id="conv-1",
)
)
assert len(FakeClient.instances) == 1
assert FakeClient.instances[0].prompts == ["turn 0", "turn 1"]
assert backend.sessions["conv-1"]["turns"] == 2
async def test_history_is_seeded_into_store_and_resumed(cwd: Path) -> None:
store = InMemorySessionStore()
backend = _backend(cwd, store)
history = [
{"role": "user", "content": "first"},
{"role": "assistant", "content": [{"type": "text", "text": "reply"}]},
{"role": "user", "content": "second"},
]
await _drain(
backend.complete(
agent=backend.agent, messages=history, conversation_id="conv-2"
)
)
client = FakeClient.instances[0]
assert client.options.resume is not None
key = {
"project_key": project_key_for_directory(str(cwd)),
"session_id": client.options.resume,
}
entries = store.get_entries(key)
assert messages_from_entries(entries) == history[:2]
assert client.prompts == ["second"]
async def test_known_session_id_is_resumed_without_seeding(cwd: Path) -> None:
store = InMemorySessionStore()
backend = _backend(cwd, store)
history = [
{"role": "user", "content": "first"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "second"},
]
await _drain(
backend.complete(
agent=backend.agent,
messages=history,
conversation_id="conv-3",
session_id="known-sid",
)
)
assert FakeClient.instances[0].options.resume == "known-sid"
assert store.size == 0
async def test_stateless_caller_hits_same_session_next_turn(cwd: Path) -> None:
backend = _backend(cwd, InMemorySessionStore())
first = [{"role": "user", "content": "hi"}]
await _drain(backend.complete(agent=backend.agent, messages=first))
follow_up = [
*first,
{
"role": "assistant",
"content": [
{"type": "text", "text": "calling "},
{"type": "text", "text": "done"},
],
},
{"role": "user", "content": "more"},
]
assert list(backend.sessions) == [fingerprint(follow_up[:-1])]
await _drain(backend.complete(agent=backend.agent, messages=follow_up))
assert len(FakeClient.instances) == 1
assert FakeClient.instances[0].prompts == ["hi", "more"]
async def test_mcp_deny_and_usage_sink(cwd: Path) -> None:
seen: list[UsageEvent] = []
async def sink(event: UsageEvent) -> None:
seen.append(event)
backend = _backend(
cwd,
InMemorySessionStore(),
sink,
expose_mcps=(ExposedMcp(name="firefly", deny=("delete_*",)),),
)
await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": "x"}],
conversation_id="c",
)
)
assert FakeClient.instances[0].options.disallowed_tools == [
"mcp__firefly__delete_account"
]
assert len(seen) == 1
assert seen[0].usage.input_tokens == 3
assert seen[0].conversation_id == "c"
assert seen[0].session_id == "fresh-session"
async def test_skill_sets_become_sorted_plugins(cwd: Path) -> None:
sets = cwd / "skills"
for name in ("zeta", "общие"):
(sets / name / "demo").mkdir(parents=True)
(sets / name / "demo" / "SKILL.md").write_text("---\nname: demo\n---\n")
backend = _backend(
cwd, InMemorySessionStore(), skill_sets=(sets / "zeta", sets / "общие")
)
await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": "x"}],
conversation_id="c",
)
)
plugins = FakeClient.instances[0].options.plugins
assert [p["type"] for p in plugins] == ["local", "local"]
paths = [Path(p["path"]) for p in plugins]
assert [p.name for p in paths] == ["zeta", "общие"]
for path in paths:
assert (path / ".claude-plugin" / "plugin.json").exists()
assert (path / "skills" / "demo" / "SKILL.md").exists()
assert FakeClient.instances[0].options.skills == "all"
async def test_prompt_sources_are_assembled(cwd: Path) -> None:
(cwd / "a.md").write_text("alpha\n")
(cwd / "b.md").write_text("\nbeta\n\n")
backend = _backend(
cwd,
InMemorySessionStore(),
prompts=Prompts(deep=(("role", cwd / "a.md"), cwd / "b.md")),
)
await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": "x"}],
conversation_id="c",
)
)
assert (
FakeClient.instances[0].options.system_prompt
== "<role>\nalpha\n</role>\n\nbeta\n"
)
async def test_runner_user_lands_in_wrapper(cwd: Path) -> None:
import os
import pwd
me = pwd.getpwuid(os.getuid())
backend = _backend(cwd, InMemorySessionStore())
backend._runner = RunnerConfig(user=me.pw_name, home=cwd)
backend._uid, backend._gid = me.pw_uid, me.pw_gid
await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": "x"}],
conversation_id="c",
)
)
opts = FakeClient.instances[0].options
assert opts.env["HOME"] == str(cwd)
assert opts.env["CLAUDE_CONFIG_DIR"] == str(cwd / ".claude")
wrapper = Path(opts.cli_path).read_text()
assert f"UID = {me.pw_uid}" in wrapper
assert "os.setuid(UID)" in wrapper
async def test_close_disconnects(cwd: Path) -> None:
backend = _backend(cwd, InMemorySessionStore())
async with backend:
await _drain(
backend.complete(
agent=backend.agent,
messages=[{"role": "user", "content": "x"}],
conversation_id="c",
)
)
assert FakeClient.instances[0].connected is False
assert backend.sessions == {}
+697
View File
@@ -0,0 +1,697 @@
import asyncio
import tempfile
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
import pytest
from claude_agent_sdk import (
AssistantMessage,
InMemorySessionStore,
ResultMessage,
StreamEvent,
TextBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
project_key_for_directory,
)
from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions
from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend
from beaver_gateway.core.bus import EventBus
from beaver_gateway.core.conversations import Conversations, ConversationTexts, parse_at
from beaver_gateway.core.registry import AgentRegistry
from beaver_gateway.core.sessions import SessionPool
from beaver_gateway.core.transcript import (
build_entries,
close_open_tool_uses,
open_tool_uses,
render_messages,
strip_tool_entries,
window_entries,
)
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.storage import Database
from beaver_gateway.storage.models import (
Conversation,
ConversationBinding,
InjectQueueItem,
)
class StubFrontend(Frontend):
def __init__(
self,
name: str,
kinds: tuple[str, ...],
agents: dict[str, str] | None = None,
*,
home: bool = False,
) -> None:
self.name = name
self.kinds = kinds
self.agents = agents or {}
self.home = home
self.materialized: list[str] = []
self.conversations = None
def configure(self, runtime) -> None:
pass
async def serve(self) -> None:
pass
def agent_for(self, kind: str) -> str | None:
return self.agents.get(kind)
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
if not self.home:
return None
self.materialized.append(conv.external_id)
return await self.conversations.bind(
conv, frontend=self.name, external_id=f"{self.name}:{conv.external_id}"
)
class ScriptedClient:
instances: list["ScriptedClient"] = []
hold: asyncio.Event | None = None
def __init__(self, options: Any) -> None:
self.options = options
self.prompts: list[str] = []
self.session_id = options.resume or str(uuid.uuid4())
self.interrupted = False
self.connected = False
ScriptedClient.instances.append(self)
async def connect(self) -> None:
self.connected = True
async def query(self, prompt: str) -> None:
self.prompts.append(prompt)
async def receive_response(self):
prompt = self.prompts[-1]
yield StreamEvent(
uuid="u", session_id="s", event={"type": "message_start", "message": {}}
)
yield StreamEvent(
uuid="u",
session_id="s",
event={
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
)
yield StreamEvent(
uuid="u",
session_id="s",
event={
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": f"ok:{prompt}"},
},
)
yield StreamEvent(
uuid="u", session_id="s", event={"type": "content_block_stop", "index": 0}
)
yield AssistantMessage(content=[TextBlock(text=f"ok:{prompt}")], model="m")
hold = ScriptedClient.hold
if hold is not None and not self.interrupted:
await hold.wait()
cut = self.interrupted
self.interrupted = False
yield ResultMessage(
subtype="error_during_execution" if cut else "success",
duration_ms=1,
duration_api_ms=1,
is_error=cut,
num_turns=1,
session_id=self.session_id,
stop_reason="end_turn",
total_cost_usd=0.0,
usage={"input_tokens": 1, "output_tokens": 1},
)
async def interrupt(self) -> None:
self.interrupted = True
if ScriptedClient.hold is not None:
ScriptedClient.hold.set()
async def disconnect(self) -> None:
self.connected = False
class World:
def __init__(self, root: Path) -> None:
ScriptedClient.instances.clear()
ScriptedClient.hold = None
self.root = root
self.db = Database(f"sqlite:///{root / 'w.db'}")
self.store = InMemorySessionStore()
self.agent = ClaudeAgent(
name="a",
model="m",
system_prompt="hi",
cwd=root,
kinds=("master", "branch", "job", "fork"),
gateway_tools=("say",),
options=ClaudeOptions(effort="low"),
)
self.deep_agent = ClaudeAgent(
name="d", model="m", system_prompt="deep", cwd=root
)
self.markdown = StubFrontend("markdown", ("deep",), {"deep": "d"}, home=True)
self.api = StubFrontend("api", ("master", "branch", "deep", "job"))
self.pool = SessionPool(rss_limit=1 << 40, max_live=100)
self.backend, self.deep_backend = (
ClaudeSdkBackend(
agent=agent,
mcp_internal_urls={},
session_store=self.store,
client_factory=ScriptedClient,
work_dir=root / "work",
pool=self.pool,
)
for agent in (self.agent, self.deep_agent)
)
self.bus = EventBus()
self.conversations = Conversations(
db=self.db,
agents=AgentRegistry([self.agent, self.deep_agent]),
backends={"a": self.backend, "d": self.deep_backend},
bus=self.bus,
pool=self.pool,
store=self.store,
texts=ConversationTexts(),
frontends=[self.api, self.markdown],
idle_interval=3600,
)
self.markdown.conversations = self.conversations
self.api.conversations = self.conversations
async def setup(self) -> "World":
await self.db.create_all()
return self
def key(self, session_id: str) -> dict[str, str]:
return {
"project_key": project_key_for_directory(str(self.root)),
"session_id": session_id,
}
async def statuses(self, conv) -> list[tuple[str, str]]:
rows = await self.conversations.queue.recent(conv.id)
return [(r.priority, r.status) for r in reversed(rows)]
async def settle(self, conv, expected: int, timeout: float = 5.0) -> None:
deadline = asyncio.get_running_loop().time() + timeout
while asyncio.get_running_loop().time() < deadline:
rows = await self.conversations.queue.recent(conv.id)
if sum(1 for r in rows if r.status == "done") >= expected:
return
await asyncio.sleep(0.02)
msg = f"queue did not settle: {await self.statuses(conv)}"
raise AssertionError(msg)
@pytest.fixture
async def world() -> World:
root = Path(tempfile.mkdtemp(prefix="beaver-conv-"))
w = await World(root).setup()
yield w
await w.conversations.stop()
await w.pool.close_all()
await w.db.dispose()
async def test_two_messages_run_one_at_a_time_in_order(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
ScriptedClient.hold = asyncio.Event()
await world.conversations.post(conv, "first")
await asyncio.sleep(0.2)
await world.conversations.post(conv, "second")
await asyncio.sleep(0.2)
assert await world.statuses(conv) == [("user", "running"), ("user", "queued")]
live = world.pool.get(conv.external_id)
assert live is not None and live.busy
assert (await world.conversations.get(conv.external_id)).running_turn is not None
ScriptedClient.hold.set()
await world.settle(conv, 2)
assert len(ScriptedClient.instances) == 1
assert ScriptedClient.instances[0].prompts == ["first", "second"]
row = await world.conversations.get(conv.external_id)
assert row.running_turn is None
assert row.session_id == ScriptedClient.instances[0].session_id
async def test_urgent_interrupts_and_goes_first(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
ScriptedClient.hold = asyncio.Event()
await world.conversations.post(conv, "first")
await asyncio.sleep(0.2)
await world.conversations.post(conv, "second")
await world.conversations.inject(conv, "ALERT", urgency="urgent", origin="крон")
await world.settle(conv, 2)
client = ScriptedClient.instances[0]
assert await world.statuses(conv) == [
("user", "interrupted"),
("user", "done"),
("urgent", "done"),
]
assert await world.statuses(conv) == [
("user", "interrupted"),
("user", "done"),
("urgent", "done"),
]
assert client.prompts[0] == "first"
assert client.prompts[1].startswith("[инжект: крон")
assert client.prompts[1].endswith("ALERT")
assert client.prompts[2] == "second"
async def test_normal_injects_ride_with_the_next_user_message(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
await world.conversations.inject(
conv, "vault changed", urgency="normal", origin="watch"
)
await asyncio.sleep(0.2)
assert await world.statuses(conv) == [("normal", "queued")]
await world.conversations.post(conv, "hello")
await world.settle(conv, 2)
prompts = ScriptedClient.instances[0].prompts
assert len(prompts) == 1
assert prompts[0].startswith("hello\n\n[инжекты")
assert "- [watch] vault changed" in prompts[0]
async def test_inject_turn_reply_is_not_routed(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
seen: list[dict[str, Any]] = []
async def collect() -> None:
async for event in world.bus.stream(conversation_id=conv.external_id):
seen.append(event)
task = asyncio.create_task(collect())
world.conversations._normal_window = 0.05
await world.conversations.inject(conv, "tick", urgency="normal", origin="крон")
await world.settle(conv, 1)
await asyncio.sleep(0.05)
task.cancel()
types = [e["type"] for e in seen]
assert "turn.start" in types and "turn.end" in types
assert "reply" not in types
assert all(e.get("origin") == "inject" for e in seen if e["type"] == "turn.start")
async def test_spawn_seeds_first_user_message(world: World) -> None:
conv = await world.conversations.spawn(
kind="branch", agent="a", seed="brief", text="do X", title="t"
)
await world.settle(conv, 1)
prompt = ScriptedClient.instances[0].prompts[0]
assert prompt.startswith("[сид: brief] branch «t», ")
assert prompt.endswith("\n\ndo X")
assert ScriptedClient.instances[0].options.system_prompt == "hi"
async def test_fork_leaves_original_untouched(world: World) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": "one"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": "x"}],
},
{"role": "assistant", "content": "done one"},
{"role": "user", "content": "two"},
{"role": "assistant", "content": "done two"},
{"role": "user", "content": "three"},
{"role": "assistant", "content": "done three"},
]
entries = build_entries(history, session_id=sid, cwd=str(world.root), model="m")
await world.store.append(world.key(sid), entries)
conv = await world.conversations.create(
kind="branch", agent="a", origin="test", session_id=sid
)
before = [dict(e) for e in await world.store.load(world.key(sid))]
result = await world.conversations.fork(
conv, "summarize", window=2, strip_tools=True
)
after = await world.store.load(world.key(sid))
assert after == before
assert result.text == "ok:summarize"
assert result.conversation.status == "closed"
assert result.conversation.parent_id == conv.id
fork_entries = await world.store.load(world.key(result.conversation.session_id))
assert not {e["uuid"] for e in fork_entries} & {e["uuid"] for e in before}
texts = [e["message"]["content"] for e in fork_entries]
assert texts[0] == "two"
assert len(fork_entries) == 4
assert fork_entries[0]["parentUuid"] is None
assert all(
b["type"] == "text"
for e in fork_entries[1:]
for b in e["message"]["content"]
if isinstance(e["message"]["content"], list)
)
client = ScriptedClient.instances[0]
assert client.options.resume == result.conversation.session_id
assert client.options.mcp_servers == {}
assert world.pool.get(result.conversation.external_id) is None
async def test_copy_seed_forks_parent_with_window(world: World) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": f"q{i}"}
if i % 2 == 0
else {"role": "assistant", "content": f"a{i}"}
for i in range(8)
]
await world.store.append(
world.key(sid),
build_entries(history, session_id=sid, cwd=str(world.root), model="m"),
)
parent = await world.conversations.create(
kind="master", agent="a", origin="test", session_id=sid
)
child = await world.conversations.spawn(
kind="branch", agent="a", seed="copy", parent=parent, window=1
)
assert child.session_id and child.session_id != sid
copied = await world.store.load(world.key(child.session_id))
assert [e["message"]["content"] for e in copied] == [
"q6",
[{"type": "text", "text": "a7"}],
]
await world.settle(child, 1)
assert ScriptedClient.instances[0].options.resume == child.session_id
async def test_merge_injects_summary_into_parent(world: World) -> None:
sid = str(uuid.uuid4())
await world.store.append(
world.key(sid),
build_entries(
[{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}],
session_id=sid,
cwd=str(world.root),
model="m",
),
)
master = await world.conversations.create(kind="master", agent="a", origin="test")
branch = await world.conversations.create(
kind="branch", agent="a", origin="test", parent=master, session_id=sid
)
result = await world.conversations.merge(branch)
assert result.text.startswith("ok:")
assert (await world.conversations.get(branch.external_id)).status == "merged"
assert await world.statuses(master) == [("normal", "queued")]
item = (await world.conversations.queue.recent(master.id))[0]
assert item.origin == "слив" and item.text == result.text
async def test_recover_closes_open_tool_use_and_injects_interrupted(
world: World,
) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": "go"},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "t9",
"name": "Bash",
"input": {"command": "sleep"},
}
],
},
]
await world.store.append(
world.key(sid),
build_entries(history, session_id=sid, cwd=str(world.root), model="m"),
)
conv = await world.conversations.create(
kind="master", agent="a", origin="test", session_id=sid
)
async with world.db.session() as session:
row = await session.get(type(conv), conv.id)
row.running_turn = "turn_dead"
session.add(row)
session.add(
InjectQueueItem(
conversation_id=conv.id,
priority="user",
origin="user",
text="go",
status="running",
turn_id="turn_dead",
)
)
await session.commit()
cut = await world.conversations.recover()
assert [c.external_id for c in cut] == [conv.external_id]
entries = await world.store.load(world.key(sid))
assert not open_tool_uses(entries)
tail = entries[-1]
assert tail["type"] == "user"
assert tail["message"]["content"][0] == {
"type": "tool_result",
"tool_use_id": "t9",
"content": "прервано",
"is_error": True,
}
assert tail["parentUuid"] == entries[-2]["uuid"]
assert (await world.conversations.get(conv.external_id)).running_turn is None
assert await world.statuses(conv) == [("user", "interrupted"), ("normal", "queued")]
note = (await world.conversations.queue.recent(conv.id))[0]
assert (
"turn_dead" in note.text
and "оборван" in note.text
and "1 незакрытых" in note.text
)
await asyncio.sleep(0.2)
assert ScriptedClient.instances == []
async def test_read_and_bindings(world: World) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": "q1"},
{
"role": "assistant",
"content": [
{"type": "text", "text": "a1"},
{"type": "tool_use", "id": "t", "name": "Read", "input": {}},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}],
},
{"role": "assistant", "content": "a1b"},
{"role": "user", "content": "q2"},
{"role": "assistant", "content": "a2"},
]
await world.store.append(
world.key(sid),
build_entries(history, session_id=sid, cwd=str(world.root), model="m"),
)
conv = await world.conversations.create(
kind="deep", agent="d", origin="test", session_id=sid
)
assert (
await world.conversations.read(conv)
== "user:\nq1\n\nassistant:\na1\n\na1b\n(tools: Read)\n\nuser:\nq2\n\nassistant:\na2"
)
assert (
await world.conversations.read(conv, window=1) == "user:\nq2\n\nassistant:\na2"
)
await world.conversations.bind(conv, frontend="markdown", external_id="a.md")
await world.conversations.bind(conv, frontend="markdown", external_id="b.md")
bindings = await world.conversations.bindings(conv)
assert [(b.external_id, b.visible) for b in bindings] == [
("a.md", False),
("b.md", True),
]
found = await world.conversations.find_bound(
frontend="markdown", external_id="b.md"
)
assert found is not None and found.id == conv.id
assert (
await world.conversations.find_bound(frontend="markdown", external_id="a.md")
is None
)
async def test_schedule_rows_and_parse_at(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
row = await world.conversations.schedule(conv, "+15m", "push X")
delta = (row.execute_at.replace(tzinfo=UTC) - datetime.now(UTC)).total_seconds()
assert 14 * 60 < delta <= 15 * 60
assert [s.text for s in await world.conversations.schedules(conv)] == ["push X"]
assert parse_at("2026-09-01T10:00:00+02:00") == datetime(
2026, 9, 1, 8, 0, tzinfo=UTC
)
with pytest.raises(ValueError, match="Invalid isoformat"):
parse_at("tomorrow")
def test_transcript_helpers() -> None:
history = [
{"role": "user", "content": "q1"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "hm", "signature": "s"},
{"type": "tool_use", "id": "t", "name": "Read", "input": {}},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}],
},
{"role": "assistant", "content": "a1"},
{"role": "user", "content": "q2"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "t2", "name": "Bash", "input": {}}],
},
]
entries = build_entries(history, session_id="s", cwd="/v", model="m")
assert [b["id"] for _, b in open_tool_uses(entries)] == ["t2"]
fixes = close_open_tool_uses(entries)
assert len(fixes) == 1 and fixes[0]["sessionId"] == "s" and fixes[0]["cwd"] == "/v"
assert not open_tool_uses([*entries, *fixes])
stripped = strip_tool_entries(entries)
assert [e["message"]["content"] for e in stripped] == [
"q1",
[{"type": "text", "text": "a1"}],
"q2",
]
assert (
stripped[0]["parentUuid"] is None
and stripped[1]["parentUuid"] == stripped[0]["uuid"]
)
windowed = window_entries(entries, window=1)
assert windowed[0]["message"]["content"] == "q2" and len(windowed) == 2
assert (
render_messages(
[{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}]
)
== "user:\nq\n\nassistant:\na"
)
async def test_agent_kinds_gate_create(world: World) -> None:
with pytest.raises(ValueError, match="does not serve kind 'deep'"):
await world.conversations.create(kind="deep", agent="a", origin="test")
with pytest.raises(ValueError, match="does not serve kind 'master'"):
await world.conversations.create(kind="master", agent="d", origin="test")
with pytest.raises(ValueError, match="'d' does not serve kind 'master'"):
await world.conversations.spawn(kind="master", seed="clean", agent="d")
conv = await world.conversations.create(kind="deep", agent="d", origin="test")
assert conv.agent_name == "d"
async def test_frontend_kinds_gate_bind(world: World) -> None:
master = await world.conversations.create(kind="master", agent="a", origin="test")
with pytest.raises(ValueError, match="'markdown' does not show kind 'master'"):
await world.conversations.bind(master, frontend="markdown", external_id="m.md")
with pytest.raises(ValueError, match="unknown frontend 'telegram'"):
await world.conversations.bind(master, frontend="telegram", external_id="1")
await world.conversations.bind(master, frontend="api", external_id="x")
async def test_spawn_defaults_agent_and_materializes(world: World) -> None:
master = await world.conversations.create(kind="master", agent="a", origin="test")
deep = await world.conversations.spawn(
kind="deep", seed="brief", text="dig", title="t", parent=master, origin="mcp"
)
assert deep.agent_name == "d"
assert world.markdown.materialized == [deep.external_id]
bindings = await world.conversations.bindings(deep)
assert [(b.frontend, b.external_id) for b in bindings] == [
("markdown", f"markdown:{deep.external_id}")
]
branch = await world.conversations.spawn(kind="branch", seed="clean", parent=master)
assert branch.agent_name == "a"
assert world.markdown.materialized == [deep.external_id]
with pytest.raises(ValueError, match="no default agent for kind 'job'"):
await world.conversations.spawn(kind="job", seed="clean")
await world.settle(deep, 1)
def test_agent_kinds_follow_prompts(tmp_path: Path) -> None:
from beaver_gateway.agents.claude import Prompts
plain = ClaudeAgent(name="p", model="m", system_prompt="hi", cwd=tmp_path)
assert plain.kinds == ("deep",)
dispatcher = ClaudeAgent(
name="x", model="m", cwd=tmp_path, prompts=Prompts(master=("a.md",), fork=())
)
assert dispatcher.kinds == ("master", "fork")
assert dispatcher.prompt_for("master") == ("a.md",)
assert dispatcher.prompt_for("deep") is None
with pytest.raises(ValueError, match=r"serves \['deep'\] without a prompt"):
ClaudeAgent(
name="y",
model="m",
cwd=tmp_path,
kinds=("master", "deep"),
prompts=Prompts(master=()),
)
async def test_observer_publishes_tool_results_for_the_panel(world: World) -> None:
seen: list[dict[str, Any]] = []
async def collect() -> None:
async for event in world.bus.stream(conversation_id="c1"):
seen.append(event)
task = asyncio.create_task(collect())
await asyncio.sleep(0)
observe = world.conversations._observer("c1", "t1", "user")
observe(
AssistantMessage(
content=[ToolUseBlock(id="toolu_1", name="Bash", input={"command": "ls"})],
model="m",
parent_tool_use_id="toolu_0",
)
)
observe(
UserMessage(
content=[
ToolResultBlock(tool_use_id="toolu_1", content="a\nb", is_error=False)
],
parent_tool_use_id="toolu_0",
)
)
observe(UserMessage(content="plain text, no tool result"))
await asyncio.sleep(0.05)
task.cancel()
assert [e["type"] for e in seen] == ["tool", "tool.result"]
tool, result = seen
assert tool["tool_use_id"] == "toolu_1" and tool["parent_tool_use_id"] == "toolu_0"
assert result["tool_use_id"] == "toolu_1"
assert result["parent_tool_use_id"] == "toolu_0"
assert result["turn_id"] == "t1" and result["is_error"] is False
assert result["content"] == "a\nb"
+61
View File
@@ -0,0 +1,61 @@
import sqlite3
import tempfile
from pathlib import Path
from sqlmodel import select
from beaver_gateway.storage import Database
from beaver_gateway.storage.models import Conversation
async def test_create_all_adds_missing_columns() -> None:
path = Path(tempfile.mkdtemp(prefix="beaver-migrate-")) / "old.db"
raw = sqlite3.connect(path)
raw.execute(
"CREATE TABLE conversations (id INTEGER PRIMARY KEY, frontend VARCHAR NOT NULL, "
"external_id VARCHAR NOT NULL, agent_name VARCHAR NOT NULL, "
"created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)"
)
raw.execute(
"INSERT INTO conversations VALUES (1, 'markdown', 'x', 'a', '2026-01-01', '2026-01-01')"
)
raw.commit()
raw.close()
db = Database(f"sqlite:///{path}")
await db.create_all()
async with db.session() as session:
conv = (await session.exec(select(Conversation))).one()
assert conv.session_id is None
conv.session_id = "sid"
session.add(conv)
await session.commit()
await db.dispose()
async def test_create_all_backfills_defaults_for_old_rows() -> None:
path = Path(tempfile.mkdtemp(prefix="beaver-migrate-")) / "old.db"
raw = sqlite3.connect(path)
raw.execute(
"CREATE TABLE conversations (id INTEGER PRIMARY KEY, frontend VARCHAR NOT NULL, "
"external_id VARCHAR NOT NULL, agent_name VARCHAR NOT NULL, session_id VARCHAR, "
"created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)"
)
raw.execute(
"INSERT INTO conversations VALUES (1, 'markdown', 'x', 'a', NULL, '2026-01-01', '2026-01-01')"
)
raw.commit()
raw.close()
db = Database(f"sqlite:///{path}")
await db.create_all()
async with db.session() as session:
conv = (await session.exec(select(Conversation))).one()
assert (conv.kind, conv.status, conv.pending_question, conv.flags) == (
"deep",
"open",
False,
{},
)
assert conv.running_turn is None
await db.dispose()
+228
View File
@@ -0,0 +1,228 @@
import asyncio
import tempfile
from pathlib import Path
import frontmatter
import httpx
import pytest
from beaver_gateway.core.auth import TokenStore
from beaver_gateway.core.conversation_store import load_messages
from beaver_gateway.core.gateway_tools import _tools
from beaver_gateway.core.registry import McpRegistry
from beaver_gateway.frontends.anthropic import AnthropicMessagesFrontend
from beaver_gateway.frontends.api import ApiFrontend
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.frontends.markdown import MarkdownFrontend
from test_conversations import ScriptedClient, World
AUTH = {"Authorization": "Bearer tok"}
class Stack:
def __init__(self, world: World) -> None:
self.world = world
self.vault = world.root / "vault"
self.api = ApiFrontend(master_agent="a")
self.markdown = MarkdownFrontend(vault_path=self.vault, default_agent="d")
self.anthropic = AnthropicMessagesFrontend()
frontends = [self.api, self.anthropic, self.markdown]
world.conversations._frontends = frontends
self.runtime = GatewayRuntime(
agents=world.conversations._agents,
mcps=McpRegistry([]),
backends=world.conversations._backends,
token_store=TokenStore(bootstrap={"t": "tok"}),
db=world.db,
conversations=world.conversations,
bus=world.bus,
pool=world.pool,
frontends=tuple(frontends),
)
for fe in frontends:
fe.configure(self.runtime)
self.mirror = asyncio.create_task(self.markdown.mirror.run())
def client(self, fe) -> httpx.AsyncClient:
return httpx.AsyncClient(
transport=httpx.ASGITransport(app=fe._app), base_url="http://t"
)
async def close(self) -> None:
self.mirror.cancel()
async def file_of(self, conv_id: str) -> Path:
conv = await self.world.conversations.get(conv_id)
path = await self.markdown.mirror.bound_path(conv)
assert path is not None
return path
async def wait_file(self, path: Path, needle: str, timeout: float = 5.0) -> str:
deadline = asyncio.get_running_loop().time() + timeout
while asyncio.get_running_loop().time() < deadline:
if path.exists() and needle in path.read_text(encoding="utf-8"):
return path.read_text(encoding="utf-8")
await asyncio.sleep(0.02)
msg = f"{needle!r} never showed up in {path}"
raise AssertionError(msg)
@pytest.fixture
async def stack() -> Stack:
root = Path(tempfile.mkdtemp(prefix="beaver-routing-"))
world = await World(root).setup()
await world.conversations.start()
s = Stack(world)
yield s
await s.close()
await world.conversations.stop()
await world.pool.close_all()
await world.db.dispose()
async def test_api_rejects_deep_with_dispatcher(stack: Stack) -> None:
async with stack.client(stack.api) as c:
r = await c.post(
"/api/conversations", json={"kind": "deep", "agent": "a"}, headers=AUTH
)
assert r.status_code == 400
assert "does not serve kind 'deep'" in r.json()["error"]
r = await c.post("/api/conversations", json={"kind": "job"}, headers=AUTH)
assert r.status_code == 400
assert "no default agent" in r.json()["error"]
r = await c.get("/api/agents", headers=AUTH)
agents = {a["name"]: a["kinds"] for a in r.json()["agents"]}
assert agents == {"a": ["master", "branch", "job", "fork"], "d": ["deep"]}
homes = {f["name"]: f["default_agents"] for f in r.json()["frontends"]}
assert homes == {
"api": {"master": "a"},
"anthropic": {},
"markdown": {"deep": "d"},
}
async def test_api_spawn_deep_lands_in_vault(stack: Stack) -> None:
async with stack.client(stack.api) as c:
r = await c.post("/api/conversations", json={"kind": "master"}, headers=AUTH)
assert r.status_code == 201 and r.json()["agent"] == "a"
master = r.json()["id"]
r = await c.post(
"/api/conversations",
json={"kind": "deep", "seed": "brief", "text": "dig", "title": "Тема"},
headers=AUTH,
)
assert r.status_code == 201, r.text
body = r.json()
assert body["agent"] == "d"
assert [b["frontend"] for b in body["bindings"]] == ["markdown"]
rel = body["bindings"][0]["external_id"]
assert rel.endswith("_Тема.md") and rel.startswith("_logs/d/")
r = await c.post(
f"/api/conversations/{master}/bind",
json={"frontend": "markdown", "external_id": "x.md"},
headers=AUTH,
)
assert r.status_code == 400
path = stack.vault / rel
text = await stack.wait_file(path, "ok:[сид: brief]")
post = frontmatter.loads(text)
assert post.metadata == {"agent": "d", "conversation_id": body["id"]}
assert post.content.startswith("### User:\n\n[сид: brief] deep «Тема», ")
assert post.content.rstrip().endswith("### User:")
async def test_markdown_chat_rejects_dispatcher(stack: Stack) -> None:
async with stack.client(stack.markdown) as c:
r = await c.post(
"/chat",
json={"filename": "x.md", "agent": "a", "content": "### User:\n\nhi\n"},
headers=AUTH,
)
assert r.status_code == 400
assert "does not serve kind 'deep'" in r.json()["detail"]
async def test_spawn_tool_reports_bad_pair(stack: Stack) -> None:
master = await stack.world.conversations.create(
kind="master", agent="a", origin="test"
)
spawn = next(
t
for t in _tools(stack.world.conversations, master.external_id)
if t.name == "spawn"
)
out = await spawn.handler({"kind": "deep", "agent": "a"})
assert (
out.get("is_error")
and "does not serve kind 'deep'" in out["content"][0]["text"]
)
out = await spawn.handler({"kind": "deep", "seed": "brief", "text": "go"})
assert not out.get("is_error")
conv_id = out["content"][0]["text"].split()[-1]
path = await stack.file_of(conv_id)
assert path.exists()
async def test_anthropic_turns_become_one_deep_conversation(stack: Stack) -> None:
async with stack.client(stack.anthropic) as c:
r = await c.post(
"/v1/messages",
json={"model": "a", "messages": [{"role": "user", "content": "hi"}]},
headers=AUTH,
)
assert r.status_code == 400
first = [{"role": "user", "content": "first question"}]
r = await c.post(
"/v1/messages", json={"model": "d", "messages": first}, headers=AUTH
)
assert r.status_code == 200, r.text
reply = r.json()["content"][0]["text"]
assert reply == "ok:first question"
convs = await stack.world.conversations.find(kind="deep")
assert len(convs) == 1
conv = convs[0]
assert conv.title == "first question" and conv.frontend == "anthropic"
path = await stack.file_of(conv.external_id)
text = await stack.wait_file(path, "ok:first question")
assert frontmatter.loads(text).metadata == {
"agent": "d",
"conversation_id": conv.external_id,
}
second = [
*first,
{"role": "assistant", "content": reply},
{"role": "user", "content": "second"},
]
r = await c.post(
"/v1/messages",
json={"model": "d", "messages": second, "stream": True},
headers=AUTH,
)
assert r.status_code == 200
assert "ok:second" in r.text
assert len(await stack.world.conversations.find(kind="deep")) == 1
assert len(ScriptedClient.instances) == 1
assert ScriptedClient.instances[0].prompts == ["first question", "second"]
text = await stack.wait_file(path, "ok:second")
assert text.count("### User:") == 3 and text.count("### Assistant:") == 2
bindings = await stack.world.conversations.bindings(conv)
assert [b.frontend for b in bindings if b.visible] == ["markdown", "anthropic"]
async with stack.world.db.session() as session:
stored = await load_messages(session, conversation_id=conv.id)
assert [m["role"] for m in stored] == ["user", "assistant", "user", "assistant"]
assert stored[-1]["content"] == [{"type": "text", "text": "ok:second"}]
r = await c.post(
"/v1/messages",
json={
"model": "a",
"messages": [
*second,
{"role": "assistant", "content": "ok:second"},
{"role": "user", "content": "x"},
],
},
headers=AUTH,
)
assert r.status_code == 400 and "runs on 'd'" in r.json()["detail"]
+64
View File
@@ -0,0 +1,64 @@
import os
import tempfile
from pathlib import Path
import pytest
from claude_agent_sdk.testing import run_session_store_conformance
from sqlalchemy import delete
from beaver_gateway.storage import Database, PostgresSessionStore, TranscriptEntry
POSTGRES_URL = os.environ.get("BEAVER_TEST_DATABASE_URL")
async def test_sqlite_conformance() -> None:
root = Path(tempfile.mkdtemp(prefix="beaver-store-"))
counter = 0
async def make_store() -> PostgresSessionStore:
nonlocal counter
counter += 1
db = Database(f"sqlite:///{root / f'{counter}.db'}")
await db.create_all()
return PostgresSessionStore(db)
await run_session_store_conformance(make_store)
@pytest.mark.skipif(POSTGRES_URL is None, reason="BEAVER_TEST_DATABASE_URL not set")
async def test_postgres_conformance() -> None:
assert POSTGRES_URL is not None
db = Database(POSTGRES_URL)
await db.create_all()
async def make_store() -> PostgresSessionStore:
async with db.session() as session:
await session.execute(delete(TranscriptEntry))
await session.commit()
return PostgresSessionStore(db)
try:
await run_session_store_conformance(make_store)
finally:
await db.dispose()
async def test_append_dedups_by_uuid_and_keeps_floats() -> None:
db = Database(f"sqlite:///{tempfile.mkdtemp(prefix='beaver-store-')}/d.db")
await db.create_all()
store = PostgresSessionStore(db)
key = {"project_key": "p", "session_id": "s"}
batch = [
{"type": "x", "uuid": "a", "n": 1.5},
{"type": "y", "n": 2},
{"type": "x", "uuid": "a", "n": 999},
]
await store.append(key, batch)
await store.append(key, batch)
loaded = await store.load(key)
assert loaded == [
{"type": "x", "uuid": "a", "n": 1.5},
{"type": "y", "n": 2},
{"type": "y", "n": 2},
]
assert isinstance(loaded[0]["n"], float)
+129
View File
@@ -0,0 +1,129 @@
from beaver_gateway.core.transcript import (
CLI_VERSION,
build_entries,
messages_from_entries,
)
HISTORY = [
{"role": "user", "content": "Write бобёр to notes.txt"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "hm", "signature": "sig"},
{
"type": "tool_use",
"id": "toolu_1",
"name": "Write",
"input": {"file_path": "notes.txt"},
},
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_1",
"content": "ok",
"is_error": False,
}
],
},
{"role": "assistant", "content": [{"type": "text", "text": "Done."}]},
{"role": "user", "content": [{"type": "text", "text": "thanks"}]},
{"role": "assistant", "content": "np"},
]
def test_round_trip() -> None:
entries = build_entries(HISTORY, session_id="sid", cwd="/vault", model="claude-x")
assert messages_from_entries(entries) == [
*HISTORY[:5],
{"role": "assistant", "content": [{"type": "text", "text": "np"}]},
]
def test_chain_and_shape() -> None:
entries = build_entries(HISTORY, session_id="sid", cwd="/vault", model="claude-x")
assert [e["type"] for e in entries] == [
"user",
"assistant",
"assistant",
"user",
"assistant",
"user",
"assistant",
]
assert entries[0]["parentUuid"] is None
for prev, cur in zip(entries, entries[1:], strict=False):
assert cur["parentUuid"] == prev["uuid"]
assert len({e["uuid"] for e in entries}) == len(entries)
user = entries[0]
assert user["promptSource"] == "sdk"
assert user["entrypoint"] == "sdk-py"
assert user["permissionMode"] == "bypassPermissions"
assert user["sessionId"] == "sid"
assert user["cwd"] == "/vault"
assert user["version"] == CLI_VERSION
thinking, tool_use = entries[1], entries[2]
assert thinking["message"]["id"] == tool_use["message"]["id"]
assert thinking["requestId"] == tool_use["requestId"]
assert tool_use["message"]["stop_reason"] == "tool_use"
assert tool_use["message"]["content"] == [
{
"type": "tool_use",
"id": "toolu_1",
"name": "Write",
"input": {"file_path": "notes.txt"},
}
]
result = entries[3]
assert result["sourceToolAssistantUUID"] == tool_use["uuid"]
assert result["promptId"] == user["promptId"]
assert result["toolUseResult"] == "ok"
assert result["message"]["content"][0]["is_error"] is False
final = entries[4]
assert final["message"]["stop_reason"] == "end_turn"
assert final["message"]["usage"]["input_tokens"] == 0
assert entries[5]["promptId"] != user["promptId"]
def test_tool_results_split_and_merged_back() -> None:
history = [
{"role": "user", "content": "go"},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": "a", "name": "Read", "input": {}},
{"type": "tool_use", "id": "b", "name": "Read", "input": {}},
],
},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "a", "content": "1"},
{
"type": "tool_result",
"tool_use_id": "b",
"content": [{"type": "text", "text": "2"}],
},
],
},
]
entries = build_entries(history, session_id="s", cwd="/", model="m")
assert [e["type"] for e in entries] == [
"user",
"assistant",
"assistant",
"user",
"user",
]
assert entries[3]["parentUuid"] == entries[1]["uuid"]
assert entries[4]["parentUuid"] == entries[2]["uuid"]
assert entries[4]["toolUseResult"] == "2"
assert messages_from_entries(entries) == history
Generated
+46 -42
View File
@@ -267,11 +267,13 @@ dependencies = [
{ name = "anthropic" },
{ name = "anyio" },
{ name = "argon2-cffi" },
{ name = "claude-agent-sdk" },
{ name = "fastapi" },
{ name = "fastmcp" },
{ name = "greenlet" },
{ name = "itsdangerous" },
{ name = "jinja2" },
{ name = "psutil" },
{ name = "psycopg", extra = ["binary"] },
{ name = "pydantic" },
{ name = "pydantic-settings" },
@@ -283,11 +285,9 @@ dependencies = [
[package.optional-dependencies]
local = [
{ name = "claude-code-api", version = "0.1.0", source = { editable = "../claude-code-api" } },
{ name = "raycast-api", version = "0.1.0", source = { editable = "../raycast-api" } },
]
prod = [
{ name = "claude-code-api", version = "0.1.0", source = { git = "https://git.kotikot.com/beaver/claude-code-api.git#b0d7c6d32d4de04862e109002459ec0c47f01b2f" } },
{ name = "raycast-api", version = "0.1.0", source = { git = "https://git.kotikot.com/beaver/raycast-api.git#e73894c8e435da5c0709f92f69f11bcd0dab9afe" } },
]
@@ -308,13 +308,13 @@ requires-dist = [
{ name = "anthropic", specifier = ">=0.103.0" },
{ name = "anyio", specifier = ">=4.13.0" },
{ name = "argon2-cffi", specifier = ">=25.1.0" },
{ name = "claude-code-api", marker = "extra == 'local'", editable = "../claude-code-api" },
{ name = "claude-code-api", marker = "extra == 'prod'", git = "https://git.kotikot.com/beaver/claude-code-api.git" },
{ name = "claude-agent-sdk", specifier = ">=0.2.146" },
{ name = "fastapi", specifier = ">=0.136.1" },
{ name = "fastmcp", specifier = ">=3.3.1" },
{ name = "greenlet", specifier = ">=3.5.0" },
{ name = "itsdangerous", specifier = ">=2.2.0" },
{ name = "jinja2", specifier = ">=3.1.6" },
{ name = "psutil", specifier = ">=7.2.2" },
{ name = "psycopg", extras = ["binary"], specifier = ">=3.3.4" },
{ name = "pydantic", specifier = ">=2.13.4" },
{ name = "pydantic-settings", specifier = ">=2.14.1" },
@@ -417,37 +417,22 @@ wheels = [
]
[[package]]
name = "claude-code-api"
version = "0.1.0"
source = { git = "https://git.kotikot.com/beaver/claude-code-api.git#b0d7c6d32d4de04862e109002459ec0c47f01b2f" }
resolution-markers = [
"python_full_version >= '3.14'",
"python_full_version < '3.14'",
]
name = "claude-agent-sdk"
version = "0.2.146"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "ptyprocess" },
{ name = "anyio" },
{ name = "jsonschema" },
{ name = "mcp" },
{ name = "sniffio" },
]
[[package]]
name = "claude-code-api"
version = "0.1.0"
source = { editable = "../claude-code-api" }
resolution-markers = [
"python_full_version >= '3.14'",
"python_full_version < '3.14'",
]
dependencies = [
{ name = "ptyprocess" },
]
[package.metadata]
requires-dist = [{ name = "ptyprocess", specifier = ">=0.7" }]
[package.metadata.requires-dev]
dev = [
{ name = "pre-commit", specifier = ">=4.0" },
{ name = "pytest", specifier = ">=8" },
{ name = "pytest-asyncio", specifier = ">=0.23" },
sdist = { url = "https://files.pythonhosted.org/packages/46/08/a79f73aad2156f4518ae6e0fffbcdb403cd5347d62b994b005a0c9bb12fb/claude_agent_sdk-0.2.146.tar.gz", hash = "sha256:9e402548466018421cfd0725a7601e14d9e3e879276a6a4a877c5a2ff6633546", size = 344704, upload-time = "2026-08-27T22:25:46.687Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e4/e2/7618d7ffd3d4580d53d1147e7baba005c3bf4359d8bb8ab9beab62b5a0ec/claude_agent_sdk-0.2.146-py3-none-macosx_11_0_arm64.whl", hash = "sha256:7bcb91dd34c2e0e5b4f006648d457bbfba9aa1e44cf2e5d02abb441dcbbaff98", size = 84087716, upload-time = "2026-08-27T22:25:50.846Z" },
{ url = "https://files.pythonhosted.org/packages/1f/e6/004b957547e32f6df4948c3a5a3c06974c6c19539c357ebe40a1a23941d3/claude_agent_sdk-0.2.146-py3-none-macosx_11_0_x86_64.whl", hash = "sha256:9285b3dd900376b045a06d2d6f5e19bfd96a3f6cb88d21166bb840aff14ee505", size = 88567306, upload-time = "2026-08-27T22:25:54.876Z" },
{ url = "https://files.pythonhosted.org/packages/37/fa/1c1cb847ff8a92f6ec9a0176a984ec1902cf6eab6c027fd45059d211a96d/claude_agent_sdk-0.2.146-py3-none-manylinux_2_17_aarch64.whl", hash = "sha256:ad536806e6c6fc1972ae3950b1bb171f1494e6f47002c24dbc68f19af1c34997", size = 94195787, upload-time = "2026-08-27T22:25:59.367Z" },
{ url = "https://files.pythonhosted.org/packages/60/f9/48aa74e5a85ebe11b916a99b053cc64f47bea8409ec2b9ada5b5a1169a7a/claude_agent_sdk-0.2.146-py3-none-manylinux_2_17_x86_64.whl", hash = "sha256:1f2bae0759f3962be2714c786c1a6961824c5ce3d4b7f44f0828191bd44ccebe", size = 94385743, upload-time = "2026-08-27T22:26:03.218Z" },
{ url = "https://files.pythonhosted.org/packages/0d/d4/ff986dc35330146cba60a1bb6ee3f521cdec25ccc0306ef359b9d0cb59c4/claude_agent_sdk-0.2.146-py3-none-win_amd64.whl", hash = "sha256:dc2c4773dc02492b52c74f957a059ecbc0aec7063bd1a16c039b78613194b3cf", size = 97049829, upload-time = "2026-08-27T22:26:08.701Z" },
]
[[package]]
@@ -1403,6 +1388,34 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" },
]
[[package]]
name = "psutil"
version = "7.2.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" },
{ url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" },
{ url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" },
{ url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" },
{ url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" },
{ url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" },
{ url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" },
{ url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" },
{ url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" },
{ url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" },
{ url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" },
{ url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" },
{ url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" },
{ url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" },
{ url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" },
{ url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" },
{ url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" },
{ url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" },
{ url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" },
{ url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" },
]
[[package]]
name = "psycopg"
version = "3.3.4"
@@ -1449,15 +1462,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/eb/e6/5fff07a70d1f945ed90ae131c3bd76cab32beff7c58c6db15ad5820b6d1f/psycopg_binary-3.3.4-cp314-cp314-win_amd64.whl", hash = "sha256:c37e024c07308cd06cf3ec51bfd0e7f6157585a4d84d1bce4a7f5f7913719bf8", size = 3666849, upload-time = "2026-05-01T23:31:51.165Z" },
]
[[package]]
name = "ptyprocess"
version = "0.7.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" },
]
[[package]]
name = "py-key-value-aio"
version = "0.4.4"