refactor: split flat core into capability packages, layer the conversations service, English defaults for every model-facing text
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
"""Backend adapters.
|
||||
|
||||
Each backend wraps a provider SDK (``raycast-api``, ``claude-agent-sdk``)
|
||||
and yields the unified :class:`~beaver_gateway.core.events.MessageStreamEvent`
|
||||
and yields the unified :class:`~beaver_gateway.events.stream.MessageStreamEvent`
|
||||
family. The Anthropic-style frontend serialises events straight to SSE.
|
||||
"""
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Backend protocol.
|
||||
|
||||
A backend turns an Anthropic-style turn (``messages`` + agent definition)
|
||||
into a stream of :class:`~beaver_gateway.core.events.MessageStreamEvent`
|
||||
into a stream of :class:`~beaver_gateway.events.stream.MessageStreamEvent`
|
||||
records. The frontend serializes whatever comes out straight to SSE, so
|
||||
backends are the only place where provider quirks are translated.
|
||||
|
||||
@@ -12,7 +12,7 @@ subclassing - to keep them swappable in tests with bare async generators.
|
||||
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
|
||||
:class:`~beaver_gateway.backends.capture.TurnCapture` the backend fills
|
||||
after the stream closes).
|
||||
"""
|
||||
|
||||
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
|
||||
from anthropic.types import MessageParam
|
||||
|
||||
from beaver_gateway.agents.base import BaseAgent
|
||||
from beaver_gateway.core.events import MessageStreamEvent
|
||||
from beaver_gateway.events.stream import MessageStreamEvent
|
||||
|
||||
|
||||
class Backend(Protocol):
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
"""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
|
||||
model_usage: dict[str, Any] | None = None
|
||||
"""``ResultMessage.model_usage`` verbatim: per-model tokens, cost, web searches."""
|
||||
|
||||
context_tokens: int = 0
|
||||
"""Input of the last API call in the turn (fresh + cached + written to
|
||||
cache) - the context size the model actually ran with, what Claude Code
|
||||
shows as the context. The token fields above are sums over every API
|
||||
call of the turn and grow with the number of tool calls."""
|
||||
|
||||
|
||||
@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."""
|
||||
@@ -3,7 +3,7 @@
|
||||
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
|
||||
:class:`~beaver_gateway.backends.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
|
||||
@@ -67,9 +67,18 @@ from claude_agent_sdk import (
|
||||
project_key_for_directory,
|
||||
)
|
||||
|
||||
from beaver_gateway.core import policy as policy_mod
|
||||
from beaver_gateway.core import prompt as prompt_assembly
|
||||
from beaver_gateway.core.events import (
|
||||
from beaver_gateway.agents import policy as policy_mod
|
||||
from beaver_gateway.agents import prompts as prompt_assembly
|
||||
from beaver_gateway.backends.capture import TurnCapture, TurnUsage
|
||||
from beaver_gateway.backends.sessions import Session, SessionClient, SessionPool
|
||||
from beaver_gateway.backends.transcript import (
|
||||
build_entries,
|
||||
close_open_tool_uses,
|
||||
fingerprint,
|
||||
text_of,
|
||||
)
|
||||
from beaver_gateway.conversations.kinds import as_kind
|
||||
from beaver_gateway.events.stream import (
|
||||
StopReason,
|
||||
build_content_block_stop,
|
||||
build_input_json_delta,
|
||||
@@ -83,15 +92,6 @@ from beaver_gateway.core.events import (
|
||||
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
|
||||
@@ -107,8 +107,8 @@ if TYPE_CHECKING:
|
||||
|
||||
from beaver_gateway.agents.base import BaseAgent
|
||||
from beaver_gateway.agents.claude import ClaudeAgent
|
||||
from beaver_gateway.core.events import MessageStreamEvent
|
||||
from beaver_gateway.core.kinds import Kind
|
||||
from beaver_gateway.conversations.kinds import Kind
|
||||
from beaver_gateway.events.stream import MessageStreamEvent
|
||||
|
||||
|
||||
_log = logging.getLogger("beaver_gateway.backends.claude_sdk")
|
||||
@@ -281,13 +281,17 @@ class ClaudeSdkBackend:
|
||||
def live(self, key: str) -> Session | None:
|
||||
return self._pool.get(key)
|
||||
|
||||
async def repair_session(self, session_id: str) -> int:
|
||||
async def repair_session(
|
||||
self, session_id: str, *, text: str = "interrupted"
|
||||
) -> 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))
|
||||
fixes = close_open_tool_uses(
|
||||
cast("list[Mapping[str, Any]]", entries), text=text
|
||||
)
|
||||
if fixes:
|
||||
await self._store.append(cast("Any", key), cast("Any", fixes))
|
||||
_log.warning(
|
||||
|
||||
@@ -42,7 +42,7 @@ from raycast_api import Message as RaycastMessage
|
||||
from raycast_api import RemoteTool, Tool, ToolCall
|
||||
|
||||
from beaver_gateway.agents.raycast import RaycastAgent
|
||||
from beaver_gateway.core.events import (
|
||||
from beaver_gateway.events.stream import (
|
||||
StopReason,
|
||||
build_content_block_stop,
|
||||
build_input_json_delta,
|
||||
@@ -65,7 +65,7 @@ if TYPE_CHECKING:
|
||||
from raycast_api import ChatStreamChunk, Client
|
||||
|
||||
from beaver_gateway.agents.base import BaseAgent
|
||||
from beaver_gateway.core.events import MessageStreamEvent
|
||||
from beaver_gateway.events.stream import MessageStreamEvent
|
||||
else:
|
||||
from collections.abc import Mapping
|
||||
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
"""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.backends.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
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
"""Scratch for policy rules (``core/policy``); dies with the process."""
|
||||
|
||||
@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,545 @@
|
||||
"""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__
|
||||
|
||||
_cli_version = str(__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 = "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 ""
|
||||
Reference in New Issue
Block a user