diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 99c41309363..7a03ac04533 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -338,15 +338,57 @@ def _logged_fields(value: object) -> Mapping[str, object]: return MappingProxyType({}) +_CALLER_CREDENTIAL_FIELDS: Final = frozenset( + { + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_web_identity_token", + "azure_ad_token", + "azure_password", + "azure_scope", + "azure_username", + "client_id", + "client_secret", + "gigachat_access_token", + "oci_fingerprint", + "oci_key", + "oci_key_file", + "oci_tenancy", + "oci_user", + "tenant_id", + "token", + "user_config", + "vertex_ai_credentials", + "vertex_credentials", + "watsonx_credentials", + "wx_credentials", + "zen_api_key", + } +) + + +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.""" + return not ( + _CALLER_CREDENTIAL_FIELDS.isdisjoint(caller_body) + and _CALLER_CREDENTIAL_FIELDS.isdisjoint(_logged_fields(caller_body.get("extra_body"))) + ) + + 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 the caller sent with the request, the caller's own Anthropic OAuth token, or the - caller's own GitHub Copilot connection.""" + 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.""" litellm_params: Final = _logged_fields(logged_call.get("litellm_params")) - credential_fields: Final = _logged_fields(litellm_params.get("proxy_server_request")).get("credential_fields") + 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) + 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 or github_copilot_user_session_from(litellm_params) is not None ) diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure.py b/tests/integration/observability/test_prometheus_provider_auth_failure.py index be5ac95200c..1825ac75651 100644 --- a/tests/integration/observability/test_prometheus_provider_auth_failure.py +++ b/tests/integration/observability/test_prometheus_provider_auth_failure.py @@ -32,6 +32,24 @@ RECOVERS: Final = "provider-rejects-once" HEALTHY: Final = "provider-healthy" AGENT_365_GUARDRAIL: Final = f"agent-365-{RUN}" CALLER_OAUTH_TOKEN: Final = f"sk-ant-oat01-caller-{RUN}" +AZURE_API_VERSION: Final = "2024-02-01" +CALLER_AZURE_CELLS: Final = ( + ("chat", False), + ("embeddings", False), + ("responses", False), + ("messages", False), + ("responses", True), +) +CALLER_CREDENTIAL_CELLS: Final = ( + ("caller-watsonx-token", "token"), + ("caller-gigachat-access-token", "gigachat_access_token"), +) +TOKEN_EXCHANGE_CELLS: Final = ("token-exchange", "token-exchange-credential") +TOKEN_EXCHANGE_SETTINGS: Final = { + "token_exchange_endpoint": "https://login.microsoftonline.com/00000000-0000-4000-8000-0000000c0b01/oauth2/v2.0/token", + "client_id": "synthetic-obo-client-id", + "client_secret": "synthetic-obo-client-secret", +} FAILURES: Final = "litellm_deployment_failure_responses_total" TOTALS: Final = "litellm_deployment_total_requests_total" @@ -106,7 +124,8 @@ class Provider: self._release: Final = release def __call__(self, request: Request) -> Reply: - model: Final = str(json.loads(request.body)["model"]) + sent: Final = json.loads(request.body) + model: Final = str(sent.get("model") or sent["model_id"]) with self._lock: self._seen[model] += 1 attempt: Final = self._seen[model] @@ -145,6 +164,49 @@ def anthropic_deployment(cell: str, api_base: str) -> dict[str, JsonValue]: } +def azure_deployment(cell: str, api_base: str, credential: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name(cell), + "litellm_params": { + "model": f"azure/{REJECTS}", + "api_base": api_base, + "api_version": AZURE_API_VERSION, + **credential, + }, + "model_info": {"id": identity(cell)}, + } + + +def caller_credential_deployment(cell: str, field: str, api_base: str) -> dict[str, JsonValue]: + provider: Final = "watsonx" if field == "token" else "gigachat" + project: Final[dict[str, JsonValue]] = {"project_id": f"project-{RUN}"} if provider == "watsonx" else {} + return { + "model_name": name(cell), + "litellm_params": { + "model": f"{provider}/{REJECTS}", + "api_base": api_base, + field: f"deployment-credential-{cell}", + **project, + }, + "model_info": {"id": identity(cell)}, + } + + +def token_exchange_deployment(cell: str) -> dict[str, JsonValue]: + settings: Final[dict[str, JsonValue]] = ( + {"litellm_credential_name": f"m365-obo-{RUN}"} if cell.endswith("-credential") else {**TOKEN_EXCHANGE_SETTINGS} + ) + return { + "model_name": name(cell), + "litellm_params": {"model": "microsoft_365_copilot/chat", **settings}, + "model_info": {"id": identity(cell)}, + } + + +def caller_azure_cell(opt_in: str, route: str, inside_extra_body: bool) -> str: + return f"{opt_in}-caller-azure-token-{route}{'-extra-body' if inside_extra_body else ''}" + + def config(directory: Path, upstream: Wire) -> Path: api_base: Final = f"{upstream.url}/v1" configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) @@ -158,6 +220,17 @@ def config(directory: Path, upstream: Wire) -> Path: deployment("recovers", RECOVERS, api_base), anthropic_deployment("anthropic-own-key", upstream.url), anthropic_deployment("caller-oauth", upstream.url), + azure_deployment("azure-own-token", upstream.url, {"azure_ad_token": "deployment-entra-token-azure-own-token"}), + *( + azure_deployment( + caller_azure_cell("deployment", route, inside_extra_body), + upstream.url, + {"configurable_clientside_auth_params": ["azure_ad_token"]}, + ) + for route, inside_extra_body in CALLER_AZURE_CELLS + ), + *(caller_credential_deployment(cell, field, upstream.url) for cell, field in CALLER_CREDENTIAL_CELLS), + *(token_exchange_deployment(cell) for cell in TOKEN_EXCHANGE_CELLS), { "model_name": name("caller-credential"), "litellm_params": { @@ -168,6 +241,10 @@ def config(directory: Path, upstream: Wire) -> Path: "model_info": {"id": identity("caller-credential")}, }, ] + configuration["credential_list"] = [ + *configuration.get("credential_list", []), + {"credential_name": f"m365-obo-{RUN}", "credential_values": {**TOKEN_EXCHANGE_SETTINGS}, "credential_info": {}}, + ] configuration["guardrails"] = [ { "guardrail_name": AGENT_365_GUARDRAIL, @@ -213,6 +290,36 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: yield Rig(owned, upstream) +def opted_in_config(directory: Path, upstream: Wire) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [ + *configuration["model_list"], + *( + azure_deployment(caller_azure_cell("proxy-wide", route, inside_extra_body), upstream.url, {}) + for route, inside_extra_body in CALLER_AZURE_CELLS + ), + ] + configuration["general_settings"]["allow_client_side_credentials"] = True + configuration["litellm_settings"]["callbacks"] = ["prometheus"] + configuration.setdefault("router_settings", {}).update({"num_retries": 0, "disable_cooldowns": True}) + path: Final = directory / "prometheus-caller-credentials-opted-in.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def opted_in(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("prometheus-caller-credentials-opted-in") + released: Final = threading.Event() + released.set() + with ( + wire_server(Provider(released)) as upstream, + gateway_from_environment() as shared, + owned_proxy(shared, directory, {}, config=opted_in_config(directory, upstream)) as owned, + ): + yield Rig(owned, upstream) + + def series(samples: Sequence[Sample], metric: str, **labels: str) -> tuple[Sample, ...]: return tuple( sample @@ -464,6 +571,85 @@ def test_a_provider_401_on_an_anthropic_oauth_token_the_caller_sent_stays_out_of }, recorded(samples, requested_model=name("caller-oauth")) +def test_a_provider_401_on_an_azure_deployments_own_ad_token_is_recorded(rig: Rig) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "azure-own-token") + assert response.status_code == 401, response.text + assert [request.headers.get("authorization") for request in rig.upstream.drain()] == [ + "Bearer deployment-entra-token-azure-own-token" + ] + assert_recorded_as_a_deployment_failure(rig, "azure-own-token", "401") + + +def caller_azure_request(cell: str, route: str, credential: dict[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]: + if route == "embeddings": + return "/v1/embeddings", {"model": name(cell), "input": cell, **credential} + if route == "responses": + return "/v1/responses", {"model": name(cell), "input": cell, **credential} + if route == "messages": + return "/v1/messages", {"model": name(cell), "max_tokens": 16, "messages": messages(cell), **credential} + return "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), **credential} + + +@pytest.mark.parametrize( + ("route", "inside_extra_body"), + CALLER_AZURE_CELLS, + ids=[f"{route}{'-extra-body' if inside_extra_body else ''}" for route, inside_extra_body in CALLER_AZURE_CELLS], +) +@pytest.mark.parametrize("opt_in", ("deployment", "proxy-wide")) +def test_a_provider_401_on_an_azure_ad_token_the_caller_sent_stays_out_of_the_deployment_metrics( + request: pytest.FixtureRequest, opt_in: str, route: str, inside_extra_body: bool +) -> None: + opted: Final[Rig] = request.getfixturevalue("rig" if opt_in == "deployment" else "opted_in") + cell: Final = caller_azure_cell(opt_in, route, inside_extra_body) + token: Final = f"caller-entra-token-{uuid.uuid4().hex}" + credential: Final[dict[str, JsonValue]] = ( + {"extra_body": {"azure_ad_token": token}} if inside_extra_body else {"azure_ad_token": token} + ) + opted.upstream.drain() + response: Final = opted.gateway.request("POST", *caller_azure_request(cell, route, credential)) + assert response.status_code == 401, response.text + assert [reached.headers.get("authorization") for reached in opted.upstream.drain()] == [ + None if inside_extra_body else f"Bearer {token}" + ] + samples: Final = until(opted.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name(cell)) + + +@pytest.mark.parametrize(("cell", "field"), CALLER_CREDENTIAL_CELLS, ids=[cell for cell, _ in CALLER_CREDENTIAL_CELLS]) +def test_a_provider_401_on_a_provider_credential_the_caller_sent_stays_out_of_the_deployment_metrics( + rig: Rig, cell: str, field: str +) -> None: + caller_credential: Final = f"caller-credential-{uuid.uuid4().hex}" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), field: caller_credential} + ) + assert response.status_code == 401, response.text + assert [reached.headers.get("authorization") for reached in rig.upstream.drain()] == [f"Bearer {caller_credential}"] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name(cell)) + + +@pytest.mark.parametrize("cell", TOKEN_EXCHANGE_CELLS) +def test_a_token_exchange_deployment_rejecting_a_caller_without_an_identity_token_stays_out_of_the_deployment_metrics( + rig: Rig, cell: str +) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, cell) + assert response.status_code == 401, response.text + assert "IdP-issued access token" in response.text, response.text + assert rig.upstream.drain() == () + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name(cell)) + + def _gated_key(scenario: Scenario, cell: str) -> str | None: if cell == "invalid-key": return f"sk-not-a-virtual-key-{uuid.uuid4().hex}" diff --git a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py index cfaaae7724d..8ae38e854f3 100644 --- a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py @@ -19,6 +19,34 @@ from litellm.llms.github_copilot.per_user_auth import GithubCopilotUserSession from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +PROVIDER_CREDENTIAL_FIELDS = [ + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_web_identity_token", + "azure_ad_token", + "azure_password", + "azure_scope", + "azure_username", + "client_id", + "client_secret", + "gigachat_access_token", + "oci_fingerprint", + "oci_key", + "oci_key_file", + "oci_tenancy", + "oci_user", + "tenant_id", + "token", + "user_config", + "vertex_ai_credentials", + "vertex_credentials", + "watsonx_credentials", + "wx_credentials", + "zen_api_key", +] + + @pytest.fixture(scope="function") def prometheus_logger(): """Create a PrometheusLogger instance for testing.""" @@ -193,6 +221,25 @@ class TestSkipMetricsValidation: ({"litellm_params": {"proxy_server_request": {"credential_fields": ("headers",)}}}, 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), + ( + { + "litellm_params": { + "proxy_server_request": {"body": {"extra_body": {"azure_ad_token": "caller-token"}}} + } + }, + True, + ), + ( + { + "litellm_params": { + "azure_ad_token": "deployment-token", + "proxy_server_request": {"body": {"model": "azure-deployment"}}, + } + }, + False, + ), + ({"litellm_params": {"proxy_server_request": {"body": {"extra_body": "azure_ad_token"}}}}, False), ], ids=[ "caller-supplied-api-key", @@ -201,6 +248,10 @@ class TestSkipMetricsValidation: "caller-supplied-headers-only", "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", ], ) def test_skip_a_provider_401_on_a_credential_the_caller_owns(self, prometheus_logger, logged_call, expected): @@ -212,6 +263,46 @@ class TestSkipMetricsValidation: kwargs = {"exception": exception, **logged_call} assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected + @pytest.mark.parametrize("field", PROVIDER_CREDENTIAL_FIELDS) + @pytest.mark.parametrize("inside_extra_body", [False, True], ids=["body", "extra-body"]) + def test_skip_a_provider_401_on_a_provider_credential_field_the_caller_sent( + self, prometheus_logger, field, inside_extra_body + ): + sent = {field: "caller-credential"} + body = {"model": "shared-deployment", **({"extra_body": sent} if inside_extra_body else sent)} + kwargs = { + "exception": litellm.AuthenticationError(message="Unauthorized", llm_provider="azure", model="gpt-4o"), + "litellm_params": {"proxy_server_request": {"body": body}}, + } + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is True + + @pytest.mark.parametrize("field", PROVIDER_CREDENTIAL_FIELDS) + def test_count_a_provider_401_on_a_provider_credential_field_the_deployment_owns(self, prometheus_logger, field): + kwargs = { + "exception": litellm.AuthenticationError(message="Unauthorized", llm_provider="azure", model="gpt-4o"), + "litellm_params": { + field: "deployment-credential", + "proxy_server_request": {"body": {"model": "shared-deployment"}}, + }, + } + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is False + + @pytest.mark.parametrize( + "token_exchange_endpoint,expected", + [("https://login.example.com/oauth2/v2.0/token", True), ("", False), (None, False)], + ids=["token-exchange-deployment", "blank-endpoint", "no-endpoint"], + ) + def test_skip_a_provider_401_on_a_token_exchange_deployment( + self, prometheus_logger, token_exchange_endpoint, expected + ): + kwargs = { + "exception": litellm.AuthenticationError( + message="Unauthorized", llm_provider="microsoft_365_copilot", model="copilot" + ), + "litellm_params": {"token_exchange_endpoint": token_exchange_endpoint, "client_secret": "deployment"}, + } + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected + def test_no_skip_for_valid_request(self, prometheus_logger): assert prometheus_logger._should_skip_metrics_for_invalid_key() is False