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/_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() 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(