From b5a51e72cd2f286c1b335414e8a2b91815424810 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 4 Aug 2026 13:16:34 -0700 Subject: [PATCH] feat(otel): resolve a request's trace destinations from its identity --- litellm/integrations/otel/plumbing/context.py | 19 +- litellm/proxy/auth/user_api_key_auth.py | 51 + litellm/proxy/litellm_pre_call_utils.py | 168 ++- litellm/types/utils.py | 16 + ...test_initialize_dynamic_callback_params.py | 53 + .../proxy/auth/test_user_api_key_auth.py | 80 ++ .../proxy/test_litellm_pre_call_utils.py | 988 +++++++++++------- 7 files changed, 983 insertions(+), 392 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 19b36c0b967..5290d87ae9d 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -2,7 +2,7 @@ from collections.abc import Mapping from contextvars import ContextVar, Token -from typing import Final +from typing import TYPE_CHECKING, Final from opentelemetry import baggage from opentelemetry.context import Context, get_current @@ -18,6 +18,9 @@ from opentelemetry.trace.propagation.tracecontext import ( TraceContextTextMapPropagator, ) +if TYPE_CHECKING: + from litellm.integrations.otel.model.destination import OtelDestination + _PROPAGATOR: Final = TraceContextTextMapPropagator() # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the @@ -37,6 +40,20 @@ _PROPAGATOR: Final = TraceContextTextMapPropagator() # request task, so there is nothing to leak. _request_root_span: Final["ContextVar[Span | None]"] = ContextVar("litellm_otel_request_root_span", default=None) +_request_destinations: 'ContextVar[tuple["OtelDestination", ...]]' = ContextVar( + "litellm_otel_request_destinations", default=() +) + + +def set_request_destinations(destinations: 'tuple["OtelDestination", ...]') -> None: + """Anchor the admin-resolved destinations for this request.""" + _request_destinations.set(tuple(destinations)) + + +def request_destinations() -> 'tuple["OtelDestination", ...]': + """Destinations the request fans out to, or empty when none were resolved.""" + return _request_destinations.get() + def set_request_root_span(span: Span) -> None: """Anchor the request's root (server) span for explicit child parenting. diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4baa7b99a4f..03eee8b4c78 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1009,6 +1009,51 @@ async def _resolve_jwt_to_virtual_key( return None +async def _hoist_request_destinations(request: Request, user_api_key_dict: UserAPIKeyAuth) -> None: + """Resolve admin-owned OTEL destinations for this request and anchor them. + + Runs after the auth builder, while we are still inside the request task, so + the ``ContextVar`` is visible to every ``SpanProcessor.on_end`` that fires + for spans this request opens. Stashes the same list on ``request.state`` so + ``_apply_admin_logging_exporters`` can reuse it without a second DB pass. + + Best-effort: a resolver failure must not break the request. The contextvar + is left at its default (empty tuple), so the fan-out processor no-ops. Idempotent: + it fires early in the builder and again as an outer catch-all; the second call + skips once ``request.state`` holds the result (a failed first call leaves it unset). + """ + if getattr(getattr(request, "state", None), "otel_destinations", None) is not None: + return + try: + from litellm.integrations.otel.model.destination import OtelDestination + from litellm.integrations.otel.plumbing.context import ( + set_request_destinations, + ) + from litellm.proxy.litellm_pre_call_utils import ( + _resolve_logging_exporters, + ) + + destinations_raw, _backends = await _resolve_logging_exporters(user_api_key_dict) + destinations = tuple( + OtelDestination( + callback_name=item.get("callback_name"), + endpoint=item.get("endpoint", ""), + headers=item.get("headers") or {}, + resource_attributes=item.get("resource_attributes") or {}, + protocol=item.get("protocol"), + ) + for item in destinations_raw + if isinstance(item, dict) and item.get("endpoint") + ) + set_request_destinations(destinations) + try: + request.state.otel_destinations = destinations_raw + except Exception: # noqa: BLE001 # request.state mirror is best-effort; the ContextVar is the source of truth + pass + except Exception as exc: # noqa: BLE001 # destination hoist is best-effort telemetry setup; it must never fail auth + verbose_proxy_logger.debug("OTel V2: hoist destination resolution failed: %s", exc) + + def _ensure_parent_otel_span_on_request_state(request: Request) -> None: """Idempotently create the OTEL SERVER span and stash it on ``request.state.parent_otel_span``. Safe to call multiple times. @@ -1405,6 +1450,8 @@ async def _user_api_key_auth_builder( valid_token = auto_registered api_key = valid_token.token or "" + await _hoist_request_destinations(request, valid_token) + # Check if model has zero cost - if so, skip all budget checks model = _get_model_from_request_context( request_data=request_data, @@ -1742,6 +1789,8 @@ async def _user_api_key_auth_builder( user_obj: LiteLLM_UserTable | None = None valid_token_dict: dict = {} if valid_token is not None: + valid_token.parent_otel_span = parent_otel_span + await _hoist_request_destinations(request, valid_token) # Got Valid Token from Cache, DB # Run checks for # 1. If token can call model @@ -2646,6 +2695,8 @@ async def user_api_key_auth( raise user_api_key_auth_obj.budget_reservation = None + await _hoist_request_destinations(request, user_api_key_auth_obj) + # A body that never parsed is authenticated (so the trace carries identity # and this ``auth`` span) but not authorized: there is no model to check it # against, and budget reservation would increment live spend counters that diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index f83061a15ce..577f8000b60 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4,7 +4,7 @@ import json import re import time from collections import OrderedDict -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -148,8 +148,11 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: + from litellm.integrations.otel.model.destination import OtelDestination + from litellm.models.credentials import CredentialItem from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.types.proxy.policy_engine import PolicyMatchContext + from litellm.types.utils import OtelDestinationParams ProxyConfig = _ProxyConfig else: @@ -709,6 +712,160 @@ class KeyAndTeamLoggingSettings: return None +async def _effective_org_id(user_api_key_dict: UserAPIKeyAuth) -> str | None: + """The org this request belongs to, falling back to the team's org when the token + carries none. Team keys frequently have no ``org_id`` on the token, so without this + an org-scoped destination would be invisible at request time even though the write + gate (which loads the team) accepted it. Mirrors the fallback in ``_check_org_budget``. + """ + if user_api_key_dict.org_id is not None: + return user_api_key_dict.org_id + team_id = user_api_key_dict.team_id + if team_id is None: + return None + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_checks import get_team_object + + if proxy_server.prisma_client is None: + return None + try: + team_obj = await get_team_object( + team_id=team_id, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + parent_otel_span=getattr(user_api_key_dict, "parent_otel_span", None), + ) + except HTTPException: + return None + return getattr(team_obj, "organization_id", None) + + +async def _resolve_logging_exporters( + user_api_key_dict: UserAPIKeyAuth, +) -> "tuple[tuple[OtelDestinationParams, ...], tuple[str, ...]]": + """Resolve the destinations this request fans out to, as (destinations, backends). + + ``credential_info.access`` is the sole routing determinant: a destination is + selected when its ``access`` grants the caller's team/org. Empty access grants no + one, so an empty-access destination never fires (proxy-wide requires + ``access.global``). Each survivor is built via ``destination_for_credential`` and deduped on + (endpoint, headers, resource attributes). Returns ([], []) when nothing is selected + (default-deny). + """ + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy.management_endpoints.logging_exporter_access import ( + access_grants, + destination_for_credential, + identity_scope, + parse_credential_info, + ) + + # Admin-owned destinations are an OTEL v2 feature; the LITELLM_OTEL_V2 flag is the + # sole activation gate. With the flag off, registering a destination resolves to + # nothing (no backend is activated for the request) until the admin sets the flag. + if not is_otel_v2_enabled(): + return (), () + + if not any( + (info := parse_credential_info(credential.credential_info)) is not None and info.credential_type == "logging" + for credential in litellm.credential_list + ): + return (), () + + team_id = user_api_key_dict.team_id + org_id = await _effective_org_id(user_api_key_dict) + team_ids, org_ids = identity_scope(team_id, org_id) + + def _selected(credential: "CredentialItem") -> bool: + info = parse_credential_info(credential.credential_info) + if info is None or info.credential_type != "logging": + return False + return access_grants(info.access, team_ids, org_ids) + + built = tuple( + result + for credential in litellm.credential_list + if _selected(credential) + if (result := destination_for_credential(credential)) is not None + ) + deduped = { + ( + destination.endpoint, + tuple(sorted(destination.headers.items())), + tuple(sorted(destination.resource_attributes.items())), + ): ( + backend, + destination, + ) + for backend, destination in built + } + destinations: tuple[OtelDestinationParams, ...] = tuple( + { + "callback_name": backend, + "endpoint": destination.endpoint, + "headers": destination.headers, + "resource_attributes": destination.resource_attributes, + "protocol": destination.protocol, + } + for backend, destination in deduped.values() + ) + backends = tuple(dict.fromkeys(backend for backend, _ in deduped.values())) + return destinations, backends + + +def _request_destination_from_raw(item: object) -> "OtelDestination | None": + from litellm.integrations.otel.model.destination import OtelDestination + + if isinstance(item, OtelDestination): + return item + if not isinstance(item, dict) or not item.get("endpoint"): + return None + try: + return OtelDestination.model_validate(item) + except PydanticValidationError: + return None + + +def _set_request_otel_destinations(destinations: Sequence[object]) -> None: + from litellm.integrations.otel.plumbing.context import set_request_destinations + + set_request_destinations( + tuple(destination for item in destinations if (destination := _request_destination_from_raw(item)) is not None) + ) + + +async def _apply_admin_logging_exporters( + user_api_key_dict: UserAPIKeyAuth, + cached_destinations: "Sequence[object] | None" = None, +) -> None: + """Anchor the resolved fan-out destinations on the request context. + + The destinations are set on a server-only ContextVar (never on ``data``), so + they are neither request-shaped nor reachable by the provider body; the OTEL v2 + router and the fan-out processor both read them from that ContextVar. Default-deny + means an identity no destination's access grants gets no per-tenant destination here. + + ``cached_destinations`` -- when ``user_api_key_auth`` already resolved the + destinations on this request (the FastAPI path), reuse the result instead of + running the resolver a second time. The SDK path passes ``None`` and the + resolver runs here. + + An empty resolution is published too, matching ``_hoist_request_destinations``. + Returning early instead would leave a previous message's destinations standing on + a ContextVar this request never overwrites: a stateful MCP session runs every + message on the task its ``initialize`` spawned, so a revoked grant would keep + exporting for the life of that session. + """ + if cached_destinations is not None: + destinations = tuple(cached_destinations) + else: + try: + destinations, _backends = await _resolve_logging_exporters(user_api_key_dict) + except Exception: # noqa: BLE001 # best-effort telemetry setup must never break the request + return + _set_request_otel_destinations(destinations) + + def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> TeamCallbackMetadata | None: @@ -2004,6 +2161,11 @@ async def add_litellm_data_to_request( ) # Team Callbacks controls + data.pop("otel_destinations", None) + for _carrier_key in (_metadata_variable_name, "litellm_metadata"): + carrier = data.get(_carrier_key) + if isinstance(carrier, dict): + carrier.pop("otel_destinations", None) callback_settings_obj: Final = _get_dynamic_logging_metadata( user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) @@ -2012,13 +2174,15 @@ async def add_litellm_data_to_request( data["failure_callback"] = callback_settings_obj.failure_callback if callback_settings_obj.callback_vars is not None: - # unpack callback_vars in data for k, v in callback_settings_obj.callback_vars.items(): data[k] = v # Callbacks that must not honour request-supplied credentials read this # proxy-owned field instead of the raw request kwargs. data[TRUSTED_CALLBACK_VARS_FIELD] = callback_settings_obj.callback_vars + cached = getattr(getattr(request, "state", None), "otel_destinations", None) + await _apply_admin_logging_exporters(user_api_key_dict, cached_destinations=cached) + # Add disabled callbacks from key metadata if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: disabled_callbacks: Final = user_api_key_dict.metadata["litellm_disabled_callbacks"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 354857f8d72..05311ce940e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3199,6 +3199,22 @@ OPENAI_RESPONSE_HEADERS: Final = [ ] +class OtelDestinationParams(TypedDict, total=False): + """A resolved, admin-owned OTLP destination carried server-side only. + + Populated by the proxy from the exporters assigned to a request's identity + chain; never read from a request body or metadata. The v2 logger validates and + exports through it. ``callback_name`` is the OTEL backend this destination + belongs to, so fan-out routes each destination to the right backend's logger. + """ + + callback_name: str + endpoint: str + headers: Mapping[str, str] + resource_attributes: Mapping[str, str] + protocol: str | None + + class StandardCallbackDynamicParams(TypedDict, total=False): # Langfuse dynamic params langfuse_public_key: str | None diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py index 0dca4f3a1b1..42d06729e1c 100644 --- a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -80,6 +80,59 @@ def test_resolves_plain_values_from_metadata(): assert params.get("langfuse_host") == "https://test.langfuse.com" +def test_otel_destinations_not_carried_on_dynamic_params(): + """OTEL destination routing no longer travels through dynamic params or request + data at all (Y6): the admin-resolved destinations are anchored on a server-only + ContextVar and the v2 router reads them from there. Even a value under the internal + ``litellm_metadata`` key is NOT surfaced on the dynamic params, so the + request-carried carrier stays removed. Re-introducing the carrier fails this.""" + destinations = [ + { + "callback_name": "langfuse_otel", + "endpoint": "https://cloud.langfuse.com/api/public/otel", + "headers": {"Authorization": "Basic ADMIN"}, + } + ] + + params = initialize_standard_callback_dynamic_params( + {"litellm_metadata": {"otel_destinations": destinations}} + ) + + assert params.get("otel_destinations") is None + + +def test_otel_destinations_top_level_kwarg_is_ignored(): + """A top-level ``otel_destinations`` kwarg is intentionally NOT read. The proxy + stashes admin-resolved destinations under ``litellm_metadata`` to keep unknown + keys out of the body forwarded to the provider; reading the top-level key would + re-open that surface and is therefore ignored.""" + params = initialize_standard_callback_dynamic_params( + {"otel_destinations": [{"callback_name": "langfuse_otel"}]} + ) + assert params.get("otel_destinations") is None + + +def test_otel_destinations_never_read_from_request_metadata(): + """A request body/metadata must not be able to inject OTEL destinations: + otel_destinations is deliberately absent from the request-read whitelist, so a + value nested in metadata is ignored. Guards the trust boundary.""" + kwargs = { + "metadata": { + "otel_destinations": [ + { + "callback_name": "langfuse_otel", + "endpoint": "https://attacker.example/api/public/otel", + "headers": {"Authorization": "Basic ATTACKER"}, + } + ] + } + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("otel_destinations") is None + + def test_litellm_params_metadata_overrides_metadata(): kwargs = { "metadata": { diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 60d9689dc0b..0e11885c5f6 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5012,6 +5012,63 @@ async def test_builder_succeeds_when_db_lookup_returns_valid_token(): mock_return.assert_awaited_once() +@pytest.mark.asyncio +async def test_builder_hoists_destinations_before_post_lookup_auth_checks(): + valid_token = UserAPIKeyAuth(api_key="sk-db-lookup-test", token="hashed-valid") + get_key_object = AsyncMock(return_value=valid_token) + + async def _assert_hoisted_first(*args, **kwargs): + assert mock_hoist.await_count == 1 + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._hoist_request_destinations", + new_callable=AsyncMock, + ) as mock_hoist, + patch( + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + side_effect=_assert_hoisted_first, + ) as mock_enforce, + ): + result = await _run_builder_with_key_lookup(get_key_object) + + assert result is valid_token + mock_hoist.assert_awaited_once() + mock_enforce.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_hoist_destinations_resolver_failure_never_breaks_auth(): + """Destination resolution is best-effort telemetry setup: if the resolver raises, + _hoist_request_destinations must swallow it and leave the ContextVar at its empty + default so auth proceeds and the fan-out processor no-ops. A raise here would take + down every request.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.integrations.otel.plumbing.context import request_destinations + from litellm.proxy.auth.user_api_key_auth import _hoist_request_destinations + + request = MagicMock() + request.state = MagicMock() + valid_token = UserAPIKeyAuth(api_key="sk-x", token="hashed") + + with patch( + "litellm.proxy.litellm_pre_call_utils._resolve_logging_exporters", + new_callable=AsyncMock, + side_effect=RuntimeError("resolver blew up"), + ): + # must not raise + await _hoist_request_destinations(request, valid_token) + + assert request_destinations() == () + + def _mint_cli_session_token(monkeypatch, *, user_id="cli-admin"): """Mint a CLI session token for a PROXY_ADMIN user so auth resolves on the admin early-return path (no prisma/common_checks needed).""" @@ -5967,3 +6024,26 @@ async def test_unlicensed_jwt_auth_is_forbidden_not_unauthorized(): assert error.code == "403" assert "enterprise" in error.message.lower() +@pytest.mark.asyncio +async def test_hoist_request_destinations_idempotent(monkeypatch): + """The hoist fires early in the auth builder and again as an outer catch-all; the + second call must not re-run the resolver once request.state holds the result.""" + import litellm.proxy.litellm_pre_call_utils as pcu + from litellm.proxy.auth.user_api_key_auth import _hoist_request_destinations + + calls = {"n": 0} + + async def fake_resolve(_uapk): + calls["n"] += 1 + return ((), ()) + + monkeypatch.setattr(pcu, "_resolve_logging_exporters", fake_resolve) + request = MagicMock() + request.state = SimpleNamespace() + uapk = UserAPIKeyAuth(api_key="x", token="x") + + await _hoist_request_destinations(request, uapk) + await _hoist_request_destinations(request, uapk) + + assert calls["n"] == 1 + assert request.state.otel_destinations == () diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index e31058f402e..cac93b88cd0 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -34,9 +34,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( ) from litellm.types.utils import CredentialItem -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path def test_check_if_token_is_service_account(): @@ -44,9 +42,7 @@ def test_check_if_token_is_service_account(): Test that only keys with `service_account_id` in metadata are considered service accounts """ # Test case 1: Service account token - service_account_token = UserAPIKeyAuth( - api_key="test-key", metadata={"service_account_id": "test-service-account"} - ) + service_account_token = UserAPIKeyAuth(api_key="test-key", metadata={"service_account_id": "test-service-account"}) assert check_if_token_is_service_account(service_account_token) == True # Test case 2: Regular user token @@ -54,9 +50,7 @@ def test_check_if_token_is_service_account(): assert check_if_token_is_service_account(regular_token) == False # Test case 3: Token with other metadata - other_metadata_token = UserAPIKeyAuth( - api_key="test-key", metadata={"user_id": "test-user"} - ) + other_metadata_token = UserAPIKeyAuth(api_key="test-key", metadata={"user_id": "test-user"}) assert check_if_token_is_service_account(other_metadata_token) == False @@ -103,15 +97,11 @@ class TestGetMetadataVariableName: def test_returns_litellm_metadata_for_bedrock_invoke(self): # GH#30629: bedrock passthrough must use litellm_metadata # to prevent key-level tags from leaking into provider body - request = self._make_request( - "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke" - ) + request = self._make_request("/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke") assert _get_metadata_variable_name(request) == "litellm_metadata" def test_returns_litellm_metadata_for_bedrock_converse(self): - request = self._make_request( - "/bedrock/model/us.anthropic.claude-sonnet-4-6/converse" - ) + request = self._make_request("/bedrock/model/us.anthropic.claude-sonnet-4-6/converse") assert _get_metadata_variable_name(request) == "litellm_metadata" @@ -119,9 +109,7 @@ def test_get_enforced_params_for_service_account_settings(): """ Test that service account enforced params are only added to service account keys """ - service_account_token = UserAPIKeyAuth( - api_key="test-key", metadata={"service_account_id": "test-service-account"} - ) + service_account_token = UserAPIKeyAuth(api_key="test-key", metadata={"service_account_id": "test-service-account"}) general_settings_with_service_account_settings = { "service_account_settings": {"enforced_params": ["metadata.service"]}, } @@ -131,9 +119,7 @@ def test_get_enforced_params_for_service_account_settings(): ) assert result == ["metadata.service"] - regular_token = UserAPIKeyAuth( - api_key="test-key", metadata={"enforced_params": ["user"]} - ) + regular_token = UserAPIKeyAuth(api_key="test-key", metadata={"enforced_params": ["user"]}) result = _get_enforced_params( general_settings=general_settings_with_service_account_settings, user_api_key_dict=regular_token, @@ -146,9 +132,7 @@ def test_get_enforced_params_for_service_account_settings(): [ ( {"enforced_params": ["param1", "param2"]}, - UserAPIKeyAuth( - api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" - ), + UserAPIKeyAuth(api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"), ["param1", "param2"], ), ( @@ -174,9 +158,7 @@ def test_get_enforced_params_for_service_account_settings(): ), ], ) -def test_get_enforced_params( - general_settings, user_api_key_dict, expected_enforced_params -): +def test_get_enforced_params(general_settings, user_api_key_dict, expected_enforced_params): from litellm.proxy.litellm_pre_call_utils import _get_enforced_params enforced_params = _get_enforced_params(general_settings, user_api_key_dict) @@ -292,9 +274,7 @@ async def test_add_litellm_data_to_request_strips_admin_injection_slots(): populated = updated["metadata"] assert populated["user_api_key_metadata"] == real_admin_metadata assert populated["user_api_key_team_metadata"] == real_admin_metadata - assert "_pipeline_managed_guardrails" not in populated or populated[ - "_pipeline_managed_guardrails" - ] != ["evaded"] + assert "_pipeline_managed_guardrails" not in populated or populated["_pipeline_managed_guardrails"] != ["evaded"] other = updated.get("litellm_metadata") or {} assert other.get("user_api_key_metadata") in (None, {}, real_admin_metadata) @@ -463,9 +443,7 @@ async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_str snapshot_body = updated["proxy_server_request"]["body"] assert snapshot_body is not None snapshot_metadata = snapshot_body.get("metadata") or {} - assert "user_api_key_user_id" not in snapshot_metadata or ( - snapshot_metadata["user_api_key_user_id"] != "victim" - ) + assert "user_api_key_user_id" not in snapshot_metadata or (snapshot_metadata["user_api_key_user_id"] != "victim") @pytest.mark.asyncio @@ -520,9 +498,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_secret_fields( ) # secret_fields must exist on the live data dict - assert ( - "secret_fields" in updated - ), "secret_fields must still be present on the live data dict" + assert "secret_fields" in updated, "secret_fields must still be present on the live data dict" assert "raw_headers" in updated["secret_fields"] # But the body snapshot must NOT contain secret_fields @@ -581,8 +557,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r snapshot_body = updated["proxy_server_request"]["body"] assert "proxy_server_request" not in snapshot_body, ( - "proxy_server_request must be excluded from its own body snapshot " - "to prevent the body from self-referencing" + "proxy_server_request must be excluded from its own body snapshot to prevent the body from self-referencing" ) @@ -998,23 +973,18 @@ async def test_add_litellm_data_to_request_strips_client_redaction_bypass_contro assert "turn_off_message_logging" not in (updated.get("litellm_params") or {}).get("metadata", {}) assert "turn_off_message_logging" not in updated["metadata"] assert "turn_off_message_logging" not in (updated.get("litellm_metadata") or {}) + assert "litellm-disable-message-redaction" not in {header.lower() for header in updated["metadata"]["headers"]} assert "litellm-disable-message-redaction" not in { - header.lower() for header in updated["metadata"]["headers"] - } - assert "litellm-disable-message-redaction" not in { - header.lower() - for header in updated["metadata"]["requester_metadata"].get("headers", {}) + header.lower() for header in updated["metadata"]["requester_metadata"].get("headers", {}) } assert "litellm-disable-message-redaction" not in { header.lower() for header in updated["proxy_server_request"]["headers"] } assert "litellm-disable-message-redaction" not in { - header.lower() - for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] + header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } assert "litellm-disable-message-redaction" not in { - header.lower() - for header in (updated.get("litellm_metadata") or {}).get("headers", {}) + header.lower() for header in (updated.get("litellm_metadata") or {}).get("headers", {}) } @@ -1084,12 +1054,7 @@ async def test_add_litellm_data_to_request_admin_callback_vars_turn_off_message_ dynamic_params = initialize_standard_callback_dynamic_params(updated) assert dynamic_params.get("turn_off_message_logging") == "False" - assert ( - should_redact_message_logging( - {"standard_callback_dynamic_params": dynamic_params} - ) - is False - ) + assert should_redact_message_logging({"standard_callback_dynamic_params": dynamic_params}) is False finally: litellm.turn_off_message_logging = original_turn_off_message_logging @@ -1160,12 +1125,7 @@ async def test_add_litellm_data_to_request_admin_callback_vars_turn_off_message_ dynamic_params = initialize_standard_callback_dynamic_params(updated) assert dynamic_params.get("turn_off_message_logging") == "True" - assert ( - should_redact_message_logging( - {"standard_callback_dynamic_params": dynamic_params} - ) - is True - ) + assert should_redact_message_logging({"standard_callback_dynamic_params": dynamic_params}) is True finally: litellm.turn_off_message_logging = original_turn_off_message_logging @@ -1206,9 +1166,7 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o "headers": {"litellm-disable-message-redaction": "true"}, "turn_off_message_logging": False, }, - "litellm_metadata": json.dumps( - {"headers": {"LiteLLM-Disable-Message-Redaction": "true"}} - ), + "litellm_metadata": json.dumps({"headers": {"LiteLLM-Disable-Message-Redaction": "true"}}), }, request=request_mock, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", **auth_kwargs), @@ -1221,23 +1179,18 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o assert updated["turn_off_message_logging"] is False assert updated["metadata"]["turn_off_message_logging"] is False + assert "litellm-disable-message-redaction" in {header.lower() for header in updated["metadata"]["headers"]} assert "litellm-disable-message-redaction" in { - header.lower() for header in updated["metadata"]["headers"] - } - assert "litellm-disable-message-redaction" in { - header.lower() - for header in updated["metadata"]["requester_metadata"].get("headers", {}) + header.lower() for header in updated["metadata"]["requester_metadata"].get("headers", {}) } assert "litellm-disable-message-redaction" in { header.lower() for header in updated["proxy_server_request"]["headers"] } assert "litellm-disable-message-redaction" in { - header.lower() - for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] + header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } assert "litellm-disable-message-redaction" in { - header.lower() - for header in (updated.get("litellm_metadata") or {}).get("headers", {}) + header.lower() for header in (updated.get("litellm_metadata") or {}).get("headers", {}) } @@ -1527,9 +1480,7 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart(): request_mock.client.host = "127.0.0.1" # Simulate multipart data (metadata as string) - metadata_dict = { - "tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"] - } + metadata_dict = {"tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"]} stringified_metadata = json.dumps(metadata_dict) data = { @@ -1857,23 +1808,15 @@ def test_key_dynamic_logging_settings(): # Test with langfuse logging key_with_langfuse = UserAPIKeyAuth( api_key="test-key", - metadata={ - "logging": [{"callback_name": "langfuse", "callback_type": "success"}] - }, + metadata={"logging": [{"callback_name": "langfuse", "callback_type": "success"}]}, team_metadata={}, ) - result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - key_with_langfuse - ) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) assert result == [{"callback_name": "langfuse", "callback_type": "success"}] # Test with no logging metadata - key_without_logging = UserAPIKeyAuth( - api_key="test-key", metadata={}, team_metadata={} - ) - result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - key_without_logging - ) + key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_without_logging) assert result is None @@ -1885,35 +1828,23 @@ def test_team_dynamic_logging_settings(): key_with_team_arize = UserAPIKeyAuth( api_key="test-key", metadata={}, - team_metadata={ - "logging": [{"callback_name": "arize", "callback_type": "failure"}] - }, - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_with_team_arize + team_metadata={"logging": [{"callback_name": "arize", "callback_type": "failure"}]}, ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_arize) assert result == [{"callback_name": "arize", "callback_type": "failure"}] # Test with langfuse team logging key_with_team_langfuse = UserAPIKeyAuth( api_key="test-key", metadata={}, - team_metadata={ - "logging": [{"callback_name": "langfuse", "callback_type": "success"}] - }, - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_with_team_langfuse + team_metadata={"logging": [{"callback_name": "langfuse", "callback_type": "success"}]}, ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_langfuse) assert result == [{"callback_name": "langfuse", "callback_type": "success"}] # Test with no team logging metadata - key_without_team_logging = UserAPIKeyAuth( - api_key="test-key", metadata={}, team_metadata={} - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_without_team_logging - ) + key_without_team_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_without_team_logging) assert result is None @@ -1994,9 +1925,7 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): mock_proxy_config = MagicMock() # Call the function - result = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config - ) + result = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config) # Verify the result assert result is not None @@ -2012,9 +1941,7 @@ def test_add_team_callback_rejects_env_reference(): AddTeamCallback( callback_name="langfuse", callback_type="success", - callback_vars={ - "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY_TEMP" - }, + callback_vars={"langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY_TEMP"}, ) assert "os.environ/" in str(exc_info.value) @@ -2045,9 +1972,7 @@ def test_get_dynamic_logging_metadata_ignores_env_reference_from_key_metadata( team_metadata={}, ) - result = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=MagicMock() - ) + result = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) assert result is None @@ -2058,16 +1983,12 @@ def test_get_num_retries_from_request(): """ # Test case 1: Header is present with valid integer string headers_with_retries = {"x-litellm-num-retries": "3"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_retries - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_retries) assert result == 3 # Test case 2: Header is not present headers_without_retries = {"Content-Type": "application/json"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_without_retries - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_without_retries) assert result is None # Test case 3: Empty headers dictionary @@ -2082,9 +2003,7 @@ def test_get_num_retries_from_request(): # Test case 5: Header present with large number headers_with_large_number = {"x-litellm-num-retries": "100"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_large_number - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_large_number) assert result == 100 # Test case 6: Multiple headers with num retries header @@ -2108,9 +2027,7 @@ def test_get_num_retries_from_request(): # Test case 9: Header present with negative number headers_with_negative = {"x-litellm-num-retries": "-1"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_negative - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_negative) assert result == -1 @@ -2262,9 +2179,7 @@ def test_add_user_api_key_auth_to_request_metadata(): ), ], ) -def test_add_headers_to_llm_call_by_model_group( - data, model_group_settings, expected_headers_added -): +def test_add_headers_to_llm_call_by_model_group(data, model_group_settings, expected_headers_added): """ Test LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group method @@ -2285,9 +2200,7 @@ def test_add_headers_to_llm_call_by_model_group( "X-Custom-Header": "custom-value", } - user_api_key_dict = UserAPIKeyAuth( - api_key="test-key", user_id="test-user", org_id="test-org" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="test-user", org_id="test-org") # Mock the model_group_settings original_model_group_settings = getattr(litellm, "model_group_settings", None) @@ -2305,7 +2218,6 @@ def test_add_headers_to_llm_call_by_model_group( "add_headers_to_llm_call", return_value=expected_returned_headers if expected_headers_added else {}, ) as mock_add_headers: - # Make a copy of original data to verify it's not mutated unexpectedly original_data = copy.deepcopy(data) @@ -2362,7 +2274,6 @@ def test_add_headers_to_llm_call_by_model_group_empty_headers_returned(): "add_headers_to_llm_call", return_value={}, # Return empty dict ) as mock_add_headers: - result = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=headers, user_api_key_dict=user_api_key_dict ) @@ -2410,7 +2321,6 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data(): "add_headers_to_llm_call", return_value=new_headers, ) as mock_add_headers: - result = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=headers, user_api_key_dict=user_api_key_dict ) @@ -2527,13 +2437,9 @@ async def test_add_litellm_metadata_from_request_headers(): general_settings = {} # Create mock select_data_generator with correct signature - def mock_select_data_generator( - response=None, user_api_key_dict=None, request_data=None - ): + def mock_select_data_generator(response=None, user_api_key_dict=None, request_data=None): async def mock_generator(): - yield "data: " + json.dumps( - {"choices": [{"delta": {"content": "Hello"}}]} - ) + "\n\n" + yield "data: " + json.dumps({"choices": [{"delta": {"content": "Hello"}}]}) + "\n\n" yield "data: [DONE]\n\n" return mock_generator() @@ -2560,21 +2466,19 @@ async def test_add_litellm_metadata_from_request_headers(): await asyncio.sleep(3) # Check if standard_logging_object was set - assert ( - test_logger.standard_logging_object is not None - ), "standard_logging_object should be populated after LLM request" + assert test_logger.standard_logging_object is not None, ( + "standard_logging_object should be populated after LLM request" + ) # Verify the logging object contains expected metadata standard_logging_obj = test_logger.standard_logging_object - print( - f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}" - ) + print(f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}") SPEND_LOGS_METADATA = standard_logging_obj["metadata"]["spend_logs_metadata"] - assert SPEND_LOGS_METADATA == dict( - json.loads(headers["x-litellm-spend-logs-metadata"]) - ), "spend_logs_metadata should be the same as the headers" + assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), ( + "spend_logs_metadata should be the same as the headers" + ) finally: litellm.callbacks = original_callbacks @@ -2725,11 +2629,7 @@ def test_add_litellm_metadata_from_request_headers_generic_session_id_header(): def test_add_litellm_metadata_from_anthropic_user_id_sets_session_id(): - data = { - "metadata": { - "user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01" - } - } + data = {"metadata": {"user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01"}} LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( headers={}, data=data, _metadata_variable_name="metadata" ) @@ -2845,9 +2745,7 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers assert ( - get_chain_id_from_headers( - {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} - ) + get_chain_id_from_headers({"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"}) == "e96634a3-fa28-4083-b354-55542e2dca01" ) # Short / non-alphanumeric values should be ignored @@ -3014,19 +2912,13 @@ def test_get_internal_user_header_from_mapping_returns_expected_header(): {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - mappings - ) + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert header_name == "X-OpenWebUI-User-Id" def test_get_internal_user_header_from_mapping_none_when_absent(): - mappings = [ - {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} - ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - mappings - ) + mappings = [{"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}] + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert header_name is None single = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} @@ -3047,9 +2939,7 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): ] } - result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( - general_settings, user_api_key_dict, headers - ) + result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(general_settings, user_api_key_dict, headers) assert result is user_api_key_dict assert user_api_key_dict.user_id == "internal-user-123" @@ -3065,9 +2955,7 @@ def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchan assert user_api_key_dict.user_id is None general_settings = { - "user_header_mappings": [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} - ] + "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] } result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, {"Other": "value"} @@ -3087,9 +2975,7 @@ def test_get_sanitized_user_information_from_key_includes_guardrails_metadata(): metadata={"guardrails": ["presidio", "aporia"], "other_field": "value"}, ) - result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) assert result["user_api_key_auth_metadata"] is not None assert "guardrails" in result["user_api_key_auth_metadata"] @@ -3118,9 +3004,7 @@ def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): team_max_budget=1000.0, ) - sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) assert sanitized["user_api_key_spend"] == 1.5 assert sanitized["user_api_key_max_budget"] == 10.0 @@ -3129,9 +3013,7 @@ def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): assert sanitized["user_api_key_team_spend"] == 250.75 assert sanitized["user_api_key_team_max_budget"] == 1000.0 - logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - dict(sanitized) - ) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(dict(sanitized)) assert logging_metadata["user_api_key_user_spend"] == 25.5 assert logging_metadata["user_api_key_user_max_budget"] == 100.0 @@ -3148,12 +3030,8 @@ def test_user_and_team_spend_and_budget_default_to_none_in_standard_logging_meta user_api_key_dict = UserAPIKeyAuth(api_key="test-key-hash") - sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) - logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - dict(sanitized) - ) + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(dict(sanitized)) assert logging_metadata["user_api_key_user_spend"] is None assert logging_metadata["user_api_key_user_max_budget"] is None @@ -3485,22 +3363,16 @@ async def test_embedding_header_forwarding_with_model_group(): # Verify that only x- prefixed headers (except x-stainless) were forwarded forwarded_headers = updated_data["headers"] - assert ( - "X-Custom-Header" in forwarded_headers - ), "X-Custom-Header should be forwarded" + assert "X-Custom-Header" in forwarded_headers, "X-Custom-Header should be forwarded" assert forwarded_headers["X-Custom-Header"] == "custom-value" assert "X-Request-ID" in forwarded_headers, "X-Request-ID should be forwarded" assert forwarded_headers["X-Request-ID"] == "test-request-123" # Verify that authorization header was NOT forwarded (sensitive header) - assert ( - "Authorization" not in forwarded_headers - ), "Authorization header should not be forwarded" + assert "Authorization" not in forwarded_headers, "Authorization header should not be forwarded" # Verify that Content-Type was NOT forwarded (doesn't start with x-) - assert ( - "Content-Type" not in forwarded_headers - ), "Content-Type should not be forwarded" + assert "Content-Type" not in forwarded_headers, "Content-Type should not be forwarded" # Verify original data fields are preserved assert updated_data["model"] == "local-openai/text-embedding-3-small" @@ -3556,9 +3428,9 @@ async def test_embedding_header_forwarding_without_model_group_config(): ) # Verify that headers were NOT added since model is not in forward list - assert ( - "headers" not in updated_data or updated_data.get("headers") is None - ), "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + assert "headers" not in updated_data or updated_data.get("headers") is None, ( + "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + ) # Verify original data fields are preserved assert updated_data["model"] == "text-embedding-ada-002" @@ -3612,9 +3484,7 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry = get_attachment_registry() attachment_registry._attachments = [ PolicyAttachment(policy="global-baseline", scope="*"), # applies to all - PolicyAttachment( - policy="healthcare", teams=["healthcare-team"] - ), # applies to healthcare team + PolicyAttachment(policy="healthcare", teams=["healthcare-team"]), # applies to healthcare team ] attachment_registry._initialized = True @@ -3683,9 +3553,9 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po ) # Verify that 'policies' was removed from the request body - assert ( - "policies" not in data - ), "'policies' should be removed from request body to prevent forwarding to LLM provider" + assert "policies" not in data, ( + "'policies' should be removed from request body to prevent forwarding to LLM provider" + ) # Verify that other fields are preserved assert "model" in data @@ -3730,9 +3600,7 @@ async def test_api_created_global_policy_applies_to_new_key_without_restart(): "runtime-global-policy", Policy(guardrails=PolicyGuardrails(add=["runtime-guardrail"])), ) - attachment_registry.add_attachment( - PolicyAttachment(policy="runtime-global-policy", scope="*") - ) + attachment_registry.add_attachment(PolicyAttachment(policy="runtime-global-policy", scope="*")) await add_guardrails_from_policy_engine( data=data, @@ -3829,9 +3697,7 @@ async def test_bearer_token_not_in_debug_logs(): from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ProxyConfig - secret_token = ( - "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" - ) + secret_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" mock_request = MagicMock(spec=Request) mock_request.headers = { @@ -3877,8 +3743,7 @@ async def test_bearer_token_not_in_debug_logs(): log_output = log_capture.getvalue() assert secret_token not in log_output, ( - f"Bearer token leaked in debug logs. " - f"Found token in log output:\n{log_output[:500]}" + f"Bearer token leaked in debug logs. Found token in log output:\n{log_output[:500]}" ) @@ -4043,9 +3908,7 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, @@ -4056,9 +3919,7 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): } }, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" assert data["api_key"] == "key-hotel-rec-vision" assert data["api_version"] == "2024-06-01" @@ -4071,9 +3932,7 @@ def test_apply_overrides_project_default(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, @@ -4084,9 +3943,7 @@ def test_apply_overrides_project_default(setup_test_credentials): } }, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-rec-app.openai.azure.com/" assert data["api_key"] == "key-hotel-rec" @@ -4098,17 +3955,13 @@ def test_apply_overrides_team_model_specific(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-westus.openai.azure.com/" assert data["api_key"] == "key-hotel-westus" @@ -4120,17 +3973,13 @@ def test_apply_overrides_team_default(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" @@ -4143,9 +3992,7 @@ def test_apply_overrides_no_config(setup_test_credentials): team_metadata={}, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4161,17 +4008,9 @@ def test_apply_overrides_clientside_credentials_take_precedence( } user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://my-custom-endpoint.openai.azure.com/" assert data["api_key"] == "my-custom-key" @@ -4181,15 +4020,9 @@ def test_apply_overrides_missing_credential_name(setup_test_credentials): data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4199,17 +4032,9 @@ def test_apply_overrides_api_version_only_if_present(setup_test_credentials): data = {"model": "gpt-3.5"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" assert "api_version" not in data @@ -4220,15 +4045,9 @@ def test_apply_overrides_no_model_in_data(setup_test_credentials): data = {"messages": [{"role": "user", "content": "hello"}]} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": {"azure": {"litellm_credentials": "some-cred"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "some-cred"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data @@ -4240,9 +4059,7 @@ def test_apply_overrides_none_metadata(setup_test_credentials): team_metadata=None, project_metadata=None, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data @@ -4251,15 +4068,9 @@ def test_apply_overrides_clientside_api_version_preserved(setup_test_credentials data = {"model": "gpt-4-vision", "api_version": "2025-01-01"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) # api_base and api_key should be set from credential assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" assert data["api_key"] == "key-hotel-rec-vision" @@ -4272,9 +4083,7 @@ def test_resolve_non_dict_model_config_ignored(): result = _resolve_credential_from_model_config("gpt-4", "not-a-dict", None) assert result is None - result = _resolve_credential_from_model_config( - "gpt-4", None, ["also", "not", "a", "dict"] - ) + result = _resolve_credential_from_model_config("gpt-4", None, ["also", "not", "a", "dict"]) assert result is None # Valid config still works alongside invalid one @@ -4292,9 +4101,7 @@ def test_resolve_pre_alias_model_name_fallback(): "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, } # Post-alias name doesn't match, but pre-alias does (team scope) - result = _resolve_credential_from_model_config( - "azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4") assert result == "team-gpt4" # Same test for project scope @@ -4314,15 +4121,11 @@ def test_resolve_post_alias_name_takes_priority(): "gpt-4o-team-1": {"azure": {"litellm_credentials": "post-alias-cred"}}, } # Team scope - result = _resolve_credential_from_model_config( - "gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4") assert result == "post-alias-cred" # Project scope - result = _resolve_credential_from_model_config( - "gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4") assert result == "post-alias-cred" @@ -4354,15 +4157,9 @@ def test_apply_overrides_feature_flag_disabled_by_default(): data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4522,9 +4319,7 @@ async def test_team_guardrail_merges_with_global_policy(): policy_registry = get_policy_registry() policy_registry._policies = { "global-policy": Policy( - guardrails=PolicyGuardrails( - add=["policy-guardrail-1", "policy-guardrail-2"] - ), + guardrails=PolicyGuardrails(add=["policy-guardrail-1", "policy-guardrail-2"]), ), } policy_registry._initialized = True @@ -4545,18 +4340,10 @@ async def test_team_guardrail_merges_with_global_policy(): guardrails = data["metadata"].get("guardrails", []) - assert ( - "team-direct-guardrail" in guardrails - ), f"Team guardrail missing from merged list: {guardrails}" - assert ( - "policy-guardrail-1" in guardrails - ), f"policy-guardrail-1 missing: {guardrails}" - assert ( - "policy-guardrail-2" in guardrails - ), f"policy-guardrail-2 missing: {guardrails}" - assert len(guardrails) == len( - set(guardrails) - ), f"Duplicates in guardrails list: {guardrails}" + assert "team-direct-guardrail" in guardrails, f"Team guardrail missing from merged list: {guardrails}" + assert "policy-guardrail-1" in guardrails, f"policy-guardrail-1 missing: {guardrails}" + assert "policy-guardrail-2" in guardrails, f"policy-guardrail-2 missing: {guardrails}" + assert len(guardrails) == len(set(guardrails)), f"Duplicates in guardrails list: {guardrails}" # Verify get_guardrail_from_metadata returns the merged list even # when litellm_metadata is present (the bug: it returned [] before fix) @@ -4567,9 +4354,9 @@ async def test_team_guardrail_merges_with_global_policy(): dummy = _DummyGuardrail(guardrail_name="team-direct-guardrail") returned = dummy.get_guardrail_from_metadata(data) - assert ( - "team-direct-guardrail" in returned - ), f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + assert "team-direct-guardrail" in returned, ( + f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + ) finally: policy_registry._policies = {} @@ -4622,9 +4409,7 @@ def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata(): } result = dummy.get_guardrail_from_metadata(data) - assert result == [ - "my-guardrail" - ], f"Expected guardrails from litellm_metadata fallback, got: {result}" + assert result == ["my-guardrail"], f"Expected guardrails from litellm_metadata fallback, got: {result}" def _build_request_mock_with_headers(headers: dict) -> Request: @@ -4651,9 +4436,7 @@ class TestApplyClientTagPolicyPreAuth: """ def test_merges_header_tags_into_metadata(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme,env:prod"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme,env:prod"}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -4670,9 +4453,7 @@ class TestApplyClientTagPolicyPreAuth: assert data["metadata"]["tags"] == ["tenant:acme", "env:prod"] def test_unions_header_tags_with_existing_metadata_tags(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme,env:prod"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme,env:prod"}) data = { "model": "gpt-3.5-turbo", "metadata": {"tags": ["env:prod", "team:platform"]}, @@ -4697,9 +4478,7 @@ class TestApplyClientTagPolicyPreAuth: # (inside common_checks) enforces per-tag budgets on whatever tags # it sees in request_data, including body tags. The helper only # adds header tags to metadata.tags. - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = { "model": "gpt-3.5-turbo", "tags": ["root-tag"], @@ -4726,9 +4505,7 @@ class TestApplyClientTagPolicyPreAuth: ] def test_uses_litellm_metadata_when_present(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = { "model": "gpt-3.5-turbo", "litellm_metadata": {"foo": "bar"}, @@ -4823,9 +4600,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid": return 0.50 return fallback_spend @@ -4860,9 +4635,7 @@ class TestApplyClientTagPolicyPreAuth: from litellm.proxy.auth.auth_checks import _tag_max_budget_check from litellm.proxy.utils import ProxyLogging - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -4882,9 +4655,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -4919,9 +4690,7 @@ class TestApplyClientTagPolicyPreAuth: "/v1/messages", ], ) - async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route( - self, route - ): + async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route(self, route): """Regression: on LITELLM_METADATA_ROUTES (bedrock, /v1/messages, ...), common_checks pre-seeds ``litellm_metadata`` and writes key tags there before ``_tag_max_budget_check`` reads from the same key. The auth wrapper @@ -4935,9 +4704,7 @@ class TestApplyClientTagPolicyPreAuth: from litellm.proxy.auth.auth_checks import common_checks from litellm.proxy.utils import ProxyLogging - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = {"model": "us.anthropic.claude-sonnet-4-6"} valid_token = UserAPIKeyAuth( token="test-token", @@ -4962,9 +4729,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -5132,9 +4897,7 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:engineering": return 0.50 return fallback_spend @@ -5185,9 +4948,7 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:engineering": return 0.05 return fallback_spend @@ -5268,9 +5029,7 @@ def test_resolve_provider_from_deployment_falls_back_to_pre_alias(): router.get_deployment_by_model_group_name.side_effect = lookup - result = _resolve_provider_from_deployment( - router, "post-alias-name", pre_alias_model_name="pre-alias-name" - ) + result = _resolve_provider_from_deployment(router, "post-alias-name", pre_alias_model_name="pre-alias-name") assert result == "bedrock" @@ -5335,17 +5094,9 @@ def test_apply_overrides_no_router_keeps_legacy_behaviour(setup_test_credentials data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict, llm_router=None + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict, llm_router=None) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" @@ -5371,9 +5122,7 @@ def test_apply_overrides_provider_prefix_in_model_skips_router_lookup( ) router = MagicMock() - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict, llm_router=router - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict, llm_router=router) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" router.get_deployment_by_model_group_name.assert_not_called() @@ -5447,6 +5196,359 @@ async def test_add_litellm_data_to_request_agentic_cli_drop_params( assert updated.get("drop_params") == expected_drop_params +@pytest.fixture +def _seeded_logging_credentials(): + from litellm.models.credentials import CredentialItem + + original = litellm.credential_list + litellm.credential_list = [ + CredentialItem( + credential_name="langfuse-eu", + credential_values={ + "langfuse_host": "https://cloud.langfuse.com", + "langfuse_public_key": "pk-eu", + "langfuse_secret_key": "sk-eu", + }, + credential_info={ + "credential_type": "logging", + "description": "langfuse_otel", + "access": {"teams": ["team-x"]}, + }, + ), + CredentialItem( + credential_name="arize-prod", + credential_values={ + "arize_space_id": "S", + "arize_api_key": "K", + "arize_project_name": "tenant-arize", + }, + credential_info={ + "credential_type": "logging", + "description": "arize", + "access": {"teams": ["team-az"]}, + }, + ), + CredentialItem( + credential_name="generic-org", + credential_values={"otel_endpoint": "http://collector.internal/v1/traces"}, + credential_info={ + "credential_type": "logging", + "description": "generic", + "access": {"orgs": ["org-1"]}, + }, + ), + CredentialItem( + credential_name="empty-deny", + credential_values={"otel_endpoint": "http://never/v1/traces"}, + credential_info={ + "credential_type": "logging", + "description": "generic", + "access": {}, + }, + ), + # A provider credential that must never resolve as a logging destination. + CredentialItem( + credential_name="openai-key", + credential_values={"api_key": "sk-openai"}, + credential_info={"custom_llm_provider": "openai"}, + ), + ] + # Admin-owned destinations are gated on LITELLM_OTEL_V2; the resolver no-ops when + # the flag is off, so exercise these tests with the feature actually enabled. + from litellm.integrations.otel.model.config import is_otel_v2_enabled + + prev_flag = os.environ.get("LITELLM_OTEL_V2") + os.environ["LITELLM_OTEL_V2"] = "true" + is_otel_v2_enabled.cache_clear() + try: + yield + finally: + litellm.credential_list = original + if prev_flag is None: + os.environ.pop("LITELLM_OTEL_V2", None) + else: + os.environ["LITELLM_OTEL_V2"] = prev_flag + is_otel_v2_enabled.cache_clear() + + +def _auth(token="hashed-key", org_id=None, team_id="team-x"): + return UserAPIKeyAuth(api_key="hashed-key", token=token, org_id=org_id, team_id=team_id) + + +def _patch_identity(monkeypatch, *, team_org_id=None, **_ignored): + """Connect a prisma client and route the resolver's only remaining DB lookup. + + Selection is access-only, read from ``litellm.credential_list``. The sole lookup + left is ``_effective_org_id`` resolving the team's organization when the token + carries no ``org_id``, so ``get_team_object`` returns just ``organization_id``. + """ + from types import SimpleNamespace + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth import auth_checks + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr( + auth_checks, + "get_team_object", + AsyncMock(return_value=SimpleNamespace(organization_id=team_org_id)), + ) + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_team_access(_seeded_logging_credentials, monkeypatch): + """A destination whose access grants the caller's team fires for it, and only it.""" + from litellm.proxy.litellm_pre_call_utils import _resolve_logging_exporters + + _patch_identity(monkeypatch) + destinations, backends = await _resolve_logging_exporters(_auth(team_id="team-x")) + assert {d["endpoint"] for d in destinations} == {"https://cloud.langfuse.com/api/public/otel"} + assert backends == ("langfuse_otel",) + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_org_access(_seeded_logging_credentials, monkeypatch): + """An org-scoped destination fires for a caller in that org.""" + from litellm.proxy.litellm_pre_call_utils import _resolve_logging_exporters + + _patch_identity(monkeypatch) + destinations, backends = await _resolve_logging_exporters(_auth(team_id="team-none", org_id="org-1")) + assert {d["endpoint"] for d in destinations} == {"http://collector.internal/v1/traces"} + assert backends == ("generic",) + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_org_fallback_from_team(_seeded_logging_credentials, monkeypatch): + """When the token carries no org_id, the team's organization grants org-scoped + destinations via the ``_effective_org_id`` fallback.""" + from litellm.proxy.litellm_pre_call_utils import _resolve_logging_exporters + + _patch_identity(monkeypatch, team_org_id="org-1") + destinations, _ = await _resolve_logging_exporters(_auth(team_id="team-none", org_id=None)) + assert {d["endpoint"] for d in destinations} == {"http://collector.internal/v1/traces"} + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_carries_arize_project(_seeded_logging_credentials, monkeypatch): + from litellm.proxy.litellm_pre_call_utils import _resolve_logging_exporters + + _patch_identity(monkeypatch) + destinations, _ = await _resolve_logging_exporters(_auth(team_id="team-az")) + + assert destinations == ( + { + "callback_name": "arize", + "endpoint": "https://otlp.arize.com/v1", + "headers": {"space_id": "S", "api_key": "K"}, + "resource_attributes": { + "model_id": "tenant-arize", + "arize.project.name": "tenant-arize", + }, + # Arize's own endpoint is gRPC, so the destination carries no transport + # override and the backend's intrinsic default applies. + "protocol": None, + }, + ) + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_empty_without_access(_seeded_logging_credentials, monkeypatch): + """An identity no destination's access grants gets nothing; empty access is deny-all.""" + from litellm.proxy.litellm_pre_call_utils import _resolve_logging_exporters + + _patch_identity(monkeypatch) + destinations, backends = await _resolve_logging_exporters(_auth(team_id="team-none")) + assert destinations == () and backends == () + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_skips_provider_creds(_seeded_logging_credentials, monkeypatch): + """A provider credential (not credential_type=logging) is never a destination, + even for a team that resolves a real one.""" + from litellm.proxy.litellm_pre_call_utils import _resolve_logging_exporters + + _patch_identity(monkeypatch) + destinations, backends = await _resolve_logging_exporters(_auth(team_id="team-az")) + assert backends == ("arize",) + + +@pytest.mark.asyncio +async def test_apply_admin_logging_exporters_stamps_and_activates(_seeded_logging_credentials, monkeypatch): + from litellm.integrations.otel.plumbing.context import ( + _request_destinations, + request_destinations, + ) + from litellm.proxy.litellm_pre_call_utils import _apply_admin_logging_exporters + + _patch_identity(monkeypatch) + token = _request_destinations.set(()) + try: + await _apply_admin_logging_exporters(_auth()) + + context_destinations = request_destinations() + assert len(context_destinations) == 1 + assert context_destinations[0].callback_name == "langfuse_otel" + assert context_destinations[0].endpoint == "https://cloud.langfuse.com/api/public/otel" + finally: + _request_destinations.reset(token) + + +@pytest.mark.asyncio +async def test_apply_admin_logging_exporters_swallows_resolver_failure(monkeypatch): + """Telemetry setup is best-effort: when the pre-call resolver raises a + non-``HTTPException``, ``_apply_admin_logging_exporters`` swallows it so the request + proceeds with no exception escaping and no destinations anchored.""" + import litellm.proxy.litellm_pre_call_utils as pcu + from litellm.integrations.otel.plumbing.context import ( + _request_destinations, + request_destinations, + ) + + async def _boom(_uapk): + raise RuntimeError("cache backend exploded") + + monkeypatch.setattr(pcu, "_resolve_logging_exporters", _boom) + token = _request_destinations.set(()) + try: + await pcu._apply_admin_logging_exporters(_auth(), cached_destinations=None) + assert request_destinations() == () + finally: + _request_destinations.reset(token) + + +@pytest.mark.asyncio +async def test_empty_resolution_clears_a_previous_messages_destinations(monkeypatch): + """Regression: an empty resolution must be published, not skipped. + + A stateful streamable-HTTP MCP session runs every message on the task its + ``initialize`` POST spawned, so the destination ContextVar is shared across + messages on that task. Returning early on an empty resolution leaves the previous + message's destinations standing, and a revoked grant keeps exporting for the life + of the session. ``_hoist_request_destinations`` already sets unconditionally; this + is the same contract on the other path. + """ + import litellm.proxy.litellm_pre_call_utils as pcu + from litellm.integrations.otel.model.destination import OtelDestination + from litellm.integrations.otel.plumbing.context import ( + _request_destinations, + request_destinations, + ) + + stale = (OtelDestination(endpoint="http://revoked.internal/v1/traces", callback_name="generic"),) + + async def _grants_nothing(_uapk): + return (), () + + monkeypatch.setattr(pcu, "_resolve_logging_exporters", _grants_nothing) + token = _request_destinations.set(stale) + try: + await pcu._apply_admin_logging_exporters(_auth(), cached_destinations=None) + assert request_destinations() == (), "a revoked identity must not inherit the previous message's destinations" + finally: + _request_destinations.reset(token) + + +@pytest.mark.asyncio +async def test_client_cannot_control_otel_destinations(_seeded_logging_credentials): + """Y3 spoofing guard: a client cannot control OTEL export destinations. + + A request injects ``otel_destinations`` at the top level AND inside + ``litellm_metadata`` pointing at an attacker endpoint, for an identity no + destination grants. Destinations are admin-owned and resolved server-side, so the + client value is wiped before the resolver runs; default-deny then adds nothing. + The attacker endpoint must appear nowhere in the outgoing request. This drives + the full ``add_litellm_data_to_request`` so the wipe-then-resolve ORDER is under + test, not just the resolver in isolation (the resolver never reads client data, + but the guard that a client value cannot survive lives in the wipe). + """ + from types import SimpleNamespace + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + # No auth-boundary cache: force the resolver to run for real (default-deny). + request_mock.state = SimpleNamespace() + + attacker = [ + { + "callback_name": "langfuse_otel", + "endpoint": "https://attacker.example/otel", + "headers": {"Authorization": "Basic stolen"}, + "resource_attributes": {}, + } + ] + data = { + "model": "gpt-3.5-turbo", + "otel_destinations": attacker, + "litellm_metadata": {"otel_destinations": attacker}, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + # No routing-relevant carrier retains the client value: the top level, the + # request metadata key, and litellm_metadata (the only carrier the dynamic-params + # reader consults). proxy_server_request.body is a verbatim audit echo of the + # client's own request and is never read for trace routing, so it is not checked. + assert "otel_destinations" not in updated + assert "otel_destinations" not in (updated.get("metadata") or {}) + assert "otel_destinations" not in (updated.get("litellm_metadata") or {}) + + # The reader that feeds the gen-AI-span tracer sees nothing (default-deny). + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params, + ) + + dynamic_params = initialize_standard_callback_dynamic_params(updated) + assert not dynamic_params.get("otel_destinations") + + +@pytest.mark.asyncio +async def test_apply_admin_logging_exporters_registers_on_failure(_seeded_logging_credentials, monkeypatch): + """An admin-owned destination must capture a FAILED upstream call, not only a + successful one. + + The destination sink is one process-wide logger, so it has to sit on the failure + list as well as the success list; registering it on success alone means a + 401/timeout never reaches the destination and the trace lands with no error + gen-AI span. Registration is idempotent. + """ + import litellm + from litellm.integrations.otel.destination_logger import admin_destination_logger + from litellm.integrations.otel.logger import publish_global_otel_v2_provider + + sink = admin_destination_logger() + publish_global_otel_v2_provider([], lambda provider: None) + publish_global_otel_v2_provider([], lambda provider: None) + for bucket in (litellm._async_success_callback, litellm._async_failure_callback): + assert sum(1 for callback in bucket if callback is sink) == 1 + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_merges_metadata_tags_on_responses_route(): """Regression for #31584: user-supplied metadata.tags must be merged into @@ -5752,9 +5854,7 @@ def test_get_sanitized_user_information_from_key_drops_callback_config(): }, ) - result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) auth_metadata = result["user_api_key_auth_metadata"] assert "logging" not in auth_metadata @@ -5766,6 +5866,15 @@ def test_get_sanitized_user_information_from_key_drops_callback_config(): assert "logging" in (user_api_key_dict.metadata or {}) +def test_otel_destination_params_declares_resource_attributes(): + """The resolver populates ``resource_attributes`` and the auth hoist reads it, so + the ``OtelDestinationParams`` TypedDict must declare it (else strict type-checking + rejects the resolver's dict).""" + from litellm.types.utils import OtelDestinationParams + + assert "resource_attributes" in OtelDestinationParams.__annotations__ + + def test_team_alias_targeting_deleted_team_deployment_keeps_requested_model(monkeypatch): """ Regression: a team's model_aliases can point at the internal routing key @@ -5794,9 +5903,7 @@ def test_team_alias_targeting_deleted_team_deployment_keeps_requested_model(monk ) with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): - _update_model_if_team_alias_exists( - data=test_data, user_api_key_dict=user_api_key_dict - ) + _update_model_if_team_alias_exists(data=test_data, user_api_key_dict=user_api_key_dict) assert test_data.get("model") == "gpt-4" @@ -5820,9 +5927,7 @@ def test_team_alias_targeting_live_team_deployment_still_rewrites(monkeypatch): ) with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): - _update_model_if_team_alias_exists( - data=test_data, user_api_key_dict=user_api_key_dict - ) + _update_model_if_team_alias_exists(data=test_data, user_api_key_dict=user_api_key_dict) assert test_data.get("model") == "model_name_team-1_live-uuid" @@ -6583,3 +6688,108 @@ async def test_add_litellm_data_to_request_caller_tags_empty_when_caller_sends_n assert updated["metadata"]["tags"] == ["key-supplied"] assert updated["metadata"]["caller_tags"] == () + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_short_circuits_without_destinations(monkeypatch): + """With no logging destination in the registry, the resolver returns empty and does NOT + run the per-request org/team lookup, so a proxy not using admin-owned destinations pays + nothing on the auth path for team-scoped keys.""" + from litellm.proxy import litellm_pre_call_utils as pcu + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai", + credential_values={"api_key": "sk"}, + credential_info={"custom_llm_provider": "openai"}, + ) + ], + ) + lookups = {"org": 0} + + async def _spy_effective_org_id(user_api_key_dict): + lookups["org"] += 1 + return None + + monkeypatch.setattr(pcu, "_effective_org_id", _spy_effective_org_id) + + key = UserAPIKeyAuth(api_key="k", team_id="t1") + destinations, backends = await pcu._resolve_logging_exporters(key) + + assert destinations == () and backends == () + assert lookups["org"] == 0 # the org/team lookup was skipped entirely + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_runs_lookup_when_a_destination_exists(monkeypatch): + """The short-circuit must not skip resolution when a destination exists: a global + destination is still resolved for a team-scoped key, and the org lookup runs.""" + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy import litellm_pre_call_utils as pcu + + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="d-global", + credential_values={"otel_endpoint": "https://collector/v1/traces"}, + credential_info={"credential_type": "logging", "description": "generic", "access": {"global": True}}, + ) + ], + ) + lookups = {"org": 0} + + async def _spy_effective_org_id(user_api_key_dict): + lookups["org"] += 1 + return None + + monkeypatch.setattr(pcu, "_effective_org_id", _spy_effective_org_id) + + key = UserAPIKeyAuth(api_key="k", team_id="t1") + destinations, backends = await pcu._resolve_logging_exporters(key) + + assert lookups["org"] == 1 # a destination exists, so the resolver runs the lookup + assert "generic" in backends # global access grants the team key + + +@pytest.mark.asyncio +async def test_resolve_logging_exporters_noop_when_flag_off(monkeypatch): + """LITELLM_OTEL_V2 is the sole activation gate: with the flag off, the resolver + returns nothing even when a global destination is registered, so no backend is + activated for the request and an existing v1 deployment is unaffected.""" + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.proxy import litellm_pre_call_utils as pcu + + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="d-global", + credential_values={"otel_endpoint": "https://collector/v1/traces"}, + credential_info={"credential_type": "logging", "description": "generic", "access": {"global": True}}, + ) + ], + ) + + async def _boom_org(user_api_key_dict): + raise AssertionError("resolver must short-circuit before any DB lookup when the flag is off") + + monkeypatch.setattr(pcu, "_effective_org_id", _boom_org) + + key = UserAPIKeyAuth(api_key="k", team_id="t1") + destinations, backends = await pcu._resolve_logging_exporters(key) + is_otel_v2_enabled.cache_clear() + + assert destinations == () + assert backends == ()