diff --git a/agent-flow/README.md b/agent-flow/README.md index e7b8e5b5ada9..df718e195c34 100644 --- a/agent-flow/README.md +++ b/agent-flow/README.md @@ -2,7 +2,7 @@ `agent-flow` is a torch-like framework for composing agent-backed layers. Each layer implements `forward`, and each `AgentLayer.forward` performs one -Claude Code or Codex SDK execution. Modules can be composed directly or used +logical Claude Code or Codex invocation. Modules can be composed directly or used to build small orchestration harnesses. ## Prerequisites @@ -66,6 +66,70 @@ Run it directly: python examples/quick_start.py ``` +## Backend behavior + +Define custom tools with `from agent_flow import tool`. Both backends accept +these definitions and existing Claude SDK tool objects. Python shorthand and +`TypedDict` inputs are normalized to JSON Schema. Explicit JSON Schemas are +normalized to the root object shape MCP tools require (a string `type`, a +`properties` key, and a root `$ref` moved into `allOf`) so both backends +advertise the same schema. Resource links and embedded text resources become +readable tool output; image results retain their MIME type. Unsupported result +types return explicit tool errors. + +Declare required calls on the layer, independently of its backend: + +```python +config = AgentLayerConfig( + backend=BackendConfig(kind="codex", model=CODEX_DEFAULT_MODEL, tools=tools), + required_tools=("append_reviewer_progress", "update_status"), +) +``` + +Every listed tool must be attempted by the root agent during that invocation. +If calls are missing, the framework sends one corrective SDK turn on the same +session, including for stateless layers. It requests only the remaining calls. +If they are still missing, `RequiredToolCallError` fails the invocation. This +policy counts attempts, not successful writes; tool handlers and workflow +validation remain responsible for the outputs. Token accounting includes both +turns, while context occupancy remains the final snapshot. + +`BackendConfig.extra_mcp_servers` adds STDIO or HTTP MCP servers to one session. +Codex translates these into thread configuration without editing global config. +`AgentLayerConfig.disallowed_tools` removes matching framework tools and maps +supported native controls, including shell, web search, image tools, and +`mcp__SERVER__TOOL` exclusions. Codex rejects names it cannot enforce, such as a +generic `Write` ban. Tool exclusions are distinct from filesystem permissions. + +Native hooks remain backend-specific. Claude accepts its SDK callback objects. +For Codex, configure and trust native hooks through Codex itself; per-client +hook injection is rejected because its trust cannot be verified. Prefer +`required_tools` for workflow progress/status requirements on either backend. +Claude receives tool annotations; Codex dynamic tools have no equivalent native +annotation fields. + +Codex reports subagent activity, structured rate limits, cache-write tokens, and +current context usage after a turn. `UsageInfo.estimated_thread_cost_usd` is an +optional cumulative estimate, separate from reported `cost_usd`. Codex has no +pre-input context query or resolved available-agent inventory API; inspection +does not send a `ping` or another model turn to approximate either. + +The adapters target the SDK versions pinned in `pyproject.toml`. The Codex +transport uses the SDK's public process/request APIs and one isolated +notification observer to include child turns and account events. This observer +is covered by transport tests against the pinned SDK. + +## Tests + +```bash +pip install -e '.[test]' +python -m pytest tests +pre-commit run -a +``` + +Live Claude context-budget checks require local account configuration and are +opt-in: `python -m pytest tests --run-live-backends -m live_backend`. + ## Workflows Ready-to-run multi-agent workflows ship as subpackages of diff --git a/agent-flow/agent_flow/__init__.py b/agent-flow/agent_flow/__init__.py index 5157132fcf1b..be0a192e3bac 100644 --- a/agent-flow/agent_flow/__init__.py +++ b/agent-flow/agent_flow/__init__.py @@ -5,9 +5,14 @@ BackendConfig, SessionConfig, ) -from .hooks import called_required_tool_this_turn, require_tool_call_stop_hook +from .hooks import ( + RequiredToolCallError, + called_required_tool_this_turn, + require_tool_call_stop_hook, +) from .layers import AgentLayer from .module import Module, Sequential +from .tools import ToolDefinition, tool from .types import ( AgentRequest, AgentResponse, @@ -33,13 +38,16 @@ "CompactBoundaryEvent", "Module", "RateLimitWarningEvent", + "RequiredToolCallError", "Sequential", "ServerToolCallEvent", "SessionConfig", "SessionInitEvent", "ThinkingEvent", "ToolCallEvent", + "ToolDefinition", "UsageInfo", "called_required_tool_this_turn", "require_tool_call_stop_hook", + "tool", ] diff --git a/agent-flow/agent_flow/backends/claude_code.py b/agent-flow/agent_flow/backends/claude_code.py index f62bcda9269b..4063384b5410 100644 --- a/agent-flow/agent_flow/backends/claude_code.py +++ b/agent-flow/agent_flow/backends/claude_code.py @@ -12,7 +12,8 @@ from claude_agent_sdk import ( ClaudeAgentOptions, ClaudeSDKClient, - PermissionResultAllow, + SdkMcpTool, + ToolAnnotations, create_sdk_mcp_server, ) from claude_agent_sdk.types import ( @@ -28,6 +29,7 @@ ToolUseBlock, ) +from ..tools import normalize_tool, tool_result_to_claude from ..types import ( AgentTextEvent, CompactBoundaryEvent, @@ -420,6 +422,29 @@ def _claude_backend_version() -> str: _REASONING_EFFORT = "max" +def _sdk_tools(tools: list) -> list[SdkMcpTool]: + """Translate framework or legacy tool definitions at the SDK boundary.""" + result = [] + for original in tools: + definition = normalize_tool(original) + + async def handler(arguments, definition=definition): + return tool_result_to_claude(await definition.handler(arguments)) + + result.append( + SdkMcpTool( + name=definition.name, + description=definition.description, + input_schema=definition.input_schema, + handler=handler, + annotations=ToolAnnotations(**definition.annotations) + if definition.annotations + else None, + ) + ) + return result + + class ClaudeCodeBackend(Backend): def version(self) -> str: return _claude_backend_version() @@ -453,12 +478,9 @@ async def create_client( if tools: mcp_servers["agent-tools"] = create_sdk_mcp_server( name="agent-tools", - tools=tools, + tools=_sdk_tools(tools), ) - async def _approve_tool(tool_name, tool_input, context): - return PermissionResultAllow() - options = ClaudeAgentOptions( tools=ToolsPreset(type="preset", preset="claude_code"), system_prompt=SystemPromptPreset( @@ -472,7 +494,6 @@ async def _approve_tool(tool_name, tool_input, context): cwd=cwd or Path.cwd(), sandbox={"enabled": False}, permission_mode="bypassPermissions", - can_use_tool=_approve_tool, hooks=hooks, disallowed_tools=list(disallowed_tools or []), ) diff --git a/agent-flow/agent_flow/backends/codex.py b/agent-flow/agent_flow/backends/codex.py index 06f5d4cb1d13..14a7f1e44a64 100644 --- a/agent-flow/agent_flow/backends/codex.py +++ b/agent-flow/agent_flow/backends/codex.py @@ -1,58 +1,30 @@ from __future__ import annotations import asyncio -import importlib -import importlib.util -import inspect import os import shutil import subprocess -import sys from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError from importlib.metadata import version as _pkg_version from pathlib import Path -from typing import Any, AsyncIterator, Awaitable, Callable - -from ..types import ( - AgentTextEvent, - CompactBoundaryEvent, - RateLimitWarningEvent, - ServerToolCallEvent, - SessionInitEvent, - ThinkingEvent, - ToolCallEvent, - UsageInfo, +from typing import Any, AsyncIterator + +from openai_codex.client import CodexConfig +from openai_codex.generated.v2_all import ( + AskForApproval, + AskForApprovalValue, + FunctionDynamicToolSpec, + SandboxMode, + ThreadStartParams, ) -from .base import Backend, BackendClient, BackendEvent, ResultEvent - -# Dispatch table keyed by thread id. Populated when a CodexClient is created -# with dynamic tools, drained when it is torn down. The AppServerClient's -# approval handler reads this to route ``item/tool/call`` server requests to -# the right handler without having to know about threads. -_TOOL_HANDLERS: dict[str, dict[str, Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]]] = {} - - -def _module_or_none(name: str) -> Any | None: - module = sys.modules.get(name) - if module is not None: - return module - if name in sys.modules: - return None - if importlib.util.find_spec(name) is None: - return None - return importlib.import_module(name) - -def _module(name: str) -> Any: - module = sys.modules.get(name) - if module is not None: - return module - return importlib.import_module(name) - - -def _symbol(module_name: str, symbol_name: str) -> Any: - return getattr(_module(module_name), symbol_name) +from ..tools import normalize_tool +from ..types import SessionInitEvent +from .base import Backend, BackendClient +from .codex_config import build_session_config +from .codex_events import CodexClient +from .codex_transport import CodexTransport def _resolve_codex_bin() -> str: @@ -61,631 +33,56 @@ def _resolve_codex_bin() -> str: if not Path(override).is_file(): raise FileNotFoundError(f"CODEX_BIN={override!r} does not point to a file.") return override + try: + from codex_cli_bin import bundled_codex_path + except ImportError: + system_bin = shutil.which("codex") + if system_bin: + return system_bin + raise FileNotFoundError("Install the pinned Codex CLI runtime or set CODEX_BIN.") from None + return str(bundled_codex_path()) - codex_cli_bin = _module_or_none("codex_cli_bin") - bundled_codex_path = ( - getattr(codex_cli_bin, "bundled_codex_path", None) if codex_cli_bin is not None else None - ) - if bundled_codex_path is not None: - return str(bundled_codex_path()) - - system_bin = shutil.which("codex") - if system_bin is not None: - return system_bin - - raise FileNotFoundError( - "Could not locate a Codex CLI binary. Install the pinned runtime " - "(the 'openai-codex-cli-bin' package shipped with 'openai-codex'), " - "install the Codex CLI on PATH, or set CODEX_BIN to the binary path." - ) - - -def _usage_from_token_breakdown(breakdown: Any) -> UsageInfo: - return UsageInfo( - input_tokens=getattr(breakdown, "input_tokens", None), - output_tokens=getattr(breakdown, "output_tokens", None), - cache_read_tokens=getattr(breakdown, "cached_input_tokens", None), - total_tokens=getattr(breakdown, "total_tokens", None), - ) - - -def _get(value: Any, *names: str, default: Any = None) -> Any: - for name in names: - if isinstance(value, dict) and name in value: - return value[name] - if not isinstance(value, dict) and hasattr(value, name): - return getattr(value, name) - return default - - -def _plain(value: Any) -> Any: - if value is None or isinstance(value, (str, int, float, bool)): - return value - if isinstance(value, dict): - return {str(k): _plain(v) for k, v in value.items()} - if isinstance(value, (list, tuple, set)): - return [_plain(v) for v in value] - model_dump = getattr(value, "model_dump", None) - if callable(model_dump): +def _dynamic_tool_spec(tool: Any) -> dict[str, Any]: + definition = normalize_tool(tool) + return FunctionDynamicToolSpec( + type="function", + name=definition.name, + description=definition.description, + input_schema=definition.input_schema, + defer_loading=False, + ).model_dump(by_alias=True, mode="json", exclude_none=True) + + +async def _session_init(transport: CodexTransport, cwd: Path) -> SessionInitEvent | None: + async def optional(method, params): try: - return model_dump(by_alias=True, exclude_none=True) - except TypeError: - return model_dump() - - if hasattr(value, "__dict__"): - return { - k: _plain(v) for k, v in vars(value).items() if not k.startswith("_") and v is not None - } - - enum_value = getattr(value, "value", None) - if enum_value is not None: - return enum_value - return str(value) - - -def _as_text(value: Any) -> str: - if value is None: - return "" - enum_value = getattr(value, "value", None) - if enum_value is not None: - value = enum_value - return str(value) - - -def _root(item: Any) -> Any: - return item.root if hasattr(item, "root") else item - - -def _item_id(item: Any, root: Any | None = None) -> str | None: - root = root if root is not None else _root(item) - value = _get(root, "id", default=_get(item, "id")) - return value if isinstance(value, str) else None - - -_THREAD_ITEM_KINDS = { - "AgentMessageThreadItem": "agentMessage", - "CollabAgentToolCallThreadItem": "collabAgentToolCall", - "CollabToolCallThreadItem": "collabAgentToolCall", - "CommandExecutionThreadItem": "commandExecution", - "ContextCompactionThreadItem": "contextCompaction", - "CompactedThreadItem": "compacted", - "DynamicToolCallThreadItem": "dynamicToolCall", - "FileChangeThreadItem": "fileChange", - "ImageViewThreadItem": "imageView", - "McpToolCallThreadItem": "mcpToolCall", - "PlanThreadItem": "plan", - "ReasoningThreadItem": "reasoning", - "WebSearchThreadItem": "webSearch", -} - - -def _thread_item_kind(root: Any) -> str: - raw = _get(root, "type", "kind") - if raw is not None: - return _as_text(raw) - return _THREAD_ITEM_KINDS.get(root.__class__.__name__, "") - - -def _notification_kind(payload: Any) -> str: - raw = _get(payload, "method", "type", "kind") - if raw is not None: - return _as_text(raw) - - name = payload.__class__.__name__ - if name == "ItemCompletedNotification": - return "item/completed" - if name == "ItemStartedNotification": - return "item/started" - if name == "ThreadTokenUsageUpdatedNotification": - return "thread/tokenUsage/updated" - if name == "TurnCompletedNotification": - return "turn/completed" - if "AgentMessage" in name and "Delta" in name: - return "item/agentMessage/delta" - if "Reasoning" in name and "Delta" in name: - return "item/reasoning/delta" - if name in {"WarningNotification", "ConfigWarningNotification"}: - return "warning" - if name == "ErrorNotification": - return "error" - return name - - -def _arguments(value: Any) -> dict[str, Any]: - return value if isinstance(value, dict) else {} - - -def _reasoning_text(root: Any) -> str: - parts: list[str] = [] - for field in ("summary", "content", "text"): - value = _get(root, field) - if isinstance(value, (list, tuple)): - parts.extend(str(part).strip() for part in value if str(part).strip()) - elif isinstance(value, str) and value.strip(): - parts.append(value.strip()) - return "\n".join(parts).strip() - - -def _collab_tool_input(root: Any) -> dict[str, Any]: - body: dict[str, Any] = {} - for source, target in ( - ("prompt", "prompt"), - ("model", "model"), - ("reasoning_effort", "reasoning_effort"), - ("reasoningEffort", "reasoning_effort"), - ("sender_thread_id", "sender_thread_id"), - ("senderThreadId", "sender_thread_id"), - ("receiver_thread_ids", "receiver_thread_ids"), - ("receiverThreadIds", "receiver_thread_ids"), - ("agents_states", "agents_states"), - ("agentsStates", "agents_states"), - ("status", "status"), - ): - value = _get(root, source) - if value is not None: - body[target] = _plain(value) - return body - - -def _tool_event_from_item(root: Any, item: Any) -> BackendEvent | None: - kind = _thread_item_kind(root) - tool_use_id = _item_id(item, root) - if kind == "commandExecution": - command = _get(root, "command", default="") - return ToolCallEvent(name="Bash", input={"command": command}, tool_use_id=tool_use_id) - - if kind in {"mcpToolCall", "dynamicToolCall"}: - name = _get(root, "tool", default="tool") - return ToolCallEvent( - name=_as_text(name), - input=_arguments(_get(root, "arguments")), - tool_use_id=tool_use_id, - ) - - if kind in {"collabAgentToolCall", "collabToolCall"}: - name = _get(root, "tool", default="collab_agent") - return ToolCallEvent( - name=_as_text(name), - input=_collab_tool_input(root), - tool_use_id=tool_use_id, - ) - - if kind == "fileChange": - body = {"changes": _plain(_get(root, "changes", default=[]))} - status = _get(root, "status") - if status is not None: - body["status"] = _as_text(status) - return ToolCallEvent(name="FileChange", input=body, tool_use_id=tool_use_id) - - if kind == "webSearch": - action = _get(root, "action") - if action is not None: - body = _plain(action) - if not isinstance(body, dict): - body = {"action": body} - else: - body = {} - query = _get(root, "query") - if isinstance(query, str): - body.setdefault("query", query) - return ServerToolCallEvent(name="web_search", input=body, tool_use_id=tool_use_id) - - if kind == "imageView": - path = _get(root, "path") - body = {"path": path} if isinstance(path, str) else {} - return ServerToolCallEvent(name="view_image", input=body, tool_use_id=tool_use_id) - - return None - - -def _status_is(root: Any, *statuses: str) -> bool: - status = _get(root, "status") - return _as_text(status) in statuses - - -def _collect_item_errors( - root: Any, item: Any, errors: list[str], permission_denials: list[Any] -) -> None: - kind = _thread_item_kind(root) - if kind not in { - "collabAgentToolCall", - "collabToolCall", - "commandExecution", - "dynamicToolCall", - "fileChange", - "mcpToolCall", - }: - return - - if _status_is(root, "declined"): - permission_denials.append( - { - "kind": kind, - "id": _item_id(item, root), - "item": _plain(root), - } - ) - return - - if _status_is(root, "failed"): - error = _get(root, "error") - if error is None: - errors.append(f"{kind} failed") - else: - errors.append(_as_text(_plain(error))) - - -def _compact_boundary_from_item(root: Any) -> CompactBoundaryEvent: - trigger = _get(root, "trigger") - if not isinstance(trigger, str): - trigger = None - pre_tokens = _get(root, "pre_tokens", "preTokens") - if not isinstance(pre_tokens, int): - pre_tokens = None - return CompactBoundaryEvent(trigger=trigger, pre_tokens=pre_tokens) - - -def _rate_limit_warning_from_error(error: Any) -> RateLimitWarningEvent | None: - info = _get(error, "codex_error_info", "codexErrorInfo") - message = _as_text(_get(error, "message", default=error)) - haystack = f"{_as_text(info)} {message}".lower() - if "usagelimit" not in haystack and "rate limit" not in haystack: - return None - return RateLimitWarningEvent(status="rejected", rate_limit_type=_as_text(info) or None) - - -def _is_final_answer_phase(phase: Any, final_answer: Any) -> bool: - value = _as_text(phase) - return value in {"final_answer", "finalAnswer"} or value == _as_text(final_answer) - - -def _extract_final_response(items: list[object]) -> str: - MessagePhase = _symbol("openai_codex.generated.v2_all", "MessagePhase") - - last_unknown: str | None = None - for item in reversed(items): - root = _root(item) - if _thread_item_kind(root) == "agentMessage": - if _is_final_answer_phase(_get(root, "phase"), MessagePhase.final_answer): - return root.text - if _get(root, "phase") is None and last_unknown is None: - last_unknown = root.text - return last_unknown or "" - - -def _turn_completed_successfully(turn: Any) -> bool: - status = _get(turn, "status") - return status is None or _as_text(status) == "completed" - - -def _turn_failure_message(turn: Any) -> str: - error = _get(turn, "error") - if error is not None: - return _as_text(_get(error, "message", default=error)) - status = _as_text(_get(turn, "status", default="unknown")) - return f"turn ended with status {status!r}" - - -def _int_or_none(value: Any) -> int | None: - if isinstance(value, bool): - return None - return value if isinstance(value, int) else None - - -def _number_or_none(value: Any) -> int | float | None: - if isinstance(value, bool): - return None - return value if isinstance(value, (int, float)) else None - - -def _turn_duration_ms(turn: Any) -> int | None: - duration_ms = _int_or_none(_get(turn, "duration_ms", "durationMs")) - if duration_ms is not None: - return duration_ms - - started_at = _number_or_none(_get(turn, "started_at", "startedAt")) - completed_at = _number_or_none(_get(turn, "completed_at", "completedAt")) - if started_at is None or completed_at is None or completed_at < started_at: - return None - return int((completed_at - started_at) * 1000) - - -def _turn_num_turns(turn: Any) -> int: - num_turns = _int_or_none(_get(turn, "num_turns", "numTurns", "turn_count", "turnCount")) - return num_turns if num_turns is not None else 1 - - -def _usage_with_turn_metadata(usage: UsageInfo | None, turn: Any) -> UsageInfo | None: - duration_ms = _turn_duration_ms(turn) - if usage is None: - usage = UsageInfo() - usage.num_turns = _turn_num_turns(turn) - if duration_ms is None: - return usage - usage.duration_ms = duration_ms - return usage - - -def _usage_from_thread_token_usage(token_usage: Any) -> UsageInfo | None: - total = _get(token_usage, "total") - if total is None: - return None - - usage = _usage_from_token_breakdown(total) - context_tokens = _int_or_none(_get(total, "total_tokens", "totalTokens")) - context_window = _int_or_none(_get(token_usage, "model_context_window", "modelContextWindow")) - - usage.context_tokens = context_tokens - usage.context_window = context_window - if context_tokens is not None and context_window: - usage.context_percentage = 100.0 * context_tokens / context_window - return usage - - -class CodexClient(BackendClient): - def __init__(self, thread, session_init: SessionInitEvent | None = None) -> None: - self._thread = thread - self._session_init = session_init - self._session_init_emitted = False - - async def list_available_skills(self) -> list[str] | None: - """Skill names from the session init built at client creation. - - ``_build_session_init_event`` already asked the Codex client for - ``skills_list`` when this client was created, so the answer is in - hand before any turn — the event that ``send_message`` replays is - only the same cached value. ``None`` when that call came back - empty or failed, which is the backend declining to say rather - than reporting an empty environment. - """ - if self._session_init is None: + return await asyncio.wait_for(transport.request(method, params), timeout=5) + except Exception: return None - return list(self._session_init.skills) - - async def send_message(self, message: str) -> AsyncIterator[BackendEvent]: - TextInput = _symbol("openai_codex", "TextInput") - - turn = await self._thread.turn(TextInput(message)) - items: list[object] = [] - latest_usage: UsageInfo | None = None - errors: list[str] = [] - permission_denials: list[Any] = [] - emitted_tool_items: set[str] = set() - emitted_compaction_items: set[str] = set() - - if self._session_init is not None and not self._session_init_emitted: - self._session_init_emitted = True - yield self._session_init - - async for event in turn.stream(): - payload = event.payload - kind = _notification_kind(payload) - if kind in {"item/started", "item/completed"}: - item = _get(payload, "item") - if item is None: - continue - if kind == "item/completed": - items.append(item) - root = _root(item) - item_id = _item_id(item, root) - - tool_event = _tool_event_from_item(root, item) - if tool_event is not None and item_id not in emitted_tool_items: - if item_id is not None: - emitted_tool_items.add(item_id) - yield tool_event - - if ( - _thread_item_kind(root) in {"contextCompaction", "compacted"} - and item_id not in emitted_compaction_items - ): - if item_id is not None: - emitted_compaction_items.add(item_id) - yield _compact_boundary_from_item(root) - - if kind != "item/completed": - continue - - _collect_item_errors(root, item, errors, permission_denials) - if _thread_item_kind(root) == "reasoning": - text = _reasoning_text(root) - if text: - yield ThinkingEvent(text=text) - elif _thread_item_kind(root) == "agentMessage": - text = (root.text or "").strip() - if text: - yield AgentTextEvent(text=text) - elif kind == "item/agentMessage/delta": - continue - elif kind == "item/reasoning/delta": - continue - elif kind == "thread/tokenUsage/updated": - usage = _usage_from_thread_token_usage(payload.token_usage) - if usage is not None: - latest_usage = usage - elif kind == "warning": - warning = _rate_limit_warning_from_error(_get(payload, "message", default=payload)) - if warning is not None: - yield warning - elif kind == "error": - warning = _rate_limit_warning_from_error(_get(payload, "error", default=payload)) - if warning is not None: - yield warning - elif kind == "turn/completed" and payload.turn.id == turn.id: - turn_error = getattr(payload.turn, "error", None) - if turn_error is not None: - warning = _rate_limit_warning_from_error(turn_error) - if warning is not None: - yield warning - message = getattr(turn_error, "message", None) or str(turn_error) - raise RuntimeError(f"Codex turn failed: {message}") - if not _turn_completed_successfully(payload.turn): - raise RuntimeError(f"Codex turn failed: {_turn_failure_message(payload.turn)}") - yield ResultEvent( - text=_extract_final_response(items), - usage=_usage_with_turn_metadata(latest_usage, payload.turn), - is_error=bool(errors), - errors=errors, - permission_denials=permission_denials, - ) - - -def _dynamic_tool_spec(tool: Any) -> dict[str, Any]: - """Serialize an ``SdkMcpTool`` into a Codex ``DynamicToolSpec`` dict.""" - return { - "name": tool.name, - "description": tool.description, - "inputSchema": tool.input_schema, - "deferLoading": False, - } - - -def _mcp_to_codex_content(result: Any) -> tuple[list[dict[str, Any]], bool]: - """Translate an MCP tool result to Codex ``DynamicToolCallResponse`` parts. - - Tools decorated with ``claude_agent_sdk.tool`` return ``{"content": [{type: - "text", text: "..."}], "is_error"?: bool}``. Codex expects content items - shaped as ``{type: "inputText", text: "..."}`` plus a top-level ``success`` - flag. - """ - items: list[dict[str, Any]] = [] - success = True - if isinstance(result, dict): - for item in result.get("content") or []: - if not isinstance(item, dict): - continue - kind = item.get("type") - if kind == "text": - items.append({"type": "inputText", "text": item.get("text", "")}) - elif kind == "image": - url = item.get("data") or item.get("image_url") or "" - items.append({"type": "inputImage", "imageUrl": url}) - if result.get("is_error"): - success = False - if not items: - items = [{"type": "inputText", "text": ""}] - return items, success - - -def _handle_dynamic_tool_call(params: dict[str, Any]) -> dict[str, Any]: - thread_id = params.get("threadId") or "" - tool_name = params.get("tool") or "" - arguments = params.get("arguments") - if not isinstance(arguments, dict): - arguments = {} - - handlers = _TOOL_HANDLERS.get(thread_id, {}) - handler = handlers.get(tool_name) - if handler is None: - return { - "contentItems": [ - { - "type": "inputText", - "text": (f"Tool {tool_name!r} is not registered for this thread."), - } - ], - "success": False, - } - - try: - result = asyncio.run(handler(arguments)) - except Exception as exc: # noqa: BLE001 — surface any tool error to model - return { - "contentItems": [ - { - "type": "inputText", - "text": f"Tool {tool_name!r} failed: {exc}", - } - ], - "success": False, - } - - items, success = _mcp_to_codex_content(result) - return {"contentItems": items, "success": success} - - -async def _call_optional_client_method( - client: Any, names: tuple[str, ...], payload: dict[str, Any] -) -> Any: - for name in names: - method = getattr(client, name, None) - if method is None: - continue - result = method(payload) - if inspect.isawaitable(result): - return await result - return result - return None - - -def _as_list(value: Any) -> list[Any]: - if value is None: - return [] - if isinstance(value, list): - return value - if isinstance(value, tuple): - return list(value) - return [value] - - -def _extract_codex_skill_names(response: Any) -> list[str]: - names: list[str] = [] - for entry in _as_list(_get(response, "data", default=response)): - for skill in _as_list(_get(entry, "skills")): - if _get(skill, "enabled") is False: - continue - name = _get(skill, "name") - if isinstance(name, str): - names.append(name) - return names - - -def _extract_codex_plugin_names(response: Any) -> list[str]: - names: list[str] = [] - - def walk(value: Any, depth: int = 0) -> None: - if value is None or depth > 4: - return - if isinstance(value, (list, tuple)): - for item in value: - walk(item, depth + 1) - return - - name = _get(value, "name", "plugin_name", "pluginName") - if isinstance(name, str): - names.append(name) - - for key in ("plugins", "items", "entries", "data"): - child = _get(value, key) - if child is not None: - walk(child, depth + 1) - walk(response) - return names - - -async def _build_session_init_event(client: Any, cwd: Path) -> SessionInitEvent | None: - try: - skills_response = await _call_optional_client_method( - client, - ("skills_list", "skill_list"), - {"cwds": [str(cwd)], "forceReload": False}, - ) - plugins_response = await _call_optional_client_method( - client, ("plugin_list", "plugins_list"), {} - ) - except Exception: - return None - - event = SessionInitEvent( - skills=_extract_codex_skill_names(skills_response), - plugins=_extract_codex_plugin_names(plugins_response), - agents=[], + skills_response, plugins_response = await asyncio.gather( + optional("skills/list", {"cwds": [str(cwd)], "forceReload": False}), + optional("plugin/list", {}), ) - if event.skills or event.plugins or event.agents: - return event - return None + skills = [] + if skills_response is not None: + for entry in skills_response.get("data", []): + for skill in entry.get("skills", []): + if skill.get("enabled") is not False and isinstance(skill.get("name"), str): + skills.append(skill["name"]) + plugins = [] + if plugins_response is not None: + for marketplace in plugins_response.get("marketplaces", []): + for plugin in marketplace.get("plugins", []): + if plugin.get("installed") and plugin.get("enabled") is not False: + if isinstance(plugin.get("name"), str): + plugins.append(plugin["name"]) + if skills_response is None and not plugins: + return None + # The protocol has no resolved available-agent inventory API. Keep this + # unknown instead of claiming every file under .codex/agents is loaded. + return SessionInitEvent(skills=skills, plugins=plugins) _CLI_VERSION_TIMEOUT_S = 5.0 @@ -743,90 +140,10 @@ def _codex_backend_version() -> str: _REASONING_EFFORT = "max" -_SDK_PATCHED = False - - -def _relax_service_tier_on_module(module: Any) -> int: - """Rewrite ``ServiceTier`` references to ``str`` on every pydantic model in ``module``. - - Returns the number of models that were rebuilt. - """ - OldEnum = getattr(module, "ServiceTier", None) - if OldEnum is None: - return 0 - - import types - from typing import Union, get_args, get_origin - - from pydantic import BaseModel - - def _replace(annotation: Any) -> Any: - if annotation is OldEnum: - return str - origin = get_origin(annotation) - if origin is None: - return annotation - args = get_args(annotation) - new_args = tuple(_replace(a) for a in args) - if new_args == args: - return annotation - # ``X | Y`` produces ``types.UnionType``; ``Union[X, Y]`` produces - # ``typing.Union``. Both must be rebuilt as a typing.Union so that - # pydantic can re-derive the schema. - if origin is Union or origin is types.UnionType: - return Union[new_args] # type: ignore[valid-type] - try: - return origin[new_args] - except TypeError: - return annotation - - affected: list[type] = [] - for name in dir(module): - cls = getattr(module, name, None) - if not (isinstance(cls, type) and issubclass(cls, BaseModel)): - continue - changed = False - for finfo in cls.model_fields.values(): - new_ann = _replace(finfo.annotation) - if new_ann is not finfo.annotation: - finfo.annotation = new_ann - changed = True - if changed: - affected.append(cls) - - for cls in affected: - try: - cls.model_rebuild(force=True) - except Exception: - pass - return len(affected) - - -def _patch_codex_sdk_service_tier() -> None: - """Relax the SDK's ``ServiceTier`` enum to ``str`` on response models. - - The rust protocol (``codex-rs/app-server-protocol/src/protocol/v2/thread.rs``) - declares ``service_tier: Option`` on responses, but the Python SDK - was generated with a strict ``ServiceTier`` enum that only contains - ``fast`` and ``flex``. The server returns ``"priority"`` (the request value - for the Fast tier per ``ServiceTier::request_value`` in - ``codex-rs/protocol/src/config_types.rs``), which the strict enum rejects. - """ - global _SDK_PATCHED - if _SDK_PATCHED: - return - try: - v2_all = importlib.import_module("openai_codex.generated.v2_all") - except ImportError: - return - _relax_service_tier_on_module(v2_all) - _SDK_PATCHED = True - class CodexBackend(Backend): def __init__(self) -> None: - self._codex = None - self._default_approval_handler: Callable[..., dict[str, Any]] | None = None + self._transport: CodexTransport | None = None def version(self) -> str: return _codex_backend_version() @@ -835,49 +152,16 @@ def reasoning_effort(self) -> str: return _REASONING_EFFORT async def __aenter__(self) -> "CodexBackend": - _patch_codex_sdk_service_tier() - AsyncCodex = _symbol("openai_codex", "AsyncCodex") - CodexConfig = _symbol("openai_codex.client", "CodexConfig") - - codex_bin = _resolve_codex_bin() - - self._codex = AsyncCodex( - config=CodexConfig( - codex_bin=codex_bin, - experimental_api=True, - ) + self._transport = CodexTransport( + CodexConfig(codex_bin=_resolve_codex_bin(), experimental_api=True) ) - await self._codex.__aenter__() - self._install_server_request_handler() + await self._transport.start() return self async def __aexit__(self, *args: object) -> None: - if self._codex is not None: - await self._codex.__aexit__(None, None, None) - self._codex = None - self._default_approval_handler = None - - def _install_server_request_handler(self) -> None: - """Route tool-call requests to the dynamic handler and auto-accept all approvals. - - ``item/tool/call`` requests are routed to ``_handle_dynamic_tool_call``. Combined with - ``approval_policy=never`` and ``sandbox=danger_full_access`` in ``create_client``, the - agent never prompts the user: any ``*/requestApproval`` JSON-RPC request is answered with - ``{"decision": "accept"}`` regardless of which item kind raised it (commandExecution, - fileChange, or anything new the SDK adds later). - """ - sync_client = self._codex._client._sync - existing = sync_client._approval_handler - self._default_approval_handler = existing - - def _handler(method: str, params: dict[str, Any] | None) -> dict[str, Any]: - if method == "item/tool/call": - return _handle_dynamic_tool_call(params or {}) - if method.endswith("/requestApproval"): - return {"decision": "accept"} - return existing(method, params) - - sync_client._approval_handler = _handler + if self._transport is not None: + await self._transport.close() + self._transport = None @asynccontextmanager async def create_client( @@ -890,47 +174,80 @@ async def create_client( extra_mcp_servers: dict[str, Any] | None = None, cwd: Path | None = None, ) -> AsyncIterator[BackendClient]: - # ``disallowed_tools`` and ``extra_mcp_servers`` are currently - # Claude-Code-only concepts (they map onto - # ``ClaudeAgentOptions.disallowed_tools`` and - # ``ClaudeAgentOptions.mcp_servers``). The Codex backend has no - # analogous tool-filtering / external-MCP-server hooks, so we - # accept and ignore them to keep ``Backend.create_client`` - # uniform. - if self._codex is None: + if self._transport is None: raise RuntimeError("CodexBackend must be entered before creating clients.") - - AsyncThread = _symbol("openai_codex.api", "AsyncThread") - _params_dict = _symbol("openai_codex.client", "_params_dict") - AskForApproval = _symbol("openai_codex.generated.v2_all", "AskForApproval") - AskForApprovalValue = _symbol("openai_codex.generated.v2_all", "AskForApprovalValue") - SandboxMode = _symbol("openai_codex.generated.v2_all", "SandboxMode") - ThreadStartParams = _symbol("openai_codex.generated.v2_all", "ThreadStartParams") - - session_cwd = cwd or Path.cwd() + session_cwd = (cwd or Path.cwd()).resolve() + config, definitions = build_session_config( + extra_mcp_servers=extra_mcp_servers, + hooks=hooks, + disallowed_tools=disallowed_tools, + tools=tools, + cwd=session_cwd, + ) + mcp_overrides = config.get("mcp_servers", {}) + if any("disabled_tools" in server for server in mcp_overrides.values()): + # Codex replaces arrays when merging thread overrides. Preserve + # exclusions inherited from the native config for this session's cwd. + # A failed read must prevent starting a thread with weakened bans. + native_config = await asyncio.wait_for( + self._transport.request( + "config/read", {"cwd": str(session_cwd), "includeLayers": True} + ), + timeout=5, + ) + inherited_configs = [native_config["config"]] + # Starting an unrestricted session can trust a project and activate + # its config. Conservatively retain those exclusions as well. + inherited_configs.extend( + layer["config"] + for layer in native_config.get("layers") or [] + if layer["name"]["type"] == "project" and layer.get("disabledReason") is not None + ) + for name, server in mcp_overrides.items(): + if "disabled_tools" not in server: + continue + inherited = [] + for inherited_config in inherited_configs: + native_servers = inherited_config.get("mcp_servers") or {} + excluded = native_servers.get(name, {}).get("disabled_tools") + if excluded is None: + continue + if not isinstance(excluded, list) or not all( + isinstance(tool_name, str) for tool_name in excluded + ): + raise ValueError( + f"Native MCP disabled_tools for {name!r} must be a list of strings" + ) + inherited.extend(excluded) + server["disabled_tools"] = list( + dict.fromkeys([*inherited, *server["disabled_tools"]]) + ) + config.update(model_reasoning_effort=_REASONING_EFFORT, model_context_window=1000000) params = ThreadStartParams( model=model, developer_instructions=system_prompt or None, - config={ - "model_reasoning_effort": _REASONING_EFFORT, - "model_context_window": 1000000, - }, + config=config, cwd=str(session_cwd), sandbox=SandboxMode.danger_full_access, approval_policy=AskForApproval(root=AskForApprovalValue.never), ) - start_payload = _params_dict(params) - if tools: - start_payload["dynamicTools"] = [_dynamic_tool_spec(t) for t in tools] - - await self._codex._ensure_initialized() - session_init = await _build_session_init_event(self._codex._client, session_cwd) - started = await self._codex._client.thread_start(start_payload) - thread_id = started.thread.id - - if tools: - _TOOL_HANDLERS[thread_id] = {t.name: t.handler for t in tools} + payload = params.model_dump(by_alias=True, mode="json", exclude_none=True) + if definitions: + payload["dynamicTools"] = [_dynamic_tool_spec(tool) for tool in definitions] + session_init = await _session_init(self._transport, session_cwd) + started = await self._transport.request("thread/start", payload) + thread_id = started["thread"]["id"] + self._transport.register_tools(thread_id, definitions) try: - yield CodexClient(AsyncThread(self._codex, thread_id), session_init=session_init) + yield CodexClient(self._transport, thread_id, session_init=session_init) finally: - _TOOL_HANDLERS.pop(thread_id, None) + self._transport.unregister_tools(thread_id) + # Unload the thread, while retaining the saved session just as the + # native runtime does. No global MCP config is changed. + try: + await asyncio.wait_for( + self._transport.request("thread/unsubscribe", {"threadId": thread_id}), + timeout=5, + ) + except Exception: + pass diff --git a/agent-flow/agent_flow/backends/codex_config.py b/agent-flow/agent_flow/backends/codex_config.py new file mode 100644 index 000000000000..84f9196c327e --- /dev/null +++ b/agent-flow/agent_flow/backends/codex_config.py @@ -0,0 +1,292 @@ +"""Translate session options to the configuration supported by Codex 0.154.0. + +These overrides belong to ``thread/start.config``. They never edit the user's +Codex configuration. Unsupported options fail at registration rather than being +silently ignored, especially when they describe tool restrictions. +""" + +from __future__ import annotations + +import copy +import json +import math +from collections.abc import Mapping +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit + +_MCP_COMMON = { + "enabled", + "required", + "enabled_tools", + "disabled_tools", + "startup_timeout_sec", + "startup_timeout_ms", + "tool_timeout_sec", + "supports_parallel_tool_calls", + "default_tools_approval_mode", + "tools", + "omit_tools_from", + "environment_id", + "name", +} +_MCP_STDIO = {"command", "args", "env", "env_vars", "cwd"} +_MCP_HTTP = { + "url", + "http_headers", + "env_http_headers", + "http_headers_helper", + "bearer_token_env_var", + "auth", + "oauth", + "oauth_resource", + "scopes", +} +_SHELL_TOOLS = {"Bash", "shell", "shell_command", "exec_command", "write_stdin"} +_APPROVAL_MODES = {"auto", "prompt", "writes", "approve"} + + +def _mapping(value: Any, label: str) -> dict[str, Any]: + if not isinstance(value, Mapping) or not all(isinstance(k, str) for k in value): + raise TypeError(f"{label} must be a mapping with string keys") + return copy.deepcopy(dict(value)) + + +def _string(value: Any, label: str) -> None: + if not isinstance(value, str) or not value: + raise ValueError(f"{label} must be a nonempty string") + + +def _strings(value: Any, label: str) -> None: + if not isinstance(value, list) or not all(isinstance(v, str) for v in value): + raise TypeError(f"{label} must be a list of strings") + + +def _string_map(value: Any, label: str) -> None: + values = _mapping(value, label) + if not all(isinstance(v, str) for v in values.values()): + raise TypeError(f"{label} must contain string values") + + +def _nonnegative(value: Any, label: str, *, integer: bool = False) -> None: + valid_type = type(value) is int if integer else type(value) in (int, float) + if not valid_type or not math.isfinite(value) or value < 0: + kind = "integer" if integer else "number" + raise ValueError(f"{label} must be a finite nonnegative {kind}") + + +def _unknown_keys(value: dict[str, Any], allowed: set[str], label: str) -> None: + unknown = value.keys() - allowed + if unknown: + raise ValueError(f"Unsupported {label} fields: {', '.join(sorted(unknown))}") + + +def _exact_name(value: str, label: str, *, allow_all: bool = False) -> None: + _string(value, label) + if allow_all and value == "*": + return + if any(c in value for c in "*?()"): + raise ValueError(f"{label} requires an exact tool name" + (" or '*'" if allow_all else "")) + + +def _mcp_server(name: str, value: Any, cwd: Path | None) -> dict[str, Any]: + _string(name, "MCP server name") + if name == "agent-tools": + raise ValueError("MCP server name 'agent-tools' is reserved for framework tools") + server = _mapping(value, f"MCP server {name!r}") + transport = server.pop("type", "stdio" if "command" in server else "http") + if transport not in {"stdio", "http", "streamable-http"}: + raise ValueError( + f"Codex does not support MCP transport {transport!r}; use STDIO or streamable HTTP" + ) + if "headers" in server: + if "http_headers" in server: + raise ValueError(f"MCP server {name!r} specifies both headers and http_headers") + server["http_headers"] = server.pop("headers") + stdio = transport == "stdio" + _unknown_keys(server, _MCP_COMMON | (_MCP_STDIO if stdio else _MCP_HTTP), "MCP server") + _string(server.get("command" if stdio else "url"), "MCP command" if stdio else "MCP URL") + if stdio: + if "cwd" in server: + path = Path(server["cwd"]) + server["cwd"] = str(path if path.is_absolute() else (cwd or Path.cwd()) / path) + elif cwd is not None: + server["cwd"] = str(cwd.absolute()) + else: + url = urlsplit(server["url"]) + if url.scheme not in {"http", "https"} or not url.netloc: + raise ValueError("An HTTP MCP server URL must use http:// or https:// with a host") + + for key in { + "args", + "enabled_tools", + "disabled_tools", + "scopes", + "omit_tools_from", + } & server.keys(): + _strings(server[key], f"MCP {key}") + for key in {"env", "http_headers", "env_http_headers"} & server.keys(): + _string_map(server[key], f"MCP {key}") + for key in {"enabled", "required", "supports_parallel_tool_calls"} & server.keys(): + if not isinstance(server[key], bool): + raise TypeError(f"MCP {key} must be a boolean") + for key in {"tool_timeout_sec", "startup_timeout_sec", "startup_timeout_ms"} & server.keys(): + _nonnegative(server[key], f"MCP {key}", integer=key.endswith("_ms")) + if "startup_timeout_sec" in server and "startup_timeout_ms" in server: + raise ValueError("Specify only one MCP startup timeout unit") + for key in { + "bearer_token_env_var", + "http_headers_helper", + "oauth_resource", + "environment_id", + "name", + } & server.keys(): + _string(server[key], f"MCP {key}") + if "auth" in server and server["auth"] not in {"oauth", "chatgpt"}: + raise ValueError("MCP auth must be 'oauth' or 'chatgpt'") + if "default_tools_approval_mode" in server: + if server["default_tools_approval_mode"] not in _APPROVAL_MODES: + raise ValueError("Unsupported MCP default_tools_approval_mode") + if "omit_tools_from" in server and set(server["omit_tools_from"]) - { + "direct", + "deferred", + "code_mode", + }: + raise ValueError("Unsupported MCP omit_tools_from surface") + if "env_vars" in server: + if not isinstance(server["env_vars"], list): + raise TypeError("MCP env_vars must be a list") + for variable in server["env_vars"]: + if isinstance(variable, str): + _string(variable, "MCP environment variable") + else: + variable = _mapping(variable, "MCP environment variable") + _unknown_keys(variable, {"name", "source"}, "MCP environment variable") + _string(variable.get("name"), "MCP environment variable name") + if "source" in variable: + _string(variable["source"], "MCP environment variable source") + if "oauth" in server: + oauth = _mapping(server["oauth"], "MCP OAuth settings") + _unknown_keys(oauth, {"client_id", "callback_port", "callback_url"}, "MCP OAuth") + for key in {"client_id", "callback_url"} & oauth.keys(): + _string(oauth[key], f"MCP OAuth {key}") + if "callback_port" in oauth: + _nonnegative(oauth["callback_port"], "MCP OAuth callback_port", integer=True) + if oauth["callback_port"] > 65535: + raise ValueError("MCP OAuth callback_port exceeds 65535") + if "tools" in server: + for tool_name, settings in _mapping(server["tools"], "MCP tools").items(): + _string(tool_name, "MCP tool name") + settings = _mapping(settings, "MCP tool settings") + _unknown_keys(settings, {"approval_mode", "output_token_limit"}, "MCP tool") + if "approval_mode" in settings and settings["approval_mode"] not in _APPROVAL_MODES: + raise ValueError("Unsupported MCP tool approval_mode") + if "output_token_limit" in settings: + _nonnegative(settings["output_token_limit"], "MCP output_token_limit", integer=True) + if not settings["output_token_limit"]: + raise ValueError("MCP output_token_limit must be positive") + return server + + +def _reject_unverifiable_hooks(hooks: dict[str, Any]) -> None: + for groups in hooks.values(): + if isinstance(groups, list) and any(not isinstance(group, Mapping) for group in groups): + raise TypeError( + "Codex cannot execute Claude SDK HookMatcher callbacks. Use required_tools " + "for shared required-tool enforcement." + ) + raise ValueError( + "Codex cannot verify trust of per-client native hooks. Configure and trust native hooks " + "through Codex before starting agent-flow, or use required_tools for shared required-tool " + "enforcement. Untrusted session hooks would be silently skipped by Codex." + ) + + +def build_session_config( + *, + extra_mcp_servers: dict[str, Any] | None = None, + hooks: dict[str, Any] | None = None, + disallowed_tools: list[str] | None = None, + tools: list[Any] | None = None, + cwd: Path | None = None, +) -> tuple[dict[str, Any], list[Any]]: + """Build per-thread Codex config and remove prohibited framework tools. + + Restrictions accept framework tool names, ``mcp__SERVER__TOOL`` names, + ``apps.APP_ID.TOOL`` names, and built-ins with native disable controls. + ``mcp__SERVER__*`` and ``apps.APP_ID.*`` disable an entire server/app. + Shell tool aliases disable the native shell family. A tool restriction is + not a filesystem policy: use a sandbox to prohibit writes through all tools. + Native hooks must be configured and trusted through Codex itself; this + adapter cannot verify trust of per-client hooks and rejects them. + """ + config: dict[str, Any] = {} + servers = { + name: _mcp_server(name, value, cwd) + for name, value in _mapping(extra_mcp_servers or {}, "MCP servers").items() + } + if servers: + config["mcp_servers"] = servers + if hooks: + _reject_unverifiable_hooks(_mapping(hooks, "Codex hooks")) + allowed = list(tools or []) + registered_names = {t.name for t in allowed} + if disallowed_tools is not None: + _strings(disallowed_tools, "disallowed_tools") + for name in disallowed_tools or []: + _string(name, "Disallowed tool") + dynamic_name = name.removeprefix("mcp__agent-tools__") + if ( + name in registered_names + or dynamic_name in registered_names + or name.startswith("mcp__agent-tools__") + ): + _exact_name(dynamic_name, "Framework tool restriction", allow_all=True) + allowed = [t for t in allowed if t.name not in {dynamic_name, name}] + if dynamic_name == "*": + allowed = [] + continue + if name in _SHELL_TOOLS: + config.setdefault("features", {})["shell_tool"] = False + elif name in {"AskUserQuestion", "request_user_input"}: + config.setdefault("tools", {})["experimental_request_user_input"] = {"enabled": False} + elif name in {"WebSearch", "WebFetch", "web_search"}: + config["web_search"] = "disabled" + elif name in {"view_image", "image_generation"}: + config.setdefault("features", {})[name] = False + elif name in {"update_plan", "TodoWrite"}: + config.setdefault("tools", {})["update_plan"] = {"enabled": False} + elif name.startswith("mcp__") and name.count("__") >= 2: + server_name, tool_name = name[5:].split("__", 1) + _exact_name(server_name, "Disallowed MCP server") + _exact_name(tool_name, "Disallowed MCP tool", allow_all=True) + if server_name == "codex_apps": + raise ValueError("Restrict Codex app tools using apps.APP_ID.TOOL or apps.APP_ID.*") + server = config.setdefault("mcp_servers", {}).setdefault(server_name, {}) + if tool_name == "*": + server["enabled"] = False + else: + disabled = server.setdefault("disabled_tools", []) + if tool_name not in disabled: + disabled.append(tool_name) + elif name.startswith("apps.") and len(name.split(".", 2)) == 3: + _, app_id, tool_name = name.split(".", 2) + _exact_name(app_id, "Disallowed app ID") + _exact_name(tool_name, "Disallowed app tool", allow_all=True) + if app_id == "_default": + raise ValueError("App restrictions require an explicit app ID, not _default") + app = config.setdefault("apps", {}).setdefault(app_id, {}) + if tool_name == "*": + app["enabled"] = False + else: + app.setdefault("tools", {})[tool_name] = {"enabled": False} + else: + raise ValueError( + f"Codex cannot enforce disallowed tool {name!r}. Use a registered framework tool " + "name, mcp__SERVER__TOOL, apps.APP_ID.TOOL, or a supported native tool control." + ) + try: + json.dumps(config, allow_nan=False) + except (TypeError, ValueError) as exc: + raise TypeError("Codex session configuration must be JSON serializable") from exc + return config, allowed diff --git a/agent-flow/agent_flow/backends/codex_events.py b/agent-flow/agent_flow/backends/codex_events.py new file mode 100644 index 000000000000..4bf0deed2c45 --- /dev/null +++ b/agent-flow/agent_flow/backends/codex_events.py @@ -0,0 +1,420 @@ +"""Translate the pinned Codex app-server protocol into shared backend events.""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import aclosing +from dataclasses import replace +from typing import Any, AsyncIterator, Protocol + +from ..types import ( + AgentTextEvent, + CompactBoundaryEvent, + RateLimitWarningEvent, + ServerToolCallEvent, + SessionInitEvent, + ThinkingEvent, + ToolCallEvent, + UsageInfo, +) +from .base import BackendClient, BackendEvent, ResultEvent + + +class CodexTransport(Protocol): + def stream_turn(self, thread_id: str, message: str) -> AsyncIterator[Any]: ... + + async def request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: ... + + +def _payload(notification: Any) -> dict[str, Any]: + """Normalize SDK models once, including forward-compatible notifications.""" + payload = notification.payload + if isinstance(payload, dict): + return payload + if hasattr(payload, "params"): + return payload.params + return payload.model_dump(by_alias=True, mode="json", exclude_unset=True) + + +def _integer(value: Any) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +_TOKEN_FIELDS = { + "inputTokens": "input_tokens", + "outputTokens": "output_tokens", + "cachedInputTokens": "cache_read_tokens", + "cacheWriteInputTokens": "cache_creation_tokens", + "totalTokens": "total_tokens", +} + + +def _usage(token_usage: dict[str, Any], baseline: dict[str, int]) -> UsageInfo: + # Thread totals are cumulative; ResultEvent describes this send_message call. + fields: dict[str, Any] = {} + for source, target in _TOKEN_FIELDS.items(): + current = _integer(token_usage.get("total", {}).get(source)) + if current is not None: + previous = baseline.get(source, 0) + fields[target] = current - previous if current >= previous else current + context = _integer(token_usage.get("last", {}).get("totalTokens")) + window = _integer(token_usage.get("modelContextWindow")) + return UsageInfo( + **fields, + context_tokens=context, + context_window=window, + context_percentage=100.0 * context / window if context is not None and window else None, + ) + + +def _tool_event(item: dict[str, Any]) -> ToolCallEvent | ServerToolCallEvent | None: + kind = item.get("type") + item_id = item.get("id") + if kind == "commandExecution": + return ToolCallEvent("Bash", {"command": item.get("command", "")}, item_id) + if kind in {"dynamicToolCall", "mcpToolCall"}: + arguments = item.get("arguments") + name = item["tool"] + if kind == "mcpToolCall" and item.get("server"): + name = f"mcp__{item['server']}__{name}" + return ToolCallEvent(name, arguments if isinstance(arguments, dict) else {}, item_id) + if kind == "collabAgentToolCall": + fields = { + "prompt": "prompt", + "model": "model", + "reasoningEffort": "reasoning_effort", + "senderThreadId": "sender_thread_id", + "receiverThreadIds": "receiver_thread_ids", + "agentsStates": "agents_states", + "status": "status", + } + return ToolCallEvent( + item["tool"], + { + target: item[source] + for source, target in fields.items() + if item.get(source) is not None + }, + item_id, + ) + if kind == "fileChange": + body = {"changes": item.get("changes", [])} + if item.get("status") is not None: + body["status"] = item["status"] + return ToolCallEvent("FileChange", body, item_id) + if kind == "webSearch": + action = item.get("action") + body = dict(action) if isinstance(action, dict) else {} + if item.get("query") is not None: + body.setdefault("query", item["query"]) + return ServerToolCallEvent("web_search", body, item_id) + if kind == "imageView": + return ServerToolCallEvent("view_image", {"path": item.get("path")}, item_id) + return None + + +def _reasoning_text(item: dict[str, Any]) -> str: + parts = [] + for field in ("summary", "content"): + value = item.get(field, []) + if isinstance(value, str): + value = [value] + parts.extend(part.strip() for part in value if isinstance(part, str) and part.strip()) + return "\n".join(parts) + + +def _rate_limit_events(snapshot: dict[str, Any]) -> list[RateLimitWarningEvent]: + """Use an explicit local 80% warning threshold with native reset/limit data.""" + prefix = snapshot.get("limitId") or snapshot.get("limitName") or "codex" + reached = snapshot.get("rateLimitReachedType") + rejected = reached is not None or snapshot.get("spendControlReached") is True + events = [] + for name in ("primary", "secondary", "individualLimit"): + window = snapshot.get(name) + if not isinstance(window, dict): + continue + percent = window.get("usedPercent") + if name == "individualLimit" and _integer(window.get("remainingPercent")) is not None: + percent = 100 - window["remainingPercent"] + if _integer(percent) is None: + continue + utilization = percent / 100 + status = ( + "rejected" + if rejected or utilization >= 1 + else "allowed_warning" + if utilization >= 0.8 + else "allowed" + ) + events.append( + RateLimitWarningEvent( + status, f"{prefix}/{name}", _integer(window.get("resetsAt")), utilization + ) + ) + if rejected and not events: + events.append(RateLimitWarningEvent("rejected", reached or f"{prefix}/spend_control")) + return events + + +def _merge_available(previous: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]: + merged = dict(previous) + for name, value in update.items(): + if isinstance(value, dict): + merged[name] = _merge_available(previous.get(name) or {}, value) + elif value is not None: + merged[name] = value + return merged + + +class CodexClient(BackendClient): + def __init__( + self, + transport: CodexTransport, + thread_id: str, + session_init: SessionInitEvent | None = None, + ) -> None: + self._transport = transport + self._thread_id = thread_id + self._session_init = session_init + self._initialized = False + self._totals: dict[str, int] = {} + self._context: UsageInfo | None = None + self._parents: dict[str, str] = {} + self._labels: dict[str, str] = {} + self._parent_tools: dict[str, str] = {} + self._rate_limits: dict[str, dict[str, Any]] = {} + self._last_rate_events: dict[str | None, RateLimitWarningEvent] = {} + + async def list_available_skills(self) -> list[str] | None: + return list(self._session_init.skills) if self._session_init is not None else None + + async def get_context_usage(self) -> UsageInfo | None: + # The protocol has no pre-turn context query. Never send a probe turn. + return replace(self._context) if self._context is not None else None + + async def _optional_request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: + try: + return await asyncio.wait_for(self._transport.request(method, params), timeout=5) + except Exception: + # Account APIs are unavailable for some billing routes, including API keys. + return {} + + def _rate_updates(self, payload: dict[str, Any], *, complete: bool = False): + snapshots = dict(payload.get("rateLimitsByLimitId") or {}) + main = payload.get("rateLimits") + if isinstance(main, dict): + snapshots[main.get("limitId") or "codex"] = main + for key, update in snapshots.items(): + previous = {} if complete else self._rate_limits.get(key, {}) + # Rolling updates explicitly use null for unavailable metadata. + snapshot = _merge_available(previous, update) + snapshot.setdefault("limitId", key) + self._rate_limits[key] = snapshot + for event in _rate_limit_events(snapshot): + if self._last_rate_events.get(event.rate_limit_type) != event: + self._last_rate_events[event.rate_limit_type] = event + yield event + + def _rate_recovery_needs_snapshot(self, payload: dict[str, Any]) -> bool: + update = payload.get("rateLimits") or {} + previous = self._rate_limits.get(update.get("limitId") or "codex", {}) + return any( + previous.get(field) and field in update and update[field] is None + for field in ("rateLimitReachedType", "spendControlReached") + ) + + def _belongs_to_session(self, thread_id: str) -> bool: + seen: set[str] = set() + while thread_id != self._thread_id: + if thread_id in seen or thread_id not in self._parents: + return False + seen.add(thread_id) + thread_id = self._parents[thread_id] + return True + + def _register_thread(self, thread: dict[str, Any]) -> None: + thread_id = thread["id"] + if thread.get("parentThreadId"): + self._parents[thread_id] = thread["parentThreadId"] + nickname, role = thread.get("agentNickname"), thread.get("agentRole") + if nickname or role: + self._labels[thread_id] = ( + f"{nickname} ({role})" if nickname and role else nickname or role + ) + + def _attribute(self, event: BackendEvent, thread_id: str) -> BackendEvent: + if thread_id != self._thread_id: + event.parent_tool_use_id = self._parent_tools.get(thread_id) + event.agent_label = self._labels.get(thread_id, thread_id) + return event + + async def send_message(self, message: str) -> AsyncIterator[BackendEvent]: + if not self._initialized: + self._initialized = True + if self._session_init is not None: + yield self._session_init + initial_limits = await self._optional_request("account/rateLimits/read", {}) + for event in self._rate_updates(initial_limits, complete=True): + yield event + + baseline = dict(self._totals) + latest_usage: UsageInfo | None = None + final_text = "" + fallback_text = "" + partial_text: list[str] = [] + permission_denials: list[Any] = [] + diagnostics: list[dict[str, Any]] = [] + main_turn_id: str | None = None + emitted: set[tuple[str, str, str]] = set() + pending: list[tuple[str, BackendEvent]] = [] + + def flush_pending(*, final: bool = False): + remaining = [] + for thread_id, event in pending: + if self._belongs_to_session(thread_id) and ( + final or thread_id in self._parent_tools + ): + yield self._attribute(event, thread_id) + else: + remaining.append((thread_id, event)) + pending[:] = remaining + + async with aclosing(self._transport.stream_turn(self._thread_id, message)) as stream: + async for notification in stream: + method, payload = notification.method, _payload(notification) + thread_id = payload.get("threadId") + if method == "account/rateLimits/updated": + recovery = self._rate_recovery_needs_snapshot(payload) + for event in self._rate_updates(payload): + yield event + if recovery: + # Null in a sparse update is unavailable, not proof of recovery. + snapshot = await self._optional_request("account/rateLimits/read", {}) + for event in self._rate_updates(snapshot, complete=True): + yield event + continue + if method == "thread/started": + self._register_thread(payload["thread"]) + for event in flush_pending(): + yield event + continue + if method == "turn/started" and thread_id == self._thread_id: + main_turn_id = payload["turn"]["id"] + continue + if thread_id == self._thread_id and main_turn_id is not None: + event_turn_id = ( + payload.get("turn", {}).get("id") + if method == "turn/completed" + else payload.get("turnId") + ) + if event_turn_id is not None and event_turn_id != main_turn_id: + continue + if method in {"item/started", "item/completed"}: + item = payload["item"] + kind = item.get("type") + if kind == "collabAgentToolCall": + for receiver in item.get("receiverThreadIds", []): + self._parents.setdefault( + receiver, item.get("senderThreadId", thread_id) + ) + if item.get("tool") == "spawnAgent": + self._parent_tools[receiver] = item["id"] + for event in flush_pending(): + yield event + + events: list[BackendEvent] = [] + tool = _tool_event(item) + if tool is not None: + events.append(tool) + if kind == "contextCompaction" and thread_id == self._thread_id: + events.append( + CompactBoundaryEvent( + item.get("trigger"), _integer(item.get("preTokens")) + ) + ) + if method == "item/completed": + if item.get("status") == "declined" and self._belongs_to_session(thread_id): + permission_denials.append( + {"kind": kind, "id": item.get("id"), "item": item} + ) + # A failed command/tool can be recovered by the agent. Only turn failure + # makes the request fail; permission denials remain separately visible. + if kind == "reasoning": + text = _reasoning_text(item) + if text: + events.append(ThinkingEvent(text)) + elif kind == "agentMessage": + text = (item.get("text") or "").strip() + if text: + events.append(AgentTextEvent(text)) + if thread_id == self._thread_id: + partial_text.append(text) + if item.get("phase") == "final_answer": + final_text = text + elif item.get("phase") is None: + fallback_text = text + for event in events: + key = (thread_id, item.get("id"), type(event).__name__) + if key in emitted: + continue + emitted.add(key) + if thread_id == self._thread_id: + yield event + else: + pending.append((thread_id, event)) + for event in flush_pending(): + yield event + elif method == "thread/tokenUsage/updated" and thread_id == self._thread_id: + token_usage = payload["tokenUsage"] + latest_usage = _usage(token_usage, baseline) + self._totals = { + key: value + for key, value in token_usage.get("total", {}).items() + if key in _TOKEN_FIELDS and _integer(value) is not None + } + self._context = UsageInfo( + context_tokens=latest_usage.context_tokens, + context_window=latest_usage.context_window, + context_percentage=latest_usage.context_percentage, + ) + elif method == "error" and thread_id == self._thread_id: + diagnostics.append(payload) + elif method == "turn/completed" and thread_id == self._thread_id: + turn = payload["turn"] + for event in flush_pending(final=True): + yield event + if turn.get("error") is not None or turn.get("status") != "completed": + error = turn.get("error") or { + "message": f"Turn ended: {turn.get('status')}" + } + info = error.get("codexErrorInfo") + if info in ("usageLimitExceeded", "rateLimitExceeded"): + yield RateLimitWarningEvent("rejected", info) + detail = { + "threadId": self._thread_id, + "turnId": turn.get("id"), + "error": error, + "diagnostics": diagnostics, + "partialText": partial_text, + } + raise RuntimeError( + f"Codex turn failed: {json.dumps(detail, ensure_ascii=False)}" + ) + usage = latest_usage or UsageInfo() + usage.num_turns = 1 + usage.duration_ms = _integer(turn.get("durationMs")) + cost = await self._optional_request( + "account/usage/read", {"threadId": self._thread_id} + ) + thread_usage = cost.get("threadUsage") or {} + estimate = _integer(thread_usage.get("estimatedUsageUsdMicros")) + if thread_usage.get("threadId") == self._thread_id and estimate is not None: + usage.estimated_thread_cost_usd = estimate / 1_000_000 + yield ResultEvent( + text=final_text or fallback_text, + usage=usage, + permission_denials=permission_denials, + ) + return + raise RuntimeError("Codex notification stream ended before the main turn completed") diff --git a/agent-flow/agent_flow/backends/codex_transport.py b/agent-flow/agent_flow/backends/codex_transport.py new file mode 100644 index 000000000000..12fe661b676d --- /dev/null +++ b/agent-flow/agent_flow/backends/codex_transport.py @@ -0,0 +1,236 @@ +"""Codex SDK transport and per-session tool dispatch. + +The SDK's public turn streams omit account notifications and descendant turns. +One isolated override observes decoded notifications before the SDK routes them; +all process management and JSON-RPC requests still use the pinned SDK. Register +observers before starting a turn so even an immediate completion is retained. +""" + +from __future__ import annotations + +import asyncio +import contextvars +from dataclasses import dataclass +from typing import Any, AsyncIterator + +from openai_codex.client import CodexClient as SdkClient +from openai_codex.client import CodexConfig +from openai_codex.models import Notification +from pydantic import RootModel + +from ..tools import ToolDefinition, normalize_tool, tool_result_to_codex + + +class _Response(RootModel[dict[str, Any]]): + pass + + +def notification_data(event: Notification) -> dict[str, Any]: + """Normalize typed and unknown SDK notifications at the transport boundary.""" + payload = event.payload + if hasattr(payload, "params"): + return payload.params + if hasattr(payload, "model_dump"): + return payload.model_dump(by_alias=True, mode="json", exclude_unset=True) + if isinstance(payload, dict): + return payload + return {} + + +class _ObservedSdkClient(SdkClient): + def __init__(self, config, approval_handler, observer) -> None: + super().__init__(config=config, approval_handler=approval_handler) + self._observer = observer + + def _coerce_notification(self, method: str, params: object) -> Notification: + event = super()._coerce_notification(method, params) + self._observer(event) + return event + + +@dataclass +class _ToolBinding: + tools: dict[str, ToolDefinition] + context: contextvars.Context + + +class CodexTransport: + def __init__(self, config: CodexConfig) -> None: + self._loop = asyncio.get_running_loop() + self._client = _ObservedSdkClient(config, self._server_request, self._receive) + self._bindings: dict[str, _ToolBinding] = {} + self._parents: dict[str, str] = {} + self._threads: dict[str, Notification] = {} + self._queues: set[asyncio.Queue] = set() + self._active_threads: set[str] = set() + self._monitor: asyncio.Task | None = None + self._closed = False + self._failure: BaseException | None = None + + async def start(self) -> None: + try: + await asyncio.to_thread(self._client.start) + await asyncio.to_thread(self._client.initialize) + except BaseException: + await self.close() + raise + self._monitor = asyncio.create_task(self._drain_global_notifications()) + + async def close(self) -> None: + if self._closed: + return + self._closed = True + await asyncio.to_thread(self._client.close) + if self._monitor is not None: + self._monitor.cancel() + await asyncio.gather(self._monitor, return_exceptions=True) + self._broadcast(RuntimeError("Codex transport closed.")) + self._bindings.clear() + self._parents.clear() + self._threads.clear() + + async def _drain_global_notifications(self) -> None: + # The observer already delivered these events. Drain the SDK's global + # queue to bound memory and to propagate transport failures to waiters. + try: + while True: + await asyncio.to_thread(self._client.next_notification) + except asyncio.CancelledError: + raise + except Exception as exc: + self._failure = exc + self._broadcast(exc) + + def _receive(self, event: Notification) -> None: + # Called on the SDK reader thread. Record ancestry before a subsequent + # dynamic-tool request, which may arrive before the event loop wakes. + data = notification_data(event) + if event.method == "thread/started": + thread = data.get("thread", {}) + thread_id = thread.get("id") + parent_id = thread.get("parentThreadId") + if isinstance(thread_id, str): + self._threads[thread_id] = event + if isinstance(parent_id, str): + self._parents[thread_id] = parent_id + if not self._closed: + self._loop.call_soon_threadsafe(self._broadcast, event) + + def _broadcast(self, event: Notification | BaseException) -> None: + for queue in tuple(self._queues): + queue.put_nowait(event) + + def register_tools(self, thread_id: str, tools: list) -> None: + definitions = [normalize_tool(tool) for tool in tools] + self._bindings[thread_id] = _ToolBinding( + {tool.name: tool for tool in definitions}, contextvars.copy_context() + ) + + def unregister_tools(self, thread_id: str) -> None: + self._bindings.pop(thread_id, None) + + def _binding(self, thread_id: str) -> _ToolBinding | None: + seen: set[str] = set() + while thread_id and thread_id not in seen: + seen.add(thread_id) + if thread_id in self._bindings: + return self._bindings[thread_id] + thread_id = self._parents.get(thread_id, "") + return None + + def _server_request(self, method: str, params: dict | None) -> dict: + if method.endswith("/requestApproval"): + if method == "item/permissions/requestApproval": + return {"permissions": (params or {}).get("permissions", {}), "scope": "turn"} + return {"decision": "accept"} + if method != "item/tool/call": + # Interactive input belongs to the shared ask_human tool. Do not + # invent an answer to other kinds of server-initiated requests. + return {} + params = params or {} + tool_name = params.get("tool", "") + binding = self._binding(params.get("threadId", "")) + tool = binding.tools.get(tool_name) if binding else None + if tool is None: + return self._tool_error(f"Tool {tool_name!r} is not registered for this thread.") + arguments = params.get("arguments") + if not isinstance(arguments, dict): + return self._tool_error(f"Tool {tool_name!r} requires an object argument.") + + async def invoke() -> dict: + return tool_result_to_codex(await tool.handler(arguments)) + + try: + future = binding.context.copy().run( + asyncio.run_coroutine_threadsafe, invoke(), self._loop + ) + return future.result() + except Exception as exc: + return self._tool_error(f"Tool {tool_name!r} failed: {exc}") + + @staticmethod + def _tool_error(message: str) -> dict: + return {"contentItems": [{"type": "inputText", "text": message}], "success": False} + + async def request(self, method: str, params: dict | None) -> dict[str, Any]: + if self._closed: + raise RuntimeError("Codex transport is closed.") + if self._failure is not None: + raise RuntimeError("Codex transport failed.") from self._failure + response = await asyncio.to_thread( + self._client.request, method, params, response_model=_Response + ) + return response.root + + async def stream_turn(self, thread_id: str, message: str) -> AsyncIterator[Notification]: + if thread_id in self._active_threads: + raise RuntimeError("Concurrent turns on the same Codex session are not supported.") + self._active_threads.add(thread_id) + queue: asyncio.Queue = asyncio.Queue() + self._queues.add(queue) + completed = False + turn_id = None + try: + # Rehydrate lineage for descendants surviving a previous turn. + for event in tuple(self._threads.values()): + queue.put_nowait(event) + start_request = asyncio.create_task( + self.request( + "turn/start", + {"threadId": thread_id, "input": [{"type": "text", "text": message}]}, + ) + ) + try: + started = await asyncio.shield(start_request) + except asyncio.CancelledError: + # Cancelling a Python waiter does not cancel a JSON-RPC request. + # Recover the newly created turn so the finally block can stop it. + try: + started = await asyncio.shield(start_request) + turn_id = started["turn"]["id"] + except Exception: + pass + raise + turn_id = started["turn"]["id"] + while True: + event = await queue.get() + if isinstance(event, BaseException): + raise event + if event.method == "turn/completed": + data = notification_data(event) + if data.get("turn", {}).get("id") == turn_id: + completed = True + yield event + if completed: + break + finally: + self._queues.discard(queue) + self._active_threads.discard(thread_id) + if turn_id is not None and not completed and not self._closed: + try: + await asyncio.wait_for( + self.request("turn/interrupt", {"threadId": thread_id, "turnId": turn_id}), + timeout=5, + ) + except Exception: + pass diff --git a/agent-flow/agent_flow/config.py b/agent-flow/agent_flow/config.py index d0c61d5956c5..76d81617bddc 100644 --- a/agent-flow/agent_flow/config.py +++ b/agent-flow/agent_flow/config.py @@ -9,7 +9,7 @@ SessionMode = Literal["stateless", "persistent"] CLAUDE_CODE_DEFAULT_MODEL = os.environ.get("CLAUDE_CODE_DEFAULT_MODEL", "claude-opus-5") -CODEX_DEFAULT_MODEL = os.environ.get("CODEX_DEFAULT_MODEL", "gpt-5.6-sol") +CODEX_DEFAULT_MODEL = os.environ.get("CODEX_DEFAULT_MODEL", "gpt-6-astra") @dataclass(frozen=True) @@ -17,20 +17,14 @@ class BackendConfig: kind: BackendKind model: str tools: list[Any] | None = None - # Backend-specific hook configuration. Currently only the ``claude-code`` - # backend consumes this — it is forwarded verbatim to the Claude Agent - # SDK's ``ClaudeAgentOptions(hooks=...)`` (a dict keyed by event name, - # values are lists of ``HookMatcher``). Other backends ignore it. + # Native SDK hook configuration. Claude accepts HookMatcher callbacks. + # Codex requires hooks to be configured and trusted in its native config; + # per-client callback injection is rejected. Use AgentLayerConfig.required_tools + # for the shared bounded tool-call policy on either backend. hooks: dict[str, Any] | None = None - # Extra MCP servers to make available to the agent, keyed by the - # server name the model sees (tools become ``mcp____``). - # Currently only the ``claude-code`` backend consumes this — values - # are forwarded verbatim into ``ClaudeAgentOptions.mcp_servers`` - # alongside the in-process ``agent-tools`` server built from - # ``tools``. Typical use: wiring an HTTP MCP server into a single - # agent's session via - # ``{"knowledge-base": {"type": "http", "url": "..."}}``. Other backends - # accept and ignore the field. + # Additional per-session MCP servers, in Claude-style STDIO/HTTP form. + # Codex translates these into thread configuration without writing global + # config. Unsupported transport/auth options fail during validation. extra_mcp_servers: dict[str, Any] | None = None # Working directory the backend runs the agent in. Forwarded to the # SDK as the session ``cwd`` (Claude Code's @@ -96,6 +90,11 @@ class AgentLayerConfig: # ``WebFetch``, ...). Empty means "no extra bans", which is the # historical behaviour for every existing layer. disallowed_tools: tuple[str, ...] = () + # Every listed tool must be attempted by the root agent during one + # layer invocation. Missing calls trigger one corrective SDK turn in + # the same session; failure after that raises RequiredToolCallError. + # Calls from previous invocations or subagents do not satisfy this. + required_tools: tuple[str, ...] = () # Called once per backend event as ``on_activity(kind, event)``, where # ``kind`` is a short tag (``tool``, ``text``, ``thinking``, ...). # diff --git a/agent-flow/agent_flow/console.py b/agent-flow/agent_flow/console.py index 130231174192..55dc3cdfa43c 100644 --- a/agent-flow/agent_flow/console.py +++ b/agent-flow/agent_flow/console.py @@ -287,6 +287,11 @@ def _format_usage_body(usage: UsageInfo) -> str: if usage.cost_usd is not None: parts.append(f"[info]cost[/info] [target]${usage.cost_usd:.4f}[/target]") + if usage.estimated_thread_cost_usd is not None: + parts.append( + "[info]estimated thread cost[/info] " + f"[target]${usage.estimated_thread_cost_usd:.4f}[/target]" + ) if usage.num_turns is not None: parts.append(f"[info]turns[/info] [target]{usage.num_turns}[/target]") if usage.duration_ms is not None: diff --git a/agent-flow/agent_flow/hooks.py b/agent-flow/agent_flow/hooks.py index 91595271d5a1..3e802aa3a077 100644 --- a/agent-flow/agent_flow/hooks.py +++ b/agent-flow/agent_flow/hooks.py @@ -1,24 +1,14 @@ -"""Reusable hook helpers for the claude-code backend. +"""Shared required-tool policy and legacy Claude-specific hook helpers. -The Claude Agent SDK supports a ``hooks`` option on ``ClaudeAgentOptions`` -that lets callers intercept events such as ``Stop`` (the agent trying to end -its turn). This module builds common hook patterns on top of that API. - -The primary helper is :func:`require_tool_call_stop_hook`, which produces a -``Stop`` hook that blocks the agent from ending a turn until it has called -at least one of a set of required tools. This upgrades prompt-only "please -call X before stopping" instructions into a runtime guarantee. - -Example:: - - from agent_flow import BackendConfig, CLAUDE_CODE_DEFAULT_MODEL - from agent_flow.hooks import require_tool_call_stop_hook - - hooks = require_tool_call_stop_hook(["append_planner_progress"]) - backend = BackendConfig( - kind="claude-code", model=CLAUDE_CODE_DEFAULT_MODEL, tools=planner_tools, hooks=hooks - ) +Use ``AgentLayerConfig(required_tools=("append_progress", "update_status"))`` +to require every listed tool in either backend. ``AgentLayer`` observes +root-agent calls, then sends one corrective SDK turn in the same session +if calls are missing. It raises ``RequiredToolCallError`` if that bounded +correction still leaves required calls missing. +``require_tool_call_stop_hook`` remains a Claude-native alternative for +existing callers. It inspects a transcript and requires ANY of its listed +tools, using Claude's Stop hook rather than a framework-initiated turn. """ from __future__ import annotations @@ -27,7 +17,45 @@ from pathlib import Path from typing import Any, Iterable -from claude_agent_sdk.types import HookMatcher +from .types import ToolCallEvent + + +class RequiredToolCallError(RuntimeError): + """A layer finished its corrective turn without all required tool calls.""" + + +class RequiredToolPolicy: + """Track ALL required root-agent calls within one logical invocation. + + This counts attempted calls, matching the legacy transcript hook's + contract; it does not claim the tool's operation succeeded. Unlike the + legacy helper's ANY matching, workflows require every configured tool. + """ + + def __init__(self, required_tools: Iterable[str]) -> None: + self._remaining = dict.fromkeys(required_tools) + + @property + def missing(self) -> tuple[str, ...]: + return tuple(self._remaining) + + def record(self, event: ToolCallEvent) -> None: + # Labels alone can identify a child when no invocation ID was + # supplied upstream. Neither form of child activity can discharge + # a root agent's obligation to write its own progress/status. + if event.parent_tool_use_id is not None or event.agent_label is not None: + return + for required in self.missing: + if _matches_required(event.name, {required}): + del self._remaining[required] + + def correction_message(self) -> str: + names = ", ".join(f"`{name}`" for name in self.missing) + return ( + f"You must call every remaining required tool before completing this request: {names}. " + "Call each now using the work already completed in this session, then finish. " + "Do not repeat tool calls whose requirement was already satisfied." + ) def _iter_transcript_entries(transcript_path: str | Path) -> list[dict]: @@ -152,6 +180,8 @@ def require_tool_call_stop_hook( one per stop event, which is enough in practice while keeping the guarantee bounded. """ + from claude_agent_sdk.types import HookMatcher + required = list(required_tool_names) if not required: raise ValueError("required_tool_names must not be empty") diff --git a/agent-flow/agent_flow/layers.py b/agent-flow/agent_flow/layers.py index e6082212cd1c..d9d9ff70a46b 100644 --- a/agent-flow/agent_flow/layers.py +++ b/agent-flow/agent_flow/layers.py @@ -1,10 +1,10 @@ from __future__ import annotations from contextlib import AsyncExitStack, asynccontextmanager +from dataclasses import replace from typing import Any, AsyncIterator, Callable import anyio -from claude_agent_sdk import tool from .backends import create_backend from .backends.base import Backend, BackendClient, ResultEvent @@ -24,9 +24,11 @@ print_tool_call, print_user_prompt, ) +from .hooks import RequiredToolCallError, RequiredToolPolicy from .logger import Logger, get_logger from .module import Module from .runtime import PortalRunner, build_request +from .tools import tool from .types import ( AgentRequest, AgentResponse, @@ -43,6 +45,37 @@ PromptBuilder = Callable[[str], AgentRequest] +def _merge_turn_usage(previous: UsageInfo | None, current: UsageInfo | None) -> UsageInfo | None: + """Sum per-SDK-turn accounting and keep the final context snapshot. + + Backends normalize cumulative upstream counters to turn deltas before + yielding ResultEvent. Context occupancy is a snapshot, not an expense, + and must not be summed across the initial and corrective turns. + """ + if previous is None: + return current + if current is None: + return previous + additive = ( + "input_tokens", + "output_tokens", + "cache_creation_tokens", + "cache_read_tokens", + "total_tokens", + "cost_usd", + "num_turns", + "duration_ms", + ) + totals = {} + for field in additive: + before, after = getattr(previous, field), getattr(current, field) + totals[field] = None if before is None and after is None else (before or 0) + (after or 0) + estimate = current.estimated_thread_cost_usd + if estimate is None: + estimate = previous.estimated_thread_cost_usd + return replace(current, **totals, estimated_thread_cost_usd=estimate) + + def _read_stdin(prompt: str = "> ") -> str: """Indirection so tests can monkeypatch the blocking stdin read.""" return input(prompt) @@ -296,21 +329,25 @@ async def ask_human(args): return ask_human + def _client_options(self, system_prompt: str | None) -> dict[str, Any]: + """Resolve the same client configuration for either session mode.""" + return { + "system_prompt": system_prompt or "", + "model": self.config.backend.model, + "tools": self._resolve_tools(), + "hooks": self.config.backend.hooks, + "disallowed_tools": self._resolve_disallowed_tools(), + "extra_mcp_servers": self.config.backend.extra_mcp_servers, + "cwd": self.config.backend.cwd, + } + @asynccontextmanager async def _temporary_client( self, system_prompt: str | None ) -> AsyncIterator[tuple[Backend, BackendClient]]: backend = create_backend(self.config.backend) async with backend: - async with backend.create_client( - system_prompt=system_prompt or "", - model=self.config.backend.model, - tools=self._resolve_tools(), - hooks=self.config.backend.hooks, - disallowed_tools=self._resolve_disallowed_tools(), - extra_mcp_servers=self.config.backend.extra_mcp_servers, - cwd=self.config.backend.cwd, - ) as client: + async with backend.create_client(**self._client_options(system_prompt)) as client: yield backend, client async def _ensure_persistent_client( @@ -322,9 +359,7 @@ async def _ensure_persistent_client( baseline only on the session's very first turn (later turns of a persistent session already carry conversation history). """ - if self._persistent_system_prompt is None: - self._persistent_system_prompt = system_prompt - elif self._persistent_system_prompt != system_prompt: + if self._persistent_client is not None and self._persistent_system_prompt != system_prompt: raise ValueError("Persistent AgentLayer sessions require a stable system prompt.") created = False @@ -333,18 +368,11 @@ async def _ensure_persistent_client( stack = AsyncExitStack() await stack.__aenter__() client = await stack.enter_async_context( - backend.create_client( - system_prompt=system_prompt or "", - model=self.config.backend.model, - tools=self._resolve_tools(), - hooks=self.config.backend.hooks, - disallowed_tools=self._resolve_disallowed_tools(), - extra_mcp_servers=self.config.backend.extra_mcp_servers, - cwd=self.config.backend.cwd, - ) + backend.create_client(**self._client_options(system_prompt)) ) self._persistent_client = client self._persistent_client_stack = stack + self._persistent_system_prompt = system_prompt created = True return self._persistent_client, created @@ -364,6 +392,7 @@ async def _execute( sink = logger.console result_text = "" usage: UsageInfo | None = None + policy = RequiredToolPolicy(self.config.required_tools) def observe(kind: str, event: Any) -> None: """Forward one event to the config's observer, if any. @@ -380,40 +409,62 @@ def observe(kind: str, event: Any) -> None: except Exception: # noqa: BLE001 - an observer cannot fail the run pass - async for event in client.send_message(request.content): - if isinstance(event, ResultEvent): - result_text = event.text - usage = event.usage - elif isinstance(event, ToolCallEvent): - observe("tool", event) - if self.config.print_activity: - print_tool_call(self.layer_name, event, sink) - elif isinstance(event, ServerToolCallEvent): - observe("server_tool", event) - if self.config.print_activity: - print_server_tool_call(self.layer_name, event, sink) - elif isinstance(event, ThinkingEvent): - observe("thinking", event) - if self.config.print_activity: - print_thinking(self.layer_name, event, sink) - elif isinstance(event, AgentTextEvent): - observe("text", event) - if self.config.print_activity: - print_agent_text(self.layer_name, event, sink) - elif isinstance(event, SessionInitEvent): - observe("session", event) - if self.config.print_activity: - print_session_init(self.layer_name, event, sink) - elif isinstance(event, RateLimitWarningEvent): - observe("rate_limit", event) - if self.config.print_activity: - print_rate_limit(self.layer_name, event, sink) - elif isinstance(event, CompactBoundaryEvent): - observe("compact", event) - if self.config.print_activity: - print_compact_boundary(self.layer_name, event, sink) - else: - raise TypeError(f"Unexpected backend event: {type(event).__name__}") + # A correction is a second SDK turn, after the original turn has + # completed. Keep the same client even for a stateless invocation: + # its history contains the work needed to finish missing calls. + # The policy survives this correction, but never another invocation. + message = request.content + for attempt in range(2): + turn_usage: UsageInfo | None = None + async for event in client.send_message(message): + if isinstance(event, ResultEvent): + if event.is_error: + raise RuntimeError( + "; ".join(event.errors) or event.text or "Agent turn failed" + ) + result_text = event.text + turn_usage = event.usage + elif isinstance(event, ToolCallEvent): + policy.record(event) + observe("tool", event) + if self.config.print_activity: + print_tool_call(self.layer_name, event, sink) + elif isinstance(event, ServerToolCallEvent): + observe("server_tool", event) + if self.config.print_activity: + print_server_tool_call(self.layer_name, event, sink) + elif isinstance(event, ThinkingEvent): + observe("thinking", event) + if self.config.print_activity: + print_thinking(self.layer_name, event, sink) + elif isinstance(event, AgentTextEvent): + observe("text", event) + if self.config.print_activity: + print_agent_text(self.layer_name, event, sink) + elif isinstance(event, SessionInitEvent): + observe("session", event) + if self.config.print_activity: + print_session_init(self.layer_name, event, sink) + elif isinstance(event, RateLimitWarningEvent): + observe("rate_limit", event) + if self.config.print_activity: + print_rate_limit(self.layer_name, event, sink) + elif isinstance(event, CompactBoundaryEvent): + observe("compact", event) + if self.config.print_activity: + print_compact_boundary(self.layer_name, event, sink) + else: + raise TypeError(f"Unexpected backend event: {type(event).__name__}") + usage = _merge_turn_usage(usage, turn_usage) + if not policy.missing: + break + if attempt: + missing = ", ".join(policy.missing) + raise RequiredToolCallError( + f"{self.layer_name} did not call required tools after one corrective turn: " + f"{missing}" + ) + message = policy.correction_message() return AgentResponse( content=result_text, @@ -558,20 +609,12 @@ def fetch_baseline_context_usage(self) -> UsageInfo | None: model call), so it works wherever the CLI is available. Backends without a pre-input breakdown return ``None``. - Some backends only populate the breakdown after the session has - streamed at least once; when the first query comes back empty we - bootstrap with a minimal ``ping`` turn and re-query. That makes - this the one inspection helper that can cost a turn — - ``fetch_available_skills`` never does. + This inspection never sends a model turn. A backend without a + pre-input breakdown cannot provide a baseline by adding input. """ async def run() -> UsageInfo | None: async with self._temporary_client(self.config.system_prompt) as (_, client): - usage = await client.get_context_usage() - if usage is not None and usage.context_percentage is not None: - return usage - async for _ in client.send_message("ping"): - pass return await client.get_context_usage() return anyio.run(run) diff --git a/agent-flow/agent_flow/tools.py b/agent-flow/agent_flow/tools.py new file mode 100644 index 000000000000..e0c7a48856ee --- /dev/null +++ b/agent-flow/agent_flow/tools.py @@ -0,0 +1,323 @@ +"""Backend-independent tool definitions and MCP result conversion. + +Tool annotations describe behavior; they never grant permission to execute a tool. +Resources become readable text on both backends. Content that cannot be represented +is rejected explicitly so that a tool cannot appear to succeed with missing output. +""" + +from __future__ import annotations + +import copy +import json +import types +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from enum import Enum +from typing import ( + Annotated, + Any, + Literal, + NotRequired, + Required, + Union, + get_args, + get_origin, + get_type_hints, + is_typeddict, +) + +ToolHandler = Callable[[dict[str, Any]], Awaitable[Any]] + + +@dataclass +class ToolDefinition: + """A tool that can be translated into either backend's native API.""" + + name: str + description: str + input_schema: dict[str, Any] + handler: ToolHandler + annotations: dict[str, Any] | None = None + meta: dict[str, Any] | None = None + + +def _mapping(value: Any, *, label: str) -> dict[str, Any]: + if isinstance(value, Mapping): + return dict(value) + if callable(getattr(value, "model_dump", None)): + return value.model_dump(mode="json", by_alias=True, exclude_none=True) + raise TypeError(f"{label} must be a mapping, got {type(value).__name__}") + + +def _json_schema(value: Any, active_types: set[type]) -> dict[str, Any]: + if value is Any: + return {} + if value is None or value is type(None): + return {"type": "null"} + primitive_types = {str: "string", int: "integer", float: "number", bool: "boolean"} + if isinstance(value, type) and value in primitive_types: + return {"type": primitive_types[value]} + + origin = get_origin(value) + args = get_args(value) + if origin in (Annotated, Required, NotRequired): + return _json_schema(args[0], active_types) + if origin in (Union, types.UnionType): + return {"anyOf": [_json_schema(arg, active_types) for arg in args]} + if origin is Literal: + return {"enum": list(args)} + if isinstance(value, type) and issubclass(value, Enum): + return {"enum": [member.value for member in value]} + if origin is list or value is list: + return {"type": "array", "items": _json_schema(args[0], active_types) if args else {}} + if origin is tuple or value is tuple: + if not args: + return {"type": "array"} + if len(args) == 2 and args[1] is Ellipsis: + return {"type": "array", "items": _json_schema(args[0], active_types)} + return { + "type": "array", + "prefixItems": [_json_schema(arg, active_types) for arg in args], + "minItems": len(args), + "maxItems": len(args), + } + if origin in (dict, Mapping) or value is dict: + if args and args[0] not in (str, Any): + raise TypeError("Tool schema dictionary keys must be strings") + return { + "type": "object", + "additionalProperties": _json_schema(args[1], active_types) if args else {}, + } + if is_typeddict(value): + if value in active_types: + raise TypeError("Recursive TypedDict tool schemas require an explicit JSON Schema") + active_types.add(value) + try: + hints = get_type_hints(value, include_extras=True) + properties = {key: _json_schema(hint, active_types) for key, hint in hints.items()} + required = [] + for key, hint in hints.items(): + qualifier = _field_qualifier(hint) + if qualifier is Required or ( + qualifier is not NotRequired and key in value.__required_keys__ + ): + required.append(key) + return {"type": "object", "properties": properties, "required": required} + finally: + active_types.remove(value) + raise TypeError(f"Unsupported Python type in tool schema: {value!r}") + + +def _field_qualifier(hint: Any) -> Any: + while get_origin(hint) is Annotated: + hint = get_args(hint)[0] + return get_origin(hint) + + +def _mcp_object_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Give an explicit JSON Schema the root shape MCP tool schemas require. + + Both backends advertise an object schema with a ``properties`` key. The Claude SDK + also relies on that shape to recognise a JSON Schema: a dict without a string ``type`` + and a ``properties`` key is re-read as Python shorthand, so keywords such as ``$ref`` + become required parameters. Root constraints that do not fit the shape move into + ``allOf``, which keeps their meaning across JSON Schema drafts. + """ + constraints = [] + if "$ref" in schema: + constraints.append({"$ref": schema.pop("$ref")}) + if "type" in schema and not isinstance(schema["type"], str): + constraints.append({"type": schema.pop("type")}) + if constraints: + schema["allOf"] = constraints + list(schema.get("allOf", [])) + schema.setdefault("type", "object") + schema.setdefault("properties", {}) + return schema + + +def normalize_input_schema(schema: Any) -> dict[str, Any]: + """Copy JSON Schema into the MCP object shape, or convert shorthand/TypedDict input.""" + if isinstance(schema, Mapping): + # A field named "type" can also occur in shorthand, e.g. {"type": str}. + schema_type = schema.get("type") + is_json_schema = isinstance(schema_type, (str, list)) or any( + isinstance(schema.get(key), expected) + for key, expected in ( + ("$ref", str), + ("$defs", Mapping), + ("$schema", str), + ("properties", Mapping), + ("additionalProperties", (bool, Mapping)), + ("required", list), + ("allOf", list), + ("anyOf", list), + ("oneOf", list), + ) + ) + if is_json_schema: + normalized = _mcp_object_schema(copy.deepcopy(dict(schema))) + else: + properties = {key: _json_schema(hint, set()) for key, hint in schema.items()} + normalized = { + "type": "object", + "properties": properties, + "required": [ + key for key, hint in schema.items() if _field_qualifier(hint) is not NotRequired + ], + } + elif is_typeddict(schema): + normalized = _json_schema(schema, set()) + else: + raise TypeError("Tool input_schema must be a JSON Schema, shorthand mapping, or TypedDict") + # Fail at registration instead of in a backend's transport encoder. + try: + json.dumps(normalized, allow_nan=False) + except (TypeError, ValueError) as exc: + raise TypeError("Tool input_schema must contain JSON-serializable values") from exc + return normalized + + +def normalize_tool(value: Any) -> ToolDefinition: + """Convert a framework or legacy SDK tool without importing either backend SDK.""" + name = value.name + description = value.description + if not isinstance(name, str) or not name: + raise ValueError("A tool must have a nonempty string name") + if not isinstance(description, str): + raise TypeError("A tool description must be a string") + if not callable(value.handler): + raise TypeError(f"Tool {name!r} must have a callable handler") + annotations = getattr(value, "annotations", None) + if annotations is not None: + annotations = _mapping(annotations, label="Tool annotations") + meta = getattr(value, "meta", None) + if meta is not None: + meta = _mapping(meta, label="Tool metadata") + if annotations and annotations.get("maxResultSizeChars") is not None: + meta = dict(meta or {}) + meta.setdefault("anthropic/maxResultSizeChars", annotations["maxResultSizeChars"]) + return ToolDefinition( + name=name, + description=description, + input_schema=normalize_input_schema(value.input_schema), + handler=value.handler, + annotations=copy.deepcopy(annotations), + meta=copy.deepcopy(meta), + ) + + +def tool( + name: str, + description: str, + input_schema: Any, + annotations: Any = None, + *, + meta: Mapping[str, Any] | None = None, +) -> Callable[[ToolHandler], ToolDefinition]: + """Define an async tool using the same call shape as the Claude SDK decorator.""" + + def decorate(handler: ToolHandler) -> ToolDefinition: + return normalize_tool( + ToolDefinition(name, description, input_schema, handler, annotations, meta) + ) + + return decorate + + +def normalize_tool_result(result: Any) -> dict[str, Any]: + """Normalize MCP blocks and error flags, preserving images and top-level metadata.""" + result = _mapping(result, label="Tool result") + raw_content = result.get("content", []) + if raw_content is None: + raw_content = [] + if not isinstance(raw_content, (list, tuple)): + raise TypeError("Tool result content must be a list of content blocks") + content = [] + for raw_item in raw_content: + item = _mapping(raw_item, label="Tool result content block") + kind = item.get("type") + if kind == "text": + if not isinstance(item.get("text"), str): + raise ValueError("Tool text content requires a string text field") + content.append(item) + elif kind == "image": + content.append(item) + elif kind == "resource_link": + parts = [str(item[key]) for key in ("name", "uri", "description") if item.get(key)] + if not parts: + raise ValueError("A tool resource link needs a name, URI, or description") + content.append({"type": "text", "text": "\n".join(parts)}) + elif kind == "resource": + resource = _mapping(item.get("resource"), label="Embedded tool resource") + if not isinstance(resource.get("text"), str): + raise ValueError( + "Binary embedded tool resources are not supported; return a resource link" + ) + content.append({"type": "text", "text": resource["text"]}) + else: + raise ValueError(f"Unsupported tool result content type: {kind!r}") + if "structuredContent" in result: + # Dynamic tools and the Claude SDK's in-process adapter do not forward + # structured content; include it as JSON text so it remains visible. + content.append( + { + "type": "text", + "text": json.dumps(result.pop("structuredContent"), ensure_ascii=False), + } + ) + result["content"] = content + camel_case_error = result.pop("isError", False) + result["is_error"] = bool(result.get("is_error") or camel_case_error) + return result + + +def _image_url(item: dict[str, Any]) -> str: + data = item.get("data") + if isinstance(data, str) and data: + if data.startswith("data:"): + return data + mime_type = item.get("mimeType") + if not isinstance(mime_type, str) or not mime_type.startswith("image/"): + raise ValueError("A base64 tool image requires an image/* mimeType") + return f"data:{mime_type};base64,{data}" + image_url = item.get("image_url") + if isinstance(image_url, Mapping): + image_url = image_url.get("url") + if isinstance(image_url, str) and image_url: + return image_url + raise ValueError("A tool image requires base64 data and mimeType, or image_url") + + +def tool_result_to_codex(result: Any) -> dict[str, Any]: + """Translate MCP results into the Codex dynamic-tool response wire format.""" + normalized = normalize_tool_result(result) + items = [] + for item in normalized["content"]: + if item["type"] == "image": + items.append({"type": "inputImage", "imageUrl": _image_url(item)}) + else: + items.append({"type": "inputText", "text": item["text"]}) + return { + "contentItems": items or [{"type": "inputText", "text": ""}], + "success": not normalized["is_error"], + } + + +def tool_result_to_claude(result: Any) -> dict[str, Any]: + """Translate results into Claude SDK MCP format, retaining image bytes and MIME.""" + normalized = normalize_tool_result(result) + for index, item in enumerate(normalized["content"]): + if item["type"] != "image": + continue + url = _image_url(item) + if not url.startswith("data:") or ";base64," not in url: + raise ValueError( + "Claude SDK tool images require base64 data; remote image URLs are unsupported" + ) + header, data = url.split(";base64,", 1) + normalized["content"][index] = { + **item, + "data": data, + "mimeType": header.removeprefix("data:"), + } + return normalized diff --git a/agent-flow/agent_flow/types.py b/agent-flow/agent_flow/types.py index 786ca1a9f3ed..564b9aa0aa1b 100644 --- a/agent-flow/agent_flow/types.py +++ b/agent-flow/agent_flow/types.py @@ -24,6 +24,8 @@ class UsageInfo: context_tokens: int | None = None context_window: int | None = None context_percentage: float | None = None + # Optional cumulative thread estimate; never add this to per-turn cost_usd. + estimated_thread_cost_usd: float | None = None @dataclass diff --git a/agent-flow/agent_flow/workflows/agent_team/README.md b/agent-flow/agent_flow/workflows/agent_team/README.md index 78beda29be8c..4514bc93fc41 100644 --- a/agent-flow/agent_flow/workflows/agent_team/README.md +++ b/agent-flow/agent_flow/workflows/agent_team/README.md @@ -166,8 +166,9 @@ reads directly to drive control flow: - `append_qa_progress` (decisions: `APPROVE` / `REJECT`, plus `weighted_score`) The Coder and Reviewer additionally have `update_status` and -`read_status` tools. `update_status` is a required-tool call so the stop -hook enforces that `status.md` is refreshed alongside `progress.yaml`. +`read_status` tools. Both backends require the progress and `update_status` +calls. Missing calls get one corrective turn on the same session; a second +omission fails the invocation instead of silently continuing with stale state. The Coder, Reviewer, and QA each have a `read_human_feedback` tool that returns the entries in `progress.yaml`'s `human_feedback` list. QA is @@ -201,7 +202,7 @@ regardless of backend — the claude-code roles and the codex-backed PlanDrafter / Reviewer alike — so no dynamically configured server is ever created. -Each role then runs with `tools=None`, no required-tool stop hooks, and +Each role then runs with `tools=None`, no required-tool policy, and no `ask_human`; only the backend's own built-in tools remain (reading, editing, and running commands are not MCP servers, so they are unaffected). diff --git a/agent-flow/agent_flow/workflows/agent_team/progress.py b/agent-flow/agent_flow/workflows/agent_team/progress.py index e0dbab419751..8042ec1bd3e3 100644 --- a/agent-flow/agent_flow/workflows/agent_team/progress.py +++ b/agent-flow/agent_flow/workflows/agent_team/progress.py @@ -71,11 +71,11 @@ from typing import Any import yaml -from claude_agent_sdk import tool from rich.syntax import Syntax from agent_flow.console import print_layer_panel from agent_flow.logger import get_logger +from agent_flow.tools import tool PLAN_STAGE = "plan_stage" BUILD_STAGE = "build_stage" diff --git a/agent-flow/agent_flow/workflows/agent_team/status.py b/agent-flow/agent_flow/workflows/agent_team/status.py index 042c8ae285d3..ed903a9aabd2 100644 --- a/agent-flow/agent_flow/workflows/agent_team/status.py +++ b/agent-flow/agent_flow/workflows/agent_team/status.py @@ -40,11 +40,11 @@ from pathlib import Path from typing import Any -from claude_agent_sdk import tool from rich.markdown import Markdown from agent_flow.console import print_layer_panel from agent_flow.logger import get_logger +from agent_flow.tools import tool EMPTY_STATUS_PLACEHOLDER = ( "# (status.md is empty — no rolling state yet)\n" diff --git a/agent-flow/agent_flow/workflows/agent_team/workflow.py b/agent-flow/agent_flow/workflows/agent_team/workflow.py index df4938cc3a13..bf8f97a2aa13 100644 --- a/agent-flow/agent_flow/workflows/agent_team/workflow.py +++ b/agent-flow/agent_flow/workflows/agent_team/workflow.py @@ -10,7 +10,6 @@ AgentLayerConfig, BackendConfig, SessionConfig, - require_tool_call_stop_hook, ) from agent_flow.console import print_message, print_rule from agent_flow.logger import get_logger @@ -49,22 +48,6 @@ _REPLAN_STAGES = (STAGE_REPLAN, STAGE_REPLAN_REVIEWER, STAGE_REPLAN_HUMAN) -def _compose_required_tools_hooks(required_tools: list[str]) -> dict | None: - """Compose stop hooks that require *every* listed tool to be called. - - ``require_tool_call_stop_hook`` enforces "at least one of the listed - names was called". Stacking one such hook per tool — each independent — - yields AND semantics: every per-tool hook must allow the stop, so all - listed tools must have been called this turn. - """ - if not required_tools: - return None - merged: dict[str, list] = {"Stop": []} - for name in required_tools: - merged["Stop"].extend(require_tool_call_stop_hook([name])["Stop"]) - return merged - - def _make_agent( name: str, system_prompt: str, @@ -75,13 +58,13 @@ def _make_agent( session_mode: str = "persistent", human_input_enabled: bool = False, ) -> AgentLayer: - hooks = _compose_required_tools_hooks(required_tools or []) return AgentLayer( AgentLayerConfig( name=name, system_prompt=system_prompt, - backend=BackendConfig(kind=backend_kind, model=model, tools=tools, hooks=hooks), + backend=BackendConfig(kind=backend_kind, model=model, tools=tools), session=SessionConfig(mode=session_mode), + required_tools=tuple(required_tools or ()), human_input_enabled=human_input_enabled, ) ) @@ -341,7 +324,7 @@ def _build_agent( """Construct a role agent, honoring the in-process-tools mode. In no-in-process-MCP mode every role runs with ``tools=None``, no - required-tool Stop hooks, and no ``ask_human`` — so neither backend + required-tool policy, and no ``ask_human`` — so neither backend registers a dynamically configured MCP server. In MCP mode this is a passthrough to ``_make_agent``. """ diff --git a/agent-flow/agent_flow/workflows/perf_analyze/progress.py b/agent-flow/agent_flow/workflows/perf_analyze/progress.py index 435d4711f658..7bbf80460971 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/progress.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/progress.py @@ -57,11 +57,11 @@ from typing import Any import yaml -from claude_agent_sdk import tool from rich.syntax import Syntax from agent_flow.console import print_layer_panel from agent_flow.logger import get_logger +from agent_flow.tools import tool ANALYSIS_STAGE = "analysis" _AGENTS = ("benchmarker", "projector", "analyzer", "reporter") diff --git a/agent-flow/agent_flow/workflows/perf_analyze/workflow.py b/agent-flow/agent_flow/workflows/perf_analyze/workflow.py index aa9859991c1e..ef3cec87c943 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/workflow.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/workflow.py @@ -12,7 +12,6 @@ AgentLayerConfig, BackendConfig, SessionConfig, - require_tool_call_stop_hook, ) from agent_flow.console import print_message, print_rule from agent_flow.logger import get_logger @@ -101,22 +100,6 @@ def _progress_has_entries(path: Path) -> bool: return bool(data[ANALYSIS_STAGE]) -def _compose_required_tools_hooks(required_tools: list[str]) -> dict | None: - """Compose stop hooks that require *every* listed tool to be called. - - ``require_tool_call_stop_hook`` enforces "at least one of the listed - names was called". Stacking one such hook per tool — each independent - — yields AND semantics: every per-tool hook must allow the stop, so - all listed tools must have been called this turn. - """ - if not required_tools: - return None - merged: dict[str, list] = {"Stop": []} - for name in required_tools: - merged["Stop"].extend(require_tool_call_stop_hook([name])["Stop"]) - return merged - - def _make_agent( name: str, system_prompt: str, @@ -126,7 +109,6 @@ def _make_agent( model: str = CLAUDE_CODE_DEFAULT_MODEL, session_mode: str = "persistent", ) -> AgentLayer: - hooks = _compose_required_tools_hooks(required_tools or []) return AgentLayer( AgentLayerConfig( name=name, @@ -135,9 +117,9 @@ def _make_agent( kind=backend_kind, model=model, tools=tools, - hooks=hooks, ), session=SessionConfig(mode=session_mode), + required_tools=tuple(required_tools or ()), ) ) diff --git a/agent-flow/agent_flow/workflows/perf_optimize/progress.py b/agent-flow/agent_flow/workflows/perf_optimize/progress.py index fd98a22c77cb..64704cf43be5 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/progress.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/progress.py @@ -50,11 +50,11 @@ from typing import Any import yaml -from claude_agent_sdk import tool from rich.syntax import Syntax from agent_flow.console import print_layer_panel from agent_flow.logger import get_logger +from agent_flow.tools import tool OPTIMIZATION_STAGE = "optimization" _AGENTS = ( diff --git a/agent-flow/agent_flow/workflows/perf_optimize/workflow.py b/agent-flow/agent_flow/workflows/perf_optimize/workflow.py index c6a5027c1eb0..1c3baca92f43 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/workflow.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/workflow.py @@ -19,7 +19,6 @@ AgentLayerConfig, BackendConfig, SessionConfig, - require_tool_call_stop_hook, ) from agent_flow.console import print_message, print_rule from agent_flow.logger import get_logger @@ -95,22 +94,6 @@ def _progress_has_entries(path: Path) -> bool: return bool(data[OPTIMIZATION_STAGE]) -def _compose_required_tools_hooks(required_tools: list[str]) -> dict | None: - """Compose stop hooks that require *every* listed tool to be called. - - ``require_tool_call_stop_hook`` enforces "at least one of the listed - names was called". Stacking one such hook per tool — each independent - — yields AND semantics: every per-tool hook must allow the stop, so - all listed tools must have been called this turn. - """ - if not required_tools: - return None - merged: dict[str, list] = {"Stop": []} - for name in required_tools: - merged["Stop"].extend(require_tool_call_stop_hook([name])["Stop"]) - return merged - - def _make_agent( name: str, system_prompt: str, @@ -121,7 +104,6 @@ def _make_agent( session_mode: str = "persistent", cwd: Path | None = None, ) -> AgentLayer: - hooks = _compose_required_tools_hooks(required_tools or []) return AgentLayer( AgentLayerConfig( name=name, @@ -130,10 +112,10 @@ def _make_agent( kind=backend_kind, model=model, tools=tools, - hooks=hooks, cwd=cwd, ), session=SessionConfig(mode=session_mode), + required_tools=tuple(required_tools or ()), ) ) diff --git a/agent-flow/pyproject.toml b/agent-flow/pyproject.toml index 5cd5b1fc25cf..1fd595cc2812 100644 --- a/agent-flow/pyproject.toml +++ b/agent-flow/pyproject.toml @@ -8,8 +8,8 @@ dependencies = [ "anyio>=4.0", "pyyaml>=6.0", "rich>=13.0", - "claude-agent-sdk==0.2.143", - "openai-codex==0.147.0", + "claude-agent-sdk==0.2.152", + "openai-codex==0.154.0", ] [project.optional-dependencies] diff --git a/agent-flow/tests/conftest.py b/agent-flow/tests/conftest.py index b58e00f701ab..5d5867caa694 100644 --- a/agent-flow/tests/conftest.py +++ b/agent-flow/tests/conftest.py @@ -3,7 +3,31 @@ import sys from pathlib import Path +import pytest + ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) + + +def pytest_addoption(parser): + parser.addoption( + "--run-live-backends", + action="store_true", + help="Run integration checks that launch installed agent CLIs and use local account config.", + ) + + +def pytest_configure(config): + config.addinivalue_line("markers", "live_backend: requires a configured local agent CLI") + + +def pytest_collection_modifyitems(config, items): + if not config.getoption("--run-live-backends"): + skip = pytest.mark.skip( + reason="requires --run-live-backends and configured local agent CLIs" + ) + for item in items: + if "live_backend" in item.keywords: + item.add_marker(skip) diff --git a/agent-flow/tests/helpers.py b/agent-flow/tests/helpers.py index 0741e60e56cc..77196a6d4b4e 100644 --- a/agent-flow/tests/helpers.py +++ b/agent-flow/tests/helpers.py @@ -16,12 +16,16 @@ def __init__( error: Exception | None = None, usage: UsageInfo | None = None, skills: list[str] | None = None, + turns: list[dict] | None = None, + context_usage: UsageInfo | None = None, ) -> None: self.text = text self.tool_calls = tool_calls or [] self.error = error self.usage = usage self.skills = skills + self.turns = turns + self.context_usage = context_usage self.messages: list[str] = [] self.send_count = 0 self.closed = False @@ -29,14 +33,24 @@ def __init__( async def list_available_skills(self) -> list[str] | None: return self.skills + async def get_context_usage(self) -> UsageInfo | None: + return self.context_usage + async def send_message(self, message: str): self.messages.append(message) self.send_count += 1 - if self.error is not None: - raise self.error - for tool_call in self.tool_calls: - yield tool_call - yield ResultEvent(text=self.text, usage=self.usage) + turn = self.turns[min(self.send_count - 1, len(self.turns) - 1)] if self.turns else {} + error = turn.get("error", self.error) + if error is not None: + raise error + for event in turn.get("events", turn.get("tool_calls", self.tool_calls)): + yield event + yield ResultEvent( + text=turn.get("text", self.text), + usage=turn.get("usage", self.usage), + is_error=turn.get("is_error", False), + errors=turn.get("errors", []), + ) class FakeBackend(Backend): @@ -74,6 +88,8 @@ async def create_client( error=plan.get("error"), usage=plan.get("usage"), skills=plan.get("skills"), + turns=plan.get("turns"), + context_usage=plan.get("context_usage"), ) client.system_prompt = system_prompt client.model = model diff --git a/agent-flow/tests/test_agent_layer.py b/agent-flow/tests/test_agent_layer.py index 550b43cae776..63459135790d 100644 --- a/agent-flow/tests/test_agent_layer.py +++ b/agent-flow/tests/test_agent_layer.py @@ -4,6 +4,7 @@ from unittest.mock import patch import anyio +import pytest from agent_flow import ( AgentLayer, @@ -643,3 +644,229 @@ def observer(kind, event): def test_no_observer_by_default(): """Every existing layer keeps the behaviour it had.""" assert _config().on_activity is None + + +@pytest.mark.parametrize("backend_kind", ["claude-code", "codex"]) +@pytest.mark.parametrize("session_mode", ["stateless", "persistent"]) +async def test_required_tools_correct_on_same_client_and_keep_satisfied_calls( + backend_kind, + session_mode, +): + """One logical invocation can finish remaining calls in a second SDK turn.""" + backend = FakeBackend( + [ + { + "turns": [ + {"text": "initial", "tool_calls": [ToolCallEvent("append_progress", {})]}, + { + "text": "corrected", + "tool_calls": [ToolCallEvent("mcp__agent-tools__update_status", {})], + }, + ] + } + ] + ) + observed = [] + config = replace( + _config(session_mode=session_mode, print_activity=False), + backend=BackendConfig(kind=backend_kind, model="test-model"), + required_tools=("append_progress", "update_status"), + on_activity=lambda kind, event: observed.append((kind, event)), + ) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + assert await layer.aforward("do the work") == "corrected" + assert backend.create_client_calls == 1 + client = backend.clients[0] + assert client.closed + assert client.messages[0] == "do the work" + assert len(client.messages) == 2 + assert "`update_status`" in client.messages[1] + assert "append_progress" not in client.messages[1] + assert [event.name for kind, event in observed if kind == "tool"] == [ + "append_progress", + "mcp__agent-tools__update_status", + ] + + +async def test_required_tools_do_not_retry_a_satisfied_turn(): + backend = FakeBackend([{"tool_calls": [ToolCallEvent("append_progress", {})]}]) + config = replace(_config(print_activity=False), required_tools=("append_progress",)) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + assert await layer.aforward("record progress") == "ok" + assert backend.clients[0].send_count == 1 + + +async def test_required_tools_raise_after_one_correction_and_close_client(): + from agent_flow.hooks import RequiredToolCallError + + backend = FakeBackend([{"tool_calls": [ToolCallEvent("append_progress", {})]}]) + config = replace( + _config(print_activity=False), + required_tools=("append_progress", "update_status"), + ) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + with pytest.raises(RequiredToolCallError, match="one corrective turn: update_status"): + await layer.aforward("record progress and status") + assert backend.clients[0].send_count == 2 + assert backend.clients[0].closed + + +@pytest.mark.parametrize( + "child_fields", + [ + {"parent_tool_use_id": "spawn-1"}, + {"agent_label": "Reviewer"}, + ], +) +async def test_required_tools_exclude_subagent_calls(child_fields): + from agent_flow.hooks import RequiredToolCallError + + backend = FakeBackend( + [ + { + "tool_calls": [ToolCallEvent("append_progress", {}, **child_fields)], + } + ] + ) + config = replace(_config(print_activity=False), required_tools=("append_progress",)) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + with pytest.raises(RequiredToolCallError, match="append_progress"): + await layer.aforward("record your own progress") + assert backend.clients[0].send_count == 2 + + +@pytest.mark.parametrize( + "turn", + [ + {"error": RuntimeError("backend failed")}, + {"is_error": True, "errors": ["backend failed"]}, + ], +) +async def test_required_tools_propagate_backend_errors_without_correction(turn): + backend = FakeBackend([{"turns": [turn]}]) + config = replace(_config(print_activity=False), required_tools=("append_progress",)) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + with pytest.raises(RuntimeError, match="backend failed"): + await layer.aforward("record progress") + assert backend.clients[0].send_count == 1 + assert backend.clients[0].closed + + +async def test_required_tools_aggregate_turn_expense_but_keep_last_context(): + from agent_flow.logger import get_logger + + from .helpers import FakeClient + + client = FakeClient( + turns=[ + { + "text": "initial", + "usage": UsageInfo( + input_tokens=100, + output_tokens=20, + cache_creation_tokens=30, + cache_read_tokens=40, + total_tokens=120, + cost_usd=0.1, + num_turns=2, + duration_ms=1000, + context_tokens=120, + context_window=1000, + context_percentage=12, + ), + }, + { + "text": "corrected", + "tool_calls": [ToolCallEvent("append_progress", {})], + "usage": UsageInfo( + input_tokens=150, + output_tokens=10, + cache_creation_tokens=5, + cache_read_tokens=50, + total_tokens=160, + cost_usd=0.2, + num_turns=1, + duration_ms=200, + context_tokens=280, + context_window=1000, + context_percentage=28, + ), + }, + ] + ) + config = replace(_config(print_activity=False), required_tools=("append_progress",)) + async with AgentLayer(config) as layer: + response = await layer._execute( + client, AgentRequest("work", metadata={"task": 1}), get_logger() + ) + assert response.content == "corrected" + assert response.metadata == {"task": 1} + assert response.usage == UsageInfo( + input_tokens=250, + output_tokens=30, + cache_creation_tokens=35, + cache_read_tokens=90, + total_tokens=280, + cost_usd=pytest.approx(0.3), + num_turns=3, + duration_ms=1200, + context_tokens=280, + context_window=1000, + context_percentage=28, + ) + + +@pytest.mark.parametrize("known_usage", [None, UsageInfo(context_tokens=25, context_window=100)]) +def test_fetch_baseline_never_sends_a_model_turn(known_usage): + backend = FakeBackend([{"context_usage": known_usage}]) + with patch("agent_flow.layers.create_backend", return_value=backend): + with AgentLayer(_config()) as layer: + assert layer.fetch_baseline_context_usage() is known_usage + assert backend.clients[0].send_count == 0 + assert backend.clients[0].closed + + +@pytest.mark.parametrize("session_mode", ["stateless", "persistent"]) +async def test_session_modes_resolve_the_same_client_options(session_mode, tmp_path): + backend = FakeBackend() + config = replace( + _config(session_mode=session_mode, print_activity=False, human_input_enabled=True), + backend=BackendConfig( + kind="claude-code", + model="test-model", + tools=["sentinel-tool"], + hooks={"Stop": []}, + extra_mcp_servers={"test": {"command": "test"}}, + cwd=tmp_path, + ), + disallowed_tools=("Write",), + ) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + await layer.aforward("work") + client = backend.clients[0] + assert client.system_prompt == "You are helpful." + assert client.model == "test-model" + assert client.tools[0].name == "ask_human" + assert client.tools[1:] == ["sentinel-tool"] + assert client.hooks == {"Stop": []} + assert client.disallowed_tools == ["Write", "AskUserQuestion"] + assert client.extra_mcp_servers == {"test": {"command": "test"}} + assert client.cwd == tmp_path + + +@pytest.mark.parametrize("final_estimate, expected", [(0.3, 0.3), (None, 0.2), (0.0, 0.0)]) +def test_usage_carries_latest_available_thread_cost_estimate(final_estimate, expected): + from agent_flow.layers import _merge_turn_usage + + usage = _merge_turn_usage( + UsageInfo(input_tokens=10, estimated_thread_cost_usd=0.2), + UsageInfo(input_tokens=20, estimated_thread_cost_usd=final_estimate), + ) + assert usage.input_tokens == 30 + assert usage.estimated_thread_cost_usd == expected diff --git a/agent-flow/tests/test_backends.py b/agent-flow/tests/test_backends.py index c429e0307db4..fa237f557d01 100644 --- a/agent-flow/tests/test_backends.py +++ b/agent-flow/tests/test_backends.py @@ -2,10 +2,12 @@ import sys from pathlib import Path -from types import ModuleType, SimpleNamespace +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock +import claude_agent_sdk +import jsonschema import pytest from claude_agent_sdk import types as sdk_types from claude_agent_sdk.types import ( @@ -24,13 +26,7 @@ from agent_flow.backends import create_backend from agent_flow.backends.base import ResultEvent from agent_flow.backends.claude_code import ClaudeCodeBackend, ClaudeCodeClient -from agent_flow.backends.codex import ( - CodexBackend, - CodexClient, - _dynamic_tool_spec, - _extract_final_response, - _mcp_to_codex_content, -) +from agent_flow.backends.codex import CodexBackend from agent_flow.types import ( AgentTextEvent, CompactBoundaryEvent, @@ -78,265 +74,6 @@ def _make_tool_use_block(name: str, input: dict, id: str = "tool-1") -> ToolUseB return ToolUseBlock(id=id, name=name, input=input) -def _install_codex_sdk_modules(): - package = ModuleType("openai_codex") - client_module = ModuleType("openai_codex.client") - api_module = ModuleType("openai_codex.api") - generated = ModuleType("openai_codex.generated.v2_all") - - class TextInput: - def __init__(self, text): - self.text = text - - class AsyncCodex: - def __init__(self, config): - self.config = config - - async def __aenter__(self): - return self - - async def __aexit__(self, *args): - return None - - class CodexConfig: - def __init__(self, **kwargs): - self.kwargs = kwargs - - class AgentMessageThreadItem: - def __init__(self, text, phase=None): - self.text = text - self.phase = phase - - class AskForApproval: - def __init__(self, root): - self.root = root - - class AskForApprovalValue: - untrusted = "untrusted" - never = "never" - - class CommandExecutionThreadItem: - def __init__(self, command, id="cmd-1", status="completed"): - self.id = id - self.command = command - self.status = status - - class CollabAgentToolCallThreadItem: - def __init__( - self, - tool, - id="collab-1", - status="completed", - prompt=None, - model=None, - reasoning_effort=None, - sender_thread_id="sender-1", - receiver_thread_ids=None, - agents_states=None, - ): - self.id = id - self.tool = tool - self.status = status - self.prompt = prompt - self.model = model - self.reasoning_effort = reasoning_effort - self.sender_thread_id = sender_thread_id - self.receiver_thread_ids = receiver_thread_ids or [] - self.agents_states = agents_states or {} - - class ContextCompactionThreadItem: - def __init__(self, id="compact-1", trigger=None, pre_tokens=None): - self.id = id - self.trigger = trigger - self.pre_tokens = pre_tokens - - class DynamicToolCallThreadItem: - def __init__(self, tool, arguments, id="dyn-1", status="completed"): - self.id = id - self.tool = tool - self.arguments = arguments - self.status = status - - class FileChangeThreadItem: - def __init__(self, changes, id="file-1", status="completed", error=None): - self.id = id - self.changes = changes - self.status = status - self.error = error - - class ItemCompletedNotification: - def __init__(self, item): - self.item = item - - class ItemStartedNotification: - def __init__(self, item): - self.item = item - - class McpToolCallThreadItem: - def __init__(self, tool, arguments, id="mcp-1", status="completed", error=None): - self.id = id - self.tool = tool - self.arguments = arguments - self.status = status - self.error = error - - class ReasoningThreadItem: - def __init__(self, summary=None, content=None, id="reason-1"): - self.id = id - self.summary = summary or [] - self.content = content or [] - - class WebSearchThreadItem: - def __init__(self, query, action=None, id="web-1"): - self.id = id - self.query = query - self.action = action - - class MessagePhase: - final_answer = "final_answer" - - class SandboxMode: - workspace_write = "workspace_write" - danger_full_access = "danger_full_access" - - class ThreadItem: - pass - - class TurnCompletedNotification: - def __init__(self, turn): - self.turn = turn - - class TurnStatus: - completed = "completed" - failed = "failed" - in_progress = "inProgress" - interrupted = "interrupted" - - class TokenUsageBreakdown: - def __init__( - self, - input_tokens=0, - output_tokens=0, - cached_input_tokens=0, - total_tokens=0, - reasoning_output_tokens=0, - ): - self.input_tokens = input_tokens - self.output_tokens = output_tokens - self.cached_input_tokens = cached_input_tokens - self.total_tokens = total_tokens - self.reasoning_output_tokens = reasoning_output_tokens - - class ThreadTokenUsage: - def __init__(self, last, total, model_context_window=None): - self.last = last - self.total = total - self.model_context_window = model_context_window - - class ThreadTokenUsageUpdatedNotification: - def __init__(self, thread_id, token_usage, turn_id): - self.thread_id = thread_id - self.token_usage = token_usage - self.turn_id = turn_id - - class ThreadStartParams: - def __init__( - self, - model=None, - base_instructions=None, - developer_instructions=None, - config=None, - cwd=None, - sandbox=None, - approval_policy=None, - **kwargs, - ): - self.model = model - self.base_instructions = base_instructions - self.developer_instructions = developer_instructions - self.config = config - self.cwd = cwd - self.sandbox = sandbox - self.approval_policy = approval_policy - - def _params_dict(params): - out = {} - if params.model is not None: - out["model"] = params.model - if params.base_instructions is not None: - out["baseInstructions"] = params.base_instructions - if params.developer_instructions is not None: - out["developerInstructions"] = params.developer_instructions - if params.config is not None: - out["config"] = params.config - if params.cwd is not None: - out["cwd"] = params.cwd - if params.sandbox is not None: - out["sandbox"] = params.sandbox - if params.approval_policy is not None: - out["approvalPolicy"] = params.approval_policy - return out - - class AsyncThread: - def __init__(self, codex, thread_id): - self.codex = codex - self.thread_id = thread_id - - package.TextInput = TextInput - package.AsyncCodex = AsyncCodex - client_module.CodexConfig = CodexConfig - client_module._params_dict = _params_dict - api_module.AsyncThread = AsyncThread - generated.AgentMessageThreadItem = AgentMessageThreadItem - generated.AskForApproval = AskForApproval - generated.AskForApprovalValue = AskForApprovalValue - generated.CollabAgentToolCallThreadItem = CollabAgentToolCallThreadItem - generated.CommandExecutionThreadItem = CommandExecutionThreadItem - generated.ContextCompactionThreadItem = ContextCompactionThreadItem - generated.DynamicToolCallThreadItem = DynamicToolCallThreadItem - generated.FileChangeThreadItem = FileChangeThreadItem - generated.ItemCompletedNotification = ItemCompletedNotification - generated.ItemStartedNotification = ItemStartedNotification - generated.McpToolCallThreadItem = McpToolCallThreadItem - generated.MessagePhase = MessagePhase - generated.ReasoningThreadItem = ReasoningThreadItem - generated.SandboxMode = SandboxMode - generated.ThreadItem = ThreadItem - generated.TurnCompletedNotification = TurnCompletedNotification - generated.WebSearchThreadItem = WebSearchThreadItem - generated.TokenUsageBreakdown = TokenUsageBreakdown - generated.ThreadTokenUsage = ThreadTokenUsage - generated.ThreadTokenUsageUpdatedNotification = ThreadTokenUsageUpdatedNotification - generated.ThreadStartParams = ThreadStartParams - generated.TurnStatus = TurnStatus - - sys.modules["openai_codex"] = package - sys.modules["openai_codex.client"] = client_module - sys.modules["openai_codex.api"] = api_module - sys.modules["openai_codex.generated.v2_all"] = generated - - return { - "TextInput": TextInput, - "AgentMessageThreadItem": AgentMessageThreadItem, - "CollabAgentToolCallThreadItem": CollabAgentToolCallThreadItem, - "CommandExecutionThreadItem": CommandExecutionThreadItem, - "ContextCompactionThreadItem": ContextCompactionThreadItem, - "DynamicToolCallThreadItem": DynamicToolCallThreadItem, - "FileChangeThreadItem": FileChangeThreadItem, - "ItemCompletedNotification": ItemCompletedNotification, - "ItemStartedNotification": ItemStartedNotification, - "McpToolCallThreadItem": McpToolCallThreadItem, - "MessagePhase": MessagePhase, - "ReasoningThreadItem": ReasoningThreadItem, - "TurnCompletedNotification": TurnCompletedNotification, - "WebSearchThreadItem": WebSearchThreadItem, - "TokenUsageBreakdown": TokenUsageBreakdown, - "ThreadTokenUsage": ThreadTokenUsage, - "ThreadTokenUsageUpdatedNotification": ThreadTokenUsageUpdatedNotification, - "TurnStatus": TurnStatus, - } - - class TestCreateBackend: def test_factory_returns_expected_backend_types(self): assert create_backend("claude-code").__class__.__name__ == "ClaudeCodeBackend" @@ -987,7 +724,79 @@ async def receive_response(): assert result.permission_denials == [] +# A recursive explicit JSON Schema: the form ``normalize_input_schema`` points users +# to when a TypedDict refers to itself. +_RECURSIVE_NODE_SCHEMA = { + "$ref": "#/$defs/node", + "$defs": { + "node": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "children": {"type": "array", "items": {"$ref": "#/$defs/node"}}, + }, + "required": ["name"], + } + }, +} + + class TestClaudeBackendCreateClient: + async def test_framework_tools_keep_annotations_and_independent_handlers(self): + from agent_flow.tools import tool + + @tool("resource", "Read a resource", {"path": str}, annotations={"readOnlyHint": True}) + async def resource(args): + return {"content": [{"type": "resource", "resource": {"text": args["path"]}}]} + + @tool("image", "Read an image", {}, annotations={"maxResultSizeChars": 1024}) + async def image(args): + return {"content": [{"type": "image", "data": "YWJj", "mimeType": "image/png"}]} + + resource_sdk, image_sdk = cc_mod._sdk_tools([resource, image]) + assert resource_sdk.input_schema["properties"]["path"] == {"type": "string"} + assert resource_sdk.annotations.model_dump(by_alias=True)["readOnlyHint"] is True + assert image_sdk.annotations.maxResultSizeChars == 1024 + assert (await resource_sdk.handler({"path": "document"}))["content"] == [ + {"type": "text", "text": "document"} + ] + assert (await image_sdk.handler({}))["content"] == [ + {"type": "image", "data": "YWJj", "mimeType": "image/png"} + ] + + @pytest.mark.parametrize( + "schema, accepted, rejected", + [ + ( + {"type": "object", "additionalProperties": {"type": "string"}}, + {"key": "value"}, + {"key": 1}, + ), + ( + _RECURSIVE_NODE_SCHEMA, + {"name": "root", "children": [{"name": "leaf"}]}, + {"name": "root", "children": [{"children": []}]}, + ), + ], + ) + def test_explicit_json_schema_survives_sdk_schema_builder(self, schema, accepted, rejected): + from agent_flow.tools import tool + + @tool("explicit", "Explicit JSON Schema", schema) + async def explicit(args): + return {"content": []} + + [sdk_tool] = cc_mod._sdk_tools([explicit]) + # ``_build_input_schema`` is the pinned SDK's wire-schema builder. A dict without a + # string ``type`` and a ``properties`` key is re-read as Python shorthand, so the + # schema's own keywords would be advertised, and then validated on every call, as + # required parameters. The framework must hand it a shape it passes through as is. + advertised = claude_agent_sdk._build_input_schema(sdk_tool) + assert advertised is sdk_tool.input_schema + jsonschema.validate(instance=accepted, schema=advertised) + with pytest.raises(jsonschema.ValidationError): + jsonschema.validate(instance=rejected, schema=advertised) + async def _capture_options(self, monkeypatch, **kwargs): # Stand-in for ``ClaudeSDKClient`` that just records the options # ``create_client`` would have launched the real SDK with. @@ -1027,12 +836,10 @@ async def test_create_client_uses_bypass_permission_mode(self, monkeypatch): options = await self._capture_options(monkeypatch) assert options.permission_mode == "bypassPermissions" - async def test_create_client_can_use_tool_always_allows(self, monkeypatch): - # Defense-in-depth: even if a tool ends up being routed through - # the ``can_use_tool`` hook, it must approve every call. + async def test_create_client_omits_shadowed_permission_callback(self, monkeypatch): + # The pinned SDK bypasses this callback in bypassPermissions mode. options = await self._capture_options(monkeypatch) - result = await options.can_use_tool("Bash", {"command": "rm -rf /"}, None) - assert isinstance(result, cc_mod.PermissionResultAllow) + assert options.can_use_tool is None async def test_create_client_no_extra_mcp_servers_by_default(self, monkeypatch): # Without ``extra_mcp_servers``, ``mcp_servers`` stays empty @@ -1047,7 +854,14 @@ async def test_create_client_extra_mcp_servers_merged_alongside_agent_tools(self # ``agent-tools`` server is layered on top when ``tools`` is set. options = await self._capture_options( monkeypatch, - tools=[object()], + tools=[ + SimpleNamespace( + name="test", + description="test", + input_schema={"type": "object", "properties": {}}, + handler=AsyncMock(), + ) + ], extra_mcp_servers={ "knowledge-base": {"type": "http", "url": "https://example.test/mcp"}, }, @@ -1091,803 +905,6 @@ async def test_create_client_rejects_reserved_agent_tools_key(self, monkeypatch) pass -class TestCodexBackend: - async def test_client_maps_command_and_final_response(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - cmd_item = sdk["CommandExecutionThreadItem"]("echo hello") - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=cmd_item)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hello")] - - thread.turn.assert_awaited_once() - assert events[0] == ToolCallEvent( - name="Bash", input={"command": "echo hello"}, tool_use_id="cmd-1" - ) - assert isinstance(events[1], ResultEvent) - - async def test_client_captures_token_usage_from_notifications(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - total = sdk["TokenUsageBreakdown"]( - input_tokens=100000, - output_tokens=50000, - cached_input_tokens=20000, - total_tokens=170000, - ) - usage = sdk["ThreadTokenUsage"]( - last=total, - total=total, - model_context_window=400000, - ) - yield SimpleNamespace( - payload=sdk["ThreadTokenUsageUpdatedNotification"]( - thread_id="t-1", - token_usage=usage, - turn_id="turn-1", - ) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"]( - SimpleNamespace(id="turn-1", duration_ms=2500) - ) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - assert len(events) == 1 - result = events[0] - assert isinstance(result, ResultEvent) - assert result.usage is not None - assert result.usage.input_tokens == 100000 - assert result.usage.output_tokens == 50000 - assert result.usage.cache_read_tokens == 20000 - assert result.usage.total_tokens == 170000 - assert result.usage.context_tokens == 170000 - assert result.usage.context_window == 400000 - assert result.usage.context_percentage == pytest.approx(42.5) - assert result.usage.num_turns == 1 - assert result.usage.duration_ms == 2500 - - async def test_client_surfaces_turn_duration_without_token_usage(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"]( - SimpleNamespace(id="turn-1", duration_ms=3200) - ) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - assert len(events) == 1 - result = events[0] - assert isinstance(result, ResultEvent) - assert result.usage is not None - assert result.usage.num_turns == 1 - assert result.usage.duration_ms == 3200 - - async def test_client_emits_final_answer_as_agent_text_event(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - draft = sdk["AgentMessageThreadItem"]("draft") - final = sdk["AgentMessageThreadItem"]( - "final answer", phase=sdk["MessagePhase"].final_answer - ) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=draft)) - ) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=final)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - text_events = [e for e in events if isinstance(e, AgentTextEvent)] - assert [e.text for e in text_events] == ["draft", "final answer"] - result_events = [e for e in events if isinstance(e, ResultEvent)] - assert len(result_events) == 1 - assert result_events[0].text == "final answer" - - async def test_client_ignores_token_deltas_and_emits_completed_message(self): - sdk = _install_codex_sdk_modules() - - class AgentMessageDeltaNotification: - def __init__(self, itemId, delta): - self.itemId = itemId - self.delta = delta - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - yield SimpleNamespace(payload=AgentMessageDeltaNotification("msg-1", "A")) - yield SimpleNamespace(payload=AgentMessageDeltaNotification("msg-1", "B")) - final = sdk["AgentMessageThreadItem"]("AB", phase=sdk["MessagePhase"].final_answer) - final.id = "msg-1" - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=final)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - text_events = [e for e in events if isinstance(e, AgentTextEvent)] - assert [e.text for e in text_events] == ["AB"] - - async def test_client_maps_codex_reasoning_to_thinking_event(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - reasoning = sdk["ReasoningThreadItem"]( - summary=["checking tests"], content=["raw detail"] - ) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=reasoning)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - thinking = [e for e in events if isinstance(e, ThinkingEvent)] - assert len(thinking) == 1 - assert thinking[0].text == "checking tests\nraw detail" - - async def test_client_maps_codex_web_search_to_server_tool_event(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - search = sdk["WebSearchThreadItem"]( - query="codex app-server", - action={"type": "search", "query": "codex app-server"}, - id="web-7", - ) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=search)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - server_calls = [e for e in events if isinstance(e, ServerToolCallEvent)] - assert len(server_calls) == 1 - assert server_calls[0].name == "web_search" - assert server_calls[0].tool_use_id == "web-7" - assert server_calls[0].input == { - "type": "search", - "query": "codex app-server", - } - - async def test_client_maps_codex_collab_agent_tool_calls_and_failures(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - collab = sdk["CollabAgentToolCallThreadItem"]( - tool="spawnAgent", - id="collab-7", - status="failed", - prompt="inspect repo", - model="gpt-test", - reasoning_effort="high", - receiver_thread_ids=["child-1"], - ) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=collab)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - tool_call = next(e for e in events if isinstance(e, ToolCallEvent)) - assert tool_call.name == "spawnAgent" - assert tool_call.tool_use_id == "collab-7" - assert tool_call.input["prompt"] == "inspect repo" - assert tool_call.input["model"] == "gpt-test" - assert tool_call.input["reasoning_effort"] == "high" - assert tool_call.input["receiver_thread_ids"] == ["child-1"] - - result = next(e for e in events if isinstance(e, ResultEvent)) - assert result.is_error is True - assert result.errors == ["collabAgentToolCall failed"] - - async def test_client_emits_compact_boundary_for_context_compaction_item(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - compact = sdk["ContextCompactionThreadItem"](trigger="auto", pre_tokens=150000) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=compact)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - boundaries = [e for e in events if isinstance(e, CompactBoundaryEvent)] - assert len(boundaries) == 1 - assert boundaries[0].trigger == "auto" - assert boundaries[0].pre_tokens == 150000 - - async def test_client_maps_codex_file_changes_and_denials(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - file_change = sdk["FileChangeThreadItem"]( - changes=[{"path": "app.py", "kind": "update"}], - status="declined", - id="file-9", - ) - yield SimpleNamespace( - payload=sdk["ItemCompletedNotification"](SimpleNamespace(root=file_change)) - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [event async for event in client.send_message("hi")] - - file_event = next(e for e in events if isinstance(e, ToolCallEvent)) - assert file_event.name == "FileChange" - assert file_event.tool_use_id == "file-9" - assert file_event.input == { - "changes": [{"path": "app.py", "kind": "update"}], - "status": "declined", - } - result = next(e for e in events if isinstance(e, ResultEvent)) - assert result.permission_denials[0]["kind"] == "fileChange" - assert result.permission_denials[0]["id"] == "file-9" - - async def test_client_emits_rate_limit_warning_before_turn_error(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - error = SimpleNamespace( - message="usage limit reached", - codex_error_info="UsageLimitExceeded", - ) - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"]( - SimpleNamespace( - id="turn-1", - error=error, - ) - ) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - events = [] - with pytest.raises(RuntimeError, match="usage limit reached"): - async for event in client.send_message("hi"): - events.append(event) - - warnings = [e for e in events if isinstance(e, RateLimitWarningEvent)] - assert len(warnings) == 1 - assert warnings[0].status == "rejected" - assert warnings[0].rate_limit_type == "UsageLimitExceeded" - - async def test_client_raises_when_turn_completes_with_error(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"]( - SimpleNamespace( - id="turn-1", - error=SimpleNamespace(message="model does not exist"), - ) - ) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - - with pytest.raises(RuntimeError, match="Codex turn failed: model does not exist"): - async for _ in client.send_message("hi"): - pass - - async def test_client_raises_when_turn_completes_interrupted(self): - sdk = _install_codex_sdk_modules() - - thread = MagicMock() - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"]( - SimpleNamespace( - id="turn-1", - status=sdk["TurnStatus"].interrupted, - ) - ) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - - client = CodexClient(thread) - - with pytest.raises(RuntimeError, match="status 'interrupted'"): - async for _ in client.send_message("hi"): - pass - - async def test_enter_passes_resolved_binary_to_app_server_config(self, monkeypatch): - _install_codex_sdk_modules() - - monkeypatch.setattr(codex_mod, "_resolve_codex_bin", lambda: "/bin/codex") - monkeypatch.setattr(CodexBackend, "_install_server_request_handler", lambda self: None) - - backend = CodexBackend() - async with backend: - config = backend._codex.config - - assert config.kwargs["codex_bin"] == "/bin/codex" - assert config.kwargs["experimental_api"] is True - assert "launch_args_override" not in config.kwargs - - def test_extract_final_response_prefers_final_answer_phase(self): - sdk = _install_codex_sdk_modules() - - items = [ - sdk["AgentMessageThreadItem"]("draft"), - sdk["AgentMessageThreadItem"]("final", phase=sdk["MessagePhase"].final_answer), - ] - - assert _extract_final_response(items) == "final" - - async def _run_create_client(self, system_prompt: str) -> dict[str, Any]: - _install_codex_sdk_modules() - - backend = CodexBackend() - captured: dict[str, Any] = {} - - async def fake_thread_start(payload): - captured["payload"] = payload - return SimpleNamespace(thread=SimpleNamespace(id="t-1")) - - backend._codex = SimpleNamespace( - _ensure_initialized=AsyncMock(), - _client=SimpleNamespace(thread_start=fake_thread_start), - ) - - async with backend.create_client(system_prompt=system_prompt, model="gpt-5.4"): - pass - - return captured["payload"] - - async def test_create_client_uses_codex_default_system_prompt(self): - # Always defer to Codex's bundled base prompt — never override it - # via base_instructions. An empty user prompt must not become an - # empty developer_instructions either, otherwise Codex forwards an - # empty instructions payload to the Responses API and the request - # fails with 400 "Instructions are required". - payload = await self._run_create_client(system_prompt="") - assert "baseInstructions" not in payload - assert "developerInstructions" not in payload - - async def test_create_client_appends_user_prompt_via_developer_message(self): - # User-provided prompts ride on developer_instructions so they - # layer on top of Codex's default base prompt, mirroring how the - # Claude Code backend appends to its preset. - payload = await self._run_create_client(system_prompt="be helpful") - assert "baseInstructions" not in payload - assert payload["developerInstructions"] == "be helpful" - - async def test_create_client_overrides_model_context_window_to_1m(self): - # The bundled codex CLI's models.json caps gpt-5.5 at 272k; without - # this override codex auto-compacts at ~258k even though the API - # serves the full 1M context. - payload = await self._run_create_client(system_prompt="hi") - assert payload["config"]["model_context_window"] == 1000000 - - async def test_create_client_requests_max_reasoning_effort(self): - # Codex threads run at the top reasoning tier the CLI exposes. - # ``max`` sits above ``xhigh`` in the CLI's effort ladder - # (minimal/low/medium/high/xhigh/max/ultra); the generated SDK enum - # still stops at ``xhigh`` but passes unknown values through, so the - # string must reach ``thread/start`` verbatim. - payload = await self._run_create_client(system_prompt="hi") - assert payload["config"]["model_reasoning_effort"] == "max" - - async def test_create_client_bypasses_all_approvals_in_thread_start(self): - # The Codex backend runs in fully autonomous mode: no sandbox - # restrictions and the model never asks for approval. The mock - # captures the params that would be sent to ``thread/start``. - payload = await self._run_create_client(system_prompt="hi") - assert payload["sandbox"] == "danger_full_access" - assert payload["approvalPolicy"].root == "never" - - async def test_create_client_accepts_extra_mcp_servers_as_noop(self): - # ``extra_mcp_servers`` is currently a Claude-Code-only concept; - # the Codex backend must accept the kwarg without erroring and - # without leaking it into ``thread/start``. - _install_codex_sdk_modules() - backend = CodexBackend() - captured: dict[str, Any] = {} - - async def fake_thread_start(payload): - captured["payload"] = payload - return SimpleNamespace(thread=SimpleNamespace(id="t-1")) - - backend._codex = SimpleNamespace( - _ensure_initialized=AsyncMock(), - _client=SimpleNamespace(thread_start=fake_thread_start), - ) - - async with backend.create_client( - system_prompt="hi", - model="gpt-5.4", - extra_mcp_servers={ - "knowledge-base": { - "type": "http", - "url": "https://example.test/mcp", - }, - }, - ): - pass - - # Codex thread/start payload must not carry MCP-server config. - payload = captured["payload"] - assert "mcp_servers" not in payload - assert "mcpServers" not in payload - - async def test_create_client_emits_codex_session_init_snapshot(self): - sdk = _install_codex_sdk_modules() - - backend = CodexBackend() - - async def fake_thread_start(payload): - return SimpleNamespace(thread=SimpleNamespace(id="t-1")) - - async def fake_skills_list(payload): - return SimpleNamespace( - data=[ - SimpleNamespace( - skills=[ - SimpleNamespace(name="skill-a", enabled=True), - SimpleNamespace(name="skill-disabled", enabled=False), - ] - ) - ] - ) - - async def fake_plugin_list(payload): - return SimpleNamespace( - plugins=[ - SimpleNamespace(name="plugin-a"), - SimpleNamespace(name="plugin-b"), - ] - ) - - backend._codex = SimpleNamespace( - _ensure_initialized=AsyncMock(), - _client=SimpleNamespace( - skills_list=fake_skills_list, - plugin_list=fake_plugin_list, - thread_start=fake_thread_start, - ), - ) - - async with backend.create_client(system_prompt="hi", model="gpt-5.4") as client: - thread = client._thread - turn = MagicMock() - turn.id = "turn-1" - - async def stream(): - yield SimpleNamespace( - payload=sdk["TurnCompletedNotification"](SimpleNamespace(id="turn-1")) - ) - - turn.stream = stream - thread.turn = AsyncMock(return_value=turn) - events = [event async for event in client.send_message("hi")] - - init = next(e for e in events if isinstance(e, SessionInitEvent)) - assert init.skills == ["skill-a"] - assert init.plugins == ["plugin-a", "plugin-b"] - - async def test_client_lists_skills_from_the_creation_time_session_init(self): - # ``skills_list`` already ran when the client was created, so the - # answer is in hand — reading it must not start a turn. - thread = MagicMock() - thread.turn = AsyncMock() - client = CodexClient( - thread, - session_init=SessionInitEvent(skills=["skill-a"], plugins=["plugin-a"]), - ) - - assert await client.list_available_skills() == ["skill-a"] - thread.turn.assert_not_awaited() - - async def test_client_reports_no_skill_list_without_a_session_init(self): - assert await CodexClient(MagicMock()).list_available_skills() is None - - -class TestCodexApprovalBypass: - def _make_backend_with_handler(self): - # Wire a CodexBackend up to a fake AsyncCodex whose sync client - # exposes ``_approval_handler`` so we can exercise the override - # ``_install_server_request_handler`` installs. - - captured_methods: list[tuple[str, dict[str, Any] | None]] = [] - - def existing_handler(method, params): - captured_methods.append((method, params)) - return {"sentinel": "fallthrough"} - - sync_client = SimpleNamespace(_approval_handler=existing_handler) - backend = CodexBackend() - backend._codex = SimpleNamespace(_client=SimpleNamespace(_sync=sync_client)) - backend._install_server_request_handler() - return backend, sync_client, captured_methods - - def test_install_handler_auto_accepts_command_approval(self): - _, sync_client, captured = self._make_backend_with_handler() - result = sync_client._approval_handler( - "item/commandExecution/requestApproval", - {"command": "rm -rf /tmp/everything"}, - ) - assert result == {"decision": "accept"} - # Auto-accept must short-circuit the SDK's default handler. - assert captured == [] - - def test_install_handler_auto_accepts_file_change_approval(self): - _, sync_client, captured = self._make_backend_with_handler() - result = sync_client._approval_handler( - "item/fileChange/requestApproval", - {"changes": [{"path": "x", "kind": "delete"}]}, - ) - assert result == {"decision": "accept"} - assert captured == [] - - def test_install_handler_auto_accepts_unknown_future_approval(self): - # Future SDK versions could add e.g. ``item/applyPatch/requestApproval`` - # — the bypass is intentionally generic on the ``/requestApproval`` - # suffix so new approval kinds are auto-accepted without code changes. - _, sync_client, captured = self._make_backend_with_handler() - result = sync_client._approval_handler( - "item/applyPatch/requestApproval", - {"diff": "..."}, - ) - assert result == {"decision": "accept"} - assert captured == [] - - def test_install_handler_routes_dynamic_tool_call(self, monkeypatch): - _, sync_client, captured = self._make_backend_with_handler() - - captured_tool_params: list[dict[str, Any]] = [] - - def fake_handle(params): - captured_tool_params.append(params) - return {"contentItems": [], "success": True} - - monkeypatch.setattr(codex_mod, "_handle_dynamic_tool_call", fake_handle) - result = sync_client._approval_handler( - "item/tool/call", - {"threadId": "t-1", "tool": "x"}, - ) - assert result == {"contentItems": [], "success": True} - assert captured_tool_params == [{"threadId": "t-1", "tool": "x"}] - # Dynamic-tool dispatch must not also fall through to the SDK - # default handler. - assert captured == [] - - def test_install_handler_falls_through_for_unrelated_methods(self): - _, sync_client, captured = self._make_backend_with_handler() - result = sync_client._approval_handler("some/other/method", {"k": "v"}) - # Unknown non-approval methods defer to the SDK's default handler. - assert result == {"sentinel": "fallthrough"} - assert captured == [("some/other/method", {"k": "v"})] - - -class TestCodexDynamicTools: - def test_dynamic_tool_spec_serializes_sdk_tool(self): - tool = SimpleNamespace( - name="do_thing", - description="Does a thing.", - input_schema={"type": "object", "properties": {}}, - ) - - assert _dynamic_tool_spec(tool) == { - "name": "do_thing", - "description": "Does a thing.", - "inputSchema": {"type": "object", "properties": {}}, - "deferLoading": False, - } - - def test_mcp_to_codex_content_maps_text_and_errors(self): - items, success = _mcp_to_codex_content( - { - "content": [{"type": "text", "text": "ok"}], - } - ) - assert items == [{"type": "inputText", "text": "ok"}] - assert success is True - - items, success = _mcp_to_codex_content( - { - "content": [{"type": "text", "text": "boom"}], - "is_error": True, - } - ) - assert items == [{"type": "inputText", "text": "boom"}] - assert success is False - - def test_mcp_to_codex_content_emits_placeholder_when_empty(self): - items, success = _mcp_to_codex_content({"content": []}) - assert items == [{"type": "inputText", "text": ""}] - assert success is True - - def test_handle_dynamic_tool_call_dispatches_to_registered_handler(self): - calls: list[dict[str, Any]] = [] - - async def handler(args): - calls.append(args) - return {"content": [{"type": "text", "text": f"got {args['x']}"}]} - - codex_mod._TOOL_HANDLERS["t-1"] = {"my_tool": handler} - try: - response = codex_mod._handle_dynamic_tool_call( - { - "threadId": "t-1", - "turnId": "turn-1", - "callId": "call-1", - "tool": "my_tool", - "arguments": {"x": 7}, - } - ) - finally: - codex_mod._TOOL_HANDLERS.pop("t-1", None) - - assert calls == [{"x": 7}] - assert response == { - "contentItems": [{"type": "inputText", "text": "got 7"}], - "success": True, - } - - def test_handle_dynamic_tool_call_reports_unknown_tool(self): - response = codex_mod._handle_dynamic_tool_call( - { - "threadId": "nobody", - "tool": "ghost", - "arguments": {}, - } - ) - - assert response["success"] is False - assert "not registered" in response["contentItems"][0]["text"] - - def test_handle_dynamic_tool_call_captures_handler_exceptions(self): - async def handler(_args): - raise RuntimeError("explode") - - codex_mod._TOOL_HANDLERS["t-1"] = {"broken": handler} - try: - response = codex_mod._handle_dynamic_tool_call( - { - "threadId": "t-1", - "tool": "broken", - "arguments": {}, - } - ) - finally: - codex_mod._TOOL_HANDLERS.pop("t-1", None) - - assert response["success"] is False - assert "explode" in response["contentItems"][0]["text"] - - class TestResolveCodexBin: def test_env_override_wins(self, tmp_path, monkeypatch): fake = tmp_path / "codex" @@ -1939,65 +956,6 @@ def test_raises_when_nothing_found(self, monkeypatch): sys.modules.pop("codex_cli_bin", None) -class TestRelaxServiceTier: - """``_relax_service_tier_on_module`` rewrites ``ServiceTier`` references to ``str``. - - This lets the SDK accept new wire values like ``"priority"`` that the generated - enum hasn't been regenerated to include yet. - """ - - def _build_fake_sdk_module(self) -> ModuleType: - from enum import Enum - - from pydantic import BaseModel, ConfigDict, Field - from typing_extensions import Annotated - - module = ModuleType("fake_codex_sdk") - - class ServiceTier(Enum): - fast = "fast" - flex = "flex" - - class ThreadStartResponse(BaseModel): - model_config = ConfigDict(populate_by_name=True) - model: str - service_tier: Annotated[ServiceTier | None, Field(alias="serviceTier")] = None - - class TurnEvent(BaseModel): - model_config = ConfigDict(populate_by_name=True) - service_tier: ServiceTier | None = None - - class Unrelated(BaseModel): - value: str - - module.ServiceTier = ServiceTier - module.ThreadStartResponse = ThreadStartResponse - module.TurnEvent = TurnEvent - module.Unrelated = Unrelated - return module - - def test_priority_value_rejected_before_patch(self): - module = self._build_fake_sdk_module() - with pytest.raises(Exception): - module.ThreadStartResponse.model_validate({"model": "m", "serviceTier": "priority"}) - - def test_relax_accepts_arbitrary_string(self): - module = self._build_fake_sdk_module() - affected = codex_mod._relax_service_tier_on_module(module) - # Two models referenced ServiceTier; Unrelated did not. - assert affected == 2 - - ok = module.ThreadStartResponse.model_validate({"model": "m", "serviceTier": "priority"}) - assert ok.service_tier == "priority" - - ok2 = module.TurnEvent.model_validate({"service_tier": "anything"}) - assert ok2.service_tier == "anything" - - def test_no_servicetier_attr_returns_zero(self): - module = ModuleType("empty") - assert codex_mod._relax_service_tier_on_module(module) == 0 - - class TestClaudeBackendVersion: def _reset_cache(self, monkeypatch): monkeypatch.setattr(cc_mod, "_VERSION_CACHE", None) @@ -2134,15 +1092,3 @@ def test_cli_version_returns_empty_on_nonzero_exit(self, monkeypatch): lambda *a, **k: SimpleNamespace(returncode=1, stdout="", stderr="boom"), ) assert codex_mod._codex_cli_version() == "" - - -@pytest.fixture(autouse=True) -def cleanup_fake_sdk_modules(): - yield - for name in [ - "openai_codex", - "openai_codex.api", - "openai_codex.client", - "openai_codex.generated.v2_all", - ]: - sys.modules.pop(name, None) diff --git a/agent-flow/tests/test_codex_backend.py b/agent-flow/tests/test_codex_backend.py new file mode 100644 index 000000000000..b39603756063 --- /dev/null +++ b/agent-flow/tests/test_codex_backend.py @@ -0,0 +1,404 @@ +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest +from openai_codex.generated.v2_all import ( + DynamicToolSpec, + FunctionDynamicToolSpec, + ThreadStartParams, +) +from pydantic import TypeAdapter + +from agent_flow.backends import codex as codex_module +from agent_flow.backends.codex import CodexBackend, _dynamic_tool_spec +from agent_flow.tools import tool + + +class FakeTransport: + """Record backend requests without launching a Codex process.""" + + def __init__(self, config=None): + self.config = config + self.started = False + self.closed = False + self.requests = [] + self.registered = {} + self.unregistered = [] + self.thread_count = 0 + self.config_read_response = {"config": {}} + + async def start(self): + self.started = True + + async def close(self): + self.closed = True + + async def request(self, method, params): + self.requests.append((method, copy.deepcopy(params))) + if method == "config/read": + if isinstance(self.config_read_response, Exception): + raise self.config_read_response + return self.config_read_response + if method == "thread/start": + self.thread_count += 1 + return {"thread": {"id": f"thread-{self.thread_count}"}} + if method == "skills/list": + return { + "data": [ + {"skills": [{"name": "available"}, {"name": "disabled", "enabled": False}]} + ] + } + if method == "plugin/list": + return {"marketplaces": [{"plugins": [{"name": "installed", "installed": True}]}]} + return {} + + def register_tools(self, thread_id, tools): + self.registered[thread_id] = list(tools) + + def unregister_tools(self, thread_id): + self.unregistered.append(thread_id) + self.registered.pop(thread_id, None) + + +@pytest.fixture +async def backend(monkeypatch): + monkeypatch.setattr(codex_module, "CodexTransport", FakeTransport) + monkeypatch.setattr(codex_module, "_resolve_codex_bin", lambda: "/unused/codex") + async with CodexBackend() as instance: + yield instance + + +@tool("read", "Read data", {"path": str}, annotations={"readOnlyHint": True}) +async def read_tool(args): + return {"content": [{"type": "text", "text": args["path"]}]} + + +@tool("write", "Write data", {"path": str}) +async def write_tool(args): + return {"content": [{"type": "text", "text": "written"}]} + + +def test_dynamic_tool_spec_validates_against_pinned_sdk_discriminated_union(): + payload = _dynamic_tool_spec(read_tool) + function = FunctionDynamicToolSpec.model_validate(payload) + union = TypeAdapter(DynamicToolSpec).validate_python(payload) + assert function.type == union.root.type == "function" + assert payload["inputSchema"] == { + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + } + assert payload["deferLoading"] is False + + +async def test_backend_start_and_close_use_configured_runtime(monkeypatch): + monkeypatch.setattr(codex_module, "CodexTransport", FakeTransport) + monkeypatch.setattr(codex_module, "_resolve_codex_bin", lambda: "/configured/codex") + instance = CodexBackend() + async with instance: + transport = instance._transport + assert transport.started + assert transport.config.codex_bin == "/configured/codex" + assert transport.config.experimental_api is True + assert transport.closed + assert instance._transport is None + + +async def test_session_options_use_developer_instructions_and_preserve_native_defaults( + backend, tmp_path +): + transport = backend._transport + async with backend.create_client( + "Project instructions", "gpt-5.4", tools=[read_tool], cwd=tmp_path + ) as client: + payload = next(params for method, params in transport.requests if method == "thread/start") + validated = ThreadStartParams.model_validate(payload) + assert validated.model == "gpt-5.4" + assert payload["cwd"] == str(tmp_path.resolve()) + assert payload["developerInstructions"] == "Project instructions" + assert "baseInstructions" not in payload + assert payload["approvalPolicy"] == "never" + assert payload["sandbox"] == "danger-full-access" + assert payload["config"]["model_reasoning_effort"] == backend.reasoning_effort() + assert payload["dynamicTools"] == [_dynamic_tool_spec(read_tool)] + assert transport.registered == {"thread-1": [read_tool]} + assert await client.list_available_skills() == ["available"] + assert transport.registered == {} + assert transport.unregistered == ["thread-1"] + assert ("thread/unsubscribe", {"threadId": "thread-1"}) in transport.requests + + +async def test_empty_prompt_does_not_replace_runtime_instructions(backend, tmp_path): + async with backend.create_client("", "gpt-5.4", cwd=tmp_path): + payload = next( + params for method, params in backend._transport.requests if method == "thread/start" + ) + assert "developerInstructions" not in payload + assert "baseInstructions" not in payload + assert "dynamicTools" not in payload + + +async def test_tool_restrictions_affect_both_advertised_tools_and_dispatch_registration( + backend, tmp_path +): + async with backend.create_client( + "", + "gpt-5.4", + tools=[read_tool, write_tool], + disallowed_tools=["mcp__agent-tools__write", "Bash", "WebSearch"], + cwd=tmp_path, + ): + transport = backend._transport + payload = next(params for method, params in transport.requests if method == "thread/start") + assert [entry["name"] for entry in payload["dynamicTools"]] == ["read"] + assert [entry.name for entry in transport.registered["thread-1"]] == ["read"] + assert payload["config"]["features"]["shell_tool"] is False + assert payload["config"]["web_search"] == "disabled" + + +async def test_concurrent_sessions_keep_mcp_overrides_local_and_do_not_mutate_inputs( + backend, tmp_path +): + first_servers = { + "docs": {"type": "stdio", "command": "python", "args": ["first.py"], "env": {"TEAM": "one"}} + } + second_servers = { + "docs": {"type": "http", "url": "https://example.com/mcp", "headers": {"X-Team": "two"}} + } + originals = copy.deepcopy((first_servers, second_servers)) + async with backend.create_client( + "first", "gpt-5.4", extra_mcp_servers=first_servers, cwd=tmp_path + ): + async with backend.create_client( + "second", "gpt-5.4", extra_mcp_servers=second_servers, cwd=tmp_path + ): + requests = [ + params for method, params in backend._transport.requests if method == "thread/start" + ] + first, second = (params["config"]["mcp_servers"]["docs"] for params in requests) + assert first == { + "command": "python", + "args": ["first.py"], + "env": {"TEAM": "one"}, + "cwd": str(tmp_path), + } + assert second == {"url": "https://example.com/mcp", "http_headers": {"X-Team": "two"}} + assert set(backend._transport.registered) == {"thread-1", "thread-2"} + assert set(backend._transport.registered) == {"thread-1"} + assert (first_servers, second_servers) == originals + assert not any(method.startswith("config/") for method, _ in backend._transport.requests) + + +async def test_mcp_restrictions_preserve_native_exclusions( + backend: CodexBackend, tmp_path: Path +) -> None: + transport = backend._transport + transport.config_read_response = { + "config": { + "mcp_servers": { + "files": {"disabled_tools": ["delete"], "command": "native-files"}, + "other": {"disabled_tools": ["admin"]}, + }, + "model": "native-model", + } + } + native_config = copy.deepcopy(transport.config_read_response) + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__write"], cwd=tmp_path / "nested" / ".." + ): + assert transport.requests[0] == ( + "config/read", + {"cwd": str(tmp_path.resolve()), "includeLayers": True}, + ) + payload = next(params for method, params in transport.requests if method == "thread/start") + assert payload["config"]["mcp_servers"] == { + "files": {"disabled_tools": ["delete", "write"]} + } + assert "model" not in payload["config"] + assert transport.config_read_response == native_config + + +async def test_mcp_exclusions_preserve_disabled_project_layers_that_may_become_trusted( + backend: CodexBackend, tmp_path: Path +) -> None: + transport = backend._transport + transport.config_read_response = { + "config": {"mcp_servers": {"files": {"disabled_tools": ["delete"]}}}, + "layers": [ + { + "name": {"type": kind}, + "config": {"mcp_servers": {"files": {"disabled_tools": disabled}}}, + "disabledReason": reason, + } + for kind, disabled, reason in [ + ("project", ["delete", "project-delete"], "not trusted"), + ("project", ["superseded-project-ban"], None), + ("user", ["superseded-user-ban"], None), + ("user", ["disabled-user-ban"], "disabled"), + ] + ], + } + native_config = copy.deepcopy(transport.config_read_response) + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__write"], cwd=tmp_path + ): + payload = next(params for method, params in transport.requests if method == "thread/start") + assert payload["config"]["mcp_servers"] == { + "files": {"disabled_tools": ["delete", "project-delete", "write"]} + } + assert transport.config_read_response == native_config + + +@pytest.mark.parametrize("session_disabled", [[], ["write", "delete", "write"]]) +async def test_extra_mcp_exclusions_merge_with_native_without_mutating_inputs( + backend: CodexBackend, tmp_path: Path, session_disabled: list[str] +) -> None: + transport = backend._transport + transport.config_read_response = { + "config": {"mcp_servers": {"files": {"disabled_tools": ["delete", "delete"]}}} + } + servers = {"files": {"command": "files", "disabled_tools": session_disabled}} + original = copy.deepcopy(servers) + async with backend.create_client("", "gpt-5.4", extra_mcp_servers=servers, cwd=tmp_path): + payload = next(params for method, params in transport.requests if method == "thread/start") + expected = ["delete", "write"] if session_disabled else ["delete"] + assert payload["config"]["mcp_servers"]["files"]["disabled_tools"] == expected + assert servers == original + assert transport.config_read_response["config"]["mcp_servers"]["files"]["disabled_tools"] == [ + "delete", + "delete", + ] + + +@pytest.mark.parametrize( + "native_config", + [ + {}, + {"mcp_servers": {}}, + {"mcp_servers": {"files": {}}}, + {"mcp_servers": {"files": {"disabled_tools": None}}}, + ], +) +async def test_mcp_exclusions_allow_missing_native_lists( + backend: CodexBackend, tmp_path: Path, native_config: dict[str, object] +) -> None: + transport = backend._transport + transport.config_read_response = {"config": native_config} + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__write"], cwd=tmp_path + ): + payload = next(params for method, params in transport.requests if method == "thread/start") + assert payload["config"]["mcp_servers"] == {"files": {"disabled_tools": ["write"]}} + + +async def test_mcp_exclusions_read_each_session_cwd_without_leaking_overrides( + backend: CodexBackend, tmp_path: Path +) -> None: + transport = backend._transport + first_cwd, second_cwd = tmp_path / "first", tmp_path / "second" + transport.config_read_response = { + "config": {"mcp_servers": {"files": {"disabled_tools": ["delete"]}}} + } + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__write"], cwd=first_cwd + ): + transport.config_read_response = { + "config": {"mcp_servers": {"files": {"disabled_tools": ["rename"]}}} + } + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__upload"], cwd=second_cwd + ): + configs = [ + params["config"]["mcp_servers"]["files"]["disabled_tools"] + for method, params in transport.requests + if method == "thread/start" + ] + assert configs == [["delete", "write"], ["rename", "upload"]] + assert [params for method, params in transport.requests if method == "config/read"] == [ + {"cwd": str(first_cwd.resolve()), "includeLayers": True}, + {"cwd": str(second_cwd.resolve()), "includeLayers": True}, + ] + + +@pytest.mark.parametrize("failure", [RuntimeError("config unavailable"), TimeoutError()]) +async def test_mcp_exclusions_fail_closed_when_native_config_cannot_be_read( + backend: CodexBackend, tmp_path: Path, failure: Exception +) -> None: + transport = backend._transport + transport.config_read_response = failure + with pytest.raises(type(failure)): + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__write"], cwd=tmp_path + ): + pytest.fail("unreadable native exclusions must not create a client") + assert [method for method, _ in transport.requests] == ["config/read"] + assert transport.registered == {} + + +@pytest.mark.parametrize("native_disabled", ["delete", ["delete", 1]]) +async def test_mcp_exclusions_reject_malformed_native_lists( + backend: CodexBackend, tmp_path: Path, native_disabled: object +) -> None: + transport = backend._transport + transport.config_read_response = { + "config": {"mcp_servers": {"files": {"disabled_tools": native_disabled}}} + } + with pytest.raises(ValueError, match="disabled_tools.*list of strings"): + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__write"], cwd=tmp_path + ): + pytest.fail("invalid native exclusions must not create a client") + assert [method for method, _ in transport.requests] == ["config/read"] + + +async def test_whole_mcp_server_restriction_does_not_read_native_config( + backend: CodexBackend, tmp_path: Path +) -> None: + async with backend.create_client( + "", "gpt-5.4", disallowed_tools=["mcp__files__*"], cwd=tmp_path + ): + requests = backend._transport.requests + assert all(method != "config/read" for method, _ in requests) + payload = next(params for method, params in requests if method == "thread/start") + assert payload["config"]["mcp_servers"] == {"files": {"enabled": False}} + + +async def test_session_exception_unregisters_tools_and_unsubscribes(backend, tmp_path): + with pytest.raises(ValueError, match="application failure"): + async with backend.create_client("", "gpt-5.4", tools=[read_tool], cwd=tmp_path): + raise ValueError("application failure") + assert backend._transport.registered == {} + assert ("thread/unsubscribe", {"threadId": "thread-1"}) in backend._transport.requests + + +@pytest.mark.parametrize( + "options, message", + [ + ({"disallowed_tools": ["arbitrary-native-tool"]}, "cannot enforce"), + ( + {"disallowed_tools": ["mcp__files__write", "arbitrary-native-tool"]}, + "cannot enforce", + ), + ( + {"hooks": {"Stop": [{"hooks": [{"type": "command", "command": "true"}]}]}}, + "cannot verify trust", + ), + ( + {"extra_mcp_servers": {"remote": {"type": "sse", "url": "https://example.com/sse"}}}, + "MCP transport", + ), + ], +) +async def test_unsupported_options_fail_before_any_session_rpc(backend, tmp_path, options, message): + with pytest.raises((TypeError, ValueError), match=message): + async with backend.create_client("", "gpt-5.4", cwd=tmp_path, **options): + pytest.fail("invalid options must not create a client") + assert backend._transport.requests == [] + + +async def test_client_requires_entered_backend(): + with pytest.raises(RuntimeError, match="must be entered"): + async with CodexBackend().create_client("", "gpt-5.4", cwd=Path.cwd()): + pytest.fail("an unopened backend must not create a client") diff --git a/agent-flow/tests/test_codex_config.py b/agent-flow/tests/test_codex_config.py new file mode 100644 index 000000000000..0e9c4dfb5df8 --- /dev/null +++ b/agent-flow/tests/test_codex_config.py @@ -0,0 +1,206 @@ +import copy +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from agent_flow.backends.codex_config import build_session_config + + +def test_empty_configuration_does_not_set_native_defaults(): + assert build_session_config() == ({}, []) + assert build_session_config(hooks={}) == ({}, []) + + +def test_stdio_mcp_uses_session_cwd_and_keeps_environment(): + servers = { + "files": { + "type": "stdio", + "command": "python", + "args": ["-m", "files"], + "env": {"ROOT": "/workspace"}, + "env_vars": ["TOKEN", {"name": "HOST_TOKEN"}], + "startup_timeout_sec": 15, + "tool_timeout_sec": 30, + } + } + original = copy.deepcopy(servers) + config, _ = build_session_config(extra_mcp_servers=servers, cwd=Path("/workspace")) + assert config["mcp_servers"]["files"] == { + "command": "python", + "args": ["-m", "files"], + "env": {"ROOT": "/workspace"}, + "env_vars": ["TOKEN", {"name": "HOST_TOKEN"}], + "startup_timeout_sec": 15, + "tool_timeout_sec": 30, + "cwd": "/workspace", + } + assert servers == original + + +def test_stdio_explicit_relative_cwd_resolves_against_session_cwd(): + config, _ = build_session_config( + extra_mcp_servers={"files": {"command": "files", "cwd": "child"}}, + cwd=Path("/workspace"), + ) + assert config["mcp_servers"]["files"]["cwd"] == "/workspace/child" + + +def test_http_mcp_translates_headers_and_preserves_auth(): + server = { + "type": "http", + "url": "https://example.test/mcp", + "headers": {"X-Client": "agent-flow"}, + "env_http_headers": {"X-Token": "TOKEN"}, + "bearer_token_env_var": "API_TOKEN", + "auth": "oauth", + "oauth": {"client_id": "client", "callback_port": 8123}, + "scopes": ["read"], + "enabled_tools": ["read"], + "tools": {"read": {"approval_mode": "auto", "output_token_limit": 500}}, + } + config, _ = build_session_config(extra_mcp_servers={"docs": server}) + expected = {k: v for k, v in server.items() if k not in {"headers", "type"}} + expected["http_headers"] = server["headers"] + assert config["mcp_servers"]["docs"] == expected + + +@pytest.mark.parametrize( + "server, message", + [ + ({"type": "sse", "url": "https://example.test"}, "transport"), + ({"type": "sdk", "instance": object()}, "transport"), + ({"command": "server", "url": "https://example.test"}, "Unsupported"), + ({"url": "file:///tmp/server"}, "http://"), + ({"url": "https://example.test", "headers": {}, "http_headers": {}}, "both headers"), + ({"command": "server", "env": {"TOKEN": 7}}, "string values"), + ({"command": "server", "args": "argument"}, "list of strings"), + ({"command": "server", "invented_option": True}, "Unsupported"), + ({"command": "server", "enabled": "false"}, "boolean"), + ({"command": "server", "startup_timeout_sec": float("nan")}, "finite"), + ({"command": "server", "startup_timeout_sec": -1}, "nonnegative"), + ({"command": "server", "startup_timeout_sec": 1, "startup_timeout_ms": 1}, "one MCP"), + ({"url": "https://example.test", "auth": "none"}, "auth"), + ({"url": "https://example.test", "oauth": {"callback_port": 70000}}, "65535"), + ({"command": "server", "env_vars": [{"name": "TOKEN", "typo": "value"}]}, "Unsupported"), + ({"command": "server", "tools": {"x": {"output_token_limit": 0}}}, "positive"), + ({"command": "server", "default_tools_approval_mode": "never"}, "approval_mode"), + ], +) +def test_invalid_mcp_options_fail_instead_of_being_dropped(server, message): + with pytest.raises((TypeError, ValueError), match=message): + build_session_config(extra_mcp_servers={"external": server}) + + +def test_framework_server_name_is_reserved(): + with pytest.raises(ValueError, match="reserved"): + build_session_config(extra_mcp_servers={"agent-tools": {"command": "server"}}) + + +def test_dynamic_restrictions_accept_bare_and_claude_namespaced_names(): + one, two, three = [SimpleNamespace(name=name) for name in ["one", "two", "three"]] + tools = [one, two, three] + config, allowed = build_session_config( + tools=tools, disallowed_tools=["one", "mcp__agent-tools__two", "one"] + ) + assert config == {} + assert allowed == [three] + assert tools == [one, two, three] + + +def test_framework_namespace_wildcard_removes_all_dynamic_tools(): + assert build_session_config( + tools=[SimpleNamespace(name="one")], disallowed_tools=["mcp__agent-tools__*"] + ) == ({}, []) + + +def test_mcp_restrictions_merge_without_mutating_original_config(): + servers = {"files": {"command": "files", "disabled_tools": ["remove"]}} + original = copy.deepcopy(servers) + config, _ = build_session_config( + extra_mcp_servers=servers, + disallowed_tools=["mcp__files__write", "mcp__files__remove", "mcp__other__*"], + ) + assert config["mcp_servers"]["files"]["disabled_tools"] == ["remove", "write"] + assert config["mcp_servers"]["other"] == {"enabled": False} + assert servers == original + + +@pytest.mark.parametrize("name", ["Bash", "shell", "shell_command", "exec_command", "write_stdin"]) +def test_shell_aliases_disable_native_shell_family(name): + assert build_session_config(disallowed_tools=[name])[0] == {"features": {"shell_tool": False}} + + +def test_builtin_controls_match_pinned_native_configuration(): + config, _ = build_session_config( + disallowed_tools=[ + "AskUserQuestion", + "WebFetch", + "view_image", + "image_generation", + "TodoWrite", + ] + ) + assert config == { + "tools": { + "experimental_request_user_input": {"enabled": False}, + "update_plan": {"enabled": False}, + }, + "web_search": "disabled", + "features": {"view_image": False, "image_generation": False}, + } + + +def test_app_controls_require_explicit_app_ids(): + config, _ = build_session_config(disallowed_tools=["apps.calendar.create", "apps.mail.*"]) + assert config == { + "apps": { + "calendar": {"tools": {"create": {"enabled": False}}}, + "mail": {"enabled": False}, + } + } + with pytest.raises(ValueError, match="APP_ID"): + build_session_config(disallowed_tools=["mcp__codex_apps__create"]) + + +@pytest.mark.parametrize("name", ["Write", "Edit", "apply_patch", "Bash(rm:*)", "unknown"]) +def test_restrictions_without_native_enforcement_fail_explicitly(name): + with pytest.raises(ValueError, match="cannot enforce"): + build_session_config(disallowed_tools=[name]) + + +def test_mcp_patterns_are_not_misrepresented_as_exact_restrictions(): + with pytest.raises(ValueError, match="exact tool name"): + build_session_config(disallowed_tools=["mcp__files__write*"]) + + +@pytest.mark.parametrize( + "name", ["mcp__agent-tools__write*", "apps.calendar.create*", "mcp__*__write", "apps.*.create"] +) +def test_other_restriction_patterns_fail_explicitly(name): + with pytest.raises(ValueError, match="exact tool name"): + build_session_config(disallowed_tools=[name]) + + +def test_default_app_settings_cannot_be_used_as_a_global_app_ban(): + with pytest.raises(ValueError, match="explicit app ID"): + build_session_config(disallowed_tools=["apps._default.*"]) + + +def test_claude_python_callbacks_require_shared_policy(): + with pytest.raises(TypeError, match="HookMatcher.*required_tools"): + build_session_config(hooks={"Stop": [SimpleNamespace(hooks=[lambda: None])]}) + + +def test_native_hooks_fail_if_session_trust_cannot_be_verified(): + with pytest.raises(ValueError, match="cannot verify trust"): + build_session_config( + hooks={ + "Stop": [{"hooks": [{"type": "command", "command": "validate-stop"}]}], + } + ) + + +def test_embedded_native_trust_claims_do_not_bypass_trust_validation(): + with pytest.raises(ValueError, match="cannot verify trust"): + build_session_config(hooks={"state": {"key": {"trusted_hash": "arbitrary"}}}) diff --git a/agent-flow/tests/test_codex_events.py b/agent-flow/tests/test_codex_events.py new file mode 100644 index 000000000000..8a2f4edc7270 --- /dev/null +++ b/agent-flow/tests/test_codex_events.py @@ -0,0 +1,504 @@ +from __future__ import annotations + +import pytest +from openai_codex.generated.v2_all import ( + AccountRateLimitsUpdatedNotification, + ThreadTokenUsageUpdatedNotification, + TurnCompletedNotification, +) +from openai_codex.models import Notification, UnknownNotification + +from agent_flow.backends.base import ResultEvent +from agent_flow.backends.codex_events import CodexClient +from agent_flow.types import ( + AgentTextEvent, + CompactBoundaryEvent, + RateLimitWarningEvent, + ServerToolCallEvent, + SessionInitEvent, + ThinkingEvent, + ToolCallEvent, +) + + +def notice(method, **payload): + models = { + "thread/tokenUsage/updated": ThreadTokenUsageUpdatedNotification, + "account/rateLimits/updated": AccountRateLimitsUpdatedNotification, + "turn/completed": TurnCompletedNotification, + } + model = models.get(method) + return Notification( + method, model.model_validate(payload) if model is not None else UnknownNotification(payload) + ) + + +def item(kind, item_id, *, thread="main", method="item/completed", **fields): + return notice( + method, + threadId=thread, + turnId="turn", + item={"id": item_id, "type": kind, **fields}, + ) + + +def completed(*, thread="main", **fields): + return notice( + "turn/completed", + threadId=thread, + turn={"id": "turn", "status": "completed", "items": [], "itemsView": "full", **fields}, + ) + + +def tokens(total, last=None, *, thread="main"): + def breakdown(multiplier): + return { + "inputTokens": 80 * multiplier, + "outputTokens": 20 * multiplier, + "cachedInputTokens": 30 * multiplier, + "cacheWriteInputTokens": 10 * multiplier, + "totalTokens": 100 * multiplier, + "reasoningOutputTokens": 5 * multiplier, + } + + return notice( + "thread/tokenUsage/updated", + threadId=thread, + turnId="turn", + tokenUsage={ + "total": breakdown(total), + "last": breakdown(last if last is not None else total), + "modelContextWindow": 1000, + }, + ) + + +def started(thread, parent, nickname=None, role=None): + return notice( + "thread/started", + thread={ + "id": thread, + "parentThreadId": parent, + "agentNickname": nickname, + "agentRole": role, + }, + ) + + +class Transport: + def __init__(self, *turns, responses=None): + self.turns = list(turns) + self.responses = responses or {} + self.requests = [] + self.messages = [] + + async def stream_turn(self, thread_id, message): + self.messages.append((thread_id, message)) + for event in self.turns.pop(0): + yield event + + async def request(self, method, params): + self.requests.append((method, params)) + value = self.responses.get(method, {}) + if isinstance(value, Exception): + raise value + return value + + +async def collect(client, message="hello"): + return [event async for event in client.send_message(message)] + + +async def test_persistent_usage_is_per_send_context_is_last_and_cache_write_is_mapped(): + transport = Transport( + [tokens(2, 1), completed(durationMs=1234)], + [tokens(3, 1), tokens(5, 2), tokens(900, thread="unrelated"), completed(durationMs=90)], + ) + client = CodexClient(transport, "main") + assert await client.get_context_usage() is None + assert transport.messages == [] + first = (await collect(client))[-1].usage + second = (await collect(client))[-1].usage + assert (first.total_tokens, second.total_tokens) == (200, 300) + assert second.input_tokens == 240 + assert second.output_tokens == 60 + assert second.cache_read_tokens == 90 + assert second.cache_creation_tokens == 30 + assert (first.context_tokens, second.context_tokens) == (100, 200) + assert second.context_percentage == 20 + assert second.context_window == 1000 + assert (first.duration_ms, second.duration_ms) == (1234, 90) + assert (first.num_turns, second.num_turns) == (1, 1) + context = await client.get_context_usage() + context.context_tokens = 99 + assert (await client.get_context_usage()).context_tokens == 200 + + +async def test_token_counter_reset_never_reports_negative_usage(): + client = CodexClient(Transport([tokens(4), completed()], [tokens(1), completed()]), "main") + await collect(client) + assert (await collect(client))[-1].usage.total_tokens == 100 + + +async def test_session_inventory_is_emitted_once_without_an_extra_model_turn(): + init = SessionInitEvent(skills=["review"], plugins=["plugin"], agents=["explorer"]) + transport = Transport([completed()], [completed()]) + client = CodexClient(transport, "main", init) + assert await client.list_available_skills() == ["review"] + assert transport.messages == [] + assert (await collect(client))[0] is init + assert not any(isinstance(event, SessionInitEvent) for event in await collect(client)) + + +async def test_tools_text_reasoning_and_compaction_are_deduplicated(): + client = CodexClient( + Transport( + [ + item("commandExecution", "bash", method="item/started", command="ls"), + item("commandExecution", "bash", command="ls", status="completed"), + item("dynamicToolCall", "dynamic", tool="update_status", arguments={"ok": True}), + item("mcpToolCall", "mcp", server="kb", tool="lookup", arguments={"term": "x"}), + item("fileChange", "file", changes=[{"path": "x.py"}], status="completed"), + item("webSearch", "web", query="query", action={"type": "search"}), + item("imageView", "img", path="/tmp/image.png"), + item("contextCompaction", "compact", method="item/started"), + item("contextCompaction", "compact"), + item("reasoning", "reason", summary=["thinking"], content=["detail"]), + notice("item/agentMessage/delta", threadId="main", delta="ignored"), + item("agentMessage", "comment", text="working", phase="commentary"), + item("agentMessage", "final", text="answer", phase="final_answer"), + completed(), + ] + ), + "main", + ) + events = await collect(client) + assert [e.name for e in events if isinstance(e, ToolCallEvent)] == [ + "Bash", + "update_status", + "mcp__kb__lookup", + "FileChange", + ] + assert [e.name for e in events if isinstance(e, ServerToolCallEvent)] == [ + "web_search", + "view_image", + ] + assert len([e for e in events if isinstance(e, CompactBoundaryEvent)]) == 1 + assert [e.text for e in events if isinstance(e, ThinkingEvent)] == ["thinking\ndetail"] + assert [e.text for e in events if isinstance(e, AgentTextEvent)] == ["working", "answer"] + assert events[-1].text == "answer" + + +async def test_recovered_tool_failure_does_not_fail_successful_turn(): + events = await collect( + CodexClient( + Transport( + [ + item("commandExecution", "bad", command="false", status="failed"), + item("fileChange", "denied", changes=[], status="declined"), + item("agentMessage", "answer", text="recovered"), + completed(), + ] + ), + "main", + ) + ) + result = events[-1] + assert not result.is_error + assert result.errors == [] + assert result.permission_denials[0]["id"] == "denied" + assert result.text == "recovered" + + +async def test_child_before_parent_race_preserves_labels_and_spawn_attribution(): + events = await collect( + CodexClient( + Transport( + [ + # A child can emit before its parent's spawn call completes. + item("agentMessage", "child-text", thread="child", text="found it"), + started("grandchild", "child", "Ada", "worker"), + item("reasoning", "grand-reason", thread="grandchild", summary=["checking"]), + started("child", "main", "Grace", "explorer"), + item("commandExecution", "child-tool", thread="child", command="rg foo"), + item( + "collabAgentToolCall", + "nested-spawn", + thread="child", + tool="spawnAgent", + senderThreadId="child", + receiverThreadIds=["grandchild"], + ), + item( + "collabAgentToolCall", + "spawn", + tool="spawnAgent", + senderThreadId="main", + receiverThreadIds=["child"], + ), + tokens(80, thread="child"), + completed(thread="child"), + started("unrelated", "elsewhere", "Other", "worker"), + item("agentMessage", "other", thread="unrelated", text="invisible"), + completed(thread="unrelated"), + tokens(1), + item("agentMessage", "main-final", text="done", phase="final_answer"), + completed(), + ] + ), + "main", + ) + ) + child_text = next(e for e in events if isinstance(e, AgentTextEvent) and e.text == "found it") + assert (child_text.parent_tool_use_id, child_text.agent_label) == ("spawn", "Grace (explorer)") + child_tool = next(e for e in events if isinstance(e, ToolCallEvent) and e.name == "Bash") + assert (child_tool.parent_tool_use_id, child_tool.agent_label) == ("spawn", "Grace (explorer)") + grandchild = next(e for e in events if isinstance(e, ThinkingEvent)) + assert (grandchild.parent_tool_use_id, grandchild.agent_label) == ( + "nested-spawn", + "Ada (worker)", + ) + assert not any(isinstance(e, AgentTextEvent) and e.text == "invisible" for e in events) + results = [e for e in events if isinstance(e, ResultEvent)] + assert len(results) == 1 + assert results[0].text == "done" + assert results[0].usage.total_tokens == 100 + + +async def test_child_without_spawn_event_has_label_but_no_invented_invocation_id(): + events = await collect( + CodexClient( + Transport( + [ + started("child", "main", role="worker"), + item("agentMessage", "text", thread="child", text="child"), + completed(), + ] + ), + "main", + ) + ) + text = next(e for e in events if isinstance(e, AgentTextEvent)) + assert text.agent_label == "worker" + assert text.parent_tool_use_id is None + + +async def test_structured_rate_limits_preserve_sparse_metadata_and_report_thresholds(): + initial = { + "rateLimits": { + "limitId": "codex", + "primary": {"usedPercent": 10, "resetsAt": 100}, + "secondary": {"usedPercent": 85, "resetsAt": 200}, + } + } + update = notice( + "account/rateLimits/updated", + rateLimits={"limitId": "codex", "primary": {"usedPercent": 100, "resetsAt": 110}}, + ) + events = await collect( + CodexClient( + Transport( + [update, update, completed()], responses={"account/rateLimits/read": initial} + ), + "main", + ) + ) + limits = [e for e in events if isinstance(e, RateLimitWarningEvent)] + assert [(e.rate_limit_type, e.status, e.utilization, e.resets_at) for e in limits] == [ + ("codex/primary", "allowed", 0.1, 100), + ("codex/secondary", "allowed_warning", 0.85, 200), + ("codex/primary", "rejected", 1.0, 110), + ] + + +async def test_native_spend_control_limit_reached_needs_no_text_matching(): + events = await collect( + CodexClient( + Transport( + [ + notice( + "account/rateLimits/updated", + rateLimits={"rateLimitReachedType": "workspace_member_credits_depleted"}, + ), + completed(), + ] + ), + "main", + ) + ) + rate = next(e for e in events if isinstance(e, RateLimitWarningEvent)) + assert rate.status == "rejected" + assert rate.rate_limit_type == "workspace_member_credits_depleted" + + +async def test_sparse_window_update_preserves_previously_reported_reset_time(): + events = await collect( + CodexClient( + Transport( + [ + notice( + "account/rateLimits/updated", + rateLimits={"primary": {"usedPercent": 90}}, + ), + completed(), + ], + responses={ + "account/rateLimits/read": { + "rateLimits": {"primary": {"usedPercent": 10, "resetsAt": 123}} + } + }, + ), + "main", + ) + ) + rates = [e for e in events if isinstance(e, RateLimitWarningEvent)] + assert rates[-1].status == "allowed_warning" + assert rates[-1].utilization == 0.9 + assert rates[-1].resets_at == 123 + + +async def test_nullable_sparse_reached_state_requires_full_snapshot_to_confirm_recovery(): + class RecoveryTransport(Transport): + async def request(self, method, params): + if method == "account/rateLimits/read": + recovered = bool(self.requests) + self.requests.append((method, params)) + return { + "rateLimits": { + "limitId": "codex", + "primary": {"usedPercent": 10, "resetsAt": 100}, + "rateLimitReachedType": None if recovered else "rate_limit_reached", + } + } + return await super().request(method, params) + + events = await collect( + CodexClient( + RecoveryTransport( + [ + notice( + "account/rateLimits/updated", + rateLimits={"limitId": "codex", "rateLimitReachedType": None}, + ), + completed(), + ] + ), + "main", + ) + ) + rates = [e for e in events if isinstance(e, RateLimitWarningEvent)] + assert [e.status for e in rates] == ["rejected", "allowed"] + + +async def test_sparse_null_does_not_clear_reached_when_snapshot_cannot_be_refetched(): + class UnavailableRecoveryTransport(Transport): + async def request(self, method, params): + if self.requests: + raise RuntimeError("temporarily unavailable") + return await super().request(method, params) + + events = await collect( + CodexClient( + UnavailableRecoveryTransport( + [ + notice( + "account/rateLimits/updated", + rateLimits={"spendControlReached": None}, + ), + completed(), + ], + responses={ + "account/rateLimits/read": {"rateLimits": {"spendControlReached": True}} + }, + ), + "main", + ) + ) + assert [e.status for e in events if isinstance(e, RateLimitWarningEvent)] == ["rejected"] + + +async def test_another_turn_on_same_thread_cannot_finish_active_turn(): + events = await collect( + CodexClient( + Transport( + [ + notice("turn/started", threadId="main", turn={"id": "active"}), + tokens(900), + completed(), + completed(id="active"), + ] + ), + "main", + ) + ) + assert len(events) == 1 + assert events[0].usage.total_tokens is None + + +@pytest.mark.parametrize("estimate", [0, 1_250_000, None]) +async def test_cost_is_optional_cumulative_thread_estimate_separate_from_turn_cost(estimate): + transport = Transport( + [completed()], + responses={ + "account/usage/read": { + "threadUsage": {"threadId": "main", "estimatedUsageUsdMicros": estimate} + } + }, + ) + result = (await collect(CodexClient(transport, "main")))[-1] + assert result.usage.estimated_thread_cost_usd == ( + estimate / 1_000_000 if estimate is not None else None + ) + assert result.usage.cost_usd is None + assert ("account/usage/read", {"threadId": "main"}) in transport.requests + + +async def test_unavailable_account_apis_do_not_fail_completed_turn(): + events = await collect( + CodexClient( + Transport( + [completed()], + responses={ + "account/usage/read": RuntimeError("API key billing"), + "account/rateLimits/read": RuntimeError("API key billing"), + }, + ), + "main", + ) + ) + assert events[-1].usage.estimated_thread_cost_usd is None + + +async def test_failure_keeps_native_details_http_status_ids_and_partial_output(): + client = CodexClient( + Transport( + [ + item("agentMessage", "partial", text="I reached this step", phase="commentary"), + notice( + "error", + threadId="main", + error={"message": "retry failed", "additionalDetails": "attempt 3"}, + ), + completed( + status="failed", + error={ + "message": "upstream failure", + "codexErrorInfo": {"httpConnectionFailed": {"httpStatusCode": 503}}, + "additionalDetails": "request-123", + }, + ), + ] + ), + "main", + ) + with pytest.raises(RuntimeError) as exc: + await collect(client) + for detail in ("upstream failure", "503", "request-123", "I reached this step", "attempt 3"): + assert detail in str(exc.value) + assert '"threadId": "main"' in str(exc.value) + assert '"turnId": "turn"' in str(exc.value) + + +async def test_stream_ending_without_root_completion_cannot_report_success(): + with pytest.raises(RuntimeError, match="before the main turn completed"): + await collect(CodexClient(Transport([completed(thread="child")]), "main")) diff --git a/agent-flow/tests/test_codex_transport.py b/agent-flow/tests/test_codex_transport.py new file mode 100644 index 000000000000..60bbb244f0ab --- /dev/null +++ b/agent-flow/tests/test_codex_transport.py @@ -0,0 +1,350 @@ +from __future__ import annotations + +import asyncio +import contextvars +import queue +import threading +from collections.abc import Callable +from typing import Any + +import pytest +from openai_codex.client import CodexConfig +from openai_codex.generated.v2_all import ThreadStartedNotification +from openai_codex.models import Notification, UnknownNotification + +from agent_flow.backends import codex_transport as transport_module +from agent_flow.backends.base import ResultEvent +from agent_flow.backends.codex_events import CodexClient +from agent_flow.backends.codex_transport import CodexTransport, notification_data +from agent_flow.tools import tool +from agent_flow.types import AgentTextEvent, RateLimitWarningEvent + + +class FakeSdkClient(transport_module._ObservedSdkClient): + """Use real SDK notification parsing with an in-memory request/reader seam.""" + + def __init__(self, config, approval_handler, observer): + super().__init__(config, approval_handler, observer) + self.server_request = approval_handler + self.notifications: queue.Queue[Notification | BaseException] = queue.Queue() + self.turn_requests: queue.Queue[dict[str, Any]] = queue.Queue() + self.requests: list[tuple[str, dict]] = [] + self.on_request: Callable[[str, dict], dict | None] | None = None + self.started = False + self.initialized = False + self.closed = False + + def start(self): + self.started = True + + def initialize(self): + self.initialized = True + + def close(self): + self.closed = True + self.notifications.put(EOFError("fake SDK closed")) + + def next_notification(self): + event = self.notifications.get(timeout=5) + if isinstance(event, BaseException): + raise event + return event + + def emit(self, method: str, params: dict): + event = self._coerce_notification(method, params) + self.notifications.put(event) + return event + + def request(self, method, params, *, response_model): + self.requests.append((method, params)) + if method == "turn/start": + self.turn_requests.put(params) + response = self.on_request(method, params) if self.on_request is not None else None + if response is None: + response = ( + {"turn": {"id": f"turn-{params['threadId']}"}} if method == "turn/start" else {} + ) + return response_model.model_validate(response) + + async def next_turn(self): + return await asyncio.to_thread(self.turn_requests.get, True, 2) + + +@pytest.fixture +async def transport(monkeypatch): + monkeypatch.setattr(transport_module, "_ObservedSdkClient", FakeSdkClient) + instance = CodexTransport(CodexConfig(codex_bin="/unused/codex")) + await instance.start() + try: + yield instance + finally: + await instance.close() + + +def _turn(thread_id: str, status: str = "completed"): + return { + "threadId": thread_id, + "turn": {"id": f"turn-{thread_id}", "items": [], "itemsView": "full", "status": status}, + } + + +def _thread(thread_id: str, parent: str): + return { + "thread": { + "id": thread_id, + "parentThreadId": parent, + "sessionId": "session-one", + "cliVersion": "0.154.0", + "createdAt": 0, + "updatedAt": 0, + "cwd": "/tmp", + "ephemeral": False, + "historyMode": "legacy", + "modelProvider": "openai", + "preview": "", + "source": "appServer", + "status": {"type": "idle"}, + "turns": [], + } + } + + +def _text(thread_id: str, text: str): + return { + "threadId": thread_id, + "turnId": f"turn-{thread_id}", + "completedAtMs": 100, + "item": { + "id": f"message-{thread_id}", + "type": "agentMessage", + "text": text, + "phase": "final_answer", + }, + } + + +async def _collect(stream): + return [event async for event in stream] + + +def test_unknown_sdk_notifications_expose_original_params(): + params = {"threadId": "child", "newField": {"value": 2}} + event = Notification("future/notification", UnknownNotification(params=params)) + assert notification_data(event) == params + + +async def test_completion_before_start_response_and_global_events_are_retained(transport): + sdk = transport._client + + def complete_before_response(method, params): + if method != "turn/start": + return None + sdk.emit("thread/started", _thread("child", "root")) + sdk.emit("account/rateLimits/updated", {"rateLimits": {"primary": {"usedPercent": 90}}}) + sdk.emit("item/completed", _text("child", "child result")) + sdk.emit("turn/completed", _turn("root")) + return {"turn": {"id": "turn-root"}} + + sdk.on_request = complete_before_response + events = await asyncio.wait_for(_collect(transport.stream_turn("root", "hello")), timeout=2) + assert [event.method for event in events] == [ + "thread/started", + "account/rateLimits/updated", + "item/completed", + "turn/completed", + ] + assert notification_data(events[0])["thread"]["parentThreadId"] == "root" + assert isinstance(events[0].payload, ThreadStartedNotification) + assert not any(method == "turn/interrupt" for method, _ in sdk.requests) + + +async def test_concurrent_clients_filter_other_roots_and_keep_account_updates(transport): + sdk = transport._client + clients = {thread: CodexClient(transport, thread) for thread in ("alpha", "beta")} + tasks = { + thread: asyncio.create_task(_collect(client.send_message("work"))) + for thread, client in clients.items() + } + try: + started = [await sdk.next_turn(), await sdk.next_turn()] + assert {params["threadId"] for params in started} == {"alpha", "beta"} + for thread in clients: + sdk.emit("turn/started", _turn(thread, "inProgress")) + sdk.emit("account/rateLimits/updated", {"rateLimits": {"primary": {"usedPercent": 90}}}) + sdk.emit("item/completed", _text("beta", "beta only")) + sdk.emit("item/completed", _text("alpha", "alpha only")) + sdk.emit("turn/completed", _turn("alpha")) + sdk.emit("turn/completed", _turn("beta")) + results = await asyncio.wait_for(asyncio.gather(*tasks.values()), timeout=2) + for thread, events in zip(tasks, results): + assert [event.text for event in events if isinstance(event, AgentTextEvent)] == [ + f"{thread} only" + ] + assert [event.text for event in events if isinstance(event, ResultEvent)] == [ + f"{thread} only" + ] + warnings = [event for event in events if isinstance(event, RateLimitWarningEvent)] + assert len(warnings) == 1 + assert warnings[0].utilization == 0.9 + finally: + for task in tasks.values(): + task.cancel() + await asyncio.gather(*tasks.values(), return_exceptions=True) + + +async def test_sdk_eof_wakes_waiting_turn_and_subsequent_requests_fail(transport): + task = asyncio.create_task(_collect(transport.stream_turn("root", "work"))) + await transport._client.next_turn() + transport._client.notifications.put(EOFError("reader reached EOF")) + with pytest.raises(EOFError, match="reader reached EOF"): + await asyncio.wait_for(task, timeout=2) + with pytest.raises(RuntimeError, match="transport failed"): + await transport.request("thread/read", {"threadId": "root"}) + + +async def test_cancelled_turn_interrupts_and_releases_the_session(transport): + sdk = transport._client + observed = asyncio.Event() + + async def consume(): + async for _event in transport.stream_turn("root", "work"): + observed.set() + + task = asyncio.create_task(consume()) + await sdk.next_turn() + sdk.emit("turn/started", _turn("root", "inProgress")) + await asyncio.wait_for(observed.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert ("turn/interrupt", {"threadId": "root", "turnId": "turn-root"}) in sdk.requests + assert not transport._queues + + retry = asyncio.create_task(_collect(transport.stream_turn("root", "retry"))) + await sdk.next_turn() + sdk.emit("turn/completed", _turn("root")) + await asyncio.wait_for(retry, timeout=2) + + +async def test_cancellation_while_start_response_is_pending_interrupts_created_turn(transport): + release_response = threading.Event() + + def wait_before_response(method, params): + if method == "turn/start": + if not release_response.wait(timeout=2): + raise TimeoutError("test did not release the turn/start response") + return None + + transport._client.on_request = wait_before_response + task = asyncio.create_task(_collect(transport.stream_turn("root", "work"))) + try: + await transport._client.next_turn() + task.cancel() + release_response.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=2) + assert ( + "turn/interrupt", + {"threadId": "root", "turnId": "turn-root"}, + ) in transport._client.requests + assert not transport._queues + finally: + release_response.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +async def test_same_session_rejects_concurrent_turn_without_disturbing_first(transport): + first = asyncio.create_task(_collect(transport.stream_turn("root", "first"))) + await transport._client.next_turn() + with pytest.raises(RuntimeError, match="Concurrent turns"): + await _collect(transport.stream_turn("root", "second")) + transport._client.emit("turn/completed", _turn("root")) + assert await asyncio.wait_for(first, timeout=2) + + +async def test_consecutive_client_turns_release_transport_before_returning(transport): + def complete(method, params): + if method == "turn/start": + transport._client.emit("item/completed", _text("root", "done")) + transport._client.emit("turn/completed", _turn("root")) + + transport._client.on_request = complete + client = CodexClient(transport, "root") + for prompt in ("initial", "correct missing required calls"): + events = await _collect(client.send_message(prompt)) + assert isinstance(events[-1], ResultEvent) + assert not transport._active_threads + assert not transport._queues + + +async def test_dynamic_tool_runs_on_owner_loop_with_registration_context_and_child_ancestry( + transport, +): + current_workflow = contextvars.ContextVar("test_codex_workflow", default="unset") + owner_loop = asyncio.get_running_loop() + observed = [] + + @tool("read", "Read data", {"query": str}) + async def read(args): + observed.append((asyncio.get_running_loop(), current_workflow.get(), args)) + await asyncio.sleep(0) + return {"content": [{"type": "text", "text": "result"}]} + + token = current_workflow.set("workflow-one") + transport.register_tools("root", [read]) + current_workflow.reset(token) + + def request_from_reader(): + transport._client.emit( + "thread/started", {"thread": {"id": "child", "parentThreadId": "root"}} + ) + transport._client.emit( + "thread/started", {"thread": {"id": "grandchild", "parentThreadId": "child"}} + ) + return transport._client.server_request( + "item/tool/call", + {"threadId": "grandchild", "tool": "read", "arguments": {"query": "value"}}, + ) + + response = await asyncio.wait_for(asyncio.to_thread(request_from_reader), timeout=2) + assert response == {"contentItems": [{"type": "inputText", "text": "result"}], "success": True} + assert observed == [(owner_loop, "workflow-one", {"query": "value"})] + + transport.unregister_tools("root") + unavailable = await asyncio.to_thread( + transport._client.server_request, + "item/tool/call", + {"threadId": "grandchild", "tool": "read", "arguments": {}}, + ) + assert unavailable["success"] is False + assert "not registered" in unavailable["contentItems"][0]["text"] + + +@pytest.mark.parametrize( + "scenario", ["unknown", "wrong_arguments", "handler_error", "unsupported_result"] +) +async def test_dynamic_tool_failures_return_explicit_unsuccessful_results(transport, scenario): + @tool("work", "Do work", {}) + async def work(args): + if scenario == "handler_error": + raise ValueError("backend unavailable") + return {"content": [{"type": "audio", "data": "YWJj", "mimeType": "audio/wav"}]} + + transport.register_tools("root", [work]) + params = { + "threadId": "root", + "tool": "missing" if scenario == "unknown" else "work", + "arguments": [] if scenario == "wrong_arguments" else {}, + } + response = await asyncio.wait_for( + asyncio.to_thread(transport._client.server_request, "item/tool/call", params), timeout=2 + ) + assert response["success"] is False + expected = { + "unknown": "not registered", + "wrong_arguments": "requires an object", + "handler_error": "backend unavailable", + "unsupported_result": "Unsupported tool result", + }[scenario] + assert expected in response["contentItems"][0]["text"] diff --git a/agent-flow/tests/test_console.py b/agent-flow/tests/test_console.py index 653bdaa5f08e..24182acdd8c6 100644 --- a/agent-flow/tests/test_console.py +++ b/agent-flow/tests/test_console.py @@ -13,6 +13,7 @@ SessionInitEvent, ThinkingEvent, ToolCallEvent, + UsageInfo, ) @@ -31,6 +32,16 @@ def _capture_console(monkeypatch) -> Console: return fake +def test_usage_labels_thread_estimate_separately_from_actual_cost(monkeypatch): + fake = _capture_console(monkeypatch) + console_module.print_agent_completed( + "planner", usage=UsageInfo(cost_usd=0.25, estimated_thread_cost_usd=1.75) + ) + output = fake._buffer_text.getvalue() # type: ignore[attr-defined] + assert "cost $0.2500" in output + assert "estimated thread cost $1.7500" in output + + def test_main_agent_tool_call_has_no_subagent_badge(monkeypatch): fake = _capture_console(monkeypatch) console_module.print_tool_call( diff --git a/agent-flow/tests/test_hooks.py b/agent-flow/tests/test_hooks.py index ad882ef16b5c..31c73f0bac48 100644 --- a/agent-flow/tests/test_hooks.py +++ b/agent-flow/tests/test_hooks.py @@ -289,3 +289,34 @@ async def test_hook_uses_custom_reason(self, tmp_path): "decision": "block", "reason": "please log progress!", } + + +class TestSharedRequiredToolPolicy: + def test_all_policy_and_legacy_any_helper_remain_distinct(self, tmp_path): + from agent_flow.hooks import RequiredToolPolicy + from agent_flow.types import ToolCallEvent + + required = ("append_progress", "update_status") + transcript = _write_transcript( + tmp_path / "turn.jsonl", + [ + _user("work"), + _assistant_tool_use("mcp__agent-tools__append_progress"), + ], + ) + assert called_required_tool_this_turn(transcript, required) + policy = RequiredToolPolicy(required) + policy.record(ToolCallEvent("mcp__agent-tools__append_progress", {})) + assert policy.missing == ("update_status",) + policy.record(ToolCallEvent("update_status", {})) + assert policy.missing == () + + def test_qualified_requirement_does_not_match_another_server(self): + from agent_flow.hooks import RequiredToolPolicy + from agent_flow.types import ToolCallEvent + + policy = RequiredToolPolicy(("mcp__trusted__append_progress",)) + policy.record(ToolCallEvent("mcp__other__append_progress", {})) + assert policy.missing == ("mcp__trusted__append_progress",) + policy.record(ToolCallEvent("mcp__trusted__append_progress", {})) + assert policy.missing == () diff --git a/agent-flow/tests/test_sessions.py b/agent-flow/tests/test_sessions.py index a60222f4a59e..0c148ae94c10 100644 --- a/agent-flow/tests/test_sessions.py +++ b/agent-flow/tests/test_sessions.py @@ -147,3 +147,51 @@ async def test_async_context_manager_releases_persistent_client_and_backend(): assert backend.create_client_calls == 1 assert backend.client_exit_count == 1 assert backend.exit_count == 1 + + +async def test_required_tools_do_not_reuse_previous_persistent_invocation_calls(): + from dataclasses import replace + + from agent_flow.hooks import RequiredToolCallError + from agent_flow.types import ToolCallEvent + + backend = FakeBackend( + [ + { + "turns": [ + {"tool_calls": [ToolCallEvent("append_progress", {})]}, + {"tool_calls": []}, + {"tool_calls": []}, + ] + } + ] + ) + config = replace(_config(), print_activity=False, required_tools=("append_progress",)) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer(config) as layer: + assert await layer.aforward("first task") == "ok" + with pytest.raises(RequiredToolCallError, match="append_progress"): + await layer.aforward("second task") + assert backend.clients[0].closed + assert backend.create_client_calls == 1 + assert backend.clients[0].send_count == 3 + + +async def test_persistent_session_requires_stable_prompt_when_first_was_none(): + from dataclasses import replace + + backend = FakeBackend() + prompts = iter([None, "new instructions"]) + config = replace(_config(), system_prompt=None, print_activity=False) + with patch("agent_flow.layers.create_backend", return_value=backend): + async with AgentLayer( + config, + prompt_builder=lambda content: AgentRequest( + content, + system_prompt=next(prompts), + ), + ) as layer: + await layer.aforward("first") + with pytest.raises(ValueError, match="stable system prompt"): + await layer.aforward("second") + assert backend.create_client_calls == 1 diff --git a/agent-flow/tests/test_tools.py b/agent-flow/tests/test_tools.py new file mode 100644 index 000000000000..c3e3cef4f3a6 --- /dev/null +++ b/agent-flow/tests/test_tools.py @@ -0,0 +1,363 @@ +from __future__ import annotations + +import copy +from enum import Enum +from types import SimpleNamespace +from typing import Annotated, Any, Literal, NotRequired, Required, TypedDict +from unittest.mock import AsyncMock + +import pytest + +from agent_flow.tools import ( + ToolDefinition, + normalize_input_schema, + normalize_tool, + normalize_tool_result, + tool, + tool_result_to_claude, + tool_result_to_codex, +) + + +class Location(TypedDict): + path: str + line: NotRequired[int] + + +class SearchArgs(TypedDict, total=False): + query: Required[str] + locations: list[Location] + limit: int | None + modes: list[Literal["text", "image"]] + + +class ExtendedSearchArgs(SearchArgs): + exact: bool + + +class RecursiveArgs(TypedDict): + child: NotRequired[RecursiveArgs] + + +def test_shorthand_schema_supports_primitives_and_a_field_named_type(): + schema = normalize_input_schema({"type": str, "count": int, "ratio": float, "ok": bool}) + assert schema == { + "type": "object", + "properties": { + "type": {"type": "string"}, + "count": {"type": "integer"}, + "ratio": {"type": "number"}, + "ok": {"type": "boolean"}, + }, + "required": ["type", "count", "ratio", "ok"], + } + + +def test_shorthand_fields_can_use_json_schema_keyword_names(): + schema = normalize_input_schema({"properties": list[str], "required": NotRequired[bool]}) + assert schema["properties"] == { + "properties": {"type": "array", "items": {"type": "string"}}, + "required": {"type": "boolean"}, + } + assert schema["required"] == ["properties"] + + +def test_typeddict_required_and_optional_fields_survive_future_annotations(): + schema = normalize_input_schema(ExtendedSearchArgs) + assert schema["required"] == ["query", "exact"] + assert schema["properties"]["locations"] == { + "type": "array", + "items": { + "type": "object", + "properties": {"path": {"type": "string"}, "line": {"type": "integer"}}, + "required": ["path"], + }, + } + assert schema["properties"]["limit"] == {"anyOf": [{"type": "integer"}, {"type": "null"}]} + assert schema["properties"]["modes"] == {"type": "array", "items": {"enum": ["text", "image"]}} + + +def test_schema_supports_annotations_mapping_and_tuples(): + schema = normalize_input_schema( + { + "metadata": dict[str, Any], + "numbers": tuple[int, ...], + "pair": tuple[str, int], + "label": Annotated[str, "Label"], + "optional": Annotated[NotRequired[bool], "Optional flag"], + } + ) + assert schema["properties"]["metadata"] == {"type": "object", "additionalProperties": {}} + assert schema["properties"]["numbers"] == {"type": "array", "items": {"type": "integer"}} + assert schema["properties"]["pair"] == { + "type": "array", + "prefixItems": [{"type": "string"}, {"type": "integer"}], + "minItems": 2, + "maxItems": 2, + } + assert schema["properties"]["label"] == {"type": "string"} + assert "optional" not in schema["required"] + + +def test_enum_schema(): + class Mode(Enum): + READ = "read" + WRITE = "write" + + assert normalize_input_schema({"mode": Mode})["properties"]["mode"] == { + "enum": ["read", "write"] + } + + +def test_explicit_json_schema_is_copied_without_reinterpretation(): + original = { + "type": "object", + "properties": {"names": {"type": "array", "items": {"type": "string"}}}, + } + pristine = copy.deepcopy(original) + normalized = normalize_input_schema(original) + assert normalized == pristine + assert normalized is not original + normalized["properties"]["names"]["items"]["type"] = "integer" + assert original == pristine + + +@pytest.mark.parametrize( + "original, expected", + [ + ( + {"type": "object", "additionalProperties": {"type": "string"}}, + {"type": "object", "properties": {}, "additionalProperties": {"type": "string"}}, + ), + ( + {"$ref": "#/$defs/arguments", "$defs": {"arguments": {"type": "object"}}}, + { + "type": "object", + "properties": {}, + "allOf": [{"$ref": "#/$defs/arguments"}], + "$defs": {"arguments": {"type": "object"}}, + }, + ), + ( + { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}}, + "allOf": [{"required": ["name"]}], + }, + { + "type": "object", + "properties": {"name": {"type": "string"}}, + "allOf": [{"type": ["object", "null"]}, {"required": ["name"]}], + }, + ), + ( + {"anyOf": [{"required": ["a"]}, {"required": ["b"]}]}, + { + "type": "object", + "properties": {}, + "anyOf": [{"required": ["a"]}, {"required": ["b"]}], + }, + ), + ], +) +def test_explicit_json_schema_is_normalized_to_mcp_object_shape(original, expected): + # The Claude SDK only recognises a JSON Schema when it has a string ``type`` and a + # ``properties`` key; anything else is re-read as Python shorthand. Root keywords that + # do not fit that shape move into ``allOf`` with their meaning intact. + pristine = copy.deepcopy(original) + assert normalize_input_schema(original) == expected + assert original == pristine + + +@pytest.mark.parametrize( + "schema, message", + [ + ({"bad": object}, "Unsupported Python type"), + ({"bad": dict[int, str]}, "keys must be strings"), + (RecursiveArgs, "Recursive TypedDict"), + ({"type": "object", "default": object()}, "JSON-serializable"), + (str, "must be a JSON Schema"), + ], +) +def test_invalid_schema_fails_during_registration(schema, message): + with pytest.raises(TypeError, match=message): + normalize_input_schema(schema) + + +async def test_framework_decorator_produces_callable_handler_without_sdk_dependency(): + @tool("echo", "Return the input", {"text": str}, annotations={"readOnlyHint": True}) + async def echo(args): + return {"content": [{"type": "text", "text": args["text"]}]} + + assert isinstance(echo, ToolDefinition) + assert echo.annotations == {"readOnlyHint": True} + assert echo.input_schema["properties"] == {"text": {"type": "string"}} + assert await echo.handler({"text": "hello"}) == {"content": [{"type": "text", "text": "hello"}]} + + +async def test_legacy_sdk_tool_conversion_preserves_annotations_and_handler(): + from claude_agent_sdk import ToolAnnotations + from claude_agent_sdk import tool as claude_tool + + @claude_tool( + "read", + "Read data", + SearchArgs, + annotations=ToolAnnotations(readOnlyHint=True, maxResultSizeChars=500), + ) + async def read(args): + return {"content": [{"type": "text", "text": args["query"]}]} + + normalized = normalize_tool(read) + assert normalized.handler is read.handler + assert normalized.input_schema["required"] == ["query"] + assert normalized.annotations["readOnlyHint"] is True + assert normalized.meta == {"anthropic/maxResultSizeChars": 500} + assert await normalized.handler({"query": "value"}) == { + "content": [{"type": "text", "text": "value"}] + } + + +def test_normalizing_tool_does_not_mutate_annotations(): + original = SimpleNamespace( + name="read", + description="Read", + input_schema={}, + handler=AsyncMock(), + annotations={"readOnlyHint": True}, + ) + normalized = normalize_tool(original) + normalized.annotations["readOnlyHint"] = False + assert original.annotations == {"readOnlyHint": True} + + +@pytest.mark.parametrize("flag", ["is_error", "isError"]) +def test_error_results_preserve_both_mcp_spellings(flag): + result = {"content": [{"type": "text", "text": "failed"}], flag: True} + assert tool_result_to_codex(result) == { + "contentItems": [{"type": "inputText", "text": "failed"}], + "success": False, + } + assert tool_result_to_claude(result)["is_error"] is True + + +def test_resources_are_readable_instead_of_empty_success(): + result = { + "content": [ + { + "type": "resource_link", + "name": "Report", + "uri": "file:///report.txt", + "description": "Findings", + }, + {"type": "resource", "resource": {"uri": "memo://one", "text": "The result is 42."}}, + ] + } + assert tool_result_to_codex(result) == { + "contentItems": [ + {"type": "inputText", "text": "Report\nfile:///report.txt\nFindings"}, + {"type": "inputText", "text": "The result is 42."}, + ], + "success": True, + } + assert tool_result_to_claude(result)["content"][1] == { + "type": "text", + "text": "The result is 42.", + } + + +def test_images_use_data_urls_in_codex_and_mcp_bytes_in_claude(): + result = {"content": [{"type": "image", "data": "YWJj", "mimeType": "image/png"}]} + assert tool_result_to_codex(result) == { + "contentItems": [{"type": "inputImage", "imageUrl": "data:image/png;base64,YWJj"}], + "success": True, + } + assert tool_result_to_claude(result)["content"] == result["content"] + + +def test_image_data_url_round_trip(): + result = {"content": [{"type": "image", "image_url": "data:image/webp;base64,YWJj"}]} + converted = tool_result_to_claude(result)["content"][0] + assert converted["data"] == "YWJj" + assert converted["mimeType"] == "image/webp" + assert ( + tool_result_to_codex(result)["contentItems"][0]["imageUrl"] == "data:image/webp;base64,YWJj" + ) + + +def test_remote_image_url_is_supported_by_codex_and_explicitly_rejected_by_claude(): + result = {"content": [{"type": "image", "image_url": "https://example.com/image.png"}]} + assert ( + tool_result_to_codex(result)["contentItems"][0]["imageUrl"] + == "https://example.com/image.png" + ) + with pytest.raises(ValueError, match="remote image URLs"): + tool_result_to_claude(result) + + +@pytest.mark.parametrize( + "content, message", + [ + ( + [{"type": "audio", "data": "YWJj", "mimeType": "audio/wav"}], + "Unsupported tool result content", + ), + ( + [{"type": "resource", "resource": {"uri": "binary://one", "blob": "YWJj"}}], + "Binary embedded", + ), + ([{"type": "image", "data": "YWJj"}], "mimeType"), + ([{"type": "image"}], "requires base64 data"), + ([{"type": "text", "text": None}], "requires a string"), + ([{"type": "resource_link"}], "resource link needs"), + ], +) +def test_unrepresentable_content_raises_instead_of_disappearing(content, message): + with pytest.raises(ValueError, match=message): + tool_result_to_codex({"content": content}) + + +def test_pydantic_mcp_result_conversion(): + from mcp.types import CallToolResult, TextContent + + result = CallToolResult(content=[TextContent(type="text", text="failure")], isError=True) + assert tool_result_to_codex(result) == { + "contentItems": [{"type": "inputText", "text": "failure"}], + "success": False, + } + + +def test_structured_content_remains_visible(): + result = {"content": [], "structuredContent": {"answer": 42}} + assert tool_result_to_codex(result)["contentItems"] == [ + {"type": "inputText", "text": '{"answer": 42}'} + ] + + +def test_result_normalization_is_idempotent_and_does_not_mutate_input(): + original = { + "content": [], + "structuredContent": {"answer": 42}, + "isError": True, + "is_error": True, + } + pristine = copy.deepcopy(original) + normalized = normalize_tool_result(original) + assert normalize_tool_result(normalized) == normalized + assert "isError" not in normalized + assert original == pristine + + +def test_empty_result_retains_its_error_status(): + assert tool_result_to_codex({"content": [], "isError": True}) == { + "contentItems": [{"type": "inputText", "text": ""}], + "success": False, + } + + +@pytest.mark.parametrize( + "result", [None, "text", {"content": "text"}, {"content": ""}, {"content": ["text"]}] +) +def test_invalid_result_shape_is_rejected(result): + with pytest.raises(TypeError): + normalize_tool_result(result) diff --git a/agent-flow/tests/workflows/agent_team/test_workflow.py b/agent-flow/tests/workflows/agent_team/test_workflow.py index 990bd536a6be..edaf744c7ca4 100644 --- a/agent-flow/tests/workflows/agent_team/test_workflow.py +++ b/agent-flow/tests/workflows/agent_team/test_workflow.py @@ -1707,50 +1707,25 @@ def test_every_registered_tool_is_documented_in_the_role_prompt(tmp_path): def test_coder_and_reviewer_required_tools_include_status_update(tmp_path): - """The composed Stop hook must enforce both progress and status calls.""" + """Every role declares all required calls, including Codex-backed roles.""" module = _load_module() workflow = _make_workflow(module, tmp_path) try: - # Each per-tool hook is a separate matcher; with two required tools - # we expect two matchers stacked together (composition = AND). - coder_hooks = workflow.coder.config.backend.hooks - reviewer_hooks = workflow.reviewer.config.backend.hooks - plan_drafter_hooks = workflow.plan_drafter.config.backend.hooks - plan_reviewer_hooks = workflow.plan_reviewer.config.backend.hooks - qa_hooks = workflow.qa.config.backend.hooks - - assert coder_hooks is not None - assert len(coder_hooks["Stop"]) == 2 - assert reviewer_hooks is not None - assert len(reviewer_hooks["Stop"]) == 2 - - # PlanDrafter, PlanReviewer, and QA each have a single required - # tool — one matcher. - assert plan_drafter_hooks is not None - assert len(plan_drafter_hooks["Stop"]) == 1 - assert plan_reviewer_hooks is not None - assert len(plan_reviewer_hooks["Stop"]) == 1 - assert qa_hooks is not None - assert len(qa_hooks["Stop"]) == 1 + expected = { + "coder": ("append_coder_progress", "update_status"), + "reviewer": ("append_reviewer_progress", "update_status"), + "plan_drafter": ("append_plan_drafter_progress",), + "plan_reviewer": ("append_plan_reviewer_progress",), + "qa": ("append_qa_progress",), + } + for role, names in expected.items(): + config = getattr(workflow, role).config + assert config.required_tools == names + assert config.backend.hooks is None finally: workflow.close() -def test_compose_required_tools_hooks_handles_empty_and_single(tmp_path): - """Helper returns ``None`` for empty and one matcher for a single tool.""" - module = _load_module() - - assert module._compose_required_tools_hooks([]) is None - - single = module._compose_required_tools_hooks(["append_plan_drafter_progress"]) - assert single is not None - assert len(single["Stop"]) == 1 - - pair = module._compose_required_tools_hooks(["append_coder_progress", "update_status"]) - assert pair is not None - assert len(pair["Stop"]) == 2 - - def test_qa_uses_stateless_session(tmp_path): """QA's session must be stateless so each iteration starts fresh.""" module = _load_module() @@ -2306,6 +2281,7 @@ def test_resume_from_qa_skips_coder_and_reviewer(tmp_path): assert state.done is True +@pytest.mark.live_backend def test_claude_agents_baseline_context_under_20_percent(tmp_path): """Each Claude-backed team agent's pre-input context must stay < 20%. @@ -3334,6 +3310,7 @@ def test_no_mcp_mode_builds_agents_without_tools(tmp_path): for agent in (wf.plan_drafter, wf.plan_reviewer, wf.coder, wf.reviewer, wf.qa): assert agent.config.backend.tools is None assert agent.config.backend.hooks is None + assert agent.config.required_tools == () assert agent.config.human_input_enabled is False # The backend split itself is untouched: only the tools are dropped. assert wf.plan_drafter.config.backend.kind == "codex" @@ -3349,7 +3326,7 @@ def test_mcp_mode_is_the_default_and_keeps_tools(tmp_path): try: assert wf.use_in_process_tools is True assert wf.plan_drafter.config.backend.tools is not None - assert wf.plan_drafter.config.backend.hooks is not None + assert wf.plan_drafter.config.required_tools == ("append_plan_drafter_progress",) assert wf.plan_drafter.config.human_input_enabled is True finally: wf.close() @@ -3364,8 +3341,10 @@ def test_no_mcp_mode_resets_rebuild_agents_without_tools(tmp_path): wf._reset_reviewer() assert wf.coder.config.backend.tools is None assert wf.coder.config.backend.hooks is None + assert wf.coder.config.required_tools == () assert wf.reviewer.config.backend.tools is None assert wf.reviewer.config.backend.hooks is None + assert wf.reviewer.config.required_tools == () finally: wf.close() diff --git a/agent-flow/tests/workflows/modeling_bringup/test_workflow.py b/agent-flow/tests/workflows/modeling_bringup/test_workflow.py index a59ffa511a6e..0f37f24edc5d 100644 --- a/agent-flow/tests/workflows/modeling_bringup/test_workflow.py +++ b/agent-flow/tests/workflows/modeling_bringup/test_workflow.py @@ -818,6 +818,7 @@ def _fmt(cells: tuple[str, ...]) -> str: print(f"skipped: {'; '.join(unavailable)}") +@pytest.mark.live_backend @pytest.mark.parametrize( "include_slurm_environment", [False, True], diff --git a/agent-flow/tests/workflows/perf_analyze/test_workflow.py b/agent-flow/tests/workflows/perf_analyze/test_workflow.py index 65bd04ea1836..aaa4e47035b2 100644 --- a/agent-flow/tests/workflows/perf_analyze/test_workflow.py +++ b/agent-flow/tests/workflows/perf_analyze/test_workflow.py @@ -479,8 +479,9 @@ def test_all_agents_use_claude_code_backend(tmp_path): ): assert layer.config.backend.kind == "claude-code" assert layer.config.backend.model == CLAUDE_CODE_DEFAULT_MODEL - # Each role is gated by a required-tool stop hook. - assert layer.config.backend.hooks is not None + # Each role must record progress through the shared policy. + assert layer.config.required_tools == (f"append_{layer.config.name.lower()}_progress",) + assert layer.config.backend.hooks is None finally: workflow.close() diff --git a/agent-flow/tests/workflows/perf_optimize/test_workflow.py b/agent-flow/tests/workflows/perf_optimize/test_workflow.py index 000df12755bf..9987f9a800c6 100644 --- a/agent-flow/tests/workflows/perf_optimize/test_workflow.py +++ b/agent-flow/tests/workflows/perf_optimize/test_workflow.py @@ -2725,7 +2725,8 @@ def test_all_agents_use_claude_code_backend_with_scoped_sessions(tmp_path): layer = getattr(workflow, role) assert layer.config.backend.kind == "claude-code", role assert layer.config.backend.model == CLAUDE_CODE_DEFAULT_MODEL, role - assert layer.config.backend.hooks is not None, role + assert layer.config.required_tools == (f"append_{role}_progress",), role + assert layer.config.backend.hooks is None, role # The judges are stateless (fresh eyes per verdict); the # optimizer's persistent session is additionally reset per # item by the orchestrator (covered by