Merge 39ba956e71 into 5d66f2d8ff

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -2337,6 +2337,100 @@ async def test_custom_code_guardrail(
         )
 
 
+_GUARDRAIL_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
+
+
+def _metadata_fields(value: object) -> Mapping[str, object]:
+    try:
+        return _GUARDRAIL_METADATA_ADAPTER.validate_python(value)
+    except ValidationError:
+        return {}
+
+
+_WITHHELD_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_request_route", "agent_id", "user_api_key_auth"})
+
+_PROXY_RESOLVED_IDENTITY_FIELDS: Final[frozenset[str]] = frozenset(
+    {
+        "requester_ip_address",
+        "user_agent",
+        "headers",
+        "agent_id",
+        "actor_agent_id",
+        "target_agent_id",
+        "billing_agent_id",
+        "agent_execution_mode",
+        "verified_human_user_id",
+        "user_api_end_user_max_budget",
+    }
+)
+
+_SECRET_HEADER_MARKERS: Final = ("authorization", "cookie", "key", "token", "secret", "password", "credential")
+
+
+def _is_secret_header(name: str) -> bool:
+    lowered: Final = name.lower()
+    return any(marker in lowered for marker in _SECRET_HEADER_MARKERS)
+
+
+def _is_json_safe(value: object) -> bool:
+    try:
+        json.dumps(value)
+    except (TypeError, ValueError):
+        return False
+    return True
+
+
+def _forwarded_identity_value(key: str, value: object) -> object:
+    if key != "headers":
+        return value
+    return {name: item for name, item in _metadata_fields(value).items() if not _is_secret_header(name)}
+
+
+def _is_proxy_identity_field(key: str) -> bool:
+    """Whether a metadata field attributes the authenticated caller rather than the request.
+
+    Pre-call processing owns the whole ``user_api_key_`` namespace, the end user included, and
+    fills the agent, connection and header fields from the key and the connection right after,
+    so a caller value never reaches an LLM route. The same fields are reserved here.
+    """
+    return key in _PROXY_RESOLVED_IDENTITY_FIELDS or key == "user_api_key" or key.startswith("user_api_key_")
+
+
+def _guardrail_request_metadata(
+    caller: object, proxy: object, agent_id: str | None = None
+) -> dict[str, object]:  # mutable-ok: guardrails annotate the metadata they are handed
+    """Caller metadata with the identity fields taken from the authenticated key.
+
+    The result is handed to the guardrail as ``request_data["metadata"]``, which providers
+    write their own detection details into, so it cannot be a read-only mapping.
+
+    The caller's bucket goes through the sanitizer the LLM routes use, so this endpoint drops
+    the same fields they do, and what that leaves is then overlaid with the identity the proxy
+    resolved.
+
+    The request route is withheld from both sides. The tool policy guardrail reads the request
+    itself when a route is present, instead of the tools replayed to it, so a route here turns
+    its checks off: the caller must not be able to set one, and this endpoint's own route
+    describes the evaluation rather than the call being evaluated.
+
+    The agent is taken from ``agent_id``, the one the key authenticated with, and omitted when the
+    key is bound to none. Pre-call processing keeps the request's own agent when the key has none,
+    so the processed metadata is not a trustworthy source for it either.
+    """
+    from litellm.proxy.litellm_pre_call_utils import strip_caller_supplied_identity
+
+    sanitized: Final = dict(_metadata_fields(caller))
+    strip_caller_supplied_identity(sanitized)
+    caller_fields: Final = tuple((key, value) for key, value in sanitized.items() if not _is_proxy_identity_field(key))
+    key_identity: Final = tuple(
+        (key, _forwarded_identity_value(key, value))
+        for key, value in _metadata_fields(proxy).items()
+        if _is_proxy_identity_field(key) and key not in _WITHHELD_FIELDS and _is_json_safe(value)
+    )
+    bound_agent: Final = (("agent_id", agent_id),) if agent_id is not None else ()
+    return dict((*caller_fields, *key_identity, *bound_agent))
+
+
 def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse:
     return TestCustomCodeGuardrailResponse(
         success=False,
@@ -2503,6 +2597,13 @@ async def apply_guardrail(
 
         response_handler: Final = GUARDRAIL_RESPONSE_HANDLERS[request.response_format]
         request_data: Final = dict(response_handler.prepare_request_data(request))
+        caller_metadata: Final = request_data.get("metadata")
+        metadata: Final = _guardrail_request_metadata(
+            caller_metadata,
+            _GUARDRAIL_METADATA_ADAPTER.validate_python(data).get("metadata"),
+            agent_id=user_api_key_dict.invoked_agent_id or user_api_key_dict.agent_id,
+        )
+        request_data["metadata"] = metadata
         input_type: Final = response_handler.resolve_input_type(request, active_guardrail)
         guardrailed_inputs: Final = await active_guardrail.apply_guardrail(
             inputs={"texts": [request.text]},

litellm/proxy/litellm_pre_call_utils.py

@@ -865,6 +865,29 @@ def strip_client_guardrail_information(data: dict[str, object]) -> None:  # muta
             )
 
 
+def strip_caller_supplied_identity(
+    metadata: dict[str, object],  # mutable-ok: strips in place on the metadata bucket it is handed
+    *,
+    allow_client_message_redaction_opt_out: bool = False,
+) -> None:
+    """Drop from a caller-supplied metadata bucket the fields pre-call processing owns.
+
+    The proxy fills a family of ``user_api_key_*`` fields into one of the two metadata dicts,
+    so a caller pre-populating any of them on the other dict would have its forged values
+    surface in guardrails, spend tracking, audit logs and identity resolution. Stripping by
+    prefix covers ``user_api_key_*`` fields added later without per-key maintenance.
+
+    Shared with ``/apply_guardrail``, which accepts caller metadata the same way and so has to
+    take identity from the authenticated key rather than from the body.
+    """
+    _strip_untrusted_request_header_controls(
+        metadata.get("headers"),
+        allow_client_message_redaction_opt_out=allow_client_message_redaction_opt_out,
+    )
+    for key in [k for k in metadata if k.startswith("user_api_key_") or k in _UNTRUSTED_METADATA_CONTROL_FIELDS]:
+        metadata.pop(key, None)
+
+
 def _strip_router_reserved_metadata(
     data: dict[str, object],  # mutable-ok: strips in place on the request body the pre-call pipeline threads through
 ) -> None:
@@ -2842,14 +2865,10 @@ async def add_litellm_data_to_request(
     for _meta_key in ("metadata", "litellm_metadata"):
         _user_meta = data.get(_meta_key)
         if isinstance(_user_meta, dict):
-            _strip_untrusted_request_header_controls(
-                _user_meta.get("headers"),
-                allow_client_message_redaction_opt_out=(_allow_client_message_redaction_opt_out),
+            strip_caller_supplied_identity(
+                _user_meta,
+                allow_client_message_redaction_opt_out=_allow_client_message_redaction_opt_out,
             )
-            for _k in [
-                k for k in _user_meta if k.startswith("user_api_key_") or k in _UNTRUSTED_METADATA_CONTROL_FIELDS
-            ]:
-                _user_meta.pop(_k, None)
 
     # Strip pricing overrides AFTER the litellm_metadata string-to-dict parse
     # above, for the same reason as the user_api_key_* strip — JSON-string

tests/integration/authorization/test_guardrail_submissions.py

@@ -233,7 +233,11 @@ def _assert_guardrail_request(
     assert payload["litellm_version"] is None or isinstance(payload["litellm_version"], str), payload
     request_data: Final = object_value(payload["request_data"])
     if model is None:
-        assert request_data == {}, payload
+        assert set(request_data) == {"metadata"}, payload
+        direct_metadata: Final = object_value(request_data["metadata"])
+        assert direct_metadata.get("user_api_key_alias") == key_alias, direct_metadata
+        assert "user_api_key_auth" not in direct_metadata, direct_metadata
+        assert "authorization" not in object_value(direct_metadata.get("headers") or {}), direct_metadata
     else:
         assert set(request_data).issubset(
             {

tests/integration/observability/test_apply_guardrail_endpoint.py

@@ -116,7 +116,13 @@ def _assert_guardrail_request(request: Request, text: str, response: httpx.Respo
         "tool_calls": None,
         "model": None,
     }, response.text
-    assert payload["request_data"] == {}, response.text
+    direct_request_data: Final = payload["request_data"]
+    assert isinstance(direct_request_data, dict), response.text
+    assert set(direct_request_data) == {"metadata"}, response.text
+    direct_metadata: Final = direct_request_data["metadata"]
+    assert isinstance(direct_metadata, dict), response.text
+    assert "user_api_key_auth" not in direct_metadata, response.text
+    assert "authorization" not in (direct_metadata.get("headers") or {}), response.text
     request_headers: Final = payload["request_headers"]
     assert request_headers is None or isinstance(request_headers, dict), response.text
     if isinstance(request_headers, dict):

tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py

@@ -63,7 +63,7 @@ async def test_apply_guardrail_endpoint_returns_correct_response(
         # Verify the guardrail was called with correct parameters
         mock_guardrail.apply_guardrail.assert_called_once_with(
             inputs={"texts": ["Test text with PII"]},
-            request_data={},
+            request_data={"metadata": {}},
             input_type="request",
         )
 
@@ -198,5 +198,5 @@ async def test_apply_guardrail_endpoint_without_optional_params(mock_proxy_loggi
 
         # Verify the guardrail was called with correct parameters
         mock_guardrail.apply_guardrail.assert_called_once_with(
-            inputs={"texts": ["Test text"]}, request_data={}, input_type="request"
+            inputs={"texts": ["Test text"]}, request_data={"metadata": {}}, input_type="request"
         )

tests/unit/proxy/guardrails/test_guardrail_endpoints.py

@@ -4,7 +4,7 @@ from collections.abc import Callable
 from datetime import datetime
 from types import MappingProxyType
 from typing import Dict, Final, List, Literal, Optional
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, Mock
 
 import httpx
 import pytest
@@ -25,6 +25,7 @@ from litellm.proxy.guardrails.guardrail_endpoints import (
     TestCustomCodeGuardrailRequest,
     UpdateGuardrailRequest,
     _get_guardrails_list_response,
+    _guardrail_request_metadata,
     apply_guardrail,
     approve_guardrail_submission,
     create_guardrail,
@@ -2169,18 +2170,20 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
     }
 
 
-def _patch_apply_guardrail_env(mocker, guardrail_result):
-    mock_guardrail = mocker.Mock()
+def _patch_apply_guardrail_env(
+    mocker: MockerFixture, guardrail_result: object
+) -> Mock:
+    mock_guardrail: Final = mocker.Mock()
     mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result)
 
-    mock_registry = mocker.Mock()
+    mock_registry: Final = mocker.Mock()
     mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
     mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
 
-    mock_logging_obj = mocker.Mock()
+    mock_logging_obj: Final = mocker.Mock()
     mock_logging_obj.async_success_handler = AsyncMock()
     mock_logging_obj.model_call_details = {}
-    mock_processor = mocker.Mock()
+    mock_processor: Final = mocker.Mock()
     mock_processor.common_processing_pre_call_logic = AsyncMock(
         return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
     )
@@ -2189,7 +2192,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result):
         return_value=mock_processor,
     )
 
-    mock_proxy_logging = mocker.Mock()
+    mock_proxy_logging: Final = mocker.Mock()
     mock_proxy_logging.post_call_success_hook = AsyncMock()
     mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
     mocker.patch("litellm.proxy.proxy_server.general_settings", {})
@@ -2483,23 +2486,248 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
     )
 
 
+def test_the_authenticated_key_identity_replaces_the_caller_version() -> None:
+    metadata: Final = _guardrail_request_metadata(
+        {"forbidden_topics": ["tax"], "user_api_key_alias": "someone-else"},
+        {"user_api_key_alias": "billing-app", "user_api_key_user_id": "u-1"},
+    )
+    assert metadata == {
+        "forbidden_topics": ["tax"],
+        "user_api_key_alias": "billing-app",
+        "user_api_key_user_id": "u-1",
+    }
+
+
+def test_the_bare_key_hash_cannot_come_from_the_caller() -> None:
+    metadata: Final = _guardrail_request_metadata({"user_api_key": "spoof"}, {"user_api_key": "hashed"})
+    assert metadata == {"user_api_key": "hashed"}
+
+
+def test_the_end_user_comes_from_the_proxy_not_the_caller() -> None:
+    """LLM routes resolve the end user and drop whatever the body said, so this endpoint does too."""
+    metadata: Final = _guardrail_request_metadata(
+        {"user_api_key_end_user_id": "someone-elses-customer"},
+        {"user_api_key_end_user_id": "customer-42"},
+    )
+    assert metadata == {"user_api_key_end_user_id": "customer-42"}
+
+
+def test_an_end_user_the_proxy_did_not_resolve_is_dropped() -> None:
+    metadata: Final = _guardrail_request_metadata({"user_api_key_end_user_id": "someone-elses-customer"}, {})
+    assert metadata == {}
+
+
+def test_no_request_route_reaches_the_guardrail() -> None:
+    """A route makes the tool policy guardrail read the request instead of the replayed tools."""
+    metadata: Final = _guardrail_request_metadata(
+        {"user_api_key_request_route": "/v1/chat/completions", "topic": "tax"},
+        {"user_api_key_request_route": "/apply_guardrail"},
+    )
+    assert metadata == {"topic": "tax"}
+
+
+def test_an_unbound_agent_is_not_taken_from_the_request() -> None:
+    """Pre-call processing keeps the request's own agent_id when the key has none, so the proxy's
+    metadata already carries a caller-chosen value and reserving the field is not enough."""
+    metadata: Final = _guardrail_request_metadata(
+        {"agent_id": "someone-elses-agent"}, {"agent_id": "someone-elses-agent"}
+    )
+    assert metadata == {}
+
+
+def test_a_bound_agent_comes_from_the_authenticated_key() -> None:
+    metadata: Final = _guardrail_request_metadata(
+        {"agent_id": "someone-elses-agent"}, {"agent_id": "someone-elses-agent"}, agent_id="billing-agent"
+    )
+    assert metadata == {"agent_id": "billing-agent"}
+
+
+def test_identity_fields_outside_the_key_namespace_are_reserved_too() -> None:
+    """Pre-call processing fills every one of these from the key and the connection, so an LLM
+    route never carries a caller value for them and neither may this endpoint."""
+    metadata: Final = _guardrail_request_metadata(
+        {
+            "requester_ip_address": "203.0.113.9",
+            "user_agent": "forged/1.0",
+            "agent_id": "someone-elses-agent",
+            "actor_agent_id": "someone-elses-actor",
+            "target_agent_id": "someone-elses-target",
+            "billing_agent_id": "someone-elses-payer",
+            "agent_execution_mode": "forged",
+            "verified_human_user_id": "someone-else",
+            "topic": "tax",
+        },
+        {
+            "requester_ip_address": "198.51.100.4",
+            "user_agent": "litellm-cli/2.3",
+            "actor_agent_id": "support-agent",
+            "target_agent_id": "billing-agent",
+            "billing_agent_id": "support-agent",
+            "agent_execution_mode": "managed",
+            "verified_human_user_id": "user-7",
+        },
+        agent_id="billing-agent",
+    )
+    assert metadata == {
+        "topic": "tax",
+        "requester_ip_address": "198.51.100.4",
+        "user_agent": "litellm-cli/2.3",
+        "actor_agent_id": "support-agent",
+        "target_agent_id": "billing-agent",
+        "billing_agent_id": "support-agent",
+        "agent_execution_mode": "managed",
+        "verified_human_user_id": "user-7",
+        "agent_id": "billing-agent",
+    }
+
+
+def test_a_reserved_identity_field_the_proxy_left_empty_is_dropped() -> None:
+    """Reserving the names is what keeps a forged value out when the proxy resolved none."""
+    metadata: Final = _guardrail_request_metadata(
+        {
+            "actor_agent_id": "someone-elses-actor",
+            "target_agent_id": "someone-elses-target",
+            "billing_agent_id": "someone-elses-payer",
+            "agent_execution_mode": "forged",
+            "verified_human_user_id": "someone-else",
+            "topic": "tax",
+        },
+        {},
+    )
+    assert metadata == {"topic": "tax"}
+
+
+def test_the_end_user_budget_is_not_a_caller_field() -> None:
+    metadata: Final = _guardrail_request_metadata(
+        {"user_api_end_user_max_budget": 999.0, "topic": "tax"},
+        {"user_api_end_user_max_budget": 10.0},
+    )
+    assert metadata == {"topic": "tax", "user_api_end_user_max_budget": 10.0}
+
+
+def test_the_authenticated_key_object_is_left_behind() -> None:
+    """A guardrail vendor receives this encoded as JSON, which the auth object does not survive."""
+    metadata: Final = _guardrail_request_metadata(
+        {},
+        {
+            "user_api_key_auth": UserAPIKeyAuth(key_alias="billing-app"),
+            "user_api_key_metadata": {"tags": {"a", "b"}},
+            "user_api_key_alias": "billing-app",
+        },
+    )
+    assert metadata == {"user_api_key_alias": "billing-app"}
+
+
+def test_credential_headers_do_not_reach_the_guardrail() -> None:
+    """The header bucket carries provider credentials that the logging redaction does not name."""
+    metadata: Final = _guardrail_request_metadata(
+        {},
+        {
+            "headers": {
+                "anthropic-api-key": "sk-ant-secret",
+                "x-amz-security-token": "aws-session-token",
+                "authorization": "Bearer sk-caller",
+                "user-agent": "curl/8.7",
+            }
+        },
+    )
+    assert metadata == {"headers": {"user-agent": "curl/8.7"}}
+
+
+def test_a_guardrail_can_annotate_what_it_is_handed() -> None:
+    """OpenAI moderation and Presidio write their detection details into this mapping."""
+    metadata: Final = _guardrail_request_metadata({}, {"user_api_key_alias": "billing-app"})
+    metadata["moderation_model"] = "omni-moderation-latest"
+    assert metadata == {"user_api_key_alias": "billing-app", "moderation_model": "omni-moderation-latest"}
+
+
+class _MetadataRecordingGuardrail(CustomGuardrail):
+    """A guardrail registered the way the proxy registers one, keeping what the endpoint hands it."""
+
+    def __init__(self) -> None:
+        super().__init__(guardrail_name="identity-probe")
+        self.calls: Final[list[dict[str, object]]] = []
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[object] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.calls.append(dict(request_data))
+        return inputs
+
+
+async def _metadata_the_guardrail_was_handed(
+    caller_metadata: dict[str, object] | None, auth: UserAPIKeyAuth
+) -> dict[str, object]:
+    """Drive the real endpoint with a real registered guardrail and report what reached it.
+
+    A registered guardrail also runs as an ordinary callback on this request, so the endpoint's
+    own call is the first one it receives."""
+    import litellm
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+    from litellm.proxy.guardrails.guardrail_endpoints import router
+    from litellm.proxy.proxy_server import openai_exception_handler
+
+    guardrail: Final = _MetadataRecordingGuardrail()
+    litellm.logging_callback_manager.add_litellm_callback(guardrail)
+    body: Final[dict[str, object]] = {"guardrail_name": "identity-probe", "text": "hello"}
+    if caller_metadata is not None:
+        body["metadata"] = caller_metadata
+    app: Final = FastAPI()
+    app.include_router(router)
+    app.add_exception_handler(ProxyException, openai_exception_handler)
+    app.dependency_overrides[user_api_key_auth] = lambda: auth
+    try:
+        async with httpx.AsyncClient(
+            transport=httpx.ASGITransport(app=app), base_url="http://guardrails.test"
+        ) as client:
+            response: Final = await client.post("/guardrails/apply_guardrail", json=body)
+        assert response.status_code == 200, response.text
+    finally:
+        litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail)
+    assert guardrail.calls, "the endpoint never reached the registered guardrail"
+    return dict(guardrail.calls[0].get("metadata") or {})
+
+
 @pytest.mark.asyncio
-async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
-    """Without metadata, request_data stays empty (backward-compatible)."""
-    mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
-
-    request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
-    await apply_guardrail(
-        fastapi_request=mocker.Mock(),
-        request=request,
-        user_api_key_dict=UserAPIKeyAuth(),
+async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key() -> None:
+    """Identity fields come from the proxy's own metadata, never from the caller, so a request
+    cannot impersonate another key or probe its policy."""
+    forwarded: Final = await _metadata_the_guardrail_was_handed(
+        {
+            "forbidden_topics": ["tax"],
+            "user_api_key_alias": "someone-else",
+            "user_api_key_user_email": "victim@example.com",
+            "user_api_key": "hash-of-someone-else",
+        },
+        UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"),
     )
+    json.dumps(forwarded)
+    assert forwarded["forbidden_topics"] == ["tax"]
+    assert forwarded.get("user_api_key_alias") == "billing-app"
+    assert "victim@example.com" not in forwarded.values()
+    assert forwarded.get("user_api_key") != "hash-of-someone-else"
 
-    mock_guardrail.apply_guardrail.assert_awaited_once_with(
-        inputs={"texts": ["hello"]},
-        request_data={},
-        input_type="request",
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata() -> None:
+    """A caller that sends no metadata still arrives with the key's identity attached, because
+    pre-call processing fills it in."""
+    forwarded: Final = await _metadata_the_guardrail_was_handed(None, UserAPIKeyAuth(key_alias="billing-app"))
+    assert forwarded.get("user_api_key_alias") == "billing-app"
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_prefers_the_agent_the_key_was_invoked_as() -> None:
+    """A key bound to an agent can be invoked as another, and the invoked one is who is acting."""
+    auth: Final = UserAPIKeyAuth(agent_id="support-agent").model_copy(
+        update={"invoked_agent_id": "billing-agent"}
     )
+    forwarded: Final = await _metadata_the_guardrail_was_handed(None, auth)
+    assert forwarded.get("agent_id") == "billing-agent"
 
 
 @pytest.mark.asyncio