feat(runtime/codex): execute agent MCP/function tools via the shim

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 ""

veadk/runtime/codex/runtime.py

@@ -45,8 +45,9 @@
 )
 
 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.skills import sync_skills_to_codex_home
+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
 
@@ -83,7 +84,8 @@ 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)
         # Expose the agent's skills to Codex by materializing them under
         # `$CODEX_HOME/skills/`, where Codex's native skill system discovers
@@ -93,6 +95,15 @@ async def run_async(
         except Exception as e:  # noqa: BLE001
             logger.warning(f"codex: skill sync skipped: {e}")
 
+        # 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.
@@ -143,6 +154,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)

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}")