diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index a4a43630ba0..bc12ccd42d3 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -174,9 +174,9 @@ def _guardrail_translations_from( def _guardrail_translation_modules() -> Iterator[str]: - """Yield every module that can declare guardrail translation handlers, the optional MCP one first.""" - yield _MCP_GUARDRAIL_TRANSLATION_MODULE + """Yield every module that can declare guardrail translation handlers, the optional MCP one last.""" yield from _bundled_guardrail_translation_modules() + yield _MCP_GUARDRAIL_TRANSLATION_MODULE def _discover( diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 64d5af0f772..ffb788be497 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -134,7 +134,7 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N data["litellm_metadata"] = user_metadata -_UNSCANNED_WARNING_KEYS: Final = 256 +_UNSCANNED_WARNING_KEYS: Final = 4096 def _resolved_call_type(call_type: str | None) -> CallTypes | None: @@ -154,22 +154,29 @@ def _warn_left_unscanned_once( ) -> None: if _resolved_call_type(call_type) is not None: verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; %s.", + "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; %s. " + "Add a guardrail translation handler for that call type.", guardrail_name, request_route, call_type, consequence, ) return - unscannable: Final = ( - f"call type '{call_type}' is not one litellm can scan" if call_type else "its call type could not be resolved" + unscannable, remedy = ( + ( + f"call type '{call_type}' is not one litellm can scan", + "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES", + ) + if call_type + else ("its call type could not be resolved", "Add the route to API_ROUTE_TO_CALL_TYPES") ) verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but %s, so no guardrail can run on that route; %s.", + "Guardrail '%s' selected for route '%s' but %s, so no guardrail can run on that route; %s. %s.", guardrail_name, request_route, unscannable, consequence, + remedy, ) @@ -368,7 +375,8 @@ class UnifiedLLMGuardrails(CustomLogger): call_type = logging_call_type mappings: Final = load_guardrail_translation_mappings() - if _resolved_call_type(call_type) not in mappings: + resolved_call_type: Final = _resolved_call_type(call_type) + if resolved_call_type is None or resolved_call_type not in mappings: _warn_left_unscanned( guardrail_to_apply=guardrail_to_apply, user_api_key_dict=user_api_key_dict, @@ -377,7 +385,7 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[resolved_call_type]()) try: response = await endpoint_translation.process_output_response( @@ -479,7 +487,7 @@ class UnifiedLLMGuardrails(CustomLogger): has no in-stream error frame) the exception is re-raised so the proxy can report it with a real HTTP status. """ - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: + if _resolved_call_type(call_type) in A2A_CALL_TYPES: yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data)) return if stream_started and endpoint_translation is not None: @@ -1188,14 +1196,15 @@ class UnifiedLLMGuardrails(CustomLogger): yield item # Stream has ended - do final processing with all collected chunks - if call_type is not None and CallTypes(call_type) in mappings: + final_call_type: Final = _resolved_call_type(call_type) + if final_call_type is not None and final_call_type in mappings: verbose_proxy_logger.debug( "Processing final streaming response with all %s chunks for guardrail %s", len(responses_so_far), guardrail_to_apply.guardrail_name, ) - endpoint_translation = mappings[CallTypes(call_type)]() + endpoint_translation = mappings[final_call_type]() # When buffering, snapshot the original chunks before moderation. # A shallow copy suffices: end-of-stream diff --git a/tests/test_litellm/llms/test_guardrail_translation_discovery.py b/tests/test_litellm/llms/test_guardrail_translation_discovery.py index 3b4446e4da5..1879b15f5fd 100644 --- a/tests/test_litellm/llms/test_guardrail_translation_discovery.py +++ b/tests/test_litellm/llms/test_guardrail_translation_discovery.py @@ -10,6 +10,9 @@ import pytest import litellm.llms as llms_package from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) from litellm.types.utils import CallTypes OPENAI_CHAT_TRANSLATION_MODULE = "litellm.llms.openai.chat.guardrail_translation" @@ -184,3 +187,19 @@ def test_a_broken_mcp_dependency_is_reported_and_retried(caplog): assert CallTypes.call_mcp_tool in recovered assert not llms_package.guardrail_translation_discovery.unavailable + + +class _StandInMCPHandler: + """A bundled package's handler that claims the call type the MCP package owns.""" + + +def test_the_mcp_package_wins_a_handler_collision_with_a_bundled_package(monkeypatch): + """MCP handlers are the specialised ones, so a bundled package declaring the same call type must not shadow them.""" + colliding = ModuleType("litellm.llms.colliding_stub.guardrail_translation") + colliding.guardrail_translation_mappings = {CallTypes.call_mcp_tool: _StandInMCPHandler} + monkeypatch.setitem(sys.modules, colliding.__name__, colliding) + monkeypatch.setattr(llms_package, "_bundled_guardrail_translation_modules", lambda: iter((colliding.__name__,))) + + discovery = llms_package.discover_guardrail_translations() + + assert discovery.mappings[CallTypes.call_mcp_tool] is MCPGuardrailTranslationHandler diff --git a/tests/test_litellm/proxy/guardrails/conftest.py b/tests/test_litellm/proxy/guardrails/conftest.py new file mode 100644 index 00000000000..8a6341440f4 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/conftest.py @@ -0,0 +1,13 @@ +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, +) + + +@pytest.fixture(autouse=True) +def _forget_unscanned_warnings(): + """The unscanned warning fires once per route and reason, so every guardrail test starts with nothing remembered.""" + unified_module._warn_left_unscanned_once.cache_clear() + yield + unified_module._warn_left_unscanned_once.cache_clear() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 65fab6e1270..cc3437c46f3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -85,14 +85,6 @@ def _patch_translation_mappings(monkeypatch, mappings): monkeypatch.setattr(unified_module, "load_guardrail_translation_mappings", lambda: mappings) -@pytest.fixture(autouse=True) -def _forget_unscanned_warnings(): - """The unscanned warning fires once per route and reason, so each test starts with nothing remembered.""" - unified_module._warn_left_unscanned_once.cache_clear() - yield - unified_module._warn_left_unscanned_once.cache_clear() - - @pytest.fixture(autouse=True) def _inject_mcp_handler_mapping(monkeypatch): """Inject MCP handler mapping so the unified guardrail can run inside tests.""" @@ -2356,6 +2348,7 @@ class TestUnscannedStreamIsAnnounced: assert len(chunks) == 3 assert any( "no guardrail translation handler" in message + and "Add a guardrail translation handler for that call type." in message and "recording-guardrail" in message and "/chat/completions" in message for message in warnings @@ -2374,6 +2367,7 @@ class TestUnscannedStreamIsAnnounced: assert len(chunks) == 3 assert any( "call type could not be resolved" in message + and "Add the route to API_ROUTE_TO_CALL_TYPES." in message and "recording-guardrail" in message for message in warnings ), warnings @@ -2473,6 +2467,7 @@ class TestUnscannedRequestIsAnnounced: assert returned["messages"] == [{"role": "user", "content": "hello world"}] assert any( "call type 'not_a_call_type' is not one litellm can scan" in message + and "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES." in message and "skipping pre-call scanning" in message for message in self._warnings(caplog) ), self._warnings(caplog) @@ -2494,6 +2489,7 @@ class TestUnscannedRequestIsAnnounced: assert returned["messages"] == [{"role": "user", "content": "hello world"}] assert any( "call type 'not_a_call_type' is not one litellm can scan" in message + and "Map the route to a CallTypes member in API_ROUTE_TO_CALL_TYPES." in message and "skipping during-call scanning" in message for message in self._warnings(caplog) ), self._warnings(caplog)