diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9e9553361aa..32b332f22e1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -272,6 +272,9 @@ _DESTINATION_ONLY_CALLBACK_VARS: Final[frozenset[str]] = frozenset({CAPTURE_MESS _CALLBACK_CREDENTIAL_KEYS: Final = frozenset(StandardCallbackDynamicParams.__annotations__) | frozenset( {TRUSTED_CALLBACK_VARS_FIELD} ) +_PROXY_USER_INFORMATION_HEADERS: Final = frozenset( + f"x-litellm-{key}" for key in StandardLoggingUserAPIKeyMetadata.__annotations__ +) service_logger_obj: Final = ServiceLogging() # used for tracking latency on OTEL # Bounded dedup for stale-alias warnings (FIFO eviction when over cap). _MAX_STALE_ALIAS_WARNING_KEYS: Final = 10_000 @@ -1343,17 +1346,22 @@ def _request_entries(value: object) -> tuple[object, ...]: return (value,) +def _carries_a_caller_header(headers: object) -> bool: + return any(name.lower() not in _PROXY_USER_INFORMATION_HEADERS for name in _request_fields(headers)) + + def _caller_credential_fields(data: Mapping[str, object]) -> tuple[str, ...]: """Names of the request fields through which the caller can change how the provider authenticates the call: ``api_key``, ``organization``, and ``headers``, ``extra_headers`` or a ``provider_specific_header`` entry's - ``extra_headers`` that carry any header.""" + ``extra_headers`` that carry any header besides the key and user details the proxy adds under + ``add_user_information_to_llm_headers``.""" carried: Final = { "api_key": "api_key" in data, - "extra_headers": bool(_request_fields(data.get("extra_headers"))), - "headers": bool(_request_fields(data.get("headers"))), + "extra_headers": _carries_a_caller_header(data.get("extra_headers")), + "headers": _carries_a_caller_header(data.get("headers")), "organization": "organization" in data, "provider_specific_header": any( - _request_fields(_request_fields(entry).get("extra_headers")) + _carries_a_caller_header(_request_fields(entry).get("extra_headers")) for entry in _request_entries(data.get("provider_specific_header")) ), } diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure.py b/tests/integration/observability/test_prometheus_provider_auth_failure.py index bb7c9559cc8..48bc5b2f24b 100644 --- a/tests/integration/observability/test_prometheus_provider_auth_failure.py +++ b/tests/integration/observability/test_prometheus_provider_auth_failure.py @@ -86,6 +86,7 @@ REJECTED_CELLS: Final = ( "plain-provider-specific-header", "forwarded-provider-header", "forwarded-plain-header", + "forwarded-user-information", ) GATED_CELLS: Final = ("invalid-key", "missing-key", "expired-key", "blocked-key", "no-model-access") CALLER_SCOPE_CELLS: Final = ("forwarded-gateway-key", "caller-organization", "forwarded-organization") @@ -293,8 +294,10 @@ def config(directory: Path, upstream: Wire) -> Path: name("forwarded-provider-header"), name("forwarded-plain-header"), name("forwarded-gateway-key"), + name("forwarded-user-information"), ] } + configuration["litellm_settings"]["add_user_information_to_llm_headers"] = True configuration["general_settings"]["forward_llm_provider_auth_headers"] = True configuration["general_settings"]["forward_openai_org_id"] = True configuration.setdefault("router_settings", {}).update({"num_retries": 0, "disable_cooldowns": True}) @@ -662,6 +665,19 @@ def test_a_provider_401_on_the_deployments_own_key_stays_out_when_the_proxy_forw assert_kept_out_of_the_deployment_metrics(rig, "forwarded-plain-header") +def test_a_provider_401_on_the_deployments_own_key_counts_when_the_proxy_forwards_only_its_own_user_information( + rig: Rig, +) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "forwarded-user-information") + assert response.status_code == 401, response.text + reached: Final = rig.upstream.drain() + assert [ + (request.headers["authorization"], "x-litellm-user_api_key_hash" in request.headers) for request in reached + ] == [("Bearer sk-deployment-credential-forwarded-user-information", True)] + assert_recorded_as_a_deployment_failure(rig, "forwarded-user-information", "401") + + def test_a_provider_401_on_a_gateway_key_the_proxy_forwards_stays_out_of_a_healthy_deployments_metrics( rig: Rig, ) -> None: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 776c04f1479..787c268d051 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -978,6 +978,31 @@ async def test_credential_fields_name_every_field_that_carries_a_caller_header_o general_settings: dict[str, object], expected: tuple[str, ...], ) -> None: + updated: Final = await _request_data_added_by_the_proxy(sent, request_headers, general_settings) + + assert updated["proxy_server_request"]["credential_fields"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_headers,expected", + [({}, ()), ({"x-caller-trace": "trace-1"}, ("headers",))], + ids=["user-information-only", "user-information-and-caller-header"], +) +async def test_credential_fields_ignore_the_user_information_the_proxy_adds_to_forwarded_headers( + monkeypatch: pytest.MonkeyPatch, request_headers: dict[str, str], expected: tuple[str, ...] +) -> None: + monkeypatch.setattr(litellm, "add_user_information_to_llm_headers", True) + + updated: Final = await _request_data_added_by_the_proxy({}, request_headers, FORWARD_PROVIDER_AUTH_HEADERS) + + assert updated["headers"]["x-litellm-user_api_key_user_id"] == "test-user" + assert updated["proxy_server_request"]["credential_fields"] == expected + + +async def _request_data_added_by_the_proxy( + sent: dict[str, object], request_headers: dict[str, str], general_settings: dict[str, object] +) -> dict: from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request request_mock: Final = MagicMock(spec=Request) @@ -991,7 +1016,7 @@ async def test_credential_fields_name_every_field_that_carries_a_caller_header_o request_mock.client.host = "127.0.0.1" data: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], **sent} - updated: Final = await add_litellm_data_to_request( + return await add_litellm_data_to_request( data=data, request=request_mock, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", user_id="test-user"), @@ -1000,8 +1025,6 @@ async def test_credential_fields_name_every_field_that_carries_a_caller_header_o version="test-version", ) - assert updated["proxy_server_request"]["credential_fields"] == expected - def test_initial_snapshot_refresh_clears_a_previous_guardrail_checkpoint() -> None: from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot