fix(auto-router compression): restrict both hops to real compression guardrails

The two policy fields are operator-supplied names and nothing else constrained them.
The routing hop calls apply_guardrail directly, which hands the guardrail the
conversation and POSTs it to whatever service backs that guardrail, and the model hop
is added to metadata["guardrails"], which runs it even when it is not default_on. So
naming an ordinary guardrail turned either hop into a way to invoke it and ship prompt
content to it. Both hops now refuse a name that does not resolve to an active
compression guardrail, and say so in the log rather than failing quietly.

litellm/proxy/guardrails/auto_router_compression.py

@@ -135,19 +135,36 @@ def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None:
     return None
 
 
+def _compression_guardrail_classes() -> tuple[type, ...]:
+    """The registered guardrail classes whose provider compresses prompts."""
+    from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
+
+    return tuple(cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS)
+
+
+def is_compression_guardrail(guardrail: object) -> bool:
+    """Whether `guardrail` is an instance of a compression guardrail provider.
+
+    Both hops are validated through here. The two policy fields are operator-supplied
+    names and nothing else constrains them, so without this a name that resolves to an
+    ordinary guardrail would be handed the conversation and invoked: the routing hop
+    calls `apply_guardrail` directly, which POSTs the content wherever that guardrail
+    sends it, and the model hop is added to `metadata["guardrails"]`, which runs it even
+    when it is not `default_on`.
+    """
+    classes: Final = _compression_guardrail_classes()
+    return bool(classes) and isinstance(guardrail, classes)
+
+
 def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]:
     """Every currently-active guardrail whose type is a compression guardrail."""
     import litellm
     from litellm.integrations.custom_guardrail import CustomGuardrail
-    from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
 
-    compression_classes: Final = tuple(
-        cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS
-    )
-    if not compression_classes:
+    if not _compression_guardrail_classes():
         return ()
     active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail)
-    return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name)
+    return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name)
 
 
 async def arm_pre_call(
@@ -192,7 +209,18 @@ async def arm_pre_call(
         )
     )
 
-    if policy.model is not None:
+    # Only a name that resolves to a real compression guardrail may be armed: this adds
+    # it to `metadata["guardrails"]`, which runs it even when it is not `default_on`.
+    armed_model_hop: Final = policy.model is not None and any(
+        guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails()
+    )
+    if policy.model is not None and not armed_model_hop:
+        verbose_proxy_logger.warning(
+            "AutoRouter compression: '%s' is not an active compression guardrail; the model hop is uncompressed",
+            policy.model,
+        )
+
+    if armed_model_hop:
         _model_hop_armed.set(True)
         _, metadata = get_or_create_metadata_bucket(data)
         requested: Final = metadata.get("guardrails")
@@ -249,6 +277,16 @@ async def messages_for_routing(
         )
         return _as_routing_messages(messages)
 
+    # apply_guardrail below hands this guardrail the conversation and it POSTs the
+    # content to whatever service backs it, so the name has to be a compression
+    # guardrail rather than any guardrail the operator happened to name.
+    if not is_compression_guardrail(guardrail):
+        verbose_proxy_logger.warning(
+            "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages",
+            policy.routing,
+        )
+        return _as_routing_messages(messages)
+
     inputs: Final[GenericGuardrailAPIInputs] = {
         "structured_messages": _as_routing_messages(messages)  # pyright: ignore[reportAssignmentType]  # plain dicts, not AllMessageValues; see headroom.py's own use of this shape
     }

tests/test_litellm/proxy/guardrails/test_auto_router_compression.py

@@ -186,15 +186,33 @@ class _RecordingCompressionGuardrail(CustomGuardrail):
 
 
 @pytest.fixture
-def registered_guardrail():
+def registered_guardrail(monkeypatch):
     import litellm
+    from litellm.proxy.guardrails import guardrail_registry
 
+    # Registered under a compression provider name: both hops refuse a name that does
+    # not resolve to one, so a bare callback would (correctly) never be used.
+    monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail)
     guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress")
     litellm.logging_callback_manager.add_litellm_callback(guardrail)
     yield guardrail
     litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail)
 
 
+class _NonCompressionGuardrail(CustomGuardrail):
+    """A guardrail that is not a compression provider, e.g. a PII or content filter."""
+
+    def __init__(self, guardrail_name: str):
+        super().__init__(guardrail_name=guardrail_name)
+        self.called = False
+
+    async def apply_guardrail(
+        self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None
+    ) -> GenericGuardrailAPIInputs:
+        self.called = True
+        return inputs
+
+
 class TestArmPreCall:
     @pytest.mark.asyncio
     async def test_no_router_is_noop(self):
@@ -266,7 +284,13 @@ class TestArmPreCall:
             litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail)
 
     @pytest.mark.asyncio
-    async def test_model_side_guardrail_is_requested_even_when_not_default_on(self):
+    async def test_model_side_guardrail_is_requested_even_when_not_default_on(self, monkeypatch):
+        import litellm
+        from litellm.proxy.guardrails import guardrail_registry
+
+        monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail)
+        active = _RecordingCompressionGuardrail(guardrail_name="headroom-b")
+        litellm.logging_callback_manager.add_litellm_callback(active)
         router = _FakeRouter(
             [
                 {
@@ -280,8 +304,11 @@ class TestArmPreCall:
             ]
         )
         data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
-        await arm_pre_call(data=data, llm_router=router)
-        assert data["metadata"]["guardrails"] == ["headroom-b"]
+        try:
+            await arm_pre_call(data=data, llm_router=router)
+            assert data["metadata"]["guardrails"] == ["headroom-b"]
+        finally:
+            litellm.logging_callback_manager.remove_callback_from_all_lists(active)
 
     @pytest.mark.asyncio
     async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self):
@@ -347,6 +374,27 @@ class TestMessagesForRouting:
         assert result == [{"role": "user", "content": "[COMPRESSED] my ssn is [REDACTED]"}]
         assert registered_guardrail.request_data_seen[0]["messages"] == masked
 
+    @pytest.mark.asyncio
+    async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch):
+        """Regression (security): the policy fields are operator-supplied names that
+        nothing else constrains. apply_guardrail hands the guardrail the conversation
+        and it POSTs that content to whatever service backs it, so naming an ordinary
+        guardrail must not turn the routing hop into a way to ship prompts there."""
+        import litellm
+
+        other = _NonCompressionGuardrail(guardrail_name="pii-filter")
+        litellm.logging_callback_manager.add_litellm_callback(other)
+        try:
+            policy = AutoRouterCompressionPolicy(routing="pii-filter", model=None)
+            messages = [{"role": "user", "content": "my ssn is 123-45-6789"}]
+
+            result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={})
+
+            assert other.called is False
+            assert result == messages
+        finally:
+            litellm.logging_callback_manager.remove_callback_from_all_lists(other)
+
     @pytest.mark.asyncio
     async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail):
         """Regression: a real compression guardrail writes its stats onto whatever

tests/test_litellm/test_router.py

@@ -9686,7 +9686,12 @@ class TestAutoRouterCompressionDecoupling:
         return router, strategy
 
     @pytest.fixture
-    def registered_guardrail(self):
+    def registered_guardrail(self, monkeypatch):
+        from litellm.proxy.guardrails import guardrail_registry
+
+        # Registered under a compression provider name: both hops refuse a name that
+        # does not resolve to one, so a bare callback would never be used.
+        monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", self._CompressingGuardrail)
         guardrail = self._CompressingGuardrail(guardrail_name="fake-compress")
         litellm.logging_callback_manager.add_litellm_callback(guardrail)
         yield guardrail