From 819fe84726a3c07bac1e647575889a729ecd0894 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:08:57 -0700 Subject: [PATCH 1/5] [None][chore] agent-flow: bump backend SDK pins and default Codex model Pin claude-agent-sdk to 0.2.152 and openai-codex to 0.154.0, matching the versions installed in the dev environment, and move the default Codex model from gpt-5.6-sol to gpt-6-astra. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- agent-flow/agent_flow/config.py | 2 +- agent-flow/pyproject.toml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/agent-flow/agent_flow/config.py b/agent-flow/agent_flow/config.py index d0c61d5956c5..34827cd6a314 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) 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] From a328e29cea9e310b22d0099722be53f5b501fafe Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:34:59 -0700 Subject: [PATCH 2/5] [None][feat] agent-flow: backend-independent tools and required-tool policy Move tool definition and MCP result conversion into a shared `tools` module so both backends consume one definition format. Python shorthand and `TypedDict` inputs normalize to JSON Schema, resource links and embedded text resources become readable tool output, and unsupported result types return explicit tool errors instead of appearing to succeed with missing output. Replace the per-workflow stack of Claude `Stop` hooks with `AgentLayerConfig.required_tools`, which works on either backend. `AgentLayer` observes root-agent tool calls, sends one corrective turn on the same session when calls are missing, and raises `RequiredToolCallError` if that bounded correction still leaves them missing. Subagent and prior-invocation calls do not discharge the requirement. Token accounting sums both turns, while context occupancy stays the final snapshot rather than being summed. `require_tool_call_stop_hook` remains available as a Claude-native alternative. Split the Codex backend into `codex_config`, `codex_events`, and `codex_transport`. Session options translate into `thread/start.config` without editing the user's Codex configuration, unsupported tool restrictions fail at registration rather than being silently ignored, and the transport observes decoded notifications so descendant turns and account events reach the shared event stream. Drop the `ping` bootstrap turn from context-budget inspection so inspection never costs a model turn, and report `UsageInfo.estimated_thread_cost_usd` separately from per-turn `cost_usd`. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- agent-flow/README.md | 63 +- agent-flow/agent_flow/__init__.py | 10 +- agent-flow/agent_flow/backends/claude_code.py | 33 +- agent-flow/agent_flow/backends/codex.py | 911 ++----------- .../agent_flow/backends/codex_config.py | 292 ++++ .../agent_flow/backends/codex_events.py | 420 ++++++ .../agent_flow/backends/codex_transport.py | 236 ++++ agent-flow/agent_flow/config.py | 25 +- agent-flow/agent_flow/console.py | 5 + agent-flow/agent_flow/hooks.py | 70 +- agent-flow/agent_flow/layers.py | 175 ++- agent-flow/agent_flow/tools.py | 302 +++++ agent-flow/agent_flow/types.py | 2 + .../agent_flow/workflows/agent_team/README.md | 7 +- .../workflows/agent_team/progress.py | 2 +- .../agent_flow/workflows/agent_team/status.py | 2 +- .../workflows/agent_team/workflow.py | 23 +- .../workflows/perf_analyze/progress.py | 2 +- .../workflows/perf_analyze/workflow.py | 20 +- .../workflows/perf_optimize/progress.py | 2 +- .../workflows/perf_optimize/workflow.py | 20 +- agent-flow/tests/conftest.py | 24 + agent-flow/tests/helpers.py | 26 +- agent-flow/tests/test_agent_layer.py | 227 ++++ agent-flow/tests/test_backends.py | 1176 +---------------- agent-flow/tests/test_codex_backend.py | 221 ++++ agent-flow/tests/test_codex_config.py | 206 +++ agent-flow/tests/test_codex_events.py | 504 +++++++ agent-flow/tests/test_codex_transport.py | 350 +++++ agent-flow/tests/test_console.py | 11 + agent-flow/tests/test_hooks.py | 31 + agent-flow/tests/test_sessions.py | 48 + agent-flow/tests/test_tools.py | 321 +++++ .../workflows/agent_team/test_workflow.py | 55 +- .../modeling_bringup/test_workflow.py | 1 + .../workflows/perf_analyze/test_workflow.py | 5 +- .../workflows/perf_optimize/test_workflow.py | 3 +- 37 files changed, 3656 insertions(+), 2175 deletions(-) create mode 100644 agent-flow/agent_flow/backends/codex_config.py create mode 100644 agent-flow/agent_flow/backends/codex_events.py create mode 100644 agent-flow/agent_flow/backends/codex_transport.py create mode 100644 agent-flow/agent_flow/tools.py create mode 100644 agent-flow/tests/test_codex_backend.py create mode 100644 agent-flow/tests/test_codex_config.py create mode 100644 agent-flow/tests/test_codex_events.py create mode 100644 agent-flow/tests/test_codex_transport.py create mode 100644 agent-flow/tests/test_tools.py diff --git a/agent-flow/README.md b/agent-flow/README.md index e7b8e5b5ada9..166df5665ce1 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,67 @@ 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. 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..b7b145dd62e5 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,42 @@ 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, + ) + 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 34827cd6a314..76d81617bddc 100644 --- a/agent-flow/agent_flow/config.py +++ b/agent-flow/agent_flow/config.py @@ -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..1bd1721b34ab --- /dev/null +++ b/agent-flow/agent_flow/tools.py @@ -0,0 +1,302 @@ +"""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 normalize_input_schema(schema: Any) -> dict[str, Any]: + """Copy JSON Schema or convert a shorthand dictionary/TypedDict to JSON Schema.""" + 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 = 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/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..96417fe150b0 100644 --- a/agent-flow/tests/test_backends.py +++ b/agent-flow/tests/test_backends.py @@ -2,7 +2,7 @@ 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 @@ -24,13 +24,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 +72,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" @@ -988,6 +723,28 @@ async def receive_response(): 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"} + ] + 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 +784,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 +802,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 +853,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 +904,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 +1040,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..2ef52803eb6e --- /dev/null +++ b/agent-flow/tests/test_codex_backend.py @@ -0,0 +1,221 @@ +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 + + 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 == "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_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"), + ( + {"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..32f820a39bfe --- /dev/null +++ b/agent-flow/tests/test_tools.py @@ -0,0 +1,321 @@ +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"] + } + + +@pytest.mark.parametrize( + "original", + [ + {"type": "object", "properties": {"names": {"type": "array", "items": {"type": "string"}}}}, + {"type": "object", "additionalProperties": {"type": "string"}}, + {"$ref": "#/$defs/arguments", "$defs": {"arguments": {"type": "object"}}}, + ], +) +def test_explicit_json_schema_is_copied_without_reinterpretation(original): + pristine = copy.deepcopy(original) + normalized = normalize_input_schema(original) + assert normalized == pristine + assert normalized is not original + if "properties" in normalized: + normalized["properties"]["names"]["items"]["type"] = "integer" + 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 tool as claude_tool + from mcp.types import ToolAnnotations + + @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 From e346b80ce202192c94043d94bcd7133ad986c0cb Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:18:10 -0700 Subject: [PATCH 3/5] [None][fix] agent-flow: preserve inherited Codex MCP exclusions Merge inherited native and project MCP exclusions with session restrictions before starting a Codex thread. Abort session creation if native exclusions cannot be read, and add regression coverage for merging, isolation, and failure handling. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- agent-flow/agent_flow/backends/codex.py | 38 +++++ agent-flow/tests/test_codex_backend.py | 183 ++++++++++++++++++++++++ 2 files changed, 221 insertions(+) diff --git a/agent-flow/agent_flow/backends/codex.py b/agent-flow/agent_flow/backends/codex.py index b7b145dd62e5..14a7f1e44a64 100644 --- a/agent-flow/agent_flow/backends/codex.py +++ b/agent-flow/agent_flow/backends/codex.py @@ -184,6 +184,44 @@ async def create_client( 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, diff --git a/agent-flow/tests/test_codex_backend.py b/agent-flow/tests/test_codex_backend.py index 2ef52803eb6e..b39603756063 100644 --- a/agent-flow/tests/test_codex_backend.py +++ b/agent-flow/tests/test_codex_backend.py @@ -27,6 +27,7 @@ def __init__(self, config=None): self.registered = {} self.unregistered = [] self.thread_count = 0 + self.config_read_response = {"config": {}} async def start(self): self.started = True @@ -36,6 +37,10 @@ async def close(self): 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}"}} @@ -186,6 +191,180 @@ async def test_concurrent_sessions_keep_mcp_overrides_local_and_do_not_mutate_in 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): @@ -198,6 +377,10 @@ async def test_session_exception_unregisters_tools_and_unsubscribes(backend, tmp "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", From 699e2aa0b0310ea45662e865669b4d2c9d16e881 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:37:32 -0700 Subject: [PATCH 4/5] [None][fix] agent-flow: normalize explicit tool schemas to the MCP object shape Explicit JSON Schemas were copied verbatim into both backends. The pinned claude-agent-sdk only passes a dict through when it has a string "type" and a "properties" key; anything else is re-read as Python shorthand, so schemas such as {"type": "object", "additionalProperties": ...} or {"$ref": ..., "$defs": ...} advertised their own keywords as required parameters and every call failed input validation. normalize_input_schema now gives explicit schemas the root object shape MCP tools require: a string "type", a "properties" key, and any root "$ref" or list-valued "type" moved into "allOf" so the meaning is unchanged. Both backends advertise the same schema. Tests cover the normalized shapes and run a generated schema through the SDK's own schema builder and validator, the seam that hid this. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- agent-flow/README.md | 9 +++-- agent-flow/agent_flow/tools.py | 25 ++++++++++++- agent-flow/tests/test_backends.py | 52 ++++++++++++++++++++++++++ agent-flow/tests/test_tools.py | 62 ++++++++++++++++++++++++++----- 4 files changed, 133 insertions(+), 15 deletions(-) diff --git a/agent-flow/README.md b/agent-flow/README.md index 166df5665ce1..df718e195c34 100644 --- a/agent-flow/README.md +++ b/agent-flow/README.md @@ -70,9 +70,12 @@ python examples/quick_start.py 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. Resource links and embedded -text resources become readable tool output; image results retain their MIME -type. Unsupported result types return explicit tool errors. +`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: diff --git a/agent-flow/agent_flow/tools.py b/agent-flow/agent_flow/tools.py index 1bd1721b34ab..e0c7a48856ee 100644 --- a/agent-flow/agent_flow/tools.py +++ b/agent-flow/agent_flow/tools.py @@ -114,8 +114,29 @@ def _field_qualifier(hint: Any) -> Any: 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 or convert a shorthand dictionary/TypedDict to JSON Schema.""" + """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") @@ -134,7 +155,7 @@ def normalize_input_schema(schema: Any) -> dict[str, Any]: ) ) if is_json_schema: - normalized = copy.deepcopy(dict(schema)) + normalized = _mcp_object_schema(copy.deepcopy(dict(schema))) else: properties = {key: _json_schema(hint, set()) for key, hint in schema.items()} normalized = { diff --git a/agent-flow/tests/test_backends.py b/agent-flow/tests/test_backends.py index 96417fe150b0..fa237f557d01 100644 --- a/agent-flow/tests/test_backends.py +++ b/agent-flow/tests/test_backends.py @@ -6,6 +6,8 @@ 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 ( @@ -722,6 +724,23 @@ 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 @@ -745,6 +764,39 @@ async def image(args): {"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. diff --git a/agent-flow/tests/test_tools.py b/agent-flow/tests/test_tools.py index 32f820a39bfe..0f060ce53ce8 100644 --- a/agent-flow/tests/test_tools.py +++ b/agent-flow/tests/test_tools.py @@ -109,21 +109,63 @@ class Mode(Enum): } +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", + "original, expected", [ - {"type": "object", "properties": {"names": {"type": "array", "items": {"type": "string"}}}}, - {"type": "object", "additionalProperties": {"type": "string"}}, - {"$ref": "#/$defs/arguments", "$defs": {"arguments": {"type": "object"}}}, + ( + {"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_copied_without_reinterpretation(original): +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) - normalized = normalize_input_schema(original) - assert normalized == pristine - assert normalized is not original - if "properties" in normalized: - normalized["properties"]["names"]["items"]["type"] = "integer" + assert normalize_input_schema(original) == expected assert original == pristine From 85b6c81223a92d0f953279882804adf74c809e0f Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:09:30 -0700 Subject: [PATCH 5/5] [None][fix] agent-flow: import ToolAnnotations from claude_agent_sdk in tool test test_legacy_sdk_tool_conversion_preserves_annotations_and_handler built its annotations from mcp.types.ToolAnnotations, but maxResultSizeChars is a Claude Code field rather than an MCP hint, and only claude_agent_sdk.ToolAnnotations declares it. On mcp 1.x the import happened to work because that class set extra="allow"; mcp 2.x moved the wire types to the mcp-types distribution and dropped it, so the field was discarded at construction and normalize_tool() saw no annotation to fold into meta. Import the annotation type from claude_agent_sdk, which is also what backends/claude_code.py constructs, so the test exercises the type the backend actually uses and passes on either mcp major. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- agent-flow/tests/test_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent-flow/tests/test_tools.py b/agent-flow/tests/test_tools.py index 0f060ce53ce8..c3e3cef4f3a6 100644 --- a/agent-flow/tests/test_tools.py +++ b/agent-flow/tests/test_tools.py @@ -196,8 +196,8 @@ async def echo(args): 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 - from mcp.types import ToolAnnotations @claude_tool( "read",