From 07e84618e5914befcdb01af6a6795c83fb9cf2d8 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Thu, 5 Nov 2026 19:04:58 -0800 Subject: [PATCH 1/3] Add Claude Code and Pi workers --- examples/workers/agent_tools.py | 6 +- examples/workers/approvals.py | 6 +- examples/workers/streaming.py | 6 +- pyproject.toml | 11 + src/ai/providers/_optional.py | 16 +- src/ai/workers/__init__.py | 10 +- src/ai/workers/_adapters/claude.py | 631 ++++++++++++ src/ai/workers/_adapters/pi.py | 1111 ++++++++++++++++++++++ src/ai/workers/_adapters/pi_extension.js | 88 ++ src/ai/workers/_kinds.py | 75 +- src/ai/workers/_worker.py | 4 +- tests/workers/_adapters/test_claude.py | 317 ++++++ tests/workers/_adapters/test_pi.py | 554 +++++++++++ tests/workers/test_live.py | 137 +++ uv.lock | 41 +- 15 files changed, 2994 insertions(+), 19 deletions(-) create mode 100644 src/ai/workers/_adapters/claude.py create mode 100644 src/ai/workers/_adapters/pi.py create mode 100644 src/ai/workers/_adapters/pi_extension.js create mode 100644 tests/workers/_adapters/test_claude.py create mode 100644 tests/workers/_adapters/test_pi.py create mode 100644 tests/workers/test_live.py diff --git a/examples/workers/agent_tools.py b/examples/workers/agent_tools.py index 3885abe7..f56bef4c 100644 --- a/examples/workers/agent_tools.py +++ b/examples/workers/agent_tools.py @@ -3,7 +3,7 @@ The CLI calls `lookup_order`. The worker hands the call to the agent, the agent runs the tool, and the CLI gets the result and goes on. - uv run python examples/workers/agent_tools.py [codex] + uv run python examples/workers/agent_tools.py [claude-code|codex|pi] """ import asyncio @@ -14,7 +14,9 @@ from ai import workers KINDS: dict[str, Callable[..., workers.Worker]] = { + "claude-code": workers.ClaudeCode, "codex": workers.Codex, + "pi": workers.Pi, } @@ -47,4 +49,4 @@ async def main(kind: str) -> None: print([m.role for m in run.messages]) -asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "codex")) +asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "claude-code")) diff --git a/examples/workers/approvals.py b/examples/workers/approvals.py index 348a47b2..0d71f625 100644 --- a/examples/workers/approvals.py +++ b/examples/workers/approvals.py @@ -3,7 +3,7 @@ With `approve=True`, every tool call the CLI makes becomes a `ToolApproval` hook. Answer it with `ai.resolve_hook`, the way any tool approval is. - uv run python examples/workers/approvals.py [codex] + uv run python examples/workers/approvals.py [claude-code|codex|pi] """ import asyncio @@ -14,7 +14,9 @@ from ai import workers KINDS: dict[str, Callable[..., workers.Worker]] = { + "claude-code": workers.ClaudeCode, "codex": workers.Codex, + "pi": workers.Pi, } @@ -39,4 +41,4 @@ async def main(kind: str) -> None: print() -asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "codex")) +asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "claude-code")) diff --git a/examples/workers/streaming.py b/examples/workers/streaming.py index a5239cd3..feb01d65 100644 --- a/examples/workers/streaming.py +++ b/examples/workers/streaming.py @@ -1,6 +1,6 @@ """A worker where a model goes: stream it, then continue the conversation. -uv run python examples/workers/streaming.py [codex] +uv run python examples/workers/streaming.py [claude-code|codex|pi] """ import asyncio @@ -11,7 +11,9 @@ from ai import workers KINDS: dict[str, Callable[..., workers.Worker]] = { + "claude-code": workers.ClaudeCode, "codex": workers.Codex, + "pi": workers.Pi, } @@ -49,4 +51,4 @@ async def main(kind: str) -> None: print(reply.text) -asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "codex")) +asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "claude-code")) diff --git a/pyproject.toml b/pyproject.toml index 2bbe0863..a22eab04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ [project.optional-dependencies] anthropic = ["anthropic>=0.83.0,<2.0.0"] +claude-code = ["claude-agent-sdk>=0.2.144,<0.3"] mcp = ["mcp>=2.0.0,<3.0.0"] openai = ["openai>=2.34.0,<4.0.0"] otel = [ @@ -69,6 +70,7 @@ bump = true [dependency-groups] dev = [ "anthropic>=0.83.0", + "claude-agent-sdk>=0.2.144,<0.3", "mcp>=2.0.0,<3.0.0", "python-dotenv>=1.2.1", "pytest>=8.0", @@ -82,6 +84,7 @@ dev = [ "vercel-oidc>=0.8.0", "opentelemetry-sdk>=1.43.0", "opentelemetry-exporter-otlp-proto-http>=1.43.0", + "pytest-timeout>=2.3", "typesafe-sdk>=0.7.2", ] @@ -115,6 +118,14 @@ exclude = [ [tool.pytest.ini_options] asyncio_mode = "auto" testpaths = ["tests"] +# Live worker tests drive real agent CLIs, so they are opt-in: +# `uv run pytest -m live tests/workers` (a later -m overrides this one). +addopts = "-m 'not live'" +# A live test drives a real agent; hanging must be a failure. +timeout = 300 +markers = [ + "live: drives a real Claude Code, Codex or Pi (deselect with -m 'not live')", +] filterwarnings = [ "ignore:'asyncio\\.DefaultEventLoopPolicy' is deprecated.*:DeprecationWarning:async_solipsism\\.loop", ] diff --git a/src/ai/providers/_optional.py b/src/ai/providers/_optional.py index d3f48140..5ee15265 100644 --- a/src/ai/providers/_optional.py +++ b/src/ai/providers/_optional.py @@ -12,17 +12,25 @@ def import_optional_sdk( - module_name: str, *, provider: str, extra: str + module_name: str, + *, + provider: str | None = None, + extra: str, + feature: str | None = None, ) -> ModuleType: - """Import an optional upstream SDK or raise a helpful installation error.""" + """Import an optional upstream SDK or raise a helpful installation error. + + ``feature`` names what needs the SDK when it is not a provider. + """ root_module = module_name.partition(".")[0] + feature = feature or f"the {provider} provider" try: return importlib.import_module(module_name) except ModuleNotFoundError as exc: if exc.name not in {module_name, root_module}: raise raise ai_errors.InstallationError( - f"could not import `{root_module}`, which is required to use the " - f"{provider} provider, you can install it with `pip install " + f"could not import `{root_module}`, which is required to use " + f"{feature}, you can install it with `pip install " f'"ai[{extra}]"` or `uv add "ai[{extra}]"`' ) from exc diff --git a/src/ai/workers/__init__.py b/src/ai/workers/__init__.py index 058a3b95..2b66dac4 100644 --- a/src/ai/workers/__init__.py +++ b/src/ai/workers/__init__.py @@ -1,4 +1,4 @@ -"""Coding-agent CLIs as workers. +"""Coding-agent CLIs as workers: Claude Code, Codex and Pi. Experimental: not part of the stable API, may change or be removed. @@ -7,7 +7,7 @@ import ai from ai import workers - async with workers.Codex(model="gpt-5.5") as worker: + async with workers.ClaudeCode(model="sonnet") as worker: async with ai.stream(worker, [ai.user_message("Fix the tests")]) as s: async for event in s: ... @@ -20,7 +20,7 @@ A worker is a pydantic model: `worker.model_dump()` saves the conversation, `workers.Worker.model_validate(data)` continues it, `worker.model_copy()` is -a new worker on the same settings, and `Codex(session_id=...)` resumes +a new worker on the same settings, and `ClaudeCode(session_id=...)` resumes a conversation the CLI has stored. """ @@ -31,13 +31,15 @@ TranscriptMismatchError, WorkerError, ) -from ._kinds import Codex +from ._kinds import ClaudeCode, Codex, Pi from ._worker import Worker __all__ = [ + "ClaudeCode", "Codex", "ExecutableMissingError", "NotAuthenticatedError", + "Pi", "SessionInfo", "TranscriptMismatchError", "Worker", diff --git a/src/ai/workers/_adapters/claude.py b/src/ai/workers/_adapters/claude.py new file mode 100644 index 00000000..7ca08b65 --- /dev/null +++ b/src/ai/workers/_adapters/claude.py @@ -0,0 +1,631 @@ +"""Claude Code, driven through claude-agent-sdk. + +The session id we hand the CLI IS its own session id, so `claude --resume +` opens the same conversation, here or in the CLI's own UI. + +The comments record what was MEASURED against a real `claude` CLI (see +VERIFIED_VERSION), not what the API suggested. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import json +import os +import pathlib +import uuid +import warnings +from collections.abc import AsyncIterator +from typing import Any + +import claude_agent_sdk + +from ...types import events as events_ +from ...types import messages as messages_ +from ...types import usage as usage_ +from .. import _bridge +from .. import _errors as errors +from . import _handoff, base + +KIND = "claude-code" +INSTALL = "npm install -g @anthropic-ai/claude-code" +VERIFIED_VERSION = "2.1.289" +"""The CLI the live tests last passed against. Older ones get a warning.""" + +APPROVAL_TIMEOUT = 24 * 3600 +"""Seconds the CLI waits for an approval. Its own default (60s) would cut a +person short: the hook would time out and the call would be denied.""" + +AUTH_HINT = ( + "claude-code uses your own login: run `claude` once and sign in, or " + "pass env={'ANTHROPIC_API_KEY': ...}." +) + + +class ClaudeAdapter: + def __init__( + self, + *, + cwd: str, + env: dict[str, str], + executable: str, + model: str | None, + effort: str | None, + mcp_servers: dict[str, Any], + approve: base.Approve | None, + allowed_tools: list[str] | None, + disallowed_tools: list[str] | None, + max_turns: int | None, + ) -> None: + # The CLI keys its store by the real path of its directory. + self._cwd = os.path.realpath(cwd) + self._env = env + self._executable = executable + self._model = model + self._effort = effort + self._mcp_servers = mcp_servers + self._approve = approve + self._allowed_tools = allowed_tools + self._disallowed_tools = disallowed_tools + self._max_turns = max_turns + self._instructions: str | None = None + self._client: Any = None + self._session_id: str | None = None + # We interrupted the turn. The CLI reports an error result after an + # interrupt, and a stop we asked for must not raise as a failure. + self._interrupted = False + # A turn is running: prompted, no ResultMessage yet. A steer outside + # one would start a turn of its own (measured). + self._running = False + self.version: str | None = None + + async def start(self) -> None: + self.version = await base.probe( + KIND, + self._executable, + install=INSTALL, + verified=VERIFIED_VERSION, + cwd=self._cwd, + env=self._env, + ) + + async def close(self) -> None: + client, self._client = self._client, None + if client is not None: + with contextlib.suppress(Exception): + await client.disconnect() + + # -- the conversation ----------------------------------------------------- + + async def open( + self, + *, + session_id: str | None, + fork: bool, + seed: list[messages_.Message] | None, + instructions: str | None, + ) -> str: + self._instructions = instructions + if seed: + session_id = await self._stage(seed) + elif session_id is not None: + await self._require_known(session_id) + if fork: + try: + # The SDK remaps every UUID and rewrites the chain. + result = await asyncio.to_thread( + claude_agent_sdk.fork_session, session_id, self._cwd + ) + except Exception as exc: + raise errors.WorkerError( + f"{KIND}: cannot fork session {session_id}: {exc}" + ) from exc + session_id = result.session_id + if session_id is None: + session_id = str(uuid.uuid4()) + await self._connect(session_id, resume=None) + else: + await self._connect(session_id, resume=session_id) + self._session_id = session_id + return session_id + + async def _stage(self, history: list[messages_.Message]) -> str: + """Write a past into the CLI's store under a fresh id; run nothing.""" + session_id = str(uuid.uuid4()) + _handoff.warn_unsupported_files(history, KIND) + dropped = _handoff.reasoning_count(history) + if dropped: + warnings.warn( + f"{dropped} reasoning part(s) in the history were not carried " + "into claude-code: Anthropic's safeguards refuse replayed " + "model reasoning ([reasoning_extraction]). Text and tool " + "calls were carried in full.", + stacklevel=2, + ) + # Where the CLI keeps transcripts: + # `$CLAUDE_CONFIG_DIR/projects//.jsonl`. + config = ( + self._env.get("CLAUDE_CONFIG_DIR") + or os.environ.get("CLAUDE_CONFIG_DIR") + or str(pathlib.Path.home() / ".claude") + ) + key = claude_agent_sdk.project_key_for_directory(self._cwd) + path = pathlib.Path(config) / "projects" / key / f"{session_id}.jsonl" + records = _handoff.to_claude_records(history, session_id, self._cwd) + text = "".join( + json.dumps(r, separators=(",", ":")) + "\n" for r in records + ) + + def write() -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + + await asyncio.to_thread(write) + return session_id + + async def _require_known(self, session_id: str) -> None: + """Refuse an unknown session before connecting to it. + + Connecting to a ghost surfaces as claude-agent-sdk's own error + somewhere downstream. + """ + try: + info = await asyncio.to_thread( + claude_agent_sdk.get_session_info, session_id, self._cwd + ) + except Exception as exc: + raise errors.WorkerError( + f"{KIND}: cannot read session {session_id}: {exc}" + ) from exc + if info is None: + raise errors.WorkerError( + f"{KIND}: no session {session_id} in {self._cwd}" + ) + + async def _connect(self, session_id: str, *, resume: str | None) -> None: + options: dict[str, Any] = { + "cwd": self._cwd, + # Raw Anthropic stream events, which map 1:1 onto Start/Delta/End + # triples. Without them the SDK yields only whole messages. + "include_partial_messages": True, + "session_id": session_id if resume is None else None, + "resume": resume, + } + if self._env: + # claude-agent-sdk spawns the CLI itself: pass the env along. + options["env"] = dict(self._env) + if self._model: + options["model"] = self._model + if self._instructions: + # Added to Claude Code's own system prompt, never replacing it. + options["system_prompt"] = { + "type": "preset", + "preset": "claude_code", + "append": self._instructions, + } + if self._mcp_servers: + options["mcp_servers"] = self._mcp_servers + if self._allowed_tools is not None: + options["tools"] = self._allowed_tools + if self._disallowed_tools: + options["disallowed_tools"] = self._disallowed_tools + if self._max_turns is not None: + options["max_turns"] = self._max_turns + if self._effort: + # Reasoning effort, with adaptive thinking so the budget follows. + options["effort"] = self._effort + options["thinking"] = {"type": "adaptive"} + if self._approve is not None: + # A PreToolUse hook sees EVERY tool call, read-only ones included + # (measured). The permission callback alone only sees what the + # CLI's policy asks about; it stays as a safety net. + options["hooks"] = { + "PreToolUse": [ + claude_agent_sdk.HookMatcher( + matcher=None, + hooks=[self._pre_tool_use], + timeout=APPROVAL_TIMEOUT, + ) + ] + } + options["can_use_tool"] = self._can_use_tool + else: + # No gate: an always-allow callback would freeze the agent at + # its next tool call once nobody answers (measured). + options["permission_mode"] = "bypassPermissions" + client = claude_agent_sdk.ClaudeSDKClient( + options=claude_agent_sdk.ClaudeAgentOptions(**options) + ) + try: + await client.connect() + except claude_agent_sdk.ClaudeSDKError as exc: + raise _translate(exc) from exc + self._client = client + + # -- approval ------------------------------------------------------------- + + async def _pre_tool_use( + self, data: Any, tool_use_id: str | None, ctx: Any + ) -> Any: + assert self._approve is not None + granted, reason = await self._approve( + tool_use_id or str(data.get("tool_use_id") or uuid.uuid4()), + str(data.get("tool_name") or ""), + dict(data.get("tool_input") or {}), + ) + output: dict[str, Any] = { + "hookEventName": "PreToolUse", + "permissionDecision": "allow" if granted else "deny", + } + if not granted: + output["permissionDecisionReason"] = reason + return {"hookSpecificOutput": output} + + async def _can_use_tool( + self, + name: str, + tool_input: dict[str, Any], + ctx: claude_agent_sdk.ToolPermissionContext, + ) -> claude_agent_sdk.PermissionResult: + assert self._approve is not None + granted, reason = await self._approve( + str(getattr(ctx, "tool_use_id", None) or uuid.uuid4()), + name, + dict(tool_input), + ) + if granted: + return claude_agent_sdk.PermissionResultAllow() + return claude_agent_sdk.PermissionResultDeny(message=reason) + + # -- turns ---------------------------------------------------------------- + + async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: + base.check_files(prompt, KIND, ("image/", "application/pdf")) + client = self._client + assert client is not None + # An interrupted turn leaves its trailing result unread: the CLI + # still emits one after the interrupt. Read it first, or THIS turn + # would settle on the previous turn's ending. + await self._drain_interrupted() + content = await _content(prompt) + names: dict[str, str] = {} + open_blocks: dict[int, str] = {} + try: + if isinstance(content, str): + await client.query(content) + else: + await client.query(_one_message(content)) + except claude_agent_sdk.ClaudeSDKError as exc: + raise _translate(exc) from exc + self._running = True + try: + # StreamStart once the CLI answers: its first message proves it + # read the prompt. + started = False + async for message in _guard(client.receive_response()): + if not started: + started = True + yield events_.StreamStart() + if isinstance(message, claude_agent_sdk.StreamEvent): + # Text and reasoning stream here, token by token. The + # AssistantMessage that follows repeats them, so only + # tool calls are taken from it. + for event in _stream_events(message.event, open_blocks): + yield event + elif isinstance(message, claude_agent_sdk.AssistantMessage): + for block in message.content: + if isinstance(block, claude_agent_sdk.ToolUseBlock): + names[block.id] = block.name + for event in base.builtin_call( + block.id, + block.name, + json.dumps(block.input), + started=set(), + ): + yield event + elif isinstance(message, claude_agent_sdk.UserMessage): + content_blocks = ( + message.content + if isinstance(message.content, list) + else [] + ) + for block in content_blocks: + if isinstance(block, claude_agent_sdk.ToolResultBlock): + yield base.builtin_result( + block.tool_use_id, + names.get(block.tool_use_id, ""), + block.content, + is_error=bool(block.is_error), + ) + elif isinstance(message, claude_agent_sdk.ResultMessage): + self._running = False + interrupted, self._interrupted = self._interrupted, False + yield events_.StreamEnd( + usage=_usage(message), + finish_reason=_finish_reason( + message, interrupted=interrupted + ), + ) + finally: + self._running = False + + async def _drain_interrupted(self) -> None: + if not self._interrupted or self._client is None: + return + self._interrupted = False + with contextlib.suppress(Exception): + async with asyncio.timeout(30): + async for message in self._client.receive_response(): + if isinstance(message, claude_agent_sdk.ResultMessage): + break + + async def steer(self, text: str) -> None: + # `query()` writes straight to the CLI's stdin, so a message can + # land while the turn is still running. + if not self._running or self._client is None: + raise errors.WorkerError(f"{KIND}: no turn is running to steer") + await self._client.query(text) + + async def interrupt(self) -> None: + # Only a running turn: the CLI answers an interrupt with a result, + # which the next turn reads first (see `_drain_interrupted`). + if self._client is None or not self._running: + return + self._interrupted = True + with contextlib.suppress(Exception): + await self._client.interrupt() + + # -- the store ------------------------------------------------------------ + + async def history(self, session_id: str) -> list[messages_.Message]: + """Read the CLI's own transcript, every content block included.""" + await self._require_known(session_id) + stored = await asyncio.to_thread( + claude_agent_sdk.get_session_messages, session_id, self._cwd + ) + out: list[messages_.Message] = [] + for record in stored: + out.extend(_stored_to_messages(record)) + return _handoff.model_view(out) + + async def list_sessions(self) -> list[base.SessionInfo]: + infos = await asyncio.to_thread( + claude_agent_sdk.list_sessions, self._cwd + ) + return [ + base.SessionInfo( + session_id=i.session_id, + title=i.custom_title or i.summary or i.first_prompt, + cwd=i.cwd, + created_at=i.created_at, + updated_at=i.last_modified, + ) + for i in infos + ] + + def tui_command( + self, session_id: str | None + ) -> tuple[list[str], dict[str, str], str | None]: + running_as = session_id or str(uuid.uuid4()) + argv = [self._executable] + argv += ( + ["--resume", session_id] + if session_id + else ["--session-id", running_as] + ) + if self._model: + argv += ["--model", self._model] + if self._instructions: + argv += ["--append-system-prompt", self._instructions] + servers = { + k: v + for k, v in self._mcp_servers.items() + if k != _bridge.SERVER_NAME + } + if servers: + argv += ["--mcp-config", json.dumps({"mcpServers": servers})] + if self._allowed_tools is not None: + argv += ["--tools", ",".join(self._allowed_tools)] + if self._disallowed_tools: + argv += ["--disallowedTools", ",".join(self._disallowed_tools)] + # Measured: the TUI otherwise updates itself in the background and + # announces "Restart to apply" mid-session. + env = {"DISABLE_AUTOUPDATER": "1", **self._env} + return argv, env, running_as + + +async def _content(prompt: base.Prompt) -> str | list[dict[str, Any]]: + """Return a prompt as the CLI takes it: a string, or content blocks. + + Images and PDFs travel as base64 blocks (measured: both read correctly). + """ + if len(prompt) == 1 and isinstance(prompt[0], messages_.TextPart): + return prompt[0].text + blocks: list[dict[str, Any]] = [] + for part in prompt: + if isinstance(part, messages_.TextPart): + blocks.append({"type": "text", "text": part.text}) + elif part.media_type.startswith("text/"): + blocks.append({"type": "text", "text": await base.as_text(part)}) + else: + kind = ( + "image" if part.media_type.startswith("image/") else "document" + ) + blocks.append( + { + "type": kind, + "source": { + "type": "base64", + "media_type": part.media_type, + "data": base64.b64encode(await base.data(part)).decode( + "ascii" + ), + }, + } + ) + return blocks + + +async def _one_message( + content: list[dict[str, Any]], +) -> AsyncIterator[dict[str, Any]]: + """Yield one stream-json user message, the way `query()` sends a str.""" + yield { + "type": "user", + "message": {"role": "user", "content": content}, + "parent_tool_use_id": None, + "session_id": "default", + } + + +def _usage(message: Any) -> usage_.Usage: + raw = message.usage or {} + return usage_.Usage( + input_tokens=int(raw.get("input_tokens", 0)), + output_tokens=int(raw.get("output_tokens", 0)), + cache_read_tokens=int(raw.get("cache_read_input_tokens", 0)), + cache_write_tokens=int(raw.get("cache_creation_input_tokens", 0)), + # Cost has no home in Usage; the CLI's own number rides in `raw`. + raw={"cost_usd": message.total_cost_usd, "claude": raw}, + ) + + +def _finish_reason(message: Any, *, interrupted: bool) -> str: + subtype = message.subtype or "" + if interrupted: + # We asked for this. The CLI reports an error result after an + # interrupt; it is not a failure. + return "cancelled" + if subtype == "error_max_turns": + return "length" + if getattr(message, "is_error", False): + detail = str(getattr(message, "result", None) or subtype) + raise errors.turn_failed(KIND, detail, AUTH_HINT) + return "stop" + + +def _translate(exc: BaseException) -> BaseException: + """Turn a claude-agent-sdk exception into a worker error.""" + if isinstance(exc, claude_agent_sdk.CLINotFoundError): + return errors.ExecutableMissingError(KIND, "claude", INSTALL) + if isinstance(exc, claude_agent_sdk.ProcessError): + code = getattr(exc, "exit_code", None) + return errors.WorkerError(f"{KIND} crashed (exit {code}): {exc}") + if isinstance(exc, claude_agent_sdk.CLIConnectionError): + return errors.WorkerError(f"{KIND}: lost the CLI connection: {exc}") + if isinstance(exc, claude_agent_sdk.CLIJSONDecodeError): + return errors.WorkerError(f"{KIND}: unparseable CLI output: {exc}") + return exc + + +async def _guard(stream: Any) -> Any: + """Re-raise claude-agent-sdk failures as worker errors, mid-stream too.""" + try: + async for item in stream: + yield item + except claude_agent_sdk.ClaudeSDKError as exc: + raise _translate(exc) from exc + + +def _stream_events( + raw: dict[str, Any], open_blocks: dict[int, str] +) -> list[events_.Event]: + """One raw Anthropic stream event -> text and reasoning events. + + `open_blocks` tracks which content-block index is text or thinking, so + the closing event emits the matching End. + """ + kind = raw.get("type") + index = raw.get("index", 0) + out: list[events_.Event] = [] + block_id = f"b{index}" + if kind == "content_block_start": + block_type = (raw.get("content_block") or {}).get("type") + if block_type == "text": + open_blocks[index] = "text" + out.append(events_.TextStart(block_id=block_id)) + elif block_type == "thinking": + open_blocks[index] = "thinking" + out.append(events_.ReasoningStart(block_id=block_id)) + elif kind == "content_block_delta": + delta = raw.get("delta") or {} + if delta.get("type") == "text_delta" and delta.get("text"): + out.append( + events_.TextDelta(chunk=delta["text"], block_id=block_id) + ) + elif delta.get("type") == "thinking_delta" and delta.get("thinking"): + out.append( + events_.ReasoningDelta( + chunk=delta["thinking"], block_id=block_id + ) + ) + elif kind == "content_block_stop": + opened = open_blocks.pop(index, None) + if opened == "text": + out.append(events_.TextEnd(block_id=block_id)) + elif opened == "thinking": + out.append(events_.ReasoningEnd(block_id=block_id)) + return out + + +def _stored_to_messages(record: Any) -> list[messages_.Message]: + """One stored transcript record -> messages, shaped like a live turn's. + + Tool results become their own `tool` message; `_handoff.model_view` + then folds them into the model view. + """ + raw = record.message if isinstance(record.message, dict) else {} + role = raw.get("role") or record.type + content = raw.get("content") + if isinstance(content, str): + blocks: list[Any] = [{"type": "text", "text": content}] + else: + blocks = content if isinstance(content, list) else [] + parts: list[Any] = [] + tool_results: list[Any] = [] + for block in blocks: + if not isinstance(block, dict): + continue + kind = block.get("type") + if kind == "text" and block.get("text"): + parts.append(messages_.TextPart(text=block["text"])) + elif kind in ("image", "document"): + found = block.get("source") + source: dict[str, Any] = found if isinstance(found, dict) else {} + if source.get("type") == "base64" and source.get("data"): + fallback = "image/png" if kind == "image" else "application/pdf" + parts.append( + messages_.FilePart( + data=str(source["data"]), + media_type=str(source.get("media_type") or fallback), + ) + ) + elif source.get("type") == "url" and source.get("url"): + parts.append(messages_.FilePart.from_url(str(source["url"]))) + elif kind == "thinking" and block.get("thinking"): + parts.append(messages_.ReasoningPart(text=block["thinking"])) + elif kind == "tool_use": + parts.append( + messages_.ToolCallPart( + tool_call_id=block.get("id", ""), + tool_name=block.get("name", ""), + tool_args=json.dumps(block.get("input") or {}), + ) + ) + elif kind == "tool_result": + tool_results.append( + messages_.ToolResultPart( + tool_call_id=block.get("tool_use_id", ""), + tool_name="", + result=block.get("content"), + result_kind="error" if block.get("is_error") else "json", + ) + ) + out: list[messages_.Message] = [] + if parts and role in ("user", "assistant"): + out.append(messages_.Message(role=role, parts=parts)) + if tool_results: + out.append(messages_.Message(role="tool", parts=tool_results)) + return out diff --git a/src/ai/workers/_adapters/pi.py b/src/ai/workers/_adapters/pi.py new file mode 100644 index 00000000..6c64c731 --- /dev/null +++ b/src/ai/workers/_adapters/pi.py @@ -0,0 +1,1111 @@ +"""Pi, driven through `pi --mode rpc`. + +One pi process per conversation: an RPC process holds one session. The id we +hand pi (`--session-id`) IS its session id, so `pi --session-id ` opens +the same conversation in pi's own TUI. + +pi has no permission prompt. Approval is an extension loaded into every pi +a worker starts (`pi_extension.js`): its `tool_call` handler asks through an +RPC dialog, and this adapter answers it. The protocol shapes here were +measured against a real pi (see VERIFIED_VERSION). +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import datetime +import json +import os +import pathlib +import re +import uuid +import warnings +from collections.abc import AsyncIterator +from typing import Any + +from ...types import events as events_ +from ...types import messages as messages_ +from ...types import usage as usage_ +from .. import _bridge +from .. import _errors as errors +from . import _handoff, base + +KIND = "pi" +INSTALL = "npm install -g @earendil-works/pi-coding-agent" +VERIFIED_VERSION = "1.3.0" +"""The CLI the live tests last passed against. Older ones get a warning.""" +EXTENSION = str(pathlib.Path(__file__).with_name("pi_extension.js")) +APPROVE = "ai-python:approve" +"""The dialog title the extension asks approvals under.""" +READY = "ai-python-ready" +"""The command the extension registers last: proof that it loaded.""" +CONFIG_VAR = "AI_PYTHON_PI" +"""The environment variable that configures the extension.""" +SESSION_ID = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9._-]{0,198}[A-Za-z0-9])?\Z") +"""pi's rule for a session id: letters, digits, `.`, `_` and `-`, starting +and ending with a letter or digit. Never a path.""" +DIALOGS = ("select", "confirm", "input", "editor") +"""Extension UI methods that wait for an answer; the rest are notices.""" +CRASHED = "__crashed__" +CLOSED = "__closed__" + +AUTH_HINT = ( + "pi uses its own login (`/login`, auth.json) or a provider key in the " + "environment: pass env={'ANTHROPIC_API_KEY': ...}." +) + + +class PiAdapter: + def __init__( + self, + *, + cwd: str, + env: dict[str, str], + executable: str, + model: str | None, + effort: str | None, + mcp_servers: dict[str, Any], + approve: base.Approve | None, + allowed_tools: list[str] | None, + disallowed_tools: list[str] | None, + ) -> None: + self._cwd = cwd + self._env = env + self._executable = executable + self._model = model + self._effort = effort + self._mcp_servers = mcp_servers + self._approve = approve + self._allowed_tools = allowed_tools + self._disallowed_tools = disallowed_tools + self._instructions: str | None = None + self._agent: _Agent | None = None + # We ended the turn: it settles as cancelled. + self._interrupted = False + # A prompt was accepted and no `agent_settled` came yet. + self._running = False + self.version: str | None = None + + async def start(self) -> None: + self.version = await base.probe( + KIND, + self._executable, + install=INSTALL, + verified=VERIFIED_VERSION, + cwd=self._cwd, + env=self._env, + ) + + async def close(self) -> None: + agent, self._agent = self._agent, None + self._running = False + if agent is not None: + await agent.close() + with contextlib.suppress(Exception): + await agent.process.terminate() + + def _args(self) -> list[str]: + """How pi runs: headless and in the TUI alike.""" + args = ["-e", EXTENSION] + if self._model: + args += ["--model", self._model] + if self._effort: + args += ["--thinking", self._effort] + if self._instructions: + args += ["--append-system-prompt", self._instructions] + if self._allowed_tools is not None: + args += ( + ["--tools", ",".join(self._allowed_tools)] + if self._allowed_tools + else ["--no-tools"] + ) + if self._disallowed_tools: + args += ["--exclude-tools", ",".join(self._disallowed_tools)] + return args + + def _extension_env( + self, *, approve: bool, servers: dict[str, Any] + ) -> dict[str, str]: + config: dict[str, Any] = {"approve": approve} + if servers: + config["mcp"] = mcp_config(servers) + return {**self._env, CONFIG_VAR: json.dumps(config)} + + # -- the conversation ----------------------------------------------------- + + async def open( + self, + *, + session_id: str | None, + fork: bool, + seed: list[messages_.Message] | None, + instructions: str | None, + ) -> str: + self._instructions = instructions + if seed: + session_id = await self._stage(seed) + args = ["--session-id", session_id] + elif session_id is not None: + # Checked first: `--session-id` CREATES a session it cannot + # find, so a typo would silently start an empty conversation. + found = await self._find(session_id) + if fork: + session_id = str(uuid.uuid4()) + # By path, not id: pi matches a partial id, a path is exact. + args = ["--fork", found["path"], "--session-id", session_id] + else: + args = ["--session-id", session_id] + else: + session_id = str(uuid.uuid4()) + args = ["--session-id", session_id] + await self._spawn(session_id, args) + return session_id + + async def _spawn(self, session_id: str, args: list[str]) -> None: + """Start pi on one conversation and wait until it answers. + + A pi that cannot start (an unknown model, a bad flag) exits at once, + and the handshake reports its last words. + """ + process = await base.spawn( + [self._executable, "--mode", "rpc", *args, *self._args()], + cwd=self._cwd, + env=self._extension_env( + approve=self._approve is not None, servers=self._mcp_servers + ), + ) + agent = _Agent(process, on_dialog=self._on_dialog) + agent.start() + try: + state = await agent.call({"type": "get_state"}, timeout=60) + reported = (state or {}).get("sessionId") + if reported and reported != session_id: + raise errors.WorkerError( + f"{KIND} opened session {reported}, not {session_id}" + ) + # The approval gate must be there before any work is accepted: + # a gate that failed to load is a gate that fails open. + listed = await agent.call({"type": "get_commands"}, timeout=60) + names = { + c.get("name") for c in (listed or {}).get("commands") or [] + } + if READY not in names: + raise errors.WorkerError( + f"{KIND} did not load the SDK's extension ({EXTENSION}); " + "refusing to run without it" + ) + except BaseException: + await agent.close() + with contextlib.suppress(Exception): + await process.terminate() + raise + self._agent = agent + + async def _stage(self, history: list[messages_.Message]) -> str: + """Write a session file with this past; start nothing on it.""" + session_id = str(uuid.uuid4()) + _handoff.warn_unsupported_files(history, KIND) + dropped = _handoff.reasoning_count(history) + if dropped: + # pi drops thinking another model produced, and replaying it as + # user text is refused by Anthropic's safeguards. + warnings.warn( + f"{dropped} reasoning part(s) in the history were not carried " + "into pi: a model only reads its own reasoning. Text, tool " + "calls and images were carried in full.", + stacklevel=2, + ) + directory = self._session_dir() + now = datetime.datetime.now(datetime.UTC) + stamp = now.isoformat(timespec="milliseconds").replace("+00:00", "Z") + lines: list[dict[str, Any]] = [ + { + "type": "session", + "version": 3, + "id": session_id, + "timestamp": stamp, + "cwd": os.path.realpath(self._cwd), + } + ] + parent: str | None = None + # The model the conversation was on, as pi records a switch: pi + # restores a session's model from its last reply's author. Measured: + # without it, a seeded conversation ran on pi's default model. + written: dict[str, Any] = next( + ( + found + for m in reversed(history) + if m.role == "assistant" + and isinstance( + found := (m.provider_metadata or {}).get(KIND), dict + ) + ), + {}, + ) + changes: list[dict[str, Any]] = [] + if written.get("provider") and written.get("model"): + changes.append( + { + "type": "model_change", + "provider": written["provider"], + "modelId": written["model"], + } + ) + if written.get("thinkingLevel"): + changes.append( + { + "type": "thinking_level_change", + "thinkingLevel": written["thinkingLevel"], + } + ) + entries = [ + *changes, + *( + {"type": "message", "message": message} + for message in _handoff.to_pi_messages( + history, timestamp=int(now.timestamp() * 1000) + ) + ), + ] + for entry in entries: + entry_id = uuid.uuid4().hex[:8] + lines.append( + { + **entry, + "id": entry_id, + "parentId": parent, + "timestamp": stamp, + } + ) + parent = entry_id + name = stamp.replace(":", "-").replace(".", "-") + path = directory / f"{name}_{session_id}.jsonl" + text = "".join(json.dumps(line) + "\n" for line in lines) + + def write() -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + + await asyncio.to_thread(write) + return session_id + + # -- approval ------------------------------------------------------------- + + async def _on_dialog(self, agent: _Agent, request: dict[str, Any]) -> None: + if request.get("method") not in DIALOGS: + return # a notice: nothing waits for an answer + reply: dict[str, Any] = { + "type": "extension_ui_response", + "id": request.get("id"), + } + if ( + request.get("method") != "input" + or request.get("title") != APPROVE + or self._approve is None + or self._interrupted + ): + # Another extension asking a person, who is not there; or the + # turn is being stopped, and `abort` would wait on the question. + agent.pending.pop(str(request.get("id")), None) + await agent.send({**reply, "cancelled": True}) + return + try: + payload = json.loads(request.get("placeholder") or "{}") + except json.JSONDecodeError: + payload = {} + granted, reason = await self._approve( + str(payload.get("toolCallId") or ""), + str(payload.get("toolName") or ""), + dict(payload.get("input") or {}), + ) + decision = ( + {"decision": "allow"} + if granted + else {"decision": "deny", "reason": reason} + ) + # Answered from here on: `dismiss` must not withdraw it too. + agent.pending.pop(str(request.get("id")), None) + await agent.send({**reply, "value": json.dumps(decision)}) + + # -- turns ---------------------------------------------------------------- + + async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: + base.check_files(prompt, KIND, ("image/",)) + agent = self._agent + assert agent is not None + if agent.dead is not None: + raise errors.WorkerError(f"{KIND} crashed: {agent.dead}") + # What the last turn left behind (a stopped turn's tail) is not this + # turn's: `abort` answers only once pi is idle, so it has all arrived. + while not agent.events.empty(): + agent.events.get_nowait() + self._interrupted = False + command = await _prompt_command(prompt) + try: + accepted = await agent.call(command) + except _CommandError as exc: + raise errors.turn_failed(KIND, str(exc), AUTH_HINT) from exc + if (accepted or {}).get("disposition") == "handled": + # An extension command consumed it. It may start a run of its own + # (measured: then pi is already streaming, and that run is ours). + state = await agent.call({"type": "get_state"}) + if not (state or {}).get("isStreaming"): + yield events_.StreamStart() + yield events_.StreamEnd(finish_reason="stop") + return + self._running = True + usage = usage_.Usage() + cost = 0.0 + started: set[str] = set() + streamed: set[str] = set() + open_blocks: dict[str, str] = {} + message_no = 0 + last: dict[str, Any] = {} + retry_error = "" + settled: dict[str, Any] = {} + try: + yield events_.StreamStart() + while True: + event = await agent.events.get() + kind = event.get("type") + if kind == CRASHED: + raise errors.WorkerError( + f"{KIND} crashed: {event.get('detail')}" + ) + if kind == CLOSED: + raise errors.WorkerError( + f"{KIND}: the worker closed while its turn ran" + ) + message = event.get("message") or {} + if kind == "message_start" and message.get("role") == ( + "assistant" + ): + message_no += 1 + elif kind == "message_update": + for out in _stream_events( + event.get("assistantMessageEvent") or {}, + f"{message_no}", + open_blocks, + started, + ): + if isinstance( + out, events_.TextStart | events_.ReasoningStart + ): + streamed.add(out.block_id) + yield out + elif kind == "message_end" and message.get("role") == ( + "assistant" + ): + last = message + raw = message.get("usage") or {} + usage = usage + _usage(raw) + cost += float((raw.get("cost") or {}).get("total") or 0) + for out in _close_blocks(open_blocks): + yield out + if message.get("stopReason") != "error": + # What the message holds that was never streamed (a + # run an extension command started, measured). + for out in _unstreamed( + message.get("content"), + f"{message_no}", + streamed, + started, + ): + yield out + elif kind == "tool_execution_end" and not event.get( + "parentToolCallId" + ): + # A nested call (a codemode script's) is not in the + # transcript: only the tool that made it reports it. + call_id = str(event.get("toolCallId") or "") + name = str(event.get("toolName") or "") + yield base.builtin_result( + call_id, + name, + _tool_result(event.get("result") or {}), + is_error=bool(event.get("isError")), + ) + elif kind == "auto_retry_end" and not event.get("success"): + retry_error = str(event.get("finalError") or "") + elif kind == "agent_settled": + settled = event + break + finally: + self._running = False + interrupted = self._interrupted or bool(settled.get("aborted")) + self._interrupted = False + stop = last.get("stopReason") + if stop == "error" and not interrupted: + raise errors.turn_failed( + KIND, + str(last.get("errorMessage") or retry_error), + AUTH_HINT, + ) + yield events_.StreamEnd( + # Cost has no home in Usage; pi's own number rides in `raw`. + usage=usage.model_copy(update={"raw": {"cost_usd": cost}}), + finish_reason=( + "cancelled" + if interrupted or stop == "aborted" + else "length" + if stop == "length" + else "stop" + ), + provider_metadata=_written_by(last), + ) + + async def steer(self, text: str) -> None: + """Queue guidance for the running turn. + + pi delivers it once the current message's tool calls are done, + before its next model call. + """ + if not self._running or self._agent is None: + raise errors.WorkerError(f"{KIND}: no turn is running to steer") + try: + await self._agent.call({"type": "steer", "message": text}) + except _CommandError as exc: + raise errors.WorkerError(f"{KIND}: steer rejected: {exc}") from exc + + async def interrupt(self) -> None: + self._interrupted = True + agent = self._agent + if agent is None or agent.dead is not None: + return + # A tool call waiting on approval keeps pi busy, and `abort` only + # answers once pi is idle (measured: the turn ran on). Withdraw the + # question first: the call is denied, and the abort lands at once. + await agent.dismiss() + with contextlib.suppress(Exception): + await agent.call({"type": "abort"}, timeout=30) + + # -- the store ------------------------------------------------------------ + + def _session_dir(self) -> pathlib.Path: + """Where pi keeps this directory's sessions, decided as pi does. + + The environment, then project settings, then global ones; else one + directory per working directory under the agent dir. + """ + env = {**os.environ, **self._env} + cwd = os.path.realpath(self._cwd) + agent = pathlib.Path( + os.path.expanduser(env.get("PI_CODING_AGENT_DIR") or "~/.pi/agent") + ) + custom = ( + env.get("PI_CODING_AGENT_SESSION_DIR") + or _setting(pathlib.Path(cwd, ".pi", "settings.json")) + or _setting(agent / "settings.json") + ) + if custom: + return pathlib.Path(cwd, os.path.expanduser(custom)) + safe = cwd.lstrip("/\\") + for c in "/\\:": + safe = safe.replace(c, "-") + return agent / "sessions" / f"--{safe}--" + + def _scan(self, *, titles: bool) -> list[dict[str, Any]]: + """Describe the sessions of this directory, newest first.""" + directory = self._session_dir() + cwd = os.path.realpath(self._cwd) + out: list[dict[str, Any]] = [] + try: + paths = sorted(directory.glob("*.jsonl")) + except OSError: + return [] + for path in paths: + with contextlib.suppress(Exception): + info = _describe(path, (cwd, self._cwd), titles=titles) + if info is not None: + out.append(info) + return sorted(out, key=lambda s: -s["updated"]) + + async def _find(self, session_id: str) -> dict[str, Any]: + if not SESSION_ID.match(session_id): + raise errors.WorkerError( + f"{KIND}: {session_id!r} is not a pi session id" + ) + found = next( + ( + s + for s in await asyncio.to_thread(self._scan, titles=False) + if s["id"] == session_id + ), + None, + ) + if found is None: + raise errors.WorkerError( + f"{KIND}: no session {session_id} in {self._cwd}" + ) + return found + + async def history(self, session_id: str) -> list[messages_.Message]: + """Read the active branch of the session file. + + Every message on it, root to latest entry, including what a + compaction summarized away, which pi keeps on disk. + """ + found = await self._find(session_id) + text = await asyncio.to_thread(pathlib.Path(found["path"]).read_text) + stored: list[messages_.Message] = [] + for record in _active_branch(text): + stored.extend(_stored_to_messages(record)) + return _handoff.model_view(stored) + + async def list_sessions(self) -> list[base.SessionInfo]: + return [ + base.SessionInfo( + session_id=s["id"], + title=s.get("title"), + cwd=s.get("cwd"), + created_at=s.get("created"), + updated_at=s.get("updated"), + ) + for s in await asyncio.to_thread(self._scan, titles=True) + ] + + def tui_command( + self, session_id: str | None + ) -> tuple[list[str], dict[str, str], str | None]: + running_as = session_id or str(uuid.uuid4()) + argv = [self._executable, "--session-id", running_as, *self._args()] + servers = { + k: v + for k, v in self._mcp_servers.items() + if k != _bridge.SERVER_NAME + } + # The person at the TUI is the gate: the extension asks nothing. + return ( + argv, + self._extension_env(approve=False, servers=servers), + running_as, + ) + + +class _CommandError(Exception): + """pi answered a command with `success: false`.""" + + +class _Agent: + """One `pi --mode rpc` process: its commands, its events, its dialogs. + + Responses resolve the command that asked (by id). Session events go on + `events`, in order. A dialog is answered in a task of its own: an + approval can take minutes, and the reader must keep reading meanwhile, + including the reply to an `abort` sent from inside it. + """ + + def __init__(self, process: base.Process, *, on_dialog: Any) -> None: + self.process = process + self.events: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self.dead: str | None = None + self._on_dialog = on_dialog + self._waiters: dict[str, asyncio.Future[Any]] = {} + self._pump: asyncio.Task[None] | None = None + self._dialogs: set[asyncio.Task[None]] = set() + # Dialogs pi still waits on, by request id: what `dismiss` withdraws. + self.pending: dict[str, asyncio.Task[None]] = {} + + def start(self) -> None: + self._pump = asyncio.create_task(self._read()) + + async def call( + self, command: dict[str, Any], *, timeout: float = 120 + ) -> Any: + if self.dead is not None: + raise errors.WorkerError(f"{KIND} crashed: {self.dead}") + request_id = uuid.uuid4().hex + future: asyncio.Future[Any] = asyncio.get_running_loop().create_future() + self._waiters[request_id] = future + try: + await self.send({**command, "id": request_id}) + async with asyncio.timeout(timeout): + return await future + finally: + self._waiters.pop(request_id, None) + + async def send(self, record: dict[str, Any]) -> None: + try: + await self.process.write(json.dumps(record) + "\n") + except Exception as exc: + raise errors.WorkerError( + f"{KIND}: could not write to pi: {exc}" + ) from exc + + def _track(self, request_id: str, task: asyncio.Task[None]) -> None: + self.pending[request_id] = task + + def forget(_: asyncio.Task[None]) -> None: + if self.pending.get(request_id) is task: + del self.pending[request_id] + + task.add_done_callback(forget) + + async def dismiss(self) -> None: + """Withdraw every question pi still waits on. + + The approval deciding it is cancelled, and pi hears `cancelled`: the + extension then blocks the call. + """ + for request_id, task in list(self.pending.items()): + self.pending.pop(request_id, None) + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + with contextlib.suppress(Exception): + await self.send( + { + "type": "extension_ui_response", + "id": request_id, + "cancelled": True, + } + ) + + async def close(self) -> None: + if self._pump is not None: + self._pump.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._pump + self._pump = None + for task in list(self._dialogs): + task.cancel() + for task in list(self._dialogs): + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + self._dialogs.clear() + self._fail("connection closed") + # A turn reading events must hear it too: the producer is gone. + self.events.put_nowait({"type": CLOSED}) + + def _fail(self, detail: str) -> None: + self.dead = self.dead or detail + for future in self._waiters.values(): + if not future.done(): + future.set_exception( + errors.WorkerError(f"{KIND} crashed: {detail}") + ) + self._waiters.clear() + + async def _read(self) -> None: + try: + while True: + line = await self.process.readline() + if not line: + detail = await self.process.death_note() + self._fail(detail) + self.events.put_nowait({"type": CRASHED, "detail": detail}) + return + try: + record = json.loads(line) + except json.JSONDecodeError: + continue # stdout is protocol only; tolerate noise + if not isinstance(record, dict): + continue + kind = record.get("type") + if kind == "response": + self._resolve(record) + elif kind == "extension_ui_request": + task = asyncio.create_task(self._on_dialog(self, record)) + self._dialogs.add(task) + task.add_done_callback(self._dialogs.discard) + if record.get("method") in DIALOGS: + self._track(str(record.get("id")), task) + else: + self.events.put_nowait(record) + except asyncio.CancelledError: + raise + except Exception as exc: + self._fail(str(exc)) + self.events.put_nowait({"type": CRASHED, "detail": str(exc)}) + + def _resolve(self, record: dict[str, Any]) -> None: + future = self._waiters.get(str(record.get("id"))) + if future is None or future.done(): + return + if record.get("success") is False: + future.set_exception(_CommandError(str(record.get("error") or ""))) + else: + future.set_result(record.get("data")) + + +async def _prompt_command(prompt: base.Prompt) -> dict[str, Any]: + """Return a prompt as pi's `prompt` command: text, and images as base64.""" + texts: list[str] = [] + images: list[dict[str, Any]] = [] + for part in prompt: + if isinstance(part, messages_.TextPart): + texts.append(part.text) + elif part.media_type.startswith("text/"): + texts.append(await base.as_text(part)) + else: + images.append( + { + "type": "image", + "data": base64.b64encode(await base.data(part)).decode( + "ascii" + ), + "mimeType": part.media_type, + } + ) + command: dict[str, Any] = {"type": "prompt", "message": "\n\n".join(texts)} + if images: + command["images"] = images + return command + + +def _stream_events( + update: dict[str, Any], + message: str, + open_blocks: dict[str, str], + started: set[str], +) -> list[events_.Event]: + """One `assistantMessageEvent` -> events. + + Block ids are `.`: pi numbers content blocks per + assistant message, and a turn holds several. Tool-call arguments travel + as ONE delta, the final arguments, at `toolcall_end`: streamed fragments + are not guaranteed to add up to what pi settles on. + """ + kind = update.get("type") or "" + block = f"{message}.{update.get('contentIndex', 0)}" + out: list[events_.Event] = [] + if kind == "text_start": + open_blocks[block] = "text" + out.append(events_.TextStart(block_id=block)) + elif kind == "text_delta" and update.get("delta"): + out.append(events_.TextDelta(chunk=update["delta"], block_id=block)) + elif kind == "thinking_start": + open_blocks[block] = "thinking" + out.append(events_.ReasoningStart(block_id=block)) + elif kind == "thinking_delta" and update.get("delta"): + out.append( + events_.ReasoningDelta(chunk=update["delta"], block_id=block) + ) + elif kind in ("text_end", "thinking_end"): + out += _close_blocks(open_blocks, only=block) + elif kind == "toolcall_start" and update.get("id"): + started.add(update["id"]) + out.append( + events_.BuiltinToolStart( + tool_call_id=update["id"], + tool_name=update.get("toolName") or "", + ) + ) + elif kind == "toolcall_end": + call = _tool_call(update.get("toolCall") or {}) + out += base.builtin_call( + call.tool_call_id, call.tool_name, call.tool_args, started=started + ) + return out + + +def _unstreamed( + content: Any, message: str, streamed: set[str], started: set[str] +) -> list[events_.Event]: + """Return the events of a finished message's unstreamed blocks.""" + out: list[events_.Event] = [] + blocks = content if isinstance(content, list) else [] + for index, block in enumerate(blocks): + if not isinstance(block, dict): + continue + block_id = f"{message}.{index}" + kind = block.get("type") + if kind == "text" and block.get("text") and block_id not in streamed: + out += [ + events_.TextStart(block_id=block_id), + events_.TextDelta(chunk=block["text"], block_id=block_id), + events_.TextEnd(block_id=block_id), + ] + elif ( + kind == "thinking" + and block.get("thinking") + and block_id not in streamed + ): + out += [ + events_.ReasoningStart(block_id=block_id), + events_.ReasoningDelta( + chunk=block["thinking"], block_id=block_id + ), + events_.ReasoningEnd(block_id=block_id), + ] + elif kind == "toolCall" and str(block.get("id")) not in started: + call = _tool_call(block) + out += base.builtin_call( + call.tool_call_id, + call.tool_name, + call.tool_args, + started=started, + ) + return out + + +def _close_blocks( + open_blocks: dict[str, str], only: str | None = None +) -> list[events_.Event]: + """End the text and reasoning blocks still open (an aborted message's).""" + out: list[events_.Event] = [] + for block in [only] if only is not None else list(open_blocks): + opened = open_blocks.pop(block, None) + if opened == "text": + out.append(events_.TextEnd(block_id=block)) + elif opened == "thinking": + out.append(events_.ReasoningEnd(block_id=block)) + return out + + +def _tool_call(block: dict[str, Any]) -> messages_.ToolCallPart: + return messages_.ToolCallPart( + tool_call_id=str(block.get("id") or ""), + tool_name=str(block.get("name") or ""), + tool_args=json.dumps(block.get("arguments") or {}), + ) + + +def _assistant_parts(content: Any) -> list[messages_.Part]: + parts: list[messages_.Part] = [] + for block in content if isinstance(content, list) else []: + if not isinstance(block, dict): + continue + kind = block.get("type") + if kind == "text" and block.get("text"): + parts.append(messages_.TextPart(text=block["text"])) + elif kind == "thinking" and block.get("thinking"): + parts.append(messages_.ReasoningPart(text=block["thinking"])) + elif kind == "toolCall": + parts.append(_tool_call(block)) + return parts + + +def _content_parts(content: Any) -> list[messages_.Part]: + """Return a user message's parts: from a string, or text and images.""" + if isinstance(content, str): + return [messages_.TextPart(text=content)] if content else [] + parts: list[messages_.Part] = [] + for block in content if isinstance(content, list) else []: + if not isinstance(block, dict): + continue + if block.get("type") == "text" and block.get("text"): + parts.append(messages_.TextPart(text=block["text"])) + elif block.get("type") == "image" and block.get("data"): + parts.append( + messages_.FilePart( + data=str(block["data"]), + media_type=str(block.get("mimeType") or "image/png"), + ) + ) + return parts + + +def _tool_result(result: Any) -> Any: + """Return what a tool returned: its text, when that is all it is. + + pi tools return content blocks; text-only content (nearly every call) + reads as the string it is, anything else stays the blocks themselves. + """ + content = result.get("content") if isinstance(result, dict) else result + if isinstance(content, list) and all( + isinstance(b, dict) and b.get("type") == "text" for b in content + ): + content = "".join(b.get("text") or "" for b in content) + return content + + +def _active_branch(text: str) -> list[dict[str, Any]]: + """Return the session file's active branch: root to latest entry. + + A session is a tree (`id`/`parentId`); the leaf is the entry appended + last, and abandoned branches are not part of the conversation. + """ + entries: dict[str, dict[str, Any]] = {} + leaf: str | None = None + for line in text.splitlines(): + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if ( + isinstance(entry, dict) + and entry.get("id") + and entry.get("type") != "session" + ): + entries[str(entry["id"])] = entry + leaf = str(entry["id"]) + path: list[dict[str, Any]] = [] + seen: set[str] = set() + while leaf is not None and leaf in entries and leaf not in seen: + seen.add(leaf) + path.append(entries[leaf]) + parent = entries[leaf].get("parentId") + leaf = str(parent) if parent is not None else None + return list(reversed(path)) + + +def _stored_to_messages(entry: dict[str, Any]) -> list[messages_.Message]: + """One session entry -> messages, CLI tools as tool calls. + + Everything pi sends the model is kept: what an extension injected + (`custom_message`) and a `!` command's output (`bashExecution`) reach it + as user messages, and so they are user messages here; a branch summary + is the context pi carried over from an abandoned branch. + """ + kind = entry.get("type") + if kind == "custom_message": + parts = _content_parts(entry.get("content")) + return [messages_.Message(role="user", parts=parts)] if parts else [] + if kind == "branch_summary" and entry.get("summary"): + text = str(entry["summary"]) + return [ + messages_.Message( + role="user", parts=[messages_.TextPart(text=text)] + ) + ] + if kind != "message": + return [] + message = entry.get("message") or {} + role = message.get("role") + if role in ("custom", "user"): + parts = _content_parts(message.get("content")) + return [messages_.Message(role="user", parts=parts)] if parts else [] + if role == "bashExecution" and not message.get("excludeFromContext"): + # pi's own rendering of it for the model. + text = ( + f"Ran `{message.get('command', '')}`\n```\n" + f"{message.get('output', '')}\n```" + ) + return [ + messages_.Message( + role="user", parts=[messages_.TextPart(text=text)] + ) + ] + if role == "assistant": + parts = _assistant_parts(message.get("content")) + if not parts: + return [] + return [ + messages_.Message( + role="assistant", + parts=parts, + provider_metadata=_written_by(message), + ) + ] + if role == "toolResult": + result = messages_.ToolResultPart( + tool_call_id=str(message.get("toolCallId") or ""), + tool_name=str(message.get("toolName") or ""), + result=_tool_result(message), + result_kind="error" if message.get("isError") else "json", + ) + return [messages_.Message(role="tool", parts=[result])] + return [] + + +def _written_by(message: dict[str, Any]) -> dict[str, Any] | None: + """Which provider and model wrote an assistant message, as pi says. + + Kept on the message (`provider_metadata["pi"]`) so a conversation + seeded into another pi continues on the model it was on. + """ + written = { + k: message[k] + for k in ("api", "provider", "model", "thinkingLevel") + if isinstance(message.get(k), str) + } + if not written or written.get("provider") == "ai-python": + return None + return {KIND: written} + + +def _usage(raw: dict[str, Any]) -> usage_.Usage: + """Return pi's per-message usage. pi's `input` excludes cached input.""" + + def count(name: str) -> int: + value = raw.get(name) + return value if isinstance(value, int) else 0 + + reasoning = raw.get("reasoning") + return usage_.Usage( + input_tokens=count("input") + count("cacheRead") + count("cacheWrite"), + output_tokens=count("output"), + cache_read_tokens=count("cacheRead"), + cache_write_tokens=count("cacheWrite"), + reasoning_tokens=reasoning if isinstance(reasoning, int) else None, + ) + + +def mcp_config(servers: dict[str, Any]) -> dict[str, Any]: + """MCP servers, in the shape claude takes, as pi registers them. + + pi validates the config, so only the keys it knows are passed. + """ + keys = ("command", "args", "env", "cwd", "url", "headers") + return { + name: {k: server[k] for k in keys if server.get(k) is not None} + for name, server in servers.items() + } + + +def _setting(path: pathlib.Path) -> str | None: + """Return the `sessionDir` a pi settings file names, if any.""" + try: + value = json.loads(path.read_text()).get("sessionDir") + except Exception: + return None + return value if isinstance(value, str) and value else None + + +def _describe( + path: pathlib.Path, cwds: tuple[str, ...], *, titles: bool +) -> dict[str, Any] | None: + """Describe one session file of this directory, or None. + + A session file can come from anyone (a project's settings choose the + directory), and its id becomes a CLI argument: one that breaks pi's + rule is not a session. + """ + with path.open() as f: + header = json.loads(f.readline() or "{}") + if header.get("type") != "session" or not SESSION_ID.match( + str(header.get("id") or "") + ): + return None + if header.get("cwd") not in cwds: + return None + created: int | None = None + with contextlib.suppress(ValueError): + when = datetime.datetime.strptime( + str(header.get("timestamp") or "")[:19], "%Y-%m-%dT%H:%M:%S" + ).replace(tzinfo=datetime.UTC) + created = int(when.timestamp() * 1000) + info: dict[str, Any] = { + "id": header["id"], + "path": str(path), + "cwd": header.get("cwd"), + "created": created, + "updated": int(path.stat().st_mtime * 1000), + } + if not titles: + return info + name: str | None = None + first: str | None = None + for line in f: + try: + entry = json.loads(line) + except ValueError: + continue + if entry.get("type") == "session_info": + name = entry.get("name") or None + elif first is None and entry.get("type") == "message": + message = entry.get("message") or {} + if message.get("role") == "user": + content = message.get("content") + if isinstance(content, list): + content = " ".join( + b.get("text", "") + for b in content + if isinstance(b, dict) and b.get("type") == "text" + ) + first = (content or "").strip()[:200] or None + info["title"] = name or first + return info diff --git a/src/ai/workers/_adapters/pi_extension.js b/src/ai/workers/_adapters/pi_extension.js new file mode 100644 index 00000000..164712c0 --- /dev/null +++ b/src/ai/workers/_adapters/pi_extension.js @@ -0,0 +1,88 @@ +// Loaded into every `pi` the ai-python SDK launches (`pi -e `). +// It does nothing on its own: AI_PYTHON_PI, a JSON object in the +// environment, says which of its three jobs this process needs. +// +// - "approve": ask the SDK about every tool call. pi has no permission +// channel of its own; an extension's `tool_call` handler is the gate, and +// an RPC dialog (`ctx.ui.input`) is how it reaches the SDK. The answer is +// a JSON decision: allow (optionally with rewritten input) or deny. +// - "mcp": MCP servers, registered for this session only. +// - "gateway": the Vercel AI Gateway, reached with `Authorization: Bearer` +// rather than the `x-api-key` header pi's provider sends, so a sandbox +// firewall that rewrites `Authorization` can inject the credential. +// `baseUrl` moves the provider when the gateway is not the default one. +const APPROVE = "ai-python:approve"; +// Registered LAST: its presence in `get_commands` is how the SDK knows every +// part above loaded — above all the approval gate. No marker, no session. +const READY = "ai-python-ready"; + +export default function (pi) { + // Unparseable configuration fails the load (and so the session) rather + // than reading as "no approval asked for". + const config = JSON.parse(process.env.AI_PYTHON_PI || "{}") || {}; + + if (config.approve) { + pi.on("tool_call", async (event, ctx) => { + const answer = await ctx.ui.input( + APPROVE, + JSON.stringify({ + toolCallId: event.toolCallId, + parentToolCallId: event.parentToolCallId, + toolName: event.toolName, + input: event.input, + }), + ); + let decision = null; + try { + decision = JSON.parse(answer ?? ""); + } catch { + decision = null; + } + if (decision && decision.decision === "allow") { + if (decision.input && typeof decision.input === "object") { + // pi reads the arguments from this object after the handler: + // rewrite it in place. + for (const key of Object.keys(event.input)) delete event.input[key]; + Object.assign(event.input, decision.input); + } + return undefined; + } + return { + block: true, + // No answer: the SDK withdrew the question (the turn was stopped). + reason: + answer === undefined + ? "the approval was withdrawn" + : (decision && decision.reason) || "denied by policy", + ...(decision && decision.stop ? { terminate: true } : {}), + }; + }); + } + + for (const [name, server] of Object.entries(config.mcp || {})) { + pi.registerMcpServer(name, server); + } + + const gateway = config.gateway; + if (gateway) { + if (gateway.baseUrl) { + pi.registerProvider("vercel-ai-gateway", { baseUrl: gateway.baseUrl }); + } + pi.on("before_provider_headers", (event, ctx) => { + if (ctx.model?.provider !== "vercel-ai-gateway") return; + for (const key of Object.keys(event.headers)) { + const lower = key.toLowerCase(); + if (lower === "x-api-key" || lower === "authorization") { + event.headers[key] = null; + } + } + event.headers["authorization"] = + "Bearer " + (process.env.AI_GATEWAY_API_KEY || ""); + }); + } + + pi.registerCommand(READY, { + description: "Marks the ai-python SDK's extension as loaded", + handler: async () => {}, + }); +} diff --git a/src/ai/workers/_kinds.py b/src/ai/workers/_kinds.py index 69285dc0..e12c875f 100644 --- a/src/ai/workers/_kinds.py +++ b/src/ai/workers/_kinds.py @@ -1,4 +1,4 @@ -"""The workers, one per CLI. +"""The workers: Claude Code, Codex and Pi. Each adds its CLI's own options to `Worker` and builds its adapter. """ @@ -9,8 +9,50 @@ import pydantic +from ..providers import _optional from . import _bridge, _worker -from ._adapters import base, codex +from ._adapters import base, codex, pi + + +class ClaudeCode(_worker.Worker): + """Claude Code, through `claude-agent-sdk`. Needs the `claude-code` extra. + + `allowed_tools` limits Claude Code's own tools (`["Read", "Grep"]`, or + `[]` for none), and `disallowed_tools` takes some away. `max_turns` caps + the agent's steps in one turn; a turn that hits it ends with + `finish_reason="length"`. + """ + + kind: Literal["claude-code"] = "claude-code" + allowed_tools: list[str] | None = None + disallowed_tools: list[str] | None = None + max_turns: int | None = None + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: base.Approve | None, + ) -> base.Adapter: + _optional.import_optional_sdk( + "claude_agent_sdk", + extra="claude-code", + feature="ai.workers.ClaudeCode", + ) + # Imported here: the adapter needs claude_agent_sdk at import time. + from ._adapters import claude # noqa: PLC0415 + + return claude.ClaudeAdapter( + cwd=self.cwd, + env=dict(self.env), + executable=self.executable or "claude", + model=self.model, + effort=self.effort, + mcp_servers=dict(mcp_servers), + approve=approve, + allowed_tools=self.allowed_tools, + disallowed_tools=self.disallowed_tools, + max_turns=self.max_turns, + ) class Codex(_worker.Worker): @@ -55,3 +97,32 @@ def _make_adapter( sandbox=self.sandbox, config=config, ) + + +class Pi(_worker.Worker): + """Pi, over `pi --mode rpc`. + + `allowed_tools` limits pi's own tools (`["read", "grep"]`, or `[]` for + none), and `disallowed_tools` takes some away. + """ + + kind: Literal["pi"] = "pi" + allowed_tools: list[str] | None = None + disallowed_tools: list[str] | None = None + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: base.Approve | None, + ) -> base.Adapter: + return pi.PiAdapter( + cwd=self.cwd, + env=dict(self.env), + executable=self.executable or "pi", + model=self.model, + effort=self.effort, + mcp_servers=dict(mcp_servers), + approve=approve, + allowed_tools=self.allowed_tools, + disallowed_tools=self.disallowed_tools, + ) diff --git a/src/ai/workers/_worker.py b/src/ai/workers/_worker.py index 89aa3994..4670fa65 100644 --- a/src/ai/workers/_worker.py +++ b/src/ai/workers/_worker.py @@ -1,6 +1,6 @@ """`ai.workers.Worker`: a coding-agent CLI that goes where a model goes. -A worker is ONE conversation with one coding-agent CLI, such as Codex. It is +A worker is ONE conversation with one CLI: Claude Code, Codex or Pi. It is a pydantic model: its fields are its settings, the CLI's `session_id` and the conversation so far (`messages`). Dump it, store it, and validate it again later; it continues the same conversation. @@ -112,7 +112,7 @@ class _Runtime: class Worker(pydantic.BaseModel): - """Base class of the workers, such as `Codex`. + """Base class of `ClaudeCode`, `Codex` and `Pi`. `Worker.model_validate(data)` builds the kind that `data["kind"]` names. """ diff --git a/tests/workers/_adapters/test_claude.py b/tests/workers/_adapters/test_claude.py new file mode 100644 index 00000000..6854d5e4 --- /dev/null +++ b/tests/workers/_adapters/test_claude.py @@ -0,0 +1,317 @@ +"""The claude adapter against a scripted claude-agent-sdk client, offline.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncIterator +from pathlib import Path +from typing import Any + +import pytest + +claude_agent_sdk = pytest.importorskip("claude_agent_sdk") + +from ai import workers # noqa: E402 +from ai.types import events # noqa: E402 +from ai.types import messages as m # noqa: E402 +from ai.workers._adapters import claude # noqa: E402 + + +def _stream(event: dict[str, Any]) -> Any: + return claude_agent_sdk.StreamEvent(uuid="u", session_id="s", event=event) + + +def _result(**fields: Any) -> Any: + return claude_agent_sdk.ResultMessage( + subtype=fields.pop("subtype", "success"), + duration_ms=1, + duration_api_ms=1, + is_error=fields.pop("is_error", False), + num_turns=1, + session_id="s", + usage={"input_tokens": 7, "output_tokens": 3}, + total_cost_usd=0.01, + **fields, + ) + + +TURN = [ + _stream( + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking"}, + } + ), + _stream( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "plan"}, + } + ), + _stream({"type": "content_block_stop", "index": 0}), + claude_agent_sdk.AssistantMessage( + content=[ + claude_agent_sdk.ToolUseBlock( + id="t1", name="Bash", input={"command": "ls"} + ) + ], + model="claude", + ), + claude_agent_sdk.UserMessage( + content=[ + claude_agent_sdk.ToolResultBlock(tool_use_id="t1", content="a.py") + ] + ), + _stream( + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text"}, + } + ), + _stream( + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "one file"}, + } + ), + _stream({"type": "content_block_stop", "index": 1}), + claude_agent_sdk.AssistantMessage( + content=[claude_agent_sdk.TextBlock(text="one file")], model="claude" + ), + _result(), +] + + +class Client: + """Answers like `ClaudeSDKClient`: a query, then the scripted reply.""" + + def __init__(self, *replies: list[Any]) -> None: + self.replies = list(replies) + self.queries: list[Any] = [] + self.interrupted = False + + async def query(self, prompt: Any) -> None: + self.queries.append(prompt) + + async def receive_response(self) -> AsyncIterator[Any]: + for message in self.replies.pop(0): + yield message + + async def interrupt(self) -> None: + self.interrupted = True + + async def disconnect(self) -> None: + pass + + +def _adapter(tmp_path: Path, **options: Any) -> claude.ClaudeAdapter: + return claude.ClaudeAdapter( + cwd=str(tmp_path), + env={"CLAUDE_CONFIG_DIR": str(tmp_path / "config")}, + executable="claude", + model=options.get("model"), + effort=None, + mcp_servers=options.get("mcp_servers", {}), + approve=None, + allowed_tools=None, + disallowed_tools=None, + max_turns=None, + ) + + +async def _turn(adapter: claude.ClaudeAdapter, *prompt: Any) -> list[Any]: + async with asyncio.timeout(5): + return [e async for e in adapter.turn(list(prompt))] + + +async def test_a_turn_streams_in_the_model_view(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + adapter._client = Client(TURN) + seen = await _turn(adapter, m.TextPart(text="what is here?")) + assert [type(e).__name__ for e in seen] == [ + "StreamStart", + "ReasoningStart", + "ReasoningDelta", + "ReasoningEnd", + "BuiltinToolStart", + "BuiltinToolDelta", + "BuiltinToolEnd", + "BuiltinToolResult", + "TextStart", + "TextDelta", + "TextEnd", + "StreamEnd", + ] + assert adapter._client.queries == ["what is here?"] + end = seen[-1] + assert end.finish_reason == "stop" + assert end.usage.input_tokens == 7 + assert end.usage.raw["cost_usd"] == 0.01 + delta = next(e for e in seen if isinstance(e, events.BuiltinToolDelta)) + assert json.loads(delta.chunk) == {"command": "ls"} + result = next(e for e in seen if isinstance(e, events.BuiltinToolResult)) + assert (result.result.tool_name, result.result.result) == ("Bash", "a.py") + + +async def test_files_travel_as_base64_blocks(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + adapter._client = Client([_result()]) + png = m.FilePart(data=b"\x89PNG", media_type="image/png") + await _turn(adapter, m.TextPart(text="look"), png) + (sent,) = adapter._client.queries + (record,) = [record async for record in sent] + message = record["message"] + assert message["content"][1] == { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw==", + }, + } + + +async def test_an_unknown_file_is_refused(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + adapter._client = Client([_result()]) + zipped = m.FilePart(data=b"PK", media_type="application/zip") + with pytest.raises(workers.WorkerError, match="application/zip"): + await _turn(adapter, zipped) + + +async def test_finish_reasons(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + adapter._client = Client([_result(subtype="error_max_turns")]) + assert (await _turn(adapter, m.TextPart(text="x")))[-1].finish_reason == ( + "length" + ) + adapter._client = Client( + [_result(is_error=True, result="Invalid API key · Please run /login")] + ) + with pytest.raises(workers.NotAuthenticatedError): + await _turn(adapter, m.TextPart(text="x")) + adapter._client = Client([_result(is_error=True, result="overloaded")]) + with pytest.raises(workers.WorkerError, match="overloaded"): + await _turn(adapter, m.TextPart(text="x")) + + +async def test_an_interrupted_turn_is_cancelled_not_failed( + tmp_path: Path, +) -> None: + adapter = _adapter(tmp_path) + # The CLI's error result after the interrupt is read first, then the + # next turn settles on its own result. + client = Client([_result(is_error=True, result="interrupted")], [_result()]) + adapter._client = client + adapter._running = True # a turn was running when it was interrupted + await adapter.interrupt() + assert client.interrupted + seen = await _turn(adapter, m.TextPart(text="again")) + assert seen[-1].finish_reason == "stop" + assert client.replies == [] + + +async def test_a_seed_is_written_where_the_cli_reads(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + seed = [ + m.Message(role="user", parts=[m.TextPart(text="earlier")]), + m.Message( + role="assistant", + parts=[ + m.BuiltinToolCallPart( + tool_call_id="t1", tool_name="Read", tool_args="{}" + ), + m.BuiltinToolReturnPart( + tool_call_id="t1", tool_name="Read", result="milk" + ), + m.TextPart(text="It says milk."), + ], + ), + ] + session_id = await adapter._stage(seed) + key = claude_agent_sdk.project_key_for_directory(str(tmp_path)) + path = tmp_path / "config" / "projects" / key / f"{session_id}.jsonl" + records = [json.loads(line) for line in path.read_text().splitlines()] + assert [r["type"] for r in records] == [ + "user", + "assistant", + "user", + "assistant", + ] + assert records[1]["message"]["content"][0]["type"] == "tool_use" + assert records[2]["message"]["content"][0]["type"] == "tool_result" + + +def test_a_stored_transcript_reads_as_the_model_view() -> None: + records = [ + claude_agent_sdk.SessionMessage( + type="user", + uuid="1", + session_id="s", + message={"role": "user", "content": "q"}, + parent_tool_use_id=None, + ), + claude_agent_sdk.SessionMessage( + type="assistant", + uuid="2", + session_id="s", + message={ + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t1", + "name": "Read", + "input": {}, + } + ], + }, + parent_tool_use_id=None, + ), + claude_agent_sdk.SessionMessage( + type="user", + uuid="3", + session_id="s", + message={ + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "x"} + ], + }, + parent_tool_use_id=None, + ), + ] + from ai.workers._adapters import _handoff + + stored = [msg for r in records for msg in claude._stored_to_messages(r)] + history = _handoff.model_view(stored) + assert [x.role for x in history] == ["user", "assistant"] + call, result = history[1].parts + assert isinstance(call, m.BuiltinToolCallPart) + assert isinstance(result, m.BuiltinToolReturnPart) + assert (result.tool_name, result.result) == ("Read", "x") + + +def test_the_tui_command(tmp_path: Path) -> None: + adapter = _adapter( + tmp_path, + model="sonnet", + mcp_servers={ + "ai_sdk": {"type": "http", "url": "http://127.0.0.1:1/mcp"}, + "docs": {"command": "docs-mcp"}, + }, + ) + argv, env, minted = adapter.tui_command("abc") + assert argv[:3] == ["claude", "--resume", "abc"] + assert argv[argv.index("--model") + 1] == "sonnet" + config = json.loads(argv[argv.index("--mcp-config") + 1]) + assert list(config["mcpServers"]) == ["docs"] # no bridge in a TUI + assert env["DISABLE_AUTOUPDATER"] == "1" + assert minted == "abc" + argv, _, minted = adapter.tui_command(None) + assert argv[1] == "--session-id" and argv[2] == minted diff --git a/tests/workers/_adapters/test_pi.py b/tests/workers/_adapters/test_pi.py new file mode 100644 index 00000000..86f99ccb --- /dev/null +++ b/tests/workers/_adapters/test_pi.py @@ -0,0 +1,554 @@ +"""The pi adapter against a scripted `pi --mode rpc`, offline. + +The records here are the shapes a real pi 1.3.0 sent when measured. +""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Callable +from pathlib import Path +from typing import Any, cast + +import pytest + +from ai import workers +from ai.types import events +from ai.types import messages as m +from ai.workers._adapters import pi + +from .. import fakes + +SID = "s1" + + +def _update(kind: str, index: int = 0, **fields: Any) -> dict[str, Any]: + return { + "type": "message_update", + "assistantMessageEvent": { + "type": kind, + "contentIndex": index, + **fields, + }, + } + + +def _assistant(*content: dict[str, Any], stop: str = "stop") -> dict[str, Any]: + return { + "role": "assistant", + "content": list(content), + "usage": { + "input": 10, + "output": 5, + "cacheRead": 100, + "cacheWrite": 0, + "cost": {"total": 0.01}, + }, + "stopReason": stop, + } + + +CALL = { + "type": "toolCall", + "id": "toolu_1", + "name": "bash", + "arguments": {"command": "echo hi"}, +} + +# text -> tool call -> tool result -> text: what pi sent for +# "run `echo hi`, then say done". +TOOL_TURN: list[dict[str, Any]] = [ + {"type": "agent_start"}, + {"type": "message_start", "message": {"role": "assistant"}}, + _update("thinking_start"), + _update("thinking_delta", delta="plan"), + _update("thinking_end", content="plan"), + _update("toolcall_start", 1, id="toolu_1", toolName="bash"), + _update("toolcall_delta", 1, delta='{"command": "ec'), + _update("toolcall_end", 1, toolCall=CALL), + { + "type": "message_end", + "message": _assistant( + {"type": "thinking", "thinking": "plan"}, CALL, stop="toolUse" + ), + }, + { + "type": "tool_execution_end", + "toolCallId": "toolu_1", + "toolName": "bash", + "result": {"content": [{"type": "text", "text": "hi\n"}]}, + "isError": False, + }, + {"type": "message_start", "message": {"role": "assistant"}}, + _update("text_start"), + _update("text_delta", delta="done"), + _update("text_end", content="done"), + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "done"}), + }, + {"type": "agent_end", "messages": []}, + {"type": "agent_settled", "aborted": False}, +] + + +class Pi: + """Answers like `pi --mode rpc`: a prompt is accepted, then `turn` + plays out; `abort` settles the turn as aborted.""" + + def __init__( + self, + turn: list[dict[str, Any]] | None = None, + *, + prompt_error: str | None = None, + ) -> None: + self.process = fakes.FakeProcess(self.answer) + self.turn = turn if turn is not None else TOOL_TURN + self.prompt_error = prompt_error + self.on_prompt: Callable[[], None] | None = None + + def says(self, record: dict[str, Any]) -> None: + self.process.server_says(record) + + def respond(self, message: dict[str, Any], **fields: Any) -> None: + self.says( + { + "type": "response", + "id": message["id"], + "command": message["type"], + **fields, + } + ) + + def answer(self, message: dict[str, Any]) -> None: + kind = message.get("type") + if kind == "prompt": + if self.prompt_error is not None: + self.respond(message, success=False, error=self.prompt_error) + return + self.respond(message, success=True, data={"disposition": "started"}) + if self.on_prompt is not None: + self.on_prompt() + return + for record in self.turn: + self.says(record) + elif kind == "abort": + self.says({"type": "agent_settled", "aborted": True}) + self.respond(message, success=True) + elif kind == "get_state": + self.respond(message, success=True, data={"sessionId": SID}) + elif "id" in message: + self.respond(message, success=True) + + def sent(self, kind: str) -> list[dict[str, Any]]: + return [w for w in self.process.written if w.get("type") == kind] + + +def _adapter( + server: Pi | None = None, approve: Any = None, cwd: str = "/w", **env: str +) -> pi.PiAdapter: + adapter = pi.PiAdapter( + cwd=cwd, + env=env, + executable="pi", + model=None, + effort=None, + mcp_servers={}, + approve=approve, + allowed_tools=None, + disallowed_tools=None, + ) + if server is not None: + agent = pi._Agent( + cast("Any", server.process), on_dialog=adapter._on_dialog + ) + agent.start() + adapter._agent = agent + return adapter + + +async def _turn(adapter: pi.PiAdapter, *prompt: Any) -> list[Any]: + parts: list[m.TextPart | m.FilePart] = list(prompt) or [ + m.TextPart(text="go") + ] + async with asyncio.timeout(5): + return [e async for e in adapter.turn(parts)] + + +async def test_a_turn_streams_in_the_model_view() -> None: + server = Pi() + seen = await _turn(_adapter(server), m.TextPart(text="say done")) + assert [type(e).__name__ for e in seen] == [ + "StreamStart", + "ReasoningStart", + "ReasoningDelta", + "ReasoningEnd", + "BuiltinToolStart", + "BuiltinToolDelta", + "BuiltinToolEnd", + "BuiltinToolResult", + "TextStart", + "TextDelta", + "TextEnd", + "StreamEnd", + ] + # The arguments travel once, whole: pi's fragments are not what it + # settles on. + delta = next(e for e in seen if isinstance(e, events.BuiltinToolDelta)) + assert json.loads(delta.chunk) == {"command": "echo hi"} + result = next(e for e in seen if isinstance(e, events.BuiltinToolResult)) + assert result.result.result == "hi\n" + end = seen[-1] + assert end.finish_reason == "stop" + # Two model calls: cached input counts as input, cost rides in raw. + assert end.usage.input_tokens == 220 + assert end.usage.output_tokens == 10 + assert end.usage.cache_read_tokens == 200 + assert end.usage.raw == {"cost_usd": 0.02} + assert server.sent("prompt")[0]["message"] == "say done" + # Blocks of different model calls never share an id. + starts = [e.block_id for e in seen if isinstance(e, events.TextStart)] + assert starts == ["2.0"] + + +async def test_images_travel_as_base64() -> None: + server = Pi() + png = m.FilePart(data=b"\x89PNG", media_type="image/png") + await _turn(_adapter(server), m.TextPart(text="look"), png) + assert server.sent("prompt")[0]["images"] == [ + {"type": "image", "data": "iVBORw==", "mimeType": "image/png"} + ] + + +async def test_a_rejected_prompt_without_a_key_is_not_authenticated() -> None: + """Measured: pi refuses the prompt itself when it has no credentials.""" + server = Pi(prompt_error="No API key found for openai.\n\nUse /login") + with pytest.raises(workers.NotAuthenticatedError): + await _turn(_adapter(server)) + + +async def test_a_failed_model_call_fails_the_turn() -> None: + server = Pi( + [ + { + "type": "message_end", + "message": { + **_assistant(stop="error"), + "errorMessage": "529 overloaded", + }, + }, + {"type": "agent_settled", "aborted": False}, + ] + ) + with pytest.raises(workers.WorkerError, match="overloaded"): + await _turn(_adapter(server)) + + +def _ask(server: Pi) -> None: + server.says( + { + "type": "extension_ui_request", + "id": "ui-1", + "method": "input", + "title": pi.APPROVE, + "placeholder": json.dumps( + { + "toolCallId": "toolu_1", + "toolName": "bash", + "input": {"command": "echo hi"}, + } + ), + } + ) + + +@pytest.mark.parametrize("granted", [True, False]) +async def test_approvals_ask_and_answer(granted: bool) -> None: + server = Pi() + asked: list[Any] = [] + + async def approve( + call_id: str, name: str, tool_input: dict[str, Any] + ) -> tuple[bool, str]: + asked.append((call_id, name, tool_input)) + return granted, "not that" + + def answer(message: dict[str, Any]) -> None: + if message.get("type") == "extension_ui_response": + for record in TOOL_TURN: + server.says(record) + else: + Pi.answer(server, message) + + server.on_prompt = lambda: _ask(server) + server.process.answer = answer + seen = await _turn(_adapter(server, approve)) + assert asked == [("toolu_1", "bash", {"command": "echo hi"})] + reply = server.sent("extension_ui_response")[0] + assert reply["id"] == "ui-1" + assert json.loads(reply["value"]) == ( + {"decision": "allow"} + if granted + else {"decision": "deny", "reason": "not that"} + ) + assert isinstance(seen[-1], events.StreamEnd) + + +async def test_another_extensions_dialog_is_dismissed() -> None: + """Nobody is at a terminal to answer it; the agent must not wait.""" + server = Pi() + adapter = _adapter(server) + assert adapter._agent is not None + await adapter._on_dialog( + adapter._agent, {"id": "x", "method": "confirm", "title": "Clear?"} + ) + await adapter._on_dialog( + adapter._agent, {"id": "y", "method": "notify", "message": "hello"} + ) + assert server.sent("extension_ui_response") == [ + {"type": "extension_ui_response", "id": "x", "cancelled": True} + ] + + +async def test_an_interrupted_turn_ends_as_cancelled() -> None: + server = Pi([]) + adapter = _adapter(server) + turn = adapter.turn([m.TextPart(text="go")]) + assert isinstance(await anext(turn), events.StreamStart) + reader = asyncio.ensure_future(_rest(turn)) + await adapter.interrupt() + seen = await asyncio.wait_for(reader, 5) + assert seen[-1].finish_reason == "cancelled" + + +async def _rest(turn: Any) -> list[Any]: + return [e async for e in turn] + + +async def test_a_dead_pi_crashes_the_turn() -> None: + server = Pi([]) + server.on_prompt = lambda: server.process._lines.put_nowait("") + with pytest.raises(workers.WorkerError, match="crashed"): + await _turn(_adapter(server)) + + +async def test_a_run_an_extension_command_starts_is_followed() -> None: + """Measured: `/ask` answered `handled`, and the run the command started + followed, with pi already streaming. Its text was never streamed.""" + server = Pi( + [ + {"type": "agent_start"}, + {"type": "message_start", "message": {"role": "assistant"}}, + { + "type": "message_end", + "message": _assistant({"type": "text", "text": "PONG"}), + }, + {"type": "agent_settled", "aborted": False}, + ] + ) + + def answer(message: dict[str, Any]) -> None: + if message.get("type") == "prompt": + server.respond( + message, success=True, data={"disposition": "handled"} + ) + for record in server.turn: + server.says(record) + elif message.get("type") == "get_state": + server.respond(message, success=True, data={"isStreaming": True}) + else: + Pi.answer(server, message) + + server.process.answer = answer + seen = await _turn(_adapter(server), m.TextPart(text="/ask PONG")) + assert [e.chunk for e in seen if isinstance(e, events.TextDelta)] == [ + "PONG" + ] + + +async def test_a_command_that_starts_nothing_ends_the_turn_empty() -> None: + server = Pi() + + def answer(message: dict[str, Any]) -> None: + if message.get("type") == "prompt": + server.respond( + message, success=True, data={"disposition": "handled"} + ) + elif message.get("type") == "get_state": + server.respond(message, success=True, data={"isStreaming": False}) + else: + Pi.answer(server, message) + + server.process.answer = answer + seen = await _turn(_adapter(server), m.TextPart(text="/noop")) + assert [type(e).__name__ for e in seen] == ["StreamStart", "StreamEnd"] + + +def test_history_follows_the_active_branch_only() -> None: + lines = [ + {"type": "session", "id": "S", "cwd": "/w"}, + { + "type": "message", + "id": "a", + "parentId": None, + "message": {"role": "user", "content": "first"}, + }, + { + "type": "message", + "id": "b", + "parentId": "a", + "message": _assistant({"type": "text", "text": "abandoned"}), + }, + { + "type": "message", + "id": "c", + "parentId": "a", + "message": _assistant(CALL, stop="toolUse"), + }, + { + "type": "message", + "id": "d", + "parentId": "c", + "message": { + "role": "toolResult", + "toolCallId": "toolu_1", + "toolName": "bash", + "content": [{"type": "text", "text": "hi"}], + "isError": True, + }, + }, + {"type": "label", "id": "e", "parentId": "d"}, + ] + text = "\n".join(json.dumps(line) for line in lines) + stored = [ + msg + for entry in pi._active_branch(text) + for msg in pi._stored_to_messages(entry) + ] + assert [msg.role for msg in stored] == ["user", "assistant", "tool"] + result = stored[2].parts[0] + assert isinstance(result, m.ToolResultPart) + assert (result.result, result.result_kind) == ("hi", "error") + + +async def test_a_seed_is_listed_and_read_back(tmp_path: Path) -> None: + """The store, for real: write where pi looks, find it there.""" + agent_dir = tmp_path / "agent" + project = tmp_path / "project" + project.mkdir() + adapter = _adapter(cwd=str(project), PI_CODING_AGENT_DIR=str(agent_dir)) + seed = [ + m.Message(role="user", parts=[m.TextPart(text="the word is HERON")]), + m.Message( + role="assistant", + parts=[ + m.ReasoningPart(text="hmm"), + m.TextPart(text="Reading."), + m.BuiltinToolCallPart( + tool_call_id="c1", + tool_name="read", + tool_args='{"path": "a.txt"}', + ), + m.BuiltinToolReturnPart( + tool_call_id="c1", tool_name="read", result="milk" + ), + ], + ), + m.Message(role="assistant", parts=[m.TextPart(text="It says milk.")]), + ] + with pytest.warns(UserWarning, match="reasoning"): + session_id = await adapter._stage(seed) + infos = await adapter.list_sessions() + assert [i.session_id for i in infos] == [session_id] + assert infos[0].title == "the word is HERON" + files = list(agent_dir.glob("sessions/*/*.jsonl")) + assert len(files) == 1 + assert files[0].parent.name.startswith("--") + assert files[0].name.endswith(f"_{session_id}.jsonl") + back = await adapter.history(session_id) + assert [msg.role for msg in back] == ["user", "assistant", "assistant"] + text, call, result = back[1].parts # the reasoning was not carried + assert isinstance(text, m.TextPart) + assert isinstance(call, m.BuiltinToolCallPart) + assert json.loads(call.tool_args) == {"path": "a.txt"} + assert isinstance(result, m.BuiltinToolReturnPart) + assert result.result == "milk" + assert back[2].text == "It says milk." + with pytest.raises(workers.WorkerError): + await adapter.history("not-a-session") + + +async def test_a_seeded_pi_conversation_stays_on_its_model( + tmp_path: Path, +) -> None: + """Measured: a seeded pi conversation ran on pi's default model, not the + one it was on. History says which model wrote each reply, and seeding + writes it back the way pi records a switch.""" + agent_dir = tmp_path / "agent" + project = tmp_path / "project" + project.mkdir() + adapter = _adapter(cwd=str(project), PI_CODING_AGENT_DIR=str(agent_dir)) + stored = { + "type": "message", + "message": { + **_assistant({"type": "text", "text": "hi"}), + "api": "anthropic-messages", + "provider": "vercel-ai-gateway", + "model": "anthropic/claude-haiku-4.5", + "thinkingLevel": "high", + }, + } + seed = [ + m.Message(role="user", parts=[m.TextPart(text="hello")]), + *pi._stored_to_messages(stored), + ] + session_id = await adapter._stage(seed) + (written,) = agent_dir.glob(f"sessions/*/*_{session_id}.jsonl") + entries = [json.loads(line) for line in written.read_text().splitlines()] + assert [e["type"] for e in entries] == [ + "session", + "model_change", + "thinking_level_change", + "message", + "message", + ] + assert (entries[1]["provider"], entries[1]["modelId"]) == ( + "vercel-ai-gateway", + "anthropic/claude-haiku-4.5", + ) + # A chain, root first: pi reads it as one branch. + assert [e.get("parentId") for e in entries[1:]] == [ + None, + *[e["id"] for e in entries[1:-1]], + ] + + +def test_options_reach_pi_in_its_own_spelling() -> None: + adapter = pi.PiAdapter( + cwd="/w", + env={"K": "v"}, + executable="pi", + model="anthropic/claude-sonnet-4.6", + effort="high", + mcp_servers={ + "ai_sdk": {"type": "http", "url": "http://127.0.0.1:1/mcp"}, + "docs": {"command": "docs-mcp", "ignored": True}, + }, + approve=None, + allowed_tools=[], + disallowed_tools=["bash"], + ) + argv, env, minted = adapter.tui_command(None) + assert argv[:3] == ["pi", "--session-id", minted] + assert argv[3:5] == ["-e", pi.EXTENSION] + assert argv[argv.index("--model") + 1] == "anthropic/claude-sonnet-4.6" + assert argv[argv.index("--thinking") + 1] == "high" + assert "--no-tools" in argv + assert argv[argv.index("--exclude-tools") + 1] == "bash" + assert env["K"] == "v" + # The person at the TUI is the gate; the bridge is not there. + assert json.loads(env[pi.CONFIG_VAR]) == { + "approve": False, + "mcp": {"docs": {"command": "docs-mcp"}}, + } diff --git a/tests/workers/test_live.py b/tests/workers/test_live.py new file mode 100644 index 00000000..6c89aa8c --- /dev/null +++ b/tests/workers/test_live.py @@ -0,0 +1,137 @@ +"""Workers against the real CLIs. Opt-in: `uv run pytest -m live tests/workers`. + +Each test runs once per CLI installed here, in a fresh directory, with the +CLI's own login. They cost tokens. +""" + +from __future__ import annotations + +import importlib.util +import shutil +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import pytest + +import ai +from ai import workers +from ai.types import events + +pytestmark = pytest.mark.live + +KINDS: dict[str, tuple[str, Callable[..., workers.Worker]]] = { + "claude-code": ("claude", workers.ClaudeCode), + "codex": ("codex", workers.Codex), + "pi": ("pi", workers.Pi), +} + + +@pytest.fixture(params=list(KINDS)) +def kind(request: pytest.FixtureRequest) -> str: + kind = str(request.param) + executable, _ = KINDS[kind] + if shutil.which(executable) is None: + pytest.skip(f"{executable} is not installed") + if kind == "claude-code" and not importlib.util.find_spec( + "claude_agent_sdk" + ): + pytest.skip("claude-agent-sdk is not installed") + return kind + + +@pytest.fixture +def make(kind: str, tmp_path: Path) -> Callable[..., workers.Worker]: + """Build a worker of this kind, in a fresh directory.""" + _, cls = KINDS[kind] + return lambda **fields: cls(cwd=str(tmp_path), **fields) + + +async def test_a_turn_and_a_resume(make: Any) -> None: + async with make() as worker: + first = await ai.experimental_generate( + worker, + [ai.user_message("Reply with the word PONG and nothing else.")], + ) + assert "PONG" in first.text + assert worker.session_id + again = workers.Worker.model_validate_json(worker.model_dump_json()) + async with again: + second = await ai.experimental_generate( + again, + [ + *again.messages, + ai.user_message("Which word did you reply with? One word."), + ], + ) + assert "PONG" in second.text.upper() + assert again.session_id == worker.session_id + + +async def test_a_native_fork_keeps_the_past(make: Any) -> None: + async with make() as worker: + await ai.experimental_generate( + worker, + [ai.user_message("Remember the word HERON. Reply with OK only.")], + ) + branch = worker.model_copy(update={"fork": True}) + async with branch: + reply = await ai.experimental_generate( + branch, + [ + *branch.messages, + ai.user_message("Which word did I ask you to remember?"), + ], + ) + assert "HERON" in reply.text.upper() + assert branch.session_id != worker.session_id + + +async def test_agent_tools_reach_the_cli(make: Any) -> None: + @ai.tool + async def lookup_order(order_id: str) -> str: + """Look up the status of an order by its id.""" + return f"order {order_id}: shipped on 2026-11-02" + + prompt = ( + "Use the lookup_order tool for order A-17, then state its status in " + "one short sentence." + ) + async with make() as worker: + agent = ai.Agent(tools=[lookup_order]) + async with agent.run(worker, [ai.user_message(prompt)]) as run: + seen = [event async for event in run] + assert [m.role for m in run.messages] == [ + "user", + "assistant", + "tool", + "assistant", + ] + assert "2026" in run.messages[-1].text or "Nov" in run.messages[-1].text + hidden = [ + e + for e in seen + if isinstance(e, events.BuiltinToolStart) and "lookup" in e.tool_name + ] + assert hidden == [] # the CLI's own report of the bridge call + + +async def test_a_denied_tool_call_does_not_run( + kind: str, make: Any, tmp_path: Path +) -> None: + if kind == "codex": + pytest.skip("codex asks only when its own policy escalates") + pending: list[Any] = [] + prompt = "Create a file named notes.txt containing hello. Use a tool." + async with make(approve=True) as worker: + async with ai.Agent().run(worker, [ai.user_message(prompt)]) as run: + async for event in run: + if isinstance(event, events.PendingHookEvent): + pending.append(event.hook) + ai.resolve_hook( + event.hook, + {"granted": False, "reason": "not allowed here"}, + ) + assert pending + assert all(h.hook_type == "ToolApproval" for h in pending) + assert not (tmp_path / "notes.txt").exists() diff --git a/uv.lock b/uv.lock index d371b1aa..0d0c4a4c 100644 --- a/uv.lock +++ b/uv.lock @@ -25,6 +25,9 @@ dependencies = [ anthropic = [ { name = "anthropic" }, ] +claude-code = [ + { name = "claude-agent-sdk" }, +] mcp = [ { name = "mcp" }, ] @@ -46,6 +49,7 @@ vercel = [ dev = [ { name = "anthropic" }, { name = "async-solipsism" }, + { name = "claude-agent-sdk" }, { name = "mcp" }, { name = "mypy" }, { name = "openai" }, @@ -53,6 +57,7 @@ dev = [ { name = "opentelemetry-sdk" }, { name = "pytest" }, { name = "pytest-asyncio" }, + { name = "pytest-timeout" }, { name = "python-dotenv" }, { name = "rich" }, { name = "ruff" }, @@ -64,6 +69,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.83.0,<2.0.0" }, + { name = "claude-agent-sdk", marker = "extra == 'claude-code'", specifier = ">=0.2.144,<0.3" }, { name = "httpx2", specifier = ">=2.12.0" }, { name = "mcp", marker = "extra == 'mcp'", specifier = ">=2.0.0,<3.0.0" }, { name = "modelsdotdev", specifier = "==0.*" }, @@ -75,12 +81,13 @@ requires-dist = [ { name = "typing-extensions", specifier = ">=4.15.0" }, { name = "vercel-oidc", marker = "extra == 'vercel'", specifier = ">=0.8.0" }, ] -provides-extras = ["anthropic", "mcp", "openai", "otel", "typesafe", "vercel"] +provides-extras = ["anthropic", "claude-code", "mcp", "openai", "otel", "typesafe", "vercel"] [package.metadata.requires-dev] dev = [ { name = "anthropic", specifier = ">=0.83.0" }, { name = "async-solipsism", specifier = ">=0.9" }, + { name = "claude-agent-sdk", specifier = ">=0.2.144,<0.3" }, { name = "mcp", specifier = ">=2.0.0,<3.0.0" }, { name = "mypy", specifier = "~=2.1.0" }, { name = "openai", specifier = ">=2.14.0" }, @@ -88,6 +95,7 @@ dev = [ { name = "opentelemetry-sdk", specifier = ">=1.43.0" }, { name = "pytest", specifier = ">=8.0" }, { name = "pytest-asyncio", specifier = ">=0.24" }, + { name = "pytest-timeout", specifier = ">=2.3" }, { name = "python-dotenv", specifier = ">=1.2.1" }, { name = "rich", specifier = ">=14.2.0" }, { name = "ruff", specifier = "~=0.8.0" }, @@ -333,6 +341,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/db/8f/61959034484a4a7c527811f4721e75d02d653a35afb0b6054474d8185d4c/charset_normalizer-3.4.7-py3-none-any.whl", hash = "sha256:3dce51d0f5e7951f8bb4900c257dad282f49190fdbebecd4ba99bcc41fef404d", size = 61958, upload-time = "2026-04-02T09:28:37.794Z" }, ] +[[package]] +name = "claude-agent-sdk" +version = "0.2.178" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "jsonschema" }, + { name = "mcp" }, + { name = "sniffio" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bb/96/029c76e62b48c7eef3982873f6882d8e5f94a8eaa293392a91229cc0b87e/claude_agent_sdk-0.2.178.tar.gz", hash = "sha256:b59aca6ba61d43a6d78f0ad03a2230485a26a73d7f5e7b60e9d5e8a70cdba8f2", size = 400217, upload-time = "2026-11-04T00:52:25.514Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/cd/73bf0d3b6ec1456461e7e83e900c3bcf1ef6d767530dac1ad4590cc2ea74/claude_agent_sdk-0.2.178-py3-none-macosx_11_0_arm64.whl", hash = "sha256:7a20f490cfdd7d7cb9a70b3500cb0d9d026420d987e25a10667c638c5dfe036d", size = 138228664, upload-time = "2026-11-04T00:52:32.383Z" }, + { url = "https://files.pythonhosted.org/packages/ba/b6/62e1a2d67ff4561cf31758b4332594086979c849eede9ab21950601af24e/claude_agent_sdk-0.2.178-py3-none-macosx_11_0_x86_64.whl", hash = "sha256:c944ff5f0bac523b8c77d8d2f7913d533efc983e42eece1502fe74aa3a4a0482", size = 145514246, upload-time = "2026-11-04T00:52:39.069Z" }, + { url = "https://files.pythonhosted.org/packages/bb/10/c5d2908e1bded45de8b69b44bc5d76f431534e18cca99d192a755d04801e/claude_agent_sdk-0.2.178-py3-none-manylinux_2_17_aarch64.whl", hash = "sha256:98fcd5ea6b591d14a65580bee3df3f5fa7c75ad68a44ff81780177d2ed8db2a3", size = 149927469, upload-time = "2026-11-04T00:52:45.46Z" }, + { url = "https://files.pythonhosted.org/packages/49/5f/fb5dc93b604e8c55bd0da2465ba075b28e6ba4f0b0cda30cd15649c87e7d/claude_agent_sdk-0.2.178-py3-none-manylinux_2_17_x86_64.whl", hash = "sha256:1d5e6d2cced242b9d74658ec90d96936a7c9db5867b85705ae52ec0736060146", size = 150300143, upload-time = "2026-11-04T00:52:52.87Z" }, + { url = "https://files.pythonhosted.org/packages/46/9b/2a8cf8def2c32bd296d6380d08dafef63969ca9eb0c5674a7fa8cce4d1d6/claude_agent_sdk-0.2.178-py3-none-win_amd64.whl", hash = "sha256:b31f5b4076b6a0ce138716f03d04ab9899fa8b4bde7d7760f9862bdcf9b2786b", size = 151962195, upload-time = "2026-11-04T00:53:00.908Z" }, +] + [[package]] name = "click" version = "8.3.1" @@ -1092,6 +1119,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" }, ] +[[package]] +name = "pytest-timeout" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload-time = "2025-05-05T19:44:34.99Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload-time = "2025-05-05T19:44:33.502Z" }, +] + [[package]] name = "python-dotenv" version = "1.2.1" From 8ac61fd9fcf89ad13f9d5f4dce56a591694347ed Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 6 Nov 2026 13:36:04 -0800 Subject: [PATCH 2/3] Arrange ai.workers like providers: core files, a directory per CLI --- examples/workers/agent_tools.py | 52 -- examples/workers/approvals.py | 36 +- examples/workers/claude_code/agent_tools.py | 46 ++ examples/workers/codex/agent_tools.py | 46 ++ examples/workers/pi/agent_tools.py | 46 ++ src/ai/workers/__init__.py | 6 +- src/ai/workers/_adapter.py | 139 +++++ src/ai/workers/_adapters/_handoff.py | 504 ------------------ src/ai/workers/_adapters/base.py | 333 ------------ src/ai/workers/_bridge.py | 6 +- src/ai/workers/_convert.py | 116 ++++ src/ai/workers/_handoff.py | 218 ++++++++ src/ai/workers/_kinds.py | 128 ----- src/ai/workers/_worker.py | 19 +- src/ai/workers/claude_code/__init__.py | 3 + .../claude.py => claude_code/adapter.py} | 42 +- src/ai/workers/claude_code/handoff.py | 106 ++++ src/ai/workers/claude_code/worker.py | 49 ++ src/ai/workers/codex/__init__.py | 3 + .../{_adapters/codex.py => codex/adapter.py} | 44 +- src/ai/workers/codex/handoff.py | 86 +++ .../workers/{_adapters => codex}/jsonrpc.py | 96 +++- src/ai/workers/codex/worker.py | 54 ++ src/ai/workers/pi/__init__.py | 3 + .../{_adapters/pi.py => pi/adapter.py} | 142 ++++- .../pi_extension.js => pi/extension.js} | 0 src/ai/workers/pi/handoff.py | 148 +++++ src/ai/workers/pi/worker.py | 37 ++ .../workers/claude_code}/__init__.py | 0 .../test_adapter.py} | 12 +- .../workers/{_adapters => codex}/__init__.py | 0 .../test_codex.py => codex/test_adapter.py} | 9 +- .../{_adapters => codex}/test_jsonrpc.py | 2 +- tests/workers/fakes.py | 27 +- tests/workers/pi/__init__.py | 0 .../test_pi.py => pi/test_adapter.py} | 24 +- tests/workers/{_adapters => }/test_handoff.py | 27 +- 37 files changed, 1441 insertions(+), 1168 deletions(-) delete mode 100644 examples/workers/agent_tools.py create mode 100644 examples/workers/claude_code/agent_tools.py create mode 100644 examples/workers/codex/agent_tools.py create mode 100644 examples/workers/pi/agent_tools.py create mode 100644 src/ai/workers/_adapter.py delete mode 100644 src/ai/workers/_adapters/_handoff.py delete mode 100644 src/ai/workers/_adapters/base.py create mode 100644 src/ai/workers/_convert.py create mode 100644 src/ai/workers/_handoff.py delete mode 100644 src/ai/workers/_kinds.py create mode 100644 src/ai/workers/claude_code/__init__.py rename src/ai/workers/{_adapters/claude.py => claude_code/adapter.py} (95%) create mode 100644 src/ai/workers/claude_code/handoff.py create mode 100644 src/ai/workers/claude_code/worker.py create mode 100644 src/ai/workers/codex/__init__.py rename src/ai/workers/{_adapters/codex.py => codex/adapter.py} (95%) create mode 100644 src/ai/workers/codex/handoff.py rename src/ai/workers/{_adapters => codex}/jsonrpc.py (68%) create mode 100644 src/ai/workers/codex/worker.py create mode 100644 src/ai/workers/pi/__init__.py rename src/ai/workers/{_adapters/pi.py => pi/adapter.py} (89%) rename src/ai/workers/{_adapters/pi_extension.js => pi/extension.js} (100%) create mode 100644 src/ai/workers/pi/handoff.py create mode 100644 src/ai/workers/pi/worker.py rename {src/ai/workers/_adapters => tests/workers/claude_code}/__init__.py (100%) rename tests/workers/{_adapters/test_claude.py => claude_code/test_adapter.py} (96%) rename tests/workers/{_adapters => codex}/__init__.py (100%) rename tests/workers/{_adapters/test_codex.py => codex/test_adapter.py} (98%) rename tests/workers/{_adapters => codex}/test_jsonrpc.py (99%) create mode 100644 tests/workers/pi/__init__.py rename tests/workers/{_adapters/test_pi.py => pi/test_adapter.py} (97%) rename tests/workers/{_adapters => }/test_handoff.py (90%) diff --git a/examples/workers/agent_tools.py b/examples/workers/agent_tools.py deleted file mode 100644 index f56bef4c..00000000 --- a/examples/workers/agent_tools.py +++ /dev/null @@ -1,52 +0,0 @@ -"""A worker as the model of an `ai.Agent`: the agent's tools reach it by MCP. - -The CLI calls `lookup_order`. The worker hands the call to the agent, the -agent runs the tool, and the CLI gets the result and goes on. - - uv run python examples/workers/agent_tools.py [claude-code|codex|pi] -""" - -import asyncio -import sys -from collections.abc import Callable - -import ai -from ai import workers - -KINDS: dict[str, Callable[..., workers.Worker]] = { - "claude-code": workers.ClaudeCode, - "codex": workers.Codex, - "pi": workers.Pi, -} - - -@ai.tool -async def lookup_order(order_id: str) -> str: - """Look up the status of an order by its id.""" - return f"order {order_id}: shipped on 2026-11-02" - - -async def main(kind: str) -> None: - agent = ai.Agent(tools=[lookup_order]) - messages = [ - ai.user_message( - "Use the lookup_order tool for order A-17, then tell me its " - "status in one sentence." - ) - ] - async with KINDS[kind]() as worker: - async with agent.run(worker, messages) as run: - async for event in run: - match event: - case ai.events.ToolEnd(): # the agent runs this one - call = event.tool_call - print(f"[agent tool] {call.tool_name} {call.tool_args}") - case ai.events.BuiltinToolEnd(): # the CLI's own tool - print(f"[cli tool] {event.tool_call.tool_name}") - case ai.events.TextDelta(): - print(event.chunk, end="", flush=True) - print() - print([m.role for m in run.messages]) - - -asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "claude-code")) diff --git a/examples/workers/approvals.py b/examples/workers/approvals.py index 0d71f625..16e4fc47 100644 --- a/examples/workers/approvals.py +++ b/examples/workers/approvals.py @@ -22,23 +22,25 @@ async def main(kind: str) -> None: messages = [ai.user_message("Create notes.txt containing 'hello'.")] - async with KINDS[kind](approve=True) as worker: - async with ai.Agent().run(worker, messages) as run: - async for event in run: - match event: - case ai.events.PendingHookEvent(hook=hook): - meta = hook.metadata - answer = await asyncio.to_thread( - input, - f"allow {meta['tool']} {meta['kwargs']}? [y/N] ", - ) - granted = answer.strip().lower() == "y" - ai.resolve_hook( - hook, {"granted": granted, "reason": "user said no"} - ) - case ai.events.TextDelta(): - print(event.chunk, end="", flush=True) - print() + async with ( + KINDS[kind](approve=True) as worker, + ai.Agent().run(worker, messages) as run, + ): + async for event in run: + match event: + case ai.events.PendingHookEvent(hook=hook): + meta = hook.metadata + answer = await asyncio.to_thread( + input, + f"allow {meta['tool']} {meta['kwargs']}? [y/N] ", + ) + granted = answer.strip().lower() == "y" + ai.resolve_hook( + hook, {"granted": granted, "reason": "user said no"} + ) + case ai.events.TextDelta(): + print(event.chunk, end="", flush=True) + print() asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "claude-code")) diff --git a/examples/workers/claude_code/agent_tools.py b/examples/workers/claude_code/agent_tools.py new file mode 100644 index 00000000..7de0f79f --- /dev/null +++ b/examples/workers/claude_code/agent_tools.py @@ -0,0 +1,46 @@ +"""Claude Code as the model of an `ai.Agent`: the agent's tools reach it by MCP. + +Claude Code calls `lookup_order`. The worker hands the call to the agent, the +agent runs the tool, and Claude Code gets the result and goes on. + + uv run python examples/workers/claude_code/agent_tools.py +""" + +import asyncio + +import ai +from ai import workers + + +@ai.tool +async def lookup_order(order_id: str) -> str: + """Look up the status of an order by its id.""" + return f"order {order_id}: shipped on 2026-11-02" + + +async def main() -> None: + agent = ai.Agent(tools=[lookup_order]) + messages = [ + ai.user_message( + "Use the lookup_order tool for order A-17, then tell me its " + "status in one sentence." + ) + ] + async with ( + workers.ClaudeCode() as worker, + agent.run(worker, messages) as run, + ): + async for event in run: + match event: + case ai.events.ToolEnd(): # the agent runs this one + call = event.tool_call + print(f"[agent tool] {call.tool_name} {call.tool_args}") + case ai.events.BuiltinToolEnd(): # Claude Code's own tool + print(f"[cli tool] {event.tool_call.tool_name}") + case ai.events.TextDelta(): + print(event.chunk, end="", flush=True) + print() + print([m.role for m in run.messages]) + + +asyncio.run(main()) diff --git a/examples/workers/codex/agent_tools.py b/examples/workers/codex/agent_tools.py new file mode 100644 index 00000000..3a781bd9 --- /dev/null +++ b/examples/workers/codex/agent_tools.py @@ -0,0 +1,46 @@ +"""Codex as the model of an `ai.Agent`: the agent's tools reach it by MCP. + +Codex calls `lookup_order`. The worker hands the call to the agent, the +agent runs the tool, and Codex gets the result and goes on. + + uv run python examples/workers/codex/agent_tools.py +""" + +import asyncio + +import ai +from ai import workers + + +@ai.tool +async def lookup_order(order_id: str) -> str: + """Look up the status of an order by its id.""" + return f"order {order_id}: shipped on 2026-11-02" + + +async def main() -> None: + agent = ai.Agent(tools=[lookup_order]) + messages = [ + ai.user_message( + "Use the lookup_order tool for order A-17, then tell me its " + "status in one sentence." + ) + ] + async with ( + workers.Codex() as worker, + agent.run(worker, messages) as run, + ): + async for event in run: + match event: + case ai.events.ToolEnd(): # the agent runs this one + call = event.tool_call + print(f"[agent tool] {call.tool_name} {call.tool_args}") + case ai.events.BuiltinToolEnd(): # Codex's own tool + print(f"[cli tool] {event.tool_call.tool_name}") + case ai.events.TextDelta(): + print(event.chunk, end="", flush=True) + print() + print([m.role for m in run.messages]) + + +asyncio.run(main()) diff --git a/examples/workers/pi/agent_tools.py b/examples/workers/pi/agent_tools.py new file mode 100644 index 00000000..45f3ec5b --- /dev/null +++ b/examples/workers/pi/agent_tools.py @@ -0,0 +1,46 @@ +"""Pi as the model of an `ai.Agent`: the agent's tools reach it by MCP. + +pi calls `lookup_order`. The worker hands the call to the agent, the +agent runs the tool, and pi gets the result and goes on. + + uv run python examples/workers/pi/agent_tools.py +""" + +import asyncio + +import ai +from ai import workers + + +@ai.tool +async def lookup_order(order_id: str) -> str: + """Look up the status of an order by its id.""" + return f"order {order_id}: shipped on 2026-11-02" + + +async def main() -> None: + agent = ai.Agent(tools=[lookup_order]) + messages = [ + ai.user_message( + "Use the lookup_order tool for order A-17, then tell me its " + "status in one sentence." + ) + ] + async with ( + workers.Pi() as worker, + agent.run(worker, messages) as run, + ): + async for event in run: + match event: + case ai.events.ToolEnd(): # the agent runs this one + call = event.tool_call + print(f"[agent tool] {call.tool_name} {call.tool_args}") + case ai.events.BuiltinToolEnd(): # pi's own tool + print(f"[cli tool] {event.tool_call.tool_name}") + case ai.events.TextDelta(): + print(event.chunk, end="", flush=True) + print() + print([m.role for m in run.messages]) + + +asyncio.run(main()) diff --git a/src/ai/workers/__init__.py b/src/ai/workers/__init__.py index 2b66dac4..4ad6b7bc 100644 --- a/src/ai/workers/__init__.py +++ b/src/ai/workers/__init__.py @@ -24,15 +24,17 @@ a conversation the CLI has stored. """ -from ._adapters.base import SessionInfo +from ._adapter import SessionInfo from ._errors import ( ExecutableMissingError, NotAuthenticatedError, TranscriptMismatchError, WorkerError, ) -from ._kinds import ClaudeCode, Codex, Pi from ._worker import Worker +from .claude_code import ClaudeCode +from .codex import Codex +from .pi import Pi __all__ = [ "ClaudeCode", diff --git a/src/ai/workers/_adapter.py b/src/ai/workers/_adapter.py new file mode 100644 index 00000000..fd34f25b --- /dev/null +++ b/src/ai/workers/_adapter.py @@ -0,0 +1,139 @@ +"""The adapter protocol: what `Worker` needs from each CLI. + +An adapter drives ONE conversation with one CLI. It speaks the model view +both ways: `turn()` yields the events of one assistant message, with the +CLI's own tools as `BuiltinTool*` events, and ends with a `StreamEnd`; a +history it reads or is seeded with is a plain `list[Message]`. +""" + +from __future__ import annotations + +import asyncio +import os +import re +import warnings +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from typing import Any, Protocol + +import pydantic + +from ..types import events +from ..types import messages as messages_ +from . import _errors as errors + +Prompt = list[messages_.TextPart | messages_.FilePart] + +# `approve(tool_call_id, tool_name, tool_input)` returns `(granted, reason)`. +# +# Never raises: the worker answers a failure with a denial. +Approve = Callable[[str, str, dict[str, Any]], Awaitable[tuple[bool, str]]] + + +class SessionInfo(pydantic.BaseModel): + """A conversation in a CLI's own store.""" + + session_id: str + title: str | None = None + cwd: str | None = None + created_at: int | None = None + """Milliseconds since the epoch, when the CLI records it.""" + updated_at: int | None = None + + +class CLIAdapter(Protocol): + version: str | None + """What ` --version` printed, once started.""" + + async def start(self) -> None: + """Check the CLI runs, and start what it needs (codex: a server).""" + ... + + async def close(self) -> None: ... + + async def open( + self, + *, + session_id: str | None, + fork: bool, + seed: list[messages_.Message] | None, + instructions: str | None, + ) -> str: + """Open the conversation and return its session id. + + A new one (seeded with a past when `seed` is given), the stored one + `session_id` names, or a native fork of it. + """ + ... + + def turn(self, prompt: Prompt) -> AsyncIterator[events.Event]: + """Send a prompt; yield one assistant message, ending in StreamEnd.""" + ... + + async def steer(self, text: str) -> None: + """Send guidance to the turn that runs now.""" + ... + + async def interrupt(self) -> None: + """End the running turn. The CLI keeps its transcript.""" + ... + + async def history(self, session_id: str) -> list[messages_.Message]: + """Read a stored conversation.""" + ... + + async def list_sessions(self) -> list[SessionInfo]: ... + + def tui_command( + self, session_id: str | None + ) -> tuple[list[str], dict[str, str], str | None]: + """Return the CLI's interactive command, its env, and its session id. + + The session id is minted here for a new conversation when the CLI + lets us choose it, else None. + """ + ... + + +async def probe( + kind: str, + executable: str, + *, + install: str, + verified: str, + cwd: str, + env: Mapping[str, str], +) -> str: + """Check that the CLI runs, and return its version. + + Warns when it is older than the version the adapter was verified with: + the adapters lean on CLI behavior that is measured, not documented. + """ + try: + proc = await asyncio.create_subprocess_exec( + executable, + "--version", + cwd=cwd, + env={**os.environ, **env}, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + async with asyncio.timeout(60): + stdout, _ = await proc.communicate() + except (OSError, TimeoutError) as exc: + raise errors.ExecutableMissingError(kind, executable, install) from exc + if proc.returncode != 0: + raise errors.ExecutableMissingError(kind, executable, install) + version = (stdout.decode(errors="replace").strip().splitlines() or [""])[0] + have, want = _version(version), _version(verified) + if have is not None and want is not None and have < want: + warnings.warn( + f"{kind} {version} is older than {verified}, the version this " + "SDK was verified with; update it if something misbehaves", + stacklevel=2, + ) + return version + + +def _version(text: str | None) -> tuple[int, ...] | None: + found = re.search(r"\d+(?:\.\d+)+", text or "") + return tuple(int(n) for n in found.group(0).split(".")) if found else None diff --git a/src/ai/workers/_adapters/_handoff.py b/src/ai/workers/_adapters/_handoff.py deleted file mode 100644 index 873a0b6c..00000000 --- a/src/ai/workers/_adapters/_handoff.py +++ /dev/null @@ -1,504 +0,0 @@ -"""Seeding a CLI with a past: one rendering rule, every CLI. - -A worker given `[*past, user_message]` on its first turn starts a new -conversation from `past`. Each CLI stores it natively so the agent truly -has the context, and the rule for what becomes what is the same everywhere: - -- text: verbatim, same role. -- tool calls and results: the NATIVE tool items, original name and - arguments unaltered. Measured: the receiving agent reads them as history - and acts with its own tools; it does not try to call a tool it lacks. - Names are not translated — mapping `Write` to `fileChange` would be the - invented taxonomy this SDK refuses everywhere else. Nor are PATHS: a call - that wrote /Users/x/proj/recipe.txt says so, and that is where it was - written. Across machines the receiver must be told where things are now - (or the workspace copied to the same path) — as any colleague would be. -- reasoning: a user-role handoff note after the turn it belonged to — where - the destination accepts it. Measured: fed back as the agent's OWN thoughts - it is rejected (Claude signs thinking blocks) or silently ignored (codex - accepts a reasoning item and never reads it). As a note, codex uses it. - Claude does NOT accept it in any form: Anthropic's safeguards flag real - model reasoning replayed as user content (`[reasoning_extraction]`) — - hand-written notes passed, genuine Opus thinking was refused. So the - Claude renderer omits reasoning, and the adapter warns with the count, - because dropping context silently is the one thing this SDK never does. - Pi omits it too, for both reasons: pi drops thinking another model wrote, - and pi often runs Claude. -- files: images go to every harness, as the image input each takes (pi - takes bytes, so an image known only by URL becomes a note naming it); - PDFs to Claude as documents. Anything else (a PDF into codex, a text - file) is left out and counted (`unsupported_files`), and the adapter - warns, like it does for reasoning into Claude. -""" - -from __future__ import annotations - -import datetime -import json -import uuid -import warnings -from typing import Any - -from ...types import media -from ...types import messages as messages_ - -NOTE_PREFIX = "[Handoff from a previous session — earlier reasoning: " -AUTHOR = ("api", "provider", "model") -"""What names a reply's author in a pi session: its API and model.""" -SYSTEM_PREFIX = "[Handoff from a previous session — instructions it was given: " - - -def _text(value: Any) -> str: - return value if isinstance(value, str) else json.dumps(value) - - -def _note(reasoning: list[str]) -> str: - return ( - NOTE_PREFIX + " ".join(r.strip() for r in reasoning if r.strip()) + "]" - ) - - -def _image_url(part: messages_.FilePart) -> str: - """Return an http(s) URL as it is, anything else as a data URL.""" - return media.data_to_data_url(part.data, part.media_type) - - -def _source(part: messages_.FilePart) -> dict[str, Any]: - """Return an Anthropic content-block source for a file.""" - if isinstance(part.data, str) and media.is_downloadable_url(part.data): - return {"type": "url", "url": part.data} - return { - "type": "base64", - "media_type": part.media_type, - "data": media.data_to_base64(part.data), - } - - -def harness_view(history: list[messages_.Message]) -> list[messages_.Message]: - """Return a past with the CLI's own tools as ordinary tool calls. - - A worker's messages hold the tools a CLI ran itself as - `BuiltinToolCallPart`/`BuiltinToolReturnPart`, the way a provider's web - search is. A CLI's store holds them as tool calls, each with its result - in a `tool` message after it. The renderers below read that shape. - """ - out: list[messages_.Message] = [] - for message in history: - parts: list[messages_.Part] = [] - results: list[messages_.Part] = [] - first = True - for part in message.parts: - if results and not isinstance( - part, messages_.BuiltinToolReturnPart - ): - # What follows a result comes after it: a new message. - if parts: - out.append(_split(message, parts, first=first)) - first = False - out.append(messages_.Message(role="tool", parts=results)) - parts, results = [], [] - if isinstance(part, messages_.BuiltinToolCallPart): - parts.append( - messages_.ToolCallPart( - tool_call_id=part.tool_call_id, - tool_name=part.tool_name, - tool_args=part.tool_args or "{}", - provider_metadata=part.provider_metadata, - ) - ) - elif isinstance(part, messages_.BuiltinToolReturnPart): - results.append( - messages_.ToolResultPart( - tool_call_id=part.tool_call_id, - tool_name=part.tool_name, - result=part.result, - result_kind="error" if part.is_error else "json", - ) - ) - else: - parts.append(part) - if parts: - out.append(_split(message, parts, first=first)) - if results: - out.append(messages_.Message(role="tool", parts=results)) - return out - - -def _split( - message: messages_.Message, parts: list[messages_.Part], *, first: bool -) -> messages_.Message: - """One piece of a message: the first keeps its id, the rest get new ones.""" - if first: - return message.model_copy(update={"parts": parts}) - return messages_.Message( - role=message.role, - parts=parts, - provider_metadata=message.provider_metadata, - ) - - -def model_view(stored: list[messages_.Message]) -> list[messages_.Message]: - """Return a stored conversation the way a worker shows it. - - The reverse of `harness_view`: the CLI's tool calls become built-in - parts, and each result joins the assistant message of its call. - """ - out: list[messages_.Message] = [] - owner: dict[str, int] = {} # tool_call_id -> index in `out` - for message in stored: - if message.role == "assistant": - parts: list[messages_.Part] = [] - for part in message.parts: - if isinstance(part, messages_.ToolCallPart): - owner[part.tool_call_id] = len(out) - parts.append( - messages_.BuiltinToolCallPart( - tool_call_id=part.tool_call_id, - tool_name=part.tool_name, - tool_args=part.tool_args, - provider_metadata=part.provider_metadata, - ) - ) - else: - parts.append(part) - out.append(message.model_copy(update={"parts": parts})) - elif message.role == "tool": - for result in message.tool_results: - index = owner.get(result.tool_call_id) - if index is None: - continue - target = out[index] - out[index] = target.model_copy( - update={ - "parts": [ - *target.parts, - messages_.BuiltinToolReturnPart( - tool_call_id=result.tool_call_id, - tool_name=result.tool_name - or _name_of(target, result.tool_call_id), - result=result.result, - is_error=result.result_kind == "error", - ), - ] - } - ) - else: - out.append(message) - return out - - -def _name_of(message: messages_.Message, tool_call_id: str) -> str: - return next( - ( - p.tool_name - for p in message.parts - if isinstance(p, messages_.BuiltinToolCallPart) - and p.tool_call_id == tool_call_id - ), - "", - ) - - -def warn_unsupported_files( - messages: list[messages_.Message], kind: str -) -> None: - """Say how many files a history held that `kind` was not given.""" - left = unsupported_files(messages, kind) - if left: - warnings.warn( - f"{left} file(s) in the history were not carried into {kind}: " - + "it takes images" - + (" and PDFs" if kind == "claude-code" else "") - + " from a history. Text, tool calls and images were carried.", - stacklevel=3, - ) - - -def unsupported_files(messages: list[messages_.Message], kind: str) -> int: - """Count files a history holds that `kind` cannot take as input.""" - return sum( - 1 - for m in messages - for p in m.parts - if isinstance(p, messages_.FilePart) - and not p.media_type.startswith("image/") - and not (kind == "claude-code" and p.media_type == "application/pdf") - ) - - -def to_codex_items(messages: list[messages_.Message]) -> list[dict[str, Any]]: - """Responses API items for `thread/inject_items`.""" - items: list[dict[str, Any]] = [] - for m in harness_view(messages): - reasoning: list[str] = [] - for p in m.parts: - if isinstance(p, messages_.TextPart) and p.text.strip(): - ctype = "input_text" if m.role == "user" else "output_text" - items.append( - { - "type": "message", - "role": m.role, - "content": [{"type": ctype, "text": p.text}], - } - ) - elif ( - isinstance(p, messages_.FilePart) - and m.role == "user" - and p.media_type.startswith("image/") - ): - items.append( - { - "type": "message", - "role": "user", - "content": [ - {"type": "input_image", "image_url": _image_url(p)} - ], - } - ) - elif isinstance(p, messages_.ToolCallPart): - items.append( - { - "type": "function_call", - "call_id": p.tool_call_id or p.id, - "name": p.tool_name, - "arguments": p.tool_args or "{}", - } - ) - elif isinstance(p, messages_.ToolResultPart): - items.append( - { - "type": "function_call_output", - "call_id": p.tool_call_id, - "output": _text(p.result), - } - ) - elif isinstance(p, messages_.ReasoningPart) and p.text.strip(): - reasoning.append(p.text) - if reasoning: - items.append( - { - "type": "message", - "role": "user", - "content": [ - {"type": "input_text", "text": _note(reasoning)} - ], - } - ) - return items - - -def reasoning_count(messages: list[messages_.Message]) -> int: - return sum( - 1 - for m in messages - for p in m.parts - if isinstance(p, messages_.ReasoningPart) and p.text.strip() - ) - - -def to_claude_records( - messages: list[messages_.Message], session_id: str, cwd: str -) -> list[dict[str, Any]]: - """Build Claude Code transcript records for the CLI's store. - - Chained by parentUuid. - - Reasoning parts are omitted: see the module docstring. Callers report - the omission — `reasoning_count()` says how many. - """ - records: list[dict[str, Any]] = [] - parent: str | None = None - - def add(role: str, content: Any) -> None: - nonlocal parent - uid = str(uuid.uuid4()) - ts = ( - datetime.datetime.now(datetime.UTC) - .isoformat() - .replace("+00:00", "Z") - ) - records.append( - { - "parentUuid": parent, - "isSidechain": False, - "type": role, - "message": {"role": role, "content": content}, - "uuid": uid, - "timestamp": ts, - "sessionId": session_id, - "cwd": cwd, - } - ) - parent = uid - - for m in harness_view(messages): - blocks: list[dict[str, Any]] = [] - for p in m.parts: - if isinstance(p, messages_.TextPart) and p.text.strip(): - blocks.append({"type": "text", "text": p.text}) - elif isinstance(p, messages_.FilePart) and m.role == "user": - if p.media_type.startswith("image/"): - blocks.append({"type": "image", "source": _source(p)}) - elif p.media_type == "application/pdf": - blocks.append({"type": "document", "source": _source(p)}) - elif isinstance(p, messages_.ToolCallPart): - try: - args = json.loads(p.tool_args) if p.tool_args else {} - except json.JSONDecodeError: - args = {"raw": p.tool_args} - blocks.append( - { - "type": "tool_use", - "id": p.tool_call_id or p.id, - "name": p.tool_name, - "input": args, - } - ) - elif isinstance(p, messages_.ToolResultPart): - blocks.append( - { - "type": "tool_result", - "tool_use_id": p.tool_call_id, - "content": _text(p.result), - **( - {"is_error": True} - if p.result_kind == "error" - else {} - ), - } - ) - # ReasoningPart: omitted on purpose — Anthropic's safeguards - # refuse replayed reasoning. The adapter warns with the count. - if blocks: - # Tool results live in role="tool" messages in the IR; on disk - # the CLI keeps them as user records with tool_result blocks. - add("user" if m.role in ("user", "tool") else "assistant", blocks) - return records - - -def to_pi_messages( - messages: list[messages_.Message], *, timestamp: int -) -> list[dict[str, Any]]: - """Build pi `AgentMessage`s for a session file. - - Assistant messages pi wrote keep their author (`provider_metadata["pi"]`): - pi restores a session's model from it. Any other is attributed to a - model pi does not run (`ai-python/handoff`), so pi treats it as another - model's — it keeps its text and tool calls and normalizes the call ids - for its own provider. Consecutive assistant messages are merged, the way one - model response holds text and calls together. Reasoning is omitted: see - the module docstring. - """ - out: list[dict[str, Any]] = [] - usage = { - "input": 0, - "output": 0, - "cacheRead": 0, - "cacheWrite": 0, - "totalTokens": 0, - "cost": { - "input": 0, - "output": 0, - "cacheRead": 0, - "cacheWrite": 0, - "total": 0, - }, - } - for m in harness_view(messages): - if m.role == "tool": - for p in m.parts: - if isinstance(p, messages_.ToolResultPart): - out.append( - { - "role": "toolResult", - "toolCallId": p.tool_call_id, - "toolName": p.tool_name, - "content": [ - {"type": "text", "text": _text(p.result)} - ], - "isError": p.result_kind == "error", - "timestamp": timestamp, - } - ) - continue - blocks: list[dict[str, Any]] = [] - for p in m.parts: - if isinstance(p, messages_.TextPart) and p.text.strip(): - blocks.append({"type": "text", "text": p.text}) - elif ( - isinstance(p, messages_.FilePart) - and m.role == "user" - and p.media_type.startswith("image/") - ): - if isinstance(p.data, str) and media.is_downloadable_url( - p.data - ): - blocks.append( - {"type": "text", "text": f"[image: {p.data}]"} - ) - else: - blocks.append( - { - "type": "image", - "data": media.data_to_base64(p.data), - "mimeType": p.media_type, - } - ) - elif ( - isinstance(p, messages_.ToolCallPart) and m.role == "assistant" - ): - try: - args = json.loads(p.tool_args) if p.tool_args else {} - except json.JSONDecodeError: - args = {"raw": p.tool_args} - blocks.append( - { - "type": "toolCall", - "id": p.tool_call_id or p.id, - "name": p.tool_name, - "arguments": args if isinstance(args, dict) else {}, - } - ) - if not blocks: - continue - if m.role == "system": - # Measured (pi 1.3.0): a system message in a session file is - # pi's own prompt state, and replaying one REPLACED pi's system - # prompt and its tools. So instructions travel as what they are - # to the agent — the caller's words — never as its own output. - text = " ".join(b["text"] for b in blocks if b["type"] == "text") - blocks = [{"type": "text", "text": SYSTEM_PREFIX + text + "]"}] - if m.role in ("user", "system"): - out.append( - {"role": "user", "content": blocks, "timestamp": timestamp} - ) - continue - if out and out[-1]["role"] == "assistant": - out[-1]["content"] += blocks - else: - # A reply pi wrote keeps its author: pi restores a session's model - # from its last reply's, and a model it does not know (ours) - # put the conversation on pi's default one (measured). - author = (m.provider_metadata or {}).get("pi") or {} - known = all(isinstance(author.get(k), str) for k in AUTHOR) - out.append( - { - "role": "assistant", - "content": blocks, - **( - {k: author[k] for k in AUTHOR} - if known - else { - "api": "ai-python", - "provider": "ai-python", - "model": "handoff", - } - ), - "usage": usage, - "stopReason": "stop", - "timestamp": timestamp, - } - ) - if any(b["type"] == "toolCall" for b in out[-1]["content"]): - out[-1]["stopReason"] = "toolUse" - return out diff --git a/src/ai/workers/_adapters/base.py b/src/ai/workers/_adapters/base.py deleted file mode 100644 index 10d09016..00000000 --- a/src/ai/workers/_adapters/base.py +++ /dev/null @@ -1,333 +0,0 @@ -"""What every adapter shares: its protocol, local processes, prompt files. - -An adapter drives ONE conversation with one CLI. It speaks the model view -both ways: `turn()` yields the events of one assistant message, with the -CLI's own tools as `BuiltinTool*` events, and ends with a `StreamEnd`; a -history it reads or is seeded with is a plain `list[Message]`. -""" - -from __future__ import annotations - -import asyncio -import base64 -import contextlib -import os -import re -import warnings -from collections.abc import ( - AsyncIterator, - Awaitable, - Callable, - Mapping, - Sequence, -) -from typing import Any, Protocol - -import pydantic - -from ...models.core.helpers import files -from ...types import events, media -from ...types import messages as messages_ -from .. import _errors as errors - -Prompt = list[messages_.TextPart | messages_.FilePart] - -Approve = Callable[[str, str, dict[str, Any]], Awaitable[tuple[bool, str]]] -"""`approve(tool_call_id, tool_name, tool_input)` returns `(granted, reason)`. - -Never raises: the worker answers a failure with a denial.""" - -LINE_LIMIT = 256 * 1024 * 1024 -"""Longest line a CLI may write. The protocols are JSON lines, and one line -can be megabytes (a codex `thread/read` of a long thread). With asyncio's -default (64 KiB) such a line failed as if the CLI had crashed.""" - -TERM_GRACE = 3.0 -"""Seconds a CLI gets to exit after SIGTERM, before SIGKILL.""" - - -class SessionInfo(pydantic.BaseModel): - """A conversation in a CLI's own store.""" - - session_id: str - title: str | None = None - cwd: str | None = None - created_at: int | None = None - """Milliseconds since the epoch, when the CLI records it.""" - updated_at: int | None = None - - -class Adapter(Protocol): - version: str | None - """What ` --version` printed, once started.""" - - async def start(self) -> None: - """Check the CLI runs, and start what it needs (codex: a server).""" - ... - - async def close(self) -> None: ... - - async def open( - self, - *, - session_id: str | None, - fork: bool, - seed: list[messages_.Message] | None, - instructions: str | None, - ) -> str: - """Open the conversation and return its session id. - - A new one (seeded with a past when `seed` is given), the stored one - `session_id` names, or a native fork of it. - """ - ... - - def turn(self, prompt: Prompt) -> AsyncIterator[events.Event]: - """Send a prompt; yield one assistant message, ending in StreamEnd.""" - ... - - async def steer(self, text: str) -> None: - """Send guidance to the turn that runs now.""" - ... - - async def interrupt(self) -> None: - """End the running turn. The CLI keeps its transcript.""" - ... - - async def history(self, session_id: str) -> list[messages_.Message]: - """Read a stored conversation.""" - ... - - async def list_sessions(self) -> list[SessionInfo]: ... - - def tui_command( - self, session_id: str | None - ) -> tuple[list[str], dict[str, str], str | None]: - """Return the CLI's interactive command, its env, and its session id. - - The session id is minted here for a new conversation when the CLI - lets us choose it, else None. - """ - ... - - -# -- local processes ----------------------------------------------------------- - - -class Process: - """A CLI started by `spawn`: JSON lines on stdin and stdout.""" - - def __init__(self, proc: asyncio.subprocess.Process) -> None: - self._proc = proc - - @property - def pid(self) -> int: - return self._proc.pid - - async def write(self, data: str) -> None: - assert self._proc.stdin is not None - self._proc.stdin.write(data.encode()) - await self._proc.stdin.drain() - - async def readline(self) -> str: - assert self._proc.stdout is not None - return (await self._proc.stdout.readline()).decode() - - async def read_stderr(self) -> str: - if self._proc.stderr is None: - return "" - return (await self._proc.stderr.read()).decode(errors="replace") - - async def wait(self) -> int: - return await self._proc.wait() - - async def terminate(self) -> None: - """Stop the process group: SIGTERM, then SIGKILL.""" - if self._proc.returncode is not None: - return - with contextlib.suppress(ProcessLookupError): - os.killpg(os.getpgid(self._proc.pid), 15) - try: - async with asyncio.timeout(TERM_GRACE): - await self._proc.wait() - except TimeoutError: - with contextlib.suppress(ProcessLookupError): - os.killpg(os.getpgid(self._proc.pid), 9) - with contextlib.suppress(Exception): - await self._proc.wait() - - async def death_note(self) -> str: - """Say how the process ended, and its last words.""" - code: int | None = None - tail = "" - with contextlib.suppress(Exception): - async with asyncio.timeout(2): - code = await self.wait() - with contextlib.suppress(Exception): - async with asyncio.timeout(2): - tail = (await self.read_stderr()).strip() - how = f"exit code {code}" if code is not None else "exited" - if code in (137, -9): - how += " (SIGKILL, typically the kernel's OOM killer)" - return f"{how}; stderr: {tail[-800:]}" if tail else f"{how}; no stderr" - - -async def spawn( - argv: Sequence[str], *, cwd: str, env: Mapping[str, str] -) -> Process: - """Start a CLI in its own process group, with pipes for stdio.""" - try: - proc = await asyncio.create_subprocess_exec( - *argv, - cwd=cwd, - env={**os.environ, **env}, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - # Own the whole tree, so teardown reaches grandchildren too. - start_new_session=True, - limit=LINE_LIMIT, - ) - except OSError as exc: - raise errors.WorkerError(f"cannot start {argv[0]!r}: {exc}") from exc - return Process(proc) - - -async def probe( - kind: str, - executable: str, - *, - install: str, - verified: str, - cwd: str, - env: Mapping[str, str], -) -> str: - """Check that the CLI runs, and return its version. - - Warns when it is older than the version the adapter was verified with: - the adapters lean on CLI behavior that is measured, not documented. - """ - try: - proc = await asyncio.create_subprocess_exec( - executable, - "--version", - cwd=cwd, - env={**os.environ, **env}, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - async with asyncio.timeout(60): - stdout, _ = await proc.communicate() - except (OSError, TimeoutError) as exc: - raise errors.ExecutableMissingError(kind, executable, install) from exc - if proc.returncode != 0: - raise errors.ExecutableMissingError(kind, executable, install) - version = (stdout.decode(errors="replace").strip().splitlines() or [""])[0] - have, want = _version(version), _version(verified) - if have is not None and want is not None and have < want: - warnings.warn( - f"{kind} {version} is older than {verified}, the version this " - "SDK was verified with; update it if something misbehaves", - stacklevel=2, - ) - return version - - -def _version(text: str | None) -> tuple[int, ...] | None: - found = re.search(r"\d+(?:\.\d+)+", text or "") - return tuple(int(n) for n in found.group(0).split(".")) if found else None - - -# -- files in a prompt --------------------------------------------------------- -# -# Files are turned into bytes HERE: a URL is downloaded with the SDK's own -# SSRF-checked downloader, not handed to the CLI. Measured: Claude Code drops -# an image given by URL without a word. - - -async def data(part: messages_.FilePart) -> bytes: - """Return a file's bytes, downloading a URL first.""" - if isinstance(part.data, bytes): - return part.data - if media.is_downloadable_url(part.data): - body, _ = await files.download(part.data) - return body - return base64.b64decode(media.data_to_base64(part.data)) - - -async def as_text(part: messages_.FilePart) -> str: - """Return a text file as text a CLI reads inline, named if it has one.""" - body = (await data(part)).decode("utf-8", errors="replace") - name = part.filename or "attached file" - return f"\n{body}\n" - - -def check_files(prompt: Prompt, kind: str, accepts: Sequence[str]) -> None: - """Refuse a file the CLI cannot take, before anything runs. - - Text files (`text/*`) go everywhere, as text. `accepts` lists the other - media type prefixes the CLI takes. - """ - for part in prompt: - if isinstance(part, messages_.FilePart) and not any( - part.media_type.startswith(p) for p in ("text/", *accepts) - ): - raise errors.WorkerError( - f"{kind} does not take {part.media_type} files" - ) - - -def text_of(prompt: Prompt) -> str: - return "\n\n".join( - p.text for p in prompt if isinstance(p, messages_.TextPart) - ) - - -def builtin_call( - tool_call_id: str, - tool_name: str, - tool_args: str, - *, - started: set[str], -) -> list[events.Event]: - """Return the events of one finished CLI tool call. - - `started` holds the calls whose start was already sent: the start is - sent here otherwise, because the hydrator only completes a call it saw - start. The arguments travel as a delta: the hydrator builds a call's - arguments only from its deltas. - """ - out: list[events.Event] = [] - if tool_call_id not in started: - started.add(tool_call_id) - out.append( - events.BuiltinToolStart( - tool_call_id=tool_call_id, tool_name=tool_name - ) - ) - return [ - *out, - events.BuiltinToolDelta(tool_call_id=tool_call_id, chunk=tool_args), - events.BuiltinToolEnd( - tool_call_id=tool_call_id, - tool_call=messages_.BuiltinToolCallPart( - tool_call_id=tool_call_id, - tool_name=tool_name, - tool_args=tool_args, - ), - ), - ] - - -def builtin_result( - tool_call_id: str, tool_name: str, result: Any, *, is_error: bool -) -> events.BuiltinToolResult: - return events.BuiltinToolResult( - tool_call_id=tool_call_id, - result=messages_.BuiltinToolReturnPart( - tool_call_id=tool_call_id, - tool_name=tool_name, - result=result, - is_error=is_error, - ), - ) diff --git a/src/ai/workers/_bridge.py b/src/ai/workers/_bridge.py index f32f5a59..ef0bb5b8 100644 --- a/src/ai/workers/_bridge.py +++ b/src/ai/workers/_bridge.py @@ -26,14 +26,14 @@ if TYPE_CHECKING: from ..types import tools as tools_ +# The MCP server name the CLI sees. Claude Code calls a tool +# `mcp__ai_sdk__`. SERVER_NAME = "ai_sdk" -"""The MCP server name the CLI sees. Claude Code calls a tool -`mcp__ai_sdk__`.""" PROTOCOL_VERSION = "2025-06-18" +# `on_call(name, arguments)` returns `(text, is_error)`. OnCall = Callable[[str, dict[str, Any]], Awaitable[tuple[str, bool]]] -"""`on_call(name, arguments)` returns `(text, is_error)`.""" class Bridge: diff --git a/src/ai/workers/_convert.py b/src/ai/workers/_convert.py new file mode 100644 index 00000000..62e6f88b --- /dev/null +++ b/src/ai/workers/_convert.py @@ -0,0 +1,116 @@ +"""What every adapter converts the same way. + +Prompt files go in as bytes or text. The CLI's own tool calls come out as +`BuiltinTool*` events. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Sequence +from typing import Any + +from ..models.core.helpers import files +from ..types import events, media +from ..types import messages as messages_ +from . import _adapter +from . import _errors as errors + +# -- files in a prompt --------------------------------------------------------- +# +# Files are turned into bytes HERE: a URL is downloaded with the SDK's own +# SSRF-checked downloader, not handed to the CLI. Measured: Claude Code drops +# an image given by URL without a word. + + +async def data(part: messages_.FilePart) -> bytes: + """Return a file's bytes, downloading a URL first.""" + if isinstance(part.data, bytes): + return part.data + if media.is_downloadable_url(part.data): + body, _ = await files.download(part.data) + return body + return base64.b64decode(media.data_to_base64(part.data)) + + +async def as_text(part: messages_.FilePart) -> str: + """Return a text file as text a CLI reads inline, named if it has one.""" + body = (await data(part)).decode("utf-8", errors="replace") + name = part.filename or "attached file" + return f"\n{body}\n" + + +def check_files( + prompt: _adapter.Prompt, kind: str, accepts: Sequence[str] +) -> None: + """Refuse a file the CLI cannot take, before anything runs. + + Text files (`text/*`) go everywhere, as text. `accepts` lists the other + media type prefixes the CLI takes. + """ + for part in prompt: + if isinstance(part, messages_.FilePart) and not any( + part.media_type.startswith(p) for p in ("text/", *accepts) + ): + raise errors.WorkerError( + f"{kind} does not take {part.media_type} files" + ) + + +def text_of(prompt: _adapter.Prompt) -> str: + return "\n\n".join( + p.text for p in prompt if isinstance(p, messages_.TextPart) + ) + + +# -- the CLI's own tools as events -------------------------------------------- + + +def builtin_call( + tool_call_id: str, + tool_name: str, + tool_args: str, + *, + started: set[str], +) -> list[events.Event]: + """Return the events of one finished CLI tool call. + + `started` holds the calls whose start was already sent: the start is + sent here otherwise, because the hydrator only completes a call it saw + start. The arguments travel as a delta: the hydrator builds a call's + arguments only from its deltas. + """ + out: list[events.Event] = [] + if tool_call_id not in started: + started.add(tool_call_id) + out.append( + events.BuiltinToolStart( + tool_call_id=tool_call_id, tool_name=tool_name + ) + ) + return [ + *out, + events.BuiltinToolDelta(tool_call_id=tool_call_id, chunk=tool_args), + events.BuiltinToolEnd( + tool_call_id=tool_call_id, + tool_call=messages_.BuiltinToolCallPart( + tool_call_id=tool_call_id, + tool_name=tool_name, + tool_args=tool_args, + ), + ), + ] + + +def builtin_result( + tool_call_id: str, tool_name: str, result: Any, *, is_error: bool +) -> events.BuiltinToolResult: + return events.BuiltinToolResult( + tool_call_id=tool_call_id, + result=messages_.BuiltinToolReturnPart( + tool_call_id=tool_call_id, + tool_name=tool_name, + result=result, + is_error=is_error, + ), + ) diff --git a/src/ai/workers/_handoff.py b/src/ai/workers/_handoff.py new file mode 100644 index 00000000..62f0479d --- /dev/null +++ b/src/ai/workers/_handoff.py @@ -0,0 +1,218 @@ +"""Seeding a CLI with a past: one rendering rule, every CLI. + +A worker given `[*past, user_message]` on its first turn starts a new +conversation from `past`. Each CLI stores it natively so the agent truly +has the context, and the rule for what becomes what is the same everywhere: + +- text: verbatim, same role. +- tool calls and results: the NATIVE tool items, original name and + arguments unaltered. Measured: the receiving agent reads them as history + and acts with its own tools; it does not try to call a tool it lacks. + Names are not translated — mapping `Write` to `fileChange` would be the + invented taxonomy this SDK refuses everywhere else. Nor are PATHS: a call + that wrote /Users/x/proj/recipe.txt says so, and that is where it was + written. Across machines the receiver must be told where things are now + (or the workspace copied to the same path) — as any colleague would be. +- reasoning: a user-role handoff note after the turn it belonged to — where + the destination accepts it. Measured: fed back as the agent's OWN thoughts + it is rejected (Claude signs thinking blocks) or silently ignored (codex + accepts a reasoning item and never reads it). As a note, codex uses it. + Claude does NOT accept it in any form: Anthropic's safeguards flag real + model reasoning replayed as user content (`[reasoning_extraction]`) — + hand-written notes passed, genuine Opus thinking was refused. So the + Claude renderer omits reasoning, and the adapter warns with the count, + because dropping context silently is the one thing this SDK never does. + Pi omits it too, for both reasons: pi drops thinking another model wrote, + and pi often runs Claude. +- files: images go to every harness, as the image input each takes (pi + takes bytes, so an image known only by URL becomes a note naming it); + PDFs to Claude as documents. Anything else (a PDF into codex, a text + file) is left out and counted (`unsupported_files`), and the adapter + warns, like it does for reasoning into Claude. +""" + +from __future__ import annotations + +import json +import warnings +from typing import Any + +from ..types import messages as messages_ + + +def result_text(value: Any) -> str: + """Return a tool result as the text a CLI's store holds.""" + return value if isinstance(value, str) else json.dumps(value) + + +# -- the two views of a past --------------------------------------------------- +# +# A worker's messages hold two kinds of tool calls. The tools an `ai.Agent` +# runs (through the bridge) are `ToolCallPart`s with their results in `tool` +# messages, like a model's. The tools the CLI runs itself are built-in +# parts, with the result in the same assistant message, like a provider's +# web search. A CLI's store holds both kinds the same way: a call, then its +# result in a message of its own. `harness_view` gives the renderers that +# one shape, so each renderer handles one kind of tool call, not two; +# `model_view` is its reverse, for a history read back from a store. + + +def harness_view(history: list[messages_.Message]) -> list[messages_.Message]: + """Return a past with the CLI's own tools as ordinary tool calls. + + A worker's messages hold the tools a CLI ran itself as + `BuiltinToolCallPart`/`BuiltinToolReturnPart`, the way a provider's web + search is. A CLI's store holds them as tool calls, each with its result + in a `tool` message after it. The renderers below read that shape. + """ + out: list[messages_.Message] = [] + for message in history: + parts: list[messages_.Part] = [] + results: list[messages_.Part] = [] + first = True + for part in message.parts: + if results and not isinstance( + part, messages_.BuiltinToolReturnPart + ): + # What follows a result comes after it: a new message. + if parts: + out.append(_split(message, parts, first=first)) + first = False + out.append(messages_.Message(role="tool", parts=results)) + parts, results = [], [] + if isinstance(part, messages_.BuiltinToolCallPart): + parts.append( + messages_.ToolCallPart( + tool_call_id=part.tool_call_id, + tool_name=part.tool_name, + tool_args=part.tool_args or "{}", + provider_metadata=part.provider_metadata, + ) + ) + elif isinstance(part, messages_.BuiltinToolReturnPart): + results.append( + messages_.ToolResultPart( + tool_call_id=part.tool_call_id, + tool_name=part.tool_name, + result=part.result, + result_kind="error" if part.is_error else "json", + ) + ) + else: + parts.append(part) + if parts: + out.append(_split(message, parts, first=first)) + if results: + out.append(messages_.Message(role="tool", parts=results)) + return out + + +def _split( + message: messages_.Message, parts: list[messages_.Part], *, first: bool +) -> messages_.Message: + """One piece of a message: the first keeps its id, the rest get new ones.""" + if first: + return message.model_copy(update={"parts": parts}) + return messages_.Message( + role=message.role, + parts=parts, + provider_metadata=message.provider_metadata, + ) + + +def model_view(stored: list[messages_.Message]) -> list[messages_.Message]: + """Return a stored conversation the way a worker shows it. + + The reverse of `harness_view`: the CLI's tool calls become built-in + parts, and each result joins the assistant message of its call. + """ + out: list[messages_.Message] = [] + owner: dict[str, int] = {} # tool_call_id -> index in `out` + for message in stored: + if message.role == "assistant": + parts: list[messages_.Part] = [] + for part in message.parts: + if isinstance(part, messages_.ToolCallPart): + owner[part.tool_call_id] = len(out) + parts.append( + messages_.BuiltinToolCallPart( + tool_call_id=part.tool_call_id, + tool_name=part.tool_name, + tool_args=part.tool_args, + provider_metadata=part.provider_metadata, + ) + ) + else: + parts.append(part) + out.append(message.model_copy(update={"parts": parts})) + elif message.role == "tool": + for result in message.tool_results: + index = owner.get(result.tool_call_id) + if index is None: + continue + target = out[index] + out[index] = target.model_copy( + update={ + "parts": [ + *target.parts, + messages_.BuiltinToolReturnPart( + tool_call_id=result.tool_call_id, + tool_name=result.tool_name + or _name_of(target, result.tool_call_id), + result=result.result, + is_error=result.result_kind == "error", + ), + ] + } + ) + else: + out.append(message) + return out + + +def _name_of(message: messages_.Message, tool_call_id: str) -> str: + return next( + ( + p.tool_name + for p in message.parts + if isinstance(p, messages_.BuiltinToolCallPart) + and p.tool_call_id == tool_call_id + ), + "", + ) + + +def warn_unsupported_files( + messages: list[messages_.Message], kind: str +) -> None: + """Say how many files a history held that `kind` was not given.""" + left = unsupported_files(messages, kind) + if left: + warnings.warn( + f"{left} file(s) in the history were not carried into {kind}: " + + "it takes images" + + (" and PDFs" if kind == "claude-code" else "") + + " from a history. Text, tool calls and images were carried.", + stacklevel=3, + ) + + +def unsupported_files(messages: list[messages_.Message], kind: str) -> int: + """Count files a history holds that `kind` cannot take as input.""" + return sum( + 1 + for m in messages + for p in m.parts + if isinstance(p, messages_.FilePart) + and not p.media_type.startswith("image/") + and not (kind == "claude-code" and p.media_type == "application/pdf") + ) + + +def reasoning_count(messages: list[messages_.Message]) -> int: + return sum( + 1 + for m in messages + for p in m.parts + if isinstance(p, messages_.ReasoningPart) and p.text.strip() + ) diff --git a/src/ai/workers/_kinds.py b/src/ai/workers/_kinds.py deleted file mode 100644 index e12c875f..00000000 --- a/src/ai/workers/_kinds.py +++ /dev/null @@ -1,128 +0,0 @@ -"""The workers: Claude Code, Codex and Pi. - -Each adds its CLI's own options to `Worker` and builds its adapter. -""" - -from __future__ import annotations - -from typing import Any, Literal - -import pydantic - -from ..providers import _optional -from . import _bridge, _worker -from ._adapters import base, codex, pi - - -class ClaudeCode(_worker.Worker): - """Claude Code, through `claude-agent-sdk`. Needs the `claude-code` extra. - - `allowed_tools` limits Claude Code's own tools (`["Read", "Grep"]`, or - `[]` for none), and `disallowed_tools` takes some away. `max_turns` caps - the agent's steps in one turn; a turn that hits it ends with - `finish_reason="length"`. - """ - - kind: Literal["claude-code"] = "claude-code" - allowed_tools: list[str] | None = None - disallowed_tools: list[str] | None = None - max_turns: int | None = None - - def _make_adapter( - self, - mcp_servers: dict[str, dict[str, Any]], - approve: base.Approve | None, - ) -> base.Adapter: - _optional.import_optional_sdk( - "claude_agent_sdk", - extra="claude-code", - feature="ai.workers.ClaudeCode", - ) - # Imported here: the adapter needs claude_agent_sdk at import time. - from ._adapters import claude # noqa: PLC0415 - - return claude.ClaudeAdapter( - cwd=self.cwd, - env=dict(self.env), - executable=self.executable or "claude", - model=self.model, - effort=self.effort, - mcp_servers=dict(mcp_servers), - approve=approve, - allowed_tools=self.allowed_tools, - disallowed_tools=self.disallowed_tools, - max_turns=self.max_turns, - ) - - -class Codex(_worker.Worker): - """Codex, over `codex app-server`. - - `sandbox` is Codex's own sandbox mode: `"read-only"`, - `"workspace-write"` or `"danger-full-access"` (the default). `config` - takes Codex `-c` overrides. - """ - - kind: Literal["codex"] = "codex" - sandbox: ( - Literal["read-only", "workspace-write", "danger-full-access"] | None - ) = None - config: dict[str, Any] = pydantic.Field(default_factory=dict) - - def _make_adapter( - self, - mcp_servers: dict[str, dict[str, Any]], - approve: base.Approve | None, - ) -> base.Adapter: - config = dict(self.config) - if self.effort is not None: - config.setdefault( - "model_reasoning_effort", - "xhigh" if self.effort == "max" else self.effort, - ) - if _bridge.SERVER_NAME in mcp_servers: - # Codex's default wait for an MCP tool (60s) is too short for a - # tool an `ai.Agent` runs, maybe after an approval. - config.setdefault( - f"mcp_servers.{_bridge.SERVER_NAME}.tool_timeout_sec", - 24 * 3600, - ) - return codex.CodexAdapter( - cwd=self.cwd, - env=dict(self.env), - executable=self.executable or "codex", - model=self.model, - mcp_servers=dict(mcp_servers), - approve=approve, - sandbox=self.sandbox, - config=config, - ) - - -class Pi(_worker.Worker): - """Pi, over `pi --mode rpc`. - - `allowed_tools` limits pi's own tools (`["read", "grep"]`, or `[]` for - none), and `disallowed_tools` takes some away. - """ - - kind: Literal["pi"] = "pi" - allowed_tools: list[str] | None = None - disallowed_tools: list[str] | None = None - - def _make_adapter( - self, - mcp_servers: dict[str, dict[str, Any]], - approve: base.Approve | None, - ) -> base.Adapter: - return pi.PiAdapter( - cwd=self.cwd, - env=dict(self.env), - executable=self.executable or "pi", - model=self.model, - effort=self.effort, - mcp_servers=dict(mcp_servers), - approve=approve, - allowed_tools=self.allowed_tools, - disallowed_tools=self.disallowed_tools, - ) diff --git a/src/ai/workers/_worker.py b/src/ai/workers/_worker.py index 4670fa65..e4f462b6 100644 --- a/src/ai/workers/_worker.py +++ b/src/ai/workers/_worker.py @@ -38,18 +38,17 @@ from ..types import messages as messages_ from ..types import tools as tools_ from ..types import usage as usage_ -from . import _bridge +from . import _adapter, _bridge, _convert from . import _errors as errors -from ._adapters import base Effort = Literal["low", "medium", "high", "max"] +# Seconds to wait for more bridge calls before handing a batch to the caller. A +# CLI that calls tools in parallel sends them close together. BATCH_WINDOW = 0.05 -"""Seconds to wait for more bridge calls before handing a batch to the -caller. A CLI that calls tools in parallel sends them close together.""" +# How many times to ask again for output that does not validate. OUTPUT_RETRIES = 2 -"""How many times to ask again for output that does not validate.""" OUTPUT_INSTRUCTION = ( "\n\nRespond with ONLY a JSON value matching this schema, " @@ -96,7 +95,7 @@ class _Plan: class _Runtime: """What lives only while a worker is open. Never copied, never dumped.""" - adapter: base.Adapter + adapter: _adapter.CLIAdapter bridge: _bridge.Bridge lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) in_session: bool = False @@ -184,8 +183,8 @@ def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self: def _make_adapter( self, mcp_servers: dict[str, dict[str, Any]], - approve: base.Approve | None, - ) -> base.Adapter: + approve: _adapter.Approve | None, + ) -> _adapter.CLIAdapter: raise NotImplementedError # -- ModelLike ------------------------------------------------------------ @@ -248,7 +247,7 @@ async def stream( if rt.queue is not None: # Tool results came with a new user message. The CLI is # still in its turn: the message is guidance for it. - await rt.adapter.steer(base.text_of(plan.prompt)) + await rt.adapter.steer(_convert.text_of(plan.prompt)) else: await self._ensure_session(plan) self._start(plan.prompt, output_type) @@ -752,7 +751,7 @@ async def tui(self) -> int: return code @classmethod - async def list_sessions(cls, **fields: Any) -> list[base.SessionInfo]: + async def list_sessions(cls, **fields: Any) -> list[_adapter.SessionInfo]: """Return the conversations the CLI has stored for a directory. Takes the worker's fields: `cwd=`, `env=`, `executable=`. diff --git a/src/ai/workers/claude_code/__init__.py b/src/ai/workers/claude_code/__init__.py new file mode 100644 index 00000000..0d2ab501 --- /dev/null +++ b/src/ai/workers/claude_code/__init__.py @@ -0,0 +1,3 @@ +from .worker import ClaudeCode + +__all__ = ["ClaudeCode"] diff --git a/src/ai/workers/_adapters/claude.py b/src/ai/workers/claude_code/adapter.py similarity index 95% rename from src/ai/workers/_adapters/claude.py rename to src/ai/workers/claude_code/adapter.py index 7ca08b65..dae66044 100644 --- a/src/ai/workers/_adapters/claude.py +++ b/src/ai/workers/claude_code/adapter.py @@ -25,18 +25,18 @@ from ...types import events as events_ from ...types import messages as messages_ from ...types import usage as usage_ -from .. import _bridge +from .. import _adapter, _bridge, _convert, _handoff from .. import _errors as errors -from . import _handoff, base +from . import handoff KIND = "claude-code" INSTALL = "npm install -g @anthropic-ai/claude-code" +# The CLI the live tests last passed against. Older ones get a warning. VERIFIED_VERSION = "2.1.289" -"""The CLI the live tests last passed against. Older ones get a warning.""" +# Seconds the CLI waits for an approval. Its own default (60s) would cut a +# person short: the hook would time out and the call would be denied. APPROVAL_TIMEOUT = 24 * 3600 -"""Seconds the CLI waits for an approval. Its own default (60s) would cut a -person short: the hook would time out and the call would be denied.""" AUTH_HINT = ( "claude-code uses your own login: run `claude` once and sign in, or " @@ -54,7 +54,7 @@ def __init__( model: str | None, effort: str | None, mcp_servers: dict[str, Any], - approve: base.Approve | None, + approve: _adapter.Approve | None, allowed_tools: list[str] | None, disallowed_tools: list[str] | None, max_turns: int | None, @@ -82,7 +82,7 @@ def __init__( self.version: str | None = None async def start(self) -> None: - self.version = await base.probe( + self.version = await _adapter.probe( KIND, self._executable, install=INSTALL, @@ -153,7 +153,7 @@ async def _stage(self, history: list[messages_.Message]) -> str: ) key = claude_agent_sdk.project_key_for_directory(self._cwd) path = pathlib.Path(config) / "projects" / key / f"{session_id}.jsonl" - records = _handoff.to_claude_records(history, session_id, self._cwd) + records = handoff.to_claude_records(history, session_id, self._cwd) text = "".join( json.dumps(r, separators=(",", ":")) + "\n" for r in records ) @@ -281,8 +281,10 @@ async def _can_use_tool( # -- turns ---------------------------------------------------------------- - async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: - base.check_files(prompt, KIND, ("image/", "application/pdf")) + async def turn( + self, prompt: _adapter.Prompt + ) -> AsyncIterator[events_.Event]: + _convert.check_files(prompt, KIND, ("image/", "application/pdf")) client = self._client assert client is not None # An interrupted turn leaves its trailing result unread: the CLI @@ -318,7 +320,7 @@ async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: for block in message.content: if isinstance(block, claude_agent_sdk.ToolUseBlock): names[block.id] = block.name - for event in base.builtin_call( + for event in _convert.builtin_call( block.id, block.name, json.dumps(block.input), @@ -333,7 +335,7 @@ async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: ) for block in content_blocks: if isinstance(block, claude_agent_sdk.ToolResultBlock): - yield base.builtin_result( + yield _convert.builtin_result( block.tool_use_id, names.get(block.tool_use_id, ""), block.content, @@ -390,12 +392,12 @@ async def history(self, session_id: str) -> list[messages_.Message]: out.extend(_stored_to_messages(record)) return _handoff.model_view(out) - async def list_sessions(self) -> list[base.SessionInfo]: + async def list_sessions(self) -> list[_adapter.SessionInfo]: infos = await asyncio.to_thread( claude_agent_sdk.list_sessions, self._cwd ) return [ - base.SessionInfo( + _adapter.SessionInfo( session_id=i.session_id, title=i.custom_title or i.summary or i.first_prompt, cwd=i.cwd, @@ -436,7 +438,7 @@ def tui_command( return argv, env, running_as -async def _content(prompt: base.Prompt) -> str | list[dict[str, Any]]: +async def _content(prompt: _adapter.Prompt) -> str | list[dict[str, Any]]: """Return a prompt as the CLI takes it: a string, or content blocks. Images and PDFs travel as base64 blocks (measured: both read correctly). @@ -448,7 +450,9 @@ async def _content(prompt: base.Prompt) -> str | list[dict[str, Any]]: if isinstance(part, messages_.TextPart): blocks.append({"type": "text", "text": part.text}) elif part.media_type.startswith("text/"): - blocks.append({"type": "text", "text": await base.as_text(part)}) + blocks.append( + {"type": "text", "text": await _convert.as_text(part)} + ) else: kind = ( "image" if part.media_type.startswith("image/") else "document" @@ -459,9 +463,9 @@ async def _content(prompt: base.Prompt) -> str | list[dict[str, Any]]: "source": { "type": "base64", "media_type": part.media_type, - "data": base64.b64encode(await base.data(part)).decode( - "ascii" - ), + "data": base64.b64encode( + await _convert.data(part) + ).decode("ascii"), }, } ) diff --git a/src/ai/workers/claude_code/handoff.py b/src/ai/workers/claude_code/handoff.py new file mode 100644 index 00000000..91245a3b --- /dev/null +++ b/src/ai/workers/claude_code/handoff.py @@ -0,0 +1,106 @@ +"""Seeding Claude Code with a past: records for its transcript store. + +The rule for what becomes what is the same for every CLI: see `_handoff`. +""" + +from __future__ import annotations + +import datetime +import json +import uuid +from typing import Any + +from ...types import media +from ...types import messages as messages_ +from .. import _handoff + + +def _source(part: messages_.FilePart) -> dict[str, Any]: + """Return an Anthropic content-block source for a file.""" + if isinstance(part.data, str) and media.is_downloadable_url(part.data): + return {"type": "url", "url": part.data} + return { + "type": "base64", + "media_type": part.media_type, + "data": media.data_to_base64(part.data), + } + + +def to_claude_records( + messages: list[messages_.Message], session_id: str, cwd: str +) -> list[dict[str, Any]]: + """Build Claude Code transcript records for the CLI's store. + + Chained by parentUuid. + + Reasoning parts are omitted: see the module docstring. Callers report + the omission — `reasoning_count()` says how many. + """ + records: list[dict[str, Any]] = [] + parent: str | None = None + + def add(role: str, content: Any) -> None: + nonlocal parent + uid = str(uuid.uuid4()) + ts = ( + datetime.datetime.now(datetime.UTC) + .isoformat() + .replace("+00:00", "Z") + ) + records.append( + { + "parentUuid": parent, + "isSidechain": False, + "type": role, + "message": {"role": role, "content": content}, + "uuid": uid, + "timestamp": ts, + "sessionId": session_id, + "cwd": cwd, + } + ) + parent = uid + + for m in _handoff.harness_view(messages): + blocks: list[dict[str, Any]] = [] + for p in m.parts: + if isinstance(p, messages_.TextPart) and p.text.strip(): + blocks.append({"type": "text", "text": p.text}) + elif isinstance(p, messages_.FilePart) and m.role == "user": + if p.media_type.startswith("image/"): + blocks.append({"type": "image", "source": _source(p)}) + elif p.media_type == "application/pdf": + blocks.append({"type": "document", "source": _source(p)}) + elif isinstance(p, messages_.ToolCallPart): + try: + args = json.loads(p.tool_args) if p.tool_args else {} + except json.JSONDecodeError: + args = {"raw": p.tool_args} + blocks.append( + { + "type": "tool_use", + "id": p.tool_call_id or p.id, + "name": p.tool_name, + "input": args, + } + ) + elif isinstance(p, messages_.ToolResultPart): + blocks.append( + { + "type": "tool_result", + "tool_use_id": p.tool_call_id, + "content": _handoff.result_text(p.result), + **( + {"is_error": True} + if p.result_kind == "error" + else {} + ), + } + ) + # ReasoningPart: omitted on purpose — Anthropic's safeguards + # refuse replayed reasoning. The adapter warns with the count. + if blocks: + # Tool results live in role="tool" messages in the IR; on disk + # the CLI keeps them as user records with tool_result blocks. + add("user" if m.role in ("user", "tool") else "assistant", blocks) + return records diff --git a/src/ai/workers/claude_code/worker.py b/src/ai/workers/claude_code/worker.py new file mode 100644 index 00000000..23af2612 --- /dev/null +++ b/src/ai/workers/claude_code/worker.py @@ -0,0 +1,49 @@ +"""`ClaudeCode`: Claude Code as a worker.""" + +from __future__ import annotations + +from typing import Any, Literal + +from ...providers import _optional +from .. import _adapter, _worker + + +class ClaudeCode(_worker.Worker): + """Claude Code, through `claude-agent-sdk`. Needs the `claude-code` extra. + + `allowed_tools` limits Claude Code's own tools (`["Read", "Grep"]`, or + `[]` for none), and `disallowed_tools` takes some away. `max_turns` caps + the agent's steps in one turn; a turn that hits it ends with + `finish_reason="length"`. + """ + + kind: Literal["claude-code"] = "claude-code" + allowed_tools: list[str] | None = None + disallowed_tools: list[str] | None = None + max_turns: int | None = None + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: _adapter.Approve | None, + ) -> _adapter.CLIAdapter: + _optional.import_optional_sdk( + "claude_agent_sdk", + extra="claude-code", + feature="ai.workers.ClaudeCode", + ) + # Imported here: the adapter needs claude_agent_sdk at import time. + from . import adapter # noqa: PLC0415 + + return adapter.ClaudeAdapter( + cwd=self.cwd, + env=dict(self.env), + executable=self.executable or "claude", + model=self.model, + effort=self.effort, + mcp_servers=dict(mcp_servers), + approve=approve, + allowed_tools=self.allowed_tools, + disallowed_tools=self.disallowed_tools, + max_turns=self.max_turns, + ) diff --git a/src/ai/workers/codex/__init__.py b/src/ai/workers/codex/__init__.py new file mode 100644 index 00000000..37337a31 --- /dev/null +++ b/src/ai/workers/codex/__init__.py @@ -0,0 +1,3 @@ +from .worker import Codex + +__all__ = ["Codex"] diff --git a/src/ai/workers/_adapters/codex.py b/src/ai/workers/codex/adapter.py similarity index 95% rename from src/ai/workers/_adapters/codex.py rename to src/ai/workers/codex/adapter.py index 84dc8809..68bcb4b2 100644 --- a/src/ai/workers/_adapters/codex.py +++ b/src/ai/workers/codex/adapter.py @@ -22,19 +22,19 @@ from ...types import events as events_ from ...types import messages as messages_ from ...types import usage as usage_ -from .. import _bridge +from .. import _adapter, _bridge, _convert, _handoff from .. import _errors as errors -from . import _handoff, base, jsonrpc +from . import handoff, jsonrpc KIND = "codex" INSTALL = "npm install -g @openai/codex" +# The CLI the live tests last passed against. Older ones get a warning. VERIFIED_VERSION = "0.160.1" -"""The CLI the live tests last passed against. Older ones get a warning.""" +# Put on the turn's queue when app-server dies. CRASHED = "__crashed__" -"""Put on the turn's queue when app-server dies.""" LIST_PAGE_SIZE = 100 +# `list_sessions()` reads at most this many pages of threads. LIST_PAGES = 50 -"""`list_sessions()` reads at most this many pages of threads.""" APPROVAL_METHODS = ( "item/commandExecution/requestApproval", @@ -44,6 +44,7 @@ "applyPatchApproval", ) +# Item types that are the agent doing something, rather than speaking. TOOL_ITEMS = ( "fileChange", "commandExecution", @@ -51,7 +52,6 @@ "webSearch", "toolCall", ) -"""Item types that are the agent doing something, rather than speaking.""" AUTH_HINT = ( "codex uses your own login: run `codex login`, or pass " @@ -69,7 +69,7 @@ def __init__( executable: str, model: str | None, mcp_servers: dict[str, Any], - approve: base.Approve | None, + approve: _adapter.Approve | None, sandbox: str | None, config: dict[str, Any], ) -> None: @@ -88,7 +88,7 @@ def __init__( # CLI takes. An explicit config wins over the MCP servers. self._config = {**mcp_config(mcp_servers), **config} self._instructions: str | None = None - self._process: base.Process | None = None + self._process: jsonrpc.Process | None = None self._peer: jsonrpc.JsonRpcPeer | None = None self._thread: str | None = None self._turn: str | None = None @@ -98,7 +98,7 @@ def __init__( self.version: str | None = None async def start(self) -> None: - self.version = await base.probe( + self.version = await _adapter.probe( KIND, self._executable, install=INSTALL, @@ -106,7 +106,7 @@ async def start(self) -> None: cwd=self._cwd, env=self._env, ) - self._process = await base.spawn( + self._process = await jsonrpc.spawn( [self._executable, "app-server", *_config_args(self._config)], cwd=self._cwd, env=self._env, @@ -206,7 +206,7 @@ async def open( _handoff.warn_unsupported_files(seed, KIND) await self._rpc().call( "thread/inject_items", - {"threadId": thread, "items": _handoff.to_codex_items(seed)}, + {"threadId": thread, "items": handoff.to_codex_items(seed)}, ) return thread @@ -277,8 +277,10 @@ async def _on_request( # -- turns ---------------------------------------------------------------- - async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: - base.check_files(prompt, KIND, ("image/",)) + async def turn( + self, prompt: _adapter.Prompt + ) -> AsyncIterator[events_.Event]: + _convert.check_files(prompt, KIND, ("image/",)) # Unbounded: the reader fills it inline and must never wait on a # slow consumer. queue: asyncio.Queue[tuple[str, dict[str, Any]]] = asyncio.Queue() @@ -366,14 +368,14 @@ async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: open_reasoning.discard(item_id) yield events_.ReasoningEnd(block_id=item_id) elif kind in TOOL_ITEMS: - for event in base.builtin_call( + for event in _convert.builtin_call( item_id, _tool_name(item), json.dumps(_tool_args(item)), started=started, ): yield event - yield base.builtin_result( + yield _convert.builtin_result( item_id, _tool_name(item), _item_result(item), @@ -452,7 +454,7 @@ async def history(self, session_id: str) -> list[messages_.Message]: stored.extend(_item_to_messages(item)) return _handoff.model_view(stored) - async def list_sessions(self) -> list[base.SessionInfo]: + async def list_sessions(self) -> list[_adapter.SessionInfo]: # Asked for THIS directory, page by page: one unfiltered page of the # machine's store (measured: 3159 threads) missed threads from here. threads: list[dict[str, Any]] = [] @@ -471,7 +473,7 @@ async def list_sessions(self) -> list[base.SessionInfo]: if not cursor: break return [ - base.SessionInfo( + _adapter.SessionInfo( session_id=t.get("id", ""), title=t.get("name") or t.get("preview") or None, cwd=t.get("cwd"), @@ -507,7 +509,7 @@ def tui_command( return argv, dict(self._env), None -async def _input_items(prompt: base.Prompt) -> list[dict[str, Any]]: +async def _input_items(prompt: _adapter.Prompt) -> list[dict[str, Any]]: """Return a prompt as app-server's `UserInput` items. An image travels as a data URL, built here from its bytes. @@ -517,9 +519,11 @@ async def _input_items(prompt: base.Prompt) -> list[dict[str, Any]]: if isinstance(part, messages_.TextPart): items.append({"type": "text", "text": part.text}) elif part.media_type.startswith("text/"): - items.append({"type": "text", "text": await base.as_text(part)}) + items.append({"type": "text", "text": await _convert.as_text(part)}) else: - encoded = base64.b64encode(await base.data(part)).decode("ascii") + encoded = base64.b64encode(await _convert.data(part)).decode( + "ascii" + ) items.append( { "type": "image", diff --git a/src/ai/workers/codex/handoff.py b/src/ai/workers/codex/handoff.py new file mode 100644 index 00000000..ea637d78 --- /dev/null +++ b/src/ai/workers/codex/handoff.py @@ -0,0 +1,86 @@ +"""Seeding Codex with a past: Responses API items for `thread/inject_items`. + +The rule for what becomes what is the same for every CLI: see `_handoff`. +""" + +from __future__ import annotations + +from typing import Any + +from ...types import media +from ...types import messages as messages_ +from .. import _handoff + +NOTE_PREFIX = "[Handoff from a previous session — earlier reasoning: " + + +def _note(reasoning: list[str]) -> str: + return ( + NOTE_PREFIX + " ".join(r.strip() for r in reasoning if r.strip()) + "]" + ) + + +def _image_url(part: messages_.FilePart) -> str: + """Return an http(s) URL as it is, anything else as a data URL.""" + return media.data_to_data_url(part.data, part.media_type) + + +def to_codex_items(messages: list[messages_.Message]) -> list[dict[str, Any]]: + """Responses API items for `thread/inject_items`.""" + items: list[dict[str, Any]] = [] + for m in _handoff.harness_view(messages): + reasoning: list[str] = [] + for p in m.parts: + if isinstance(p, messages_.TextPart) and p.text.strip(): + ctype = "input_text" if m.role == "user" else "output_text" + items.append( + { + "type": "message", + "role": m.role, + "content": [{"type": ctype, "text": p.text}], + } + ) + elif ( + isinstance(p, messages_.FilePart) + and m.role == "user" + and p.media_type.startswith("image/") + ): + items.append( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_image", "image_url": _image_url(p)} + ], + } + ) + elif isinstance(p, messages_.ToolCallPart): + items.append( + { + "type": "function_call", + "call_id": p.tool_call_id or p.id, + "name": p.tool_name, + "arguments": p.tool_args or "{}", + } + ) + elif isinstance(p, messages_.ToolResultPart): + items.append( + { + "type": "function_call_output", + "call_id": p.tool_call_id, + "output": _handoff.result_text(p.result), + } + ) + elif isinstance(p, messages_.ReasoningPart) and p.text.strip(): + reasoning.append(p.text) + if reasoning: + items.append( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": _note(reasoning)} + ], + } + ) + return items diff --git a/src/ai/workers/_adapters/jsonrpc.py b/src/ai/workers/codex/jsonrpc.py similarity index 68% rename from src/ai/workers/_adapters/jsonrpc.py rename to src/ai/workers/codex/jsonrpc.py index 14837088..f5b0a8c1 100644 --- a/src/ai/workers/_adapters/jsonrpc.py +++ b/src/ai/workers/codex/jsonrpc.py @@ -1,4 +1,4 @@ -"""A small JSON-RPC 2.0 peer over a CLI's stdio. +"""The transport to `codex app-server`: its process, and JSON-RPC over stdio. Line-delimited JSON, request/response futures, and two callbacks for what a server can start. Nothing here is protocol-specific: the Codex adapter @@ -10,13 +10,19 @@ import asyncio import contextlib import json -from collections.abc import Awaitable, Callable -from typing import TYPE_CHECKING, Any +import os +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Any from .. import _errors as errors -if TYPE_CHECKING: - from . import base +# Longest line app-server may write. One JSON line can be megabytes (a +# `thread/read` of a long thread). With asyncio's default (64 KiB) such a line +# failed as if the CLI had crashed. +LINE_LIMIT = 256 * 1024 * 1024 + +# Seconds app-server gets to exit after SIGTERM, before SIGKILL. +TERM_GRACE = 3.0 NotifyHandler = Callable[[str, dict[str, Any]], Awaitable[None]] RequestHandler = Callable[[Any, str, dict[str, Any]], Awaitable[None]] @@ -30,7 +36,7 @@ def __init__(self, error: dict[str, Any]) -> None: class JsonRpcPeer: - def __init__(self, process: base.Process, *, kind: str) -> None: + def __init__(self, process: Process, *, kind: str) -> None: self._process = process self._kind = kind self._next_id = 0 @@ -208,3 +214,81 @@ def _resolve(self, message: dict[str, Any]) -> None: future.set_exception(JsonRpcError(message["error"])) else: future.set_result(message.get("result")) + + +# -- the process ------------------------------------------------------------- + + +class Process: + """A CLI started by `spawn`: JSON lines on stdin and stdout.""" + + def __init__(self, proc: asyncio.subprocess.Process) -> None: + self._proc = proc + + async def write(self, data: str) -> None: + assert self._proc.stdin is not None + self._proc.stdin.write(data.encode()) + await self._proc.stdin.drain() + + async def readline(self) -> str: + assert self._proc.stdout is not None + return (await self._proc.stdout.readline()).decode() + + async def read_stderr(self) -> str: + if self._proc.stderr is None: + return "" + return (await self._proc.stderr.read()).decode(errors="replace") + + async def wait(self) -> int: + return await self._proc.wait() + + async def terminate(self) -> None: + """Stop the process group: SIGTERM, then SIGKILL.""" + if self._proc.returncode is not None: + return + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(self._proc.pid), 15) + try: + async with asyncio.timeout(TERM_GRACE): + await self._proc.wait() + except TimeoutError: + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(self._proc.pid), 9) + with contextlib.suppress(Exception): + await self._proc.wait() + + async def death_note(self) -> str: + """Say how the process ended, and its last words.""" + code: int | None = None + tail = "" + with contextlib.suppress(Exception): + async with asyncio.timeout(2): + code = await self.wait() + with contextlib.suppress(Exception): + async with asyncio.timeout(2): + tail = (await self.read_stderr()).strip() + how = f"exit code {code}" if code is not None else "exited" + if code in (137, -9): + how += " (SIGKILL, typically the kernel's OOM killer)" + return f"{how}; stderr: {tail[-800:]}" if tail else f"{how}; no stderr" + + +async def spawn( + argv: Sequence[str], *, cwd: str, env: Mapping[str, str] +) -> Process: + """Start a CLI in its own process group, with pipes for stdio.""" + try: + proc = await asyncio.create_subprocess_exec( + *argv, + cwd=cwd, + env={**os.environ, **env}, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + # Own the whole tree, so teardown reaches grandchildren too. + start_new_session=True, + limit=LINE_LIMIT, + ) + except OSError as exc: + raise errors.WorkerError(f"cannot start {argv[0]!r}: {exc}") from exc + return Process(proc) diff --git a/src/ai/workers/codex/worker.py b/src/ai/workers/codex/worker.py new file mode 100644 index 00000000..f0293481 --- /dev/null +++ b/src/ai/workers/codex/worker.py @@ -0,0 +1,54 @@ +"""`Codex`: Codex as a worker.""" + +from __future__ import annotations + +from typing import Any, Literal + +import pydantic + +from .. import _adapter, _bridge, _worker +from . import adapter + + +class Codex(_worker.Worker): + """Codex, over `codex app-server`. + + `sandbox` is Codex's own sandbox mode: `"read-only"`, + `"workspace-write"` or `"danger-full-access"` (the default). `config` + takes Codex `-c` overrides. + """ + + kind: Literal["codex"] = "codex" + sandbox: ( + Literal["read-only", "workspace-write", "danger-full-access"] | None + ) = None + config: dict[str, Any] = pydantic.Field(default_factory=dict) + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: _adapter.Approve | None, + ) -> _adapter.CLIAdapter: + config = dict(self.config) + if self.effort is not None: + config.setdefault( + "model_reasoning_effort", + "xhigh" if self.effort == "max" else self.effort, + ) + if _bridge.SERVER_NAME in mcp_servers: + # Codex's default wait for an MCP tool (60s) is too short for a + # tool an `ai.Agent` runs, maybe after an approval. + config.setdefault( + f"mcp_servers.{_bridge.SERVER_NAME}.tool_timeout_sec", + 24 * 3600, + ) + return adapter.CodexAdapter( + cwd=self.cwd, + env=dict(self.env), + executable=self.executable or "codex", + model=self.model, + mcp_servers=dict(mcp_servers), + approve=approve, + sandbox=self.sandbox, + config=config, + ) diff --git a/src/ai/workers/pi/__init__.py b/src/ai/workers/pi/__init__.py new file mode 100644 index 00000000..af788365 --- /dev/null +++ b/src/ai/workers/pi/__init__.py @@ -0,0 +1,3 @@ +from .worker import Pi + +__all__ = ["Pi"] diff --git a/src/ai/workers/_adapters/pi.py b/src/ai/workers/pi/adapter.py similarity index 89% rename from src/ai/workers/_adapters/pi.py rename to src/ai/workers/pi/adapter.py index 6c64c731..bd44d73e 100644 --- a/src/ai/workers/_adapters/pi.py +++ b/src/ai/workers/pi/adapter.py @@ -5,7 +5,7 @@ the same conversation in pi's own TUI. pi has no permission prompt. Approval is an extension loaded into every pi -a worker starts (`pi_extension.js`): its `tool_call` handler asks through an +a worker starts (`extension.js`): its `tool_call` handler asks through an RPC dialog, and this adapter answers it. The protocol shapes here were measured against a real pi (see VERIFIED_VERSION). """ @@ -22,35 +22,43 @@ import re import uuid import warnings -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any from ...types import events as events_ from ...types import messages as messages_ from ...types import usage as usage_ -from .. import _bridge +from .. import _adapter, _bridge, _convert, _handoff from .. import _errors as errors -from . import _handoff, base +from . import handoff KIND = "pi" INSTALL = "npm install -g @earendil-works/pi-coding-agent" +# The CLI the live tests last passed against. Older ones get a warning. VERIFIED_VERSION = "1.3.0" -"""The CLI the live tests last passed against. Older ones get a warning.""" -EXTENSION = str(pathlib.Path(__file__).with_name("pi_extension.js")) +EXTENSION = str(pathlib.Path(__file__).with_name("extension.js")) +# The dialog title the extension asks approvals under. APPROVE = "ai-python:approve" -"""The dialog title the extension asks approvals under.""" +# The command the extension registers last: proof that it loaded. READY = "ai-python-ready" -"""The command the extension registers last: proof that it loaded.""" +# The environment variable that configures the extension. CONFIG_VAR = "AI_PYTHON_PI" -"""The environment variable that configures the extension.""" +# pi's rule for a session id: letters, digits, `.`, `_` and `-`, starting and +# ending with a letter or digit. Never a path. SESSION_ID = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9._-]{0,198}[A-Za-z0-9])?\Z") -"""pi's rule for a session id: letters, digits, `.`, `_` and `-`, starting -and ending with a letter or digit. Never a path.""" +# Extension UI methods that wait for an answer; the rest are notices. DIALOGS = ("select", "confirm", "input", "editor") -"""Extension UI methods that wait for an answer; the rest are notices.""" CRASHED = "__crashed__" CLOSED = "__closed__" +# Longest line pi may write. The protocol is JSON lines, and one line can be +# megabytes. With asyncio's default (64 KiB) such a line fails as if the CLI had +# crashed. +LINE_LIMIT = 256 * 1024 * 1024 + +# Seconds pi gets to exit after SIGTERM, before SIGKILL. +TERM_GRACE = 3.0 + AUTH_HINT = ( "pi uses its own login (`/login`, auth.json) or a provider key in the " "environment: pass env={'ANTHROPIC_API_KEY': ...}." @@ -67,7 +75,7 @@ def __init__( model: str | None, effort: str | None, mcp_servers: dict[str, Any], - approve: base.Approve | None, + approve: _adapter.Approve | None, allowed_tools: list[str] | None, disallowed_tools: list[str] | None, ) -> None: @@ -89,7 +97,7 @@ def __init__( self.version: str | None = None async def start(self) -> None: - self.version = await base.probe( + self.version = await _adapter.probe( KIND, self._executable, install=INSTALL, @@ -169,7 +177,7 @@ async def _spawn(self, session_id: str, args: list[str]) -> None: A pi that cannot start (an unknown model, a bad flag) exits at once, and the handshake reports its last words. """ - process = await base.spawn( + process = await spawn( [self._executable, "--mode", "rpc", *args, *self._args()], cwd=self._cwd, env=self._extension_env( @@ -264,7 +272,7 @@ async def _stage(self, history: list[messages_.Message]) -> str: *changes, *( {"type": "message", "message": message} - for message in _handoff.to_pi_messages( + for message in handoff.to_pi_messages( history, timestamp=int(now.timestamp() * 1000) ) ), @@ -331,8 +339,10 @@ async def _on_dialog(self, agent: _Agent, request: dict[str, Any]) -> None: # -- turns ---------------------------------------------------------------- - async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: - base.check_files(prompt, KIND, ("image/",)) + async def turn( + self, prompt: _adapter.Prompt + ) -> AsyncIterator[events_.Event]: + _convert.check_files(prompt, KIND, ("image/",)) agent = self._agent assert agent is not None if agent.dead is not None: @@ -421,7 +431,7 @@ async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: # transcript: only the tool that made it reports it. call_id = str(event.get("toolCallId") or "") name = str(event.get("toolName") or "") - yield base.builtin_result( + yield _convert.builtin_result( call_id, name, _tool_result(event.get("result") or {}), @@ -554,9 +564,9 @@ async def history(self, session_id: str) -> list[messages_.Message]: stored.extend(_stored_to_messages(record)) return _handoff.model_view(stored) - async def list_sessions(self) -> list[base.SessionInfo]: + async def list_sessions(self) -> list[_adapter.SessionInfo]: return [ - base.SessionInfo( + _adapter.SessionInfo( session_id=s["id"], title=s.get("title"), cwd=s.get("cwd"), @@ -584,6 +594,84 @@ def tui_command( ) +# -- the process ------------------------------------------------------------- + + +class Process: + """A CLI started by `spawn`: JSON lines on stdin and stdout.""" + + def __init__(self, proc: asyncio.subprocess.Process) -> None: + self._proc = proc + + async def write(self, data: str) -> None: + assert self._proc.stdin is not None + self._proc.stdin.write(data.encode()) + await self._proc.stdin.drain() + + async def readline(self) -> str: + assert self._proc.stdout is not None + return (await self._proc.stdout.readline()).decode() + + async def read_stderr(self) -> str: + if self._proc.stderr is None: + return "" + return (await self._proc.stderr.read()).decode(errors="replace") + + async def wait(self) -> int: + return await self._proc.wait() + + async def terminate(self) -> None: + """Stop the process group: SIGTERM, then SIGKILL.""" + if self._proc.returncode is not None: + return + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(self._proc.pid), 15) + try: + async with asyncio.timeout(TERM_GRACE): + await self._proc.wait() + except TimeoutError: + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(self._proc.pid), 9) + with contextlib.suppress(Exception): + await self._proc.wait() + + async def death_note(self) -> str: + """Say how the process ended, and its last words.""" + code: int | None = None + tail = "" + with contextlib.suppress(Exception): + async with asyncio.timeout(2): + code = await self.wait() + with contextlib.suppress(Exception): + async with asyncio.timeout(2): + tail = (await self.read_stderr()).strip() + how = f"exit code {code}" if code is not None else "exited" + if code in (137, -9): + how += " (SIGKILL, typically the kernel's OOM killer)" + return f"{how}; stderr: {tail[-800:]}" if tail else f"{how}; no stderr" + + +async def spawn( + argv: Sequence[str], *, cwd: str, env: Mapping[str, str] +) -> Process: + """Start a CLI in its own process group, with pipes for stdio.""" + try: + proc = await asyncio.create_subprocess_exec( + *argv, + cwd=cwd, + env={**os.environ, **env}, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + # Own the whole tree, so teardown reaches grandchildren too. + start_new_session=True, + limit=LINE_LIMIT, + ) + except OSError as exc: + raise errors.WorkerError(f"cannot start {argv[0]!r}: {exc}") from exc + return Process(proc) + + class _CommandError(Exception): """pi answered a command with `success: false`.""" @@ -597,7 +685,7 @@ class _Agent: including the reply to an `abort` sent from inside it. """ - def __init__(self, process: base.Process, *, on_dialog: Any) -> None: + def __init__(self, process: Process, *, on_dialog: Any) -> None: self.process = process self.events: asyncio.Queue[dict[str, Any]] = asyncio.Queue() self.dead: str | None = None @@ -730,7 +818,7 @@ def _resolve(self, record: dict[str, Any]) -> None: future.set_result(record.get("data")) -async def _prompt_command(prompt: base.Prompt) -> dict[str, Any]: +async def _prompt_command(prompt: _adapter.Prompt) -> dict[str, Any]: """Return a prompt as pi's `prompt` command: text, and images as base64.""" texts: list[str] = [] images: list[dict[str, Any]] = [] @@ -738,12 +826,12 @@ async def _prompt_command(prompt: base.Prompt) -> dict[str, Any]: if isinstance(part, messages_.TextPart): texts.append(part.text) elif part.media_type.startswith("text/"): - texts.append(await base.as_text(part)) + texts.append(await _convert.as_text(part)) else: images.append( { "type": "image", - "data": base64.b64encode(await base.data(part)).decode( + "data": base64.b64encode(await _convert.data(part)).decode( "ascii" ), "mimeType": part.media_type, @@ -795,7 +883,7 @@ def _stream_events( ) elif kind == "toolcall_end": call = _tool_call(update.get("toolCall") or {}) - out += base.builtin_call( + out += _convert.builtin_call( call.tool_call_id, call.tool_name, call.tool_args, started=started ) return out @@ -832,7 +920,7 @@ def _unstreamed( ] elif kind == "toolCall" and str(block.get("id")) not in started: call = _tool_call(block) - out += base.builtin_call( + out += _convert.builtin_call( call.tool_call_id, call.tool_name, call.tool_args, diff --git a/src/ai/workers/_adapters/pi_extension.js b/src/ai/workers/pi/extension.js similarity index 100% rename from src/ai/workers/_adapters/pi_extension.js rename to src/ai/workers/pi/extension.js diff --git a/src/ai/workers/pi/handoff.py b/src/ai/workers/pi/handoff.py new file mode 100644 index 00000000..5f4d20c6 --- /dev/null +++ b/src/ai/workers/pi/handoff.py @@ -0,0 +1,148 @@ +"""Seeding pi with a past: `AgentMessage`s for a session file. + +The rule for what becomes what is the same for every CLI: see `_handoff`. +""" + +from __future__ import annotations + +import json +from typing import Any + +from ...types import media +from ...types import messages as messages_ +from .. import _handoff + +# What names a reply's author in a pi session: its API and model. +AUTHOR = ("api", "provider", "model") +SYSTEM_PREFIX = "[Handoff from a previous session — instructions it was given: " + + +def to_pi_messages( + messages: list[messages_.Message], *, timestamp: int +) -> list[dict[str, Any]]: + """Build pi `AgentMessage`s for a session file. + + Assistant messages pi wrote keep their author (`provider_metadata["pi"]`): + pi restores a session's model from it. Any other is attributed to a + model pi does not run (`ai-python/handoff`), so pi treats it as another + model's — it keeps its text and tool calls and normalizes the call ids + for its own provider. Consecutive assistant messages are merged, the way one + model response holds text and calls together. Reasoning is omitted: see + the module docstring. + """ + out: list[dict[str, Any]] = [] + usage = { + "input": 0, + "output": 0, + "cacheRead": 0, + "cacheWrite": 0, + "totalTokens": 0, + "cost": { + "input": 0, + "output": 0, + "cacheRead": 0, + "cacheWrite": 0, + "total": 0, + }, + } + for m in _handoff.harness_view(messages): + if m.role == "tool": + for p in m.parts: + if isinstance(p, messages_.ToolResultPart): + out.append( + { + "role": "toolResult", + "toolCallId": p.tool_call_id, + "toolName": p.tool_name, + "content": [ + { + "type": "text", + "text": _handoff.result_text(p.result), + } + ], + "isError": p.result_kind == "error", + "timestamp": timestamp, + } + ) + continue + blocks: list[dict[str, Any]] = [] + for p in m.parts: + if isinstance(p, messages_.TextPart) and p.text.strip(): + blocks.append({"type": "text", "text": p.text}) + elif ( + isinstance(p, messages_.FilePart) + and m.role == "user" + and p.media_type.startswith("image/") + ): + if isinstance(p.data, str) and media.is_downloadable_url( + p.data + ): + blocks.append( + {"type": "text", "text": f"[image: {p.data}]"} + ) + else: + blocks.append( + { + "type": "image", + "data": media.data_to_base64(p.data), + "mimeType": p.media_type, + } + ) + elif ( + isinstance(p, messages_.ToolCallPart) and m.role == "assistant" + ): + try: + args = json.loads(p.tool_args) if p.tool_args else {} + except json.JSONDecodeError: + args = {"raw": p.tool_args} + blocks.append( + { + "type": "toolCall", + "id": p.tool_call_id or p.id, + "name": p.tool_name, + "arguments": args if isinstance(args, dict) else {}, + } + ) + if not blocks: + continue + if m.role == "system": + # Measured (pi 1.3.0): a system message in a session file is + # pi's own prompt state, and replaying one REPLACED pi's system + # prompt and its tools. So instructions travel as what they are + # to the agent — the caller's words — never as its own output. + text = " ".join(b["text"] for b in blocks if b["type"] == "text") + blocks = [{"type": "text", "text": SYSTEM_PREFIX + text + "]"}] + if m.role in ("user", "system"): + out.append( + {"role": "user", "content": blocks, "timestamp": timestamp} + ) + continue + if out and out[-1]["role"] == "assistant": + out[-1]["content"] += blocks + else: + # A reply pi wrote keeps its author: pi restores a session's model + # from its last reply's, and a model it does not know (ours) + # put the conversation on pi's default one (measured). + author = (m.provider_metadata or {}).get("pi") or {} + known = all(isinstance(author.get(k), str) for k in AUTHOR) + out.append( + { + "role": "assistant", + "content": blocks, + **( + {k: author[k] for k in AUTHOR} + if known + else { + "api": "ai-python", + "provider": "ai-python", + "model": "handoff", + } + ), + "usage": usage, + "stopReason": "stop", + "timestamp": timestamp, + } + ) + if any(b["type"] == "toolCall" for b in out[-1]["content"]): + out[-1]["stopReason"] = "toolUse" + return out diff --git a/src/ai/workers/pi/worker.py b/src/ai/workers/pi/worker.py new file mode 100644 index 00000000..0bde399f --- /dev/null +++ b/src/ai/workers/pi/worker.py @@ -0,0 +1,37 @@ +"""`Pi`: pi as a worker.""" + +from __future__ import annotations + +from typing import Any, Literal + +from .. import _adapter, _worker +from . import adapter + + +class Pi(_worker.Worker): + """Pi, over `pi --mode rpc`. + + `allowed_tools` limits pi's own tools (`["read", "grep"]`, or `[]` for + none), and `disallowed_tools` takes some away. + """ + + kind: Literal["pi"] = "pi" + allowed_tools: list[str] | None = None + disallowed_tools: list[str] | None = None + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: _adapter.Approve | None, + ) -> _adapter.CLIAdapter: + return adapter.PiAdapter( + cwd=self.cwd, + env=dict(self.env), + executable=self.executable or "pi", + model=self.model, + effort=self.effort, + mcp_servers=dict(mcp_servers), + approve=approve, + allowed_tools=self.allowed_tools, + disallowed_tools=self.disallowed_tools, + ) diff --git a/src/ai/workers/_adapters/__init__.py b/tests/workers/claude_code/__init__.py similarity index 100% rename from src/ai/workers/_adapters/__init__.py rename to tests/workers/claude_code/__init__.py diff --git a/tests/workers/_adapters/test_claude.py b/tests/workers/claude_code/test_adapter.py similarity index 96% rename from tests/workers/_adapters/test_claude.py rename to tests/workers/claude_code/test_adapter.py index 6854d5e4..a8d4d801 100644 --- a/tests/workers/_adapters/test_claude.py +++ b/tests/workers/claude_code/test_adapter.py @@ -15,7 +15,7 @@ from ai import workers # noqa: E402 from ai.types import events # noqa: E402 from ai.types import messages as m # noqa: E402 -from ai.workers._adapters import claude # noqa: E402 +from ai.workers.claude_code import adapter as adapter_ # noqa: E402 def _stream(event: dict[str, Any]) -> Any: @@ -109,8 +109,8 @@ async def disconnect(self) -> None: pass -def _adapter(tmp_path: Path, **options: Any) -> claude.ClaudeAdapter: - return claude.ClaudeAdapter( +def _adapter(tmp_path: Path, **options: Any) -> adapter_.ClaudeAdapter: + return adapter_.ClaudeAdapter( cwd=str(tmp_path), env={"CLAUDE_CONFIG_DIR": str(tmp_path / "config")}, executable="claude", @@ -124,7 +124,7 @@ def _adapter(tmp_path: Path, **options: Any) -> claude.ClaudeAdapter: ) -async def _turn(adapter: claude.ClaudeAdapter, *prompt: Any) -> list[Any]: +async def _turn(adapter: adapter_.ClaudeAdapter, *prompt: Any) -> list[Any]: async with asyncio.timeout(5): return [e async for e in adapter.turn(list(prompt))] @@ -286,9 +286,9 @@ def test_a_stored_transcript_reads_as_the_model_view() -> None: parent_tool_use_id=None, ), ] - from ai.workers._adapters import _handoff + from ai.workers import _handoff - stored = [msg for r in records for msg in claude._stored_to_messages(r)] + stored = [msg for r in records for msg in adapter_._stored_to_messages(r)] history = _handoff.model_view(stored) assert [x.role for x in history] == ["user", "assistant"] call, result = history[1].parts diff --git a/tests/workers/_adapters/__init__.py b/tests/workers/codex/__init__.py similarity index 100% rename from tests/workers/_adapters/__init__.py rename to tests/workers/codex/__init__.py diff --git a/tests/workers/_adapters/test_codex.py b/tests/workers/codex/test_adapter.py similarity index 98% rename from tests/workers/_adapters/test_codex.py rename to tests/workers/codex/test_adapter.py index de1fedb9..16e82303 100644 --- a/tests/workers/_adapters/test_codex.py +++ b/tests/workers/codex/test_adapter.py @@ -12,7 +12,8 @@ from ai import workers from ai.types import events from ai.types import messages as m -from ai.workers._adapters import codex, jsonrpc +from ai.workers.codex import adapter as adapter_ +from ai.workers.codex import jsonrpc from .. import fakes @@ -98,8 +99,8 @@ def note(method: str, **params: Any) -> dict[str, Any]: def _adapter( server: AppServer, approve: Any = None, **options: Any -) -> codex.CodexAdapter: - adapter = codex.CodexAdapter( +) -> adapter_.CodexAdapter: + adapter = adapter_.CodexAdapter( cwd=CWD, env={}, executable="codex", @@ -119,7 +120,7 @@ def _adapter( return adapter -async def _turn(adapter: codex.CodexAdapter, text: str = "go") -> list[Any]: +async def _turn(adapter: adapter_.CodexAdapter, text: str = "go") -> list[Any]: async with asyncio.timeout(5): return [e async for e in adapter.turn([m.TextPart(text=text)])] diff --git a/tests/workers/_adapters/test_jsonrpc.py b/tests/workers/codex/test_jsonrpc.py similarity index 99% rename from tests/workers/_adapters/test_jsonrpc.py rename to tests/workers/codex/test_jsonrpc.py index a6849316..23deb002 100644 --- a/tests/workers/_adapters/test_jsonrpc.py +++ b/tests/workers/codex/test_jsonrpc.py @@ -9,7 +9,7 @@ import pytest from ai import workers -from ai.workers._adapters import jsonrpc +from ai.workers.codex import jsonrpc from .. import fakes diff --git a/tests/workers/fakes.py b/tests/workers/fakes.py index 51ed6b6c..200e8525 100644 --- a/tests/workers/fakes.py +++ b/tests/workers/fakes.py @@ -15,7 +15,7 @@ from ai.types import events from ai.types import messages as m from ai.types import usage as usage_ -from ai.workers._adapters import base +from ai.workers import _adapter, _convert @dataclasses.dataclass @@ -56,8 +56,8 @@ def tool( ) -> list[Step]: """One of the CLI's own tool calls, start to result.""" return [ - *base.builtin_call(call_id, name, args, started=set()), - base.builtin_result(call_id, name, result, is_error=False), + *_convert.builtin_call(call_id, name, args, started=set()), + _convert.builtin_result(call_id, name, result, is_error=False), ] @@ -67,7 +67,9 @@ def reply(prompt: str) -> list[Step]: class FakeAdapter: - def __init__(self, script: Script, approve: base.Approve | None) -> None: + def __init__( + self, script: Script, approve: _adapter.Approve | None + ) -> None: self.script = script self.approve = approve self.version: str | None = None @@ -106,11 +108,13 @@ async def open( return f"{session_id}-fork" return session_id or "s1" - async def turn(self, prompt: base.Prompt) -> AsyncIterator[events.Event]: - self.turns.append(base.text_of(prompt)) + async def turn( + self, prompt: _adapter.Prompt + ) -> AsyncIterator[events.Event]: + self.turns.append(_convert.text_of(prompt)) self._resume.clear() yield events.StreamStart() - for step in self.script(base.text_of(prompt)): + for step in self.script(_convert.text_of(prompt)): if isinstance(step, Pause): await self._resume.wait() yield events.StreamEnd(finish_reason="cancelled") @@ -141,8 +145,8 @@ async def interrupt(self) -> None: async def history(self, session_id: str) -> list[m.Message]: return self.stored.get(session_id, []) - async def list_sessions(self) -> list[base.SessionInfo]: - return [base.SessionInfo(session_id="s1", title="first")] + async def list_sessions(self) -> list[_adapter.SessionInfo]: + return [_adapter.SessionInfo(session_id="s1", title="first")] def tui_command( self, session_id: str | None @@ -167,8 +171,8 @@ class FakeWorker(workers.Worker): def _make_adapter( self, mcp_servers: dict[str, dict[str, Any]], - approve: base.Approve | None, - ) -> base.Adapter: + approve: _adapter.Approve | None, + ) -> _adapter.CLIAdapter: self._fake = FakeAdapter(self._run, approve) self._servers = mcp_servers return self._fake @@ -203,7 +207,6 @@ def __init__( self._lines: asyncio.Queue[str] = asyncio.Queue() self.answer = answer or (lambda message: {}) self.written: list[dict[str, Any]] = [] - self.pid = 0 def server_says(self, message: dict[str, Any]) -> None: self._lines.put_nowait(json.dumps(message) + "\n") diff --git a/tests/workers/pi/__init__.py b/tests/workers/pi/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/workers/_adapters/test_pi.py b/tests/workers/pi/test_adapter.py similarity index 97% rename from tests/workers/_adapters/test_pi.py rename to tests/workers/pi/test_adapter.py index 86f99ccb..a3ce5e9d 100644 --- a/tests/workers/_adapters/test_pi.py +++ b/tests/workers/pi/test_adapter.py @@ -16,7 +16,7 @@ from ai import workers from ai.types import events from ai.types import messages as m -from ai.workers._adapters import pi +from ai.workers.pi import adapter as adapter_ from .. import fakes @@ -147,8 +147,8 @@ def sent(self, kind: str) -> list[dict[str, Any]]: def _adapter( server: Pi | None = None, approve: Any = None, cwd: str = "/w", **env: str -) -> pi.PiAdapter: - adapter = pi.PiAdapter( +) -> adapter_.PiAdapter: + adapter = adapter_.PiAdapter( cwd=cwd, env=env, executable="pi", @@ -160,7 +160,7 @@ def _adapter( disallowed_tools=None, ) if server is not None: - agent = pi._Agent( + agent = adapter_._Agent( cast("Any", server.process), on_dialog=adapter._on_dialog ) agent.start() @@ -168,7 +168,7 @@ def _adapter( return adapter -async def _turn(adapter: pi.PiAdapter, *prompt: Any) -> list[Any]: +async def _turn(adapter: adapter_.PiAdapter, *prompt: Any) -> list[Any]: parts: list[m.TextPart | m.FilePart] = list(prompt) or [ m.TextPart(text="go") ] @@ -251,7 +251,7 @@ def _ask(server: Pi) -> None: "type": "extension_ui_request", "id": "ui-1", "method": "input", - "title": pi.APPROVE, + "title": adapter_.APPROVE, "placeholder": json.dumps( { "toolCallId": "toolu_1", @@ -423,8 +423,8 @@ def test_history_follows_the_active_branch_only() -> None: text = "\n".join(json.dumps(line) for line in lines) stored = [ msg - for entry in pi._active_branch(text) - for msg in pi._stored_to_messages(entry) + for entry in adapter_._active_branch(text) + for msg in adapter_._stored_to_messages(entry) ] assert [msg.role for msg in stored] == ["user", "assistant", "tool"] result = stored[2].parts[0] @@ -501,7 +501,7 @@ async def test_a_seeded_pi_conversation_stays_on_its_model( } seed = [ m.Message(role="user", parts=[m.TextPart(text="hello")]), - *pi._stored_to_messages(stored), + *adapter_._stored_to_messages(stored), ] session_id = await adapter._stage(seed) (written,) = agent_dir.glob(f"sessions/*/*_{session_id}.jsonl") @@ -525,7 +525,7 @@ async def test_a_seeded_pi_conversation_stays_on_its_model( def test_options_reach_pi_in_its_own_spelling() -> None: - adapter = pi.PiAdapter( + adapter = adapter_.PiAdapter( cwd="/w", env={"K": "v"}, executable="pi", @@ -541,14 +541,14 @@ def test_options_reach_pi_in_its_own_spelling() -> None: ) argv, env, minted = adapter.tui_command(None) assert argv[:3] == ["pi", "--session-id", minted] - assert argv[3:5] == ["-e", pi.EXTENSION] + assert argv[3:5] == ["-e", adapter_.EXTENSION] assert argv[argv.index("--model") + 1] == "anthropic/claude-sonnet-4.6" assert argv[argv.index("--thinking") + 1] == "high" assert "--no-tools" in argv assert argv[argv.index("--exclude-tools") + 1] == "bash" assert env["K"] == "v" # The person at the TUI is the gate; the bridge is not there. - assert json.loads(env[pi.CONFIG_VAR]) == { + assert json.loads(env[adapter_.CONFIG_VAR]) == { "approve": False, "mcp": {"docs": {"command": "docs-mcp"}}, } diff --git a/tests/workers/_adapters/test_handoff.py b/tests/workers/test_handoff.py similarity index 90% rename from tests/workers/_adapters/test_handoff.py rename to tests/workers/test_handoff.py index 4f55dadd..7dcdf1fd 100644 --- a/tests/workers/_adapters/test_handoff.py +++ b/tests/workers/test_handoff.py @@ -19,7 +19,10 @@ ToolCallPart, ToolResultPart, ) -from ai.workers._adapters import _handoff +from ai.workers import _handoff +from ai.workers.claude_code import handoff as claude_handoff +from ai.workers.codex import handoff as codex_handoff +from ai.workers.pi import handoff as pi_handoff HISTORY = [ Message(role="user", parts=[TextPart(text="Create recipe.txt")]), @@ -76,7 +79,7 @@ def test_codex_items_follow_the_rule() -> None: - items = _handoff.to_codex_items(HISTORY) + items = codex_handoff.to_codex_items(HISTORY) kinds = [(i["type"], i.get("role")) for i in items] assert kinds == [ ("message", "user"), @@ -94,18 +97,18 @@ def test_codex_items_follow_the_rule() -> None: assert out["call_id"] == "call_1" assert out["output"] == json.dumps({"ok": True}) note = items[3]["content"][0]["text"] - assert note.startswith(_handoff.NOTE_PREFIX) + assert note.startswith(codex_handoff.NOTE_PREFIX) assert "Write is the tool" in note def test_the_model_view_seeds_the_same_items() -> None: - assert _handoff.to_codex_items(MODEL_VIEW) == _handoff.to_codex_items( - HISTORY - ) + assert codex_handoff.to_codex_items( + MODEL_VIEW + ) == codex_handoff.to_codex_items(HISTORY) def test_claude_records_follow_the_rule() -> None: - recs = _handoff.to_claude_records(MODEL_VIEW, "sid-1", "/w") + recs = claude_handoff.to_claude_records(MODEL_VIEW, "sid-1", "/w") # user prompt, assistant (text + tool_use), tool result, assistant; NO # reasoning note: Claude's safeguards refuse replayed reasoning. assert [r["type"] for r in recs] == [ @@ -133,7 +136,7 @@ def test_claude_records_follow_the_rule() -> None: def test_pi_merges_assistant_messages_and_keeps_tool_results() -> None: - rendered = _handoff.to_pi_messages( + rendered = pi_handoff.to_pi_messages( [ Message(role="user", parts=[TextPart(text="hi")]), Message(role="assistant", parts=[TextPart(text="one")]), @@ -216,11 +219,11 @@ def test_a_stored_result_without_a_name_takes_its_calls() -> None: def test_reasoning_is_never_the_agents_own_thought() -> None: """No thinking block, no reasoning item: both were measured to fail, one rejected, one silently ignored. It travels as a note, or not at all.""" - items = _handoff.to_codex_items(HISTORY) + items = codex_handoff.to_codex_items(HISTORY) assert not any(i["type"] == "reasoning" for i in items) blocks = [ b - for r in _handoff.to_claude_records(HISTORY, "s", "/w") + for r in claude_handoff.to_claude_records(HISTORY, "s", "/w") if isinstance(r["message"]["content"], list) for b in r["message"]["content"] ] @@ -229,6 +232,6 @@ def test_reasoning_is_never_the_agents_own_thought() -> None: def test_empty_parts_produce_nothing() -> None: blank = [Message(role="assistant", parts=[TextPart(text=" ")])] - assert _handoff.to_codex_items(blank) == [] + assert codex_handoff.to_codex_items(blank) == [] empty = [Message(role="assistant", parts=[ReasoningPart(text="")])] - assert _handoff.to_claude_records(empty, "s", "/w") == [] + assert claude_handoff.to_claude_records(empty, "s", "/w") == [] From 4f358b36c86d61f1708ba7b0e8cb4ec4b946b309 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Fri, 6 Nov 2026 14:14:15 -0800 Subject: [PATCH 3/3] Handle CLI tool calls in each adapter, built-in parts everywhere else --- src/ai/workers/_adapter.py | 13 +- src/ai/workers/_handoff.py | 137 --------------------- src/ai/workers/claude_code/adapter.py | 71 +++++++---- src/ai/workers/claude_code/handoff.py | 35 +++++- src/ai/workers/codex/adapter.py | 25 ++-- src/ai/workers/codex/handoff.py | 45 ++++--- src/ai/workers/pi/adapter.py | 60 +++++++--- src/ai/workers/pi/handoff.py | 139 ++++++++++++---------- tests/workers/claude_code/test_adapter.py | 8 +- tests/workers/codex/test_adapter.py | 2 +- tests/workers/pi/test_adapter.py | 15 +-- tests/workers/test_handoff.py | 97 +++++++-------- 12 files changed, 307 insertions(+), 340 deletions(-) diff --git a/src/ai/workers/_adapter.py b/src/ai/workers/_adapter.py index fd34f25b..ef5dce7d 100644 --- a/src/ai/workers/_adapter.py +++ b/src/ai/workers/_adapter.py @@ -1,9 +1,14 @@ """The adapter protocol: what `Worker` needs from each CLI. -An adapter drives ONE conversation with one CLI. It speaks the model view -both ways: `turn()` yields the events of one assistant message, with the -CLI's own tools as `BuiltinTool*` events, and ends with a `StreamEnd`; a -history it reads or is seeded with is a plain `list[Message]`. +An adapter drives ONE conversation with one CLI. Outside it, the CLI's own +tools are built-in tools, and nothing else: `turn()` yields them as +`BuiltinTool*` events, and a history it reads holds them as built-in parts, +each result in the assistant message of its call. A seed holds them the +same way, next to the calls and `tool` messages of the tools an `ai.Agent` +ran. How the CLI stores a tool call is the adapter's business alone. + +`turn()` yields the events of one assistant message and ends with a +`StreamEnd`. """ from __future__ import annotations diff --git a/src/ai/workers/_handoff.py b/src/ai/workers/_handoff.py index 62f0479d..e51f1bc6 100644 --- a/src/ai/workers/_handoff.py +++ b/src/ai/workers/_handoff.py @@ -45,143 +45,6 @@ def result_text(value: Any) -> str: return value if isinstance(value, str) else json.dumps(value) -# -- the two views of a past --------------------------------------------------- -# -# A worker's messages hold two kinds of tool calls. The tools an `ai.Agent` -# runs (through the bridge) are `ToolCallPart`s with their results in `tool` -# messages, like a model's. The tools the CLI runs itself are built-in -# parts, with the result in the same assistant message, like a provider's -# web search. A CLI's store holds both kinds the same way: a call, then its -# result in a message of its own. `harness_view` gives the renderers that -# one shape, so each renderer handles one kind of tool call, not two; -# `model_view` is its reverse, for a history read back from a store. - - -def harness_view(history: list[messages_.Message]) -> list[messages_.Message]: - """Return a past with the CLI's own tools as ordinary tool calls. - - A worker's messages hold the tools a CLI ran itself as - `BuiltinToolCallPart`/`BuiltinToolReturnPart`, the way a provider's web - search is. A CLI's store holds them as tool calls, each with its result - in a `tool` message after it. The renderers below read that shape. - """ - out: list[messages_.Message] = [] - for message in history: - parts: list[messages_.Part] = [] - results: list[messages_.Part] = [] - first = True - for part in message.parts: - if results and not isinstance( - part, messages_.BuiltinToolReturnPart - ): - # What follows a result comes after it: a new message. - if parts: - out.append(_split(message, parts, first=first)) - first = False - out.append(messages_.Message(role="tool", parts=results)) - parts, results = [], [] - if isinstance(part, messages_.BuiltinToolCallPart): - parts.append( - messages_.ToolCallPart( - tool_call_id=part.tool_call_id, - tool_name=part.tool_name, - tool_args=part.tool_args or "{}", - provider_metadata=part.provider_metadata, - ) - ) - elif isinstance(part, messages_.BuiltinToolReturnPart): - results.append( - messages_.ToolResultPart( - tool_call_id=part.tool_call_id, - tool_name=part.tool_name, - result=part.result, - result_kind="error" if part.is_error else "json", - ) - ) - else: - parts.append(part) - if parts: - out.append(_split(message, parts, first=first)) - if results: - out.append(messages_.Message(role="tool", parts=results)) - return out - - -def _split( - message: messages_.Message, parts: list[messages_.Part], *, first: bool -) -> messages_.Message: - """One piece of a message: the first keeps its id, the rest get new ones.""" - if first: - return message.model_copy(update={"parts": parts}) - return messages_.Message( - role=message.role, - parts=parts, - provider_metadata=message.provider_metadata, - ) - - -def model_view(stored: list[messages_.Message]) -> list[messages_.Message]: - """Return a stored conversation the way a worker shows it. - - The reverse of `harness_view`: the CLI's tool calls become built-in - parts, and each result joins the assistant message of its call. - """ - out: list[messages_.Message] = [] - owner: dict[str, int] = {} # tool_call_id -> index in `out` - for message in stored: - if message.role == "assistant": - parts: list[messages_.Part] = [] - for part in message.parts: - if isinstance(part, messages_.ToolCallPart): - owner[part.tool_call_id] = len(out) - parts.append( - messages_.BuiltinToolCallPart( - tool_call_id=part.tool_call_id, - tool_name=part.tool_name, - tool_args=part.tool_args, - provider_metadata=part.provider_metadata, - ) - ) - else: - parts.append(part) - out.append(message.model_copy(update={"parts": parts})) - elif message.role == "tool": - for result in message.tool_results: - index = owner.get(result.tool_call_id) - if index is None: - continue - target = out[index] - out[index] = target.model_copy( - update={ - "parts": [ - *target.parts, - messages_.BuiltinToolReturnPart( - tool_call_id=result.tool_call_id, - tool_name=result.tool_name - or _name_of(target, result.tool_call_id), - result=result.result, - is_error=result.result_kind == "error", - ), - ] - } - ) - else: - out.append(message) - return out - - -def _name_of(message: messages_.Message, tool_call_id: str) -> str: - return next( - ( - p.tool_name - for p in message.parts - if isinstance(p, messages_.BuiltinToolCallPart) - and p.tool_call_id == tool_call_id - ), - "", - ) - - def warn_unsupported_files( messages: list[messages_.Message], kind: str ) -> None: diff --git a/src/ai/workers/claude_code/adapter.py b/src/ai/workers/claude_code/adapter.py index dae66044..d4964c3d 100644 --- a/src/ai/workers/claude_code/adapter.py +++ b/src/ai/workers/claude_code/adapter.py @@ -18,7 +18,7 @@ import uuid import warnings from collections.abc import AsyncIterator -from typing import Any +from typing import Any, Literal import claude_agent_sdk @@ -387,10 +387,7 @@ async def history(self, session_id: str) -> list[messages_.Message]: stored = await asyncio.to_thread( claude_agent_sdk.get_session_messages, session_id, self._cwd ) - out: list[messages_.Message] = [] - for record in stored: - out.extend(_stored_to_messages(record)) - return _handoff.model_view(out) + return _read(stored) async def list_sessions(self) -> list[_adapter.SessionInfo]: infos = await asyncio.to_thread( @@ -574,11 +571,47 @@ def _stream_events( return out -def _stored_to_messages(record: Any) -> list[messages_.Message]: - """One stored transcript record -> messages, shaped like a live turn's. +def _read(records: list[Any]) -> list[messages_.Message]: + """Read stored transcript records as messages, like a live turn's. - Tool results become their own `tool` message; `_handoff.model_view` - then folds them into the model view. + The CLI's tools are built-in parts. On disk a tool result is a block in + the NEXT user record; here it joins the assistant message of its call. + """ + out: list[messages_.Message] = [] + calls: dict[str, tuple[messages_.Message, messages_.BuiltinToolCallPart]] + calls = {} + for record in records: + role, parts, results = _record(record) + if parts and role is not None: + message = messages_.Message(role=role, parts=parts) + out.append(message) + for part in parts: + if isinstance(part, messages_.BuiltinToolCallPart): + calls[part.tool_call_id] = (message, part) + for block in results: + found = calls.get(block.get("tool_use_id", "")) + if found is None: + continue # a result without its call: nothing to join + message, call = found + message.parts.append( + messages_.BuiltinToolReturnPart( + tool_call_id=call.tool_call_id, + tool_name=call.tool_name, + result=block.get("content"), + is_error=bool(block.get("is_error")), + ) + ) + return out + + +def _record( + record: Any, +) -> tuple[ + Literal["user", "assistant"] | None, list[Any], list[dict[str, Any]] +]: + """Return a record's role, its parts, and its tool result blocks. + + The role is None for a record that is neither user nor assistant. """ raw = record.message if isinstance(record.message, dict) else {} role = raw.get("role") or record.type @@ -588,7 +621,7 @@ def _stored_to_messages(record: Any) -> list[messages_.Message]: else: blocks = content if isinstance(content, list) else [] parts: list[Any] = [] - tool_results: list[Any] = [] + results: list[dict[str, Any]] = [] for block in blocks: if not isinstance(block, dict): continue @@ -612,24 +645,12 @@ def _stored_to_messages(record: Any) -> list[messages_.Message]: parts.append(messages_.ReasoningPart(text=block["thinking"])) elif kind == "tool_use": parts.append( - messages_.ToolCallPart( + messages_.BuiltinToolCallPart( tool_call_id=block.get("id", ""), tool_name=block.get("name", ""), tool_args=json.dumps(block.get("input") or {}), ) ) elif kind == "tool_result": - tool_results.append( - messages_.ToolResultPart( - tool_call_id=block.get("tool_use_id", ""), - tool_name="", - result=block.get("content"), - result_kind="error" if block.get("is_error") else "json", - ) - ) - out: list[messages_.Message] = [] - if parts and role in ("user", "assistant"): - out.append(messages_.Message(role=role, parts=parts)) - if tool_results: - out.append(messages_.Message(role="tool", parts=tool_results)) - return out + results.append(block) + return (role if role in ("user", "assistant") else None), parts, results diff --git a/src/ai/workers/claude_code/handoff.py b/src/ai/workers/claude_code/handoff.py index 91245a3b..15800dc4 100644 --- a/src/ai/workers/claude_code/handoff.py +++ b/src/ai/workers/claude_code/handoff.py @@ -61,9 +61,30 @@ def add(role: str, content: Any) -> None: ) parent = uid - for m in _handoff.harness_view(messages): + for m in messages: + role = "user" if m.role in ("user", "tool") else "assistant" blocks: list[dict[str, Any]] = [] + results: list[dict[str, Any]] = [] for p in m.parts: + if isinstance(p, messages_.BuiltinToolReturnPart): + # A tool the CLI ran: the worker holds its result in the + # assistant message, the store in the NEXT user record. So + # the assistant record ends here, and the results follow. + if blocks: + add(role, blocks) + blocks = [] + results.append( + { + "type": "tool_result", + "tool_use_id": p.tool_call_id, + "content": _handoff.result_text(p.result), + **({"is_error": True} if p.is_error else {}), + } + ) + continue + if results: + add("user", results) + results = [] if isinstance(p, messages_.TextPart) and p.text.strip(): blocks.append({"type": "text", "text": p.text}) elif isinstance(p, messages_.FilePart) and m.role == "user": @@ -71,7 +92,9 @@ def add(role: str, content: Any) -> None: blocks.append({"type": "image", "source": _source(p)}) elif p.media_type == "application/pdf": blocks.append({"type": "document", "source": _source(p)}) - elif isinstance(p, messages_.ToolCallPart): + elif isinstance( + p, messages_.ToolCallPart | messages_.BuiltinToolCallPart + ): try: args = json.loads(p.tool_args) if p.tool_args else {} except json.JSONDecodeError: @@ -85,6 +108,8 @@ def add(role: str, content: Any) -> None: } ) elif isinstance(p, messages_.ToolResultPart): + # A tool an `ai.Agent` ran: its result is in a `tool` + # message, which becomes a user record of its own. blocks.append( { "type": "tool_result", @@ -100,7 +125,7 @@ def add(role: str, content: Any) -> None: # ReasoningPart: omitted on purpose — Anthropic's safeguards # refuse replayed reasoning. The adapter warns with the count. if blocks: - # Tool results live in role="tool" messages in the IR; on disk - # the CLI keeps them as user records with tool_result blocks. - add("user" if m.role in ("user", "tool") else "assistant", blocks) + add(role, blocks) + if results: + add("user", results) return records diff --git a/src/ai/workers/codex/adapter.py b/src/ai/workers/codex/adapter.py index 68bcb4b2..47a176a1 100644 --- a/src/ai/workers/codex/adapter.py +++ b/src/ai/workers/codex/adapter.py @@ -448,11 +448,11 @@ async def history(self, session_id: str) -> list[messages_.Message]: "thread/read", {"threadId": session_id, "includeTurns": True, "itemsView": "full"}, ) - stored: list[messages_.Message] = [] + out: list[messages_.Message] = [] for turn in ((result or {}).get("thread") or {}).get("turns") or []: for item in turn.get("items") or []: - stored.extend(_item_to_messages(item)) - return _handoff.model_view(stored) + out.extend(_item_to_messages(item)) + return out async def list_sessions(self) -> list[_adapter.SessionInfo]: # Asked for THIS directory, page by page: one unfiltered page of the @@ -585,7 +585,11 @@ def _reasoning_text(item: dict[str, Any]) -> str: def _item_to_messages(item: dict[str, Any]) -> list[messages_.Message]: - """One stored thread item -> messages, CLI tools as tool calls.""" + """One stored thread item -> messages, shaped like a live turn's. + + A tool item holds the call and its result: one assistant message with + both, as built-in parts. + """ kind = item.get("type") if kind == "userMessage": # A user item carries `content` blocks, its `UserInput`s. @@ -619,21 +623,18 @@ def _item_to_messages(item: dict[str, Any]) -> list[messages_.Message]: ) ] if kind in TOOL_ITEMS: - call = messages_.ToolCallPart( + call = messages_.BuiltinToolCallPart( tool_call_id=item.get("id", ""), tool_name=_tool_name(item), tool_args=json.dumps(_tool_args(item)), ) - result = messages_.ToolResultPart( - tool_call_id=item.get("id", ""), + result = messages_.BuiltinToolReturnPart( + tool_call_id=call.tool_call_id, tool_name=call.tool_name, result=_item_result(item), - result_kind="error" if item.get("status") == "failed" else "json", + is_error=item.get("status") == "failed", ) - return [ - messages_.Message(role="assistant", parts=[call]), - messages_.Message(role="tool", parts=[result]), - ] + return [messages_.Message(role="assistant", parts=[call, result])] return [] diff --git a/src/ai/workers/codex/handoff.py b/src/ai/workers/codex/handoff.py index ea637d78..796f9937 100644 --- a/src/ai/workers/codex/handoff.py +++ b/src/ai/workers/codex/handoff.py @@ -14,10 +14,14 @@ NOTE_PREFIX = "[Handoff from a previous session — earlier reasoning: " -def _note(reasoning: list[str]) -> str: - return ( - NOTE_PREFIX + " ".join(r.strip() for r in reasoning if r.strip()) + "]" - ) +def _note(reasoning: list[str]) -> dict[str, Any]: + """Return the reasoning of a turn as a user-role note.""" + text = " ".join(r.strip() for r in reasoning if r.strip()) + return { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": NOTE_PREFIX + text + "]"}], + } def _image_url(part: messages_.FilePart) -> str: @@ -26,9 +30,17 @@ def _image_url(part: messages_.FilePart) -> str: def to_codex_items(messages: list[messages_.Message]) -> list[dict[str, Any]]: - """Responses API items for `thread/inject_items`.""" + """Responses API items for `thread/inject_items`. + + The items are a flat list, so each message goes in as it is: a call and + its result are two items that follow each other. That holds for a tool + an `ai.Agent` ran (a `ToolCallPart`, its result in a `tool` message) and + for a tool the CLI ran (built-in parts, the result in the same message). + A turn's reasoning goes in as a note after its calls, before their + results. + """ items: list[dict[str, Any]] = [] - for m in _handoff.harness_view(messages): + for m in messages: reasoning: list[str] = [] for p in m.parts: if isinstance(p, messages_.TextPart) and p.text.strip(): @@ -54,7 +66,9 @@ def to_codex_items(messages: list[messages_.Message]) -> list[dict[str, Any]]: ], } ) - elif isinstance(p, messages_.ToolCallPart): + elif isinstance( + p, messages_.ToolCallPart | messages_.BuiltinToolCallPart + ): items.append( { "type": "function_call", @@ -63,7 +77,12 @@ def to_codex_items(messages: list[messages_.Message]) -> list[dict[str, Any]]: "arguments": p.tool_args or "{}", } ) - elif isinstance(p, messages_.ToolResultPart): + elif isinstance( + p, messages_.ToolResultPart | messages_.BuiltinToolReturnPart + ): + if reasoning: + items.append(_note(reasoning)) + reasoning = [] items.append( { "type": "function_call_output", @@ -74,13 +93,5 @@ def to_codex_items(messages: list[messages_.Message]) -> list[dict[str, Any]]: elif isinstance(p, messages_.ReasoningPart) and p.text.strip(): reasoning.append(p.text) if reasoning: - items.append( - { - "type": "message", - "role": "user", - "content": [ - {"type": "input_text", "text": _note(reasoning)} - ], - } - ) + items.append(_note(reasoning)) return items diff --git a/src/ai/workers/pi/adapter.py b/src/ai/workers/pi/adapter.py index bd44d73e..96d6d53d 100644 --- a/src/ai/workers/pi/adapter.py +++ b/src/ai/workers/pi/adapter.py @@ -559,10 +559,7 @@ async def history(self, session_id: str) -> list[messages_.Message]: """ found = await self._find(session_id) text = await asyncio.to_thread(pathlib.Path(found["path"]).read_text) - stored: list[messages_.Message] = [] - for record in _active_branch(text): - stored.extend(_stored_to_messages(record)) - return _handoff.model_view(stored) + return _read(_active_branch(text)) async def list_sessions(self) -> list[_adapter.SessionInfo]: return [ @@ -962,7 +959,14 @@ def _assistant_parts(content: Any) -> list[messages_.Part]: elif kind == "thinking" and block.get("thinking"): parts.append(messages_.ReasoningPart(text=block["thinking"])) elif kind == "toolCall": - parts.append(_tool_call(block)) + call = _tool_call(block) + parts.append( + messages_.BuiltinToolCallPart( + tool_call_id=call.tool_call_id, + tool_name=call.tool_name, + tool_args=call.tool_args, + ) + ) return parts @@ -1030,8 +1034,44 @@ def _active_branch(text: str) -> list[dict[str, Any]]: return list(reversed(path)) +def _read(entries: list[dict[str, Any]]) -> list[messages_.Message]: + """Read session entries as messages, like a live turn's. + + The CLI's tools are built-in parts. pi keeps a tool result as a message + of its own (`toolResult`); here it joins the assistant message of its + call. + """ + out: list[messages_.Message] = [] + calls: dict[str, tuple[messages_.Message, messages_.BuiltinToolCallPart]] + calls = {} + for entry in entries: + message = entry.get("message") or {} + if entry.get("type") == "message" and message.get("role") == ( + "toolResult" + ): + found = calls.get(str(message.get("toolCallId") or "")) + if found is None: + continue # a result without its call: nothing to join + owner, call = found + owner.parts.append( + messages_.BuiltinToolReturnPart( + tool_call_id=call.tool_call_id, + tool_name=str(message.get("toolName") or call.tool_name), + result=_tool_result(message), + is_error=bool(message.get("isError")), + ) + ) + continue + for found_message in _stored_to_messages(entry): + out.append(found_message) + for part in found_message.parts: + if isinstance(part, messages_.BuiltinToolCallPart): + calls[part.tool_call_id] = (found_message, part) + return out + + def _stored_to_messages(entry: dict[str, Any]) -> list[messages_.Message]: - """One session entry -> messages, CLI tools as tool calls. + """One session entry -> messages. A tool result is `_read`'s to join. Everything pi sends the model is kept: what an extension injected (`custom_message`) and a `!` command's output (`bashExecution`) reach it @@ -1078,14 +1118,6 @@ def _stored_to_messages(entry: dict[str, Any]) -> list[messages_.Message]: provider_metadata=_written_by(message), ) ] - if role == "toolResult": - result = messages_.ToolResultPart( - tool_call_id=str(message.get("toolCallId") or ""), - tool_name=str(message.get("toolName") or ""), - result=_tool_result(message), - result_kind="error" if message.get("isError") else "json", - ) - return [messages_.Message(role="tool", parts=[result])] return [] diff --git a/src/ai/workers/pi/handoff.py b/src/ai/workers/pi/handoff.py index 5f4d20c6..65b94e39 100644 --- a/src/ai/workers/pi/handoff.py +++ b/src/ai/workers/pi/handoff.py @@ -45,66 +45,24 @@ def to_pi_messages( "total": 0, }, } - for m in _handoff.harness_view(messages): - if m.role == "tool": - for p in m.parts: - if isinstance(p, messages_.ToolResultPart): - out.append( - { - "role": "toolResult", - "toolCallId": p.tool_call_id, - "toolName": p.tool_name, - "content": [ - { - "type": "text", - "text": _handoff.result_text(p.result), - } - ], - "isError": p.result_kind == "error", - "timestamp": timestamp, - } - ) - continue - blocks: list[dict[str, Any]] = [] - for p in m.parts: - if isinstance(p, messages_.TextPart) and p.text.strip(): - blocks.append({"type": "text", "text": p.text}) - elif ( - isinstance(p, messages_.FilePart) - and m.role == "user" - and p.media_type.startswith("image/") - ): - if isinstance(p.data, str) and media.is_downloadable_url( - p.data - ): - blocks.append( - {"type": "text", "text": f"[image: {p.data}]"} - ) - else: - blocks.append( - { - "type": "image", - "data": media.data_to_base64(p.data), - "mimeType": p.media_type, - } - ) - elif ( - isinstance(p, messages_.ToolCallPart) and m.role == "assistant" - ): - try: - args = json.loads(p.tool_args) if p.tool_args else {} - except json.JSONDecodeError: - args = {"raw": p.tool_args} - blocks.append( - { - "type": "toolCall", - "id": p.tool_call_id or p.id, - "name": p.tool_name, - "arguments": args if isinstance(args, dict) else {}, - } - ) + + def result(call_id: str, name: str, value: Any, *, is_error: bool) -> None: + out.append( + { + "role": "toolResult", + "toolCallId": call_id, + "toolName": name, + "content": [ + {"type": "text", "text": _handoff.result_text(value)} + ], + "isError": is_error, + "timestamp": timestamp, + } + ) + + def add(m: messages_.Message, blocks: list[dict[str, Any]]) -> None: if not blocks: - continue + return if m.role == "system": # Measured (pi 1.3.0): a system message in a session file is # pi's own prompt state, and replaying one REPLACED pi's system @@ -116,7 +74,7 @@ def to_pi_messages( out.append( {"role": "user", "content": blocks, "timestamp": timestamp} ) - continue + return if out and out[-1]["role"] == "assistant": out[-1]["content"] += blocks else: @@ -145,4 +103,65 @@ def to_pi_messages( ) if any(b["type"] == "toolCall" for b in out[-1]["content"]): out[-1]["stopReason"] = "toolUse" + + for m in messages: + blocks: list[dict[str, Any]] = [] + for p in m.parts: + if isinstance(p, messages_.ToolResultPart): + # A tool an `ai.Agent` ran: its result is in a `tool` message. + result( + p.tool_call_id, + p.tool_name, + p.result, + is_error=p.result_kind == "error", + ) + elif isinstance(p, messages_.BuiltinToolReturnPart): + # A tool the CLI ran: the worker holds its result in the + # assistant message, pi as a message of its own. So the + # assistant message ends here, and the result follows. + add(m, blocks) + blocks = [] + result( + p.tool_call_id, p.tool_name, p.result, is_error=p.is_error + ) + elif isinstance(p, messages_.TextPart) and p.text.strip(): + blocks.append({"type": "text", "text": p.text}) + elif ( + isinstance(p, messages_.FilePart) + and m.role == "user" + and p.media_type.startswith("image/") + ): + if isinstance(p.data, str) and media.is_downloadable_url( + p.data + ): + blocks.append( + {"type": "text", "text": f"[image: {p.data}]"} + ) + else: + blocks.append( + { + "type": "image", + "data": media.data_to_base64(p.data), + "mimeType": p.media_type, + } + ) + elif ( + isinstance( + p, messages_.ToolCallPart | messages_.BuiltinToolCallPart + ) + and m.role == "assistant" + ): + try: + args = json.loads(p.tool_args) if p.tool_args else {} + except json.JSONDecodeError: + args = {"raw": p.tool_args} + blocks.append( + { + "type": "toolCall", + "id": p.tool_call_id or p.id, + "name": p.tool_name, + "arguments": args if isinstance(args, dict) else {}, + } + ) + add(m, blocks) return out diff --git a/tests/workers/claude_code/test_adapter.py b/tests/workers/claude_code/test_adapter.py index a8d4d801..85788f51 100644 --- a/tests/workers/claude_code/test_adapter.py +++ b/tests/workers/claude_code/test_adapter.py @@ -247,7 +247,8 @@ async def test_a_seed_is_written_where_the_cli_reads(tmp_path: Path) -> None: assert records[2]["message"]["content"][0]["type"] == "tool_result" -def test_a_stored_transcript_reads_as_the_model_view() -> None: +def test_a_stored_result_joins_the_message_of_its_call() -> None: + """On disk the result is in the next user record, without a name.""" records = [ claude_agent_sdk.SessionMessage( type="user", @@ -286,10 +287,7 @@ def test_a_stored_transcript_reads_as_the_model_view() -> None: parent_tool_use_id=None, ), ] - from ai.workers import _handoff - - stored = [msg for r in records for msg in adapter_._stored_to_messages(r)] - history = _handoff.model_view(stored) + history = adapter_._read(records) assert [x.role for x in history] == ["user", "assistant"] call, result = history[1].parts assert isinstance(call, m.BuiltinToolCallPart) diff --git a/tests/workers/codex/test_adapter.py b/tests/workers/codex/test_adapter.py index 16e82303..8f06b630 100644 --- a/tests/workers/codex/test_adapter.py +++ b/tests/workers/codex/test_adapter.py @@ -323,7 +323,7 @@ async def test_a_thread_from_another_directory_is_refused() -> None: ) -async def test_history_is_the_model_view() -> None: +async def test_history_has_the_cli_tools_as_builtin_parts() -> None: server = AppServer() server.threads[THREAD]["turns"] = [ { diff --git a/tests/workers/pi/test_adapter.py b/tests/workers/pi/test_adapter.py index a3ce5e9d..191c013f 100644 --- a/tests/workers/pi/test_adapter.py +++ b/tests/workers/pi/test_adapter.py @@ -421,15 +421,12 @@ def test_history_follows_the_active_branch_only() -> None: {"type": "label", "id": "e", "parentId": "d"}, ] text = "\n".join(json.dumps(line) for line in lines) - stored = [ - msg - for entry in adapter_._active_branch(text) - for msg in adapter_._stored_to_messages(entry) - ] - assert [msg.role for msg in stored] == ["user", "assistant", "tool"] - result = stored[2].parts[0] - assert isinstance(result, m.ToolResultPart) - assert (result.result, result.result_kind) == ("hi", "error") + stored = adapter_._read(adapter_._active_branch(text)) + assert [msg.role for msg in stored] == ["user", "assistant"] + call, result = stored[1].parts # the result joins its call + assert isinstance(call, m.BuiltinToolCallPart) + assert isinstance(result, m.BuiltinToolReturnPart) + assert (result.result, result.is_error) == ("hi", True) async def test_a_seed_is_listed_and_read_back(tmp_path: Path) -> None: diff --git a/tests/workers/test_handoff.py b/tests/workers/test_handoff.py index 7dcdf1fd..8746525d 100644 --- a/tests/workers/test_handoff.py +++ b/tests/workers/test_handoff.py @@ -3,7 +3,9 @@ Text verbatim; tool calls as native items with the ORIGINAL name and arguments; tool results as strings bound to their call; reasoning as a user-role note after the turn it belonged to (codex), or not at all -(claude, pi). The CLI's own tools arrive as built-in parts. +(claude, pi). A worker holds a tool an `ai.Agent` ran as a call and a +`tool` message, and a tool the CLI ran as built-in parts in one message: +each renderer takes both, and both seed the same. """ from __future__ import annotations @@ -24,7 +26,8 @@ from ai.workers.codex import handoff as codex_handoff from ai.workers.pi import handoff as pi_handoff -HISTORY = [ +# Write as a tool an `ai.Agent` ran: its result in a `tool` message. +AGENT_TOOL = [ Message(role="user", parts=[TextPart(text="Create recipe.txt")]), Message( role="assistant", @@ -54,9 +57,9 @@ Message(role="assistant", parts=[TextPart(text="Done.")]), ] -# The same conversation, the way a worker holds it. -MODEL_VIEW = [ - HISTORY[0], +# The same conversation, Write as a tool the CLI ran: built-in parts. +CLI_TOOL = [ + AGENT_TOOL[0], Message( role="assistant", parts=[ @@ -74,12 +77,12 @@ ), ], ), - HISTORY[3], + AGENT_TOOL[3], ] def test_codex_items_follow_the_rule() -> None: - items = codex_handoff.to_codex_items(HISTORY) + items = codex_handoff.to_codex_items(AGENT_TOOL) kinds = [(i["type"], i.get("role")) for i in items] assert kinds == [ ("message", "user"), @@ -101,14 +104,24 @@ def test_codex_items_follow_the_rule() -> None: assert "Write is the tool" in note -def test_the_model_view_seeds_the_same_items() -> None: - assert codex_handoff.to_codex_items( - MODEL_VIEW - ) == codex_handoff.to_codex_items(HISTORY) +def test_both_shapes_seed_the_same() -> None: + assert codex_handoff.to_codex_items(CLI_TOOL) == ( + codex_handoff.to_codex_items(AGENT_TOOL) + ) + assert pi_handoff.to_pi_messages(CLI_TOOL, timestamp=1) == ( + pi_handoff.to_pi_messages(AGENT_TOOL, timestamp=1) + ) + assert [ + (r["type"], r["message"]) + for r in claude_handoff.to_claude_records(CLI_TOOL, "s", "/w") + ] == [ + (r["type"], r["message"]) + for r in claude_handoff.to_claude_records(AGENT_TOOL, "s", "/w") + ] def test_claude_records_follow_the_rule() -> None: - recs = claude_handoff.to_claude_records(MODEL_VIEW, "sid-1", "/w") + recs = claude_handoff.to_claude_records(CLI_TOOL, "sid-1", "/w") # user prompt, assistant (text + tool_use), tool result, assistant; NO # reasoning note: Claude's safeguards refuse replayed reasoning. assert [r["type"] for r in recs] == [ @@ -125,7 +138,7 @@ def test_claude_records_follow_the_rule() -> None: result = recs[2]["message"]["content"][0] assert result["type"] == "tool_result" assert result["tool_use_id"] == "call_1" - assert _handoff.reasoning_count(MODEL_VIEW) == 1 # what is warned about + assert _handoff.reasoning_count(CLI_TOOL) == 1 # what is warned about # a real chain, in the CLI's own store shape assert recs[0]["parentUuid"] is None assert all( @@ -164,19 +177,9 @@ def test_pi_merges_assistant_messages_and_keeps_tool_results() -> None: assert rendered[2]["isError"] is True -def test_model_view_and_harness_view_are_inverse() -> None: - model = _handoff.model_view(HISTORY) - assert [m.role for m in model] == ["user", "assistant", "assistant"] - parts = model[1].parts - assert isinstance(parts[2], BuiltinToolCallPart) - assert isinstance(parts[3], BuiltinToolReturnPart) - assert parts[3].result == {"ok": True} - back = _handoff.harness_view(model) - assert [m.role for m in back] == ["user", "assistant", "tool", "assistant"] - assert isinstance(back[2].parts[0], ToolResultPart) - - def test_text_after_a_cli_tool_stays_after_its_result() -> None: + """Claude and pi keep a result in a message of its own: the message + with the call ends there, and what follows comes after the result.""" one = Message( role="assistant", parts=[ @@ -187,43 +190,35 @@ def test_text_after_a_cli_tool_stays_after_its_result() -> None: TextPart(text="It says x."), ], ) - first, result, after = _handoff.harness_view([one]) - assert (first.role, result.role, after.role) == ( + recs = claude_handoff.to_claude_records([one], "s", "/w") + assert [r["type"] for r in recs] == ["assistant", "user", "assistant"] + assert recs[1]["message"]["content"][0]["type"] == "tool_result" + assert recs[2]["message"]["content"] == [ + {"type": "text", "text": "It says x."} + ] + rendered = pi_handoff.to_pi_messages([one], timestamp=1) + assert [r["role"] for r in rendered] == [ "assistant", - "tool", + "toolResult", "assistant", - ) - assert first.id == one.id and after.id != one.id - assert after.text == "It says x." - - -def test_a_stored_result_without_a_name_takes_its_calls() -> None: - stored = [ - Message( - role="assistant", - parts=[ - ToolCallPart(tool_call_id="c", tool_name="Read", tool_args="") - ], - ), - Message( - role="tool", - parts=[ToolResultPart(tool_call_id="c", tool_name="", result="x")], - ), ] - (message,) = _handoff.model_view(stored) - result = message.parts[1] - assert isinstance(result, BuiltinToolReturnPart) - assert result.tool_name == "Read" + assert rendered[2]["content"] == [{"type": "text", "text": "It says x."}] + items = codex_handoff.to_codex_items([one]) + assert [i["type"] for i in items] == [ + "function_call", + "function_call_output", + "message", + ] def test_reasoning_is_never_the_agents_own_thought() -> None: """No thinking block, no reasoning item: both were measured to fail, one rejected, one silently ignored. It travels as a note, or not at all.""" - items = codex_handoff.to_codex_items(HISTORY) + items = codex_handoff.to_codex_items(AGENT_TOOL) assert not any(i["type"] == "reasoning" for i in items) blocks = [ b - for r in claude_handoff.to_claude_records(HISTORY, "s", "/w") + for r in claude_handoff.to_claude_records(AGENT_TOOL, "s", "/w") if isinstance(r["message"]["content"], list) for b in r["message"]["content"] ]