From 31774a7b2ee1f2c77a90f2c577ba16fbad4b85fe Mon Sep 17 00:00:00 2001 From: Silkforgeai Date: Thu, 29 Oct 2026 14:20:51 -0400 Subject: [PATCH] fix(guardrails): enforce ToolPolicyGuardrail across /v1/messages, /v1/responses, and mixed tool requests (#49100) --- .../tool_policy/tool_policy_guardrail.py | 77 +++++++++--- .../test_tool_policy_guardrail.py | 112 ++++++++++++++++++ 2 files changed, 173 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index fbf1f2576ff..d7a248f1a15 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -91,12 +91,13 @@ def _get_request_route_from_data(request_data: dict) -> str | None: def _resolve_tool_names_from_messages(messages: list[dict]) -> dict[str, str]: """ - Build a map of tool_call_id -> tool_name from assistant messages' tool_calls. + Build a map of tool_call_id -> tool_name from assistant messages' tool_calls + (OpenAI Chat Completions) and tool_use content blocks (Anthropic /v1/messages). Used to resolve which tool produced each tool result in the conversation. """ mapping: Final[dict[str, str]] = {} for msg in messages: - if msg.get("role") != "assistant": + if not isinstance(msg, dict) or msg.get("role") != "assistant": continue tool_calls = msg.get("tool_calls") or [] for tc in tool_calls: @@ -109,9 +110,63 @@ def _resolve_tool_names_from_messages(messages: list[dict]) -> dict[str, str]: fn = getattr(fn_obj, "name", None) if fn_obj else None if tc_id and fn: mapping[tc_id] = fn + content = msg.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_use": + tc_id = block.get("id") + fn = block.get("name") + if isinstance(tc_id, str) and tc_id and isinstance(fn, str) and fn: + mapping[tc_id] = fn return mapping +def _extract_executed_tool_names(request_data: dict) -> list[str]: + """ + Return the ordered list of tool names whose outputs appear in the conversation + across OpenAI Chat Completions (`role == "tool"`), Anthropic `/v1/messages` + (`type == "tool_result"` blocks), and `/v1/responses` (`type == "function_call_output"`). + """ + executed: Final[list[str]] = [] + messages: Final = request_data.get("messages") or [] + if isinstance(messages, list) and messages: + tc_id_to_name: Final = _resolve_tool_names_from_messages(messages) + for msg in messages: + if not isinstance(msg, dict): + continue + if msg.get("role") == "tool": + tool_call_id = msg.get("tool_call_id") + source_tool = tc_id_to_name.get(tool_call_id, "") if isinstance(tool_call_id, str) else "" + if source_tool: + executed.append(source_tool) + content = msg.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + tool_use_id = block.get("tool_use_id") + source_tool = tc_id_to_name.get(tool_use_id, "") if isinstance(tool_use_id, str) else "" + if source_tool: + executed.append(source_tool) + + input_items: Final = request_data.get("input") + if isinstance(input_items, list) and input_items: + call_id_to_name: Final[dict[str, str]] = {} + for item in input_items: + if isinstance(item, dict) and item.get("type") == "function_call": + call_id = item.get("call_id") or item.get("id") + fn = item.get("name") + if isinstance(call_id, str) and call_id and isinstance(fn, str) and fn: + call_id_to_name[call_id] = fn + for item in input_items: + if isinstance(item, dict) and item.get("type") == "function_call_output": + call_id = item.get("call_id") + source_tool = call_id_to_name.get(call_id, "") if isinstance(call_id, str) else "" + if source_tool: + executed.append(source_tool) + + return executed + + class ToolPolicyGuardrail(CustomGuardrail): """ Guardrail that enforces per-tool input/output policies from the in-memory @@ -146,11 +201,10 @@ class ToolPolicyGuardrail(CustomGuardrail): if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name") ] route: Final = _get_request_route_from_data(request_data) + route_tool_names: Final = extract_request_tool_names(route, request_data) if route else [] replayed_tool_names: Final = replayed_tool_search_tool_names(_REQUEST_DATA.validate_python(request_data)) - tool_names = ( - extract_request_tool_names(route, request_data) - if route and not function_tool_names - else [*function_tool_names, *replayed_tool_names] + tool_names = list( + dict.fromkeys([*function_tool_names, *route_tool_names, *replayed_tool_names]) ) else: tool_calls: Final = inputs.get("tool_calls") or [] @@ -204,17 +258,8 @@ class ToolPolicyGuardrail(CustomGuardrail): if input_type == "response": trusted_input_tools: Final = [name for name in tool_names if policy_map.get(name) == "trusted"] if trusted_input_tools: - messages: Final = request_data.get("messages") or [] - tc_id_to_name: Final = _resolve_tool_names_from_messages(messages) - untrusted_sources: Final[list[str]] = [] - for msg in messages: - if msg.get("role") != "tool": - continue - tool_call_id = msg.get("tool_call_id") - source_tool = tc_id_to_name.get(tool_call_id, "") if tool_call_id else "" - if not source_tool: - continue + for source_tool in _extract_executed_tool_names(request_data): if registry.get_output_policy(source_tool) == "untrusted": if source_tool not in untrusted_sources: untrusted_sources.append(source_tool) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py index f40465e8d7b..6a4c7b54aa1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -222,3 +222,115 @@ async def test_blocked_tool_loaded_by_a_replayed_tool_search_raises_next_to_allo ) assert exc_info.value.status_code == 400 assert exc_info.value.detail["blocked_tools"] == ["run_shell"] + + +def _registry_mock_with_output(input_policies: dict, output_policies: dict): + reg = _registry_mock(input_policies) + reg.get_output_policy.side_effect = lambda name: output_policies.get(name, "untrusted") + return reg + + +@pytest.mark.asyncio +async def test_anthropic_messages_tool_result_blocks_trusted_tool_call(guardrail): + request_data = { + "user_api_key_request_route": "/v1/messages", + "messages": [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "read_external_note", + "input": {}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": "Note contents", + } + ], + }, + ], + } + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock_with_output( + {"read_external_note": "untrusted", "send_summary": "trusted"}, + {"read_external_note": "untrusted", "send_summary": "trusted"}, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_tool_response_inputs(["send_summary"]), + request_data=request_data, + input_type="response", + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["blocked_tools"] == ["send_summary"] + assert exc_info.value.detail["untrusted_sources"] == ["read_external_note"] + + +@pytest.mark.asyncio +async def test_responses_api_function_call_output_blocks_trusted_tool_call(guardrail): + request_data = { + "user_api_key_request_route": "/v1/responses", + "input": [ + { + "type": "function_call", + "call_id": "call_1", + "name": "read_external_note", + "arguments": "{}", + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "Note contents", + }, + ], + } + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock_with_output( + {"read_external_note": "untrusted", "send_summary": "trusted"}, + {"read_external_note": "untrusted", "send_summary": "trusted"}, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_tool_response_inputs(["send_summary"]), + request_data=request_data, + input_type="response", + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["blocked_tools"] == ["send_summary"] + assert exc_info.value.detail["untrusted_sources"] == ["read_external_note"] + + +@pytest.mark.asyncio +async def test_request_stage_blocks_route_tool_when_function_tool_names_non_empty(guardrail): + request_data = { + "user_api_key_request_route": "/v1/messages", + "tools": [ + {"name": "ping", "input_schema": {"type": "object"}}, + {"type": "bash_20250124", "name": "bash"}, + ], + } + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock({"ping": "untrusted", "bash": "blocked"}), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_tool_request_inputs(["ping"]), + request_data=request_data, + input_type="request", + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["blocked_tools"] == ["bash"] +