From b7ebe6456fd62c45d2482b550908abed0ece7255 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Mon, 31 Aug 2026 08:28:58 -0700 Subject: [PATCH] fix(xai): request-local alias provenance + collision-safe wire aliasing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hardens the two #95003 alias carriers per review feedback on #95019/#95011: - _alias_reserved_tools / _rename_tool_search_bridge_for_xai now return the alias map THIS request emitted; the transport stashes it (_last_wire_aliases) and normalize_response reverses ONLY those aliases. A real user/plugin/MCP tool named hermes_tool_search is never silently dispatched as tool_search when no alias was sent. - Collision safety: if a real tool already occupies the alias name, the bridge takes hermes_tool_search_2/_3 — no duplicate wire declarations. - Legacy static reverse map retained only for normalize-only call sites that never built a request on the transport instance. - chat_completion_helpers resets provenance per request so stale maps from a prior request can't leak into the next response's dispatch. Refs #95003 --- agent/chat_completion_helpers.py | 14 ++- agent/transports/chat_completions.py | 55 ++++++++--- agent/transports/codex.py | 92 +++++++++++++++---- ..._chat_completions_xai_tool_search_alias.py | 55 ++++++++++- .../agent/transports/test_codex_transport.py | 87 ++++++++++++++++++ 5 files changed, 270 insertions(+), 33 deletions(-) diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index d0a33d3429..258e056612 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -1971,6 +1971,12 @@ def build_api_kwargs(agent, api_messages: list, tools_for_api: list | None = Non agent.provider in {"xai", "xai-oauth"} or agent._base_url_hostname == "api.x.ai" ) + # Reset request-local alias provenance for THIS request; the rewrite + # below repopulates it when it actually emits aliases. Without the + # reset, a stale map from an earlier request on the same transport + # could reverse-map a name this request never aliased. + if _ct is not None and hasattr(_ct, "_last_wire_aliases"): + _ct._last_wire_aliases = {} if _is_xai_chat and tools_for_api: try: import copy as _copy_xai @@ -1986,7 +1992,13 @@ def build_api_kwargs(agent, api_messages: list, tools_for_api: list | None = Non ) if _has_bridge: tools_for_api = _copy_xai.deepcopy(tools_for_api) - tools_for_api = _rename_tool_search_bridge_for_xai(tools_for_api) + tools_for_api, _xai_alias_map = _rename_tool_search_bridge_for_xai( + tools_for_api + ) + # Record provenance so normalize_response reverses ONLY the + # aliases this request put on the wire. + if _ct is not None: + _ct._last_wire_aliases = _xai_alias_map except Exception as exc: logger.warning( "%s⚠️ Failed to alias tool_search bridge for xAI: %s", diff --git a/agent/transports/chat_completions.py b/agent/transports/chat_completions.py index 4e557bd427..e2a6707b09 100644 --- a/agent/transports/chat_completions.py +++ b/agent/transports/chat_completions.py @@ -43,26 +43,44 @@ _XAI_TOOL_SEARCH_ALIAS = "hermes_tool_search" def _rename_tool_search_bridge_for_xai( tools: list[dict[str, Any]], -) -> list[dict[str, Any]]: +) -> tuple[list[dict[str, Any]], dict[str, str]]: """Rename the client ``tool_search`` bridge declaration to a wire alias. Only the wire name changes: descriptions, schemas, and the other two bridge names (``tool_describe`` / ``tool_call`` — not reserved by xAI) - pass through untouched. The alias is mapped back to ``tool_search`` in - ``normalize_response`` before dispatch. + pass through untouched. Returns ``(rewritten_tools, alias_map)`` where + ``alias_map`` maps each alias THIS request emits back to the original + name; the caller stashes it on the transport so ``normalize_response`` + only reverses aliases that were actually sent. If a real tool already + occupies ``hermes_tool_search``, the bridge takes a ``_2``/``_3`` + suffix instead of duplicating a wire name. """ rewritten: list[dict[str, Any]] = [] + alias_map: dict[str, str] = {} + taken = { + (tool.get("function") or {}).get("name") + for tool in tools + if isinstance(tool, dict) + } + taken.discard(None) for tool in tools: if ( isinstance(tool, dict) and (tool.get("function") or {}).get("name") == "tool_search" ): + alias = _XAI_TOOL_SEARCH_ALIAS + suffix = 2 + while alias in taken: + alias = f"{_XAI_TOOL_SEARCH_ALIAS}_{suffix}" + suffix += 1 + taken.add(alias) + alias_map[alias] = "tool_search" aliased = dict(tool) - aliased["function"] = {**tool["function"], "name": _XAI_TOOL_SEARCH_ALIAS} + aliased["function"] = {**tool["function"], "name": alias} rewritten.append(aliased) else: rewritten.append(tool) - return rewritten + return rewritten, alias_map def _static_prompt_instructions(messages: list[dict[str, Any]]) -> str: @@ -314,6 +332,13 @@ class ChatCompletionsTransport(ProviderTransport): The default path for OpenAI-compatible providers. """ + # Wire-alias provenance of the most recent request built for this + # transport: ``{alias_sent_on_wire: original_tool_name}``. ``None`` + # means no request recorded provenance (normalize-only call sites) — + # fall back to the static alias constant. An empty dict means the last + # request emitted no aliases, so no reverse rewrite may run (#95003). + _last_wire_aliases: dict[str, str] | None = None + @property def api_mode(self) -> str: return "chat_completions" @@ -962,13 +987,19 @@ class ChatCompletionsTransport(ProviderTransport): # preserve an explicit blank name for Hermes's recovery path. if tc_function is None or function_name is None: continue - # Map the xAI wire alias back to the bridge's real name. - # Unconditional is correct here: ``hermes_tool_search`` only - # exists on the wire because _rename_tool_search_bridge_for_xai - # put it there (xAI rejects the literal ``tool_search``), so - # any model call carrying the alias is a bridge invocation. - if function_name == _XAI_TOOL_SEARCH_ALIAS: - function_name = "tool_search" + # Map THIS request's wire aliases back before dispatch. + # Request-local provenance: when the paired request recorded + # its alias map, only those aliases are reversed — a real + # user/plugin/MCP tool that happens to be named + # ``hermes_tool_search`` dispatches as itself when no alias + # was emitted. The static-constant fallback covers + # normalize-only call sites with no recorded request. + _alias_map = self._last_wire_aliases + if _alias_map is None: + if function_name == _XAI_TOOL_SEARCH_ALIAS: + function_name = "tool_search" + elif function_name in _alias_map: + function_name = _alias_map[function_name] function_arguments = getattr(tc_function, "arguments", None) # Preserve provider-specific extras on the tool call. # Gemini 3 thinking models attach extra_content with diff --git a/agent/transports/codex.py b/agent/transports/codex.py index d7119af27b..ff8979edce 100644 --- a/agent/transports/codex.py +++ b/agent/transports/codex.py @@ -90,6 +90,17 @@ _RESERVED_ALIAS_TO_NAME = { for name in (*_OPENCODE_RESERVED_TOOL_NAMES, *_XAI_RESERVED_TOOL_NAMES) } +# Legacy reverse map used ONLY when normalize_response runs on a transport +# instance that never built a request (normalize-only call sites / tests). +# Production requests carry request-local provenance instead — see +# ``_last_wire_aliases`` — so a real user/plugin/MCP tool that happens to be +# named ``hermes_tool_search`` is never silently rewritten to ``tool_search`` +# unless THIS request actually emitted that alias (#95003 review contract). +_LEGACY_ALIAS_FALLBACK = { + **_RESERVED_ALIAS_TO_NAME, + "hermes_web_search": "web_search", +} + def _is_opencode_responses_backend(params: Dict[str, Any]) -> bool: """True when this Responses request targets an OpenCode endpoint. @@ -117,22 +128,41 @@ def _is_opencode_responses_backend(params: Dict[str, Any]) -> bool: def _alias_reserved_tools( response_tools: List[Dict[str, Any]], reserved_names: Tuple[str, ...], -) -> List[Dict[str, Any]]: +) -> Tuple[List[Dict[str, Any]], Dict[str, str]]: """Alias provider-reserved client function names on the wire. - Single owner for every reserved-name collision on this transport; the - reverse mapping lives in :data:`_RESERVED_ALIAS_TO_NAME`, applied in - ``normalize_response`` so Hermes dispatch never sees the alias. + Single owner for every reserved-name collision on this transport. + Returns ``(rewritten_tools, alias_map)`` where ``alias_map`` maps each + wire alias emitted by THIS request back to the original tool name. + The caller stashes the map for ``normalize_response`` so the reverse + rewrite only ever applies to aliases this request actually sent — + a legitimate user/plugin/MCP tool already named ``hermes_`` is + neither shadowed (the alias picks a ``_2``/``_3`` suffix instead of + duplicating a wire name) nor mis-dispatched on the response path. """ rewritten: List[Dict[str, Any]] = [] + alias_map: Dict[str, str] = {} + taken = { + tool.get("name") + for tool in response_tools + if isinstance(tool, dict) and tool.get("name") + } for tool in response_tools: if isinstance(tool, dict) and tool.get("name") in reserved_names: + base = f"{_RESERVED_TOOL_ALIAS_PREFIX}{tool['name']}" + alias = base + suffix = 2 + while alias in taken: + alias = f"{base}_{suffix}" + suffix += 1 + taken.add(alias) + alias_map[alias] = tool["name"] aliased = dict(tool) - aliased["name"] = f"{_RESERVED_TOOL_ALIAS_PREFIX}{tool['name']}" + aliased["name"] = alias rewritten.append(aliased) else: rewritten.append(tool) - return rewritten + return rewritten, alias_map def _xai_prefers_native_web_search() -> bool: @@ -429,6 +459,14 @@ class ResponsesApiTransport(ProviderTransport): # attribute default; mutated on the instance, not the class. _last_issuer_kind: Optional[str] = None + # Wire-alias provenance of the most recent build_kwargs call: + # ``{alias_sent_on_wire: original_tool_name}``. ``None`` means "no + # request built on this instance" (normalize-only call sites), in which + # case normalize_response falls back to the static legacy map. An empty + # dict means the last request emitted no aliases, so no reverse rewrite + # is permitted (#95003 provenance contract). + _last_wire_aliases: Optional[Dict[str, str]] = None + @property def api_mode(self) -> str: return "codex_responses" @@ -628,6 +666,12 @@ class ResponsesApiTransport(ProviderTransport): # is honored, but rename the wire tool to # ``hermes_web_search`` so Grok cannot hijack the name. The alias # is mapped back to ``web_search`` in ``normalize_response``. + # Request-local alias provenance: every wire alias THIS request + # emits is recorded here and stashed on the transport, so the + # reverse rewrite in ``normalize_response`` applies only to aliases + # that were actually sent (never to a real tool that merely shares + # an alias-shaped name). + wire_aliases: Dict[str, str] = {} if is_xai_responses and response_tools: has_client_web_search = any( isinstance(t, dict) and t.get("name") == "web_search" @@ -643,22 +687,30 @@ class ResponsesApiTransport(ProviderTransport): response_tools = filtered else: response_tools = _rename_client_web_search_for_xai(response_tools) + wire_aliases[_XAI_CLIENT_WEB_SEARCH_ALIAS] = "web_search" # OpenCode Responses backends reserve web_search / search_files as # function names (HTTP 400 "custom function name 'X' is reserved", # #85589). Alias them on the wire; normalize_response maps them back. if response_tools and _is_opencode_responses_backend(params): - response_tools = _alias_reserved_tools( + response_tools, _oc_aliases = _alias_reserved_tools( response_tools, _OPENCODE_RESERVED_TOOL_NAMES ) + wire_aliases.update(_oc_aliases) # xAI reserves ``tool_search`` for its native server-side tool and # rejects the client declaration outright (#95003). Alias it on the # wire; normalize_response maps it back before dispatch. if is_xai_responses and response_tools: - response_tools = _alias_reserved_tools( + response_tools, _xai_aliases = _alias_reserved_tools( response_tools, _XAI_RESERVED_TOOL_NAMES ) + wire_aliases.update(_xai_aliases) + + # Stash for normalize_response (same request/response pairing model + # as ``_last_issuer_kind``). An empty dict is meaningful: it means + # this request emitted NO aliases, so no reverse rewrite may run. + self._last_wire_aliases = wire_aliases # ``tools`` MUST be omitted entirely when there are no functions to # expose: the openai SDK's ``responses.stream()`` / ``responses.parse()`` @@ -896,14 +948,22 @@ class ResponsesApiTransport(ProviderTransport): if hasattr(tc, "response_item_id") and tc.response_item_id: provider_data["response_item_id"] = tc.response_item_id name = tc.function.name if hasattr(tc, "function") else getattr(tc, "name", "") - # Undo the xAI client-path wire alias so Hermes dispatches - # the real ``web_search`` tool (Firecrawl / etc.). - if name == _XAI_CLIENT_WEB_SEARCH_ALIAS: - name = "web_search" - # Undo the OpenCode reserved-name wire aliases the same way - # (hermes_web_search / hermes_search_files, #85589). - elif name in _RESERVED_ALIAS_TO_NAME: - name = _RESERVED_ALIAS_TO_NAME[name] + # Undo THIS request's wire aliases before Hermes dispatch. + # Request-local provenance: only aliases the paired + # build_kwargs call actually emitted are rewritten, so a + # legitimate tool that happens to be named + # ``hermes_tool_search`` etc. is dispatched as itself when + # no alias was sent. The static legacy map is used only for + # normalize-only call sites that never built a request on + # this transport instance. + alias_map = self._last_wire_aliases + if alias_map is None: + if name == _XAI_CLIENT_WEB_SEARCH_ALIAS: + name = "web_search" + elif name in _LEGACY_ALIAS_FALLBACK: + name = _LEGACY_ALIAS_FALLBACK[name] + elif name in alias_map: + name = alias_map[name] tool_calls.append(ToolCall( id=tc.id if hasattr(tc, "id") else (name or None), name=name, diff --git a/tests/agent/transports/test_chat_completions_xai_tool_search_alias.py b/tests/agent/transports/test_chat_completions_xai_tool_search_alias.py index b0764e1313..55bfe8ec79 100644 --- a/tests/agent/transports/test_chat_completions_xai_tool_search_alias.py +++ b/tests/agent/transports/test_chat_completions_xai_tool_search_alias.py @@ -7,6 +7,11 @@ reserved for the tool_search tool"). The fix mirrors the web_search treatment in ``transports/codex.py``: rename the bridge's wire declaration to ``hermes_tool_search`` for xAI targets and map the alias back to ``tool_search`` in ``normalize_response`` so dispatch is unchanged. + +The reverse map is request-local: ``normalize_response`` only rewrites +aliases the paired request actually emitted (stashed on the transport as +``_last_wire_aliases``), so a real user/plugin/MCP tool that happens to be +named ``hermes_tool_search`` is never silently dispatched as the bridge. """ from types import SimpleNamespace @@ -36,9 +41,10 @@ class TestRenameToolSearchBridgeForXai: "parameters": {"type": "object", "properties": {}}, }, }] - out = _rename_tool_search_bridge_for_xai(tools) + out, alias_map = _rename_tool_search_bridge_for_xai(tools) assert out[0]["function"]["name"] == _XAI_TOOL_SEARCH_ALIAS assert out[0]["function"]["name"] == "hermes_tool_search" + assert alias_map == {"hermes_tool_search": "tool_search"} def test_schema_and_description_untouched(self): fn = { @@ -46,7 +52,7 @@ class TestRenameToolSearchBridgeForXai: "description": "Search the deferred tool catalog", "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, } - out = _rename_tool_search_bridge_for_xai([{"type": "function", "function": fn}]) + out, _ = _rename_tool_search_bridge_for_xai([{"type": "function", "function": fn}]) assert out[0]["function"]["description"] == fn["description"] assert out[0]["function"]["parameters"] == fn["parameters"] @@ -57,13 +63,15 @@ class TestRenameToolSearchBridgeForXai: {"type": "function", "function": {"name": "tool_describe"}}, {"type": "function", "function": {"name": "tool_call"}}, ] - out = _rename_tool_search_bridge_for_xai(tools) + out, alias_map = _rename_tool_search_bridge_for_xai(tools) assert [t["function"]["name"] for t in out] == ["tool_describe", "tool_call"] + assert alias_map == {} def test_ordinary_tools_untouched(self): tools = [{"type": "function", "function": {"name": "web_search"}}] - out = _rename_tool_search_bridge_for_xai(tools) + out, alias_map = _rename_tool_search_bridge_for_xai(tools) assert out[0]["function"]["name"] == "web_search" + assert alias_map == {} def test_input_not_mutated(self): # The helper feeds a deep-copied list on the helper-layer path, but @@ -73,6 +81,20 @@ class TestRenameToolSearchBridgeForXai: _rename_tool_search_bridge_for_xai(tools) assert tools[0]["function"]["name"] == "tool_search" + def test_collision_with_real_hermes_tool_search_takes_suffix(self): + # A legitimate tool already using the alias name must NOT be + # shadowed and no duplicate wire names may be produced: the bridge + # takes hermes_tool_search_2 instead. + tools = [ + {"type": "function", "function": {"name": "hermes_tool_search"}}, + {"type": "function", "function": {"name": "tool_search"}}, + ] + out, alias_map = _rename_tool_search_bridge_for_xai(tools) + names = [t["function"]["name"] for t in out] + assert names == ["hermes_tool_search", "hermes_tool_search_2"] + assert len(names) == len(set(names)) + assert alias_map == {"hermes_tool_search_2": "tool_search"} + def _fake_response(tool_name): tc = SimpleNamespace( @@ -86,9 +108,34 @@ def _fake_response(tool_name): class TestNormalizeResponseMapsAliasBack: def test_alias_call_maps_back_to_bridge_name(self, transport): + transport._last_wire_aliases = {"hermes_tool_search": "tool_search"} resp = transport.normalize_response(_fake_response(_XAI_TOOL_SEARCH_ALIAS)) assert resp.tool_calls[0].name == "tool_search" def test_ordinary_call_name_preserved(self, transport): + transport._last_wire_aliases = {"hermes_tool_search": "tool_search"} resp = transport.normalize_response(_fake_response("tool_describe")) assert resp.tool_calls[0].name == "tool_describe" + + def test_no_alias_emitted_means_no_reverse_rewrite(self, transport): + # Provenance contract: if THIS request emitted no aliases, a tool + # call named hermes_tool_search is a REAL tool (user/plugin/MCP) + # and must dispatch under its own name. + transport._last_wire_aliases = {} + resp = transport.normalize_response(_fake_response("hermes_tool_search")) + assert resp.tool_calls[0].name == "hermes_tool_search" + + def test_suffixed_alias_maps_back(self, transport): + transport._last_wire_aliases = {"hermes_tool_search_2": "tool_search"} + resp = transport.normalize_response(_fake_response("hermes_tool_search_2")) + assert resp.tool_calls[0].name == "tool_search" + # And the real tool occupying the plain alias name is untouched. + resp2 = transport.normalize_response(_fake_response("hermes_tool_search")) + assert resp2.tool_calls[0].name == "hermes_tool_search" + + def test_legacy_fallback_without_provenance(self, transport): + # Normalize-only call sites (no request built on this instance) + # keep the historical unconditional mapping. + transport._last_wire_aliases = None + resp = transport.normalize_response(_fake_response(_XAI_TOOL_SEARCH_ALIAS)) + assert resp.tool_calls[0].name == "tool_search" diff --git a/tests/agent/transports/test_codex_transport.py b/tests/agent/transports/test_codex_transport.py index 4f0bdbf69b..463fea8bb4 100644 --- a/tests/agent/transports/test_codex_transport.py +++ b/tests/agent/transports/test_codex_transport.py @@ -1036,9 +1036,96 @@ class TestXaiReservedToolSearchAlias: "agent.codex_responses_adapter._normalize_codex_response", lambda resp, issuer_kind=None: (msg, "tool_calls"), ) + # Pair the response with a real request so provenance is recorded. + transport.build_kwargs( + model="grok-4.6", + messages=[{"role": "user", "content": "hi"}], + tools=list(self._TOOLS), + is_xai_responses=True, + ) + assert transport._last_wire_aliases == {"hermes_tool_search": "tool_search"} normalized = transport.normalize_response(response) assert [tc.name for tc in normalized.tool_calls] == ["tool_search"] + def _normalize_named_call(self, transport, monkeypatch, wire_name): + msg = SimpleNamespace( + content=None, + reasoning=None, + tool_calls=[ + SimpleNamespace( + id="call_1", call_id="call_1", response_item_id="fc_1", + function=SimpleNamespace(name=wire_name, arguments="{}"), + ), + ], + codex_reasoning_items=None, + codex_message_items=None, + reasoning_details=None, + ) + response = SimpleNamespace(output=[], status="completed") + monkeypatch.setattr( + "agent.codex_responses_adapter._normalize_codex_response", + lambda resp, issuer_kind=None: (msg, "tool_calls"), + ) + return transport.normalize_response(response) + + def test_no_alias_emitted_means_no_reverse_rewrite(self, transport, monkeypatch): + """Provenance contract (#95003 review): a request that emitted no + aliases must not have a real ``hermes_tool_search`` tool rewritten.""" + real_tool = {"type": "function", "function": { + "name": "hermes_tool_search", "description": "A real MCP tool.", + "parameters": {"type": "object", "properties": {}}}} + transport.build_kwargs( + model="grok-4.6", + messages=[{"role": "user", "content": "hi"}], + tools=[real_tool], + is_xai_responses=True, + ) + assert transport._last_wire_aliases == {} + normalized = self._normalize_named_call( + transport, monkeypatch, "hermes_tool_search" + ) + assert [tc.name for tc in normalized.tool_calls] == ["hermes_tool_search"] + + def test_alias_collision_takes_suffix_no_duplicates(self, transport, monkeypatch): + """A real tool already named ``hermes_tool_search`` keeps its wire + name; the bridge is suffixed and both round-trip independently.""" + tools = [ + {"type": "function", "function": { + "name": "hermes_tool_search", "description": "Real tool.", + "parameters": {"type": "object", "properties": {}}}}, + {"type": "function", "function": { + "name": "tool_search", "description": "Bridge.", + "parameters": {"type": "object", "properties": {}}}}, + ] + kw = transport.build_kwargs( + model="grok-4.6", + messages=[{"role": "user", "content": "hi"}], + tools=tools, + is_xai_responses=True, + ) + names = self._names(kw) + assert names == ["hermes_tool_search", "hermes_tool_search_2"] + assert len(names) == len(set(names)) + assert transport._last_wire_aliases == {"hermes_tool_search_2": "tool_search"} + # Bridge alias maps back; the real tool's name is untouched. + normalized = self._normalize_named_call( + transport, monkeypatch, "hermes_tool_search_2" + ) + assert [tc.name for tc in normalized.tool_calls] == ["tool_search"] + normalized2 = self._normalize_named_call( + transport, monkeypatch, "hermes_tool_search" + ) + assert [tc.name for tc in normalized2.tool_calls] == ["hermes_tool_search"] + + def test_legacy_fallback_without_provenance(self, transport, monkeypatch): + """Normalize-only call sites (no build_kwargs on this instance) keep + the historical unconditional reverse mapping.""" + assert transport._last_wire_aliases is None + normalized = self._normalize_named_call( + transport, monkeypatch, "hermes_tool_search" + ) + assert [tc.name for tc in normalized.tool_calls] == ["tool_search"] + class TestXaiWebSearchBackendPreference: """``_xai_prefers_native_web_search`` must honor web backend config."""