From d542c82f0eb92a787a34ae768c236504f57bdb15 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 20 Aug 2026 14:37:09 -0400 Subject: [PATCH 1/2] fix(otel): route Phoenix traces to per-key/team projects under otel v2 (#36706) * feat(otel): route Phoenix traces to per-key/team projects under otel v2 The v2 arize_phoenix preset read PHOENIX_PROJECT_NAME once at startup into a static resource attribute, silently dropping the per-key/team project routing v1 supported. Route it via Phoenix's x-project-name OTLP/HTTP header instead: the env var stays the global default, and a phoenix_project_name (or phoenix_project_name_override) in key/team metadata sends that key's traces to the named project. The project comes only from user_api_key_auth_metadata (server-set at auth), never from client request metadata or StandardCallbackDynamicParams, since choosing the telemetry destination is a data-exfiltration primitive. The header is appended to the exporter's static headers rather than replacing them, so the preset's Authorization survives, and it is gated to OTLP/HTTP exporters because Phoenix only reads it on /v1/traces. Also unban the bare phoenix_project_name fields from the request-body gate: the proxy integrations ignore them (only user_api_key_auth_metadata routes, and that stays banned), so rejecting them just broke SDK-style callers. * fix(otel): root project-routed Phoenix spans in their own trace Phoenix assigns a whole trace to one project by whichever span arrives first. The request's auth/db/root spans always export through the default provider without the project header, so a project-routed LLM span parented into that trace got dragged back into the default project and the header did nothing (verified against a live Phoenix instance). Detach the routed span into its own trace with a link back to the request trace, mirroring how the v1 Phoenix logger exported each request under its own local parent. * fix(otel): drain in-flight spans before shutting down evicted providers LRU eviction shut a routed provider down immediately, but an LLM span opened at pre_call stays open until the later success or failure callback; with more than 256 overlapping credential/project routes that in-flight span was silently dropped instead of exported. Refcount open spans per provider (hold at span open, release when the carrier is removed on close, carrier-map eviction, or MCP stray-carrier cleanup) and defer a retired provider's shutdown until its last open span closes. * fix(otel): take the provider hold inside route_for to close the eviction race pre_call can run on thread-pool workers, so between route_for returning a provider and the caller recording its open span, a concurrent request could overflow the LRU and shut that provider down with a zero span count, dropping the routed trace. route_for now increments the open-span count in the same locked critical section as the cache update and hands back an already-held provider; every caller releases it once its span has landed. The lock also makes the cache mutations safe under that same thread-pool concurrency. * fix(otel): skip tenant routing on deferred pre_call route_for ran before the recordable-parent check, so a thread-pool pre_call still built or LRU-touched a tenant provider and could evict an idle one even though the hold was released immediately and close re-routed. Only route when the span actually opens * add somethign * Revert "add somethign" This reverts commit 2f2cf84c5a049f0e287efd18a664479bd827fe2c. * fix(otel): cap retired tenant providers draining open spans * docs(otel): justify the retired-provider cap --- litellm/integrations/otel/logger.py | 176 ++++++++--- litellm/integrations/otel/model/metadata.py | 66 ++++- litellm/integrations/otel/plumbing/routing.py | 275 +++++++++++++++--- litellm/integrations/otel/presets/__init__.py | 42 ++- litellm/integrations/otel/presets/phoenix.py | 32 ++ litellm/proxy/auth/auth_utils.py | 6 +- .../integrations/otel/test_otel_v2_dynamic.py | 234 +++++++++++++-- .../integrations/otel/test_otel_v2_logger.py | 213 ++++++++++++++ .../proxy/auth/test_auth_utils.py | 27 +- 9 files changed, 954 insertions(+), 117 deletions(-) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index a0b5aff559f..53b9829023c 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -9,7 +9,15 @@ from typing import TYPE_CHECKING, Any, Final, cast from opentelemetry.context import Context, attach, get_current from opentelemetry.sdk._logs import LoggerProvider from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.trace import Span, Tracer, get_current_span, use_span +from opentelemetry.trace import ( + INVALID_SPAN, + Link, + Span, + Tracer, + get_current_span, + set_span_in_context, + use_span, +) import litellm from litellm._logging import verbose_logger @@ -21,6 +29,7 @@ from litellm.integrations.otel.model.config import OpenTelemetryV2Config from litellm.integrations.otel.model.metadata import ( LLMCallEvent, RequestIdentity, + auth_metadata, model_from_request_data, ) from litellm.integrations.otel.model.payloads import ( @@ -118,6 +127,12 @@ _OTEL_MODULES: Final = ( _OPEN_CALLS_MAX: Final = 10_000 +def _request_trace_links(context: Context | None) -> tuple[Link, ...] | None: + """A link back to the request trace, for a span detached into its own trace.""" + anchor: Final = get_current_span(context).get_span_context() + return (Link(anchor),) if anchor.is_valid else None + + class _LLMCallSpan: """The state carried from the ``pre_call`` boundary to span close. @@ -127,13 +142,24 @@ class _LLMCallSpan: own (worker-copied) ambient context using ``start_time_ns``. The presence of a carrier for a call at all is the proof that ``pre_call`` ran, i.e. that an upstream call was actually attempted. + + ``provider`` is the routed provider the live span was opened on (``None`` on + the default route or when creation was deferred). It is held in the tenant + cache while the span is open so LRU eviction can't shut the provider down + under it, and must be released exactly once when the carrier is removed. """ - __slots__ = ("span", "start_time_ns") + __slots__ = ("provider", "span", "start_time_ns") - def __init__(self, span: "Span | None", start_time_ns: int | None) -> None: + def __init__( + self, + span: "Span | None", + start_time_ns: int | None, + provider: "TracerProvider | None" = None, + ) -> None: self.span = span self.start_time_ns = start_time_ns + self.provider = provider class OpenTelemetryV2(CustomLogger): @@ -258,26 +284,37 @@ class OpenTelemetryV2(CustomLogger): if call_id in self._open_llm_calls: return start_time_ns: Final = to_ns(datetime.now()) - span: Span | None = None # Parent to the request's anchored root span (stable across the request), # falling back to ambient on the SDK path. Open the span live only when # that resolves to a recordable parent; otherwise defer to the close # callback (the thread-pool case, where the anchor isn't visible here). + # Do not route on the deferred path: creating or LRU-touching a tenant + # provider here would evict idle ones even though close re-routes. parent_context: Final = resolve_request_span_context() - if is_recordable_span(get_current_span(parent_context)): - span = self._emitter.start_span( + if not is_recordable_span(get_current_span(parent_context)): + self._store_open_call(call_id, _LLMCallSpan(span=None, start_time_ns=start_time_ns)) + return + # A detached route roots its own trace instead (linked to the request + # trace) — see ``TenantRoute.detached``. + route: Final = self._tenant_tracers.route_for(self.tracer, call.dynamic_params, call.auth_metadata) + try: + span: Final = self._emitter.start_span( SpanRole.LLM_CALL, call.provisional_span_name, - parent_context=parent_context, + parent_context=( + set_span_in_context(INVALID_SPAN, parent_context) if route.detached else parent_context + ), start_time_ns=start_time_ns, - tracer=self._tenant_tracers.tracer_for(self.tracer, call.dynamic_params), + tracer=route.tracer, + links=_request_trace_links(parent_context) if route.detached else None, ) - self._open_llm_calls[call_id] = _LLMCallSpan(span=span, start_time_ns=start_time_ns) - # Evict the oldest open call if the map is over budget. A call that opens - # but never closes (a stream that only fires stream events) would linger - # otherwise; the evicted span is simply dropped (never exported). - if len(self._open_llm_calls) > _OPEN_CALLS_MAX: - self._open_llm_calls.popitem(last=False) + except BaseException: + self._tenant_tracers.release(route.provider) + raise + self._store_open_call( + call_id, + _LLMCallSpan(span=span, start_time_ns=start_time_ns, provider=route.provider), + ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): if self._emit_mcp_tool_call(kwargs, start_time, end_time): @@ -371,17 +408,22 @@ class OpenTelemetryV2(CustomLogger): # otherwise linger until evicted; drop it so it's neither leaked nor closed # as a phantom LLM span. if data.identity.call_id: - self._open_llm_calls.pop(data.identity.call_id, None) - parent_context, links = resolve_mcp_span_context() - parent_context = self._seed_identity_baggage(data.identity, None, parent_context) - self._emitter.emit( - SpanRole.MCP_TOOL_CALL, - data, - parent_context=parent_context, - start_time_ns=to_ns(start_time), - end_time_ns=to_ns(end_time), - links=links, - ) + self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None)) + route: Final = self._tenant_tracers.route_for(self.tracer, None, auth_metadata(payload, kwargs)) + try: + parent_context, links = resolve_mcp_span_context() + seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context) + self._emitter.emit( + SpanRole.MCP_TOOL_CALL, + data, + parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded), + start_time_ns=to_ns(start_time), + end_time_ns=to_ns(end_time), + links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links), + tracer=route.tracer, + ) + finally: + self._tenant_tracers.release(route.provider) return True def _emit_mcp_list_tools( @@ -407,17 +449,22 @@ class OpenTelemetryV2(CustomLogger): payload, capture_content=self.config.capture_span_content ) if data.identity.call_id: - self._open_llm_calls.pop(data.identity.call_id, None) - parent_context, links = resolve_mcp_span_context() - parent_context = self._seed_identity_baggage(data.identity, None, parent_context) - self._emitter.emit( - SpanRole.MCP_LIST_TOOLS, - data, - parent_context=parent_context, - start_time_ns=to_ns(start_time), - end_time_ns=to_ns(end_time), - links=links, - ) + self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None)) + route: Final = self._tenant_tracers.route_for(self.tracer, None, auth_metadata(payload, kwargs)) + try: + parent_context, links = resolve_mcp_span_context() + seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context) + self._emitter.emit( + SpanRole.MCP_LIST_TOOLS, + data, + parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded), + start_time_ns=to_ns(start_time), + end_time_ns=to_ns(end_time), + links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links), + tracer=route.tracer, + ) + finally: + self._tenant_tracers.release(route.provider) return True def _close_llm_call( @@ -439,6 +486,36 @@ class OpenTelemetryV2(CustomLogger): carrier: Final = self._open_llm_calls.pop(call_id, None) if call_id else None if carrier is None: return None + try: + return self._finish_carrier(carrier, call, end_time) + finally: + # After the span has ended, so a release-triggered provider shutdown + # force-flushes it out rather than racing its enqueue. + self._release_carrier(carrier) + + def _store_open_call(self, call_id: str, carrier: _LLMCallSpan) -> None: + """Remember an in-flight LLM call, evicting the oldest if over budget. + + A call that opens but never closes (a stream that only fires stream + events) would linger otherwise; the evicted span is simply dropped + (never exported). + """ + self._open_llm_calls[call_id] = carrier + if len(self._open_llm_calls) > _OPEN_CALLS_MAX: + _, evicted = self._open_llm_calls.popitem(last=False) + self._release_carrier(evicted) + + def _release_carrier(self, carrier: "_LLMCallSpan | None") -> None: + """Release the routed provider a removed carrier was holding open.""" + if carrier is not None: + self._tenant_tracers.release(carrier.provider) + + def _finish_carrier( + self, + carrier: _LLMCallSpan, + call: LLMCallEvent, + end_time: datetime | float | None, + ) -> Span | None: payload: Final = call.payload if payload is None: if carrier.span is not None: @@ -462,16 +539,23 @@ class OpenTelemetryV2(CustomLogger): # The worker copied the request task's context, which carries the anchored # root span — parent to it (ambient fallback on the SDK path). Seed identity # Baggage so the span — and the SDK path, which has none — is labeled - # consistently. - parent_ctx = self._seed_identity_baggage(data.identity, data.request_model, resolve_request_span_context()) - return self._emitter.emit( - SpanRole.LLM_CALL, - data, - parent_context=parent_ctx, - start_time_ns=carrier.start_time_ns, - end_time_ns=end_time_ns, - tracer=self._tenant_tracers.tracer_for(self.tracer, call.dynamic_params), - ) + # consistently. A detached route roots its own trace instead, linked back. + route: Final = self._tenant_tracers.route_for(self.tracer, call.dynamic_params, call.auth_metadata) + try: + parent_ctx: Final = self._seed_identity_baggage( + data.identity, data.request_model, resolve_request_span_context() + ) + return self._emitter.emit( + SpanRole.LLM_CALL, + data, + parent_context=(set_span_in_context(INVALID_SPAN, parent_ctx) if route.detached else parent_ctx), + start_time_ns=carrier.start_time_ns, + end_time_ns=end_time_ns, + tracer=route.tracer, + links=_request_trace_links(parent_ctx) if route.detached else None, + ) + finally: + self._tenant_tracers.release(route.provider) # ====================================================================== # # Service hooks diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index c7cdaae0417..de6366e7dbd 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -36,8 +36,9 @@ model. They coincide on the SDK path, which is correct. from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from dataclasses import dataclass, field +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL @@ -195,6 +196,10 @@ class LLMCallEvent: # The ``standard_callback_dynamic_params`` routing the call to a per-tenant # tracer (its own exporter/endpoint), or ``None`` when the call isn't scoped. dynamic_params: Any + # The key/team config the proxy resolved at auth (``user_api_key_auth_metadata``), + # routing the call to that tenant's telemetry project. Server-set and so + # trusted, unlike ``dynamic_params``, which carries client-supplied metadata. + auth_metadata: Mapping[str, str] | None # True for synthetic proxy-gate logs (auth / rate-limit rejections): they fire # the ``pre_call`` hook but never made an upstream call, so they get no span. is_no_upstream_call: bool @@ -214,6 +219,7 @@ class LLMCallEvent: call_id=_call_id(payload, kwargs), payload=payload, dynamic_params=kwargs.get("standard_callback_dynamic_params"), + auth_metadata=auth_metadata(payload, kwargs), is_no_upstream_call=bool(kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)), provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), @@ -235,6 +241,64 @@ def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: return completion_start - api_call_start +def auth_metadata(payload: StandardLoggingPayload | None, kwargs: Mapping[str, object]) -> Mapping[str, str] | None: + """The key/team config the proxy resolved at auth, or ``None`` off the proxy. + + Read from the payload once the call closes and from ``litellm_params`` at + ``pre_call``, where no payload exists yet — the LLM-call span is *created* at + ``pre_call``, so the tracer (and therefore the destination) must be + resolvable there. Values arrive untyped, so non-string entries are dropped + rather than passed on to header builders. + """ + return next( + ( + typed + for metadata in _metadata_dicts(payload, kwargs) + if (typed := _string_entries(metadata.get("user_api_key_auth_metadata"))) + ), + None, + ) + + +def _as_str_mapping(value: object) -> Mapping[str, object] | None: + """A read-only view of ``value`` when it is a mapping, else ``None``.""" + if not isinstance(value, Mapping): + return None + return cast("Mapping[str, object]", value) # cast-ok: isinstance-guarded, JSON metadata has str keys + + +def _string_entries(value: object) -> Mapping[str, str] | None: + entries: Final = _as_str_mapping(value) + if entries is None: + return None + typed: Final = MappingProxyType({key: item for key, item in entries.items() if isinstance(item, str)}) + return typed or None + + +def _metadata_dicts( + payload: StandardLoggingPayload | None, kwargs: Mapping[str, object] +) -> Iterator[Mapping[str, object]]: + """Request metadata dicts, closed-call payload first then the live kwargs. + + ``litellm_metadata`` is the metadata field on the Anthropic-shaped routes; + litellm copies it onto ``metadata``, but both are yielded so a route that + populates only one is still covered. + """ + payload_view: Final = _as_str_mapping(payload) + if payload_view is not None: + payload_metadata: Final = _as_str_mapping(payload_view.get("metadata")) + if payload_metadata is not None: + yield payload_metadata + params: Final = _as_str_mapping(kwargs.get("litellm_params")) + if params is None: + return + yield from ( + metadata + for key in ("metadata", "litellm_metadata") + if (metadata := _as_str_mapping(params.get(key))) is not None + ) + + def _call_id(payload: StandardLoggingPayload | None, kwargs: Mapping[str, Any]) -> str | None: """The call id from the payload (when closed) or the bare kwargs (at pre_call).""" if payload is not None: diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index d9e9e84364a..f231df9e914 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -1,31 +1,45 @@ """Per-request multi-tenant tracer routing. When a request carries team/key vendor credentials in -``standard_callback_dynamic_params``, its spans must export through a -``TracerProvider`` whose OTLP headers carry those credentials. -``TenantTracerCache`` builds and caches one provider per distinct credential -set, and otherwise hands back the logger's default tracer. This lets a single -logger fan requests out to many tenants without needing a logger per tenant. +``standard_callback_dynamic_params``, or the key/team config resolved at auth +names a destination project, its spans must export through a +``TracerProvider`` whose OTLP headers carry those credentials / that project. +``TenantTracerCache`` builds and caches one provider per distinct +(credentials, project) pair, and otherwise hands back the logger's default +tracer. This lets a single logger fan requests out to many tenants without +needing a logger per tenant. """ +import threading from collections import OrderedDict from collections.abc import Mapping -from typing import Any, Final +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Final, TypeAlias +from urllib.parse import quote from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import Tracer from litellm._logging import verbose_logger -from litellm.integrations.otel.model.config import OpenTelemetryV2Config +from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config from litellm.integrations.otel.plumbing.providers import ( build_tracer_provider, get_tracer, ) -from litellm.integrations.otel.presets import dynamic_otlp_headers +from litellm.integrations.otel.presets import ( + dynamic_otlp_headers, + project_routing_headers, +) # Exporter kinds that ignore headers — never rewritten with dynamic credentials. _NON_OTLP_KINDS: Final = ("console", "in_memory", "inmemory", "memory") +# gRPC exporters still take dynamic credentials (as gRPC metadata) but not +# project headers: the routing headers backends read (Phoenix's +# ``x-project-name``) are only honored on the OTLP/HTTP endpoint. +_GRPC_KINDS: Final = ("otlp_grpc", "grpc") + # Cap on distinct credential-scoped providers held at once. ``dynamic_params`` # can be populated from request metadata, so an unbounded cache lets a caller # spawn one ``TracerProvider`` (plus its ``BatchSpanProcessor`` background @@ -34,6 +48,23 @@ _NON_OTLP_KINDS: Final = ("console", "in_memory", "inmemory", "memory") # evicted providers so their threads are reclaimed. _MAX_CACHED_PROVIDERS: Final = 256 +# Cap on providers evicted from the cache while still holding open spans, which +# are kept alive to drain instead of being shut down under them. Their only +# other bound is the logger's open-call map (10k), so without this a caller +# cycling unique credential sets across long-lived calls could pin far more +# live providers, and exporter threads, than the cache cap allows. Past this +# many, the stalest retiree is shut down and whatever it was draining is +# dropped (a shut-down ``BatchSpanProcessor`` discards spans handed to it after +# the fact), which by then means a span on a route evicted long ago. A quarter +# of the cache cap: enough that a burst of tenant churn during long-lived calls +# still drains normally, small enough that the worst case is a bounded 320 +# providers rather than one per concurrent call. +_MAX_RETIRED_PROVIDERS: Final = 64 + +_HeaderItems: TypeAlias = tuple[tuple[str, str], ...] + +_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + def _shutdown_provider(provider: TracerProvider) -> None: """Flush + stop an evicted provider's processors (reclaims their threads). @@ -49,8 +80,41 @@ def _shutdown_provider(provider: TracerProvider) -> None: verbose_logger.debug("OTel V2: error shutting down evicted provider: %s", e) +def _plain_header_string(headers: Mapping[str, str]) -> str: + return ",".join(f"{key}={value}" for key, value in headers.items()) + + +def _encoded_header_string(headers: Mapping[str, str]) -> str: + """Percent-encode values so one containing the ``k=v,k=v`` separators (e.g. + a project name with a comma) survives; ``parse_env_headers`` decodes it back. + """ + return ",".join(f"{key}={quote(value, safe='')}" for key, value in headers.items()) + + +@dataclass(frozen=True, slots=True) +class TenantRoute: + """The tracer to create a span on, plus whether it must root its own trace. + + ``detached`` is True when project routing engaged. Phoenix assigns a whole + trace to one project by whichever of its spans arrives first, so a + project-routed span parented into the request trace gets dragged into the + project of the default-exported request spans and the header does nothing. + The span must therefore start a fresh trace (with a link back to the + request trace for correlation) — which is also how the v1 Phoenix logger + behaved, exporting each request under its own Phoenix-local parent span. + """ + + tracer: Tracer + detached: bool + #: The provider ``tracer`` came from, or ``None`` on the default route. It + #: is returned already held (counted as an open span, atomically with the + #: cache update), so LRU eviction can't shut it down before the caller's + #: span lands; the caller must ``release`` it exactly once when done. + provider: TracerProvider | None = None + + class TenantTracerCache: - """Credential-scoped ``TracerProvider`` cache keyed by the dynamic headers.""" + """Credential/project-scoped ``TracerProvider`` cache keyed by the routing headers.""" def __init__( self, @@ -61,49 +125,170 @@ class TenantTracerCache: self._config = config self._callback_name = callback_name self._tracer_name = tracer_name - self._providers: OrderedDict[tuple[tuple[str, str], ...], TracerProvider] = OrderedDict() + # Guards the three mutable structures below: ``pre_call`` can run on + # thread-pool workers concurrently with the event loop, so cache + # updates, span counts, and retirement must be atomic. + self._lock: Final = threading.Lock() + self._providers: OrderedDict[tuple[_HeaderItems, _HeaderItems], TracerProvider] = OrderedDict() + self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state + # Oldest-first so an overflow of draining providers sheds the stalest. + self._retired: OrderedDict[TracerProvider, None] = OrderedDict() # mutable-ok: draining evicted providers + self._project_routable = any( + spec.owner == callback_name and spec.kind.lower() not in (*_NON_OTLP_KINDS, *_GRPC_KINDS) + for spec in config.exporters + ) + self._warned_project_unroutable = False - def tracer_for(self, default: Tracer, dynamic_params: Any) -> Tracer: - """Return the tracer for this request. + def release(self, provider: TracerProvider | None) -> None: + """Drop one open-span count; shut a retired provider down once drained. - Use ``default`` unless the request's dynamic credentials require a - credential-scoped tracer, in which case build (or reuse) one. The cache - is a bounded LRU: the least-recently-used provider is flushed and shut - down on overflow so its exporter threads don't accumulate. + ``None`` (the default route) is a no-op so callers can release a + ``TenantRoute.provider`` unconditionally. The shutdown itself runs + outside the lock: it force-flushes over the network and must not stall + every concurrently routing request. """ - headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) - if not headers: - return default - cache_key: Final = tuple(sorted(headers.items())) - provider = self._providers.get(cache_key) - if provider is not None: + if provider is None: + return + with self._lock: + remaining: Final = self._open_span_counts.get(provider, 0) - 1 + if remaining > 0: + self._open_span_counts[provider] = remaining + return + self._open_span_counts.pop(provider, None) + drained: Final = provider in self._retired + self._retired.pop(provider, None) + if drained: + _shutdown_provider(provider) + + def route_for( + self, + default: Tracer, + dynamic_params: Any, + auth_metadata: Mapping[str, str] | None = None, + ) -> TenantRoute: + """Return the tracer (and trace-detachment flag) for this request. + + Use ``default`` unless the request's dynamic credentials or its key/team + project require a scoped tracer, in which case build (or reuse) one. The + cache is a bounded LRU: the least-recently-used provider is flushed and + shut down on overflow so its exporter threads don't accumulate. + + A routed provider is returned already held — its open-span count is + incremented in the same critical section as the cache update — so a + concurrent overflow eviction can't shut it down between selection and + the caller's span start. The caller must ``release`` it exactly once. + """ + credential_headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS + project_headers: Final = self._project_headers(auth_metadata) + if not credential_headers and not project_headers: + return TenantRoute(tracer=default, detached=False) + cache_key: Final = ( + tuple(sorted(credential_headers.items())), + tuple(sorted(project_headers.items())), + ) + with self._lock: + provider: Final = self._cached_provider_locked(cache_key, credential_headers, project_headers) + self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1 + evicted: Final = self._evicted_on_overflow_locked() + if evicted is not None: + _shutdown_provider(evicted) + return TenantRoute( + tracer=get_tracer(provider, self._tracer_name), + detached=bool(project_headers), + provider=provider, + ) + + def _cached_provider_locked( + self, + cache_key: tuple[_HeaderItems, _HeaderItems], + credential_headers: Mapping[str, str], + project_headers: Mapping[str, str], + ) -> TracerProvider: + cached: Final = self._providers.get(cache_key) + if cached is not None: self._providers.move_to_end(cache_key) - else: - provider = build_tracer_provider(self._config_with_headers(headers)) - self._providers[cache_key] = provider - if len(self._providers) > _MAX_CACHED_PROVIDERS: - _, evicted = self._providers.popitem(last=False) - _shutdown_provider(evicted) - return get_tracer(provider, self._tracer_name) + return cached + built: Final = build_tracer_provider(self._routed_config(credential_headers, project_headers)) + self._providers[cache_key] = built + return built - def _config_with_headers(self, headers: Mapping[str, str]) -> OpenTelemetryV2Config: - """Clone the config, stamping ``headers`` onto the credential's own exporter. + def _evicted_on_overflow_locked(self) -> TracerProvider | None: + """Pop the LRU provider past the cap; return it if the caller must shut it down. - ``headers`` are the per-request credentials of ``self._callback_name`` (the - integration that built this cache), so they apply only to the exporter that - integration contributed (``spec.owner``). A request that carries one - tenant's Arize key must never rewrite the headers of a co-configured - Langfuse or self-hosted collector exporter, which would leak that key to a - different backend. + A provider with open spans is retired to drain instead: stopping its + processors while a span opened at ``pre_call`` is still live would + silently drop that span at end instead of exporting it. Retirees are + themselves capped, so the stalest one is shut down (and its open-span + count dropped, making its eventual ``release`` a no-op) once too many + pile up rather than letting them accumulate a thread each. """ - header_str: Final = ",".join(f"{key}={value}" for key, value in headers.items()) - header_update: Final[dict[str, str]] = {"headers": header_str} - exporters: Final = [ - ( - spec.model_copy(update=header_update) - if spec.owner == self._callback_name and spec.kind.lower() not in _NON_OTLP_KINDS - else spec + if len(self._providers) <= _MAX_CACHED_PROVIDERS: + return None + _, evicted = self._providers.popitem(last=False) + if self._open_span_counts.get(evicted, 0) == 0: + return evicted + self._retired[evicted] = None + if len(self._retired) <= _MAX_RETIRED_PROVIDERS: + return None + overflowed, _ = self._retired.popitem(last=False) + self._open_span_counts.pop(overflowed, None) + return overflowed + + def _project_headers(self, auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]: + """The per-request project-routing headers, if this cache can apply them. + + A gRPC-only exporter can't (the project header route is HTTP-only), so + the request warns once and stays on the env-configured default project. + """ + requested: Final = project_routing_headers(self._callback_name, auth_metadata) + if not requested or self._project_routable: + return requested + if not self._warned_project_unroutable: + self._warned_project_unroutable = True + verbose_logger.warning( + "OTel V2: %s key/team config names a per-request project, but its exporter " + "is not OTLP/HTTP and the project header is HTTP-only; spans stay in the " + "default project.", + self._callback_name, ) - for spec in self._config.exporters + return _NO_HEADERS + + def _routed_config( + self, + credential_headers: Mapping[str, str], + project_headers: Mapping[str, str], + ) -> OpenTelemetryV2Config: + """Clone the config, rewriting headers on the callback's own exporter. + + Both header sets apply only to the exporter ``self._callback_name`` + contributed (``spec.owner``). A request that carries one tenant's Arize + key must never rewrite the headers of a co-configured Langfuse or + self-hosted collector exporter, which would leak that key to a + different backend. + + Dynamic credentials REPLACE the exporter's headers — they are the + tenant's complete credential set. Project headers APPEND instead: the + preset's static headers carry the backend auth (Phoenix's + ``Authorization``), which must survive routing to a project. + """ + exporters: Final = [ + self._routed_exporter(spec, credential_headers, project_headers) for spec in self._config.exporters ] return self._config.model_copy(update={"exporters": exporters}) + + def _routed_exporter( + self, + spec: ExporterSpec, + credential_headers: Mapping[str, str], + project_headers: Mapping[str, str], + ) -> ExporterSpec: + kind: Final = spec.kind.lower() + if spec.owner != self._callback_name or kind in _NON_OTLP_KINDS: + return spec + base: Final = _plain_header_string(credential_headers) if credential_headers else spec.headers + routed: Final = ( + ",".join(part for part in (base, _encoded_header_string(project_headers)) if part) + if project_headers and kind not in _GRPC_KINDS + else base + ) + return spec if routed == spec.headers else spec.model_copy(update={"headers": routed}) diff --git a/litellm/integrations/otel/presets/__init__.py b/litellm/integrations/otel/presets/__init__.py index 95ac2783325..35b0584c697 100644 --- a/litellm/integrations/otel/presets/__init__.py +++ b/litellm/integrations/otel/presets/__init__.py @@ -8,7 +8,8 @@ the factory in ``litellm_logging`` can resolve a name and build a single ``OpenTelemetryV2`` instance from the result. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import Final from litellm.integrations.otel.presets.agentops import agentops_preset @@ -20,7 +21,10 @@ from litellm.integrations.otel.presets.langfuse import ( ) from litellm.integrations.otel.presets.langtrace import langtrace_preset from litellm.integrations.otel.presets.levo import levo_preset -from litellm.integrations.otel.presets.phoenix import phoenix_preset +from litellm.integrations.otel.presets.phoenix import ( + phoenix_preset, + phoenix_project_headers, +) from litellm.integrations.otel.presets.weave import weave_dynamic_headers, weave_preset from litellm.types.utils import StandardCallbackDynamicParams @@ -47,6 +51,23 @@ DYNAMIC_HEADERS_BY_CALLBACK: Final[dict[str, Callable[[StandardCallbackDynamicPa } +#: Callback name → per-request *routing* header builder, sourced from the key/team +#: config the proxy resolved at auth. Deliberately separate from +#: ``DYNAMIC_HEADERS_BY_CALLBACK``: that one is fed +#: ``StandardCallbackDynamicParams``, which is populated from client-supplied +#: request metadata. Naming a destination project is a data-exfiltration +#: primitive, so it must only ever come from server-set key/team config. +PROJECT_HEADERS_BY_CALLBACK: Final[Mapping[str, Callable[[Mapping[str, str] | None], Mapping[str, str]]]] = ( + MappingProxyType( + { + "arize_phoenix": phoenix_project_headers, + } + ) +) + +_NO_PROJECT_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + + def dynamic_otlp_headers( callback_name: str | None, dynamic_params: StandardCallbackDynamicParams | None, @@ -62,9 +83,25 @@ def dynamic_otlp_headers( return headers or None +def project_routing_headers( + callback_name: str | None, + auth_metadata: Mapping[str, str] | None, +) -> Mapping[str, str]: + """Per-request project-routing headers from trusted key/team config. + + Empty means "no per-request project" — the caller keeps its default tracer, + whose resource attributes carry the env-configured project. + """ + builder: Final = PROJECT_HEADERS_BY_CALLBACK.get(callback_name or "") + if builder is None: + return _NO_PROJECT_HEADERS + return builder(auth_metadata) + + __all__ = [ "DYNAMIC_HEADERS_BY_CALLBACK", "PRESET_BY_CALLBACK", + "PROJECT_HEADERS_BY_CALLBACK", "Preset", "agentops_preset", "arize_preset", @@ -73,5 +110,6 @@ __all__ = [ "langtrace_preset", "levo_preset", "phoenix_preset", + "project_routing_headers", "weave_preset", ] diff --git a/litellm/integrations/otel/presets/phoenix.py b/litellm/integrations/otel/presets/phoenix.py index fc1eb9f748f..eef407b6c1b 100644 --- a/litellm/integrations/otel/presets/phoenix.py +++ b/litellm/integrations/otel/presets/phoenix.py @@ -1,5 +1,7 @@ """Arize-Phoenix preset.""" +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from pydantic import AliasChoices, Field @@ -25,6 +27,36 @@ class _PhoenixSettings(BaseSettings): ) +#: Phoenix routes an OTLP/HTTP export to a project by this header, which takes +#: precedence over the ``openinference.project.name`` resource attribute the env +#: var sets. Requires arize-phoenix 15.5.0+; older collectors ignore it and the +#: spans land in the resource attribute's project. +PHOENIX_PROJECT_HEADER: Final = "x-project-name" + +#: Key/team config fields naming the target project, highest precedence first. +_PROJECT_KEYS: Final = ("phoenix_project_name_override", "phoenix_project_name") + +_NO_PROJECT: Final[Mapping[str, str]] = MappingProxyType({}) + + +def phoenix_project_headers(auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]: + """The per-request Phoenix project header for this key/team, if any. + + ``auth_metadata`` must be the key/team config the proxy resolved at auth + (``user_api_key_auth_metadata``), never client-supplied request metadata: + choosing the destination project is a data-exfiltration primitive, so a + caller must not be able to name one. Returns an empty mapping when the key + and team name no project, leaving the request on the env-configured default. + """ + if not auth_metadata: + return _NO_PROJECT + project: Final = next( + (stripped for key in _PROJECT_KEYS if (stripped := (auth_metadata.get(key) or "").strip())), + "", + ) + return MappingProxyType({PHOENIX_PROJECT_HEADER: project}) if project else _NO_PROJECT + + def phoenix_preset( *, config_overrides: OpenTelemetryV2Config | None = None, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 883f986f6fd..912a0b0ebd0 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -222,8 +222,10 @@ _SAFE_CLIENT_CALLBACK_PARAMS: Final[frozenset[str]] = frozenset( _EXTRA_BANNED_OBSERVABILITY_PARAMS: Final[frozenset[str]] = frozenset( { "posthog_api_url", - "phoenix_project_name", - "phoenix_project_name_override", + # ``phoenix_project_name`` / ``phoenix_project_name_override`` are NOT + # banned: on the proxy the Phoenix integrations only read them from + # ``user_api_key_auth_metadata`` (key/team config), so the bare request + # fields are inert and rejecting them just breaks SDK-style callers. # Server-reserved: written exclusively by add_user_api_key_auth_to_request_metadata # from the authenticated key's database record. A caller-supplied value # would survive the server merge and let an authenticated user redirect diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py index e44c56e1fdf..ca62253aa2f 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py @@ -9,7 +9,11 @@ sys.path.insert(0, os.path.abspath("../../../..")) from opentelemetry.trace import NoOpTracer from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config -from litellm.integrations.otel.presets import dynamic_otlp_headers +from litellm.integrations.otel.presets import ( + dynamic_otlp_headers, + project_routing_headers, +) +from litellm.integrations.otel.plumbing.providers import parse_headers from litellm.integrations.otel.plumbing.routing import TenantTracerCache @@ -83,10 +87,10 @@ def test_provider_cached_per_credential_set(): creds_a = {"arize_space_id": "S", "arize_api_key": "K"} creds_b = {"arize_space_id": "S2", "arize_api_key": "K2"} - cache.tracer_for(default, creds_a) - cache.tracer_for(default, creds_a) # same set → reuse, no new provider + cache.route_for(default, creds_a) + cache.route_for(default, creds_a) # same set → reuse, no new provider assert len(cache._providers) == 1 - cache.tracer_for(default, creds_b) # new set → new provider + cache.route_for(default, creds_b) # new set → new provider assert len(cache._providers) == 2 @@ -109,10 +113,11 @@ def test_provider_cache_is_bounded_and_evicts_lru(monkeypatch): def creds(space): return {"arize_space_id": space, "arize_api_key": "K"} - cache.tracer_for(default, creds("1")) - cache.tracer_for(default, creds("2")) - cache.tracer_for(default, creds("1")) # touch "1" → "2" is now LRU - cache.tracer_for(default, creds("3")) # overflow → evict "2" + # route_for returns a held provider; release models the span closing. + cache.release(cache.route_for(default, creds("1")).provider) + cache.release(cache.route_for(default, creds("2")).provider) + cache.release(cache.route_for(default, creds("1")).provider) # touch "1" → "2" is now LRU + cache.release(cache.route_for(default, creds("3")).provider) # overflow → evict "2" assert len(cache._providers) == 2 assert len(shut_down) == 1 # exactly the evicted provider was shut down @@ -121,14 +126,14 @@ def test_provider_cache_is_bounded_and_evicts_lru(monkeypatch): def test_no_dynamic_params_uses_default_tracer(): cache = _cache("arize") default = NoOpTracer() - assert cache.tracer_for(default, {}) is default + assert cache.route_for(default, {}).tracer is default assert cache._providers == {} def test_non_participating_callback_uses_default_tracer(): cache = _cache("arize_phoenix") default = NoOpTracer() - assert cache.tracer_for(default, {"arize_api_key": "K"}) is default + assert cache.route_for(default, {"arize_api_key": "K"}).tracer is default assert cache._providers == {} @@ -140,7 +145,7 @@ def test_dynamic_headers_applied_to_otlp_exporter_only(): ExporterSpec(kind="in_memory", owner="arize"), ], ) - new_cfg = cache._config_with_headers({"arize-space-id": "S", "api_key": "K"}) + new_cfg = cache._routed_config({"arize-space-id": "S", "api_key": "K"}, {}) otlp, in_mem = new_cfg.exporters assert otlp.headers == "arize-space-id=S,api_key=K" assert in_mem.headers is None # console/in_memory left untouched @@ -150,10 +155,9 @@ def test_dynamic_headers_do_not_leak_to_other_owners_exporter(): """A tenant's Arize credentials must never be stamped onto a co-configured exporter owned by a different backend (a self-hosted collector, Langfuse). - Regression for the cross-backend credential leak: ``_config_with_headers`` - used to rewrite the headers of every OTLP exporter, so one request carrying - a team's Arize key clobbered the base collector's and Langfuse's headers - with that key. + Regression for the cross-backend credential leak: the header rewrite used + to hit every OTLP exporter, so one request carrying a team's Arize key + clobbered the base collector's and Langfuse's headers with that key. """ cache = _cache( "arize", @@ -178,10 +182,206 @@ def test_dynamic_headers_do_not_leak_to_other_owners_exporter(): ), ], ) - new_cfg = cache._config_with_headers( - {"arize-space-id": "TEAMX", "api_key": "TEAMX_KEY"} + new_cfg = cache._routed_config( + {"arize-space-id": "TEAMX", "api_key": "TEAMX_KEY"}, {} ) by_owner = {e.owner: e.headers for e in new_cfg.exporters} assert by_owner["arize"] == "arize-space-id=TEAMX,api_key=TEAMX_KEY" assert by_owner[None] == "x=base-collector" assert by_owner["langfuse_otel"] == "Authorization=Basic base-langfuse" + + +# --- per-request Phoenix project routing from trusted key/team config --- # + + +def _phoenix_cache(kind="otlp_http"): + return _cache( + "arize_phoenix", + exporters=[ + ExporterSpec( + kind=kind, + endpoint="http://phoenix:6006", + headers="Authorization=Bearer phoenix-key", + owner="arize_phoenix", + ), + ], + ) + + +def test_phoenix_project_headers_precedence_and_blanks(): + assert project_routing_headers( + "arize_phoenix", {"phoenix_project_name": "team-proj"} + ) == {"x-project-name": "team-proj"} + assert project_routing_headers( + "arize_phoenix", + {"phoenix_project_name_override": "override", "phoenix_project_name": "base"}, + ) == {"x-project-name": "override"} + assert ( + project_routing_headers("arize_phoenix", {"phoenix_project_name": " "}) == {} + ) + assert project_routing_headers("arize_phoenix", None) == {} + # Only Phoenix participates in project routing. + assert project_routing_headers("arize", {"phoenix_project_name": "p"}) == {} + + +def test_project_header_appends_and_preserves_phoenix_auth(): + """Regression: routing to a project must not drop the preset's static + ``Authorization`` header — a replace would break Phoenix auth entirely.""" + cache = _phoenix_cache() + cfg = cache._routed_config({}, {"x-project-name": "team-proj"}) + (spec,) = cfg.exporters + parsed = parse_headers(spec.headers) + assert parsed["authorization"] == "Bearer phoenix-key" + assert parsed["x-project-name"] == "team-proj" + + +def test_project_name_with_header_separators_round_trips(): + cache = _phoenix_cache() + cfg = cache._routed_config({}, {"x-project-name": "my proj, prod=1"}) + (spec,) = cfg.exporters + parsed = parse_headers(spec.headers) + assert parsed["x-project-name"] == "my proj, prod=1" + assert parsed["authorization"] == "Bearer phoenix-key" + + +def test_project_header_does_not_touch_other_exporters(): + cache = _cache( + "arize_phoenix", + exporters=[ + ExporterSpec( + kind="otlp_http", + endpoint="http://collector:4318", + headers="x=base-collector", + owner=None, + ), + ExporterSpec( + kind="otlp_http", + endpoint="http://phoenix:6006", + headers="Authorization=Bearer phoenix-key", + owner="arize_phoenix", + ), + ], + ) + cfg = cache._routed_config({}, {"x-project-name": "team-proj"}) + by_owner = {e.owner: e.headers for e in cfg.exporters} + assert by_owner[None] == "x=base-collector" + assert parse_headers(by_owner["arize_phoenix"])["x-project-name"] == "team-proj" + + +def test_provider_cached_per_project(): + cache = _phoenix_cache() + default = NoOpTracer() + routed = cache.route_for(default, None, {"phoenix_project_name": "proj-a"}) + assert routed.tracer is not default + assert routed.detached is True # project spans must root their own trace + cache.route_for(default, None, {"phoenix_project_name": "proj-a"}) + assert len(cache._providers) == 1 + cache.route_for(default, None, {"phoenix_project_name": "proj-b"}) + assert len(cache._providers) == 2 + for provider in cache._providers.values(): + provider.shutdown() + + +def test_client_dynamic_params_cannot_choose_phoenix_project(): + # ``StandardCallbackDynamicParams`` is populated from client-supplied + # request metadata; the project may only come from server-set key/team + # config (the ``auth_metadata`` argument). + cache = _phoenix_cache() + default = NoOpTracer() + assert cache.route_for(default, {"phoenix_project_name": "attacker"}).tracer is default + assert ( + cache.route_for(default, {"phoenix_project_name_override": "attacker"}).tracer + is default + ) + assert cache._providers == {} + + +def test_auth_metadata_without_project_uses_default_tracer(): + cache = _phoenix_cache() + default = NoOpTracer() + assert cache.route_for(default, None, {"logging_setting": "x"}).tracer is default + assert cache._providers == {} + + +def test_grpc_exporter_gets_no_project_routing(): + # ``x-project-name`` is only honored on the OTLP/HTTP endpoint, so a + # gRPC-only Phoenix exporter stays on the default project (warned once). + cache = _phoenix_cache(kind="otlp_grpc") + default = NoOpTracer() + assert cache.route_for(default, None, {"phoenix_project_name": "proj"}).tracer is default + assert cache._providers == {} + assert cache._warned_project_unroutable is True + + +def test_eviction_defers_shutdown_while_a_span_is_open(monkeypatch): + # An LLM span opened at pre_call stays open until the close callback; LRU + # eviction in that window must not stop the provider's processors, or the + # span is silently dropped at end instead of exported. route_for itself + # takes the hold, atomically with the cache update, so a concurrent + # eviction can never shut a just-selected provider down before the caller + # records its span. + from litellm.integrations.otel.plumbing import routing as routing_mod + + monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) + shut_down = [] + monkeypatch.setattr( + routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) + ) + cache = _cache("arize") + default = NoOpTracer() + + route_a = cache.route_for(default, {"arize_space_id": "A", "arize_api_key": "K"}) + assert route_a.provider is not None + cache.route_for(default, {"arize_space_id": "B", "arize_api_key": "K"}) # evicts A + assert shut_down == [] # deferred: A is still held by route_a + cache.release(route_a.provider) + assert shut_down == [route_a.provider] + + +def test_retired_providers_are_capped(monkeypatch): + # Retiring an evicted provider keeps it, and its exporter thread, alive + # while a span is open, so retirees need a cap of their own: a caller + # cycling unique credential sets across calls that never close would + # otherwise pin one live provider per open call, far past the cache bound. + # Past the cap the stalest retiree is shut down and its later release is a + # no-op, while the ones still within the cap keep draining. + from litellm.integrations.otel.plumbing import routing as routing_mod + + monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) + monkeypatch.setattr(routing_mod, "_MAX_RETIRED_PROVIDERS", 2) + shut_down = [] + monkeypatch.setattr( + routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) + ) + cache = _cache("arize") + default = NoOpTracer() + + # Every route stays held (no release), so each one evicts and retires its + # predecessor instead of shutting it down. + routes = [ + cache.route_for(default, {"arize_space_id": str(i), "arize_api_key": "K"}) + for i in range(5) + ] + + assert len(cache._providers) == 1 + assert len(cache._retired) == 2 # capped, not one retiree per open call + assert shut_down == [routes[0].provider, routes[1].provider] + + cache.release(routes[0].provider) # already shut down: no second shutdown + assert shut_down == [routes[0].provider, routes[1].provider] + cache.release(routes[2].provider) # still draining: drains and shuts down + assert shut_down[-1] is routes[2].provider + + +def test_release_without_eviction_keeps_provider_alive(monkeypatch): + from litellm.integrations.otel.plumbing import routing as routing_mod + + shut_down = [] + monkeypatch.setattr( + routing_mod, "_shutdown_provider", lambda p: shut_down.append(p) + ) + cache = _cache("arize") + route = cache.route_for(NoOpTracer(), {"arize_space_id": "A", "arize_api_key": "K"}) + cache.release(route.provider) + assert shut_down == [] # still cached, never retired + cache.release(None) # default-route release is a no-op diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index bb2d970e9c7..e5d5b62b856 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -37,6 +37,7 @@ from litellm.integrations.otel.plumbing.context import ( # noqa: E402 set_request_root_span, ) from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 +from litellm.integrations.otel.model.config import ExporterSpec # noqa: E402 from litellm.integrations.otel.model.spans import ( # noqa: E402 LITELLM_PROXY_REQUEST_SPAN_NAME, SpanRole, @@ -2309,3 +2310,215 @@ def test_metrics_disabled_by_default_records_nothing(monkeypatch): ) ) assert _emitted_metric_names(reader) == set() + + +# --------------------------------------------------------------------------- # +# Per-request Phoenix project routing (key/team auth metadata) +# --------------------------------------------------------------------------- # + + +def _phoenix_routing_logger(capture_kind): + """A Phoenix-shaped logger whose owned exporter is a registered factory kind + that captures the exporter built per routed header set, so the test can + assert which destination each span actually exported through.""" + captured = {} + + def factory(spec): + exporter = InMemorySpanExporter() + captured[spec.headers] = exporter + return exporter + + providers.register_exporter_factory(capture_kind, factory) + cfg = OpenTelemetryV2Config( + exporters=[ + ExporterSpec( + kind=capture_kind, + endpoint="http://phoenix:6006", + headers="Authorization=Bearer phoenix-key", + owner="arize_phoenix", + ) + ] + ) + default_exporter = InMemorySpanExporter() + tracer_provider = providers.build_tracer_provider(cfg, exporter=default_exporter) + logger = OpenTelemetryV2( + config=cfg, callback_name="arize_phoenix", tracer_provider=tracer_provider + ) + return logger, default_exporter, captured + + +def test_key_team_auth_metadata_routes_llm_span_to_phoenix_project(): + """The proxy stamps the key/team config into ``user_api_key_auth_metadata``; + a ``phoenix_project_name`` there must route the LLM span through an exporter + carrying the ``x-project-name`` header while keeping the preset's auth.""" + logger, default_exporter, captured = _phoenix_routing_logger("capture_route_a") + auth_md = {"phoenix_project_name": "team-proj"} + payload = _payload(metadata={"user_api_key_auth_metadata": auth_md}) + kwargs = { + "standard_logging_object": payload, + "litellm_params": {"metadata": {"user_api_key_auth_metadata": auth_md}}, + } + _emit_llm(logger, kwargs) + + assert [s.name for s in default_exporter.get_finished_spans()] == [] + (headers,) = captured + parsed = providers.parse_headers(headers) + assert parsed["x-project-name"] == "team-proj" + assert parsed["authorization"] == "Bearer phoenix-key" + routed_spans = captured[headers].get_finished_spans() + assert len(routed_spans) == 1 + assert routed_spans[0].parent is None # own trace, so Phoenix can route it + + +def test_client_request_metadata_cannot_route_phoenix_project(): + """A bare ``phoenix_project_name`` in client request metadata (not the + server-set ``user_api_key_auth_metadata``) must be ignored: the span stays + on the default tracer and no routed exporter is ever built.""" + logger, default_exporter, captured = _phoenix_routing_logger("capture_route_b") + payload = _payload(metadata={"phoenix_project_name": "attacker-project"}) + kwargs = { + "standard_logging_object": payload, + "litellm_params": {"metadata": {"phoenix_project_name": "attacker-project"}}, + } + _emit_llm(logger, kwargs) + + assert captured == {} + assert len(default_exporter.get_finished_spans()) == 1 + + +def test_project_routing_resolves_at_pre_call_before_payload_exists(): + """Production ``pre_call`` runs before the standard logging payload exists, + so the destination project must resolve from ``litellm_params`` alone — the + span is created (and its exporter chosen) right there.""" + logger, default_exporter, captured = _phoenix_routing_logger("capture_route_c") + auth_md = {"phoenix_project_name": "team-proj"} + litellm_params = {"metadata": {"user_api_key_auth_metadata": auth_md}} + server = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + with trace.use_span(server, end_on_exit=False): + logger.log_pre_api_call( + model="gpt-4o", + messages=[], + kwargs={"litellm_call_id": "call_1", "litellm_params": litellm_params}, + ) + server.end() + assert len(captured) == 1 # routed exporter already built at pre_call + + close_kwargs = { + "standard_logging_object": _payload( + metadata={"user_api_key_auth_metadata": auth_md} + ), + "litellm_params": litellm_params, + } + asyncio.run(logger.async_log_success_event(close_kwargs, None, None, None)) + + (headers,) = captured + (routed_span,) = captured[headers].get_finished_spans() + assert routed_span.name == "chat gpt-4o" + # Phoenix pins a whole trace to one project by its first-arriving span, so + # the routed span must root its OWN trace, linked back to the request trace. + assert routed_span.parent is None + (link,) = routed_span.links + assert link.context.span_id == server.get_span_context().span_id + assert all( + s.name != "chat gpt-4o" for s in default_exporter.get_finished_spans() + ) + + +def test_evicted_provider_still_exports_span_opened_before_eviction(monkeypatch): + """LRU eviction while a routed span is still open must defer the provider + shutdown: the span opened at ``pre_call`` closes at the later success + callback and would otherwise be silently dropped instead of exported.""" + from litellm.integrations.otel.plumbing import routing as routing_mod + + monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) + logger, _default_exporter, captured = _phoenix_routing_logger("capture_evict") + md_a = {"user_api_key_auth_metadata": {"phoenix_project_name": "proj-a"}} + server = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + with trace.use_span(server, end_on_exit=False): + logger.log_pre_api_call( + model="gpt-4o", + messages=[], + kwargs={"litellm_call_id": "call_a", "litellm_params": {"metadata": md_a}}, + ) + + # A second project's full call overflows the size-1 LRU and evicts proj-a's + # provider while call_a's span is still open. + md_b = {"user_api_key_auth_metadata": {"phoenix_project_name": "proj-b"}} + _emit_llm( + logger, + { + "standard_logging_object": _payload(litellm_call_id="call_b", metadata=md_b), + "litellm_params": {"metadata": md_b}, + }, + ) + + asyncio.run( + logger.async_log_success_event( + { + "standard_logging_object": _payload(litellm_call_id="call_a", metadata=md_a), + "litellm_params": {"metadata": md_a}, + }, + None, + None, + None, + ) + ) + server.end() + + headers_a = next(h for h in captured if "proj-a" in h) + assert [s.name for s in captured[headers_a].get_finished_spans()] == ["chat gpt-4o"] + + +def test_deferred_pre_call_does_not_churn_tenant_cache(monkeypatch): + """Deferred ``pre_call`` must not create or LRU-touch a tenant provider. + + ``route_for`` used to run before the recordable-parent check, so a + thread-pool ``pre_call`` that immediately released its hold still built a + provider and could evict an idle one. Close re-routes when the span + actually opens. + """ + from litellm.integrations.otel.plumbing import routing as routing_mod + + monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) + shut_down = [] + monkeypatch.setattr(routing_mod, "_shutdown_provider", lambda p: shut_down.append(p)) + logger, _default, captured = _phoenix_routing_logger("capture_deferred_churn") + md_a = {"user_api_key_auth_metadata": {"phoenix_project_name": "proj-a"}} + server = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + with trace.use_span(server, end_on_exit=False): + _emit_llm( + logger, + { + "standard_logging_object": _payload(litellm_call_id="call_a", metadata=md_a), + "litellm_params": {"metadata": md_a}, + }, + ambient=server, + ) + assert len(logger._tenant_tracers._providers) == 1 + idle = next(iter(logger._tenant_tracers._providers.values())) + assert shut_down == [] + + md_b = {"user_api_key_auth_metadata": {"phoenix_project_name": "proj-b"}} + deferred_kwargs = { + "litellm_call_id": "call_b", + "standard_logging_object": _payload(litellm_call_id="call_b", metadata=md_b), + "litellm_params": {"metadata": md_b}, + } + logger.log_pre_api_call(model="gpt-4o", messages=[], kwargs=deferred_kwargs) + carrier = logger._open_llm_calls["call_b"] + assert carrier.span is None + assert carrier.provider is None + assert list(logger._tenant_tracers._providers.values()) == [idle] + assert shut_down == [] + assert captured and all("proj-b" not in headers for headers in captured) + + asyncio.run(logger.async_log_success_event(deferred_kwargs, None, None, None)) + server.end() + headers_b = next(h for h in captured if "proj-b" in h) + assert [s.name for s in captured[headers_b].get_finished_spans()] == ["chat gpt-4o"] diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 5becd05b8e8..3102e69bf26 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -2653,8 +2653,6 @@ class TestObservabilityCallbackBans: "posthog_api_url", "braintrust_api_key", "braintrust_project", - "phoenix_project_name", - "phoenix_project_name_override", "wandb_api_key", "weave_project_id", "gcs_bucket_name", @@ -2685,8 +2683,7 @@ class TestObservabilityCallbackBans: "langsmith_api_key", "posthog_api_url", "braintrust_project", - "phoenix_project_name", - "phoenix_project_name_override", + "user_api_key_auth_metadata", ], ) def test_observability_field_in_metadata_dict_is_rejected( @@ -2707,6 +2704,28 @@ class TestObservabilityCallbackBans: ) assert field in str(exc.value) + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) + @pytest.mark.parametrize( + "field", + ["phoenix_project_name", "phoenix_project_name_override"], + ) + def test_phoenix_project_fields_in_metadata_are_accepted(self, metadata_key, field): + # The Phoenix integrations only honor the project from + # ``user_api_key_auth_metadata`` on the proxy, so the bare metadata + # fields are inert and must not 400 SDK-style callers that send them. + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + metadata_key: {field: "client-project"}, + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) + def test_observability_field_in_litellm_params_metadata_is_rejected(self): with pytest.raises(ValueError) as exc: is_request_body_safe( From 487356733c7254944e7825bdd7c87e6d8f262e8f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 11:48:43 -0700 Subject: [PATCH 2/2] test: replace blind sleeps with deadline waits in callback and caching tests (#37660) * test: replace blind sleeps with deadline waits in callback and caching tests tests/local_testing/test_custom_callback_input.py slept a fixed 1-3s after every call and then asserted the callback handler recorded no errors. Because the handler only appends to `states` when a callback actually fires, an assert of `len(errors) == 0` passes just as happily when nothing fired at all, so the sleep was buying flakiness in exchange for a vacuous check. The async tests were worse: `time.sleep` blocks the event loop, so the success/failure tasks scheduled on it could not run before the assertion. Adds tests/_wait_helpers.py with `wait_until` / `await_until`, which poll a predicate against a deadline, and converts all 17 sites to wait on the thing the test actually cares about (the terminal state landing in `states`, or the patched log hook being called). The waits assert the callback fired, so these tests now fail on a dropped callback instead of passing silently. The three sleeps in test_caching_handler.py sat between `sync_set_cache` and `_sync_get_cache`, both fully synchronous against a local in-memory cache, so they are just deleted. * fix(test): wait on the priming call's own logging in the cache-hit test The 3s sleep in test_logging_async_cache_hit_sync_call was not waiting for the cache write, which lands before the stream iterator is exhausted. It was waiting for the priming call's success callback to drain, so the handler installed right after it only ever sees the second, cache-hit call. Waiting on a populated cache_dict let the priming call's still-pending log_success_event reach the new mock, and the test then read cache_hit off the wrong payload. Waits on the priming handler's own sync_success state instead. --- tests/_wait_helpers.py | 46 +++++++++++ tests/local_testing/test_caching_handler.py | 4 - .../test_custom_callback_input.py | 79 +++++++++++++------ 3 files changed, 103 insertions(+), 26 deletions(-) create mode 100644 tests/_wait_helpers.py diff --git a/tests/_wait_helpers.py b/tests/_wait_helpers.py new file mode 100644 index 00000000000..f67e623e3ad --- /dev/null +++ b/tests/_wait_helpers.py @@ -0,0 +1,46 @@ +"""Deadline-based waits for tests, so nothing has to guess how long a background callback takes.""" + +import asyncio +import time +from collections.abc import Callable +from typing import Final + +DEFAULT_TIMEOUT_S: Final[float] = 10.0 +DEFAULT_INTERVAL_S: Final[float] = 0.02 + + +def _fail(timeout_s: float, message: str) -> None: + raise AssertionError(f"condition not met within {timeout_s}s: {message}") + + +def wait_until( + predicate: Callable[[], bool], + *, + message: str, + timeout_s: float = DEFAULT_TIMEOUT_S, + interval_s: float = DEFAULT_INTERVAL_S, +) -> None: + deadline: Final = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(interval_s) # sleep-ok: bounded poll interval, not a blind settle + if not predicate(): + _fail(timeout_s, message) + + +async def await_until( + predicate: Callable[[], bool], + *, + message: str, + timeout_s: float = DEFAULT_TIMEOUT_S, + interval_s: float = DEFAULT_INTERVAL_S, +) -> None: + """Yields to the event loop between polls, so callbacks scheduled as tasks get a chance to run.""" + deadline: Final = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + await asyncio.sleep(interval_s) + if not predicate(): + _fail(timeout_s, message) diff --git a/tests/local_testing/test_caching_handler.py b/tests/local_testing/test_caching_handler.py index 0f4539162a2..b26334e9ee0 100644 --- a/tests/local_testing/test_caching_handler.py +++ b/tests/local_testing/test_caching_handler.py @@ -741,8 +741,6 @@ def test_sync_responses_api_caching(): # Step 1: Cache the responses API response caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs) - time.sleep(0.5) - # Step 2: Retrieve from cache cached_response = caching_handler._sync_get_cache( model=original_model, @@ -875,7 +873,6 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): } caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs) - time.sleep(0.2) cached_response = caching_handler._sync_get_cache( model=original_model, @@ -920,7 +917,6 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): } caching_handler.sync_set_cache(result=chat_completion_response, kwargs=kwargs) - time.sleep(0.2) cached_response = caching_handler._sync_get_cache( model=original_model, diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index 6a4ec9206f7..cedb5ea1a97 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -4,7 +4,6 @@ import asyncio import inspect import os import sys -import time import traceback from litellm._uuid import uuid from datetime import datetime @@ -20,6 +19,7 @@ import litellm from litellm import Cache, completion, embedding from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import LiteLLMCommonStrings +from tests._wait_helpers import await_until, wait_until # Test Scenarios (test across completion, streaming, embedding) ## 1: Pre-API-Call @@ -389,7 +389,10 @@ def test_chat_openai_stream(): continue except Exception: pass - time.sleep(1) + wait_until( + lambda: "sync_failure" in customHandler.states, + message=f"no sync_failure callback, states={customHandler.states}", + ) print(f"customHandler.errors: {customHandler.errors}") assert len(customHandler.errors) == 0 litellm.callbacks = [] @@ -430,10 +433,12 @@ async def test_async_chat_openai_stream(): ) async for chunk in response: continue - await asyncio.sleep(1) except Exception: pass - time.sleep(1) + await await_until( + lambda: "async_failure" in customHandler.states, + message=f"no async_failure callback, states={customHandler.states}", + ) print(f"customHandler.errors: {customHandler.errors}") assert len(customHandler.errors) == 0 litellm.callbacks = [] @@ -473,7 +478,10 @@ def test_chat_azure_stream(): continue except Exception: pass - time.sleep(1) + wait_until( + lambda: "sync_failure" in customHandler.states, + message=f"no sync_failure callback, states={customHandler.states}", + ) print(f"customHandler.errors: {customHandler.errors}") assert len(customHandler.errors) == 0 litellm.callbacks = [] @@ -590,7 +598,10 @@ async def test_async_chat_sagemaker_stream(): continue except Exception: pass - time.sleep(1) + await await_until( + lambda: "async_failure" in customHandler.states, + message=f"no async_failure callback, states={customHandler.states}", + ) print(f"customHandler.errors: {customHandler.errors}") assert len(customHandler.errors) == 0 litellm.callbacks = [] @@ -711,10 +722,12 @@ async def test_async_text_completion_bedrock(): async for chunk in response: continue - await asyncio.sleep(1) except Exception: pass - time.sleep(1) + await await_until( + lambda: "async_failure" in customHandler.states, + message=f"no async_failure callback, states={customHandler.states}", + ) print(f"customHandler.errors: {customHandler.errors}") assert len(customHandler.errors) == 0 litellm.callbacks = [] @@ -754,10 +767,12 @@ async def test_async_text_completion_openai_stream(): async for chunk in response: continue - await asyncio.sleep(1) except Exception: pass - time.sleep(1) + await await_until( + lambda: "async_failure" in customHandler.states, + message=f"no async_failure callback, states={customHandler.states}", + ) print(f"customHandler.errors: {customHandler.errors}") assert len(customHandler.errors) == 0 litellm.callbacks = [] @@ -816,7 +831,10 @@ def test_amazing_sync_embedding(): ) print(f"customHandler_success.errors: {customHandler_success.errors}") print(f"customHandler_success.states: {customHandler_success.states}") - time.sleep(2) + wait_until( + lambda: len(customHandler_success.states) == 3, + message=f"success states never reached pre/post/success, got {customHandler_success.states}", + ) assert len(customHandler_success.errors) == 0 assert len(customHandler_success.states) == 3 # pre, post, success # test failure callback @@ -832,7 +850,10 @@ def test_amazing_sync_embedding(): pass print(f"customHandler_failure.errors: {customHandler_failure.errors}") print(f"customHandler_failure.states: {customHandler_failure.states}") - time.sleep(2) + wait_until( + lambda: len(customHandler_failure.states) == 3, + message=f"failure states never reached pre/post/failure, got {customHandler_failure.states}", + ) assert len(customHandler_failure.errors) == 1 assert len(customHandler_failure.states) == 3 # pre, post, failure except Exception as e: @@ -939,7 +960,10 @@ def test_image_generation_openai(): print(f"customHandler_success.errors: {customHandler_success.errors}") print(f"customHandler_success.states: {customHandler_success.states}") - time.sleep(2) + wait_until( + lambda: len(customHandler_success.states) == 3, + message=f"success states never reached pre/post/success, got {customHandler_success.states}", + ) assert len(customHandler_success.errors) == 0 assert len(customHandler_success.states) == 3 # pre, post, success # test failure callback @@ -991,7 +1015,10 @@ def test_turn_off_message_logging(): mock_response="Going well!", ) - time.sleep(2) + wait_until( + lambda: "sync_success" in customHandler.states, + message=f"no sync_success callback, states={customHandler.states}", + ) assert len(customHandler.errors) == 0 @@ -1033,7 +1060,7 @@ def test_standard_logging_payload(model, turn_off_message_logging): mock_response="Going well!", ) - time.sleep(2) + wait_until(lambda: mock_client.called, message="log_success_event never fired") mock_client.assert_called_once() print( @@ -1147,7 +1174,7 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream): for chunk in response: continue - time.sleep(2) + wait_until(lambda: mock_client.called, message="log_success_event never fired") mock_client.assert_called() print( @@ -1247,7 +1274,7 @@ def test_aaastandard_logging_payload_cache_hit(): caching=True, ) - time.sleep(2) + wait_until(lambda: mock_client.called, message="log_success_event never fired") mock_client.assert_called_once() assert "standard_logging_object" in mock_client.call_args.kwargs["kwargs"] @@ -1276,6 +1303,9 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging): litellm.cache = Cache() + primingHandler = CompletionCustomHandler() + litellm.callbacks = [primingHandler] + response = litellm.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}], @@ -1285,7 +1315,10 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging): for chunk in response: print(chunk) - time.sleep(3) + wait_until( + lambda: "sync_success" in primingHandler.states, + message=f"priming call never finished logging, states={primingHandler.states}", + ) customHandler = CompletionCustomHandler() litellm.callbacks = [customHandler] litellm.success_callback = [] @@ -1303,7 +1336,7 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging): for chunk in resp: print(chunk) - time.sleep(2) + wait_until(lambda: mock_client.called, message="log_success_event never fired") mock_client.assert_called_once() assert "standard_logging_object" in mock_client.call_args.kwargs["kwargs"] @@ -1387,7 +1420,7 @@ def test_logging_standard_payload_llm_headers(stream): for chunk in resp: continue - time.sleep(2) + wait_until(lambda: mock_client.called, message="log_success_event never fired") mock_client.assert_called() standard_logging_object: StandardLoggingPayload = mock_client.call_args.kwargs[ @@ -1458,7 +1491,7 @@ async def test_standard_logging_payload_stream_usage(sync_mode): chunks = [] for chunk in resp: chunks.append(chunk) - time.sleep(2) + wait_until(lambda: mock_client.called, message="log_success_event never fired") else: resp = await litellm.acompletion( model="anthropic/claude-sonnet-4-5-20250929", @@ -1469,7 +1502,9 @@ async def test_standard_logging_payload_stream_usage(sync_mode): chunks = [] async for chunk in resp: chunks.append(chunk) - await asyncio.sleep(2) + await await_until( + lambda: mock_client.called, message="async_log_success_event never fired" + ) mock_client.assert_called_once()