diff --git a/examples/workers/__init__.py b/examples/workers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/workers/agent_tools.py b/examples/workers/agent_tools.py new file mode 100644 index 00000000..3885abe7 --- /dev/null +++ b/examples/workers/agent_tools.py @@ -0,0 +1,50 @@ +"""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 [codex] +""" + +import asyncio +import sys +from collections.abc import Callable + +import ai +from ai import workers + +KINDS: dict[str, Callable[..., workers.Worker]] = { + "codex": workers.Codex, +} + + +@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 "codex")) diff --git a/examples/workers/approvals.py b/examples/workers/approvals.py new file mode 100644 index 00000000..348a47b2 --- /dev/null +++ b/examples/workers/approvals.py @@ -0,0 +1,42 @@ +"""Approve each of the CLI's own tool calls, as `ai.Agent` tool approvals. + +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] +""" + +import asyncio +import sys +from collections.abc import Callable + +import ai +from ai import workers + +KINDS: dict[str, Callable[..., workers.Worker]] = { + "codex": workers.Codex, +} + + +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() + + +asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "codex")) diff --git a/examples/workers/streaming.py b/examples/workers/streaming.py new file mode 100644 index 00000000..a5239cd3 --- /dev/null +++ b/examples/workers/streaming.py @@ -0,0 +1,52 @@ +"""A worker where a model goes: stream it, then continue the conversation. + +uv run python examples/workers/streaming.py [codex] +""" + +import asyncio +import sys +from collections.abc import Callable + +import ai +from ai import workers + +KINDS: dict[str, Callable[..., workers.Worker]] = { + "codex": workers.Codex, +} + + +async def main(kind: str) -> None: + async with KINDS[kind]() as worker: + messages = [ + ai.user_message( + "List the files here with a shell command, then say DONE." + ) + ] + async with ai.stream(worker, messages) as s: + async for event in s: + match event: + case ai.events.TextDelta(): + print(event.chunk, end="", flush=True) + case ai.events.BuiltinToolEnd(): # the CLI's own tool + call = event.tool_call + print(f"\n[{call.tool_name}] {call.tool_args[:60]}") + print(f"\nusage: {s.message.usage}") + + # The worker keeps the conversation: send it back with what is new. + follow_up = [ + *worker.messages, + ai.user_message("How many files? Reply with a number only."), + ] + print(await ai.experimental_generate(worker, follow_up)) + + # A worker is data: save it, and continue the conversation later. + saved = worker.model_dump_json() + again = workers.Worker.model_validate_json(saved) + async with again: + reply = await ai.experimental_generate( + again, [*again.messages, ai.user_message("And the largest one?")] + ) + print(reply.text) + + +asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else "codex")) diff --git a/pyproject.toml b/pyproject.toml index 6bfd5968..2bbe0863 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,6 +176,11 @@ extend-ignore = [ "TC003", # annoying and wrong if TYPE_CHECKING for stdlib imports ] +[tool.ruff.lint.flake8-type-checking] +# pydantic evaluates field annotations at runtime, so their imports must stay +# real even under `from __future__ import annotations` +runtime-evaluated-base-classes = ["pydantic.BaseModel"] + [tool.ruff.lint.per-file-ignores] "examples/**/*.py" = [ "D", # examples are executable samples, not public API diff --git a/src/ai/__init__.py b/src/ai/__init__.py index 8fada860..c70016ad 100644 --- a/src/ai/__init__.py +++ b/src/ai/__init__.py @@ -8,6 +8,7 @@ testing, ui, util, + workers, ) from .agents import ( Agent, @@ -229,5 +230,6 @@ "ui", "user_message", "util", + "workers", "yield_from", ] diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index df063fd1..0b5933bb 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -1040,11 +1040,16 @@ async def _iterate(self) -> AsyncGenerator[events_.ToolCallResult]: def _validate_model(value: Any) -> models.ModelLike: - # Another model-like is kept as it is; a dict from JSON is an ai.Model. + # Another model-like (an ai.workers.Worker) is kept as it is. A dict + # from JSON with a worker's `kind` is a worker; any other is an ai.Model. if isinstance(value, models.ModelLike) and not isinstance( value, models.Model ): return value + if isinstance(value, dict) and "kind" in value: + from .. import workers # noqa: PLC0415 + + return workers.Worker.model_validate(value) return models.Model.model_validate(value) diff --git a/src/ai/workers/__init__.py b/src/ai/workers/__init__.py new file mode 100644 index 00000000..058a3b95 --- /dev/null +++ b/src/ai/workers/__init__.py @@ -0,0 +1,45 @@ +"""Coding-agent CLIs as workers. + +Experimental: not part of the stable API, may change or be removed. + +A worker goes where a model goes, and it keeps its conversation:: + + import ai + from ai import workers + + async with workers.Codex(model="gpt-5.5") as worker: + async with ai.stream(worker, [ai.user_message("Fix the tests")]) as s: + async for event in s: + ... + + # an ai.Agent's tools reach the CLI through MCP + async with ai.Agent(tools=[search_docs]).run(worker, messages) as run: + ... + +The CLI's own tools come back as `ai.events.BuiltinTool*` events. + +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 conversation the CLI has stored. +""" + +from ._adapters.base import SessionInfo +from ._errors import ( + ExecutableMissingError, + NotAuthenticatedError, + TranscriptMismatchError, + WorkerError, +) +from ._kinds import Codex +from ._worker import Worker + +__all__ = [ + "Codex", + "ExecutableMissingError", + "NotAuthenticatedError", + "SessionInfo", + "TranscriptMismatchError", + "Worker", + "WorkerError", +] diff --git a/src/ai/workers/_adapters/__init__.py b/src/ai/workers/_adapters/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/ai/workers/_adapters/_handoff.py b/src/ai/workers/_adapters/_handoff.py new file mode 100644 index 00000000..873a0b6c --- /dev/null +++ b/src/ai/workers/_adapters/_handoff.py @@ -0,0 +1,504 @@ +"""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 new file mode 100644 index 00000000..10d09016 --- /dev/null +++ b/src/ai/workers/_adapters/base.py @@ -0,0 +1,333 @@ +"""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/_adapters/codex.py b/src/ai/workers/_adapters/codex.py new file mode 100644 index 00000000..84dc8809 --- /dev/null +++ b/src/ai/workers/_adapters/codex.py @@ -0,0 +1,691 @@ +"""Codex, driven through `codex app-server`. + +Chosen over `codex exec --json`, which cannot be controlled: one process per +turn, stdin unused, no approval channel. app-server has `turn/steer`, +`turn/interrupt` and `item/*/requestApproval`. + +The protocol shapes here were learned by running the server and reading what +it sent (see VERIFIED_VERSION), not from the schema alone. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import importlib.metadata +import json +import os +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, jsonrpc + +KIND = "codex" +INSTALL = "npm install -g @openai/codex" +VERIFIED_VERSION = "0.160.1" +"""The CLI the live tests last passed against. Older ones get a warning.""" +CRASHED = "__crashed__" +"""Put on the turn's queue when app-server dies.""" +LIST_PAGE_SIZE = 100 +LIST_PAGES = 50 +"""`list_sessions()` reads at most this many pages of threads.""" + +APPROVAL_METHODS = ( + "item/commandExecution/requestApproval", + "item/fileChange/requestApproval", + "item/permissions/requestApproval", + "execCommandApproval", + "applyPatchApproval", +) + +TOOL_ITEMS = ( + "fileChange", + "commandExecution", + "mcpToolCall", + "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 " + "env={'OPENAI_API_KEY': ...}, or point it at another provider with " + "config={'model_provider': ..., 'model_providers..base_url': ...}." +) + + +class CodexAdapter: + def __init__( + self, + *, + cwd: str, + env: dict[str, str], + executable: str, + model: str | None, + mcp_servers: dict[str, Any], + approve: base.Approve | None, + sandbox: str | None, + config: dict[str, Any], + ) -> None: + # Codex records a thread's directory as its real path (on macOS a + # temp dir is /var/..., recorded as /private/var/..., measured). + self._cwd = os.path.realpath(cwd) + self._env = env + self._executable = executable + self._model = model + self._approve = approve + # Full access by default, like claude without a gate: a worker is + # driven by a program, and a mode that silently refuses writes reads + # as an agent that did nothing. + self._sandbox = sandbox or "danger-full-access" + # Passed to app-server as `-c key=value`, the overrides the codex + # 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._peer: jsonrpc.JsonRpcPeer | None = None + self._thread: str | None = None + self._turn: str | None = None + self._queue: asyncio.Queue[tuple[str, dict[str, Any]]] | None = None + self._usage = usage_.Usage() + self._interrupted = 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, + ) + self._process = await base.spawn( + [self._executable, "app-server", *_config_args(self._config)], + cwd=self._cwd, + env=self._env, + ) + self._peer = jsonrpc.JsonRpcPeer(self._process, kind=KIND) + self._peer.start( + self._on_notify, self._on_request, on_crash=self._on_crash + ) + try: + await self._peer.call( + "initialize", + { + "clientInfo": { + "name": "ai-python", + "title": "AI SDK for Python", + "version": importlib.metadata.version("ai"), + } + }, + timeout=60, + ) + await self._peer.notify("initialized") + except BaseException: + # A failed handshake must not leave the server running. + await self.close() + raise + + async def close(self) -> None: + if self._peer is not None: + await self._peer.close() + self._peer = None + if self._process is not None: + with contextlib.suppress(Exception): + await self._process.terminate() + self._process = None + + def _rpc(self) -> jsonrpc.JsonRpcPeer: + if self._peer is None: + raise errors.WorkerError(f"{KIND}: app-server is not running") + return self._peer + + def _thread_params(self) -> dict[str, Any]: + """How the thread runs: on start, resume and fork alike. + + A resumed or forked thread otherwise keeps the policy it was created + with, like a person's TUI asking for approvals nobody here answers. + """ + return { + "sandbox": self._sandbox, + # With a gate: ask about everything the policy can ask about. + # Without one the server never asks: a request nobody answers + # would freeze the agent (measured). + "approvalPolicy": "untrusted" if self._approve else "never", + **({"model": self._model} if self._model else {}), + **( + {"developerInstructions": self._instructions} + if self._instructions + else {} + ), + } + + # -- 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 session_id is not None: + await self._require_here(session_id) + method = "thread/fork" if fork else "thread/resume" + try: + result = await self._rpc().call( + method, {"threadId": session_id, **self._thread_params()} + ) + except jsonrpc.JsonRpcError as exc: + raise errors.WorkerError( + f"{KIND}: cannot open thread {session_id}: {exc}" + ) from exc + thread = str((result or {}).get("thread", {}).get("id") or "") + self._thread = thread or session_id + return self._thread + result = await self._rpc().call( + "thread/start", {"cwd": self._cwd, **self._thread_params()} + ) + thread = str((result or {}).get("thread", {}).get("id") or "") + if not thread: + raise errors.WorkerError( + f"{KIND}: thread/start gave no id: {result}" + ) + self._thread = thread + if seed: + # The thread's model-visible history, before the first turn. + _handoff.warn_unsupported_files(seed, KIND) + await self._rpc().call( + "thread/inject_items", + {"threadId": thread, "items": _handoff.to_codex_items(seed)}, + ) + return thread + + async def _require_here(self, session_id: str) -> dict[str, Any]: + """Refuse a thread that belongs to another directory. + + Codex's thread store is the whole machine's, so resuming a thread + from another project succeeds natively, and the agent would then + carry a transcript about files that are not here. + """ + try: + result = await self._rpc().call( + "thread/read", {"threadId": session_id} + ) + except jsonrpc.JsonRpcError as exc: + raise errors.WorkerError( + f"{KIND}: cannot read thread {session_id}: {exc}" + ) from exc + thread = (result or {}).get("thread") or {} + if not thread: + raise errors.WorkerError(f"{KIND}: no thread {session_id}") + cwd = thread.get("cwd") + if cwd is not None and cwd != self._cwd: + raise errors.WorkerError( + f"{KIND}: thread {session_id} belongs to {cwd}, not {self._cwd}" + ) + return thread + + # -- notifications and approvals ------------------------------------------ + + async def _on_notify(self, method: str, params: dict[str, Any]) -> None: + thread = params.get("threadId") or (params.get("thread") or {}).get( + "id" + ) + if thread != self._thread: + return + if method == "thread/tokenUsage/updated": + self._usage = _usage(params.get("tokenUsage") or {}) + if self._queue is not None: + self._queue.put_nowait((method, params)) + + def _on_crash(self, exc: BaseException) -> None: + if self._queue is not None: + self._queue.put_nowait((CRASHED, {"detail": str(exc)})) + self._peer = None + + async def _on_request( + self, request_id: Any, method: str, params: dict[str, Any] + ) -> None: + if method not in APPROVAL_METHODS or self._approve is None: + # Unknown, or unasked for: answer rather than hang the agent. + decision = "accept" if method in APPROVAL_METHODS else "decline" + await self._rpc().respond(request_id, {"decision": decision}) + return + granted, _ = await self._approve( + str(params.get("itemId") or params.get("callId") or ""), + _approval_tool_name(method), + { + k: v + for k, v in params.items() + if k not in ("threadId", "turnId", "itemId") and v is not None + }, + ) + # The protocol has no field for a reason: codex only hears "declined". + await self._rpc().respond( + request_id, {"decision": "accept" if granted else "decline"} + ) + + # -- turns ---------------------------------------------------------------- + + async def turn(self, prompt: base.Prompt) -> AsyncIterator[events_.Event]: + base.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() + self._queue = queue + self._interrupted = False + self._usage = usage_.Usage() + open_text: set[str] = set() + open_reasoning: set[str] = set() + started: set[str] = set() + # The last transport error codex reported while retrying: if the turn + # then fails without a message of its own, the only explanation. + last_error = "" + try: + try: + begun = await self._rpc().call( + "turn/start", + { + "threadId": self._thread, + "input": await _input_items(prompt), + }, + ) + except jsonrpc.JsonRpcError as exc: + raise errors.turn_failed(KIND, str(exc), AUTH_HINT) from exc + # The turn id comes from the RESPONSE: this server sends no + # `turn/started`, and steer and interrupt need the id. + turn = ((begun or {}).get("turn") or {}).get("id") + self._turn = turn + yield events_.StreamStart() + while True: + method, params = await queue.get() + if method == CRASHED: + raise errors.WorkerError( + f"{KIND} crashed: {params.get('detail')}" + ) + # An interrupted turn's notifications can still arrive when + # the next one starts. Without this the new turn settled on + # the OLD turn's completion and came back empty. + other = params.get("turnId") or (params.get("turn") or {}).get( + "id" + ) + if turn and other and other != turn: + continue + if method == "item/agentMessage/delta": + item_id = params.get("itemId", "") + if item_id not in open_text: + open_text.add(item_id) + yield events_.TextStart(block_id=item_id) + yield events_.TextDelta( + chunk=params.get("delta", ""), block_id=item_id + ) + elif method in ( + "item/reasoning/textDelta", + "item/reasoning/summaryTextDelta", + ): + item_id = params.get("itemId", "") + if item_id not in open_reasoning: + open_reasoning.add(item_id) + yield events_.ReasoningStart(block_id=item_id) + yield events_.ReasoningDelta( + chunk=params.get("delta", ""), block_id=item_id + ) + elif method == "item/started": + item = params.get("item") or {} + if item.get("type") in TOOL_ITEMS: + started.add(item.get("id", "")) + yield events_.BuiltinToolStart( + tool_call_id=item.get("id", ""), + tool_name=_tool_name(item), + ) + elif method == "item/completed": + item = params.get("item") or {} + item_id = item.get("id", "") + kind = item.get("type") + if kind == "agentMessage" and item_id in open_text: + open_text.discard(item_id) + yield events_.TextEnd(block_id=item_id) + elif kind == "agentMessage" and item.get("text"): + # A message that was never streamed. + yield events_.TextStart(block_id=item_id) + yield events_.TextDelta( + chunk=item["text"], block_id=item_id + ) + yield events_.TextEnd(block_id=item_id) + elif kind == "reasoning" and item_id in open_reasoning: + open_reasoning.discard(item_id) + yield events_.ReasoningEnd(block_id=item_id) + elif kind in TOOL_ITEMS: + for event in base.builtin_call( + item_id, + _tool_name(item), + json.dumps(_tool_args(item)), + started=started, + ): + yield event + yield base.builtin_result( + item_id, + _tool_name(item), + _item_result(item), + is_error=item.get("status") == "failed", + ) + elif method == "turn/completed": + completed = params.get("turn") or {} + # A failed turn says so. Reported as a normal stop with + # empty text, "401 Unauthorized" became silence. + if ( + completed.get("status") == "failed" + and not self._interrupted + ): + raise errors.turn_failed( + KIND, + _error_message(completed) or last_error, + AUTH_HINT, + ) + yield events_.StreamEnd( + usage=self._usage, + finish_reason=( + "cancelled" if self._interrupted else "stop" + ), + ) + return + elif method == "turn/failed": + raise errors.turn_failed( + KIND, _error_message(params), AUTH_HINT + ) + elif method == "error": + last_error = _error_message(params) or last_error + finally: + self._queue = None + self._turn = None + + async def steer(self, text: str) -> None: + if not self._turn: + raise errors.WorkerError(f"{KIND}: no turn is running to steer") + try: + await self._rpc().call( + "turn/steer", + { + "threadId": self._thread, + # A precondition: steering the wrong turn is worse than + # failing to steer. + "expectedTurnId": self._turn, + "input": [{"type": "text", "text": text}], + }, + ) + except jsonrpc.JsonRpcError as exc: + raise errors.WorkerError(f"{KIND}: steer rejected: {exc}") from exc + + async def interrupt(self) -> None: + self._interrupted = True + if not self._turn or self._peer is None: + return + with contextlib.suppress(Exception): + await self._peer.call( + "turn/interrupt", + {"threadId": self._thread, "turnId": self._turn}, + timeout=30, + ) + + # -- the store ------------------------------------------------------------ + + async def history(self, session_id: str) -> list[messages_.Message]: + """Read the thread's stored record, every item included.""" + await self._require_here(session_id) + result = await self._rpc().call( + "thread/read", + {"threadId": session_id, "includeTurns": True, "itemsView": "full"}, + ) + stored: 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) + + async def list_sessions(self) -> list[base.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]] = [] + cursor: str | None = None + for _ in range(LIST_PAGES): + page = await self._rpc().call( + "thread/list", + { + "cwd": self._cwd, + "limit": LIST_PAGE_SIZE, + **({"cursor": cursor} if cursor else {}), + }, + ) + threads += (page or {}).get("data") or [] + cursor = (page or {}).get("nextCursor") + if not cursor: + break + return [ + base.SessionInfo( + session_id=t.get("id", ""), + title=t.get("name") or t.get("preview") or None, + cwd=t.get("cwd"), + created_at=t.get("createdAt"), + updated_at=t.get("updatedAt"), + ) + for t in threads + if t.get("id") and t.get("cwd") in (None, self._cwd) + ] + + def tui_command( + self, session_id: str | None + ) -> tuple[list[str], dict[str, str], str | None]: + config = { + k: v + for k, v in self._config.items() + if not k.startswith(f"mcp_servers.{_bridge.SERVER_NAME}.") + } + # --disable in_app_updates: measured, the TUI otherwise opens on an + # "Update now / Skip" dialog that the first keystrokes answer. + argv = [ + self._executable, + *(["resume", session_id] if session_id else []), + "--disable", + "in_app_updates", + *_config_args(config), + ] + if self._model: + argv += ["--model", self._model] + if self._instructions: + argv += _config_args({"developer_instructions": self._instructions}) + # None: codex mints its own thread id. + return argv, dict(self._env), None + + +async def _input_items(prompt: base.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. + """ + items: list[dict[str, Any]] = [] + for part in prompt: + 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)}) + else: + encoded = base64.b64encode(await base.data(part)).decode("ascii") + items.append( + { + "type": "image", + "url": f"data:{part.media_type};base64,{encoded}", + } + ) + return items + + +def _tool_name(item: dict[str, Any]) -> str: + # Never the command line: that made every shell call a differently named + # tool. The command is in the arguments (`_tool_args`). + return ( + item.get("toolName") or item.get("server") or item.get("type") or "tool" + ) + + +def _approval_tool_name(method: str) -> str: + """Return the native tool kind an approval is about, never its args.""" + if "commandExecution" in method or "execCommand" in method: + return "commandExecution" + if "fileChange" in method or "applyPatch" in method: + return "fileChange" + return method.rsplit("/", 1)[-1] + + +def _tool_args(item: dict[str, Any]) -> dict[str, Any]: + return { + k: v + for k, v in item.items() + if k in ("command", "changes", "arguments", "query", "path", "cwd") + } + + +def _item_result(item: dict[str, Any]) -> Any: + """Return what a tool item produced, under whichever name codex uses. + + A command's output is `aggregatedOutput` (codex 0.160; reading only + `output` made every command result None), a file change's `changes`, an + MCP call's `result`. + """ + for key in ("aggregatedOutput", "output", "changes", "result"): + if item.get(key) is not None: + return item[key] + return None + + +def _reasoning_text(item: dict[str, Any]) -> str: + """Codex reports reasoning as text, or as a LIST of summary parts.""" + text = item.get("text") + if isinstance(text, str) and text: + return text + summary = item.get("summary") + if isinstance(summary, str): + return summary + if isinstance(summary, list): + return "".join(part for part in summary if isinstance(part, str)) + return "" + + +def _item_to_messages(item: dict[str, Any]) -> list[messages_.Message]: + """One stored thread item -> messages, CLI tools as tool calls.""" + kind = item.get("type") + if kind == "userMessage": + # A user item carries `content` blocks, its `UserInput`s. + text = item.get("text") + files: list[messages_.Part] = [] + if not isinstance(text, str): + blocks = [ + b for b in item.get("content") or [] if isinstance(b, dict) + ] + text = "".join(b.get("text", "") for b in blocks) + for block in blocks: + url = block.get("url") + if block.get("type") == "image" and isinstance(url, str): + files.append(messages_.FilePart.from_url(url)) + return [ + messages_.Message( + role="user", parts=[messages_.TextPart(text=text), *files] + ) + ] + if kind == "agentMessage" and item.get("text"): + return [ + messages_.Message( + role="assistant", parts=[messages_.TextPart(text=item["text"])] + ) + ] + if kind == "reasoning" and _reasoning_text(item): + return [ + messages_.Message( + role="assistant", + parts=[messages_.ReasoningPart(text=_reasoning_text(item))], + ) + ] + if kind in TOOL_ITEMS: + call = messages_.ToolCallPart( + 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", ""), + tool_name=call.tool_name, + result=_item_result(item), + result_kind="error" if item.get("status") == "failed" else "json", + ) + return [ + messages_.Message(role="assistant", parts=[call]), + messages_.Message(role="tool", parts=[result]), + ] + return [] + + +def _usage(raw: dict[str, Any]) -> usage_.Usage: + totals = raw.get("lastTurn") or raw.get("total") or raw + + def count(*names: str) -> int: + for name in names: + value = totals.get(name) + if isinstance(value, int): + return value + return 0 + + return usage_.Usage( + input_tokens=count("inputTokens", "input_tokens"), + output_tokens=count("outputTokens", "output_tokens"), + cache_read_tokens=count("cachedInputTokens", "cached_input_tokens"), + reasoning_tokens=count( + "reasoningOutputTokens", "reasoning_output_tokens" + ), + raw={"codex": raw}, + ) + + +def _error_message(payload: dict[str, Any]) -> str: + error = payload.get("error") + if isinstance(error, dict): + detail = error.get("additionalDetails") or "" + message = str(error.get("message") or "") + return f"{message} {detail}".strip() + return str(error or "") + + +def mcp_config(servers: dict[str, Any]) -> dict[str, Any]: + """MCP servers, in the shape claude takes, as codex `-c` overrides. + + `{"name": {"command": ..., "args": [...], "env": {...}}}` for a local + server, `{"name": {"url": ...}}` for a remote one. + """ + out: dict[str, Any] = {} + for name, server in servers.items(): + prefix = f"mcp_servers.{name}" + for key in ("command", "args", "url"): + if server.get(key) is not None: + out[f"{prefix}.{key}"] = server[key] + for var, value in (server.get("env") or {}).items(): + out[f"{prefix}.env.{var}"] = value + return out + + +def _config_args(config: dict[str, Any]) -> list[str]: + """`-c key=value` pairs, the values as TOML. + + JSON string, number and bool literals are valid TOML. + """ + args: list[str] = [] + for key, value in config.items(): + args += ["-c", f"{key}={json.dumps(value)}"] + return args diff --git a/src/ai/workers/_adapters/jsonrpc.py b/src/ai/workers/_adapters/jsonrpc.py new file mode 100644 index 00000000..14837088 --- /dev/null +++ b/src/ai/workers/_adapters/jsonrpc.py @@ -0,0 +1,210 @@ +"""A small JSON-RPC 2.0 peer over a CLI's stdio. + +Line-delimited JSON, request/response futures, and two callbacks for what +a server can start. Nothing here is protocol-specific: the Codex adapter +supplies the meaning. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any + +from .. import _errors as errors + +if TYPE_CHECKING: + from . import base + +NotifyHandler = Callable[[str, dict[str, Any]], Awaitable[None]] +RequestHandler = Callable[[Any, str, dict[str, Any]], Awaitable[None]] + + +class JsonRpcError(Exception): + def __init__(self, error: dict[str, Any]) -> None: + super().__init__(error.get("message") or json.dumps(error)) + self.code = error.get("code") + self.data = error.get("data") + + +class JsonRpcPeer: + def __init__(self, process: base.Process, *, kind: str) -> None: + self._process = process + self._kind = kind + self._next_id = 0 + self._waiters: dict[int, asyncio.Future[Any]] = {} + self._pump: asyncio.Task[None] | None = None + self._closed = False + self._on_crash: Callable[[BaseException], None] | None = None + # Server requests being answered. Each runs as its own task: an + # answer can take minutes (a person approving a tool call), and the + # reader must keep reading meanwhile. + self._requests: set[asyncio.Task[None]] = set() + self._answered: set[Any] = set() + + def start( + self, + on_notify: NotifyHandler, + on_request: RequestHandler, + on_crash: Callable[[BaseException], None] | None = None, + ) -> None: + # `on_crash` is for listeners. A turn waits on a stream of + # notifications, not on a future: when the process dies, failing + # the futures wakes nobody who was listening (measured: a SIGKILLed + # codex left the turn blocked until the test timed out). + self._on_crash = on_crash + self._pump = asyncio.create_task(self._read_loop(on_notify, on_request)) + + async def close(self) -> None: + self._closed = True + 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._requests): + task.cancel() + for task in list(self._requests): + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + self._requests.clear() + # A close is not a crash: `on_crash` does not fire. + self._fail_pending( + errors.WorkerError(f"{self._kind}: connection closed"), + crashed=False, + ) + + def _fail_pending( + self, exc: BaseException, *, crashed: bool = True + ) -> None: + for future in list(self._waiters.values()): + if not future.done(): + future.set_exception(exc) + self._waiters.clear() + if crashed and self._on_crash is not None: + with contextlib.suppress(Exception): + self._on_crash(exc) + + async def call( + self, method: str, params: Any = None, *, timeout: float = 120 + ) -> Any: + if self._closed: + raise errors.WorkerError(f"{self._kind}: connection is closed") + self._next_id += 1 + request_id = self._next_id + future: asyncio.Future[Any] = asyncio.get_running_loop().create_future() + self._waiters[request_id] = future + await self._send( + { + "jsonrpc": "2.0", + "id": request_id, + "method": method, + **({"params": params} if params is not None else {}), + } + ) + try: + async with asyncio.timeout(timeout): + return await future + finally: + self._waiters.pop(request_id, None) + + async def notify(self, method: str, params: Any = None) -> None: + await self._send( + { + "jsonrpc": "2.0", + "method": method, + **({"params": params} if params is not None else {}), + } + ) + + async def respond(self, request_id: Any, result: Any) -> None: + self._answered.add(request_id) + await self._send({"jsonrpc": "2.0", "id": request_id, "result": result}) + + async def _answer( + self, + on_request: RequestHandler, + request_id: Any, + method: str, + params: dict[str, Any], + ) -> None: + """Run one request handler; never leave the server without a reply.""" + try: + await on_request(request_id, method, params) + except asyncio.CancelledError: + raise + except Exception as exc: + if request_id not in self._answered: + with contextlib.suppress(Exception): + await self._send( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32603, "message": str(exc)}, + } + ) + finally: + self._answered.discard(request_id) + + async def _send(self, message: dict[str, Any]) -> None: + try: + await self._process.write(json.dumps(message) + "\n") + except Exception as exc: + raise errors.WorkerError( + f"{self._kind}: could not write to the CLI: {exc}" + ) from exc + + async def _read_loop( + self, on_notify: NotifyHandler, on_request: RequestHandler + ) -> None: + try: + while True: + line = await self._process.readline() + if not line: + note = await self._process.death_note() + self._fail_pending( + errors.WorkerError(f"{self._kind} crashed: {note}") + ) + return + try: + message = json.loads(line) + except json.JSONDecodeError: + continue # a stray non-JSON line is noise + if not isinstance(message, dict): + continue + if "method" in message and "id" not in message: + # Inline, so notifications are handled in the order sent. + await on_notify( + message["method"], message.get("params") or {} + ) + elif "method" in message: + task = asyncio.create_task( + self._answer( + on_request, + message["id"], + message["method"], + message.get("params") or {}, + ) + ) + self._requests.add(task) + task.add_done_callback(self._requests.discard) + else: + self._resolve(message) + except asyncio.CancelledError: + raise + except Exception as exc: + self._fail_pending(errors.WorkerError(f"{self._kind}: {exc}")) + + def _resolve(self, message: dict[str, Any]) -> None: + msg_id = message.get("id") + future = ( + self._waiters.pop(msg_id, None) if isinstance(msg_id, int) else None + ) + if future is None or future.done(): + return + if "error" in message: + future.set_exception(JsonRpcError(message["error"])) + else: + future.set_result(message.get("result")) diff --git a/src/ai/workers/_bridge.py b/src/ai/workers/_bridge.py new file mode 100644 index 00000000..f32f5a59 --- /dev/null +++ b/src/ai/workers/_bridge.py @@ -0,0 +1,230 @@ +"""The MCP bridge: `ai` tools, served to a coding-agent CLI. + +A worker's CLI runs its own loop and its own tools. To give it the tools an +`ai.Agent` holds, the worker serves them over MCP on localhost, and the CLI +connects to that server like any other. + +The server does not run a tool. It hands each call to `on_call` and waits +for the answer. The worker turns the call into a normal `ToolCallPart`, ends +its stream with `finish_reason="tool_call"`, and `ai.Agent` runs the tool. +The next `stream()` call brings the result, and the worker answers the +waiting request with it. While the request waits, the CLI waits too. + +Only what an MCP client needs is implemented: streamable HTTP, POST only, +plain JSON responses (no SSE), and the methods `initialize`, `ping`, +`tools/list` and `tools/call`. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ..types import tools as tools_ + +SERVER_NAME = "ai_sdk" +"""The MCP server name the CLI sees. Claude Code calls a tool +`mcp__ai_sdk__`.""" + +PROTOCOL_VERSION = "2025-06-18" + +OnCall = Callable[[str, dict[str, Any]], Awaitable[tuple[str, bool]]] +"""`on_call(name, arguments)` returns `(text, is_error)`.""" + + +class Bridge: + def __init__(self, on_call: OnCall) -> None: + self._on_call = on_call + self._tools: dict[str, tools_.Tool] = {} + self._server: asyncio.Server | None = None + self._connections: set[asyncio.Task[None]] = set() + + @property + def url(self) -> str: + if self._server is None: + raise RuntimeError("bridge is not started") + host, port = self._server.sockets[0].getsockname()[:2] + return f"http://{host}:{port}/mcp" + + @property + def names(self) -> set[str]: + return set(self._tools) + + def set_tools(self, tools: list[tools_.Tool]) -> None: + """Serve these tools from now on. + + A CLI reads the list when it connects, so a tool added after that + may stay unseen until the next connection. + """ + self._tools = { + t.name: t for t in tools if t.kind == "function" and t.spec + } + + async def start(self) -> None: + if self._server is None: + self._server = await asyncio.start_server( + self._serve, host="127.0.0.1", port=0 + ) + + async def close(self) -> None: + if self._server is None: + return + self._server.close() + for task in list(self._connections): + task.cancel() + for task in list(self._connections): + with contextlib.suppress(BaseException): + await task + with contextlib.suppress(Exception): + await self._server.wait_closed() + self._server = None + + # -- HTTP ----------------------------------------------------------------- + + async def _serve( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + task = asyncio.current_task() + assert task is not None + self._connections.add(task) + try: + while True: + request = await _read_request(reader) + if request is None: + return + method, body = request + if method != "POST": + # No SSE stream to offer, and a session needs no teardown. + status = 200 if method == "DELETE" else 405 + await _respond(writer, status, None) + continue + reply = await self._dispatch(body) + await _respond(writer, 200 if reply is not None else 202, reply) + except (ConnectionError, asyncio.IncompleteReadError): + return + finally: + self._connections.discard(task) + with contextlib.suppress(Exception): + writer.close() + + # -- JSON-RPC ------------------------------------------------------------- + + async def _dispatch(self, body: bytes) -> Any: + try: + payload = json.loads(body or b"null") + except json.JSONDecodeError: + return _error(None, -32700, "parse error") + if isinstance(payload, list): + replies = [await self._one(item) for item in payload] + kept = [r for r in replies if r is not None] + return kept or None + return await self._one(payload) + + async def _one(self, message: Any) -> Any: + if not isinstance(message, dict) or "method" not in message: + # A response from the client (we send no requests) or junk. + return None + method, rid = message["method"], message.get("id") + if rid is None: + return None # a notification: nothing to answer + params = message.get("params") or {} + if method == "initialize": + return _result( + rid, + { + "protocolVersion": params.get( + "protocolVersion", PROTOCOL_VERSION + ), + "capabilities": {"tools": {"listChanged": False}}, + "serverInfo": {"name": SERVER_NAME, "version": "0"}, + }, + ) + if method == "ping": + return _result(rid, {}) + if method == "tools/list": + return _result( + rid, + { + "tools": [ + { + "name": t.name, + "description": (t.spec and t.spec.description) + or "", + "inputSchema": (t.spec and t.spec.params) + or {"type": "object"}, + } + for t in self._tools.values() + ] + }, + ) + if method == "tools/call": + name = str(params.get("name", "")) + if name not in self._tools: + return _error(rid, -32602, f"unknown tool: {name}") + arguments = params.get("arguments") or {} + text, is_error = await self._on_call(name, dict(arguments)) + return _result( + rid, + { + "content": [{"type": "text", "text": text}], + "isError": is_error, + }, + ) + return _error(rid, -32601, f"method not found: {method}") + + +def _result(rid: Any, result: Any) -> dict[str, Any]: + return {"jsonrpc": "2.0", "id": rid, "result": result} + + +def _error(rid: Any, code: int, message: str) -> dict[str, Any]: + return { + "jsonrpc": "2.0", + "id": rid, + "error": {"code": code, "message": message}, + } + + +async def _read_request( + reader: asyncio.StreamReader, +) -> tuple[str, bytes] | None: + """Read one HTTP/1.1 request. None when the client hung up.""" + try: + head = await reader.readuntil(b"\r\n\r\n") + except asyncio.IncompleteReadError: + return None + lines = head.decode("latin-1").split("\r\n") + method = lines[0].split(" ", 1)[0].upper() + headers = {} + for line in lines[1:]: + if ":" in line: + key, value = line.split(":", 1) + headers[key.strip().lower()] = value.strip() + if headers.get("transfer-encoding", "").lower() == "chunked": + body = b"" + while True: + size = int((await reader.readuntil(b"\r\n")).strip(), 16) + if size == 0: + await reader.readuntil(b"\r\n") + break + body += await reader.readexactly(size) + await reader.readexactly(2) + return method, body + length = int(headers.get("content-length", "0") or 0) + return method, await reader.readexactly(length) if length else b"" + + +async def _respond( + writer: asyncio.StreamWriter, status: int, payload: Any +) -> None: + reason = {200: "OK", 202: "Accepted", 405: "Method Not Allowed"}[status] + body = b"" if payload is None else json.dumps(payload).encode() + head = [f"HTTP/1.1 {status} {reason}", f"Content-Length: {len(body)}"] + if payload is not None: + head.append("Content-Type: application/json") + writer.write(("\r\n".join(head) + "\r\n\r\n").encode() + body) + await writer.drain() diff --git a/src/ai/workers/_errors.py b/src/ai/workers/_errors.py new file mode 100644 index 00000000..21baddee --- /dev/null +++ b/src/ai/workers/_errors.py @@ -0,0 +1,67 @@ +"""The errors a worker raises. + +Three cases have their own class, because a caller fixes them differently. +Everything else is a `WorkerError` with a message that says what happened. +""" + +from __future__ import annotations + +from .. import errors + + +class WorkerError(errors.AIError): + """A worker or its CLI failed.""" + + +class ExecutableMissingError(WorkerError): + """The CLI is not installed, or `executable=` does not run.""" + + def __init__(self, kind: str, executable: str, install: str) -> None: + super().__init__( + f"{kind} needs {executable!r}, which does not run here. " + f"Install it with: {install}" + ) + self.kind = kind + self.executable = executable + + +class NotAuthenticatedError(WorkerError): + """The CLI has no usable credentials.""" + + def __init__(self, kind: str, detail: str, hint: str) -> None: + super().__init__(f"{kind} is not authenticated: {detail}\n{hint}") + self.kind = kind + self.detail = detail + + +class TranscriptMismatchError(WorkerError): + """The messages given to `stream()` do not continue this conversation. + + A worker holds a conversation in its CLI. It can continue it, or begin a + new one from a history, but it cannot rewrite it. Use a new worker. + """ + + +# Phrases that mean "no usable credentials", collected from what the CLIs +# print. Narrow on purpose: a 403 is authorization, not authentication. +_AUTH_MARKERS = ( + "401", + "unauthorized", + "missing bearer", + "not logged in", + "invalid_api_key", + "invalid api key", + "authentication_error", + "failed to authenticate", + "no api key", + "please run /login", + "please run `codex login`", +) + + +def turn_failed(kind: str, detail: str, hint: str) -> WorkerError: + """Name a failed turn for what the caller has to fix.""" + detail = (detail or "unknown error")[:500] + if any(marker in detail.lower() for marker in _AUTH_MARKERS): + return NotAuthenticatedError(kind, detail, hint) + return WorkerError(f"{kind} turn failed: {detail}") diff --git a/src/ai/workers/_kinds.py b/src/ai/workers/_kinds.py new file mode 100644 index 00000000..69285dc0 --- /dev/null +++ b/src/ai/workers/_kinds.py @@ -0,0 +1,57 @@ +"""The workers, one per CLI. + +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 . import _bridge, _worker +from ._adapters import base, codex + + +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, + ) diff --git a/src/ai/workers/_worker.py b/src/ai/workers/_worker.py new file mode 100644 index 00000000..89aa3994 --- /dev/null +++ b/src/ai/workers/_worker.py @@ -0,0 +1,805 @@ +"""`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 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. + +It follows the model protocol: `stream(messages, tools=, output_type=)` +yields the events of one assistant message, so `ai.stream`, +`ai.experimental_generate` and `ai.Agent.run` take a worker where they take +a model. The CLI runs its own loop and its own tools; those come back as +`BuiltinTool*` events, so an `ai.Agent` does not run them again. The tools +of an `ai.Agent` reach the CLI through the MCP bridge (see `_bridge`) and +come back as normal tool calls, for the agent to run. + +A worker has state, a model does not. `stream()` still receives the whole +conversation on every call. The worker compares it with `messages` and +sends the CLI only what is new (see `_plan`). +""" + +from __future__ import annotations + +import asyncio +import contextlib +import contextvars +import dataclasses +import json +import os +from collections.abc import AsyncGenerator, Sequence +from typing import Any, ClassVar, Literal, Self + +import pydantic +from pydantic import TypeAdapter + +from ..agents import hooks +from ..agents import runtime as runtime_ +from ..types import events as events_ +from ..types import messages as messages_ +from ..types import tools as tools_ +from ..types import usage as usage_ +from . import _bridge +from . import _errors as errors +from ._adapters import base + +Effort = Literal["low", "medium", "high", "max"] + +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.""" + +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, " + "no prose, no code fences:\n" +) + + +@dataclasses.dataclass +class _Call: + """A call from the CLI to an `ai` tool, through the bridge.""" + + tool_call_id: str + name: str + arguments: dict[str, Any] + answer: asyncio.Future[tuple[str, bool]] + + +@dataclasses.dataclass +class _Failed: + error: BaseException + + +@dataclasses.dataclass +class _Plan: + """What one `stream()` call must do, from the messages it was given.""" + + seed: list[messages_.Message] = dataclasses.field(default_factory=list) + instructions: str | None = None + results: list[messages_.ToolResultPart] = dataclasses.field( + default_factory=list + ) + system: list[str] = dataclasses.field(default_factory=list) + """System text that arrived after the conversation began.""" + prompt: list[messages_.TextPart | messages_.FilePart] = dataclasses.field( + default_factory=list + ) + transcript: list[messages_.Message] = dataclasses.field( + default_factory=list + ) + """The transcript once this call succeeds, before its answer.""" + + +@dataclasses.dataclass +class _Runtime: + """What lives only while a worker is open. Never copied, never dumped.""" + + adapter: base.Adapter + bridge: _bridge.Bridge + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + in_session: bool = False + # The running turn: its pump task and the queue it fills. + queue: asyncio.Queue[Any] | None = None + pump: asyncio.Task[None] | None = None + # Bridge calls handed to the caller, waiting for their results. + pending: dict[str, _Call] = dataclasses.field(default_factory=dict) + # CLI tool ids that are bridge calls, hidden from the stream. + hidden: set[str] = dataclasses.field(default_factory=set) + # The context of the latest `stream()` call: approvals ask in it. + context: contextvars.Context | None = None + + +class Worker(pydantic.BaseModel): + """Base class of the workers, such as `Codex`. + + `Worker.model_validate(data)` builds the kind that `data["kind"]` names. + """ + + kind: str + model: str | None = None + """The CLI's model. None: the CLI's own default.""" + cwd: str = pydantic.Field(default_factory=os.getcwd) + """Where the CLI runs and keeps its sessions.""" + env: dict[str, str] = {} + """Added to the CLI's environment: keys, base URLs, config dirs.""" + executable: str | None = None + effort: Effort | None = None + mcp_servers: dict[str, dict[str, Any]] = {} + """MCP servers for the CLI. + + `{"name": {"command": ..., "args": [...]}}` or `{"name": {"url": ...}}`. + """ + approve: bool = False + """Ask before each of the CLI's own tool calls, through `ai.hook`. + + False: the CLI's own permission settings decide (`allowed_tools`, + Codex's `sandbox`). True: each call becomes a `ToolApproval` hook, which + `ai.resolve_hook` and UIs answer as they answer any tool approval. Needs + an `ai.Agent.run()`, where hooks live. + """ + session_id: str | None = None + """The CLI's own id for the conversation. Set it to resume one.""" + fork: bool = False + """With `session_id`: fork that conversation natively on first use.""" + messages: list[messages_.Message] = [] + """The conversation so far. Send it back with the next user message.""" + + _kinds: ClassVar[dict[str, type[Worker]]] = {} + _runtime: _Runtime | None = pydantic.PrivateAttr(default=None) + + @classmethod + def __pydantic_init_subclass__(cls, **kwargs: Any) -> None: # noqa: PLW3201 + super().__pydantic_init_subclass__(**kwargs) + kind = cls.model_fields["kind"].default + if isinstance(kind, str): + Worker._kinds[kind] = cls + + @pydantic.model_validator(mode="wrap") + @classmethod + def _of_kind( + cls, data: Any, handler: pydantic.ValidatorFunctionWrapHandler + ) -> Any: + if cls is Worker and isinstance(data, dict): + kind = Worker._kinds.get(str(data.get("kind"))) + if kind is None: + raise ValueError(f"no worker of kind {data.get('kind')!r}") + return kind.model_validate(data) + return handler(data) + + def __copy__(self) -> Self: + # A copy is a new worker: it opens a CLI of its own. + copied = super().__copy__() + copied._runtime = None + return copied + + def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self: + copied = super().__deepcopy__(memo) + copied._runtime = None + return copied + + # -- subclass point ------------------------------------------------------- + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: base.Approve | None, + ) -> base.Adapter: + raise NotImplementedError + + # -- ModelLike ------------------------------------------------------------ + + @property + def id(self) -> str: + return self.model or "default" + + @property + def provider_name(self) -> str: + return self.kind + + async def stream( + self, + messages: list[messages_.Message], + *, + tools: Sequence[tools_.Tool] | None = None, + output_type: type[pydantic.BaseModel] | None = None, + params: Any = None, + ) -> AsyncGenerator[events_.Event]: + """Answer `messages` with one assistant message. + + Use `ai.stream(worker, messages)`, as with a model. The message ends + in one of two ways: + + - `finish_reason="tool_call"`: the CLI called `ai` tools from + `tools=`. The message holds them as `ToolCallPart`s. Run them and + call again with the results; the CLI waits until then. + - Anything else: the CLI's turn is over. + + `params` is ignored: the CLI chooses its own. + """ + rt = await self._open() + await rt.lock.acquire() + paused = False + try: + if self.approve: + try: + hooks.get_hook_registry() + runtime_.get_runtime() + except LookupError: + raise errors.WorkerError( + "approve=True asks through ai.hook, which needs a " + "run: use the worker in ai.Agent.run()" + ) from None + rt.context = contextvars.copy_context() + plan = self._plan(messages) + rt.bridge.set_tools( + [t for t in tools or () if t.kind == "function"] + ) + if rt.pending: + if plan.results: + self._answer(plan.results) + else: + await self._abandon() + if plan.prompt: + if rt.in_session and plan.system: + text = "\n\n".join(plan.system) + "\n\n" + plan.prompt.insert(0, messages_.TextPart(text=text)) + 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)) + else: + await self._ensure_session(plan) + self._start(plan.prompt, output_type) + if rt.queue is None: + raise ValueError( + "stream(): no new user message or tool result to send" + ) + + stamp = messages_.Message(role="assistant", parts=[]) + hydrator = events_.MessageHydrator() + + def out(event: events_.Event) -> events_.Event: + event = event.model_copy( + update={ + "message": messages_.Message( + id=stamp.id, role="assistant", parts=[] + ) + } + ) + hydrator.feed(event) + return event + + yield out(events_.StreamStart()) + # Text after the last CLI tool call: the answer, for output_type. + answer: list[str] = [] + while True: + item = await rt.queue.get() + if isinstance(item, _Failed): + self._end_turn() + raise item.error + if isinstance(item, _Call): + calls = [item] + with contextlib.suppress(TimeoutError): + async with asyncio.timeout(BATCH_WINDOW): + while True: + more = await rt.queue.get() + if isinstance(more, _Call): + calls.append(more) + else: + # Not ours to handle here: put it back. + rt.queue.put_nowait(more) + break + for call in calls: + rt.pending[call.tool_call_id] = call + part = messages_.ToolCallPart( + tool_call_id=call.tool_call_id, + tool_name=call.name, + tool_args=json.dumps(call.arguments), + ) + yield out( + events_.ToolStart( + tool_call_id=part.tool_call_id, + tool_name=part.tool_name, + ) + ) + yield out( + events_.ToolDelta( + tool_call_id=part.tool_call_id, + chunk=part.tool_args, + ) + ) + yield out( + events_.ToolEnd( + tool_call_id=part.tool_call_id, tool_call=part + ) + ) + paused = True + yield out( + events_.StreamEnd( + finish_reason="tool_call", + provider_metadata=self._metadata(None), + ) + ) + break + if isinstance(item, events_.StreamEnd): + self._end_turn() + usage = item.usage + if output_type is not None: + text, usage = await self._output( + "".join(answer), output_type, usage + ) + yield out(events_.TextStart(block_id="output")) + yield out( + events_.TextDelta(chunk=text, block_id="output") + ) + yield out(events_.TextEnd(block_id="output")) + finish = item.finish_reason + yield out( + events_.StreamEnd( + finish_reason=( + "other" if finish == "cancelled" else finish + ), + usage=usage, + provider_metadata={ + **(item.provider_metadata or {}), + **self._metadata(finish), + }, + ) + ) + break + for event in self._visible(item): + if output_type is not None and isinstance( + event, + events_.TextStart | events_.TextDelta | events_.TextEnd, + ): + # Narration is not the output. Keep only the text + # after the last tool call, and send it once parsed. + if isinstance(event, events_.TextDelta): + answer.append(event.chunk) + continue + if isinstance(event, events_.BuiltinToolStart): + answer.clear() + yield out(event) + # Committed only now: a call that failed can be sent again. + self.messages = [*plan.transcript, hydrator.message] + except (GeneratorExit, asyncio.CancelledError): + if not paused: + # Nobody reads this turn any more: end it in the CLI too. + await self._abandon() + raise + finally: + rt.lock.release() + + # -- the conversation ----------------------------------------------------- + + def _plan(self, messages: list[messages_.Message]) -> _Plan: + """Compare `messages` with the transcript; return what is new.""" + known = [m.id for m in self.messages] + given = [m.id for m in messages] + plan = _Plan() + system = [m.text for m in messages if m.role == "system"] + plan.instructions = "\n\n".join(system) or None + if known: + if given[: len(known)] != known: + raise errors.TranscriptMismatchError( + f"{self.kind}: these messages do not continue this " + "worker's conversation; start a new worker" + ) + new = messages[len(known) :] + plan.transcript = list(self.messages) + elif self.session_id is not None: + # A stored conversation the worker has not loaded: the CLI holds + # the past, so everything given is new. `load()` reads the past. + new = list(messages) + else: + # A new conversation. What comes before the last user turn is + # its past; the CLI is seeded with it before the first turn. + cut = len(messages) + while cut and messages[cut - 1].role in ("user", "system", "tool"): + cut -= 1 + plan.seed = [m for m in messages[:cut] if m.role != "system"] + new = messages[cut:] + plan.transcript = list(messages[:cut]) + for message in new: + if message.role == "tool": + plan.results.extend(message.tool_results) + elif message.role == "system": + plan.system.append(message.text) + elif message.role == "user": + plan.prompt.extend( + p + for p in message.parts + if isinstance(p, messages_.TextPart | messages_.FilePart) + ) + else: + raise errors.TranscriptMismatchError( + f"{self.kind}: a {message.role} message the worker did " + "not write; only user, system and tool messages can " + "follow its last answer (load() reads a stored past)" + ) + plan.transcript.extend(new) + return plan + + async def _ensure_session(self, plan: _Plan) -> None: + rt = self._require_runtime() + if rt.in_session: + return + self.session_id = await rt.adapter.open( + session_id=self.session_id, + fork=self.fork, + seed=plan.seed if self.session_id is None and plan.seed else None, + instructions=plan.instructions, + ) + self.fork = False + rt.in_session = True + + def _start( + self, + prompt: list[messages_.TextPart | messages_.FilePart], + output_type: type[pydantic.BaseModel] | None, + ) -> None: + rt = self._require_runtime() + if output_type is not None: + schema = json.dumps(TypeAdapter(output_type).json_schema()) + prompt = [ + *prompt, + messages_.TextPart(text=OUTPUT_INSTRUCTION + schema), + ] + queue: asyncio.Queue[Any] = asyncio.Queue() + adapter = rt.adapter + + async def pump() -> None: + try: + async for event in adapter.turn(prompt): + queue.put_nowait(event) + except asyncio.CancelledError: + raise + except BaseException as exc: + queue.put_nowait(_Failed(exc)) + + rt.queue = queue + rt.pump = asyncio.create_task(pump()) + + def _end_turn(self) -> None: + rt = self._require_runtime() + rt.queue = None + rt.pump = None + rt.hidden.clear() + + def _answer(self, results: list[messages_.ToolResultPart]) -> None: + rt = self._require_runtime() + for part in results: + call = rt.pending.pop(part.tool_call_id, None) + if call is None or call.answer.done(): + continue + value = part.model_input + if value is messages_.MODEL_INPUT_UNSET: + value = part.result + text = value if isinstance(value, str) else _dump(value) + call.answer.set_result((text, part.result_kind == "error")) + for call in rt.pending.values(): + # A call the caller did not answer: tell the CLI, do not hang. + if not call.answer.done(): + call.answer.set_result(("no result was given", True)) + rt.pending.clear() + + async def _abandon(self) -> None: + """End the running turn: the caller moved on without it.""" + rt = self._runtime + if rt is None: + return + for call in rt.pending.values(): + if not call.answer.done(): + call.answer.set_result(("cancelled by the caller", True)) + rt.pending.clear() + pump = rt.pump + if pump is None: + return + with contextlib.suppress(Exception): + await rt.adapter.interrupt() + pump.cancel() + with contextlib.suppress(BaseException): + await pump + self._end_turn() + + async def _on_call( + self, name: str, arguments: dict[str, Any] + ) -> tuple[str, bool]: + """Hand a bridge call to the caller, and wait for the result.""" + rt = self._runtime + if rt is None or rt.queue is None: + return ("no turn is running", True) + answer: asyncio.Future[tuple[str, bool]] = ( + asyncio.get_running_loop().create_future() + ) + rt.queue.put_nowait( + _Call(messages_.generate_id("call"), name, arguments, answer) + ) + return await answer + + async def _approve( + self, tool_call_id: str, tool_name: str, tool_input: dict[str, Any] + ) -> tuple[bool, str]: + """Ask about one of the CLI's own tool calls, through `ai.hook`.""" + if self._is_bridge_tool(tool_name): + # An `ai` tool: the agent gates it itself, if it is gated. + return True, "" + rt = self._runtime + if rt is None or rt.context is None: + return False, "no run to ask in; denied" + + async def ask() -> tuple[bool, str]: + approval = await hooks.hook( + f"approve_{tool_call_id}", + payload=hooks.ToolApproval, + metadata={"tool": tool_name, "kwargs": tool_input}, + tool_call_id=tool_call_id, + ) + return approval.granted, approval.reason or "denied" + + # In the caller's context: that is where the run's hooks live. + task = asyncio.create_task(ask(), context=rt.context) + try: + return await task + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + return False, "the approval was cancelled" + except Exception as exc: + return False, f"approval failed ({type(exc).__name__}); denied" + + def _is_bridge_tool(self, tool_name: str) -> bool: + """Whether a tool the CLI reports is a bridge call. + + Those reach the caller as `ToolCallPart`s, so the CLI's own report + of them is hidden. Claude Code names an MCP tool + `mcp____`; codex uses the bare name, or the server's + name when the tool name is missing (measured: codex 0.160). How pi + names them is not measured yet: the forms it may use are accepted. + """ + rt = self._runtime + if rt is None: + return False + server = _bridge.SERVER_NAME + return tool_name == server or any( + tool_name + in (n, f"mcp__{server}__{n}", f"{server}_{n}", f"{server}__{n}") + for n in rt.bridge.names + ) + + def _visible(self, event: events_.Event) -> list[events_.Event]: + """Drop StreamStart and the CLI's own report of bridge calls.""" + rt = self._require_runtime() + match event: + case events_.StreamStart(): + return [] # a worker message has its own + case events_.BuiltinToolStart(): + if self._is_bridge_tool(event.tool_name): + rt.hidden.add(event.tool_call_id) + return [] + case events_.BuiltinToolEnd(): + if event.tool_call_id in rt.hidden or self._is_bridge_tool( + event.tool_call.tool_name + ): + rt.hidden.add(event.tool_call_id) + return [] + case events_.BuiltinToolDelta() | events_.BuiltinToolResult(): + if event.tool_call_id in rt.hidden: + return [] + return [event] + + async def _output( + self, + text: str, + output_type: type[pydantic.BaseModel], + usage: usage_.Usage | None, + ) -> tuple[str, usage_.Usage | None]: + """Parse the answer, asking again when it does not validate.""" + adapter = self._require_runtime().adapter + parsed, error = _parse(text, output_type) + attempt = 0 + while error is not None and attempt < OUTPUT_RETRIES: + attempt += 1 + retry = messages_.TextPart( + text=f"Your previous reply failed validation: {error}\n" + "Reply again with ONLY the corrected JSON." + ) + chunks: list[str] = [] + async for event in adapter.turn([retry]): + if isinstance(event, events_.BuiltinToolStart): + chunks.clear() + elif isinstance(event, events_.TextDelta): + chunks.append(event.chunk) + elif isinstance(event, events_.StreamEnd) and event.usage: + usage = usage + event.usage if usage else event.usage + text = "".join(chunks) + parsed, error = _parse(text, output_type) + return (parsed.model_dump_json() if parsed else text), usage + + def _metadata(self, finish: str | None) -> dict[str, Any]: + meta: dict[str, Any] = {"kind": self.kind} + if self.session_id: + meta["session_id"] = self.session_id + if finish is not None: + meta["finish_reason"] = finish + return {"worker": meta} + + # -- lifecycle ------------------------------------------------------------ + + async def _open(self) -> _Runtime: + if self._runtime is not None: + return self._runtime + bridge = _bridge.Bridge(self._on_call) + await bridge.start() + servers = { + **self.mcp_servers, + _bridge.SERVER_NAME: {"type": "http", "url": bridge.url}, + } + adapter = self._make_adapter( + servers, self._approve if self.approve else None + ) + try: + await adapter.start() + except BaseException: + await bridge.close() + raise + self._runtime = _Runtime(adapter=adapter, bridge=bridge) + return self._runtime + + async def close(self) -> None: + """End the turn, the CLI and the bridge. The fields stay as they are. + + A closed worker opens again on its next use, and resumes its + conversation. + """ + rt = self._runtime + if rt is None: + return + await self._abandon() + self._runtime = None + try: + await rt.adapter.close() + finally: + await rt.bridge.close() + + async def __aenter__(self) -> Self: + await self._open() + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self.close() + + def _require_runtime(self) -> _Runtime: + if self._runtime is None: + raise errors.WorkerError(f"{self.kind} worker is not open") + return self._runtime + + @property + def version(self) -> str | None: + """What the CLI's `--version` printed, once the worker is open.""" + return self._runtime.adapter.version if self._runtime else None + + # -- what only a CLI has -------------------------------------------------- + + async def load(self) -> None: + """Read the stored conversation `session_id` names into `messages`. + + Also after a person worked in the CLI's TUI (see `tui()`). + """ + if self.session_id is None: + raise errors.WorkerError(f"{self.kind}: no session_id to load") + rt = await self._open() + self.messages = await rt.adapter.history(self.session_id) + + async def steer(self, text: str) -> None: + """Send guidance to the turn that runs now; the agent acts on it. + + For use from another task while `stream()` runs. The text is not + added to `messages`; `load()` reads it back from the CLI. + """ + rt = self._require_runtime() + if rt.queue is None: + raise errors.WorkerError(f"{self.kind}: no turn is running") + await rt.adapter.steer(text) + + def tui_command(self) -> tuple[list[str], dict[str, str]]: + """Return the CLI's own TUI command for this conversation. + + Run it in `cwd`, with the env added to yours:: + + argv, env = worker.tui_command() + subprocess.run(argv, cwd=worker.cwd, env={**os.environ, **env}) + + A worker with no conversation yet gets a new one, when the CLI lets + us choose its id (Claude Code, Pi). Close the worker first: one + writer per conversation. + """ + if self.fork: + raise errors.WorkerError( + f"{self.kind}: run a turn first; the fork happens on first use" + ) + rt = self._runtime + adapter = ( + rt.adapter if rt else self._make_adapter(self.mcp_servers, None) + ) + argv, env, minted = adapter.tui_command(self.session_id) + if self.session_id is None: + if minted is None: + raise errors.WorkerError( + f"{self.kind} picks its own session id: run a turn " + "first, or pass session_id=" + ) + self.session_id = minted + return argv, env + + async def tui(self) -> int: + """Hand the conversation to a person, in this terminal. + + Runs the CLI's TUI on it and returns its exit code once the person + quits. The worker then reads the conversation back (`load()`) and + goes on headless. + """ + await self.close() + argv, env = self.tui_command() + proc = await asyncio.create_subprocess_exec( + *argv, cwd=self.cwd, env={**os.environ, **env} + ) + code = await proc.wait() + await self.load() + return code + + @classmethod + async def list_sessions(cls, **fields: Any) -> list[base.SessionInfo]: + """Return the conversations the CLI has stored for a directory. + + Takes the worker's fields: `cwd=`, `env=`, `executable=`. + """ + worker = cls(**fields) + adapter = worker._make_adapter({}, None) + await adapter.start() + try: + return await adapter.list_sessions() + finally: + await adapter.close() + + +def _parse( + text: str, output_type: type[pydantic.BaseModel] +) -> tuple[pydantic.BaseModel | None, Exception | None]: + try: + return output_type.model_validate_json(_extract_json(text)), None + except Exception as exc: + return None, exc + + +def _extract_json(text: str) -> str: + """Best effort: the model was told "only JSON", but fences happen.""" + candidate = text.strip() + if candidate.startswith("```"): + parts = candidate.split("```") + if len(parts) >= 2: + candidate = parts[1].removeprefix("json").strip() + try: + json.loads(candidate) + return candidate + except json.JSONDecodeError: + pass + for opener, closer in (("{", "}"), ("[", "]")): + start, end = candidate.find(opener), candidate.rfind(closer) + if 0 <= start < end: + sliced = candidate[start : end + 1] + try: + json.loads(sliced) + return sliced + except json.JSONDecodeError: + continue + return candidate + + +def _dump(value: Any) -> str: + if isinstance(value, pydantic.BaseModel): + return value.model_dump_json() + return json.dumps(value, default=str) diff --git a/tests/workers/__init__.py b/tests/workers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/workers/_adapters/__init__.py b/tests/workers/_adapters/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/workers/_adapters/test_codex.py b/tests/workers/_adapters/test_codex.py new file mode 100644 index 00000000..de1fedb9 --- /dev/null +++ b/tests/workers/_adapters/test_codex.py @@ -0,0 +1,384 @@ +"""The codex adapter against a scripted `codex app-server`, offline.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Callable +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 codex, jsonrpc + +from .. import fakes + +THREAD = "T" +CWD = "/w" + + +class AppServer: + """Answers like `codex app-server`. `turn` plays out after turn/start.""" + + def __init__(self, turn: list[dict[str, Any]] | None = None) -> None: + self.process = fakes.FakeProcess(self.answer) + self.turn = turn or [] + self.on_turn: Callable[[], None] | None = None + self.threads: dict[str, dict[str, Any]] = { + THREAD: {"id": THREAD, "cwd": CWD, "turns": []} + } + + def says(self, method: str, **params: Any) -> None: + self.process.server_says({"method": method, "params": params}) + + def answer(self, message: dict[str, Any]) -> Any: + method = message.get("method") + params = message.get("params") or {} + if method == "turn/start": + if self.on_turn is not None: + self.on_turn() + for record in self.turn: + self.process.server_says(record) + return {"turn": {"id": "turn-1"}} + if method == "turn/interrupt": + self.says( + "turn/completed", + threadId=THREAD, + turn={"id": "turn-1", "status": "interrupted"}, + ) + return {} + if method == "thread/read": + return {"thread": self.threads.get(str(params.get("threadId")), {})} + if method in ("thread/start", "thread/fork"): + return {"thread": {"id": "new"}} + if method == "thread/resume": + return {"thread": {"id": params["threadId"]}} + if "method" in message and "id" in message: + return {} + return None + + def sent(self, method: str) -> list[dict[str, Any]]: + return [w for w in self.process.written if w.get("method") == method] + + +def note(method: str, **params: Any) -> dict[str, Any]: + return { + "method": method, + "params": {"threadId": THREAD, "turnId": "turn-1", **params}, + } + + +COMMAND = { + "type": "commandExecution", + "id": "c1", + "command": "/bin/zsh -lc 'cat secret.txt'", + "aggregatedOutput": "the word is HERON\n", + "exitCode": 0, + "status": "completed", +} + +TURN = [ + note("item/reasoning/summaryTextDelta", itemId="r1", delta="plan"), + note("item/completed", item={"type": "reasoning", "id": "r1"}), + note("item/started", item={**COMMAND, "aggregatedOutput": None}), + note("item/completed", item=COMMAND), + note("item/agentMessage/delta", itemId="a1", delta="HER"), + note("item/agentMessage/delta", itemId="a1", delta="ON"), + note("item/completed", item={"type": "agentMessage", "id": "a1"}), + note( + "thread/tokenUsage/updated", + tokenUsage={"lastTurn": {"inputTokens": 10, "outputTokens": 4}}, + ), + note("turn/completed", turn={"id": "turn-1", "status": "completed"}), +] + + +def _adapter( + server: AppServer, approve: Any = None, **options: Any +) -> codex.CodexAdapter: + adapter = codex.CodexAdapter( + cwd=CWD, + env={}, + executable="codex", + model=None, + mcp_servers=options.pop("mcp_servers", {}), + approve=approve, + sandbox=None, + config={}, + ) + peer = jsonrpc.JsonRpcPeer(cast("Any", server.process), kind="codex") + adapter._process = cast("Any", server.process) + adapter._peer = peer + adapter._thread = THREAD + peer.start( + adapter._on_notify, adapter._on_request, on_crash=adapter._on_crash + ) + return adapter + + +async def _turn(adapter: codex.CodexAdapter, text: str = "go") -> list[Any]: + async with asyncio.timeout(5): + return [e async for e in adapter.turn([m.TextPart(text=text)])] + + +async def test_a_turn_streams_in_the_model_view() -> None: + adapter = _adapter(AppServer(TURN)) + seen = await _turn(adapter) + assert [type(e).__name__ for e in seen] == [ + "StreamStart", + "ReasoningStart", + "ReasoningDelta", + "ReasoningEnd", + "BuiltinToolStart", + "BuiltinToolDelta", + "BuiltinToolEnd", + "BuiltinToolResult", + "TextStart", + "TextDelta", + "TextDelta", + "TextEnd", + "StreamEnd", + ] + result = next(e for e in seen if isinstance(e, events.BuiltinToolResult)) + # codex 0.160 reports a command's output as `aggregatedOutput`. + assert result.result.result == "the word is HERON\n" + end = seen[-1] + assert end.finish_reason == "stop" + assert end.usage.input_tokens == 10 + hydrator = events.MessageHydrator() + for event in seen: + hydrator.feed(event) + assert [type(p).__name__ for p in hydrator.message.parts] == [ + "ReasoningPart", + "BuiltinToolCallPart", + "BuiltinToolReturnPart", + "TextPart", + ] + + +async def test_a_message_that_was_never_streamed_still_arrives() -> None: + server = AppServer( + [ + note( + "item/completed", + item={"type": "agentMessage", "id": "a1", "text": "PONG"}, + ), + note( + "turn/completed", turn={"id": "turn-1", "status": "completed"} + ), + ] + ) + seen = await _turn(_adapter(server)) + assert [e.chunk for e in seen if isinstance(e, events.TextDelta)] == [ + "PONG" + ] + + +async def test_another_turns_notifications_are_not_this_turns() -> None: + server = AppServer( + [ + { + "method": "turn/completed", + "params": {"threadId": THREAD, "turn": {"id": "old"}}, + }, + { + "method": "item/agentMessage/delta", + "params": {"threadId": "other", "itemId": "x", "delta": "no"}, + }, + note( + "turn/completed", turn={"id": "turn-1", "status": "completed"} + ), + ] + ) + seen = await _turn(_adapter(server)) + assert [type(e).__name__ for e in seen] == ["StreamStart", "StreamEnd"] + + +async def test_a_failed_turn_without_a_key_is_not_authenticated() -> None: + server = AppServer( + [ + note("error", error={"message": "401 Unauthorized"}), + note("turn/completed", turn={"id": "turn-1", "status": "failed"}), + ] + ) + with pytest.raises(workers.NotAuthenticatedError): + await _turn(_adapter(server)) + + +@pytest.mark.parametrize("granted", [True, False]) +async def test_approvals_ask_and_answer(granted: bool) -> None: + server = AppServer() + asked: list[tuple[str, str, dict[str, 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, "no" + + def ask() -> None: + server.process.server_says( + { + "id": 500, + "method": "item/commandExecution/requestApproval", + "params": { + "threadId": THREAD, + "turnId": "turn-1", + "itemId": "i1", + "command": "rm -rf /", + }, + } + ) + + def complete(message: dict[str, Any]) -> Any: + if message.get("id") == 500: + server.says( + "turn/completed", + threadId=THREAD, + turn={"id": "turn-1", "status": "completed"}, + ) + return None + return AppServer.answer(server, message) + + server.on_turn = ask + server.process.answer = complete + adapter = _adapter(server, approve) + assert adapter._thread_params()["approvalPolicy"] == "untrusted" + await _turn(adapter) + assert asked == [("i1", "commandExecution", {"command": "rm -rf /"})] + answer = next(w for w in server.process.written if w.get("id") == 500) + assert answer["result"] == {"decision": "accept" if granted else "decline"} + + +async def test_without_a_gate_codex_never_asks() -> None: + adapter = _adapter(AppServer()) + assert adapter._thread_params() == { + "sandbox": "danger-full-access", + "approvalPolicy": "never", + } + + +async def test_an_interrupted_turn_ends_as_cancelled() -> None: + server = AppServer() + adapter = _adapter(server) + turn = adapter.turn([m.TextPart(text="go")]) + assert isinstance(await anext(turn), events.StreamStart) + reader = asyncio.ensure_future(_rest(turn)) + await asyncio.sleep(0) + await adapter.interrupt() + seen = await asyncio.wait_for(reader, 5) + assert seen[-1].finish_reason == "cancelled" + assert server.sent("turn/interrupt")[0]["params"] == { + "threadId": THREAD, + "turnId": "turn-1", + } + + +async def _rest(turn: Any) -> list[Any]: + return [e async for e in turn] + + +async def test_opening_a_new_thread_seeds_it() -> None: + server = AppServer() + adapter = _adapter(server) + seed = [ + m.Message(role="user", parts=[m.TextPart(text="earlier")]), + m.Message(role="assistant", parts=[m.TextPart(text="answer")]), + ] + assert ( + await adapter.open( + session_id=None, fork=False, seed=seed, instructions="Be terse." + ) + == "new" + ) + started = server.sent("thread/start")[0]["params"] + assert started["cwd"] == CWD + assert started["developerInstructions"] == "Be terse." + injected = server.sent("thread/inject_items")[0]["params"] + assert injected["threadId"] == "new" + assert [i["role"] for i in injected["items"]] == ["user", "assistant"] + + +async def test_a_fork_is_native() -> None: + server = AppServer() + adapter = _adapter(server) + forked = await adapter.open( + session_id=THREAD, fork=True, seed=None, instructions=None + ) + assert forked == "new" + assert server.sent("thread/fork")[0]["params"]["threadId"] == THREAD + + +async def test_a_thread_from_another_directory_is_refused() -> None: + server = AppServer() + server.threads["elsewhere"] = {"id": "elsewhere", "cwd": "/other"} + adapter = _adapter(server) + with pytest.raises(workers.WorkerError, match="/other"): + await adapter.open( + session_id="elsewhere", fork=False, seed=None, instructions=None + ) + + +async def test_history_is_the_model_view() -> None: + server = AppServer() + server.threads[THREAD]["turns"] = [ + { + "items": [ + { + "type": "userMessage", + "content": [{"type": "text", "text": "q"}], + }, + COMMAND, + {"type": "agentMessage", "text": "HERON"}, + ] + } + ] + history = await _adapter(server).history(THREAD) + assert [x.role for x in history] == ["user", "assistant", "assistant"] + call, result = history[1].parts + assert isinstance(call, m.BuiltinToolCallPart) + assert json.loads(call.tool_args)["command"].endswith("'cat secret.txt'") + assert isinstance(result, m.BuiltinToolReturnPart) + assert result.result == "the word is HERON\n" + + +async def test_listing_asks_for_this_directory_and_reads_every_page() -> None: + """`thread/list` is paged, and the store is the whole machine's.""" + asked: list[dict[str, Any]] = [] + pages = { + None: ({"id": "t1", "cwd": CWD}, "p2"), + "p2": ({"id": "t2", "cwd": CWD}, None), + } + server = AppServer() + + def answer(message: dict[str, Any]) -> Any: + if message.get("method") != "thread/list": + return AppServer.answer(server, message) + params = message["params"] + asked.append(params) + thread, cursor = pages[params.get("cursor")] + return {"data": [thread], "nextCursor": cursor} + + server.process.answer = answer + infos = await _adapter(server).list_sessions() + assert [i.session_id for i in infos] == ["t1", "t2"] + assert all(p["cwd"] == CWD for p in asked) + assert [p.get("cursor") for p in asked] == [None, "p2"] + + +async def test_the_tui_gets_no_bridge() -> None: + adapter = _adapter( + AppServer(), + mcp_servers={ + "ai_sdk": {"url": "http://127.0.0.1:1/mcp"}, + "docs": {"command": "docs-mcp"}, + }, + ) + argv, _, minted = adapter.tui_command(THREAD) + assert argv[:3] == ["codex", "resume", THREAD] + assert 'mcp_servers.docs.command="docs-mcp"' in argv + assert not any("ai_sdk" in arg for arg in argv) + assert minted is None # codex picks its own ids diff --git a/tests/workers/_adapters/test_handoff.py b/tests/workers/_adapters/test_handoff.py new file mode 100644 index 00000000..4f55dadd --- /dev/null +++ b/tests/workers/_adapters/test_handoff.py @@ -0,0 +1,234 @@ +"""The seeding rule, pinned without a CLI. + +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. +""" + +from __future__ import annotations + +import json + +from ai.types.messages import ( + BuiltinToolCallPart, + BuiltinToolReturnPart, + Message, + ReasoningPart, + TextPart, + ToolCallPart, + ToolResultPart, +) +from ai.workers._adapters import _handoff + +HISTORY = [ + Message(role="user", parts=[TextPart(text="Create recipe.txt")]), + Message( + role="assistant", + parts=[ + ReasoningPart(text="The user wants a file; Write is the tool."), + TextPart(text="Creating it."), + ToolCallPart( + tool_call_id="call_1", + tool_name="Write", + tool_args=json.dumps( + {"file_path": "/w/recipe.txt", "content": "two eggs"} + ), + ), + ], + ), + Message( + role="tool", + parts=[ + ToolResultPart( + tool_call_id="call_1", + tool_name="Write", + result={"ok": True}, + result_kind="json", + ) + ], + ), + Message(role="assistant", parts=[TextPart(text="Done.")]), +] + +# The same conversation, the way a worker holds it. +MODEL_VIEW = [ + HISTORY[0], + Message( + role="assistant", + parts=[ + ReasoningPart(text="The user wants a file; Write is the tool."), + TextPart(text="Creating it."), + BuiltinToolCallPart( + tool_call_id="call_1", + tool_name="Write", + tool_args=json.dumps( + {"file_path": "/w/recipe.txt", "content": "two eggs"} + ), + ), + BuiltinToolReturnPart( + tool_call_id="call_1", tool_name="Write", result={"ok": True} + ), + ], + ), + HISTORY[3], +] + + +def test_codex_items_follow_the_rule() -> None: + items = _handoff.to_codex_items(HISTORY) + kinds = [(i["type"], i.get("role")) for i in items] + assert kinds == [ + ("message", "user"), + ("message", "assistant"), # "Creating it." + ("function_call", None), # Write, name untouched + ("message", "user"), # the reasoning note, after its turn + ("function_call_output", None), + ("message", "assistant"), # "Done." + ] + call = items[2] + assert call["name"] == "Write" + assert json.loads(call["arguments"])["content"] == "two eggs" + assert call["call_id"] == "call_1" + out = items[4] + 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 "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 + ) + + +def test_claude_records_follow_the_rule() -> None: + recs = _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] == [ + "user", + "assistant", + "user", + "assistant", + ] + assistant = recs[1]["message"]["content"] + assert [b["type"] for b in assistant] == ["text", "tool_use"] + assert assistant[1]["name"] == "Write" + assert assistant[1]["input"]["content"] == "two eggs" + assert assistant[1]["id"] == "call_1" + 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 + # a real chain, in the CLI's own store shape + assert recs[0]["parentUuid"] is None + assert all( + recs[i]["parentUuid"] == recs[i - 1]["uuid"] + for i in range(1, len(recs)) + ) + assert all(r["sessionId"] == "sid-1" and r["cwd"] == "/w" for r in recs) + + +def test_pi_merges_assistant_messages_and_keeps_tool_results() -> None: + rendered = _handoff.to_pi_messages( + [ + Message(role="user", parts=[TextPart(text="hi")]), + Message(role="assistant", parts=[TextPart(text="one")]), + Message( + role="assistant", + parts=[ + BuiltinToolCallPart( + tool_call_id="c1", tool_name="Bash", tool_args="{}" + ), + BuiltinToolReturnPart( + tool_call_id="c1", + tool_name="Bash", + result={"ok": True}, + is_error=True, + ), + ], + ), + ], + timestamp=1, + ) + assert [r["role"] for r in rendered] == ["user", "assistant", "toolResult"] + assert [b["type"] for b in rendered[1]["content"]] == ["text", "toolCall"] + assert rendered[1]["stopReason"] == "toolUse" + assert rendered[2]["content"] == [{"type": "text", "text": '{"ok": true}'}] + 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: + one = Message( + role="assistant", + parts=[ + BuiltinToolCallPart(tool_call_id="c", tool_name="Read"), + BuiltinToolReturnPart( + tool_call_id="c", tool_name="Read", result="x" + ), + TextPart(text="It says x."), + ], + ) + first, result, after = _handoff.harness_view([one]) + assert (first.role, result.role, after.role) == ( + "assistant", + "tool", + "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" + + +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) + assert not any(i["type"] == "reasoning" for i in items) + blocks = [ + b + for r in _handoff.to_claude_records(HISTORY, "s", "/w") + if isinstance(r["message"]["content"], list) + for b in r["message"]["content"] + ] + assert not any(b["type"] == "thinking" for b in blocks) + + +def test_empty_parts_produce_nothing() -> None: + blank = [Message(role="assistant", parts=[TextPart(text=" ")])] + assert _handoff.to_codex_items(blank) == [] + empty = [Message(role="assistant", parts=[ReasoningPart(text="")])] + assert _handoff.to_claude_records(empty, "s", "/w") == [] diff --git a/tests/workers/_adapters/test_jsonrpc.py b/tests/workers/_adapters/test_jsonrpc.py new file mode 100644 index 00000000..a6849316 --- /dev/null +++ b/tests/workers/_adapters/test_jsonrpc.py @@ -0,0 +1,161 @@ +"""The JSON-RPC peer codex is driven over, offline.""" + +from __future__ import annotations + +import asyncio +import time +from typing import Any, cast + +import pytest + +from ai import workers +from ai.workers._adapters import jsonrpc + +from .. import fakes + + +def _peer( + process: fakes.FakeProcess, + on_request: Any = None, + on_notify: Any = None, +) -> jsonrpc.JsonRpcPeer: + async def no_notify(method: str, params: dict[str, Any]) -> None: + pass + + async def no_request( + request_id: Any, method: str, params: dict[str, Any] + ) -> None: + pass + + peer = jsonrpc.JsonRpcPeer(cast("Any", process), kind="test") + peer.start(on_notify or no_notify, on_request or no_request) + return peer + + +async def test_a_call_gets_its_answer() -> None: + process = fakes.FakeProcess(lambda message: {"echo": message["method"]}) + peer = _peer(process) + assert await peer.call("ping", timeout=1) == {"echo": "ping"} + await peer.close() + + +async def test_a_call_made_while_answering_a_request_is_answered() -> None: + """Answering a server request must not stop the peer reading.""" + process = fakes.FakeProcess() + took: list[float] = [] + done = asyncio.Event() + peer: jsonrpc.JsonRpcPeer + + async def on_request( + request_id: Any, method: str, params: dict[str, Any] + ) -> None: + started = time.monotonic() + await peer.respond(request_id, {"decision": "decline"}) + await peer.call("turn/interrupt", {}, timeout=2) + took.append(time.monotonic() - started) + done.set() + + peer = _peer(process, on_request=on_request) + process.server_says({"id": 99, "method": "approval", "params": {}}) + await asyncio.wait_for(done.wait(), 3) + assert took[0] < 0.5 + await peer.close() + + +async def test_notifications_flow_while_a_request_is_pending() -> None: + """A slow answer (a person approving) does not freeze the stream.""" + process = fakes.FakeProcess() + release = asyncio.Event() + seen: list[str] = [] + + async def on_request( + request_id: Any, method: str, params: dict[str, Any] + ) -> None: + await release.wait() + + async def on_notify(method: str, params: dict[str, Any]) -> None: + seen.append(method) + + peer = _peer(process, on_request=on_request, on_notify=on_notify) + process.server_says({"id": 1, "method": "approval", "params": {}}) + process.server_says({"method": "item/agentMessage/delta", "params": {}}) + assert await peer.call("thread/read", timeout=1) == {} + for _ in range(50): + if seen: + break + await asyncio.sleep(0.01) + assert seen == ["item/agentMessage/delta"] + release.set() + await peer.close() + + +async def test_notifications_keep_their_order() -> None: + process = fakes.FakeProcess() + seen: list[int] = [] + + async def on_notify(method: str, params: dict[str, Any]) -> None: + await asyncio.sleep(0) + seen.append(params["n"]) + + peer = _peer(process, on_notify=on_notify) + for n in range(20): + process.server_says({"method": "n", "params": {"n": n}}) + await peer.call("ping", timeout=1) + assert seen == list(range(20)) + await peer.close() + + +async def test_close_cancels_requests_still_being_answered() -> None: + process = fakes.FakeProcess() + cancelled = asyncio.Event() + started = asyncio.Event() + + async def on_request( + request_id: Any, method: str, params: dict[str, Any] + ) -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + peer = _peer(process, on_request=on_request) + process.server_says({"id": 1, "method": "approval", "params": {}}) + await asyncio.wait_for(started.wait(), 1) + await peer.close() + assert cancelled.is_set() + + +async def test_a_failing_request_handler_does_not_kill_the_peer() -> None: + process = fakes.FakeProcess() + + async def on_request( + request_id: Any, method: str, params: dict[str, Any] + ) -> None: + raise ValueError("bug in the handler") + + peer = _peer(process, on_request=on_request) + process.server_says({"id": 1, "method": "approval", "params": {}}) + assert await peer.call("ping", timeout=1) == {} + await peer.close() + with pytest.raises(workers.WorkerError): + await peer.call("ping", timeout=1) + + +async def test_a_dead_process_fails_its_callers() -> None: + process = fakes.FakeProcess(lambda message: None) # never answers + crashed: list[BaseException] = [] + peer = jsonrpc.JsonRpcPeer(cast("Any", process), kind="test") + + async def nothing(*args: Any) -> None: + pass + + peer.start(nothing, nothing, on_crash=crashed.append) + call = asyncio.ensure_future(peer.call("ping", timeout=2)) + await asyncio.sleep(0) + await process.terminate() # EOF + with pytest.raises(workers.WorkerError, match="crashed"): + await call + assert len(crashed) == 1 + await peer.close() diff --git a/tests/workers/fakes.py b/tests/workers/fakes.py new file mode 100644 index 00000000..51ed6b6c --- /dev/null +++ b/tests/workers/fakes.py @@ -0,0 +1,231 @@ +"""Test doubles for workers: a scripted adapter and an in-memory process.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import sys +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from typing import Any, Literal + +import pydantic + +from ai import workers +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 + + +@dataclasses.dataclass +class Pause: + """Wait here until the turn is interrupted (or cancelled).""" + + +@dataclasses.dataclass +class Run: + """Await a coroutine here, mid-turn: a bridge call, say.""" + + fn: Callable[[], Awaitable[None]] + + +@dataclasses.dataclass +class Ask: + """Ask the worker to approve one of the CLI's own tool calls.""" + + tool_call_id: str + tool_name: str + tool_input: dict[str, Any] = dataclasses.field(default_factory=dict) + + +Step = events.Event | Pause | Run | Ask +Script = Callable[[str], Iterable[Step]] + + +def text(chunk: str, block: str = "t") -> list[Step]: + return [ + events.TextStart(block_id=block), + events.TextDelta(chunk=chunk, block_id=block), + events.TextEnd(block_id=block), + ] + + +def tool( + call_id: str, name: str, args: str = "{}", result: Any = "ok" +) -> 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), + ] + + +def reply(prompt: str) -> list[Step]: + """Echo the prompt back.""" + return text(prompt) + + +class FakeAdapter: + def __init__(self, script: Script, approve: base.Approve | None) -> None: + self.script = script + self.approve = approve + self.version: str | None = None + self.opens: list[dict[str, Any]] = [] + self.turns: list[str] = [] + self.steers: list[str] = [] + self.interrupts = 0 + self.answers: list[tuple[bool, str]] = [] + self.closed = False + self.stored: dict[str, list[m.Message]] = {} + self._resume = asyncio.Event() + + async def start(self) -> None: + self.version = "fake 1.0" + + async def close(self) -> None: + self.closed = True + + async def open( + self, + *, + session_id: str | None, + fork: bool, + seed: list[m.Message] | None, + instructions: str | None, + ) -> str: + self.opens.append( + { + "session_id": session_id, + "fork": fork, + "seed": seed, + "instructions": instructions, + } + ) + if fork: + 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)) + self._resume.clear() + yield events.StreamStart() + for step in self.script(base.text_of(prompt)): + if isinstance(step, Pause): + await self._resume.wait() + yield events.StreamEnd(finish_reason="cancelled") + return + if isinstance(step, Run): + await step.fn() + elif isinstance(step, Ask): + assert self.approve is not None + self.answers.append( + await self.approve( + step.tool_call_id, step.tool_name, step.tool_input + ) + ) + else: + yield step + yield events.StreamEnd( + finish_reason="stop", + usage=usage_.Usage(input_tokens=3, output_tokens=2), + ) + + async def steer(self, text: str) -> None: + self.steers.append(text) + + async def interrupt(self) -> None: + self.interrupts += 1 + self._resume.set() + + 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")] + + def tui_command( + self, session_id: str | None + ) -> tuple[list[str], dict[str, str], str | None]: + return ( + [sys.executable, "-c", "pass"], + {"FAKE_TUI": "1"}, + session_id or "minted", + ) + + +class FakeWorker(workers.Worker): + """A worker whose CLI is a `FakeAdapter`, run by a script.""" + + kind: Literal["fake"] = "fake" + _script: Script = pydantic.PrivateAttr(default_factory=lambda: reply) + _fake: FakeAdapter | None = pydantic.PrivateAttr(default=None) + _servers: dict[str, dict[str, Any]] = pydantic.PrivateAttr( + default_factory=dict + ) + + def _make_adapter( + self, + mcp_servers: dict[str, dict[str, Any]], + approve: base.Approve | None, + ) -> base.Adapter: + self._fake = FakeAdapter(self._run, approve) + self._servers = mcp_servers + return self._fake + + def _run(self, prompt: str) -> Iterable[Step]: + return self._script(prompt) + + @property + def servers(self) -> dict[str, dict[str, Any]]: + return self._servers + + def script(self, script: Script) -> None: + self._script = script + + @property + def fake(self) -> FakeAdapter: + assert self._fake is not None + return self._fake + + +class FakeProcess: + """In-memory duplex stdio, for the JSON-RPC peer and the adapters. + + `answer(message)` is called for every message the client writes. For a + request, a non-None return is sent back as its result. `server_says()` + pushes a line from the "server": a notification or a request. + """ + + def __init__( + self, answer: Callable[[dict[str, Any]], Any] | None = None + ) -> None: + 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") + + async def write(self, data: str) -> None: + message = json.loads(data) + self.written.append(message) + result = self.answer(message) + if "method" in message and "id" in message and result is not None: + self.server_says({"id": message["id"], "result": result}) + + async def readline(self) -> str: + return await self._lines.get() + + async def read_stderr(self) -> str: + return "" + + async def wait(self) -> int: + return 0 + + async def terminate(self) -> None: + self._lines.put_nowait("") + + async def death_note(self) -> str: + return "exit code 0; no stderr" diff --git a/tests/workers/test_worker.py b/tests/workers/test_worker.py new file mode 100644 index 00000000..299cf600 --- /dev/null +++ b/tests/workers/test_worker.py @@ -0,0 +1,569 @@ +"""`ai.workers.Worker` over a scripted adapter: no real CLI runs here.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Iterator +from typing import Any + +import pydantic +import pytest + +import ai +from ai import workers +from ai.types import events +from ai.types import messages as m +from ai.workers import _bridge + +from . import fakes + + +@pytest.fixture +def worker() -> fakes.FakeWorker: + return fakes.FakeWorker() + + +async def post(url: str, payload: Any) -> Any: + """POST JSON the way an MCP client does, over a raw socket.""" + host, rest = url.removeprefix("http://").split(":", 1) + port, path = rest.split("/", 1) + reader, writer = await asyncio.open_connection(host, int(port)) + body = json.dumps(payload).encode() + writer.write( + ( + f"POST /{path} HTTP/1.1\r\nHost: {host}\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(body)}\r\n\r\n" + ).encode() + + body + ) + await writer.drain() + head = await reader.readuntil(b"\r\n\r\n") + length = 0 + for line in head.decode().split("\r\n"): + if line.lower().startswith("content-length:"): + length = int(line.split(":")[1]) + data = await reader.readexactly(length) if length else b"" + writer.close() + return json.loads(data) if data else None + + +async def collect(stream: Any) -> list[Any]: + return [event async for event in stream] + + +# -- the model protocol ------------------------------------------------------- + + +async def test_a_worker_is_model_like(worker: fakes.FakeWorker) -> None: + assert isinstance(worker, ai.ModelLike) + assert worker.provider_name == "fake" + assert worker.id == "default" + async with worker: + assert worker.version == "fake 1.0" + async with ai.stream(worker, [ai.user_message("hello")]) as s: + seen = await collect(s) + assert isinstance(seen[0], events.StreamStart) + assert isinstance(seen[-1], events.StreamEnd) + assert seen[-1].finish_reason == "stop" + assert s.text == "hello" + assert s.message.usage is not None and s.message.usage.input_tokens == 3 + meta = seen[-1].provider_metadata + assert meta is not None and meta["worker"]["session_id"] == "s1" + + +async def test_generate_drains_the_stream(worker: fakes.FakeWorker) -> None: + async with worker: + message = await ai.experimental_generate( + worker, [ai.user_message("hi")] + ) + assert message.text == "hi" + + +async def test_cli_tools_are_builtin(worker: fakes.FakeWorker) -> None: + worker.script( + lambda prompt: [ + *fakes.tool("c1", "Bash", '{"command": "ls"}', "a.py"), + *fakes.text("done"), + ] + ) + async with worker: + async with ai.stream(worker, [ai.user_message("go")]) as s: + kinds = [type(e).__name__ for e in await collect(s)] + assert "BuiltinToolStart" in kinds + assert "BuiltinToolResult" in kinds + assert "ToolStart" not in kinds + assert s.tool_calls == [] # nothing for an ai.Agent to run + call, result = s.message.parts[0], s.message.parts[1] + assert isinstance(call, m.BuiltinToolCallPart) + assert call.tool_name == "Bash" + assert call.tool_args == '{"command": "ls"}' + assert isinstance(result, m.BuiltinToolReturnPart) + assert result.result == "a.py" + + +async def test_a_cancelled_turn_ends_as_other(worker: fakes.FakeWorker) -> None: + worker.script(lambda prompt: [*fakes.text("a"), fakes.Pause()]) + async with worker: + stream = worker.stream([ai.user_message("go")]) + seen = [] + async for event in stream: + seen.append(event) + if isinstance(event, events.TextEnd): + await worker.fake.interrupt() + end = seen[-1] + assert isinstance(end, events.StreamEnd) + assert end.finish_reason == "other" + assert end.provider_metadata is not None + assert end.provider_metadata["worker"]["finish_reason"] == "cancelled" + + +# -- the conversation --------------------------------------------------------- + + +async def test_only_new_turns_reach_the_cli(worker: fakes.FakeWorker) -> None: + async with worker: + first = await ai.experimental_generate(worker, [ai.user_message("one")]) + assert worker.messages[-1].id == first.id + await ai.experimental_generate( + worker, [*worker.messages, ai.user_message("two")] + ) + assert worker.fake.turns == ["one", "two"] + assert len(worker.fake.opens) == 1 + assert [x.role for x in worker.messages] == [ + "user", + "assistant", + "user", + "assistant", + ] + + +async def test_a_new_worker_is_seeded_with_the_past( + worker: fakes.FakeWorker, +) -> None: + past = [ + ai.system_message("Be terse."), + ai.user_message("earlier"), + m.Message(role="assistant", parts=[m.TextPart(text="answer")]), + ] + async with worker: + await ai.experimental_generate(worker, [*past, ai.user_message("now")]) + opened = worker.fake.opens[0] + assert worker.fake.turns == ["now"] + # The system message is the CLI's instructions, not part of the seed. + assert opened["instructions"] == "Be terse." + assert [x.id for x in opened["seed"]] == [x.id for x in past[1:]] + assert opened["session_id"] is None + + +async def test_a_later_system_message_joins_the_prompt( + worker: fakes.FakeWorker, +) -> None: + async with worker: + await ai.experimental_generate(worker, [ai.user_message("one")]) + await ai.experimental_generate( + worker, + [ + *worker.messages, + ai.system_message("Now in French."), + ai.user_message("two"), + ], + ) + assert worker.fake.turns[1] == "Now in French.\n\n\n\ntwo" + + +async def test_a_rewritten_past_is_refused(worker: fakes.FakeWorker) -> None: + async with worker: + await ai.experimental_generate(worker, [ai.user_message("one")]) + with pytest.raises(workers.TranscriptMismatchError): + await ai.experimental_generate(worker, [ai.user_message("other")]) + + +async def test_a_failed_call_can_be_sent_again( + worker: fakes.FakeWorker, +) -> None: + def boom(prompt: str) -> list[fakes.Step]: + raise RuntimeError("boom") + + worker.script(boom) + messages = [ai.user_message("one")] + async with worker: + with pytest.raises(RuntimeError, match="boom"): + await ai.experimental_generate(worker, messages) + assert worker.messages == [] + worker.script(fakes.reply) + again = await ai.experimental_generate(worker, messages) + assert again.text == "one" + + +async def test_walking_away_interrupts_the_turn( + worker: fakes.FakeWorker, +) -> None: + worker.script(lambda prompt: [*fakes.text("a"), fakes.Pause()]) + async with worker: + async with ai.stream(worker, [ai.user_message("go")]) as s: + async for event in s: + if isinstance(event, events.TextDelta): + break + assert worker.fake.interrupts == 1 + + +async def test_structured_output(worker: fakes.FakeWorker) -> None: + class Finding(pydantic.BaseModel): + file: str + line: int + + worker.script( + lambda prompt: [ + *fakes.text("Let me look."), + *fakes.tool("c1", "Read"), + *fakes.text('```json\n{"file": "a.py", "line": 3}\n```', block="x"), + ] + ) + async with worker: + async with ai.stream( + worker, [ai.user_message("scan")], output_type=Finding + ) as s: + await collect(s) + assert "JSON value matching this schema" in worker.fake.turns[0] + assert s.output == Finding(file="a.py", line=3) + + +async def test_structured_output_asks_again(worker: fakes.FakeWorker) -> None: + class Count(pydantic.BaseModel): + n: int + + worker.script( + lambda prompt: fakes.text( + '{"n": 2}' if "failed validation" in prompt else "two" + ) + ) + async with worker: + async with ai.stream( + worker, [ai.user_message("how many?")], output_type=Count + ) as s: + await collect(s) + assert len(worker.fake.turns) == 2 + assert s.output == Count(n=2) + + +# -- a worker is data --------------------------------------------------------- + + +async def test_a_dumped_worker_continues_its_conversation( + worker: fakes.FakeWorker, +) -> None: + async with worker: + await ai.experimental_generate(worker, [ai.user_message("one")]) + data = json.loads(worker.model_dump_json()) + assert data["session_id"] == "s1" + again = workers.Worker.model_validate(data) + assert isinstance(again, fakes.FakeWorker) + assert [x.id for x in again.messages] == [x.id for x in worker.messages] + async with again: + await ai.experimental_generate( + again, [*again.messages, ai.user_message("two")] + ) + assert again.fake.opens[0]["session_id"] == "s1" # resumed + assert again.fake.turns == ["two"] + + +async def test_an_unknown_kind_is_refused() -> None: + with pytest.raises(pydantic.ValidationError): + workers.Worker.model_validate({"kind": "nope"}) + + +async def test_resume_by_session_id() -> None: + worker = fakes.FakeWorker(session_id="old") + async with worker: + # The CLI holds the past: what is given is all new. + await ai.experimental_generate(worker, [ai.user_message("go on")]) + assert worker.fake.opens[0]["session_id"] == "old" + assert worker.fake.opens[0]["seed"] is None + assert worker.fake.turns == ["go on"] + assert worker.session_id == "old" + + +async def test_load_reads_the_stored_past() -> None: + worker = fakes.FakeWorker(session_id="old") + stored = [ + ai.user_message("earlier"), + m.Message(role="assistant", parts=[m.TextPart(text="answer")]), + ] + async with worker: + worker.fake.stored["old"] = stored + await worker.load() + assert worker.messages == stored + await ai.experimental_generate( + worker, [*worker.messages, ai.user_message("now")] + ) + assert worker.fake.turns == ["now"] + + +async def test_a_copy_with_fork_branches_natively( + worker: fakes.FakeWorker, +) -> None: + async with worker: + await ai.experimental_generate(worker, [ai.user_message("one")]) + branch = worker.model_copy(update={"fork": True}) + async with branch: + assert branch.fake is not worker.fake # a CLI of its own + await ai.experimental_generate( + branch, [*branch.messages, ai.user_message("two")] + ) + assert branch.fake.opens[0] == { + "session_id": "s1", + "fork": True, + "seed": None, + "instructions": None, + } + assert branch.session_id == "s1-fork" + assert branch.fork is False + assert worker.session_id == "s1" + assert len(worker.messages) == 2 + assert len(branch.messages) == 4 + + +async def test_a_worker_goes_in_a_context(worker: fakes.FakeWorker) -> None: + context = ai.Context(model=worker, messages=[], tools=[]) + data = json.loads(context.model_dump_json()) + assert data["model"]["kind"] == "fake" + restored = ai.Context.model_validate(data) + assert isinstance(restored.model, fakes.FakeWorker) + + +async def test_close_keeps_the_fields(worker: fakes.FakeWorker) -> None: + async with worker: + await ai.experimental_generate(worker, [ai.user_message("one")]) + fake = worker.fake + assert fake.closed + assert worker.session_id == "s1" + async with worker: # opens again, and resumes + await ai.experimental_generate( + worker, [*worker.messages, ai.user_message("two")] + ) + assert worker.fake.opens[0]["session_id"] == "s1" + + +# -- ai.Agent tools, through the bridge --------------------------------------- + + +async def test_agent_tools_run_in_the_agent(worker: fakes.FakeWorker) -> None: + ran: list[tuple[int, int]] = [] + + @ai.tool + async def add(a: int, b: int) -> int: + """Add two numbers.""" + ran.append((a, b)) + return a + b + + answers: list[Any] = [] + + def script(prompt: str) -> Iterator[fakes.Step]: + url = worker.servers[_bridge.SERVER_NAME]["url"] + + async def call() -> None: + listed = await post( + url, {"jsonrpc": "2.0", "id": 1, "method": "tools/list"} + ) + assert [t["name"] for t in listed["result"]["tools"]] == ["add"] + answers.append( + await post( + url, + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "add", + "arguments": {"a": 1, "b": 2}, + }, + }, + ) + ) + + # What the CLI reports about its own MCP call: hidden from the + # stream, because the agent reports the call itself. + yield events.BuiltinToolStart( + tool_call_id="x1", tool_name="mcp__ai_sdk__add" + ) + yield fakes.Run(call) + yield from fakes.tool("x1", "mcp__ai_sdk__add")[1:] + text = answers[0]["result"]["content"][0]["text"] + yield from fakes.text(f"the sum is {text}") + + worker.script(script) + agent = ai.Agent(tools=[add]) + async with worker: + async with agent.run(worker, [ai.user_message("add 1 and 2")]) as run: + seen = await collect(run) + assert ran == [(1, 2)] + assert run.messages[-1].text == "the sum is 3" + names = [ + e.tool_name for e in seen if isinstance(e, events.BuiltinToolStart) + ] + assert names == [] # the CLI's own report of the MCP call is hidden + assert [x.role for x in run.messages] == [ + "user", + "assistant", + "tool", + "assistant", + ] + + +async def test_bridge_tool_names(worker: fakes.FakeWorker) -> None: + @ai.tool + async def lookup(q: str) -> str: + """Look up.""" + return q + + async with worker: + assert worker._runtime is not None + worker._runtime.bridge.set_tools([lookup.tool]) + assert worker._is_bridge_tool("lookup") # codex, with a tool name + assert worker._is_bridge_tool("ai_sdk") # codex, without one + assert worker._is_bridge_tool("mcp__ai_sdk__lookup") # claude code + assert not worker._is_bridge_tool("Bash") + + +# -- approvals, through ai.hook ----------------------------------------------- + + +@pytest.mark.parametrize("granted", [True, False]) +async def test_approvals_are_tool_approval_hooks(granted: bool) -> None: + worker = fakes.FakeWorker(approve=True) + worker.script( + lambda prompt: [ + fakes.Ask("c1", "Bash", {"command": "rm -rf build"}), + *fakes.text("done"), + ] + ) + pending: list[m.HookPart[Any]] = [] + async with worker: + async with ai.Agent().run(worker, [ai.user_message("clean")]) as run: + async for event in run: + if isinstance(event, events.PendingHookEvent): + pending.append(event.hook) + ai.resolve_hook( + event.hook, + {"granted": granted, "reason": "not today"}, + ) + answers = worker.fake.answers + assert [h.hook_type for h in pending] == ["ToolApproval"] + assert pending[0].tool_call_id == "c1" + assert pending[0].metadata == { + "tool": "Bash", + "kwargs": {"command": "rm -rf build"}, + } + assert [granted for granted, _ in answers] == [granted] + if not granted: + assert answers[0][1] == "not today" # the reason the agent reads + + +async def test_approve_needs_a_run() -> None: + worker = fakes.FakeWorker(approve=True) + async with worker: + with pytest.raises(workers.WorkerError, match=r"ai\.Agent\.run"): + await ai.experimental_generate(worker, [ai.user_message("go")]) + + +# -- steer, the TUI, the store ------------------------------------------------ + + +async def test_steer_reaches_the_running_turn( + worker: fakes.FakeWorker, +) -> None: + worker.script(lambda prompt: [*fakes.text("working"), fakes.Pause()]) + async with worker: + with pytest.raises(workers.WorkerError): + await worker.steer("too early") + async with ai.stream(worker, [ai.user_message("go")]) as s: + async for event in s: + if isinstance(event, events.TextEnd): + await worker.steer("use tabs") + await worker.fake.interrupt() + assert worker.fake.steers == ["use tabs"] + + +async def test_tui_command_mints_a_session(worker: fakes.FakeWorker) -> None: + argv, env = worker.tui_command() + assert env == {"FAKE_TUI": "1"} + assert argv[0] + assert worker.session_id == "minted" + + +async def test_tui_hands_over_and_reads_back() -> None: + worker = fakes.FakeWorker(session_id="s1") + stored = [ + ai.user_message("typed by a person"), + m.Message(role="assistant", parts=[m.TextPart(text="ok")]), + ] + async with worker: + assert await worker.tui() == 0 + # tui() closed the worker's CLI, ran the TUI, and opened again. + worker.fake.stored["s1"] = stored + await worker.load() + assert worker.messages == stored + + +async def test_list_sessions() -> None: + sessions = await fakes.FakeWorker.list_sessions(cwd="/tmp") + assert [s.session_id for s in sessions] == ["s1"] + + +# -- the bridge alone --------------------------------------------------------- + + +async def test_bridge_speaks_mcp() -> None: + async def on_call(name: str, arguments: dict[str, Any]) -> tuple[str, bool]: + return (f"{name}:{arguments['q']}", False) + + @ai.tool + async def search(q: str) -> str: + """Search.""" + return q + + bridge = _bridge.Bridge(on_call) + await bridge.start() + try: + bridge.set_tools([search.tool]) + init = await post( + bridge.url, + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2025-03-26"}, + }, + ) + assert init["result"]["protocolVersion"] == "2025-03-26" + assert ( + await post( + bridge.url, + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + is None + ) + listed = await post( + bridge.url, {"jsonrpc": "2.0", "id": 2, "method": "tools/list"} + ) + tool = listed["result"]["tools"][0] + assert tool["name"] == "search" + assert tool["inputSchema"]["properties"]["q"]["type"] == "string" + called = await post( + bridge.url, + { + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": {"name": "search", "arguments": {"q": "x"}}, + }, + ) + assert called["result"]["content"][0]["text"] == "search:x" + unknown = await post( + bridge.url, {"jsonrpc": "2.0", "id": 4, "method": "nope"} + ) + assert unknown["error"]["code"] == -32601 + finally: + await bridge.close()