From 8cc922aca51dd3102a7182ccaffe183c25db8e81 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Wed, 29 Jul 2026 00:23:14 -0700 Subject: [PATCH] perf(otel/v2): resolve destinations once per request; scope regenerate team fetch Three review findings on the auth/key hot paths, all resolver/DB redundancy rather than correctness bugs. _hoist_request_destinations fires early in the auth builder and again as an outer catch-all, so _resolve_logging_exporters ran twice per authenticated request. It now early-returns when request.state already holds the result, keeping the early hoist (for span timing) and the catch-all (a failed first call leaves state unset and still retries) while resolving once. OtelDestinationParams gains the resource_attributes key the resolver populates and the hoist reads, so the resolver's dict type-checks. regenerate_key_fn loaded the key's team via get_team_object on every team-key regeneration, but only uses it when the request sets access_group_ids or object_permission. The fetch moves back inside that guard, so an ordinary regenerate does no extra DB round-trip; an access-group/object-permission regenerate still fetches. Verified live: one chat resolves destinations once (was twice); a plain regenerate does zero team fetches (was one) while an object_permission regenerate still does one; fan-out isolation unchanged. --- litellm/proxy/auth/user_api_key_auth.py | 6 ++++- .../key_management_endpoints.py | 22 ++++++++-------- litellm/types/utils.py | 1 + .../proxy/auth/test_user_api_key_auth.py | 25 +++++++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 9 +++++++ 5 files changed, 51 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2a092a5a50c..3e4a454242e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1014,8 +1014,12 @@ async def _hoist_request_destinations(request: Request, user_api_key_dict: UserA ``_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. + 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 ( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 1bed6c34e34..34932d08751 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4751,18 +4751,18 @@ async def regenerate_key_fn( # noqa: C901 # single endpoint handling many opti detail={"error": "You are not authorized to regenerate this key"}, ) - regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None - if _key_in_db.team_id is not None: - try: - regenerate_team_table = await get_team_object( - team_id=_key_in_db.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - except HTTPException: - regenerate_team_table = None if data is not None and (data.access_group_ids or data.object_permission is not None): + regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None + if _key_in_db.team_id is not None: + try: + regenerate_team_table = await get_team_object( + team_id=_key_in_db.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + except HTTPException: + regenerate_team_table = None _regen_is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( user_api_key_dict=user_api_key_dict, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4f029e8de95..4d63985f8c0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3037,6 +3037,7 @@ class OtelDestinationParams(TypedDict, total=False): callback_name: str endpoint: str headers: Mapping[str, str] + resource_attributes: Mapping[str, str] class StandardCallbackDynamicParams(TypedDict, total=False): 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 4460d4377b7..dc9723c5eb9 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 @@ -5251,3 +5251,28 @@ async def test_temp_budget_increase_applied_for_cached_key(): cached_after = await user_api_key_cache.async_get_cache(key=hashed_token) assert cached_after.max_budget == 2.0 + + +@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 d481f49d456..1c977c4c3aa 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -5789,3 +5789,12 @@ def test_get_sanitized_user_information_from_key_drops_callback_config(): # UserAPIKeyAuth is the live auth object; the per-key callbacks are resolved # from it during pre-call, so it must not be mutated by building the log view 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__