Compare commits
6
Commits
ee2dc918ee
...
fd8cbacd7b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fd8cbacd7b | ||
|
|
bd2d2368e1 | ||
|
|
1d8d65b69a | ||
|
|
33ccc78fec | ||
|
|
50b7057fa4 | ||
|
|
e1f242a87a |
No files matched your search
@@ -11,6 +11,7 @@ wheels/
|
||||
|
||||
# Sloppy
|
||||
docs/
|
||||
t/
|
||||
|
||||
# IDE
|
||||
.idea
|
||||
|
||||
+13
-17
@@ -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
@@ -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
@@ -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"]
|
||||
@@ -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):
|
||||
|
||||
@@ -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,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.
|
||||
"""
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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 ()
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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
@@ -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}
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
@@ -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)."""
|
||||
@@ -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
|
||||
@@ -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 ""
|
||||
@@ -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."""
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
@@ -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``.
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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 ""
|
||||
@@ -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 == {}
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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"
|
||||
|
||||
Reference in new issue
Block a user