mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(responses): keep namespace tools intact through guardrail write-back
Guardrails on /v1/responses see tools flattened to Chat Completions format, where a namespace tool becomes one function per member named namespace__member. The write-back merged that flat list back positionally over the original tools, replacing the namespace tool with flat functions and breaking Codex MCP tool calls. Rebuild data['tools'] from the original Responses tools and apply only the guardrail's delta: drop members or tools it removed, append tools it added. Fixes #39183 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI
parent
558f42e304
commit
6e4c862217
2 changed files with 202 additions and 21 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"}]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue