From 4f0eb0c7cbe129b0008809dd901b30847ca58f7d Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Tue, 7 Jul 2026 10:22:23 +0800 Subject: [PATCH] feat(runtime/codex): execute agent MCP/function tools via the shim Codex presents MCP/function tools to the model as a namespace tool that Ark rejects, and it won't route a plain function_call back to a namespaced tool. So instead of configuring tools on the Codex side, the Responses shim now advertises the agent's ADK tools to the backend as plain function tools and executes them itself in a bounded tool loop, invisibly to Codex. - New tools_bridge builds (flat function specs, executors) from the agent's tools, reusing ADK end to end: BaseToolset.get_tools (incl. MCP), BaseTool._get_declaration + _function_declaration_to_tool_param for the schema, and BaseTool.run_async for execution. - proxy runs the agent's tools in a shim-internal loop (registered per turn via set_agent_tools); the tool-less path is unchanged. Codex's own non-function tools (e.g. its hosted web_search, whose schema Ark rejects) are dropped, so Codex web search is disabled by default on a chat backend. - runtime bridges the agent's tools onto the shim and closes MCP sessions when the turn ends. Note: this also removes the opt-in web-search restoration added in fix/codex-runtime-longtask (the CODEX_SHIM_MAX_TOOL_ITERS shim web loop), keeping Codex web search off by default. The retry/timeout knobs from that change are retained. Validated end to end: an Agent(runtime=codex) with an MCPToolset pointing at a live MCP server calls the tool and answers from its result, on the bundled Codex binary (no namespace, no config flag, no binary bump). --- veadk/runtime/codex/proxy.py | 146 +++++++++------------------- veadk/runtime/codex/runtime.py | 19 +++- veadk/runtime/codex/tools_bridge.py | 134 +++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 104 deletions(-) create mode 100644 veadk/runtime/codex/tools_bridge.py diff --git a/veadk/runtime/codex/proxy.py b/veadk/runtime/codex/proxy.py index a84a6339..323031d5 100644 --- a/veadk/runtime/codex/proxy.py +++ b/veadk/runtime/codex/proxy.py @@ -89,83 +89,9 @@ def _shim_timeout() -> float: return 0.0 -def _shim_max_tool_iters() -> int: - """Max shim-internal web tool round-trips per request (``0`` = disabled). - - Codex only speaks the Responses API and its hosted ``web_search`` tool is - stripped before Ark (Ark rejects its schema), so a Codex+Ark agent has no - web capability. When this is > 0, the shim translates the hosted web tools - into Ark-accepted ``function`` tools, executes the veADK builtins itself, - and loops until the model stops calling them. Default ``0`` keeps the path - byte-identical to before. Env-tunable via ``CODEX_SHIM_MAX_TOOL_ITERS`` - (recommended 6 when enabling). - """ - try: - return max(0, int(os.getenv("CODEX_SHIM_MAX_TOOL_ITERS", "0"))) - except ValueError: - return 0 - - -# Hosted web tools (Codex emits these; Ark rejects their schema) translated to -# plain Responses ``function`` tools that the shim executes against veADK's own -# builtins. Kept minimal — the builtins own the real parameter handling. -_EXECUTABLE_TOOLS: list[dict[str, Any]] = [ - { - "type": "function", - "name": "web_search", - "description": "Search the web for a query and return result documents.", - "parameters": { - "type": "object", - "properties": { - "query": {"type": "string", "description": "The search query."} - }, - "required": ["query"], - }, - }, - { - "type": "function", - "name": "web_fetch", - "description": ( - "Fetch a public web page or PDF over HTTP(S) and return its " - "readable main content." - ), - "parameters": { - "type": "object", - "properties": { - "url": {"type": "string", "description": "The http(s) URL to fetch."}, - "extract_mode": { - "type": "string", - "enum": ["markdown", "text"], - "description": "Output format; defaults to markdown.", - }, - "max_chars": { - "type": "integer", - "description": "Truncate content to at most this many chars.", - }, - }, - "required": ["url"], - }, - }, -] - -_EXECUTABLE_TOOL_NAMES = {t["name"] for t in _EXECUTABLE_TOOLS} - - -async def _run_tool(name: str, args: dict[str, Any]) -> str: - """Execute a veADK built-in web tool and return a JSON string. - - The web builtins use blocking ``requests``, so they run in a worker thread - to avoid stalling the shim's event loop. Any error is returned as a JSON - payload (never raised) so the model can recover on the next turn. - """ - try: - from veadk.tools import get_builtin_tool - - fn = get_builtin_tool(name) - result = await asyncio.to_thread(fn, **args) - return json.dumps(result, ensure_ascii=False, default=str) - except Exception as e: # noqa: BLE001 - surfaced to the model, not raised - return json.dumps({"error": str(e)}) +# Cap on shim-internal tool round-trips per turn — bounds runaway loops while +# allowing several tool calls per turn. +_AGENT_TOOL_MAX_ITERS = 8 class ResponsesShim: @@ -187,8 +113,20 @@ def __init__(self, api_base: str, api_key: str) -> None: self.url: str | None = None self._server: uvicorn.Server | None = None self._task: asyncio.Task[Any] | None = None + # The current turn's agent tools, advertised to the backend as plain + # `function` tools and executed here (invisibly to Codex). Set per turn + # via set_agent_tools(); see veadk.runtime.codex.tools_bridge. + self._agent_specs: list[dict[str, Any]] = [] + self._agent_executors: dict[str, Any] = {} self._app = self._build_app() + def set_agent_tools( + self, specs: list[dict[str, Any]], executors: dict[str, Any] + ) -> None: + """Register (or clear) the agent tools the shim should inject + execute.""" + self._agent_specs = specs or [] + self._agent_executors = executors or {} + def _build_app(self) -> FastAPI: app = FastAPI() @@ -201,25 +139,19 @@ async def responses(request: Request) -> Any: call_kwargs: dict[str, Any] = { key: body[key] for key in _PASSTHROUGH_KEYS if key in body } - # Codex injects its built-in tools (e.g. `web_search`) whose schema - # carries fields like `external_web_access` that non-OpenAI Responses - # backends (Ark) reject. Keep only standard `function` tools, which - # the bridged chat backend understands. When the tool loop is enabled - # (CODEX_SHIM_MAX_TOOL_ITERS > 0), additionally translate the stripped - # hosted web tools into Ark-accepted `function` tools that the shim - # executes itself (see the tool loop below), so a Codex+Ark agent - # regains web search/fetch. + # Advertise the agent's ADK tools to the backend as plain `function` + # tools; the shim executes them itself (see the tool loop below). + # Codex's own non-`function` tools are dropped, since Ark rejects + # their schema (e.g. the hosted `web_search`'s `external_web_access`); + # this leaves Codex's web search disabled on a chat backend. + agent_executors: dict[str, Any] = self._agent_executors if isinstance(call_kwargs.get("tools"), list): - inbound = call_kwargs["tools"] - wants_web = any( - t.get("type") in ("web_search", "web_search_preview") - for t in inbound - ) - kept = [t for t in inbound if t.get("type") == "function"] - if wants_web and _shim_max_tool_iters() > 0: - have = {t.get("name") for t in kept} - kept.extend(t for t in _EXECUTABLE_TOOLS if t["name"] not in have) + kept = [t for t in call_kwargs["tools"] if t.get("type") == "function"] + have = {t.get("name") for t in kept} + kept.extend(t for t in self._agent_specs if t.get("name") not in have) call_kwargs["tools"] = kept + elif self._agent_specs: + call_kwargs["tools"] = list(self._agent_specs) # On multi-step turns Codex replays prior assistant messages in # `input` without a `status` field, but Ark's Responses API # requires `status` on assistant messages (MissingParameter: @@ -261,7 +193,10 @@ async def responses(request: Request) -> Any: # of executable function_calls in the fresh output — the always-empty # `message` item is ignored. The loop is invisible to Codex: only the # final, tool-free turn is returned/synthesized. - max_iters = _shim_max_tool_iters() + # The shim resolves the agent's tools itself; with none registered + # the loop is disabled and the path is unchanged for tool-less runs. + exec_names = set(agent_executors) + max_iters = _AGENT_TOOL_MAX_ITERS if agent_executors else 0 iters = 0 while True: result = await litellm.aresponses(**call_kwargs) @@ -275,7 +210,7 @@ async def responses(request: Request) -> Any: it for it in (resp.get("output") or []) if it.get("type") == "function_call" - and it.get("name") in _EXECUTABLE_TOOL_NAMES + and it.get("name") in exec_names ] if not calls: break @@ -285,7 +220,7 @@ async def responses(request: Request) -> Any: args = json.loads(fc.get("arguments") or "{}") except json.JSONDecodeError: args = {} - out = await _run_tool(fc["name"], args) + out = await agent_executors[fc["name"]](args) conv.append( { "type": "function_call", @@ -308,13 +243,13 @@ async def responses(request: Request) -> Any: # If we broke at the iteration cap with an executable function_call # still pending, strip it so it never leaks to Codex (which cannot # run it and would desync the next turn / emit a null delta). - if max_iters > 0 and isinstance(resp.get("output"), list): + if exec_names and isinstance(resp.get("output"), list): resp["output"] = [ it for it in resp["output"] if not ( it.get("type") == "function_call" - and it.get("name") in _EXECUTABLE_TOOL_NAMES + and it.get("name") in exec_names ) ] @@ -521,11 +456,18 @@ def ev(payload: dict[str, Any]) -> bytes: _SHIMS: dict[tuple[str, str], ResponsesShim] = {} -async def get_shim_url(api_base: str, api_key: str) -> str: - """Return a started shim URL for the given backend, creating it if needed.""" +async def get_shim(api_base: str, api_key: str) -> ResponsesShim: + """Return a started shim for the given backend, creating it if needed.""" key = (api_base, api_key) shim = _SHIMS.get(key) if shim is None: shim = ResponsesShim(api_base=api_base, api_key=api_key) _SHIMS[key] = shim - return await shim.start() + await shim.start() + return shim + + +async def get_shim_url(api_base: str, api_key: str) -> str: + """Return a started shim URL for the given backend, creating it if needed.""" + shim = await get_shim(api_base, api_key) + return shim.url or "" diff --git a/veadk/runtime/codex/runtime.py b/veadk/runtime/codex/runtime.py index 7585cdd2..bcc0ecba 100644 --- a/veadk/runtime/codex/runtime.py +++ b/veadk/runtime/codex/runtime.py @@ -45,7 +45,8 @@ ) from veadk.runtime.base_runtime import BaseRuntime, build_system_append -from veadk.runtime.codex.proxy import get_shim_url +from veadk.runtime.codex.proxy import get_shim +from veadk.runtime.codex.tools_bridge import build_executable_tools, close_toolsets from veadk.runtime.codex.translate import build_prompt, item_to_events from veadk.utils.logger import get_logger @@ -82,9 +83,19 @@ async def run_async( "(the chat endpoint Codex is bridged onto)." ) - shim_url = await get_shim_url(api_base, api_key) + shim = await get_shim(api_base, api_key) + shim_url = shim.url or "" codex_home = _prepare_codex_home(shim_url, model) + # Bridge the agent's ADK tools (function/MCP) to the shim: it advertises + # them to the backend as plain `function` tools and executes them itself, + # so they never reach Codex (whose `namespace` MCP form Ark rejects). See + # veadk.runtime.codex.tools_bridge / proxy. + tool_specs, tool_executors, opened_toolsets = await build_executable_tools( + agent, ctx + ) + shim.set_agent_tools(tool_specs, tool_executors) + # Codex has no clean SDK channel to append to its base system prompt, so # the agent identity/instruction is folded into a leading block of the # input (a labelled preamble), not the transcript itself. @@ -135,6 +146,10 @@ async def run_async( if aclose is not None: await aclose() finally: + # Drop this turn's tools from the (shared) shim and release any MCP + # sessions opened for them. + shim.set_agent_tools([], {}) + await close_toolsets(opened_toolsets) for key, value in previous.items(): if value is None: os.environ.pop(key, None) diff --git a/veadk/runtime/codex/tools_bridge.py b/veadk/runtime/codex/tools_bridge.py new file mode 100644 index 00000000..4cbdeaf6 --- /dev/null +++ b/veadk/runtime/codex/tools_bridge.py @@ -0,0 +1,134 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bridge an agent's ADK tools to Codex's Responses shim. + +Codex presents MCP/function tools to the model as a ``type:"namespace"`` object +that a chat backend (Ark) rejects, and it will not route a plain ``function_call`` +back to a namespaced tool. So instead of configuring the tools on the Codex side, +the shim advertises them to the backend as plain ``function`` tools and executes +them itself (see :mod:`veadk.runtime.codex.proxy`), invisibly to Codex. + +This module produces, for an agent's tools, the two things the shim needs: + +- ``specs``: flat Responses ``function`` tool specs to advertise to the backend. +- ``executors``: ``{name: async (args) -> str}`` to run a tool when the model + calls it. + +Both are derived from ADK itself — tool declarations via ``BaseTool._get_declaration`` +(+ ADK's own ``_function_declaration_to_tool_param`` schema conversion) and +execution via ``BaseTool.run_async`` — so MCP tools, function tools and any other +``BaseTool`` are handled uniformly without reimplementing schema or dispatch. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, Awaitable, Callable + +from veadk.utils.logger import get_logger + +if TYPE_CHECKING: + from google.adk.agents.invocation_context import InvocationContext + from google.adk.tools.base_tool import BaseTool + from google.adk.tools.base_toolset import BaseToolset + + from veadk.agent import Agent + +logger = get_logger(__name__) + +Executor = Callable[[dict[str, Any]], Awaitable[str]] + + +async def build_executable_tools( + agent: "Agent", ctx: "InvocationContext" +) -> tuple[list[dict[str, Any]], dict[str, Executor], list["BaseToolset"]]: + """Collect the agent's ADK tools as shim-executable functions. + + Returns ``(specs, executors, opened_toolsets)`` where ``opened_toolsets`` are + the toolsets whose sessions were opened by ``get_tools`` (e.g. MCP); the + caller must ``close()`` them when the turn ends. + """ + from google.adk.models.lite_llm import _function_declaration_to_tool_param + from google.adk.tools.base_tool import BaseTool + from google.adk.tools.base_toolset import BaseToolset + from google.adk.tools.tool_context import ToolContext + + specs: list[dict[str, Any]] = [] + executors: dict[str, Executor] = {} + opened: list[BaseToolset] = [] + + def _add(tool: "BaseTool") -> None: + try: + declaration = tool._get_declaration() + except Exception as e: # noqa: BLE001 - one tool must not break the turn + logger.warning(f"codex: skipping tool with no declaration: {e}") + return + if declaration is None or not declaration.name: + return + name = declaration.name + # ADK builds the OpenAI (chat) tool param, incl. genai->JSON schema + # normalization; lift its `function` body up to the Responses flat shape. + chat_param = _function_declaration_to_tool_param(declaration) + specs.append({"type": "function", **chat_param["function"]}) + executors[name] = _make_executor(tool, ctx, ToolContext) + + for entry in getattr(agent, "tools", None) or []: + if isinstance(entry, BaseToolset): + try: + tools = await entry.get_tools() + except Exception as e: # noqa: BLE001 + logger.warning(f"codex: failed to list toolset {entry!r}: {e}") + continue + opened.append(entry) + for tool in tools: + _add(tool) + elif isinstance(entry, BaseTool): + _add(entry) + + if executors: + logger.info( + f"codex: bridging {len(executors)} agent tool(s): {list(executors)}" + ) + return specs, executors, opened + + +def _make_executor( + tool: "BaseTool", ctx: "InvocationContext", tool_context_cls: Any +) -> Executor: + async def _run(args: dict[str, Any]) -> str: + try: + result = await tool.run_async(args=args, tool_context=tool_context_cls(ctx)) + except Exception as e: # noqa: BLE001 - surfaced to the model, not raised + return json.dumps({"error": str(e)}) + if isinstance(result, str): + return result + try: + return json.dumps(result, ensure_ascii=False, default=str) + except Exception: # noqa: BLE001 + return str(result) + + return _run + + +async def close_toolsets(toolsets: list["BaseToolset"]) -> None: + """Best-effort close of toolset sessions opened during the turn.""" + for toolset in toolsets: + close = getattr(toolset, "close", None) + if close is None: + continue + try: + await close() + except Exception as e: # noqa: BLE001 + logger.warning(f"codex: failed to close toolset {toolset!r}: {e}")