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 ead2187d..16a994f3 100644 --- a/veadk/runtime/codex/runtime.py +++ b/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) 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}")