From 8284208af261bd32d78ff3fb43040117894fd358 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:51:51 -0700 Subject: [PATCH] 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. --- .../guardrails/auto_router_compression.py | 52 ++++++++++++++--- .../test_auto_router_compression.py | 56 +++++++++++++++++-- tests/test_litellm/test_router.py | 7 ++- 3 files changed, 103 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index e7f58662249..b6f32c1c46d 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/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 } diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 0b47e56cb02..676b3ba2967 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/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 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9e6a88d3433..6279c8a5404 100644 --- a/tests/test_litellm/test_router.py +++ b/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