diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 24dbd7c08d0..6d413948aa4 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,8 +28,9 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Final, Union, cast +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, TypeGuard, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam @@ -84,6 +85,82 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +_TOOL_TYPES_NOT_SENT_TO_GUARDRAIL: Final = frozenset( + {"web_search", "web_search_preview", "computer_use", "image_generation", "shell"} +) + + +def _is_str_object_mapping(value: object) -> TypeGuard[Mapping[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, Mapping) + + +def _is_object_sequence(value: object) -> TypeGuard[Sequence[object]]: # guard-ok: isinstance narrows correctly; str/bytes excluded # fmt: skip + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + + +def _namespace_members(tool: Mapping[str, object]) -> tuple[object, ...]: + members: Final = tool.get("tools") + return tuple(members) if _is_object_sequence(members) else () + + +def _is_function_member(member: object) -> bool: + return _is_str_object_mapping(member) and member.get("type") == "function" + + +def _qualified_member_name(namespace: str, member: object) -> str: + member_name: Final = member.get("name") if _is_str_object_mapping(member) else None + return f"{namespace}__{member_name or ''}" + + +def _flattened_function_names(tools: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + """Names the guardrail sees for ``tools`` once flattened to Chat Completions format.""" + top_level: Final = tuple( + str(tool.get("name") or "") for tool in tools if tool.get("type") in ("function", "custom") + ) + nested: Final = tuple( + _qualified_member_name(str(tool.get("name") or ""), member) + for tool in tools + if tool.get("type") == "namespace" + for member in _namespace_members(tool) + if _is_function_member(member) + ) + return top_level + nested + + +def _merge_namespace_tool( + tool: dict[str, object], remapped_functions: Mapping[str, dict[str, object]] +) -> dict[str, object] | None: + namespace: Final = str(tool.get("name") or "") + members: Final = _namespace_members(tool) + surviving: Final = tuple( + member + for member in members + if not _is_function_member(member) or _qualified_member_name(namespace, member) in remapped_functions + ) + if len(surviving) == len(members): + return tool + if not any(_is_function_member(member) for member in surviving): + return None + return {**tool, "tools": list(surviving)} + + +def _merge_original_tool( + tool: dict[str, object], + remapped_functions: Mapping[str, dict[str, object]], + remapped_passthrough: Sequence[dict[str, object]], +) -> dict[str, object] | None: + tool_type: Final = tool.get("type") + if tool_type in _TOOL_TYPES_NOT_SENT_TO_GUARDRAIL: + return tool + if tool_type == "function": + return remapped_functions.get(str(tool.get("name") or "")) + if tool_type == "custom": + return tool if str(tool.get("name") or "") in remapped_functions else None + if tool_type == "namespace": + return _merge_namespace_tool(tool, remapped_functions) + return tool if tool in remapped_passthrough else None + + def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) @@ -275,28 +352,33 @@ class OpenAIResponsesHandler(BaseTranslation): remapped: list[dict[str, object]], ) -> list[dict[str, object]]: """ - Merge remapped guardrailed tools with original tools that were not sent - to the guardrail (e.g. web_search, web_search_preview), preserving order. - Tools a guardrail appended (``remapped`` longer than ``original_tools``) - have no original slot and are kept so an injected tool is not dropped. + Rebuild the Responses tool list from ``original_tools`` and apply only the + guardrail's delta: tools it dropped are removed (namespace members + individually), tools it appended are kept at the end, and tools that never + reach the guardrail (web_search, computer_use, ...) stay untouched. + Namespace tools keep their wrapper and member names instead of the + ``namespace__tool`` functions the guardrail saw. """ if not original_tools: return remapped - result: Final[list[dict[str, object]]] = [] - j = 0 - for tool in original_tools: - if isinstance(tool, dict) and tool.get("type") in ( - "web_search", - "web_search_preview", - ): - result.append(tool) - else: - if j < len(remapped): - result.append(remapped[j]) - j += 1 - # Keep guardrail-appended tools that matched no original slot above. - result.extend(remapped[j:]) - return result + remapped_functions: Final = MappingProxyType( + {str(tool.get("name") or ""): tool for tool in remapped if tool.get("type") == "function"} + ) + remapped_passthrough: Final = tuple(tool for tool in remapped if tool.get("type") != "function") + original_function_names: Final = frozenset(_flattened_function_names(original_tools)) + kept: Final = tuple( + merged + for tool in original_tools + for merged in (_merge_original_tool(tool, remapped_functions, remapped_passthrough),) + if merged is not None + ) + appended: Final = tuple( + tool + for tool in remapped + if (tool.get("type") == "function" and str(tool.get("name") or "") not in original_function_names) + or (tool.get("type") != "function" and tool not in original_tools) + ) + return [*kept, *appended] def _apply_guardrailed_tools_to_data( self, diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..0fce0fde74b 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1229,3 +1229,102 @@ class TestOpenAIResponsesHandlerToolInjection: names = [t.get("name") for t in result["tools"]] assert "get_weather" in names assert "injected_tool" in names + + +class ToolDroppingGuardrail(CustomGuardrail): + """Guardrail that removes one flattened function tool by name.""" + + def __init__(self, drop: str, **kwargs): + super().__init__(**kwargs) + self.drop = drop + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = [t for t in inputs.get("tools") or [] if t["function"]["name"] != self.drop] + return inputs + + +CODEX_NAMESPACE_TOOL = { + "type": "namespace", + "name": "mcp__confluence", + "description": "Tools from the confluence MCP server", + "tools": [ + { + "type": "function", + "name": "confluence_get_page", + "description": "Get a page", + "parameters": {"type": "object", "properties": {"page_id": {"type": "string"}}}, + }, + { + "type": "function", + "name": "confluence_search", + "description": "Search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + ], +} + + +class TestOpenAIResponsesHandlerNamespaceTools: + """Regression for #39183: a guardrail on /v1/responses must not rewrite Codex + namespace MCP tools into flat ``namespace__tool`` functions.""" + + @pytest.mark.asyncio + async def test_namespace_tool_survives_tool_appending_guardrail(self): + handler = OpenAIResponsesHandler() + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [ + {"type": "function", "name": "shell", "parameters": {"type": "object", "properties": {}}}, + CODEX_NAMESPACE_TOOL, + ], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages(data, ToolAppendingGuardrail(guardrail_name="test")) + + assert [t["type"] for t in result["tools"]] == ["function", "namespace", "function"] + assert result["tools"][0]["name"] == "shell" + assert result["tools"][1] == CODEX_NAMESPACE_TOOL + assert result["tools"][2]["name"] == "injected_tool" + assert not any(t["name"].startswith("mcp__confluence__") for t in result["tools"]) + + @pytest.mark.asyncio + async def test_guardrail_can_still_drop_a_single_namespace_member(self): + handler = OpenAIResponsesHandler() + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [CODEX_NAMESPACE_TOOL], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages( + data, ToolDroppingGuardrail(drop="mcp__confluence__confluence_search", guardrail_name="test") + ) + + assert len(result["tools"]) == 1 + namespace_tool = result["tools"][0] + assert namespace_tool["type"] == "namespace" + assert namespace_tool["name"] == "mcp__confluence" + assert [t["name"] for t in namespace_tool["tools"]] == ["confluence_get_page"] + + @pytest.mark.asyncio + async def test_namespace_tool_removed_when_guardrail_drops_every_member(self): + handler = OpenAIResponsesHandler() + single_member = {**CODEX_NAMESPACE_TOOL, "tools": CODEX_NAMESPACE_TOOL["tools"][:1]} + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [{"type": "web_search_preview"}, single_member], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages( + data, ToolDroppingGuardrail(drop="mcp__confluence__confluence_get_page", guardrail_name="test") + ) + + assert result["tools"] == [{"type": "web_search_preview"}]