Merge cb8dbda690 into 18a21b7c13

litellm/proxy/common_request_processing.py

@@ -46,6 +46,7 @@ from litellm.constants import (
     MAX_LITELLM_CALL_ID_LENGTH,
     MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
     NON_INFERENCE_CALL_TYPES,
+    REDACTED_BY_LITELLM,
     RETURN_RAW_MODEL_NAME_METADATA_KEY,
     STREAM_SSE_DATA_PREFIX,
     STREAM_SSE_KEEPALIVE_PING_BYTES,
@@ -133,6 +134,8 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar
 from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
 from litellm.proxy.native_compaction import with_proxy_compaction_executor
 from litellm.proxy.route_llm_request import (
+    RouteType,
+    UnroutableModelProbe,
     endpoint_unsupported_by_model,
     route_request,
 )
@@ -143,6 +146,7 @@ from litellm.proxy.utils import (  # noqa: F401  # legacy module exports
     ProxyLogging,
     _check_and_merge_model_level_guardrails,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     check_and_merge_model_level_guardrails,
+    without_unscanned_content,
 )
 from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success
 from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
@@ -2166,6 +2170,7 @@ class ProxyBaseLLMRequestProcessing:
         self.served_model = served_model
         self._tags_before_guardrails: frozenset[str] | None = None
         self.route_type: str | None = None
+        self._routing_probe: UnroutableModelProbe | None = None
 
     @property
     def litellm_call_id(self) -> str | None:
@@ -2459,6 +2464,7 @@ class ProxyBaseLLMRequestProcessing:
         rate_limited_model: str | None = None,
         skip_guardrails: bool = False,
         logging_auth: UserAPIKeyAuth | None = None,
+        dispatch_route_type: RouteType | None = None,
     ) -> tuple[dict, LiteLLMLoggingObj]:
         start_time: Final = datetime.now()  # start before calling guardrail hooks
 
@@ -2640,12 +2646,28 @@ class ProxyBaseLLMRequestProcessing:
         prefetch_model = self.data.get("model")
         if llm_router is not None and isinstance(prefetch_model, str):
             RoutingContext(llm_router).arm_read_prefetch(prefetch_model, self.data)
-        self.data = await proxy_logging_obj.pre_call_hook(
-            user_api_key_dict=user_api_key_dict,
-            data=self.data,
-            call_type=route_type,
-            skip_guardrails=skip_guardrails,
+        routing_probe: Final = (
+            None
+            if dispatch_route_type is None
+            else UnroutableModelProbe(
+                llm_router=llm_router,
+                user_model=user_model,
+                route_type=dispatch_route_type,
+                user_api_key_dict=user_api_key_dict,
+            )
         )
+        self._routing_probe = routing_probe
+        try:
+            self.data = await proxy_logging_obj.pre_call_hook(
+                user_api_key_dict=user_api_key_dict,
+                data=self.data,
+                call_type=route_type,
+                skip_guardrails=skip_guardrails,
+                routing_probe=routing_probe,
+            )
+        finally:
+            if routing_probe is not None and routing_probe.refusal is not None:
+                self._redact_request_its_skipped_guardrails_never_scanned(routing_probe.skipped_guardrails)
         await _enforce_guardrail_added_tag_budgets(
             data=self.data,
             tags_before_guardrails=self._tags_before_guardrails,
@@ -2681,6 +2703,22 @@ class ProxyBaseLLMRequestProcessing:
 
         return self.data, logging_obj
 
+    def _redact_request_its_skipped_guardrails_never_scanned(self, skipped_guardrails: frozenset[str]) -> None:
+        """Redact the request body, its logged snapshot and the logged messages, so failure callbacks never see
+        content a skipped masking guardrail would have rewritten, and leave the skipped guardrails out of the
+        logged ``applied_guardrails``.
+        """
+        self.data.update(without_unscanned_content(self.data))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # self.data is the untyped request dict
+        refresh_proxy_server_request_body_snapshot(self.data)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # self.data is the untyped request dict
+        logging_obj: Final = self.data.get("litellm_logging_obj")  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # self.data is the untyped request dict
+        if isinstance(logging_obj, LiteLLMLoggingObj):
+            logging_obj.update_messages([{"role": "user", "content": REDACTED_BY_LITELLM}])
+            logging_obj.model_call_details["applied_guardrails"] = [
+                name
+                for name in logging_obj.model_call_details.get("applied_guardrails") or ()
+                if name not in skipped_guardrails
+            ]
+
     async def _pre_call_with_fallbacks(
         self,
         request: Request,
@@ -2697,6 +2735,7 @@ class ProxyBaseLLMRequestProcessing:
         model: str | None,
         route_type: str,
         llm_router: Router | None,
+        dispatch_route_type: RouteType | None = None,
     ) -> tuple[dict, LiteLLMLoggingObj]:
         from litellm.proxy.common_utils.proxy_rate_limit_error import (
             PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS,
@@ -2737,6 +2776,7 @@ class ProxyBaseLLMRequestProcessing:
                 model=model,
                 route_type=route_type,
                 llm_router=llm_router,
+                dispatch_route_type=dispatch_route_type,
             )
         except ProxyRateLimitError as original_exc:
             rate_limited_data: Final = self.data
@@ -2811,6 +2851,7 @@ class ProxyBaseLLMRequestProcessing:
                             route_type=route_type,
                             llm_router=llm_router,
                             rate_limited_model=original_model,
+                            dispatch_route_type=dispatch_route_type,
                         )
                     except ProxyRateLimitError as fallback_exc:
                         if is_hard_per_model_limit(fallback_exc):
@@ -3066,8 +3107,12 @@ class ProxyBaseLLMRequestProcessing:
                 model=model,
                 route_type=route_type,
                 llm_router=llm_router,
+                dispatch_route_type=route_type,
             )
             phase_event("litellm.request.pre_call_completed")
+            refusal: Final = None if self._routing_probe is None else self._routing_probe.refusal
+            if refusal is not None:
+                raise refusal
 
         _post_call_guardrails_active: Final = self._has_post_call_guardrails()
 

tests/integration/authorization/_operator_guardrail.py

@@ -0,0 +1,29 @@
+"""A pre_call guardrail loaded from the operator's own module, the way a deployment's custom guardrail is."""
+
+import os
+from collections.abc import Mapping
+from typing import Final
+
+import httpx
+from fastapi import HTTPException
+
+from litellm.caching.dual_cache import DualCache
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.utils import CallTypesLiteral
+
+
+class OperatorGuardrail(CustomGuardrail):
+    """Posts the requested model to ``OPERATOR_GUARDRAIL_URL`` and rejects with 403 when the reply says block."""
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        cache: DualCache,
+        data: Mapping[str, object],
+        call_type: CallTypesLiteral,
+    ) -> None:
+        async with httpx.AsyncClient(trust_env=False) as client:
+            response: Final = await client.post(os.environ["OPERATOR_GUARDRAIL_URL"], json={"model": data.get("model")})
+        if response.json()["action"] == "block":
+            raise HTTPException(status_code=403, detail="blocked by the operator guardrail")

tests/integration/authorization/test_unknown_model_guardrail_skip.py

@@ -0,0 +1,243 @@
+"""LIT-8011: a model name the proxy cannot route gets the same ``400 Invalid model name``
+it gets with no guardrail, and a ``default_on`` pre_call guardrail never sends it to the
+vendor. A routable model is still scanned before the provider call, and a guardrail from the
+operator's own module still runs, along with every guardrail listed before it.
+"""
+
+import json
+import uuid
+from collections.abc import Callable, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import Gateway
+from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
+
+_MODEL_TEXT: Final = "routed reply"
+_VERDICTS: Final = {
+    "allow": Reply(body=json.dumps({"action": "allow", "category": "benign", "response_detected": {}}).encode()),
+    "block": Reply(
+        body=json.dumps(
+            {"action": "block", "category": "malicious", "prompt_detected": {"dlp": True}, "response_detected": {}}
+        ).encode()
+    ),
+    "unreachable": Reply(status=503, body=b'{"error": "scanner down"}'),
+}
+
+
+def _scanner(verdict: str) -> Callable[[Request], Reply]:
+    def reply(request: Request) -> Reply:
+        assert request.target == "/v1/scan/sync/request"
+        return _VERDICTS[verdict]
+
+    return reply
+
+
+def _provider(request: Request) -> Reply:
+    if request.method == "GET" and request.target.endswith("/models"):
+        return Reply(body=json.dumps({"object": "list", "data": []}).encode())
+    assert request.target == "/v1/chat/completions", request.target
+    return Reply(
+        body=json.dumps(
+            {
+                "id": "chatcmpl_synthetic",
+                "object": "chat.completion",
+                "created": 1700000000,
+                "model": "gpt-4.1-mini",
+                "choices": [
+                    {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": _MODEL_TEXT}}
+                ],
+                "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
+            }
+        ).encode()
+    )
+
+
+def _operator_verdict(action: str) -> Callable[[Request], Reply]:
+    def reply(request: Request) -> Reply:
+        return Reply(body=json.dumps({"action": action}).encode())
+
+    return reply
+
+
+def _airs(policy_url: str, default_on: bool = True) -> dict[str, JsonValue]:
+    return {
+        "guardrail_name": "airs-" + uuid.uuid4().hex[:8],
+        "litellm_params": {
+            "guardrail": "panw_prisma_airs",
+            "mode": "pre_call",
+            "default_on": default_on,
+            "api_base": policy_url,
+            "api_key": "synthetic-panw-key",
+            "profile_name": "synthetic-profile",
+        },
+    }
+
+
+def _operator_guardrail() -> dict[str, JsonValue]:
+    return {
+        "guardrail_name": "operator-" + uuid.uuid4().hex[:8],
+        "litellm_params": {
+            "guardrail": "integration.authorization._operator_guardrail.OperatorGuardrail",
+            "mode": "pre_call",
+            "default_on": True,
+        },
+    }
+
+
+def _config(tmp_path: Path, *guardrails: Mapping[str, JsonValue]) -> Path:
+    config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    config["guardrails"] = list(guardrails)
+    path: Final = tmp_path / "pre_call_guardrails.yaml"
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+def _requests(model: str) -> tuple[tuple[str, dict[str, JsonValue]], ...]:
+    messages: Final = [{"role": "user", "content": "my card is 4111 1111 1111 1111"}]
+    return (
+        ("/v1/chat/completions", {"model": model, "messages": messages}),
+        ("/v1/chat/completions", {"model": model, "messages": messages, "stream": True}),
+        ("/v1/responses", {"model": model, "input": "my card is 4111 1111 1111 1111"}),
+        ("/v1/messages", {"model": model, "max_tokens": 16, "messages": messages}),
+        ("/v1/embeddings", {"model": model, "input": "my card is 4111 1111 1111 1111"}),
+        ("/v1/completions", {"model": model, "prompt": "my card is 4111 1111 1111 1111"}),
+    )
+
+
+@pytest.mark.parametrize("verdict", sorted(_VERDICTS))
+def test_unknown_model_gets_routing_400_without_calling_pre_call_guardrail(
+    gateway: Gateway, tmp_path: Path, verdict: str
+) -> None:
+    model: Final = "no-such-model-" + uuid.uuid4().hex[:8]
+    with wire_server(_scanner(verdict)) as policy:
+        with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url))) as candidate:
+            for target, body in _requests(model):
+                control = gateway.request("POST", target, body)
+                response = candidate.request("POST", target, body)
+
+                assert control.status_code == 400, control.text
+                assert "Invalid model name" in control.text
+                assert (response.status_code, response.json()) == (control.status_code, control.json()), target
+            assert policy.drain() == ()
+
+
+def test_routable_model_is_still_scanned_before_the_provider_call(gateway: Gateway, tmp_path: Path) -> None:
+    with wire_server(_scanner("allow")) as policy, wire_server(_provider) as upstream:
+        with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url))) as candidate:
+            with candidate.scenario() as scenario:
+                model: Final = scenario.model(
+                    model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key"
+                )
+                response: Final = candidate.request("POST", "/v1/chat/completions", dict(_requests(model)[0][1]))
+
+                assert response.status_code == 200, response.text
+                assert response.json()["choices"][0]["message"]["content"] == _MODEL_TEXT
+                scans: Final = policy.drain()
+                assert len(scans) == 1
+                assert "4111" in json.loads(scans[0].body)["contents"][0]["prompt"]
+
+
+def test_routable_model_block_verdict_still_blocks(gateway: Gateway, tmp_path: Path) -> None:
+    with wire_server(_scanner("block")) as policy, wire_server(_provider) as upstream:
+        with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url))) as candidate:
+            with candidate.scenario() as scenario:
+                model: Final = scenario.model(
+                    model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key"
+                )
+                response: Final = candidate.request("POST", "/v1/chat/completions", dict(_requests(model)[0][1]))
+
+                assert response.status_code == 400, response.text
+                assert "Invalid model name" not in response.text
+                assert len(policy.drain()) == 1
+                assert [request.target for request in upstream.drain() if request.method == "POST"] == []
+
+
+def test_unknown_model_burst_while_the_scanner_is_down_never_reaches_it(gateway: Gateway, tmp_path: Path) -> None:
+    model: Final = "no-such-model-" + uuid.uuid4().hex[:8]
+    unknown: Final = _requests(model)
+    controls: Final = tuple(gateway.request("POST", target, body) for target, body in unknown)
+    with wire_server(_scanner("unreachable")) as policy, wire_server(_provider) as upstream:
+        with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, _airs(policy.url)), workers=2) as candidate:
+            with candidate.scenario() as scenario:
+                routable: Final = scenario.model(
+                    model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key"
+                )
+                known: Final = (
+                    "/v1/chat/completions",
+                    {"model": routable, "messages": [{"role": "user", "content": "routable burst prompt"}]},
+                )
+                burst: Final = unknown * 6 + (known,) * 6
+                with ThreadPoolExecutor(max_workers=16) as pool:
+                    responses: Final = tuple(pool.map(lambda item: candidate.request("POST", *item), burst))
+
+                for index, response in enumerate(responses[: len(unknown) * 6]):
+                    control = controls[index % len(unknown)]
+                    assert (response.status_code, response.json()) == (control.status_code, control.json()), index
+                known_statuses: Final = {response.status_code for response in responses[len(unknown) * 6 :]}
+                assert 200 not in known_statuses, known_statuses
+                scanned: Final = [json.loads(scan.body)["contents"][0]["prompt"] for scan in policy.drain()]
+                assert scanned and all("routable burst prompt" in prompt for prompt in scanned), scanned
+                assert [request.target for request in upstream.drain() if request.method == "POST"] == []
+
+
+def test_opt_in_guardrail_is_named_in_the_header_only_when_it_ran(gateway: Gateway, tmp_path: Path) -> None:
+    model: Final = "no-such-model-" + uuid.uuid4().hex[:8]
+    with wire_server(_scanner("allow")) as policy, wire_server(_provider) as upstream:
+        guardrail: Final = _airs(policy.url, default_on=False)
+        opt_in: Final = {"guardrails": [guardrail["guardrail_name"]]}
+        with owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, guardrail)) as candidate:
+            with candidate.scenario() as scenario:
+                routable: Final = scenario.model(
+                    model="openai/gpt-4.1-mini", api_base=upstream.url + "/v1", api_key="synthetic-key"
+                )
+                ran: Final = candidate.request(
+                    "POST", "/v1/chat/completions", {**_requests(routable)[0][1], "metadata": opt_in}
+                )
+                assert ran.status_code == 200, ran.text
+                assert ran.headers["x-litellm-applied-guardrails"] == guardrail["guardrail_name"]
+                assert len(policy.drain()) == 1
+
+                for target, body in _requests(model):
+                    refused = candidate.request("POST", target, {**body, "metadata": opt_in})
+
+                    assert refused.status_code == 400, refused.text
+                    assert "Invalid model name" in refused.text
+                    assert "x-litellm-applied-guardrails" not in refused.headers, target
+                assert policy.drain() == ()
+
+
+def test_operator_guardrail_and_the_guardrail_before_it_still_run_for_an_unknown_model(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    model: Final = "no-such-model-" + uuid.uuid4().hex[:8]
+    with wire_server(_scanner("allow")) as policy, wire_server(_operator_verdict("block")) as operator:
+        config: Final = _config(tmp_path, _airs(policy.url), _operator_guardrail())
+        with owned_proxy(gateway, tmp_path, {"OPERATOR_GUARDRAIL_URL": operator.url}, config=config) as candidate:
+            response: Final = candidate.request("POST", "/v1/chat/completions", dict(_requests(model)[0][1]))
+
+            assert response.status_code == 403, response.text
+            assert "blocked by the operator guardrail" in response.text
+            assert len(policy.drain()) == 1
+            assert [json.loads(request.body) for request in operator.drain()] == [{"model": model}]
+
+
+def test_built_in_guardrail_after_an_allowing_operator_guardrail_is_skipped_for_an_unknown_model(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    model: Final = "no-such-model-" + uuid.uuid4().hex[:8]
+    body: Final = dict(_requests(model)[0][1])
+    control: Final = gateway.request("POST", "/v1/chat/completions", body)
+    with wire_server(_scanner("block")) as policy, wire_server(_operator_verdict("allow")) as operator:
+        config: Final = _config(tmp_path, _operator_guardrail(), _airs(policy.url))
+        with owned_proxy(gateway, tmp_path, {"OPERATOR_GUARDRAIL_URL": operator.url}, config=config) as candidate:
+            response: Final = candidate.request("POST", "/v1/chat/completions", body)
+
+            assert (response.status_code, response.json()) == (control.status_code, control.json())
+            assert [json.loads(request.body) for request in operator.drain()] == [{"model": model}]
+            assert policy.drain() == ()

tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py

@@ -743,7 +743,7 @@ async def test_message_send_reports_an_unresolvable_entra_credential_as_internal
 
     mock_proxy_logging = MagicMock()
     mock_proxy_logging.pre_call_hook = AsyncMock(
-        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data
     )
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
     downstream = AsyncMock()
@@ -1097,7 +1097,7 @@ async def test_subscribe_to_task_calls_pre_call_hook():
 
     mock_proxy_logging: Final = MagicMock(spec=ProxyLogging)
     mock_proxy_logging.pre_call_hook = AsyncMock(
-        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data
     )
     mock_proxy_logging.async_post_call_streaming_iterator_hook = _passthrough_iterator
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
@@ -1226,7 +1226,7 @@ async def test_task_method_failure_hook_uses_enriched_request_data():
 
     mock_proxy_logging = MagicMock()
     mock_proxy_logging.pre_call_hook = AsyncMock(
-        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data
     )
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
 
@@ -1293,7 +1293,7 @@ async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400()
 
     mock_proxy_logging = MagicMock()
     mock_proxy_logging.pre_call_hook = AsyncMock(
-        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None: data
     )
     mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
 

tests/unit/proxy/mcp_server/test_rest_endpoints.py

@@ -3072,7 +3072,7 @@ class TestCallToolRestAPI:
 
         pre_call_finished_at = {}
 
-        async def slow_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
+        async def slow_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None):
             await asyncio.sleep(0.05)
             pre_call_finished_at["value"] = datetime.now()
             return data
@@ -3330,10 +3330,12 @@ class TestCallToolRestAPI:
                 message="Content blocked", model="mcp-tool-call", request_data={}, guardrail_name="block-all"
             )
 
-        async def passthrough_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
+        async def passthrough_pre_call_hook(
+            user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None
+        ):
             return data
 
-        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
+        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None):
             raise guardrail_error
 
         async def fake_execute_mcp_tool(**kwargs):
@@ -3534,7 +3536,7 @@ class TestCallToolRestAPI:
         async def fake_add_litellm_data_to_request(**kwargs):
             return kwargs.get("data", {})
 
-        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False):
+        async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None):
             raise guardrail_error
 
         failure_logging = AsyncMock()

tests/unit/proxy/test_model_level_guardrails.py

@@ -601,7 +601,7 @@ async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook():
 
     captured_pre_call_guardrails: list = []
 
-    async def fake_pre_call_hook(*, user_api_key_dict, data, call_type, skip_guardrails=False):
+    async def fake_pre_call_hook(*, user_api_key_dict, data, call_type, skip_guardrails=False, routing_probe=None):
         # Snapshot the list rather than the dict: metadata is shared by
         # reference, so a merge that happens after this point would otherwise
         # show up here retroactively and the assertion would pass either way.