diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0ec28b62a3a..d35077bba47 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2844,7 +2844,6 @@ async def _authorize_authenticated_request( return None -@tracer.wrap() def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None: """Anchor the OTLP destinations this key or team overrides its traces to. @@ -2882,6 +2881,7 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc) +@tracer.wrap() async def user_api_key_auth( request: Request, api_key: str = fastapi.Security(api_key_header), diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0ced2294248..32422e6afd3 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1040,7 +1040,9 @@ def resolve_tenant_otel_destinations( the request has an outcome, so honouring the filter would mean holding every span back until the call finishes. Those entries keep today's behaviour instead, where the tenant's credentials reach the backend through per-request tracer routing and - the operator's exporter is left alone. + the operator's exporter is left alone. Its ``callback_vars`` still take part in the + merge for a backend another entry made eligible, so the destination carries the + same credentials the runtime parser resolves for that request. A backend the request disabled dynamically, through the key's ``litellm_disabled_callbacks`` or the ``x-litellm-disable-callbacks`` header in @@ -1068,12 +1070,13 @@ def resolve_tenant_otel_destinations( callback for item in entries if (callback := _get_validated_callback_metadata(item=item, source="otel-destination")) is not None - if callback.callback_type != "failure" if callback.callback_name.lower() not in disabled ) return tuple( destination - for name in dict.fromkeys(callback.callback_name for callback in callbacks) + for name in dict.fromkeys( + callback.callback_name for callback in callbacks if callback.callback_type != "failure" + ) if ( destination := destination_for( name, diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py index 1799381bada..67695d5aed8 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py @@ -2,7 +2,9 @@ import contextvars import time +from base64 import b64encode from collections.abc import Mapping +from functools import reduce from types import MappingProxyType import pytest @@ -49,8 +51,11 @@ from litellm.integrations.otel.presets.destinations import ( destination_for, ) from litellm.integrations.otel.presets.langfuse import langfuse_preset -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.litellm_pre_call_utils import resolve_tenant_otel_destinations +from litellm.proxy._types import AddTeamCallback, UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import ( + convert_key_logging_metadata_to_callback, + resolve_tenant_otel_destinations, +) from litellm.types.utils import StandardCallbackDynamicParams LANGFUSE_DEST = OtelDestination( @@ -1822,6 +1827,26 @@ class TestTenantConfigAgreement: assert [d.endpoint for d in destinations] == ["http://key.local/api/public/otel"] + def test_a_failure_entry_still_wins_the_merge_next_to_a_success_entry(self): + entries = [ + {**self._entry("http://team.local"), "callback_type": "success"}, + { + **self._entry("http://key.local", langfuse_public_key="pk-failure", langfuse_secret_key="sk-failure"), + "callback_type": "failure", + }, + ] + runtime = reduce( + lambda merged, entry: convert_key_logging_metadata_to_callback(AddTeamCallback(**entry), merged), + entries, + None, + ) + + destinations = resolve_tenant_otel_destinations(UserAPIKeyAuth(team_metadata={"logging": entries})) + + assert runtime.callback_vars["langfuse_host"] == "http://key.local" + assert [d.endpoint for d in destinations] == ["http://key.local/api/public/otel"] + assert destinations[0].headers["Authorization"] == f"Basic {b64encode(b'pk-failure:sk-failure').decode()}" + @pytest.fixture def premium(self, monkeypatch): from litellm.proxy import proxy_server 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 d44f96d95bf..ed253c4123e 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 @@ -1,7 +1,12 @@ import asyncio import json +import os +import subprocess +import sys from contextlib import contextmanager from datetime import datetime, timedelta +from pathlib import Path +from textwrap import dedent from types import SimpleNamespace from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -6872,3 +6877,126 @@ class TestLitellmReceivedAtStamping: assert result == earlier assert request.state.litellm_received_at == earlier + + +_RECORDING_DDTRACE = dedent( + ''' + import functools + import inspect + + + class _Span: + def __enter__(self): + return self + + def __exit__(self, *exc): + return None + + + class _Tracer: + def __init__(self): + self.spans = [] + + def wrap(self, name=None, **kwargs): + def decorator(f): + span_name = name or f"{f.__module__}.{f.__name__}" + if inspect.iscoroutinefunction(f): + + @functools.wraps(f) + async def async_wrapped(*args, **kw): + self.spans.append(span_name) + return await f(*args, **kw) + + return async_wrapped + + @functools.wraps(f) + def wrapped(*args, **kw): + self.spans.append(span_name) + return f(*args, **kw) + + return wrapped + + return decorator + + def trace(self, name, **kwargs): + return _Span() + + def current_span(self): + return None + + def current_root_span(self): + return None + + + tracer = _Tracer() + ''' +) + +_DDTRACE_AUTH_PROBE = dedent( + ''' + import asyncio + import json + + import ddtrace + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + proxy_server.master_key = "sk-probe" + + + async def auth(api_key): + request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"}) + request._url = URL(url="/chat/completions") + try: + await user_api_key_auth( + request=request, + api_key=api_key, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + custom_litellm_key_header=None, + ) + return "accepted" + except ProxyException: + return "rejected" + + + async def main(): + outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")] + print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans})) + + + asyncio.run(main()) + ''' +) + + +def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path): + stub_root = tmp_path / "site" + (stub_root / "ddtrace").mkdir(parents=True) + (stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE) + probe = tmp_path / "probe.py" + probe.write_text(_DDTRACE_AUTH_PROBE) + repo_root = Path(litellm.__file__).resolve().parent.parent + env = { + **os.environ, + "USE_DDTRACE": "true", + "PYTHONPATH": os.pathsep.join( + [str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p] + ), + } + + result = subprocess.run( + [sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False + ) + + assert result.returncode == 0, result.stderr[-4000:] + report = json.loads(result.stdout.strip().splitlines()[-1]) + assert report["outcomes"] == ["accepted", "rejected"] + auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" + assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span]