From 9e286fe94bf18d29b7bb56e3c2f77d114c10acc5 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:42:11 -0700 Subject: [PATCH] fix(auto-router): close review findings on per-hop compression - Suppression markers now carry the per-process token `_pre_call_marker` already uses, so a caller cannot switch off an always-on PII, content-filter or compression guardrail by naming it in its own request metadata. - Routing set to "none" with the model side compressed now classifies on the pre-compression snapshot instead of the model-side guardrail's output. - Both the proxy's pre-call arming and the router's routing hook resolve the policy through one tag-aware `policy_for_model`, so an alias with several tag-scoped markers can no longer suppress one marker's guardrail and then route under another marker's policy. - The pre-compression snapshot moved from request metadata to a ContextVar: `refresh_proxy_server_request_body_snapshot` copies metadata into `proxy_server_request.body`, which deployments persist, and the snapshot holds the prompt as it was before any masking guardrail rewrote it. - The compression selector lists Compresr guardrails too, not just Headroom. --- litellm/integrations/custom_guardrail.py | 22 ++- .../guardrails/auto_router_compression.py | 148 +++++++++------ litellm/router.py | 32 ++-- .../integrations/test_custom_guardrail.py | 35 +++- .../test_auto_router_compression.py | 177 +++++++++++++----- tests/test_litellm/test_router.py | 29 +++ .../add_model/CompressionControls.tsx | 5 +- .../add_model/buildAutoRouterCompression.ts | 7 + 8 files changed, 322 insertions(+), 133 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7f8effa2317..558e97cfc16 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -941,17 +941,29 @@ class CustomGuardrail(CustomLogger): """ return False - def _suppressed_by_auto_router_compression(self, data: dict) -> bool: - """True when an auto router's own compression policy suppresses this guardrail. + def auto_router_suppression_marker(self) -> str | None: + """The value `arm_pre_call` must write to suppress this guardrail. - Set only by litellm.proxy.guardrails.auto_router_compression.arm_pre_call, never - by the caller, so a request cannot suppress its own guardrails this way. + Carries the per-process token for the same reason `_pre_call_marker` does: a + caller controls request metadata, so a bare guardrail name there would let any + request switch off a PII, content-filter, or compression guardrail for itself. + The token is never sent to the caller, so the marker cannot be forged. """ + name: Final = self.guardrail_name + if not name: + return None + return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + + def _suppressed_by_auto_router_compression(self, data: dict[str, object]) -> bool: + """True when an auto router's own compression policy suppresses this guardrail.""" + marker: Final = self.auto_router_suppression_marker() + if marker is None: + return False for meta_key in ("metadata", "litellm_metadata"): meta = data.get(meta_key) if isinstance(meta, dict): suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, list) and self.guardrail_name in suppressed: + if isinstance(suppressed, list) and marker in suppressed: return True return False diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index c3fba937d22..7ccd1937543 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -12,34 +12,32 @@ guardrail is suppressed for that request, and only these two settings decide wha each hop sees. """ -import copy -from collections.abc import Mapping +import contextvars +from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY -from litellm.litellm_core_utils.core_helpers import ( - get_metadata_variable_name_from_kwargs, - get_or_create_metadata_bucket, -) +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.router import Router -else: - CustomGuardrail = Any - Router = Any COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# Metadata key stashing the pre-compression messages so a routing decision that -# names a different compression than the model call still compresses the -# original text, not whatever the model-side guardrail already rewrote it to. -AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: Final = "_auto_router_routing_messages_snapshot" +# The pre-compression messages, so a routing decision that does not share the model +# call's compression still classifies on the original text. Deliberately a ContextVar +# rather than a metadata key: `refresh_proxy_server_request_body_snapshot` copies +# metadata into `proxy_server_request.body`, which deployments persist to spend logs, +# and this holds the prompt as it was before any masking guardrail rewrote it. +_routing_messages_snapshot: Final[contextvars.ContextVar[tuple[Mapping[str, object], ...] | None]] = ( + contextvars.ContextVar("litellm_auto_router_routing_messages_snapshot", default=None) +) @dataclass(frozen=True, slots=True) @@ -72,30 +70,51 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( - llm_router: "Router | None", model_alias: str, team_id: str | None + llm_router: "Router | None", + model_alias: str, + team_id: str | None, + request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy declared by the auto router marker deployment `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to. - Mirrors the alias lookup in ``_check_and_merge_model_level_guardrails``: this runs - before routing has picked a strategy, so it takes the first marker deployment for - the alias rather than disambiguating by request tags. + Both the proxy's pre-call arming and the router's routing hook resolve the policy + through here, with the same tag rule, so an alias carrying several tag-scoped + markers can never suppress one marker's guardrail and then route under another + marker's policy. """ if llm_router is None: return None deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] - for deployment in deployments: - litellm_params: Final = deployment.get("litellm_params") or {} - model_field = litellm_params.get("model") - if not isinstance(model_field, str) or not model_field.startswith(AUTO_ROUTER_MODEL_PREFIX): - continue - policy = policy_from_litellm_params(litellm_params) + markers: Final = tuple( + litellm_params + for deployment in deployments + if isinstance(litellm_params := deployment.get("litellm_params") or {}, Mapping) + and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + ) + requested: Final = frozenset(request_tags) + tag_matched: Final = tuple( + params for params in markers if requested.issuperset(frozenset(params.get("tags") or ())) + ) + for params in (*tag_matched, *markers): + policy = policy_from_litellm_params(params) if policy is not None: return policy return None -def _active_compression_guardrail_names() -> frozenset[str]: - """Names of every currently-active guardrail whose type is a compression guardrail.""" +def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + for meta_key in ("metadata", "litellm_metadata"): + meta = request_kwargs.get(meta_key) + if isinstance(meta, Mapping): + team_id = meta.get("user_api_key_team_id") + if isinstance(team_id, str): + return team_id + return None + + +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 @@ -104,21 +123,20 @@ def _active_compression_guardrail_names() -> frozenset[str]: cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS ) if not compression_classes: - return frozenset() + return () active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) - return frozenset( - cb.guardrail_name for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name - ) + return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name) -async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: +async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | None") -> MutableMapping[str, object]: """Apply an auto router's compression policy, if any, before guardrails run. Suppresses every other compression guardrail, re-enables the model-side guardrail the policy names (if any) even when it isn't ``default_on``, and - snapshots the pre-compression messages so the routing decision can compress - them independently of whatever the model-side guardrail does to `data`. + snapshots the pre-compression messages so the routing decision can read them + independently of whatever the model-side guardrail does to `data`. """ + _routing_messages_snapshot.set(None) if llm_router is None: return data @@ -129,21 +147,27 @@ async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: # Read-only until a policy is confirmed: creating the metadata bucket for every # request, including the vast majority with no auto-router compression policy, # would be an unwanted side effect of merely checking for one. - metadata_key: Final = get_metadata_variable_name_from_kwargs(data) - existing_bucket: Final = data.get(metadata_key) - other_bucket: Final = data.get("metadata" if metadata_key == "litellm_metadata" else "litellm_metadata") - team_id: Final = (existing_bucket.get("user_api_key_team_id") if isinstance(existing_bucket, dict) else None) or ( - other_bucket.get("user_api_key_team_id") if isinstance(other_bucket, dict) else None - ) + from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs - policy: Final = policy_for_model(llm_router=llm_router, model_alias=model_alias, team_id=team_id) + policy: Final = policy_for_model( + llm_router=llm_router, + model_alias=model_alias, + team_id=team_id_from_request(data), + request_tags=_get_tags_from_request_kwargs(data), + ) if policy is None: return data _, metadata = get_or_create_metadata_bucket(data) - suppressed: Final = _active_compression_guardrail_names() - ({policy.model} if policy.model else set()) + # Markers carry a per-process token so a caller cannot suppress a guardrail by + # naming it in its own request metadata. + suppressed: Final = tuple( + marker + for guardrail in _active_compression_guardrails() + if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) + ) if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = sorted(suppressed) + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = list(suppressed) if policy.model is not None: requested = metadata.get("guardrails") @@ -157,44 +181,52 @@ async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) if snapshot is not None: - metadata[AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] = copy.deepcopy(snapshot) + _routing_messages_snapshot.set(tuple(dict(message) for message in snapshot)) return data +def _snapshot_messages() -> list[dict[str, Any]] | None: + snapshot: Final = _routing_messages_snapshot.get() + return None if snapshot is None else [dict(message) for message in snapshot] + + async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, messages: list[dict[str, Any]] | None, request_kwargs: Mapping[str, object], ) -> list[dict[str, Any]] | None: - """Messages to use for a routing decision, compressed per `policy.routing`. + """Messages to use for a routing decision, per `policy.routing`. - Returns None when there is no policy or the policy's routing side names no - compression, meaning the caller should route on whatever messages it already - has. The model call is untouched by this function either way: model-side - compression, if any, already ran as an ordinary pre-call guardrail before the - router was ever reached. + Returns None when the caller should route on whatever messages it already has. + The model call is untouched either way: model-side compression, if any, already + ran as an ordinary pre-call guardrail before the router was reached, so when the + two hops differ the routing decision reads the pre-compression snapshot rather + than what that guardrail left behind. """ - if policy is None or policy.routing is None: + if policy is None: + return None + + original: Final = _snapshot_messages() or messages + + if policy.routing is None: + # Explicitly no compression for routing. When the model side compressed, the + # messages in hand are its output, so fall back to the untouched snapshot. + return _snapshot_messages() if policy.model is not None else None + + if not original: return None from litellm.proxy.common_utils.registry_read_through import ( get_initialized_guardrail_with_read_through, ) - metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.get(metadata_key) - snapshot: Final = metadata.get(AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY) if isinstance(metadata, dict) else None - original: Final = snapshot if isinstance(snapshot, list) else messages - if not original: - return None - guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) if guardrail is None: verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return None + return original inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} # A throwaway request_data: apply_guardrail writes its stats onto this dict, not diff --git a/litellm/router.py b/litellm/router.py index 8149ec60ddc..bcb2e2aa7ff 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13039,11 +13039,19 @@ class Router: from litellm.proxy.guardrails.auto_router_compression import ( messages_for_routing, - policy_from_litellm_params, + policy_for_model, + team_id_from_request, ) - marker_params: Final = self._alias_marker_litellm_params(registered_model_name, selected_strategy.tags) - compression_policy: Final = policy_from_litellm_params(marker_params) if marker_params else None + # Resolved through the same tag-aware lookup the proxy's pre-call arming used, + # so an alias carrying several tag-scoped markers cannot suppress one marker's + # guardrail and then route under a different marker's policy. + compression_policy: Final = policy_for_model( + llm_router=self, + model_alias=registered_model_name, + team_id=team_id_from_request(request_kwargs), + request_tags=_get_tags_from_request_kwargs(request_kwargs), + ) # When both hops share the same compression, the model-side guardrail already # ran in the proxy's ordinary pre-call hook and compressed `messages` in place # (arm_pre_call armed it whether or not it is `default_on`); reuse that result @@ -13133,16 +13141,9 @@ class Router: return pre_routing_hook_response - def _alias_marker_litellm_params( + def _forwardable_alias_marker_params( self, model: str, strategy_tags: tuple[str, ...] - ) -> Mapping[str, object] | None: - """The auto-router marker deployment's own `litellm_params` for `model`, tag-scoped. - - Shared by `_forwardable_alias_marker_params` (forwarding api_base/api_key/... - gaps onto the routed deployment) and the auto-router compression policy lookup - (reading `auto_router_routing_compression`/`auto_router_model_compression`), so - both read the same marker row when an alias has more than one, tag-scoped marker. - """ + ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params for idx in self.model_name_to_deployment_indices.get(model, ()) @@ -13152,12 +13153,7 @@ class Router: tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags ) - return tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) - - def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] - ) -> tuple[tuple[str, object], ...]: - selected: Final = self._alias_marker_litellm_params(model, strategy_tags) + selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) if selected is None: return () return tuple( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1fb4299cb56..f590903cb74 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -532,7 +532,9 @@ class TestCustomGuardrailShouldRunGuardrail: data = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["headroom-default"], + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + always_on.auto_router_suppression_marker() + ], }, } @@ -550,10 +552,13 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) + other = CustomGuardrail(guardrail_name="some-other-guardrail", default_on=True) data = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["some-other-guardrail"], + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + other.auto_router_suppression_marker() + ], }, } @@ -562,6 +567,32 @@ class TestCustomGuardrailShouldRunGuardrail: is True ) + def test_should_run_guardrail_ignores_a_forged_suppression_marker(self): + """A caller controls request metadata, so a bare guardrail name there must not + switch off an always-on guardrail: only the per-process marker counts.""" + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + forged = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + "headroom-default", + "forged-token:headroom-default", + ], + }, + } + + assert ( + always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) + is True + ) + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): 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 b2e83e75768..b906e60bb86 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -4,28 +4,33 @@ Unit tests for litellm.proxy.guardrails.auto_router_compression. Covers: - policy_from_litellm_params: absent keys mean no policy; the "none" sentinel normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias -- arm_pre_call: no-op without a policy; suppresses active compression guardrails; - arms the model-side guardrail even when it isn't default_on; snapshots messages -- messages_for_routing: no-op without a policy or an unset routing side; compresses - via the named guardrail's apply_guardrail; never writes stats onto the caller's - own request_kwargs (regression for double-counted compression savings) +- policy_for_model: finds the auto-router marker deployment for an alias, and + picks the tag-scoped marker the request's tags actually match +- arm_pre_call: no-op without a policy; suppresses active compression guardrails + with a forgery-proof marker; arms the model-side guardrail even when it isn't + default_on; keeps the pre-compression snapshot out of persisted metadata +- messages_for_routing: no-op without a policy; routes on the pre-compression + snapshot when the two hops differ; compresses via the named guardrail's + apply_guardrail; never writes stats onto the caller's own request_kwargs + (regression for double-counted compression savings) """ +import json from typing import Any import pytest +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails import auto_router_compression from litellm.proxy.guardrails.auto_router_compression import ( - AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY, AutoRouterCompressionPolicy, arm_pre_call, messages_for_routing, policy_for_model, policy_from_litellm_params, ) -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -76,36 +81,61 @@ class _FakeRouter: return [d for d in self._deployments if d.get("model_name") == model_name] +def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: + return { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + **compression, + **({"tags": tags} if tags is not None else {}), + }, + } + + class TestPolicyForModel: def test_no_router_returns_none(self): - assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None) is None + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None def test_no_marker_deployment_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_picks_the_marker_whose_tags_the_request_carries(self): + """Regression: an alias with several tag-scoped markers must not suppress one + marker's guardrail and then route under a different marker's policy.""" router = _FakeRouter( [ - { - "model_name": "smart-router", - "litellm_params": { - "model": "auto_router/complexity_router", - "auto_router_routing_compression": "headroom-a", - "auto_router_model_compression": "none", - }, - } + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-us"}, tags=["us"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) + + eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + + assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) + assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) + + def test_untagged_marker_matches_any_request(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + policy = policy_for_model( + llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) @@ -186,10 +216,25 @@ class TestArmPreCall: data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} result = await arm_pre_call(data=data, llm_router=router) suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert "always-on-compression" in suppressed + assert suppressed == [always_on.auto_router_suppression_marker()] + # The bare name alone must never suppress: that is what a caller could forge. + assert "always-on-compression" not in suppressed + assert always_on.should_run_guardrail(data=result, event_type=GuardrailEventHooks.pre_call) is False finally: litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + @pytest.mark.asyncio + async def test_a_caller_cannot_suppress_a_guardrail_by_naming_it_in_metadata(self): + """Regression: request metadata is caller-controlled, so a bare guardrail name + there must not switch off a PII, content-filter, or compression guardrail.""" + guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + forged = { + "model": "smart-router", + "metadata": {AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["always-on-compression"]}, + } + + assert guardrail._suppressed_by_auto_router_compression(forged) is False + @pytest.mark.asyncio async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): router = _FakeRouter( @@ -209,43 +254,82 @@ class TestArmPreCall: assert result["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio - async def test_snapshots_original_messages(self): - router = _FakeRouter( - [ - { - "model_name": "smart-router", - "litellm_params": { - "model": "auto_router/complexity_router", - "auto_router_routing_compression": "headroom-a", - "auto_router_model_compression": "none", - }, - } - ] - ) - original_messages = [{"role": "user", "content": "hi"}] + async def test_snapshot_never_lands_in_persisted_metadata(self): + """Regression: refresh_proxy_server_request_body_snapshot copies metadata into + proxy_server_request.body, which deployments persist to spend logs. The + pre-compression snapshot holds the prompt before any masking guardrail ran, so + it must live outside anything that gets serialized.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] data = {"model": "smart-router", "messages": original_messages} + result = await arm_pre_call(data=data, llm_router=router) - snapshot = result["metadata"][AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] - assert snapshot == original_messages - assert snapshot is not original_messages # a copy, not the live reference + + assert "123-45-6789" not in json.dumps(result["metadata"]) + assert auto_router_compression._snapshot_messages() == original_messages + + @pytest.mark.asyncio + async def test_snapshot_is_a_copy_not_the_live_message_list(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + original_messages = [{"role": "user", "content": "hi"}] + + await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) + original_messages[0]["content"] = "mutated after the snapshot" + + assert auto_router_compression._snapshot_messages() == [{"role": "user", "content": "hi"}] + + @pytest.mark.asyncio + async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): + router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + await arm_pre_call(data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, + llm_router=router_with) + + router_without = _FakeRouter( + [{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + await arm_pre_call(data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, + llm_router=router_without) + + assert auto_router_compression._snapshot_messages() is None class TestMessagesForRouting: + @pytest.fixture(autouse=True) + def _clear_snapshot(self): + auto_router_compression._routing_messages_snapshot.set(None) + yield + auto_router_compression._routing_messages_snapshot.set(None) + @pytest.mark.asyncio async def test_no_policy_returns_none(self): assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_routing_side_unset_returns_none(self): - policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + async def test_routing_none_with_no_model_compression_returns_none(self): + """Nothing compressed either hop, so the caller's own messages are already right.""" + policy = AutoRouterCompressionPolicy(routing=None, model=None) assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_unknown_guardrail_name_returns_none(self): + async def test_routing_none_with_model_compression_routes_on_the_snapshot(self): + """Regression: with routing explicitly off and the model side compressed, the + messages in hand are the model-side guardrail's output. Routing asked for no + compression, so it must read the pre-compression snapshot instead.""" + original = [{"role": "user", "content": "the full original conversation"}] + auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original)) + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] + + result = await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) + + assert result == original + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_routes_on_the_uncompressed_messages(self): policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None) messages = [{"role": "user", "content": "hi"}] result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) - assert result is None + assert result == messages @pytest.mark.asyncio async def test_compresses_via_the_named_guardrail(self, registered_guardrail): @@ -256,15 +340,14 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_uses_the_snapshot_when_present(self, registered_guardrail): - policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) - snapshot = [{"role": "user", "content": "original"}] - request_kwargs = {"metadata": {AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: snapshot}} - # `messages` here stands in for whatever a model-side guardrail already + policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") + auto_router_compression._routing_messages_snapshot.set(({"role": "user", "content": "original"},)) + # `messages` here stands in for whatever the model-side guardrail already # rewrote `data["messages"]` to -- routing must ignore it and compress the # pristine snapshot instead. already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] result = await messages_for_routing( - policy=policy, messages=already_rewritten, request_kwargs=request_kwargs + policy=policy, messages=already_rewritten, request_kwargs={} ) assert result == [{"role": "user", "content": "[COMPRESSED] original"}] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cdae3b131ae..4c5813911ac 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10105,6 +10105,35 @@ class TestAutoRouterCompressionDecoupling: assert response.messages == original_messages assert registered_guardrail.call_count == 0 + @pytest.mark.asyncio + async def test_routing_none_routes_on_the_original_not_the_model_compressed_messages( + self, registered_guardrail + ): + """Regression: with routing explicitly off and the model side compressed, the + messages the router holds are the model-side guardrail's output. Routing asked + for no compression, so it has to classify on the pre-compression snapshot.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original_messages)) + model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + + try: + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) + finally: + auto_router_compression._routing_messages_snapshot.set(None) + + assert strategy.received_messages == original_messages + assert registered_guardrail.call_count == 0 + @pytest.mark.asyncio async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): """The same/different distinction exists so a shared choice does not pay for diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index a0a240f76b0..52ea7645034 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -5,8 +5,7 @@ import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Info } from "lucide-react"; import React from "react"; import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; -import { COMPRESSION_GUARDRAIL_PROVIDER } from "@/app/(dashboard)/cost-optimization/_components/helpers"; -import { NO_COMPRESSION } from "./buildAutoRouterCompression"; +import { isCompressionGuardrailProvider, NO_COMPRESSION } from "./buildAutoRouterCompression"; interface CompressionControlsProps { routing: string | undefined; @@ -29,7 +28,7 @@ const CompressionControls: React.FC = ({ }) => { const { data } = useGuardrails(); const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) - .filter((g) => (g.litellm_params?.guardrail ?? "").toString().toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER) + .filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail)) .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 49180d5b4ce..5afdcf2b15c 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -12,6 +12,13 @@ export const NO_COMPRESSION = "none"; +/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in + * litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */ +export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr"]; + +export const isCompressionGuardrailProvider = (provider: unknown): boolean => + typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase()); + export interface AutoRouterCompressionState { routing: string | undefined; sameAsRouting: boolean;