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

This commit is contained in:
hh
2026-08-28 01:56:39 +02:00
parent ee2dc918ee
commit e1f242a87a
28 changed files with 2154 additions and 875 deletions
+1 -1
View File
@@ -1,6 +1,6 @@
"""Backend adapters.
Each backend wraps a provider-specific SDK (``raycast-api``, ``claude-code-api``)
Each backend wraps a provider SDK (``raycast-api``, ``claude-agent-sdk``)
and yields the unified :class:`~beaver_gateway.core.events.MessageStreamEvent`
family. The Anthropic-style frontend serialises events straight to SSE.
"""
+9 -2
View File
@@ -5,8 +5,15 @@ into a stream of :class:`~beaver_gateway.core.events.MessageStreamEvent`
records. The frontend serializes whatever comes out straight to SSE, so
backends are the only place where provider quirks are translated.
Implementations are plain :class:`typing.Protocol` conformers no ABC
subclassing to keep them swappable in tests with bare async generators.
Implementations are plain :class:`typing.Protocol` conformers - no ABC
subclassing - to keep them swappable in tests with bare async generators.
``**options`` is the one extension point. Known keys, all optional and
ignored by backends that don't keep state: ``conversation_id`` (stable id
the backend may pin a live session to), ``session_id`` (backend session to
resume when nothing is live), ``capture`` (a
:class:`~beaver_gateway.core.turn_capture.TurnCapture` the backend fills
after the stream closes).
"""
from __future__ import annotations
-502
View File
@@ -1,502 +0,0 @@
"""Claude Code backend adapter.
One :class:`ClaudeCodeBackendAdapter` per :class:`ClaudeAgent`. The
underlying :class:`claude_code_api.ClaudeCodeBackend` bakes ``cwd`` /
``model`` / ``system_prompt`` / MCP wiring into a single
:class:`~claude_code_api.BackendOptions` at construction time, so a
single backend instance is conceptually bound to one agent (different
agents would mean different cwds / system prompts / exposed MCPs and
thus different live-session pools).
Per :meth:`complete` we:
* hand the full Anthropic-style ``messages`` list to
``ClaudeCodeBackend.complete`` — it does its own fingerprint-based
session lookup, so we never need to track sessions ourselves;
* turn its events into one ``message_start`` … ``message_stop``
envelope per ``complete`` call, with content-block indices increasing
monotonically across the whole turn;
* close the envelope on the ``ResultMessage``.
How the blocks inside that envelope are produced depends on the agent's
transport. By default each ``AssistantMessage`` is re-emitted as
``content_block_start`` + one delta + ``content_block_stop`` per block —
a "stream" of finished blocks, because the PTY transport has nothing
finer to offer. With ``transport="stream_json"`` *and*
``include_partial_messages``, the backend instead feeds us the real
token-level ``StreamEvent``s and we emit from those; the whole-block
records still arrive and are still kept for :class:`TurnCapture`, but
emitting from both would duplicate every block on the wire.
The per-request ``system`` parameter is intentionally **ignored** —
``BackendOptions.system_prompt`` is fixed at session-spawn time, and the
agent's ``system_prompt`` is the canonical identity of the agent.
"""
from __future__ import annotations
import json
import logging
import uuid
from collections.abc import Mapping # runtime import: isinstance in _emit_stream_event
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Self
from claude_code_api import (
AssistantMessage,
BackendOptions,
ClaudeCodeBackend,
ResultMessage,
StreamEvent,
TextBlock,
ThinkingBlock,
ToolUseBlock,
synthesize_turn_messages,
)
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.core.events import (
StopReason,
build_content_block_stop,
build_input_json_delta,
build_message_delta,
build_message_start,
build_message_stop,
build_signature_delta,
build_text_block_start,
build_text_delta,
build_thinking_block_start,
build_thinking_delta,
build_tool_use_block_start,
)
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterable
from anthropic.types import MessageParam
from beaver_gateway.agents.base import BaseAgent
from beaver_gateway.core.events import MessageStreamEvent
_log = logging.getLogger("beaver_gateway.backends.claude_code")
__all__ = ["ClaudeCodeBackendAdapter", "TurnCapture"]
@dataclass
class TurnCapture:
"""Side-channel sink for per-turn metadata.
Pass an instance via ``ClaudeCodeBackendAdapter.complete(capture=...)``.
After the stream finishes, :attr:`synthesized_messages` holds the
full assistant↔tool-result cycle (from
:func:`claude_code_api.synthesize_turn_messages`) — i.e. the exact
list of canonical Anthropic-shape messages claude-code-api stashed
the live session under. The markdown frontend uses this to write the
conversation history to its DB so a subsequent turn's prefix
fingerprint hits the same session.
Other backends (anthropic, raycast) ignore the kwarg — it lands in
their ``**options`` and is silently dropped.
"""
synthesized_messages: list[dict[str, Any]] = field(default_factory=list)
_CLAUDE_TO_ANTHROPIC_STOP: dict[str, StopReason] = {
"end_turn": "end_turn",
"tool_use": "tool_use",
"max_tokens": "max_tokens",
"stop_sequence": "stop_sequence",
"refusal": "refusal",
}
def _map_stop_reason(raw: str | None) -> StopReason:
"""Map claude-code's stop reason into Anthropic's vocabulary.
Unknown / missing values collapse to ``end_turn`` so the client sees
a clean finish rather than a wire-format error.
"""
if raw is None:
return "end_turn"
return _CLAUDE_TO_ANTHROPIC_STOP.get(raw, "end_turn")
def _build_mcp_servers(
agent: ClaudeAgent, mcp_internal_urls: Mapping[str, str]
) -> dict[str, dict[str, Any]] | None:
"""Render ``agent.expose_mcps`` into ``BackendOptions.mcp_servers``.
Each exposed MCP is a streamable-HTTP pointer at the gateway's
internal aggregator (built by :mod:`beaver_gateway.mcp.internal_app`).
``None`` keeps claude-code from materializing an ``--mcp-config``
file when the agent exposes nothing.
"""
if not agent.expose_mcps:
return None
servers: dict[str, dict[str, Any]] = {}
for em in agent.expose_mcps:
url = mcp_internal_urls.get(em.name)
if url is None:
msg = (
f"agent {agent.name!r} exposes MCP {em.name!r} "
"but no internal URL is registered for it"
)
raise ValueError(msg)
servers[em.name] = {"type": "http", "url": url}
return servers
def _build_backend_options(
agent: ClaudeAgent, mcp_internal_urls: Mapping[str, str]
) -> BackendOptions:
"""Compose the per-agent :class:`BackendOptions`.
Agent-primary fields:
* ``cwd`` / ``model`` come from the agent directly;
* ``system_prompt`` carries :attr:`BaseAgent.system_prompt`
verbatim — i.e. wire-level ``--system-prompt`` (~8.6k tokens
lighter than ``--append-system-prompt`` because claude-code's
persona/planning conventions and dynamic sections drop out;
tool schemas survive via the API ``tools=[]`` channel);
* ``append_system_prompt`` carries
:attr:`ClaudeCodeOptions.append_system_prompt`, normally
``None``. Setting it re-attaches claude-code's built-in prompt
*and* this delta — opt-in for "claude as a real coding session";
* ``allowed_tools`` follows the PLAN: when the user lists native
tools we restrict to those *plus* a per-MCP wildcard so MCP tools
stay reachable; when no native list is declared we leave
``allowed_tools`` empty (= all tools allowed by claude-code's
default);
* ``mcp_servers`` comes from :func:`_build_mcp_servers`.
Every other tunable knob is passed through from
:attr:`ClaudeAgent.options`. Our default overrides
(``wait_for_turn_duration=True``,
``dangerously_skip_permissions=True``) live on
:class:`ClaudeCodeOptions`, not here, so a user who builds
``ClaudeCodeOptions(...)`` explicitly inherits the same defaults
instead of getting whatever claude-code-api ships.
"""
allowed_tools: tuple[str, ...] = ()
if agent.available_native_tools:
mcp_wildcards = tuple(f"mcp__{em.name}" for em in agent.expose_mcps)
allowed_tools = tuple(agent.available_native_tools) + mcp_wildcards
opt = agent.options
return BackendOptions(
cwd=agent.cwd,
model=agent.model or None,
system_prompt=agent.system_prompt,
append_system_prompt=opt.append_system_prompt,
allowed_tools=allowed_tools,
mcp_servers=_build_mcp_servers(agent, mcp_internal_urls),
transport=opt.transport,
include_partial_messages=opt.include_partial_messages,
# ``None`` means "decide from the agent": an agent that exposes
# MCPs needs them on its first reply, one that doesn't shouldn't
# pay for a warm-up it gains nothing from.
warmup_turn=(
bool(agent.expose_mcps) if opt.warmup_turn is None else opt.warmup_turn
),
disallowed_tools=opt.disallowed_tools,
permission_mode=opt.permission_mode,
dangerously_skip_permissions=opt.dangerously_skip_permissions,
effort=opt.effort,
add_dir=opt.add_dir,
settings=opt.settings,
extra_args=opt.extra_args,
extra_env=opt.extra_env,
preserve_provider_env=opt.preserve_provider_env,
history_injection_mode=opt.history_injection_mode,
wait_for_turn_duration=opt.wait_for_turn_duration,
include_meta_user=opt.include_meta_user,
startup_delay=opt.startup_delay,
file_wait_timeout=opt.file_wait_timeout,
turn_duration_timeout=opt.turn_duration_timeout,
idle_session_ttl=opt.idle_session_ttl,
)
class ClaudeCodeBackendAdapter:
"""One ``claude-code-api`` backend bound to a single :class:`ClaudeAgent`.
Owns the underlying :class:`ClaudeCodeBackend`'s lifecycle through
the async-context-manager protocol so :mod:`beaver_gateway.cli` can
park it in its ``AsyncExitStack``.
"""
def __init__(
self, *, agent: ClaudeAgent, mcp_internal_urls: Mapping[str, str]
) -> None:
self._agent = agent
options = _build_backend_options(agent, mcp_internal_urls)
self._backend = ClaudeCodeBackend(options)
self._streaming = (
options.transport == "stream_json" and options.include_partial_messages
)
"""Whether the backend feeds us token-level ``StreamEvent``s.
When it does, blocks are emitted from those and the whole-block
``AssistantMessage`` records are used only for bookkeeping —
emitting from both would duplicate every block on the wire."""
@property
def agent(self) -> ClaudeAgent:
return self._agent
@property
def live_session_count(self) -> int:
return self._backend.live_session_count
@property
def live_sessions(self) -> dict[str, Any]:
"""Live claude processes keyed by claude session_id.
Pass-through to the underlying ``ClaudeCodeBackend``. The value
is a ``PtyClaudeProcess`` or a ``StreamClaudeProcess`` depending
on the agent's transport; both implement the
``captured_output()`` / ``add_output_listener`` / ``write``
surface the admin terminal consumes, so it works either way —
with the caveat that the stream transport's "terminal" is a
read-only view of the JSON event stream rather than a TUI you
can type into. Typed as ``Any`` to avoid leaking the lower
layer's type into the gateway's public surface.
"""
return self._backend.live_sessions
async def __aenter__(self) -> Self:
await self._backend.__aenter__()
return self
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
await self._backend.__aexit__(exc_type, exc, tb)
async def aclose(self) -> None:
await self._backend.aclose()
async def complete(
self,
*,
agent: BaseAgent,
messages: Iterable[MessageParam],
system: str | None = None, # noqa: ARG002 — see module docstring
capture: TurnCapture | None = None,
**options: Any, # noqa: ARG002 — no per-request knobs for claude-code yet
) -> AsyncIterator[MessageStreamEvent]:
if not isinstance(agent, ClaudeAgent):
msg = (
"ClaudeCodeBackendAdapter requires ClaudeAgent, "
f"got {type(agent).__name__}"
)
raise TypeError(msg)
if agent.name != self._agent.name:
# Adapter is per-agent; routing a different agent through it
# would mean a different cwd / system_prompt / MCP set than
# the live-session pool was spawned with.
msg = (
f"ClaudeCodeBackendAdapter bound to {self._agent.name!r} "
f"got request for {agent.name!r}"
)
raise ValueError(msg)
msgs_list: list[Mapping[str, Any]] = list(messages)
_log.info(
"complete: agent=%s n_messages=%d capture=%s live_sessions=%d",
agent.name,
len(msgs_list),
capture is not None,
self._backend.live_session_count,
)
message_id = f"msg_{uuid.uuid4().hex}"
yield build_message_start(message_id=message_id, model=agent.model)
next_index = 0
stop_reason: str | None = None
usage: Mapping[str, Any] | None = None
n_text = 0
n_thinking = 0
n_tool_use = 0
# We keep raw events so we can hand them to
# ``synthesize_turn_messages`` after the stream closes — the
# markdown frontend stores the result in its conversation
# history so the next turn's prefix matches the backend's
# session-pool fingerprint. UserMessage (tool_result) events
# are silently discarded from the wire but kept here.
raw_events: list[Any] = []
async for event in self._backend.complete(msgs_list):
raw_events.append(event)
if isinstance(event, StreamEvent):
# Only ever arrives when ``self._streaming`` — the
# backend doesn't emit these otherwise. Indices are
# already rebased onto the whole turn by claude-code-api.
for ev in _emit_stream_event(event.event):
yield ev
elif isinstance(event, AssistantMessage):
for block in event.content:
if isinstance(block, TextBlock):
n_text += 1
elif isinstance(block, ThinkingBlock):
n_thinking += 1
elif isinstance(block, ToolUseBlock):
n_tool_use += 1
if not self._streaming:
for ev in _emit_block(block, next_index):
yield ev
next_index += 1
elif isinstance(event, ResultMessage):
# ResultMessage is the terminal event from TurnManager
# — we capture its stop_reason / usage for the envelope
# below. We DO NOT break here: an early break would
# raise GeneratorExit inside claude-code-api's
# ``complete`` coroutine before it gets a chance to
# stash the live session under the post-turn
# fingerprint, so every continuation would miss the
# cache and reseed. Let the inner generator exit
# naturally instead.
stop_reason = event.stop_reason
usage = event.usage
# UserMessage (tool_result records) and SystemMessage
# (turn_duration heartbeats) carry no content for the
# /v1/messages caller — skip silently on the wire, but they
# ARE retained in ``raw_events`` for synthesis below.
if capture is not None:
capture.synthesized_messages = synthesize_turn_messages(raw_events)
_log.info(
"complete: agent=%s DONE text=%d thinking=%d tool_use=%d stop=%s synth=%d",
agent.name,
n_text,
n_thinking,
n_tool_use,
stop_reason,
len(capture.synthesized_messages) if capture is not None else 0,
)
yield build_message_delta(
stop_reason=_map_stop_reason(stop_reason), usage=_normalize_usage(usage)
)
yield build_message_stop()
def _emit_stream_event(event: Mapping[str, Any]) -> Iterable[MessageStreamEvent]:
"""Render one raw Anthropic streaming event through our own builders.
The payload is already Anthropic-shaped, so forwarding it verbatim
is tempting — but it arrives as a plain dict and the frontend
serializes pydantic models, so it would have to be validated against
the SDK's ``RawMessageStreamEvent`` union anyway. Rebuilding it from
the fields we recognize is the same amount of work and fails the
same way :func:`_emit_block` already does: a block or delta type
this gateway doesn't know is dropped rather than raising mid-stream
on a claude release that added one.
Only ``content_block_*`` events reach here — claude-code-api strips
the inner per-request message envelopes, because the caller owns the
single envelope spanning the whole turn.
"""
etype = event.get("type")
index = event.get("index")
if not isinstance(index, int):
return ()
if etype == "content_block_start":
block = event.get("content_block")
if not isinstance(block, Mapping):
return ()
btype = block.get("type")
if btype == "text":
return (build_text_block_start(index),)
if btype == "thinking":
return (build_thinking_block_start(index),)
if btype == "tool_use":
return (
build_tool_use_block_start(
index,
tool_use_id=str(block.get("id", "")),
name=str(block.get("name", "")),
),
)
return ()
if etype == "content_block_delta":
delta = event.get("delta")
if not isinstance(delta, Mapping):
return ()
dtype = delta.get("type")
if dtype == "text_delta":
return (build_text_delta(index, str(delta.get("text", ""))),)
if dtype == "thinking_delta":
return (build_thinking_delta(index, str(delta.get("thinking", ""))),)
if dtype == "signature_delta":
return (build_signature_delta(index, str(delta.get("signature", ""))),)
if dtype == "input_json_delta":
return (build_input_json_delta(index, str(delta.get("partial_json", ""))),)
return ()
if etype == "content_block_stop":
return (build_content_block_stop(index),)
return ()
def _emit_block(
block: TextBlock | ThinkingBlock | ToolUseBlock | Any, index: int
) -> Iterable[MessageStreamEvent]:
"""Render one ``claude-code`` content block as Anthropic stream events.
``ToolResultBlock`` would arrive only on user-role records — we
don't emit it here because :meth:`complete` skips ``UserMessage``.
"""
if isinstance(block, TextBlock):
return (
build_text_block_start(index),
build_text_delta(index, block.text),
build_content_block_stop(index),
)
if isinstance(block, ThinkingBlock):
return (
build_thinking_block_start(index),
build_thinking_delta(index, block.thinking),
build_signature_delta(index, block.signature),
build_content_block_stop(index),
)
if isinstance(block, ToolUseBlock):
partial = json.dumps(block.input, separators=(",", ":"), ensure_ascii=False)
return (
build_tool_use_block_start(index, tool_use_id=block.id, name=block.name),
build_input_json_delta(index, partial),
build_content_block_stop(index),
)
return ()
def _normalize_usage(usage: Mapping[str, Any] | None) -> dict[str, int] | None:
"""Coerce claude-code's ``usage`` dict to Anthropic ``MessageDeltaUsage`` shape.
claude-code copies whatever the JSONL ``usage`` record carried —
fields can be missing, strings, or ints. We pass through only the
fields ``MessageDeltaUsage`` knows about and discard the rest so an
odd ``cache_creation`` object structure doesn't fail pydantic
validation downstream.
"""
if not usage:
return None
out: dict[str, int] = {}
for key in (
"input_tokens",
"output_tokens",
"cache_creation_input_tokens",
"cache_read_input_tokens",
):
value = usage.get(key)
if isinstance(value, int):
out[key] = value
return out or None
+789
View File
@@ -0,0 +1,789 @@
"""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. Sessions are keyed by ``conversation_id`` when the frontend passes
one (markdown chats) or by a text-only fingerprint of ``messages[:-1]``
for stateless callers (``/v1/messages``). Without a live session the
adapter resumes ``session_id`` from the session store, or seeds the
incoming history into the store via ``core/transcript`` and resumes that.
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, under that uid.
"""
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, Protocol, Self, cast
import claude_agent_sdk
from claude_agent_sdk import (
AssistantMessage,
ClaudeAgentOptions,
ClaudeSDKClient,
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.transcript import build_entries
from beaver_gateway.core.turn_capture import TurnCapture, TurnUsage
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
from anthropic.types import MessageParam
from claude_agent_sdk import SessionStore
from 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",
"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_")
_REAP_INTERVAL = 60.0
_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",
}
class SessionClient(Protocol):
async def connect(self) -> None: ...
async def query(self, prompt: str) -> None: ...
def receive_response(self) -> AsyncIterator[Any]: ...
async def disconnect(self) -> None: ...
ClientFactory = "Callable[[ClaudeAgentOptions], SessionClient]"
UsageSink = "Callable[[UsageEvent], Awaitable[None]]"
@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
@dataclass
class _Live:
client: SessionClient
session_id: str | None
resumed: bool
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
last_used: float = field(default_factory=time.monotonic)
turns: int = 0
class _RunnerClient(ClaudeSDKClient):
"""``ClaudeSDKClient`` that hands the materialized resume dir to the runner uid."""
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,
) -> 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._sessions: dict[str, _Live] = {}
self._reaper: asyncio.Task[None] | None = None
self._uid = _resolve_uid(self._runner.user)
self._wrapper: Path | None = None
@property
def agent(self) -> ClaudeAgent:
return self._agent
@property
def sessions(self) -> dict[str, dict[str, Any]]:
now = time.monotonic()
return {
key: {
"session_id": live.session_id,
"idle_seconds": now - live.last_used,
"turns": live.turns,
"busy": live.lock.locked(),
}
for key, live in self._sessions.items()
}
async def __aenter__(self) -> Self:
self._reaper = asyncio.create_task(self._reap_loop())
return self
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
await self.aclose()
async def aclose(self) -> None:
if self._reaper is not None:
self._reaper.cancel()
with contextlib.suppress(BaseException):
await self._reaper
self._reaper = None
for key in list(self._sessions):
await self._close(key)
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,
**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)
live = await self._acquire(key, session_id=session_id, history=prior)
message_id = f"msg_{uuid.uuid4().hex}"
yield build_message_start(message_id=message_id, model=self._agent.model)
async with live.lock:
live.last_used = time.monotonic()
try:
turn = await self._run_turn(live, prompt)
except Exception:
if live.resumed and live.turns == 0:
_log.exception(
"resume of %s failed, reseeding from history", live.session_id
)
await self._close(key)
live = await self._acquire(key, session_id=None, history=prior)
async with live.lock:
turn = await self._run_turn(live, prompt)
else:
raise
for event in turn.events:
yield event
live.turns += 1
live.last_used = time.monotonic()
if turn.result is not None and turn.result.session_id:
live.session_id = turn.result.session_id
if conversation_id is None:
self._rekey(key, fingerprint([*history, *turn.synthesized]))
usage = _usage_of(turn.result)
if capture is not None:
capture.synthesized_messages = turn.synthesized
capture.session_id = live.session_id
capture.usage = usage
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,
)
)
if turn.result is not None and turn.result.is_error:
msg = f"claude: {turn.result.result or turn.result.subtype}"
raise RuntimeError(msg)
yield build_message_delta(
stop_reason=turn.stop_reason, usage=_wire_usage(usage)
)
yield build_message_stop()
async def _run_turn(self, live: _Live, prompt: str) -> _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 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]]
) -> _Live:
live = self._sessions.get(key)
if live is not None:
return live
resume = session_id
if resume is None and history:
resume = await self._seed(history)
live = await self._spawn(resume)
self._sessions[key] = live
return live
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,
)
key = {
"project_key": project_key_for_directory(str(self._agent.cwd)),
"session_id": session_id,
}
await self._store.append(cast("Any", key), cast("Any", entries))
_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) -> _Live:
options = self._build_options(resume)
client = self._factory(options)
await client.connect()
_log.info(
"spawned claude: agent=%s resume=%s user=%s",
self._agent.name,
resume,
self._runner.user,
)
return _Live(client=client, session_id=resume, resumed=resume is not None)
def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient:
return _RunnerClient(options, uid=self._uid)
def _build_options(self, resume: str | None) -> ClaudeAgentOptions:
agent = self._agent
opt = agent.options
env = dict(opt.env)
if self._runner.home is not None:
env["HOME"] = str(self._runner.home)
plugins = self._plugins()
system_prompt = (
prompt_assembly.assemble(agent.prompt_sources)
if agent.prompt_sources
else agent.system_prompt
)
return ClaudeAgentOptions(
model=agent.model or None,
effort=cast("Any", opt.effort),
system_prompt=system_prompt,
setting_sources=[],
strict_mcp_config=True,
mcp_servers=cast("Any", self._servers),
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,
user=self._runner.user,
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)),
)
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:
live = self._sessions.pop(old, None)
if live is None:
return
stale = self._sessions.pop(new, None)
self._sessions[new] = live
if stale is not None and stale is not live:
asyncio.get_running_loop().create_task(_disconnect(stale))
async def _close(self, key: str) -> None:
live = self._sessions.pop(key, None)
if live is not None:
await _disconnect(live)
async def _reap_loop(self) -> None:
ttl = self._agent.options.idle_session_ttl
while True:
await asyncio.sleep(_REAP_INTERVAL)
if ttl <= 0:
continue
now = time.monotonic()
for key, live in list(self._sessions.items()):
if not live.lock.locked() and now - live.last_used > ttl:
_log.info("closing idle session %s (%s)", live.session_id, key)
await self._close(key)
@dataclass
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})
env = {{
k: v for k, v in os.environ.items() if k in KEEP or k.startswith(PREFIXES)
}}
os.execve(TARGET, [TARGET, *sys.argv[1:]], env)
"""
async def _disconnect(live: _Live) -> 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_uid(user: str | None) -> int | None:
if user is None:
return None
if user.isdigit():
return int(user)
return pwd.getpwnam(user).pw_uid
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 fingerprint(messages: Iterable[Mapping[str, Any]]) -> str:
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 ""
def _prompt_text(content: Any) -> str:
text = _text_of(content)
if not text:
msg = "user message has no text content"
raise ValueError(msg)
return text
def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
out: list[dict[str, Any]] = []
for message in raw:
if isinstance(message, AssistantMessage):
out.append(
{
"role": "assistant",
"content": [_block_to_dict(b) for b in message.content],
}
)
elif isinstance(message, UserMessage):
content = message.content
if isinstance(content, list) and content:
out.append(
{"role": "user", "content": [_block_to_dict(b) for b in content]}
)
return out
def _block_to_dict(block: Any) -> dict[str, Any]:
if isinstance(block, TextBlock):
return {"type": "text", "text": block.text}
if isinstance(block, ToolUseBlock):
return {
"type": "tool_use",
"id": block.id,
"name": block.name,
"input": block.input,
}
if isinstance(block, ToolResultBlock):
result: dict[str, Any] = {
"type": "tool_result",
"tool_use_id": block.tool_use_id,
"content": block.content,
}
if block.is_error is not None:
result["is_error"] = block.is_error
return result
if isinstance(block, ThinkingBlock):
return {
"type": "thinking",
"thinking": block.thinking,
"signature": block.signature,
}
msg = f"unknown content block type: {type(block).__name__}"
raise TypeError(msg)
def _usage_of(result: ResultMessage | None) -> TurnUsage:
if result is None:
return TurnUsage()
usage = result.usage or {}
return TurnUsage(
input_tokens=_int(usage.get("input_tokens")),
output_tokens=_int(usage.get("output_tokens")),
cache_read_tokens=_int(usage.get("cache_read_input_tokens")),
cache_creation_tokens=_int(usage.get("cache_creation_input_tokens")),
cost_usd=result.total_cost_usd,
duration_ms=result.duration_ms,
num_turns=result.num_turns,
)
def _wire_usage(usage: TurnUsage) -> dict[str, int]:
return {
"input_tokens": usage.input_tokens,
"output_tokens": usage.output_tokens,
"cache_read_input_tokens": usage.cache_read_tokens,
"cache_creation_input_tokens": usage.cache_creation_tokens,
}
def _int(value: Any) -> int:
return value if isinstance(value, int) else 0
def _emit_stream_event(
event: Mapping[str, Any], index: int
) -> Iterable[MessageStreamEvent]:
etype = event.get("type")
if etype == "content_block_start":
block = event.get("content_block")
if not isinstance(block, Mapping):
return ()
btype = block.get("type")
if btype == "text":
return (build_text_block_start(index),)
if btype == "thinking":
return (build_thinking_block_start(index),)
if btype == "tool_use":
return (
build_tool_use_block_start(
index,
tool_use_id=str(block.get("id", "")),
name=str(block.get("name", "")),
),
)
return ()
if etype == "content_block_delta":
delta = event.get("delta")
if not isinstance(delta, Mapping):
return ()
dtype = delta.get("type")
if dtype == "text_delta":
return (build_text_delta(index, str(delta.get("text", ""))),)
if dtype == "thinking_delta":
return (build_thinking_delta(index, str(delta.get("thinking", ""))),)
if dtype == "signature_delta":
return (build_signature_delta(index, str(delta.get("signature", ""))),)
if dtype == "input_json_delta":
return (build_input_json_delta(index, str(delta.get("partial_json", ""))),)
return ()
if etype == "content_block_stop":
return (build_content_block_stop(index),)
return ()
def _emit_block(block: Any, index: int) -> Iterable[MessageStreamEvent]:
if isinstance(block, TextBlock):
return (
build_text_block_start(index),
build_text_delta(index, block.text),
build_content_block_stop(index),
)
if isinstance(block, ThinkingBlock):
return (
build_thinking_block_start(index),
build_thinking_delta(index, block.thinking),
build_signature_delta(index, block.signature),
build_content_block_stop(index),
)
if isinstance(block, ToolUseBlock):
partial = json.dumps(block.input, separators=(",", ":"), ensure_ascii=False)
return (
build_tool_use_block_start(index, tool_use_id=block.id, name=block.name),
build_input_json_delta(index, partial),
build_content_block_stop(index),
)
return ()
+4 -1
View File
@@ -26,11 +26,12 @@ tools we route it back to the underlying MCP in-process, append the
result as a ``tool`` message, and re-issue the stream — all inside one
Anthropic envelope (one ``message_start`` … one ``message_stop``). The
tool_use blocks DO surface to the caller (mirrors what
``ClaudeCodeBackendAdapter`` does), but tool_results stay internal.
``ClaudeSdkBackend`` does), but tool_results stay internal.
"""
from __future__ import annotations
import fnmatch
import json
import logging
import uuid
@@ -297,6 +298,8 @@ def _build_agent_catalog(
for mt in mcp_tools.get(em.name, []):
if em.tools is not None and mt.name not in em.tools:
continue
if any(fnmatch.fnmatchcase(mt.name, pat) for pat in em.deny):
continue
wire_name = f"{em.name}__{mt.name}"
routing[wire_name] = (em.name, mt.name)
local_tools.append(