fix(prometheus): keep a provider 401 on caller-sent headers out of deployment metrics

A caller can override the deployment's credential with headers it sends in the
request body: headers or provider_specific_header, or an Authorization, api-key,
x-api-key, x-goog-api-key or Ocp-Apim-Subscription-Key entry in extra_headers.
A provider 401 on such a request is the caller's failure, so it stays out of the
shared deployment's failure, request, LLM-failed and state series, as on main.
This commit is contained in:
yucheng-berri 2026-10-17 01:51:17 -07:00
parent 30746c5882
commit c7289762f5
3 changed files with 74 additions and 16 deletions

View file

@ -368,25 +368,35 @@ _CALLER_CREDENTIAL_FIELDS: Final = frozenset(
)
_CALLER_CREDENTIAL_HEADERS: Final = frozenset(
{"api-key", "authorization", "ocp-apim-subscription-key", "x-api-key", "x-goog-api-key"}
)
_CALLER_TRANSPORT_FIELDS: Final = frozenset({"api_key", "headers", "provider_specific_header"})
def _sent_a_provider_credential(caller_body: Mapping[str, object]) -> bool:
"""True when the caller's request body, or its ``extra_body``, carries a provider credential field."""
"""True when the caller's request body, or its ``extra_body``, carries a provider credential field, or its
``extra_headers`` carry a credential header."""
return not (
_CALLER_CREDENTIAL_FIELDS.isdisjoint(caller_body)
and _CALLER_CREDENTIAL_FIELDS.isdisjoint(_logged_fields(caller_body.get("extra_body")))
and _CALLER_CREDENTIAL_HEADERS.isdisjoint(
header.lower() for header in _logged_fields(caller_body.get("extra_headers"))
)
)
def _used_a_caller_credential(logged_call: Mapping[str, object]) -> bool:
"""True when the provider was called with a credential that belongs to the caller rather than the
deployment: an api_key or other provider credential field the caller sent with the request, the caller's own
Anthropic OAuth token, the caller's own GitHub Copilot connection, or the caller's identity token that a
token-exchange deployment trades in on every call."""
deployment: an api_key, headers or other provider credential field the caller sent with the request, the
caller's own Anthropic OAuth token, the caller's own GitHub Copilot connection, or the caller's identity token
that a token-exchange deployment trades in on every call."""
litellm_params: Final = _logged_fields(logged_call.get("litellm_params"))
proxy_server_request: Final = _logged_fields(litellm_params.get("proxy_server_request"))
credential_fields: Final = proxy_server_request.get("credential_fields")
logged_metadata: Final = _logged_fields(_logged_fields(logged_call.get("standard_logging_object")).get("metadata"))
return (
(isinstance(credential_fields, tuple) and "api_key" in credential_fields)
(isinstance(credential_fields, tuple) and any(field in credential_fields for field in _CALLER_TRANSPORT_FIELDS))
or _sent_a_provider_credential(_logged_fields(proxy_server_request.get("body")))
or bool(litellm_params.get("token_exchange_endpoint"))
or logged_metadata.get("used_client_oauth_token") is True

View file

@ -76,7 +76,9 @@ REJECTED_CELLS: Final = (
"fallback",
"burst",
"caller-api-key",
"caller-header",
"caller-extra-headers",
"caller-headers",
"caller-provider-specific-header",
)
GATED_CELLS: Final = ("invalid-key", "missing-key", "expired-key", "blocked-key", "no-model-access")
HEALTHY_CELLS: Final = (*GATED_CELLS, "healthy", "key-rpm-limit")
@ -530,21 +532,38 @@ def anthropic_messages(gateway: Gateway, cell: str, key: str, headers: dict[str,
return gateway.request("POST", "/v1/messages", body, key=key, headers=headers)
def test_a_provider_401_on_an_authorization_header_the_caller_sends_counts_against_the_deployment(rig: Rig) -> None:
caller_header: Final = f"Bearer sk-caller-header-{uuid.uuid4().hex}"
def caller_authorization(field: str, value: str) -> JsonValue:
if field == "provider_specific_header":
return {"custom_llm_provider": "openai", "extra_headers": {"Authorization": value}}
return {"Authorization": value}
@pytest.mark.parametrize(
"cell,field",
[
("caller-extra-headers", "extra_headers"),
("caller-headers", "headers"),
("caller-provider-specific-header", "provider_specific_header"),
],
)
def test_a_provider_401_on_an_authorization_header_the_caller_sends_stays_out_of_the_deployment_metrics(
rig: Rig, cell: str, field: str
) -> None:
caller_header: Final = f"Bearer sk-{cell}-{uuid.uuid4().hex}"
rig.upstream.drain()
response: Final = rig.gateway.request(
"POST",
"/v1/chat/completions",
{
"model": name("caller-header"),
"messages": messages("caller-header"),
"extra_headers": {"Authorization": caller_header},
},
{"model": name(cell), "messages": messages(cell), field: caller_authorization(field, caller_header)},
)
assert response.status_code == 401, response.text
assert [request.headers["authorization"] for request in rig.upstream.drain()] == [caller_header]
assert_recorded_as_a_deployment_failure(rig, "caller-header", "401")
samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell))
observed: Final = recorded(samples, model_id=identity(cell))
assert values(samples, FAILURES, model_id=identity(cell)) == [], observed
assert values(samples, TOTALS, model_id=identity(cell)) == [], observed
assert values(samples, LLM_FAILED, model_id=identity(cell)) == [], observed
assert values(samples, STATE, model_id=identity(cell)) == [], observed
def test_a_provider_401_on_an_anthropic_deployments_own_key_is_recorded(rig: Rig) -> None:

View file

@ -221,7 +221,9 @@ class TestSkipMetricsValidation:
True,
),
({"standard_logging_object": {"metadata": {"used_client_oauth_token": True}}}, True),
({"litellm_params": {"proxy_server_request": {"credential_fields": ("headers",)}}}, False),
({"litellm_params": {"proxy_server_request": {"credential_fields": ("headers",)}}}, True),
({"litellm_params": {"proxy_server_request": {"credential_fields": ("provider_specific_header",)}}}, True),
({"litellm_params": {"proxy_server_request": {"credential_fields": ()}}}, False),
({"litellm_params": {"proxy_server_request": {"credential_fields": "api_key"}}}, False),
({"standard_logging_object": {"metadata": {"used_client_oauth_token": False}}}, False),
({"litellm_params": {"proxy_server_request": {"body": {"azure_ad_token": "caller-token"}}}}, True),
@ -243,18 +245,45 @@ class TestSkipMetricsValidation:
False,
),
({"litellm_params": {"proxy_server_request": {"body": {"extra_body": "azure_ad_token"}}}}, False),
(
{
"litellm_params": {
"proxy_server_request": {"body": {"extra_headers": {"Authorization": "Bearer sk-caller"}}}
}
},
True,
),
(
{"litellm_params": {"proxy_server_request": {"body": {"extra_headers": {"API-Key": "caller-key"}}}}},
True,
),
(
{
"litellm_params": {
"proxy_server_request": {"body": {"extra_headers": {"traceparent": "00-trace-span-01"}}}
}
},
False,
),
({"litellm_params": {"proxy_server_request": {"body": {"extra_headers": "authorization"}}}}, False),
],
ids=[
"caller-supplied-api-key",
"caller-copilot-session",
"caller-anthropic-oauth-token",
"caller-supplied-headers-only",
"caller-supplied-headers",
"caller-supplied-provider-specific-header",
"no-caller-credential-fields",
"malformed-credential-fields",
"deployment-anthropic-key",
"caller-supplied-azure-ad-token",
"caller-supplied-azure-ad-token-in-extra-body",
"deployment-azure-ad-token",
"malformed-extra-body",
"caller-authorization-in-extra-headers",
"caller-api-key-header-in-extra-headers",
"traceparent-only-extra-headers",
"malformed-extra-headers",
],
)
def test_skip_a_provider_401_on_a_credential_the_caller_owns(self, prometheus_logger, logged_call, expected):