diff --git a/litellm/integrations/otel/model/destination.py b/litellm/integrations/otel/model/destination.py index 4226b4361dd..e760ba3600b 100644 --- a/litellm/integrations/otel/model/destination.py +++ b/litellm/integrations/otel/model/destination.py @@ -42,7 +42,7 @@ class OtelDestination(LiteLLMBaseModel): success_sampling_rate: float | None = Field( default=None, description=( - "Share of the request trees forwarded, 0.0..1.0, drawn once per request when its root span ends; " + "Share of the request trees forwarded, 0.0..1.0, picked by a hash of the trace id; " "``None`` forwards every one." ), ) diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 898eab5715b..fc1e9b7099d 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -1,13 +1,14 @@ """Provider / exporter factory + the Baggage span processor.""" import queue -import random import threading import time +import weakref from collections import OrderedDict from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from hashlib import sha256 from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from opentelemetry import _logs, baggage, metrics, trace from opentelemetry._logs import Logger, LoggerProvider, NoOpLoggerProvider @@ -486,15 +487,15 @@ def is_llm_call_span(span: ReadableSpan) -> bool: return GenAI.OPERATION_NAME in attributes and MCP.METHOD_NAME not in attributes -#: Request trees held back for a sampled destination until the last span of the trace +#: Request trees held back for a sampled destination until the last span of the request #: still open ends, the verdicts kept for the spans that start after that (a post-call -#: database write), and the spans in flight per trace. All bounded, so a span that never -#: ends, or a flood of requests, cannot hold spans for ever: the oldest tree is decided -#: on what it has. +#: database write), the spans in flight per request, and the request each open span +#: belongs to. All bounded, so a span that never ends, or a flood of requests, cannot +#: hold spans for ever: the oldest tree is decided on what it has. _MAX_PENDING_TREES: Final = 1024 _MAX_PENDING_SPANS_PER_TREE: Final = 512 _MAX_REMEMBERED_VERDICTS: Final = 4096 -_MAX_OPEN_TRACES: Final = 16384 +_MAX_OPEN_REQUESTS: Final = 16384 class _PendingTree: @@ -512,25 +513,35 @@ def _is_sampled(destination: "OtelDestination") -> bool: return destination.success_sampling_rate is not None or destination.error_sampling_rate is not None -#: A verdict is one request tree's draw for one destination at its rates: the exporter +#: A request is its trace id and the id of its root, the span of it whose parent is +#: remote or missing. Calls sharing a caller's trace are separate requests, each decided +#: on its own outcome. +_RequestKey: TypeAlias = tuple[int, int] + +#: A verdict is one request's keep or drop for one destination at its rates: the exporter #: ``cache_key`` leaves the rates out, so two destinations sharing an exporter with -#: different rates would otherwise share one draw. -_VerdictKey = tuple[int, object, float | None, float | None] +#: different rates would otherwise share one verdict. +_VerdictKey: TypeAlias = tuple[int, int, object, float | None, float | None] -def _verdict_key(trace_id: int, destination: "OtelDestination") -> _VerdictKey: - return (trace_id, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate) +def _verdict_key(request: _RequestKey, destination: "OtelDestination") -> _VerdictKey: + return (*request, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate) -def _keeps_tree(destination: "OtelDestination", failed: bool, draw: Callable[[], float]) -> bool: - """Whether one draw against the destination's rate keeps a request tree, read the - way the legacy Arize callback reads ``arize_success_sampling_rate`` / - ``arize_error_sampling_rate``: a tree with a failed span answers to the error - rate, an unset rate keeps everything and ``0.0`` keeps nothing.""" +def _keeps_tree(destination: "OtelDestination", failed: bool, trace_id: int) -> bool: + """Whether the destination's rate keeps a request tree of ``trace_id``: a tree with + a failed span answers to ``arize_error_sampling_rate``, the rest to + ``arize_success_sampling_rate``, an unset rate keeps everything and ``0.0`` keeps + nothing. + + The first eight bytes of a SHA-256 of the trace id are read against the rate, so + every process deciding a request of the trace with the same outcome reaches the same + verdict.""" rate: Final = destination.error_sampling_rate if failed else destination.success_sampling_rate if rate is None: return True - return rate > 0.0 and draw() <= rate + digest: Final = sha256(trace_id.to_bytes(16, "big")).digest() + return int.from_bytes(digest[:8], "big") < round(rate * 2**64) def _in_scope(span: ReadableSpan, scope: "OtelSpanScope") -> bool: @@ -667,10 +678,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): default_capture_by_backend: Mapping[str, str] = MappingProxyType({}), pending_drains: int = _MAX_PENDING_DRAINS, drain_pool: _DrainPool | None = None, - sampling_draw: Callable[[], float] | None = None, ) -> None: self._operator_sinks: Final = operator_sinks - self._draw: Final = sampling_draw if sampling_draw is not None else random.random self._excluded_db_systems: Final = excluded_db_systems self._default_capture: Final = default_capture self._default_capture_by_backend: Final = default_capture_by_backend @@ -682,24 +691,32 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count - self._pending: OrderedDict[int, _PendingTree] = OrderedDict() # mutable-ok: bounded, by trace id - self._open: OrderedDict[int, int] = OrderedDict() # mutable-ok: bounded, spans in flight by trace id + self._pending: OrderedDict[_RequestKey, _PendingTree] = OrderedDict() # mutable-ok: bounded, by request + self._open: OrderedDict[_RequestKey, set[int]] = OrderedDict() # mutable-ok: bounded, open span ids by request + self._roots: dict[int, _RequestKey] = {} # mutable-ok: holds only the spans ``_open`` holds, request by span id + self._span_roots: weakref.WeakKeyDictionary[SDKSpan, int] = weakref.WeakKeyDictionary() self._verdicts: OrderedDict[_VerdictKey, bool] = OrderedDict() # mutable-ok: bounded LRU self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains) def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None: context: Final = span.context - trace_id: Final = context.trace_id if context is not None else 0 with self._lock: - self._open[trace_id] = self._open.get(trace_id, 0) + 1 - while len(self._open) > _MAX_OPEN_TRACES: - self._open.popitem(last=False) + request: Final = self._request_of(span, trace.get_current_span(parent_context)) + if context is not None: + self._span_roots[span] = request[1] + self._roots[context.span_id] = request + self._open.setdefault(request, set()).add(context.span_id) + while len(self._open) > _MAX_OPEN_REQUESTS: + for forgotten in self._open.popitem(last=False)[1]: + self._roots.pop(forgotten, None) def on_end(self, span: ReadableSpan) -> None: suppressed: Final = suppressed_backends() attributes: Final = span.attributes or _NO_ATTRIBUTES - context: Final = span.context - trace_id: Final = context.trace_id if context is not None else 0 + with self._lock: + request: Final = self._request_of(span) + if span.context is not None: + self._roots.pop(span.context.span_id, None) failed: Final = span.status.status_code is StatusCode.ERROR for destination in request_destinations(): if ( @@ -716,11 +733,32 @@ class TenantFanOutSpanProcessor(SpanProcessor): if not _is_sampled(destination): self._forward(span, destination) continue - for held_span, held_destination in self._route(trace_id, span, destination, failed): + for held_span, held_destination in self._route(request, span, destination, failed): self._forward(held_span, held_destination) - for held_span, held_destination in self._settle(trace_id, failed): + for held_span, held_destination in self._settle(request, span, failed): self._forward(held_span, held_destination) + def _request_of(self, span: ReadableSpan, parent_span: Span | None = None) -> _RequestKey: + """The request ``span`` belongs to, read under ``_lock``: the request recorded while + it is open, else the span itself when its parent is remote or missing, else its + parent's request, recorded while the parent is open or on ``parent_span`` for a + child starting under a parent that has ended, else the parent itself.""" + context: Final = span.context + if context is None: + return (0, 0) + recorded: Final = self._roots.get(context.span_id) + if recorded is not None: + return recorded + parent: Final = span.parent + if parent is None or parent.is_remote: + return (context.trace_id, context.span_id) + parent_request: Final = self._roots.get(parent.span_id) + if parent_request is not None: + return parent_request + if isinstance(parent_span, SDKSpan): + return (context.trace_id, self._span_roots.get(parent_span, parent.span_id)) + return (context.trace_id, parent.span_id) + def _default_capture_for(self, destination: "OtelDestination") -> str: backend: Final = destination.callback_name return ( @@ -743,14 +781,14 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._release(processor) def _route( - self, trace_id: int, span: ReadableSpan, destination: "OtelDestination", failed: bool + self, request: _RequestKey, span: ReadableSpan, destination: "OtelDestination", failed: bool ) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": """The spans to forward now for ``span`` on a sampled ``destination``: the span itself when its tree's verdict is already known and keeps it, nothing while the - tree is held for the rest of its trace, or what the bounds made the fan-out decide + tree is held for the rest of its request, or what the bounds made the fan-out decide early (a tree that outgrew its cap is decided on what it has, and so is the oldest tree once too many are waiting, so a span that never ends holds nothing back for - ever). A failed span marks its tree before any such early draw. Once shutdown + ever). A failed span marks its tree before any such early decision. Once shutdown has started flushing, nothing is held: a span ending then is decided at once, since no later flush would reach it. @@ -758,58 +796,62 @@ class TenantFanOutSpanProcessor(SpanProcessor): another thread between the two would leave this span in a fresh tree. """ with self._lock: - verdict: Final = self._verdicts.get(_verdict_key(trace_id, destination)) + verdict: Final = self._verdicts.get(_verdict_key(request, destination)) if verdict is not None: return ((span, destination),) if verdict else () - tree: Final = self._pending.setdefault(trace_id, _PendingTree()) + tree: Final = self._pending.setdefault(request, _PendingTree()) tree.held.append((span, destination)) if failed: tree.failed = True if self._draining or len(tree.held) >= _MAX_PENDING_SPANS_PER_TREE: - return self._decide_locked(trace_id) + return self._decide_locked(request) if len(self._pending) > _MAX_PENDING_TREES: oldest: Final = next(iter(self._pending)) return self._decide_locked(oldest) return () - def _decide(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + def _decide(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": with self._lock: - return self._decide_locked(trace_id) + return self._decide_locked(request) - def _settle(self, trace_id: int, failed: bool) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": - """What the trace's held tree exports once the span that just ended was the last - of its trace still open. The model-call span ends in the post-call callback, - after the server span, so a tree is decided when its trace goes quiet rather than + def _settle( + self, request: _RequestKey, span: ReadableSpan, failed: bool + ) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + """What the request's held tree exports once the span that just ended was the last + of its request still open. The model-call span ends in the post-call callback, + after the server span, so a tree is decided when its request goes quiet rather than when its root ends, and a model call that fails late still answers to the error - rate. A trace whose starts this fan-out never saw (attached mid-flight, or forgotten - by the bound) is decided as each of its spans ends. + rate. A request whose starts this fan-out never saw (attached mid-flight, or + forgotten by the bound) is decided as each of its spans ends. """ with self._lock: - tree: Final = self._pending.get(trace_id) + tree: Final = self._pending.get(request) if tree is not None and failed: tree.failed = True - still_open: Final = self._open.pop(trace_id, 1) - 1 - if still_open > 0: - self._open[trace_id] = still_open + still_open: Final = self._open.pop(request, set()) + if span.context is not None: + still_open.discard(span.context.span_id) + if still_open: + self._open[request] = still_open return () - return self._decide_locked(trace_id) + return self._decide_locked(request) - def _decide_locked(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": - """Draw once per destination for the held tree; the spans the draws keep, and a - verdict remembered for the spans of this trace that start later.""" - tree: Final = self._pending.pop(trace_id, None) + def _decide_locked(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + """Decide once per destination for the held tree; the spans the verdicts keep, and + a verdict remembered for the spans of this request that start later.""" + tree: Final = self._pending.pop(request, None) if tree is None: return () - verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one draw per destination of this tree + verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one verdict per destination of this tree for _, destination in tree.held: - key = _verdict_key(trace_id, destination) + key = _verdict_key(request, destination) if key not in verdicts: - verdicts[key] = _keeps_tree(destination, tree.failed, self._draw) + verdicts[key] = _keeps_tree(destination, tree.failed, request[0]) self._verdicts[key] = verdicts[key] while len(self._verdicts) > _MAX_REMEMBERED_VERDICTS: self._verdicts.popitem(last=False) return tuple( - (span, destination) for span, destination in tree.held if verdicts[_verdict_key(trace_id, destination)] + (span, destination) for span, destination in tree.held if verdicts[_verdict_key(request, destination)] ) def _operator_already_writes( @@ -850,8 +892,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): with self._lock: self._draining = True undecided: Final = tuple(self._pending) - for trace_id in undecided: - for held_span, held_destination in self._decide(trace_id): + for request in undecided: + for held_span, held_destination in self._decide(request): self._forward(held_span, held_destination) deadline: Final = time.monotonic() + self._drain_seconds with self._lock: diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 3ce8d7ddad5..dea8bcb73a4 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -90,7 +90,7 @@ class PassThroughEndpointHandler(BaseTranslation): payload_to_check: Final = { k: v for k, v in without_server_streaming_classification(data).items() - if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") + if not k.startswith("_") and k not in ("metadata", "litellm_metadata", "litellm_logging_obj") } verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail") return safe_dumps(payload_to_check) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9a7d88b8eda..6ec1bd29beb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -36840,10 +36840,14 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36931,8 +36935,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37014,8 +37020,10 @@ "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -37186,8 +37194,10 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_12hr": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -37268,8 +37278,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37349,8 +37361,10 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_12hr": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -75046,10 +75060,14 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -76439,9 +76457,10 @@ "azure_ai/Microsoft-Decision-1": { "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", + "max_input_tokens": 32768, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" }, "gpt-rosalind-discovery": { "cache_read_input_token_cost": 5e-07, @@ -78652,5 +78671,107 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", "supports_web_search": true + }, + "gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_12hr": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-07, + "cache_creation_input_token_cost_priority": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "cache_read_input_token_cost_above_272k_tokens_batches": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-07, + "cache_read_input_token_cost_batches": 2.5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "input_cost_per_token_above_272k_tokens_batches": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_priority": 4e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_priority": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "output_cost_per_token_above_272k_tokens_batches": 3.75e-06, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.5e-05, + "output_cost_per_token_batches": 2.5e-06, + "output_cost_per_token_flex": 2.5e-06, + "output_cost_per_token_priority": 1e-05, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6.1-terra": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index a3e33929f6d..e148e28e826 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -3,8 +3,8 @@ The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a failed server must contribute a classified outcome instead of silently shrinking the list: an empty contribution with no signal makes a broken upstream indistinguishable from a healthy server with no -tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream -body crosses the trust boundary; classification is total, so any exception out of a server fetch +tools. Outcomes carry machine fields and controlled diagnostics so nothing from an upstream body +crosses the trust boundary; classification is total, so any exception out of a server fetch becomes an outcome, never a second failure. """ @@ -21,8 +21,10 @@ from urllib.request import parse_http_list, parse_keqv_list import anyio import httpx import httpx2 +from mcp import MCPError +from mcp.types import CONNECTION_CLOSED from mcp.types import Tool as MCPTool -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, PositiveInt, TypeAdapter from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -65,6 +67,16 @@ class ServerListFault(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) tag: ListFaultCategory status_code: int | None = None + sse_event_limit_bytes: PositiveInt | None = None + + @property + def public_message(self) -> str | None: + if self.sse_event_limit_bytes is None: + return None + return ( + f"Upstream MCP SSE event exceeded the {self.sse_event_limit_bytes} byte per-event limit. " + "Paginate tools/list responses so each event stays below the limit." + ) ServerOutcome: TypeAlias = ServerListOk | ServerListFault @@ -155,6 +167,18 @@ def raise_classified_list_failure( raise MCPServerListError(classify_list_exception(exc), server_name) from exc +def _sse_event_limit(exc: BaseException) -> int | None: + """The SDK replaces decoder exceptions with CONNECTION_CLOSED and a fixed diagnostic.""" + if not isinstance(exc, (httpx2.SSEError, MCPError)): + return None + if isinstance(exc, MCPError) and exc.error.code != CONNECTION_CLOSED: + return None + prefix: Final = "SSE stream failed: " if isinstance(exc, MCPError) else "" + message: Final = exc.error.message if isinstance(exc, MCPError) else str(exc) + matched: Final = re.fullmatch(prefix + r"Server-sent event exceeded the ([1-9][0-9]{0,18}) byte limit\.", message) + return int(matched[1]) if matched is not None else None + + def classify_list_exception(exc: BaseException) -> ServerListFault: """Classify a per-server listing failure into exactly one outcome. Total: an exception this function cannot recognize is the gateway's own fault (``internal``), never a re-raise.""" @@ -179,6 +203,11 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: response: Final = _find_upstream_response(exc) if response is not None: return ServerListFault(tag="upstream_error", status_code=response.status_code) + sse_event_limit: Final = next( + (limit for error in iter_exception_tree(exc) if (limit := _sse_event_limit(error)) is not None), None + ) + if sse_event_limit is not None: + return ServerListFault(tag="upstream_error", sse_event_limit_bytes=sse_event_limit) if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)): return ServerListFault(tag="timeout") if isinstance(exc, (httpx.TransportError, httpx2.TransportError)): @@ -295,7 +324,7 @@ def listing_failed_outright(outcomes: Mapping[str, ServerOutcome]) -> bool: def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: """The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST - response: category plus status code only, never upstream prose or URLs.""" + response: category, status code and controlled diagnostics, never upstream prose or URLs.""" match outcome.tag: case "ok": return {"status": "ok", "tool_count": outcome.tool_count} @@ -303,6 +332,15 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: return { "status": outcome.tag, **({"http_status": outcome.status_code} if outcome.status_code is not None else {}), + **( + { + "reason": "sse_event_too_large", + "limit_bytes": outcome.sse_event_limit_bytes, + "message": outcome.public_message, + } + if outcome.sse_event_limit_bytes is not None + else {} + ), } case _: assert_never(outcome.tag) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index ee063b81ae4..62387b7cc28 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -957,7 +957,7 @@ if MCP_AVAILABLE: status_code=list_fault_http_status(fault), detail={ "error": fault.tag, - "message": f"Failed to list tools from server {get_server_prefix(server)}", + "message": fault.public_message or f"Failed to list tools from server {get_server_prefix(server)}", }, ) from e except Exception as e: diff --git a/litellm/utils.py b/litellm/utils.py index e2820f9c682..4315ae426cf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9571,7 +9571,7 @@ class ProviderConfigManager: return get_nvidia_nim_rerank_config(model) elif litellm.LlmProviders.VERTEX_AI == provider: - return litellm.VertexAIRerankConfig() + return ProviderConfigManager._get_vertex_ai_rerank_config_cached() elif litellm.LlmProviders.FIREWORKS_AI == provider: return litellm.FireworksAIRerankConfig() elif litellm.LlmProviders.VOYAGE == provider: @@ -9590,6 +9590,11 @@ class ProviderConfigManager: return get_dashscope_family_rerank_config(provider.value) return litellm.CohereRerankConfig() + @staticmethod + @lru_cache(maxsize=1) + def _get_vertex_ai_rerank_config_cached() -> BaseRerankConfig: + return litellm.VertexAIRerankConfig() + @staticmethod def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDecisionsConfig | None: if provider == LlmProviders.PERPLEXITY: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9a7d88b8eda..6ec1bd29beb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -36840,10 +36840,14 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36931,8 +36935,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37014,8 +37020,10 @@ "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -37186,8 +37194,10 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_12hr": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -37268,8 +37278,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37349,8 +37361,10 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_12hr": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -75046,10 +75060,14 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -76439,9 +76457,10 @@ "azure_ai/Microsoft-Decision-1": { "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", + "max_input_tokens": 32768, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" }, "gpt-rosalind-discovery": { "cache_read_input_token_cost": 5e-07, @@ -78652,5 +78671,107 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", "supports_web_search": true + }, + "gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_12hr": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-07, + "cache_creation_input_token_cost_priority": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "cache_read_input_token_cost_above_272k_tokens_batches": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-07, + "cache_read_input_token_cost_batches": 2.5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "input_cost_per_token_above_272k_tokens_batches": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_priority": 4e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_priority": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "output_cost_per_token_above_272k_tokens_batches": 3.75e-06, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.5e-05, + "output_cost_per_token_batches": 2.5e-06, + "output_cost_per_token_flex": 2.5e-06, + "output_cost_per_token_priority": 1e-05, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6.1-terra": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/tests/integration/observability/test_otel_team_sampling_across_pods.py b/tests/integration/observability/test_otel_team_sampling_across_pods.py new file mode 100644 index 00000000000..dc7b295febd --- /dev/null +++ b/tests/integration/observability/test_otel_team_sampling_across_pods.py @@ -0,0 +1,723 @@ +"""A team's Arize sampling rate decides once per trace, so every pod and worker serving a call of it agrees. + +Two proxies share one database, as two pods behind a load balancer do: the first runs one worker, the second +two. Calls of one trace carry the caller's ``traceparent`` and land on different pods and workers; the team's +Arize space must receive every call of a trace or none of them. +""" + +import asyncio +import json +import re +import uuid +from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, replace +from itertools import chain, product +from pathlib import Path +from typing import Final, Literal, TypeAlias +from urllib.parse import urlparse + +import httpx +import pytest +import yaml +from integration._support.client import ( + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + until_asserted, +) +from integration._support.database import scratch_database +from integration._support.otlp_sink import Span, SpanSinks, configure_sink, recorded_requests, recorded_spans +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +MODEL: Final = "team-sampling" +SERVER: Final = 2 +CALL_ID: Final = "litellm.call_id" +OPERATION: Final = "gen_ai.operation.name" +SPACE: Final = "space-sampled" +CREDENTIALS: Final = {"arize_space_key": SPACE, "arize_api_key": "arize-sampled"} +FAILING: Final = re.compile(rb"sampling-fail") +JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +TRACES: Final = 16 +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Verdicts: TypeAlias = Mapping[str, tuple[bool, ...]] + + +def _chat_reply(stream: bool) -> Reply: + usage: Final[JsonValue] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9} + if not stream: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-sampling", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": usage, + } + ).encode() + ) + chunk: Final = {"id": "chatcmpl-sampling", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + deltas: Final[tuple[JsonValue, ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": usage}, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"), + ) + + +def _responses_reply(stream: bool) -> Reply: + response: Final[Mapping[str, JsonValue]] = { + "id": "resp_sampling", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_sampling", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, + } + if not stream: + return Reply(body=json.dumps(response).encode()) + events: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_sampling", + "output_index": 0, + "content_index": 0, + "delta": "ok", + }, + {"type": "response.completed", "sequence_number": 2, "response": dict(response)}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + ) + + +def _upstream(request: Request) -> Reply: + if FAILING.search(request.body) is not None: + return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}') + stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True + if request.target.endswith("/responses"): + return _responses_reply(stream) + return _chat_reply(stream) + + +@dataclass(frozen=True, slots=True) +class Sent: + trace_id: str + parent_span_id: str + call_id: str + + +@dataclass(frozen=True, slots=True) +class Call: + path: str + body: Mapping[str, JsonValue] + headers: Mapping[str, str] + sent: Sent + expected: int + + +def _call( + key: str, trace_id: str | None, endpoint: Endpoint = "chat", stream: bool = False, content: str = "hello" +) -> Call: + parent: Final = uuid.uuid4().hex[:16] + call_id: Final = str(uuid.uuid4()) + messages: Final[JsonValue] = [{"role": "user", "content": content}] + path, body = { + "chat": ("/v1/chat/completions", {"model": MODEL, "messages": messages, "stream": stream}), + "responses": ("/v1/responses", {"model": MODEL, "input": content, "stream": stream}), + "messages": ("/v1/messages", {"model": MODEL, "max_tokens": 16, "messages": messages, "stream": stream}), + }[endpoint] + headers: Final = { + "Authorization": f"Bearer {key}", + "x-litellm-call-id": call_id, + "connection": "close", + **({} if trace_id is None else {"traceparent": f"00-{trace_id}-{parent}-01"}), + } + expected: Final = 500 if FAILING.search(content.encode()) else 200 + return Call(path, body, headers, Sent(trace_id or "", parent if trace_id else "", call_id), expected) + + +@dataclass(frozen=True, slots=True) +class Pods: + """``first`` serves with one worker, ``second`` with two.""" + + first: OwnedProxy + second: OwnedProxy + scenario: Scenario + sinks: SpanSinks + + def team_key(self, variables: Mapping[str, str]) -> str: + team: Final = self.scenario.team() + attached: Final = self.first.gateway.request( + "POST", f"/team/{team}/callback", {"callback_name": "arize", "callback_vars": dict(variables)} + ) + assert attached.status_code == 200, attached.text + return self.scenario.key(team_id=team) + + def send(self, pod: OwnedProxy, call: Call) -> Sent: + with pod.gateway.client.stream("POST", call.path, json=call.body, headers=call.headers) as response: + text: Final = response.read().decode() + assert response.status_code == call.expected, text + return call.sent + + async def send_async(self, client: httpx.AsyncClient, call: Call) -> Sent: + async with client.stream("POST", call.path, json=call.body, headers=call.headers) as response: + text: Final = (await response.aread()).decode() + assert response.status_code == call.expected, text + return call.sent + + def on_each_worker(self, pod: OwnedProxy, call: Callable[[], Call]) -> tuple[Sent, ...]: + """One call served by each worker of ``pod``, the others suspended while it runs.""" + workers: Final = sorted(pod.worker_pids()) + assert len(workers) == pod.gateway.workers, workers + sent: Final[list[Sent]] = [] # mutable-ok: one call per worker, in worker order + for serving in workers: + with _suspended_all_but(pod, workers, serving): + sent.append(self.send(pod, call())) + return tuple(sent) + + def fence(self, key: str) -> tuple[Sent, ...]: + """Calls of an unsampled team on the sampled team's exporter, one per worker of both pods. + + Each worker sends spans to that exporter in the order they were handed to it, so all of these + arriving means every call any worker decided before them has been sent too. + """ + return (self.send(self.first, _call(key, uuid.uuid4().hex)),) + self.on_each_worker( + self.second, lambda: _call(key, uuid.uuid4().hex) + ) + + +@contextmanager +def _suspended_all_but(pod: OwnedProxy, workers: Sequence[int], serving: int) -> Generator[None]: + if len(workers) == 1: + yield + return + others: Final = [pid for pid in workers if pid != serving] + with pod.suspended(others[0]): + yield + + +def _space_id(request: Mapping[str, JsonValue]) -> JsonValue: + return {name.lower(): value for name, value in object_value(request.get("headers") or {}).items()}.get( + "arize-space-id" + ) + + +def _space_span_ids(url: str) -> frozenset[str]: + """Span ids the Arize sink received in export requests sent with the team's space id.""" + ours: Final = ( + list_value(request.get("span_ids") or []) for request in recorded_requests(url) if _space_id(request) == SPACE + ) + return frozenset(span_id for span_id in chain.from_iterable(ours) if isinstance(span_id, str)) + + +def _team_spans(pods: Pods, since: int) -> tuple[Span, ...]: + owned: Final = _space_span_ids(pods.sinks.arize) + return tuple(span for span in recorded_spans(pods.sinks.arize, since)[1] if span["span_id"] in owned) + + +def _delivered(spans: tuple[Span, ...], sent: Sequence[Sent]) -> frozenset[str]: + """Call ids of ``sent`` whose server span is among ``spans``.""" + by_parent: Final = {(call.trace_id, call.parent_span_id): call.call_id for call in sent} + return frozenset( + by_parent[(span["trace_id"], span["parent_span_id"])] + for span in spans + if span["kind"] == SERVER and (span["trace_id"], span["parent_span_id"]) in by_parent + ) + + +def _await_logged(pods: Pods, since: int, sent: Sequence[Sent]) -> None: + """Wait until every call of ``sent`` has its model span in the operator's unsampled copy.""" + expected: Final = {call.call_id for call in sent} + eventually( + lambda: ( + { + str(span["attributes"][CALL_ID]) + for span in recorded_spans(pods.sinks.operator, since)[1] + if OPERATION in span["attributes"] and CALL_ID in span["attributes"] + } + & expected + ), + lambda logged: logged == expected, + seconds=60, + ) + + +@dataclass(frozen=True, slots=True) +class Cursors: + operator: int + arize: int + + +def _cursors(pods: Pods) -> Cursors: + return Cursors(recorded_spans(pods.sinks.operator)[0], recorded_spans(pods.sinks.arize)[0]) + + +def _await_flushed(pods: Pods, cursors: Cursors, fence_key: str) -> None: + """Wait until every worker of both pods has sent the team's exporter everything it decided so far.""" + fences: Final = pods.fence(fence_key) + + def arrived() -> None: + delivered: Final = _delivered(_team_spans(pods, cursors.arize), fences) + assert {fence.call_id for fence in fences} <= delivered, "the unsampled team's calls have not arrived yet" + + until_asserted(arrived, seconds=30) + + +def _settled( + pods: Pods, + cursors: Cursors, + by_trace: Mapping[str, Sequence[Sent]], + fence_key: str, + holds: Callable[[tuple[bool, ...]], bool] = lambda kept: len(set(kept)) == 1, +) -> Verdicts: + """Per trace, which of its calls reached the team's space, once every worker has sent what it decided + and ``holds`` is true of every trace.""" + calls: Final = tuple(chain.from_iterable(by_trace.values())) + _await_logged(pods, cursors.operator, calls) + _await_flushed(pods, cursors, fence_key) + + def check() -> Verdicts: + delivered: Final = _delivered(_team_spans(pods, cursors.arize), calls) + verdicts: Final = { + trace_id: tuple(call.call_id in delivered for call in sent) for trace_id, sent in by_trace.items() + } + broken: Final = {trace_id: kept for trace_id, kept in verdicts.items() if not holds(kept)} + assert broken == {}, f"{len(broken)} of {len(verdicts)} traces were split (per call, kept): {broken}" + return verdicts + + return until_asserted(check, seconds=30) + + +def _trace_ids(cell: str, count: int = TRACES) -> tuple[str, ...]: + """Fixed RFC 4122 shaped trace ids, like the uuid4 hex a caller mints for ``traceparent``.""" + return tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"team-sampling/{cell}/{index}").hex for index in range(count)) + + +def _kept(verdicts: Verdicts) -> int: + return sum(1 for kept in verdicts.values() if all(kept)) + + +@pytest.mark.parametrize( + ("variables", "content"), + [({"arize_success_sampling_rate": "0.5"}, "hello"), ({"arize_error_sampling_rate": "0.5"}, "sampling-fail")], + ids=["success-rate", "error-rate"], +) +def test_each_trace_reaches_the_team_space_whole_or_not_at_all( + pods: Pods, variables: Mapping[str, str], content: str +) -> None: + """The ticket's case: three calls per trace load balanced over the pods.""" + sampled: Final = pods.team_key({**CREDENTIALS, **variables}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple( + pods.send(pod, _call(sampled, trace_id, content=content)) for pod in (pods.first, pods.second, pods.first) + ) + for trace_id in _trace_ids("-".join(sorted(variables))) + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES, f"a rate of 0.5 kept {_kept(verdicts)} of {TRACES} traces" + + +@pytest.mark.parametrize("client", ["sync", "async"]) +def test_every_endpoint_follows_the_verdict_of_its_trace(pods: Pods, client: str) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + plan: Final[tuple[tuple[OwnedProxy, Endpoint, bool], ...]] = ( + (pods.first, "chat", False), + (pods.second, "chat", True), + (pods.first, "responses", False), + (pods.second, "messages", True), + ) + + async def send_trace(trace_id: str) -> tuple[Sent, ...]: + clients: Final = { + id(pod): httpx.AsyncClient(base_url=str(pod.gateway.client.base_url), timeout=15, trust_env=False) + for pod in (pods.first, pods.second) + } + try: + return tuple( + await asyncio.gather( + *( + pods.send_async(clients[id(pod)], _call(sampled, trace_id, endpoint, stream)) + for pod, endpoint, stream in plan + ) + ) + ) + finally: + for opened in clients.values(): + await opened.aclose() + + by_trace: Final = { + trace_id: ( + tuple(pods.send(pod, _call(sampled, trace_id, endpoint, stream)) for pod, endpoint, stream in plan) + if client == "sync" + else asyncio.run(send_trace(trace_id)) + ) + for trace_id in _trace_ids(f"endpoints-{client}") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_the_workers_of_one_pod_agree(pods: Pods) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: pods.on_each_worker(pods.second, lambda trace_id=trace_id: _call(sampled, trace_id)) + for trace_id in _trace_ids("workers") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_replaced_worker_reaches_the_verdict_its_predecessor_reached(pods: Pods) -> None: + """A worker that restarts forgets the verdicts it remembered; the calls it serves next for the same + traces still answer the way the first ones did.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + fence_key: Final = pods.team_key(CREDENTIALS) + traces: Final = _trace_ids("replaced") + before: Final = { + trace_id: pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) for trace_id in traces + } + _await_logged(pods, cursors.operator, tuple(chain.from_iterable(before.values()))) + _await_flushed(pods, cursors, fence_key) + for victim in sorted(pods.second.worker_pids()): + pods.second.kill_worker(victim) + by_trace: Final = { + trace_id: before[trace_id] + pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) + for trace_id in traces + } + + verdicts: Final = _settled(pods, cursors, by_trace, fence_key) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_failed_call_is_kept_wherever_a_successful_call_of_its_trace_is(pods: Pods) -> None: + """One call of the trace fails on one pod and answers to the error rate, another succeeds on the other + pod and answers to the success rate: with the error rate above the success rate, no trace keeps its + successful call and drops its failed one.""" + sampled: Final = pods.team_key( + {**CREDENTIALS, "arize_success_sampling_rate": "0.3", "arize_error_sampling_rate": "0.7"} + ) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: ( + pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")), + pods.send(pods.second, _call(sampled, trace_id)), + ) + for trace_id in _trace_ids("mixed", 32) + } + + verdicts: Final = _settled( + pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[0] or not kept[1] + ) + + assert any(kept == (True, False) for kept in verdicts.values()), "no trace sat between the two rates" + assert any(kept == (True, True) for kept in verdicts.values()), "no trace sat under both rates" + + +def test_a_failed_call_answers_to_the_error_rate_on_a_pod_that_dropped_its_trace(pods: Pods) -> None: + """A pod that already dropped a successful call of the trace still keeps a later failed call of it at an + unset error rate, as the pod that only served the failed call does.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.3"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("failed-after-dropped") + succeeded: Final = {trace_id: pods.send(pods.first, _call(sampled, trace_id)) for trace_id in traces} + _await_logged(pods, cursors.operator, tuple(succeeded.values())) + by_trace: Final = { + trace_id: ( + succeeded[trace_id], + pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")), + pods.send(pods.second, _call(sampled, trace_id, content="sampling-fail")), + ) + for trace_id in traces + } + + verdicts: Final = _settled( + pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[1] and kept[2] + ) + + assert any(not kept[0] for kept in verdicts.values()), "no successful call was dropped before its trace failed" + + +def test_a_successful_call_answers_to_the_success_rate_on_a_pod_that_kept_its_failed_call(pods: Pods) -> None: + """A pod that already kept a failed call of the trace drops a later successful call of it at a zero + success rate, as the pod that only served the successful call does.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("succeeded-after-kept") + failed: Final = { + trace_id: pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")) for trace_id in traces + } + _await_logged(pods, cursors.operator, tuple(failed.values())) + by_trace: Final = { + trace_id: ( + failed[trace_id], + pods.send(pods.first, _call(sampled, trace_id)), + pods.send(pods.second, _call(sampled, trace_id)), + ) + for trace_id in traces + } + + _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False, False)) + + +def test_concurrent_calls_of_one_trace_on_one_worker_each_answer_to_their_own_outcome(pods: Pods) -> None: + """A failed and a successful call of the trace in flight at once on the single worker of the first pod + are decided apart: the failed one is kept at an unset error rate, the successful one dropped at a zero + success rate.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("concurrent-mixed") + plan: Final = tuple(product(traces, ("sampling-fail", "hello"))) + + def send(job: tuple[str, str]) -> Sent: + return pods.send(pods.first, _call(sampled, job[0], content=job[1])) + + with ThreadPoolExecutor(max_workers=8) as pool: + sent: Final = tuple(pool.map(send, plan)) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False)) + + +@pytest.mark.parametrize( + ("variables", "expected"), + [({"arize_success_sampling_rate": "1.0"}, TRACES), ({"arize_success_sampling_rate": "0.0"}, 0), ({}, TRACES)], + ids=["one", "zero", "unset"], +) +def test_a_rate_at_either_end_keeps_every_trace_or_none( + pods: Pods, variables: Mapping[str, str], expected: int +) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, **variables}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second)) + for trace_id in _trace_ids(f"edge-{'-'.join(f'{k}={v}' for k, v in sorted(variables.items()))}") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert _kept(verdicts) == expected + + +def test_a_request_without_a_traceparent_is_sampled_whole_at_the_rate(pods: Pods) -> None: + """With no caller trace the proxy starts one per request: each request's server and model spans travel + together and the rate still keeps some requests and drops others.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + sent: Final = tuple( + pods.send(pod, _call(sampled, None)) for _, pod in product(range(12), (pods.first, pods.second)) + ) + _await_logged(pods, cursors.operator, sent) + trace_of: Final = { + str(span["attributes"][CALL_ID]): span["trace_id"] + for span in recorded_spans(pods.sinks.operator, cursors.operator)[1] + if CALL_ID in span["attributes"] + } + fences: Final = pods.fence(pods.team_key(CREDENTIALS)) + + def check() -> Mapping[str, tuple[bool, bool]]: + spans: Final = _team_spans(pods, cursors.arize) + assert {fence.call_id for fence in fences} <= _delivered(spans, fences), ( + "the unsampled team's calls have not arrived yet" + ) + kinds: Final = { + call.call_id: ( + any(span["kind"] == SERVER for span in spans if span["trace_id"] == trace_of[call.call_id]), + any(OPERATION in span["attributes"] for span in spans if span["trace_id"] == trace_of[call.call_id]), + ) + for call in sent + } + partial: Final = {call_id: kept for call_id, kept in kinds.items() if kept[0] != kept[1]} + assert partial == {}, f"requests reached the space without their whole tree (server, model): {partial}" + return kinds + + kept: Final = sum(1 for server, _ in until_asserted(check, seconds=30).values() if server) + assert 0 < kept < len(sent), f"a rate of 0.5 kept {kept} of {len(sent)} requests" + + +def test_a_key_level_destination_follows_the_verdict_of_its_trace(pods: Pods) -> None: + team: Final = pods.scenario.team() + sampled: Final = pods.scenario.key( + team_id=team, + metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {**CREDENTIALS, "arize_success_sampling_rate": "0.5"}, + } + ] + }, + ) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second, pods.first)) + for trace_id in _trace_ids("key-level") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_concurrent_burst_keeps_every_trace_whole(pods: Pods) -> None: + """Four calls per trace, all in flight at once across both pods.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("burst", 24) + plan: Final = tuple(product(traces, (pods.first, pods.second, pods.first, pods.second))) + + def send(job: tuple[str, OwnedProxy]) -> Sent: + return pods.send(job[1], _call(sampled, job[0])) + + with ThreadPoolExecutor(max_workers=16) as pool: + sent: Final = list(pool.map(send, plan)) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < len(traces) + + +def test_a_stalled_team_space_during_a_burst_keeps_serving_and_each_trace_whole(pods: Pods) -> None: + """The team's Arize endpoint stops answering while a burst of calls runs over both pods: every call is + still served, and once the endpoint answers again each trace arrives whole or not at all.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("stalled") + plan: Final = tuple(product(traces, (pods.first, pods.second))) + + def send(job: tuple[str, OwnedProxy]) -> Sent: + return pods.send(job[1], _call(sampled, job[0])) + + configure_sink(pods.sinks.arize, paused=True) + try: + with ThreadPoolExecutor(max_workers=8) as pool: + sent: Final = tuple(pool.map(send, plan)) + finally: + configure_sink(pods.sinks.arize, paused=False) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < len(traces) + + +def _write_config(directory: Path, sinks: SpanSinks, provider: Wire) -> Path: + shared: Final = object_value( + JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + ) + config: Final = { + **shared, + "model_list": [ + *list_value(shared.get("model_list") or []), + { + "model_name": MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + }, + }, + ], + "litellm_settings": { + **object_value(shared.get("litellm_settings") or {}), + "callbacks": ["langfuse_otel"], + "otel_tenant_destination_mode": "additive", + "provider_url_destination_allowed_hosts": [urlparse(sinks.operator).netloc, urlparse(sinks.arize).netloc], + }, + "router_settings": {**object_value(shared.get("router_settings") or {}), "num_retries": 0}, + "general_settings": {**object_value(shared.get("general_settings") or {}), "user_api_key_cache_ttl": 2}, + } + path: Final = directory / "team-sampling.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def _two_pods(provider: Wire, sinks: SpanSinks, directory: Path) -> Generator[Pods]: + environment: Final = { + "LITELLM_OTEL_V2": "1", + "OTEL_BSP_SCHEDULE_DELAY": "300", + "MAPPER_NAMES": "genai,openinference", + "LANGFUSE_HOST": sinks.operator, + "LANGFUSE_PUBLIC_KEY": "pk-lf-operator", + "LANGFUSE_SECRET_KEY": "sk-lf-operator", + "ARIZE_SPACE_KEY": "space-operator", + "ARIZE_API_KEY": "arize-operator", + "ARIZE_HTTP_ENDPOINT": sinks.arize + "/v1/traces", + } + removed: Final = ("LITELLM_OTEL_TENANT_DESTINATION_MODE", "OTEL_TRACES_SAMPLER", "OTEL_TRACES_SAMPLER_ARG") + config: Final = _write_config(directory, sinks, provider) + for pod in ("first", "second"): + (directory / pod).mkdir() + with gateway_from_environment() as shared, scratch_database("team_sampling") as database_url: + overrides: Final = {**environment, "DATABASE_URL": database_url} + with ( + owned_proxy_process( + shared, directory / "first", overrides, config=config, remove_environment=removed + ) as first, + owned_proxy_process( + shared, + directory / "second", + overrides, + config=config, + remove_environment=removed, + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + ) as second, + replace(first.gateway, database_url=database_url).scenario() as scenario, + ): + eventually(second.worker_pids, lambda pids: len(pids) == 2, seconds=120) + yield Pods(first, second, scenario, sinks) + + +@pytest.fixture(scope="module") +def provider() -> Iterator[Wire]: + with wire_server(_upstream) as wire: + yield wire + + +@pytest.fixture(scope="module") +def pods(provider: Wire, audit_sinks: SpanSinks, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Pods]: + with _two_pods(provider, audit_sinks, tmp_path_factory.mktemp("team-sampling")) as started: + yield started diff --git a/tests/integration/observability/test_passthrough_guardrail_scan_text.py b/tests/integration/observability/test_passthrough_guardrail_scan_text.py new file mode 100644 index 00000000000..7437f696d4c --- /dev/null +++ b/tests/integration/observability/test_passthrough_guardrail_scan_text.py @@ -0,0 +1,153 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import Gateway +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +GENERIC: Final = "passthrough-scan-generic" +AIRS: Final = "passthrough-scan-airs" +OPENAI_ROUTE: Final = "/openai/chat/completions" +ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" +ROUTE_BODIES: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + OPENAI_ROUTE: {"model": "gpt-5.6", "messages": [{"role": "user", "content": "Say hi in one word"}]}, + ANTHROPIC_ROUTE: { + "model": "claude-sonnet-5-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Say hi in one word"}], + }, + } +) +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + + +def _generic_verdict(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _airs_verdict(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "allow", "category": "benign", "scan_id": "s", "report_id": "r"}).encode()) + + +def _openai_reply(_request: Request) -> Reply: + return chat_reply("chatcmpl-passthrough-scan", "gpt-5.6", "Hi", stream=False) + + +def _anthropic_reply(_request: Request) -> Reply: + return Reply(body=message_json("msg_passthrough_scan", "claude-sonnet-5-5", "Hi")) + + +def _parsed_texts(texts: JsonValue, vendor_body: dict[str, JsonValue]) -> list[JsonValue]: + assert isinstance(texts, list), vendor_body + return [json.loads(text) if isinstance(text, str) else text for text in texts] + + +def _scanned_payloads(guardrail: str, vendor_body: dict[str, JsonValue]) -> list[JsonValue]: + if guardrail == GENERIC: + return _parsed_texts(vendor_body["texts"], vendor_body) + contents: Final = vendor_body["contents"] + assert isinstance(contents, list), vendor_body + return _parsed_texts([content.get("prompt") for content in contents if isinstance(content, dict)], vendor_body) + + +@dataclass(frozen=True, slots=True) +class PassThroughRig: + owned: OwnedProxy + vendors: Mapping[str, Wire] + upstreams: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[PassThroughRig]: + directory: Final = tmp_path_factory.mktemp("passthrough-guardrail-scan") + with ExitStack() as stack: + vendors: Final = MappingProxyType( + { + GENERIC: stack.enter_context(wire_server(_generic_verdict)), + AIRS: stack.enter_context(wire_server(_airs_verdict)), + } + ) + upstreams: Final = MappingProxyType( + { + OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), + ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)), + } + ) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + path: Final = directory / "passthrough-guardrail-scan.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "guardrails": [ + { + "guardrail_name": GENERIC, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": vendors[GENERIC].url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": AIRS, + "litellm_params": { + "guardrail": "panw_prisma_airs", + "mode": "pre_call", + "default_on": False, + "api_base": vendors[AIRS].url, + "api_key": "synthetic-airs-key", + "profile_name": "integration", + }, + }, + ], + } + ) + ) + environment: Final = { + "OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url, + "OPENAI_API_KEY": "synthetic-openai-key", + "ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + } + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) + yield PassThroughRig(owned, vendors, upstreams) + + +@pytest.mark.parametrize("guardrail", [GENERIC, AIRS]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE]) +def test_pass_through_guardrail_scans_the_client_body_without_the_callers_key_metadata( + rig: PassThroughRig, route: str, guardrail: str +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + team: Final = scenario.team() + user: Final = scenario.user( + user_email=f"passthrough-scan-{uuid.uuid4().hex}@example.com", user_role="internal_user", teams=[team] + ) + key: Final = scenario.key(user_id=user, team_id=team) + body: Final = {**ROUTE_BODIES[route], "guardrails": [guardrail]} + + response: Final = gateway.request("POST", route, body, key=key) + + assert response.status_code == 200, response.text + scanned: Final = [ + _scanned_payloads(guardrail, _JSON.validate_json(request.body)) + for request in rig.vendors[guardrail].drain() + ] + assert scanned == [[body]], scanned + assert len(rig.upstreams[route].drain()) == 1 diff --git a/tests/integration/sdk/test_sync_stream_connection_release_wire.py b/tests/integration/sdk/test_sync_stream_connection_release_wire.py index a8b49d285c3..21774df66d0 100644 --- a/tests/integration/sdk/test_sync_stream_connection_release_wire.py +++ b/tests/integration/sdk/test_sync_stream_connection_release_wire.py @@ -212,3 +212,36 @@ def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None: assert _open_sessions_settling_to(wire, at_most=0) == 0, ( "a stream that raised mid-way left its TCP session open" ) + + +_STUCK: Final = _Shape( + "hosted_vllm", + "hosted_vllm/stuck-model", + _OPENAI_OPENING, + (_openai_delta("the model is stuck on this", None),) * (10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT), +) + + +def _still_generating_peer(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=(_STUCK.opening, *_STUCK.closing), pause_between_chunks=0.01) + + +def test_sync_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = litellm.completion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in stuck: + pass + assert _open_sessions_settling_to(wire, at_most=0) == 0, "a stream stuck on one chunk left its TCP session open" + + +@pytest.mark.asyncio +async def test_async_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = await litellm.acompletion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in stuck: + pass + assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, ( + "a stream stuck on one chunk left its TCP session open" + ) diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index b4485bcc5f1..60600afab84 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -1921,6 +1921,62 @@ async def test_list_tools_classifies_upstream_status_for_post_failures(failure_m assert rate_limit is None +@pytest.mark.asyncio +@pytest.mark.parametrize("paginated", (False, True)) +async def test_sse_catalog_limit_reports_the_failure_and_allows_smaller_pages(paginated: bool) -> None: + from mcp.types import Tool + + limit: Final = 1024 + tools: Final = tuple( + Tool(name=f"tool_{index}", description="private-upstream-secret " * 25, input_schema={"type": "object"}) + for index in range(2) + ) + catalog: Final = tuple(tool.model_dump(by_alias=True, exclude_none=True) for tool in tools) + assert len(json.dumps({"tools": catalog}).encode()) > limit + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + assert isinstance(payload, (JSONRPCRequest, JSONRPCNotification)) + if isinstance(payload, JSONRPCNotification): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response(200, json={ + "jsonrpc": "2.0", "id": payload.id, "result": { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {"tools": {}}, "serverInfo": {"name": "catalog", "version": "1"}, + }, + }) + assert payload.method == "tools/list" + cursor: Final = (payload.params or {}).get("cursor") + page: Final = catalog[1:] if cursor == "next-page" else catalog[:1] if paginated else catalog + result: Final = { + "tools": page, + **({"nextCursor": "next-page"} if paginated and cursor is None else {}), + } + event: Final = f"data: {json.dumps({'jsonrpc': '2.0', 'id': payload.id, 'result': result})}\n\n" + assert (len(event.encode()) < limit) == paginated + return httpx2.Response(200, headers={"content-type": "text/event-stream"}, content=event) + + def transport_context(url: str, http_client: httpx2.AsyncClient) -> _TransportContext: + return streamable_http_client(url, http_client=http_client, max_sse_event_size=limit) + + client: Final = _MockTransportClient( + respond, server_url="https://upstream.example/mcp", transport_context=transport_context, + ) + if paginated: + assert await client.list_tools(raise_on_error=True) == list(tools) + return + with pytest.raises(MCPError) as caught: + await client.list_tools(raise_on_error=True) + fault: Final = classify_list_exception(caught.value) + assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit) + assert list_fault_http_status(fault) == 502 + assert fault.public_message is not None + assert "private-upstream-secret" not in fault.public_message + + @pytest.mark.asyncio async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index c7a057cca5c..782442aad22 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -4,6 +4,7 @@ import asyncio import contextvars import threading import time +import uuid from base64 import b64encode from collections.abc import Mapping from dataclasses import replace @@ -14,10 +15,12 @@ from typing import Final import pytest from opentelemetry import trace as trace_api +from opentelemetry.context import Context from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.id_generator import RandomIdGenerator from opentelemetry.trace import Status, StatusCode import litellm @@ -1159,12 +1162,36 @@ class TestFanOut: ARIZE_TEAM_PARAMS = {"arize_space_id": "space-team", "arize_api_key": "key-team"} +def caller_context(trace_id: int) -> Context: + return trace_api.set_span_in_context( + trace_api.NonRecordingSpan( + trace_api.SpanContext(trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01)) + ) + ) + + def arize_destination(**rates: str) -> OtelDestination: destination = destination_for("arize", {**ARIZE_TEAM_PARAMS, **rates}) assert destination is not None, "the fixture must resolve for the test to mean anything" return destination +#: Trace ids whose SHA-256 reads 0.40 and 0.80 of the way through its first eight bytes, +#: so a rate of 0.5 keeps the first and drops the second. +KEPT_AT_HALF: Final = 4 +DROPPED_AT_HALF: Final = 5 + + +class _FixedTraceId(RandomIdGenerator): + """Every trace the provider starts gets ``trace_id``.""" + + def __init__(self, trace_id: int) -> None: + self._trace_id = trace_id + + def generate_trace_id(self) -> int: + return self._trace_id + + class TestTenantSampling: """A team's ``arize_success_sampling_rate`` and ``arize_error_sampling_rate`` decide how much of its traffic the fan-out delivers to its space, the way they already decide what @@ -1173,14 +1200,17 @@ class TestTenantSampling: NAMES = {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat gpt-4"} @staticmethod - def _fan_out(dest_exporter, draw=None, global_exporter=None): - provider = TracerProvider() + def _fan_out( + dest_exporter: InMemorySpanExporter, + trace_id: int | None = None, + global_exporter: InMemorySpanExporter | None = None, + ) -> TracerProvider: + provider: Final = TracerProvider() if trace_id is None else TracerProvider(id_generator=_FixedTraceId(trace_id)) if global_exporter is not None: provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "arize")) - kwargs = {"processor_factory": lambda _d: SimpleSpanProcessor(dest_exporter)} - if draw is not None: - kwargs["sampling_draw"] = draw - provider.add_span_processor(TenantFanOutSpanProcessor(**kwargs)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) return provider @staticmethod @@ -1224,7 +1254,7 @@ class TestTenantSampling: def test_an_unset_rate_exports_everything(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.99) + provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF) self._run(provider, (arize_destination(),)) @@ -1300,10 +1330,10 @@ class TestTenantSampling: def test_a_span_that_ends_after_the_root_follows_the_requests_verdict(self): """A post-call database write that starts after the whole request tree has ended - goes where the rest of the tree went, without a second draw.""" + goes where the rest of the tree went, without a second verdict.""" kept_exporter, dropped_exporter = InMemorySpanExporter(), InMemorySpanExporter() - kept = self._fan_out(kept_exporter, draw=lambda: 0.1) - dropped = self._fan_out(dropped_exporter, draw=lambda: 0.9) + kept: Final = self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF) + dropped: Final = self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF) destinations = (arize_destination(arize_success_sampling_rate="0.5"),) def late_tree(provider): @@ -1327,14 +1357,263 @@ class TestTenantSampling: } assert dropped_exporter.get_finished_spans() == () + def test_a_span_that_ends_after_a_failed_root_follows_the_failed_requests_verdict(self): + """A post-call database write after a failed request goes where that request went, + though on its own it would answer to the success rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + root_context: Final = trace_api.set_span_in_context(root) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "postgres INSERT LiteLLM_SpendLogs", + } + + def test_a_failed_call_after_a_dropped_call_of_its_trace_still_answers_to_the_error_rate(self): + """A process that dropped a successful call of the trace keeps a later failed call of + it at an unset error rate, as a process that only served the failed call does.""" + served_both, served_failed = InMemorySpanExporter(), InMemorySpanExporter() + destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),) + both: Final = self._fan_out(served_both, trace_id=DROPPED_AT_HALF) + + self._run(both, destinations) + self._run(both, destinations, failed=True) + self._run(self._fan_out(served_failed, trace_id=DROPPED_AT_HALF), destinations, failed=True) + + assert [sorted(s.name for s in exporter.get_finished_spans()) for exporter in (served_both, served_failed)] == [ + sorted(self.NAMES) + ] * 2 + + def test_a_successful_call_after_a_kept_failed_call_of_its_trace_answers_to_the_success_rate(self): + """A process that kept a failed call of the trace drops a later successful call of it + at a zero success rate, as a process that only served the successful call does.""" + served_both, served_success = InMemorySpanExporter(), InMemorySpanExporter() + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + both: Final = self._fan_out(served_both, trace_id=KEPT_AT_HALF) + + self._run(both, destinations, failed=True) + self._run(both, destinations) + self._run(self._fan_out(served_success, trace_id=KEPT_AT_HALF), destinations) + + assert len(served_both.get_finished_spans()) == len(self.NAMES), "only the failed call" + assert all(s.name in self.NAMES for s in served_both.get_finished_spans()) + assert served_success.get_finished_spans() == () + + def test_a_late_child_that_fails_after_its_request_was_dropped_is_dropped_with_it(self): + """A post-call write that fails after its successful request was dropped goes where + the request went, rather than reaching the space on its own at the error rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + root_context: Final = trace_api.set_span_in_context(root) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context) as write: + write.set_status(Status(StatusCode.ERROR, "connection reset")) + + in_fresh_context(run) + + assert dest_exporter.get_finished_spans() == () + + def test_a_failure_deep_in_a_request_marks_the_whole_request(self): + """A database span that fails under the auth span puts the whole request, root + included, on the error rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("auth /v1/chat/completions"): + with tracer.start_as_current_span("postgres SELECT LiteLLM_VerificationToken") as lookup: + lookup.set_status(Status(StatusCode.ERROR, "connection reset")) + with tracer.start_as_current_span("chat gpt-4"): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + *self.NAMES, + "postgres SELECT LiteLLM_VerificationToken", + } + + def test_a_span_two_levels_down_that_outlives_many_other_spans_still_closes_its_request(self): + """A classifier call under the route span that stays open while the worker serves + more requests than it tracks at once still counts against its own request, so the + request is decided once it ends, on its own outcome.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def other_request(): + set_request_destinations(()) + with get_tracer(provider, "litellm").start_as_current_span("POST /v1/chat/completions"): + pass + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("route gpt-4"): + with tracer.start_as_current_span("chat classifier"): + for _ in range(otel_providers._MAX_OPEN_REQUESTS + 1): + contextvars.Context().run(other_request) + with tracer.start_as_current_span("chat gpt-4"): + pass + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "chat classifier", + "chat gpt-4", + } + + def test_a_failure_two_levels_down_reaches_its_request_however_many_spans_are_open(self): + """A classifier call that fails while another request holds more spans open than + the worker tracks requests marks its own request, which is then kept at the error + rate as a whole.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + others: Final = contextvars.Context() + + def open_other_request(): + set_request_destinations(()) + tracer: Final = get_tracer(provider, "litellm") + root: Final = tracer.start_span("POST /v1/chat/completions") + parent: Final = trace_api.set_span_in_context(root) + return ( + root, + *(tracer.start_span("chat flood", context=parent) for _ in range(otel_providers._MAX_OPEN_REQUESTS)), + ) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("route gpt-4"): + with tracer.start_as_current_span("chat classifier") as classifier: + held_open: Final = others.run(open_other_request) + classifier.set_status(Status(StatusCode.ERROR, "classifier timed out")) + return held_open + + held_open: Final = in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "chat classifier", + } + for span in held_open: + span.end() + + def test_the_fan_out_forgets_a_span_once_it_ends_and_a_request_once_it_is_pushed_out(self): + """The request lookup holds only spans still open, and only those of the requests + the fan-out still counts, so neither ended spans nor a flood of open requests grow + it past the open request bound.""" + fan_out: Final = TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(InMemorySpanExporter()) + ) + provider: Final = TracerProvider() + provider.add_span_processor(fan_out) + + def run(): + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("chat gpt-4"): + pass + assert (fan_out._roots, dict(fan_out._open)) == ({}, {}) + root: Final = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) + tracer.start_span("route gpt-4", context=trace_api.set_span_in_context(root)) + return tuple( + tracer.start_span(f"POST {index}", context=trace_api.Context()) + for index in range(otel_providers._MAX_OPEN_REQUESTS) + ) + + held_open: Final = in_fresh_context(run) + + assert len(fan_out._roots) == otel_providers._MAX_OPEN_REQUESTS + for span in held_open: + span.end() + assert (fan_out._roots, dict(fan_out._open)) == ({}, {}) + + def test_a_span_starting_under_an_ended_span_below_the_root_follows_its_request(self): + """A post-call write started under the route span after its failed request ended + is kept with the request at the error rate, rather than decided on its own at the + success rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("route gpt-4") as route: + route_context: Final = trace_api.set_span_in_context(route) + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=route_context): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "postgres INSERT LiteLLM_SpendLogs", + } + + def test_concurrent_calls_of_one_trace_each_answer_to_their_own_outcome(self): + """Two calls of the caller's trace in flight at once on one process, one failing, + are decided as two processes serving one each would decide them.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + remote: Final = trace_api.SpanContext( + trace_id=KEPT_AT_HALF, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01) + ) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + caller: Final = trace_api.set_span_in_context(trace_api.NonRecordingSpan(remote)) + with tracer.start_as_current_span("POST failed", context=caller) as failed: + with tracer.start_as_current_span("POST succeeded", context=caller): + pass + failed.set_status(Status(StatusCode.ERROR, "upstream 500")) + + in_fresh_context(run) + + assert [s.name for s in dest_exporter.get_finished_spans()] == ["POST failed"] + def test_a_late_span_whose_verdict_was_forgotten_is_decided_on_its_own_not_stranded(self): """Once enough other requests have been decided to evict a request's verdict, a span of it that starts late (a post-call database write) opens and closes its own count, - so it is decided on a draw of its own as soon as it ends instead of waiting in a tree - nothing would ever close.""" + so it is decided on its own as soon as it ends instead of waiting in a tree nothing + would ever close.""" dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.1) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1353,18 +1632,17 @@ class TestTenantSampling: assert names[0] == "POST /v1/chat/completions" assert names[-1] == "postgres INSERT LiteLLM_SpendLogs", "exported as it ended, not held for a root" - def test_a_callers_traceparent_neither_steers_the_draw_nor_hides_the_root(self): - """The draw is random, not read from the trace id a caller may have chosen, and the - server span under a remote parent still closes the tree.""" + def test_a_callers_traceparent_decides_by_its_trace_id_and_still_closes_the_tree(self): + """A request continuing a caller's trace is decided on that trace id, and the server + span under a remote parent still closes the tree.""" dest_exporter = InMemorySpanExporter() - draws = iter([0.9, 0.1]) - provider = self._fan_out(dest_exporter, draw=lambda: next(draws)) + provider: Final = self._fan_out(dest_exporter) destinations = (arize_destination(arize_success_sampling_rate="0.5"),) def run(): set_request_destinations(destinations) tracer = get_tracer(provider, "litellm") - for trace_id in (1, 2): + for trace_id in (DROPPED_AT_HALF, KEPT_AT_HALF): remote = trace_api.SpanContext( trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01) ) @@ -1377,12 +1655,12 @@ class TestTenantSampling: kept = dest_exporter.get_finished_spans() assert [s.name for s in kept] == ["chat gpt-4", "POST /v1/chat/completions"] - assert all(s.context.trace_id == 2 for s in kept), "the second request, whose draw of 0.1 passed" + assert all(s.context.trace_id == KEPT_AT_HALF for s in kept) def test_a_root_that_never_ends_holds_the_tree_only_up_to_the_bound(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1397,8 +1675,8 @@ class TestTenantSampling: def test_a_flood_of_waiting_trees_decides_the_oldest_on_what_it_has(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1413,21 +1691,21 @@ class TestTenantSampling: assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat 0"] - def test_a_flood_of_open_traces_forgets_the_oldest_and_decides_it_as_its_spans_end(self): - """The fan-out counts open spans for a bounded number of traces, a bound far above - what one process keeps in flight, since every trace of the provider is counted, not - only the sampled ones. A trace pushed out of that count is decided on what it has as + def test_a_flood_of_open_requests_forgets_the_oldest_and_decides_it_as_its_spans_end(self): + """The fan-out counts open spans for a bounded number of requests, a bound far above + what one process keeps in flight, since every request of the provider is counted, not + only the sampled ones. A request pushed out of that count is decided on what it has as its next span ends, root still open, rather than held until a bound gets to it.""" dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) tracer = get_tracer(provider, "litellm") root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) tracer.start_span("auth /v1/chat/completions", context=trace_api.set_span_in_context(root)).end() - for index in range(otel_providers._MAX_OPEN_TRACES): + for index in range(otel_providers._MAX_OPEN_REQUESTS): tracer.start_span(f"POST {index}", context=trace_api.Context()) assert dest_exporter.get_finished_spans() == (), "still held before the count is forgotten" tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end() @@ -1443,14 +1721,12 @@ class TestTenantSampling: def test_shutdown_decides_what_is_still_held(self): dest_exporter = InMemorySpanExporter() - fan_out = TenantFanOutSpanProcessor( - processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter), sampling_draw=lambda: 0.0 - ) + fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) provider = TracerProvider() provider.add_span_processor(fan_out) def run(): - set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),)) + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) tracer = get_tracer(provider, "litellm") with tracer.start_as_current_span("POST /v1/chat/completions"): with tracer.start_as_current_span("chat gpt-4"): @@ -1476,14 +1752,12 @@ class TestTenantSampling: straggler.end() super().on_end(span) - fan_out = TenantFanOutSpanProcessor( - processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter), sampling_draw=lambda: 0.0 - ) + fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter)) provider = TracerProvider() provider.add_span_processor(fan_out) def run(): - set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),)) + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) tracer = get_tracer(provider, "litellm") root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end() @@ -1495,20 +1769,20 @@ class TestTenantSampling: assert {s.name for s in dest_exporter.get_finished_spans()} == {"chat gpt-4", "embed ada"} - def test_a_zero_rate_drops_even_a_draw_of_zero(self): + def test_a_zero_rate_drops_a_trace_a_half_rate_keeps(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) + provider: Final = self._fan_out(dest_exporter, trace_id=KEPT_AT_HALF) self._run(provider, (arize_destination(arize_success_sampling_rate="0.0"),)) assert dest_exporter.get_finished_spans() == () - def test_the_draw_is_compared_to_the_rate_the_way_the_legacy_callback_compares_it(self): + def test_the_trace_id_is_read_against_the_rate(self): dropped_exporter, kept_exporter = InMemorySpanExporter(), InMemorySpanExporter() - destinations = (arize_destination(arize_success_sampling_rate="0.2"),) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),) - self._run(self._fan_out(dropped_exporter, draw=lambda: 0.3), destinations) - self._run(self._fan_out(kept_exporter, draw=lambda: 0.2), destinations) + self._run(self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF), destinations) + self._run(self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF), destinations) assert dropped_exporter.get_finished_spans() == () assert {s.name for s in kept_exporter.get_finished_spans()} == self.NAMES @@ -1527,6 +1801,39 @@ class TestTenantSampling: assert all(names == self.NAMES for names in by_trace.values()) assert 0 < len(by_trace) < 64 + @pytest.mark.parametrize( + ("rates", "failed"), + [({"arize_success_sampling_rate": "0.5"}, False), ({"arize_error_sampling_rate": "0.5"}, True)], + ids=["success-rate", "error-rate"], + ) + def test_two_processes_serving_one_trace_reach_the_same_verdict(self, rates: Mapping[str, str], failed: bool): + """Two pods behind a load balancer each serve one call of a trace the caller + continued with its ``traceparent``: the team's space gets both calls or neither, + and a rate of 0.5 over UUID-shaped trace ids keeps some traces and drops others.""" + exporters: Final = (InMemorySpanExporter(), InMemorySpanExporter()) + pods: Final = tuple(self._fan_out(exporter) for exporter in exporters) + destinations: Final = (arize_destination(**rates),) + trace_ids: Final = tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"pod/{index}").int for index in range(32)) + + def run(provider: TracerProvider): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + for trace_id in trace_ids: + with tracer.start_as_current_span("POST /v1/chat/completions", context=caller_context(trace_id)): + with tracer.start_as_current_span("chat gpt-4") as span: + if failed: + span.set_status(Status(StatusCode.ERROR, "upstream 500")) + + for provider in pods: + in_fresh_context(run, provider) + + first, second = ( + {s.context.trace_id for s in exporter.get_finished_spans() if s.context is not None} + for exporter in exporters + ) + assert first == second + assert 0 < len(first) < len(trace_ids) + def test_in_additive_mode_the_operators_own_copy_is_not_sampled(self, monkeypatch): """The rates are the team's setting for the team's space; the operator's backbone keeps every span.""" @@ -1547,7 +1854,7 @@ class TestTenantSampling: assert {s.name for s in dest_exporter.get_finished_spans()} == self.NAMES - def test_two_views_of_one_exporter_each_draw_at_their_own_rate(self): + def test_two_views_of_one_exporter_each_answer_to_their_own_rate(self): """Two destinations to the same space with different rates share an exporter, not a verdict: the 1.0 view gets the tree once and the 0.0 view adds nothing.""" dest_exporter = InMemorySpanExporter() @@ -1582,7 +1889,7 @@ class TestTenantSampling: @pytest.mark.parametrize("bad", ["abc", "-0.1", "1.5", "nan"]) def test_an_unusable_rate_exports_rather_than_dropping(self, bad): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.99) + provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF) self._run(provider, (arize_destination(arize_success_sampling_rate=bad),)) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index c13cea51183..cc1fa1b8803 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,6 +1,7 @@ import asyncio import datetime import json +import queue import sys import threading import time @@ -2819,6 +2820,123 @@ def test_raise_on_model_repetition_tolerates_empty_choices( wrapper.raise_on_model_repetition() +_STUCK_CHUNK: Final = "the model is stuck on this" +_STUCK_PROVIDER_CHUNK_CAP: Final = 10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + +class _StuckProvider: + def __init__(self) -> None: + self.closed_early = threading.Event() + self.chunks_sent = 0 + + def _next_chunk(self) -> ModelResponseStream: + self.chunks_sent += 1 + return _make_chunk(_STUCK_CHUNK) + + def sync_stream(self) -> Iterator[ModelResponseStream]: + try: + yield from (self._next_chunk() for _ in range(_STUCK_PROVIDER_CHUNK_CAP)) + except GeneratorExit: + self.closed_early.set() + raise + + async def async_stream(self) -> AsyncIterator[ModelResponseStream]: + try: + for _ in range(_STUCK_PROVIDER_CHUNK_CAP): + yield self._next_chunk() + except GeneratorExit: + self.closed_early.set() + raise + + +class _StuckStreamFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: Final = queue.SimpleQueue[str]() + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + +def _stuck_stream_wrapper( + provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream], + call_type: Literal["completion", "acompletion"], + recorder: _StuckStreamFailureRecorder, +) -> CustomStreamWrapper: + logging_obj: Final = Logging( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type=call_type, + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id=f"stuck-stream-{call_type}", + function_id=f"stuck-stream-{call_type}", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="test-model", + optional_params={}, + litellm_params=get_litellm_params(custom_llm_provider="openai", acompletion=call_type == "acompletion"), + ) + return CustomStreamWrapper( + completion_stream=provider_stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + +@pytest.mark.asyncio +async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.async_stream(), "acompletion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + await wrapper.wait_for_failure_logging(timeout=5) + + assert provider.chunks_sent == chunks_sent_at_failure + assert "repeating the same chunk" in recorder.failures.get_nowait() + assert recorder.failures.empty() + + +def test_sync_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """Sync counterpart of the async test above""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.sync_stream(), "completion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + assert "repeating the same chunk" in recorder.failures.get(timeout=5) + + with pytest.raises(StopIteration): + next(wrapper) + + assert provider.chunks_sent == chunks_sent_at_failure + with pytest.raises(queue.Empty): + recorder.failures.get(timeout=0.5) + + def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): """ Test that provider-reported usage from a post-finish_reason chunk diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 0dba0d7b90f..7cbabb95931 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -19,6 +19,22 @@ def test_full_payload_guardrail_text_excludes_the_server_streaming_marker(): assert json.loads(text) == body, text +def test_full_payload_guardrail_text_excludes_the_injected_litellm_metadata(): + body: Final = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "guardrails": ["g"]} + litellm_metadata: Final = { + "user_api_key_token": "hashed-key", + "user_api_key_user_email": "someone@example.com", + "user_api_key_user_role": "proxy_admin", + } + + text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail( + {**body, "litellm_metadata": litellm_metadata}, + None, + ) + + assert json.loads(text) == body, text + + @pytest.mark.parametrize( "marker", [ diff --git a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 630b2e1eb34..8abd012f9d8 100644 --- a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -5,13 +5,19 @@ Based on the test patterns from other rerank providers and the current Vertex AI import json import os +from datetime import datetime, timedelta +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import responses +import respx +import litellm from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.types.rerank import RerankResponse +from litellm.utils import ProviderConfigManager class TestVertexAIRerankTransform: @@ -646,3 +652,83 @@ class TestVertexAIRerankTransform: second_call = mock_ensure_access_token.call_args_list[1] assert first_call.kwargs["credentials"] == "path/to/credentials.json" assert second_call.kwargs["credentials"] == "path/to/credentials.json" + + +@pytest.fixture +def clear_vertex_rerank_cache(request: pytest.FixtureRequest) -> None: + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear() + request.addfinalizer( + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("clear_vertex_rerank_cache") +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_rerank_reuses_scoped_credentials_and_refreshes_expired_tokens( + is_async: bool, +) -> None: + now: Final = datetime(2025, 1, 1) + credentials: Final = tuple( + json.dumps( + { + "type": "authorized_user", + "client_id": identity, + "client_secret": "synthetic-client-secret", + "refresh_token": f"synthetic-refresh-{identity}", + } + ) + for identity in ("first-principal", "second-principal") + ) + with ( + patch("litellm.disable_aiohttp_transport", True), + patch("google.auth._helpers.utcnow", return_value=now) as clock, + responses.RequestsMock(assert_all_requests_are_fired=False) as token_http, + respx.mock as rank_http, + ): + for token in ("initial", "other-credential", "other-project", "renewed"): + token_http.post( + "https://oauth2.googleapis.com/token", + json={"access_token": token, "expires_in": 3600, "token_type": "Bearer"}, + ) + rank_http.post( + url__regex=r"https://discoveryengine\.googleapis\.com/v1/projects/[^/]+/" + r"locations/global/rankingConfigs/default_ranking_config:rank" + ).respond(200, json={"records": [{"id": "0", "score": 0.9}]}) + + async def rank( + credential: str, project: str, expected_token: str, refresh_count: int + ) -> None: + params: Final = { + "model": "vertex_ai/semantic-ranker-default@latest", + "query": f"Query for {expected_token}", + "documents": ["A relevant document"], + "vertex_credentials": credential, + "vertex_project": project, + "caching": False, + "num_retries": 0, + } + result: Final = ( + await litellm.arerank(**params) + if is_async + else litellm.rerank(**params) + ) + assert isinstance(result, RerankResponse) + assert result.results[0]["index"] == 0 + assert len(token_http.calls) == refresh_count + request: Final = rank_http.calls.last.request + assert request.headers["Authorization"] == f"Bearer {expected_token}" + assert request.headers["X-Goog-User-Project"] == project + assert str(request.url) == ( + f"https://discoveryengine.googleapis.com/v1/projects/{project}/" + "locations/global/rankingConfigs/default_ranking_config:rank" + ) + assert json.loads(request.content)["query"] == params["query"] + + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[1], "first-project", "other-credential", 2) + await rank(credentials[0], "second-project", "other-project", 3) + await rank(credentials[0], "first-project", "initial", 3) + clock.return_value = now + timedelta(hours=2) + await rank(credentials[0], "first-project", "renewed", 4) diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index f6902c778e5..fdae575276d 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -9,9 +9,10 @@ if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 from exceptiongroup import BaseExceptionGroup import httpx +import httpx2 import pytest from mcp import MCPError -from mcp.types import ErrorData +from mcp.types import CONNECTION_CLOSED, ErrorData from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -102,6 +103,68 @@ def test_unknown_exception_is_internal(): assert classify_list_exception(ValueError("who knows")).tag == "internal" +@pytest.mark.parametrize("sdk_wrapped", (False, True)) +@pytest.mark.parametrize("nested", (False, True)) +def test_sse_event_limit_is_a_safe_upstream_fault(sdk_wrapped: bool, nested: bool) -> None: + limit: Final = 1024 + decoder_message: Final = f"Server-sent event exceeded the {limit} byte limit." + failure: Final = ( + MCPError(code=CONNECTION_CLOSED, message=f"SSE stream failed: {decoder_message}") + if sdk_wrapped else httpx2.SSEError(decoder_message) + ) + wrapped: Final = BaseExceptionGroup("private-upstream-secret", [failure]) if nested else failure + fault: Final = classify_list_exception(wrapped) + message: Final = ( + f"Upstream MCP SSE event exceeded the {limit} byte per-event limit. " + "Paginate tools/list responses so each event stays below the limit." + ) + assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit) + assert fault.public_message == message + assert list_fault_http_status(fault) == 502 + assert outcome_wire_value(fault) == { + "status": "upstream_error", "reason": "sse_event_too_large", "limit_bytes": limit, "message": message, + } + + +@pytest.mark.parametrize( + "failure", + ( + MCPError(code=CONNECTION_CLOSED, message="Connection closed"), + MCPError(code=-32603, message="SSE stream failed: Server-sent event exceeded the 1024 byte limit."), + MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the 0 byte limit."), + MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the -1 byte limit."), + MCPError(code=CONNECTION_CLOSED, message=( + "SSE stream failed: Server-sent event exceeded the 10000000000000000000 byte limit." + )), + MCPError(code=CONNECTION_CLOSED, message=( + "SSE stream failed: Server-sent event exceeded the 1024 byte limit. private-secret" + )), + httpx2.SSEError("private-secret Server-sent event exceeded the 1024 byte limit."), + httpx2.SSEError("Expected response with content type 'text/event-stream', got 'private-secret'."), + RuntimeError("SSE stream failed: Server-sent event exceeded the 1024 byte limit."), + ), +) +def test_unrelated_errors_do_not_claim_an_sse_event_limit(failure: Exception) -> None: + fault: Final = classify_list_exception(failure) + assert fault.sse_event_limit_bytes is None + assert fault.public_message is None + assert "private-secret" not in str(outcome_wire_value(fault)) + + +@pytest.mark.parametrize("status_code", (401, 403)) +def test_upstream_auth_takes_precedence_over_sse_event_limit(status_code: int) -> None: + request: Final = httpx2.Request("POST", "https://upstream.example/mcp") + response: Final = httpx2.Response(status_code, request=request) + failure: Final = BaseExceptionGroup("upstream errors", [ + httpx2.SSEError("Server-sent event exceeded the 1024 byte limit."), + httpx2.HTTPStatusError("private-secret", request=request, response=response), + ]) + fault: Final = classify_list_exception(failure) + assert fault.tag == ("auth_required" if status_code == 401 else "forbidden") + assert list_fault_http_status(fault) == status_code + assert fault.sse_event_limit_bytes is None + + def test_wire_value_carries_no_prose(): fault = ServerListFault(tag="upstream_error", status_code=500) assert outcome_wire_value(fault) == {"status": "upstream_error", "http_status": 500} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3efbe6475b2..a132ce8eada 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3,6 +3,7 @@ import asyncio import contextlib import io +import json import logging from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator from types import SimpleNamespace @@ -35,6 +36,7 @@ from litellm.llms.openai.chat.guardrail_translation.handler import ( from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) @@ -3053,3 +3055,34 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none(): ) assert len(guardrail.apply_calls) == 1 + + +@pytest.mark.parametrize("hook", ["async_pre_call_hook", "async_moderation_hook"]) +@pytest.mark.asyncio +async def test_pass_through_guardrail_text_is_the_client_body_without_the_injected_key_metadata( + monkeypatch: pytest.MonkeyPatch, hook: str +) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.pass_through: PassThroughEndpointHandler}) + guardrail: Final = RecordingGuardrail() + body: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Say hi in one word"}], + "guardrails": ["recording-guardrail"], + } + data: Final = {**body, "guardrail_to_apply": guardrail} + user_api_key_dict: Final = UserAPIKeyAuth( + token="hashed-key", user_email="someone@example.com", team_id="team-1", user_role="proxy_admin" + ) + + if hook == "async_pre_call_hook": + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=CallTypes.pass_through.value + ) + else: + await UnifiedLLMGuardrails().async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=CallTypes.pass_through.value + ) + + texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] + assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts + assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 38d4c435e9a..32777b43504 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2175,11 +2175,6 @@ "count": 1 } }, - "src/utils/dataUtils.test.ts": { - "max-nested-callbacks": { - "count": 1 - } - }, "tests/CreateKeyPage.expiredToken.test.tsx": { "@typescript-eslint/no-require-imports": { "count": 3 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx index 35434830d98..d7dfb8151fc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx @@ -80,6 +80,42 @@ describe("GeneralSettings General tab", () => { vi.mocked(deleteConfigFieldSetting).mockClear(); }); + it.each([ + ["maximum_spend_logs_retention_period", null, "7d"], + ["maximum_health_check_retention_period", "30d", "14d"], + ["maximum_daily_tag_spend_retention_period", undefined, "90d"], + ["maximum_spend_logs_cleanup_run_budget", "5m", "10m"], + ["maximum_spend_logs_cleanup_batch_timeout", "", "45s"], + ["mcp_client_id_header", "x-client-id", "X-MCP-Client.ID"], + ])("displays, updates and resets the String setting %s", async (fieldName, initialValue, updatedValue) => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + ...SETTINGS_FIXTURE, + { + field_name: fieldName, + field_type: "String", + field_value: initialValue, + field_description: "string setting", + stored_in_db: initialValue != null, + field_default_value: null, + }, + ]); + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("tab", { name: "General" })); + const input = await screen.findByRole("textbox", { name: fieldName }); + expect(input).toHaveValue(initialValue ?? ""); + fireEvent.change(input, { target: { value: updatedValue } }); + const row = screen.getByRole("row", { name: new RegExp(fieldName) }); + await user.click(within(row).getByRole("button", { name: "Update" })); + expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([["token", fieldName, updatedValue]]); + expect(within(row).getByText("In DB")).toBeInTheDocument(); + const resetIcon = row.querySelectorAll("td")[3].querySelector("svg"); + expect(resetIcon).not.toBeNull(); + await user.click(resetIcon as unknown as Element); + expect(vi.mocked(deleteConfigFieldSetting).mock.calls).toEqual([["token", fieldName]]); + expect(input).toHaveValue(""); + }); + it("updates max_ui_session_budget with its own value, not the value at its filtered index", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 5d28adb7ef3..67712d78a13 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -68,6 +68,15 @@ const SettingValueEditor: React.FC<{ /> ); } + if (setting.field_type === "String") { + return ( + onChange(setting.field_name, event.target.value)} + /> + ); + } if (setting.field_type === "Integer") { return ( ({ + ...(await importOriginal()), + copyToClipboard: vi.fn(), +})); + vi.mock("../../../common_components/DefaultProxyAdminTag", () => ({ default: ({ userId }: { userId: string }) => {userId}, })); @@ -120,16 +126,17 @@ describe("AuditLogDrawer", () => { expect(screen.queryByText(/"spend"/)).not.toBeInTheDocument(); }); - it("should copy the json of a block to the clipboard", async () => { + it.each([true, false])("confirms copied JSON only when copying succeeds (%s)", async (success) => { const user = userEvent.setup(); - const writeText = vi.fn().mockResolvedValue(undefined); - Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true }); - Object.defineProperty(window, "isSecureContext", { value: true, configurable: true }); - + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(success); render(); - await user.click(within(blockNamed("Before")).getByTitle("Copy JSON")); - await waitFor(() => expect(writeText).toHaveBeenCalledWith(JSON.stringify({ max_budget: 10 }, null, 2))); + const button = within(blockNamed("Before")).getByRole("button", { name: "Copy JSON" }); + await user.click(button); + + expect(copyToClipboard).toHaveBeenCalledWith(JSON.stringify({ max_budget: 10 }, null, 2)); + expect(button.querySelector(".lucide-check") !== null).toBe(success); }); describe("width", () => { diff --git a/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx b/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx index fae3c35b80c..eaddf1e507f 100644 --- a/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx +++ b/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { Check, Copy } from "lucide-react"; import { useState, useCallback } from "react"; import moment from "moment"; @@ -27,25 +28,9 @@ function CopyableJsonBlock({ label, value }: { label: string; value: Record { - try { - const text = JSON.stringify(value, null, 2); - if (navigator.clipboard && window.isSecureContext) { - await navigator.clipboard.writeText(text); - } else { - const el = document.createElement("textarea"); - el.value = text; - el.style.position = "fixed"; - el.style.opacity = "0"; - document.body.appendChild(el); - el.focus(); - el.select(); - document.execCommand("copy"); - document.body.removeChild(el); - } + if (await copyToClipboard(JSON.stringify(value, null, 2))) { setCopied(true); setTimeout(() => setCopied(false), 2000); - } catch (e) { - console.error("Copy failed:", e); } }, [value]); diff --git a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx index a35199c5a6e..b24667c6a2e 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx @@ -1,11 +1,17 @@ import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import type { ReactNode } from "react"; -import { afterEach, describe, expect, it, vi } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { render } from "../../../../tests/test-utils"; import type { LogEntry } from "../types"; import { DrawerHeader } from "./DrawerHeader"; +import { copyToClipboard } from "@/utils/dataUtils"; + +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: vi.fn(), +})); const logEntry = (overrides: Partial): LogEntry => ({ @@ -119,23 +125,21 @@ describe("DrawerHeader back link", () => { }); }); -describe("DrawerHeader copy", () => { - afterEach(() => { - vi.restoreAllMocks(); - Reflect.deleteProperty(navigator, "clipboard"); - Reflect.deleteProperty(document, "execCommand"); - }); - - it("copies the request id when the browser has no clipboard API, as on a plain HTTP origin", async () => { +describe("DrawerHeader copying", () => { + it.each([true, false])("confirms both ids only when copying succeeds (%s)", async (success) => { const user = userEvent.setup(); Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); - Object.defineProperty(document, "execCommand", { value: () => true, configurable: true }); - const select = vi.spyOn(HTMLTextAreaElement.prototype, "select"); - renderHeader(logEntry({}), false); + vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success); + renderHeader(logEntry({ request_id: "request-copy", litellm_call_id: "call-copy" }), false); - await user.click(screen.getByRole("button", { name: "Copy Request ID" })); + const requestButton = screen.getByRole("button", { name: "Copy Request ID" }); + const callButton = screen.getByRole("button", { name: "Copy x-litellm-call-id" }); + await user.click(requestButton); + await user.click(callButton); - expect(select.mock.contexts.map((textArea) => textArea.value)).toEqual(["170d64ea-69f0-431a-be72-332f8f78c18a"]); - expect(await screen.findByRole("button", { name: "Copied!" })).toBeInTheDocument(); + expect(copyToClipboard).toHaveBeenNthCalledWith(1, "request-copy"); + expect(copyToClipboard).toHaveBeenNthCalledWith(2, "call-copy"); + expect(requestButton).toHaveAccessibleName(success ? "Copied!" : "Copy Request ID"); + expect(callButton).toHaveAccessibleName(success ? "Copied!" : "Copy x-litellm-call-id"); }); }); diff --git a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx index 4cb923df4a4..312a0797f1d 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { type ReactNode, useState } from "react"; import { ArrowLeft, Check, ChevronDown, ChevronUp, Copy, X } from "lucide-react"; import moment from "moment"; @@ -9,7 +10,6 @@ import { AutoRouterTag } from "@/components/shared/table_cells"; import { ClassifyTag } from "./ClassifyTag"; import { SidebarToggle } from "./SidebarToggle"; import { getProviderLogoAndName } from "../../provider_info_helpers"; -import { copyToClipboard } from "@/utils/dataUtils"; import { DRAWER_HEADER_PADDING, COLOR_BORDER, @@ -184,9 +184,10 @@ function CopyableId({ const [copied, setCopied] = useState(false); const handleCopy = async () => { - if (!(await copyToClipboard(value))) return; - setCopied(true); - setTimeout(() => setCopied(false), 1200); + if (await copyToClipboard(value)) { + setCopied(true); + setTimeout(() => setCopied(false), 1200); + } }; return ( diff --git a/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx index dfd9dd54dd7..87105d3f48f 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx @@ -1,12 +1,15 @@ import React from "react"; -import { render, screen, waitFor, fireEvent, act } from "@testing-library/react"; +import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { InputCard } from "./InputCard"; import { ParsedMessage } from "./prettyMessagesTypes"; +import { copyToClipboard } from "@/utils/dataUtils"; +import { toast } from "@/lib/toast"; + +vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() })); describe("InputCard", () => { - const mockWriteText = vi.fn().mockResolvedValue(undefined); const mockMessages: ParsedMessage[] = [ { role: "user", @@ -16,13 +19,7 @@ describe("InputCard", () => { beforeEach(() => { vi.clearAllMocks(); - Object.defineProperty(navigator, "clipboard", { - value: { - writeText: mockWriteText, - }, - writable: true, - configurable: true, - }); + vi.mocked(copyToClipboard).mockResolvedValue(true); }); it("should render the InputCard component", () => { @@ -91,27 +88,16 @@ describe("InputCard", () => { expect(screen.getByText(/Cost: \$0\.001500/)).toBeInTheDocument(); }); - it("should copy last message content when copy button is clicked", async () => { - const messages: ParsedMessage[] = [ - { - role: "user", - content: "Copy this text", - }, - ]; + it.each([true, false])("delegates copying and feedback to the helper (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(success); + render(); - render(); + await user.click(screen.getByRole("button", { name: /copy/i })); - const copyButton = screen.getByRole("button", { name: /copy/i }); - - expect(copyButton).toBeInTheDocument(); - - await act(async () => { - fireEvent.click(copyButton); - }); - - await waitFor(() => { - expect(mockWriteText).toHaveBeenCalledWith("Copy this text"); - }); + expect(copyToClipboard).toHaveBeenCalledWith("Copy this text", "Input copied"); + expect(toast.success).not.toHaveBeenCalled(); }); it("should toggle collapse state when header is clicked", async () => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx b/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx index 6667b6b086e..0b0bde85b91 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx @@ -30,7 +30,10 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps) const lastMessage = nonSystemMessages.length > 0 ? nonSystemMessages[nonSystemMessages.length - 1] : null; const historyMessages = nonSystemMessages.slice(0, -1); - const handleCopy = () => copyToClipboard(lastMessage?.content, "Input copied"); + const handleCopy = () => { + const content = lastMessage?.content || ""; + void copyToClipboard(content, "Input copied"); + }; return (
({ default: ({ data }: { data: unknown }) =>
{JSON.stringify(data)}
, @@ -295,6 +296,46 @@ describe("LogDetailContent", () => { expect(screen.queryByRole("tab", { name: "Request" })).not.toBeInTheDocument(); }); + it.each([ + ["Request", true], + ["Response", true], + ["Metadata", true], + ["Request", false], + ] as const)("copies %s through the fallback and reports its outcome (%s)", async (section, success) => { + const user = userEvent.setup(); + const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard"); + const execCommandDescriptor = Object.getOwnPropertyDescriptor(document, "execCommand"); + const writeText = vi.fn<(value: string) => boolean>().mockReturnValue(success); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + Object.defineProperty(document, "execCommand", { + configurable: true, + value: (command: string) => + command === "copy" && writeText((document.activeElement as HTMLTextAreaElement).value), + }); + vi.mocked(toast.success).mockClear(); + const log = createLogEntry(); + const expected = { Request: log.messages, Response: log.response, Metadata: log.metadata }[section]; + + try { + render(); + if (section !== "Metadata") { + await user.click(screen.getByRole("tab", { name: "JSON", exact: true })); + await user.click(screen.getByRole("tab", { name: section, exact: true })); + } + const label = section === "Metadata" ? "Copy Metadata" : "Copy JSON"; + const button = screen.getByRole("button", { name: label, exact: true }); + await user.click(button); + + expect(writeText).toHaveBeenCalledWith(JSON.stringify(expected, null, 2)); + expect(button).toHaveAccessibleName(success ? "Copied!" : label); + expect(toast.success).toHaveBeenCalledTimes(success ? 1 : 0); + } finally { + if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor); + if (execCommandDescriptor) Object.defineProperty(document, "execCommand", execCommandDescriptor); + else Reflect.deleteProperty(document, "execCommand"); + } + }); + it.each(["object", "serialized", "messages only", "null captures"])( "preserves request inspection and copying for classifier logs with %s data", async (shape) => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx index 97393e24fdd..f740b6c586a 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { useState } from "react"; import { Check, ChevronDown, ChevronRight, CircleAlert, Copy, Info } from "lucide-react"; import moment from "moment"; @@ -9,7 +10,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { LogEntry } from "../types"; -import { copyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; import { PROMPT_CACHE_CREATION_TOOLTIP, PROMPT_CACHE_READ_TOOLTIP } from "@/utils/promptCacheUsage"; import GuardrailViewer from "./sections/GuardrailViewer/GuardrailViewer"; import EvalViewer from "./sections/EvalViewer/EvalViewer"; @@ -304,9 +305,10 @@ function CopyButton({ const [copied, setCopied] = useState(false); const handleCopy = async () => { - if (!(await copyToClipboard(getText()))) return; - setCopied(true); - setTimeout(() => setCopied(false), 1200); + if (await copyToClipboard(getText())) { + setCopied(true); + setTimeout(() => setCopied(false), 1200); + } }; return ( diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx index 1321a4a152a..36ab70ecf2a 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx @@ -1,3 +1,5 @@ +import userEvent from "@testing-library/user-event"; +import { copyToClipboard } from "@/utils/dataUtils"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; @@ -6,6 +8,11 @@ import { sessionSpendLogsCall } from "../../networking"; import { LogEntry } from "../types"; import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells"; +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: vi.fn(), +})); + vi.mock("../../networking", () => ({ sessionSpendLogsCall: vi.fn(), })); @@ -201,3 +208,18 @@ describe("LogDetailsDrawer session sidebar auto-router icon", () => { expect(directRow.querySelector(".lucide-waypoints")).toBeNull(); }); }); + +describe("LogDetailsDrawer copying", () => { + it.each([true, false])("confirms a session id only when copying succeeds (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success); + renderSessionDrawer(); + + const button = await screen.findByRole("button", { name: "Copy trace id" }); + await user.click(button); + + expect(copyToClipboard).toHaveBeenCalledWith("session-1"); + expect(button.querySelector(".lucide-check") !== null).toBe(success); + }); +}); diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx index dafb1748dc3..23ad629d572 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { useEffect, useMemo, useState } from "react"; import { Bot, Check, Copy, Sparkles, Wrench } from "lucide-react"; import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet"; @@ -13,7 +14,7 @@ import { useKeyboardNavigation } from "./useKeyboardNavigation"; import { LogDetailContent, GuardrailJumpLink } from "./LogDetailContent"; import { sessionSpendLogsCall } from "../../networking"; import { useQuery } from "@tanstack/react-query"; -import { copyToClipboard, getSpendString } from "@/utils/dataUtils"; +import { getSpendString } from "@/utils/dataUtils"; import { normalizeGuardrailEntries, sortSessionLogs, SessionLogSortMode } from "./utils"; import { DRAWER_WIDTH } from "./constants"; import { useLogDetails } from "@/app/(dashboard)/hooks/logDetails/useLogDetails"; @@ -288,9 +289,11 @@ export function LogDetailsDrawer({ const leftPanelDisplayId = leftPanelId.length > 14 ? `${leftPanelId.slice(0, 11)}...` : leftPanelId; const handleCopyLeftPanelId = async () => { - if (!(await copyToClipboard(leftPanelId))) return; - setCopiedLeftPanelId(true); - setTimeout(() => setCopiedLeftPanelId(false), 1200); + if (!leftPanelId) return; + if (await copyToClipboard(leftPanelId)) { + setCopiedLeftPanelId(true); + setTimeout(() => setCopiedLeftPanelId(false), 1200); + } }; if (!currentLog || !enrichedLog) return null; diff --git a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx index 3dda7c5fa65..7b4c1a987a1 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx @@ -1,12 +1,15 @@ import React from "react"; -import { render, screen, waitFor, fireEvent, act } from "@testing-library/react"; +import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { OutputCard } from "./OutputCard"; import { ParsedMessage } from "./prettyMessagesTypes"; +import { copyToClipboard } from "@/utils/dataUtils"; +import { toast } from "@/lib/toast"; + +vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() })); describe("OutputCard", () => { - const mockWriteText = vi.fn().mockResolvedValue(undefined); const mockMessage: ParsedMessage = { role: "assistant", content: "This is a test response", @@ -14,13 +17,7 @@ describe("OutputCard", () => { beforeEach(() => { vi.clearAllMocks(); - Object.defineProperty(navigator, "clipboard", { - value: { - writeText: mockWriteText, - }, - writable: true, - configurable: true, - }); + vi.mocked(copyToClipboard).mockResolvedValue(true); }); it("should render the OutputCard component", () => { @@ -53,20 +50,16 @@ describe("OutputCard", () => { expect(screen.getByText(/Cost: \$0\.002500/)).toBeInTheDocument(); }); - it("should copy message content when copy button is clicked", async () => { + it.each([true, false])("delegates copying and feedback to the helper (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(success); render(); - const copyButton = screen.getByRole("button", { name: /copy/i }); + await user.click(screen.getByRole("button", { name: /copy/i })); - expect(copyButton).toBeInTheDocument(); - - await act(async () => { - fireEvent.click(copyButton); - }); - - await waitFor(() => { - expect(mockWriteText).toHaveBeenCalledWith("This is a test response"); - }); + expect(copyToClipboard).toHaveBeenCalledWith("This is a test response", "Output copied"); + expect(toast.success).not.toHaveBeenCalled(); }); it("should not copy when message is null and copy button is clicked", async () => { @@ -80,7 +73,7 @@ describe("OutputCard", () => { await user.click(copyButton); await waitFor(() => { - expect(mockWriteText).not.toHaveBeenCalled(); + expect(copyToClipboard).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx index 1ede243d183..7daa8b02546 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx @@ -19,7 +19,11 @@ interface OutputCardProps { export function OutputCard({ message, completionTokens, outputCost }: OutputCardProps) { const [isCollapsed, setIsCollapsed] = useState(false); - const handleCopy = () => copyToClipboard(message?.content, "Output copied"); + const handleCopy = () => { + if (!message) return; + + void copyToClipboard(message.content || "", "Output copied"); + }; return (
diff --git a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx index 201c9e74ce8..d5280c35983 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx @@ -1,9 +1,12 @@ import React from "react"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import { copyToClipboard } from "@/utils/dataUtils"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { RealtimePrettyView, isRealtimeResponse } from "./RealtimePrettyView"; +vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn().mockResolvedValue(true) })); + const sampleRealtimeResponse = { usage: { total_tokens: 587, @@ -142,15 +145,20 @@ describe("isRealtimeResponse", () => { }); describe("RealtimePrettyView", () => { - const mockWriteText = vi.fn().mockResolvedValue(undefined); - beforeEach(() => { vi.clearAllMocks(); - Object.defineProperty(navigator, "clipboard", { - value: { writeText: mockWriteText }, - writable: true, - configurable: true, - }); + }); + + it("copies all transcript turns when the Clipboard API is unavailable", async () => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + render(); + + await user.click(screen.getByRole("button", { name: /copy/i })); + + expect(copyToClipboard).toHaveBeenCalledWith( + "assistant: Hello! How's your day going?\nassistant: I'm here to help with information and general questions.", + ); }); it("should render the component successfully", () => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx index 178875f1d2d..04f0a178462 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; /** * RealtimePrettyView - Structured pretty view for OpenAI Realtime API logs * Displays session config, conversation turns, and token usage @@ -8,7 +9,6 @@ import { useState } from "react"; import { ChevronDown, ChevronUp, MessageSquare, Mic, Settings, Volume2 } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; -import { copyToClipboard } from "@/utils/dataUtils"; import { SectionHeader } from "./SectionHeader"; interface RealtimeEvent { @@ -277,7 +277,7 @@ function ConversationCard({ (r.output || []).flatMap((o) => (o.content || []).map((c) => `${o.role}: ${c.transcript || c.text || ""}`)), ) .join("\n"); - copyToClipboard(transcripts); + void copyToClipboard(transcripts); }; return ( diff --git a/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx b/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx index 713a1435da3..311b536d89c 100644 --- a/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx @@ -13,13 +13,9 @@ vi.mock("lucide-react", async (importOriginal) => ({ Check: (props: ComponentProps<"svg">) => , })); -const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard"); - afterEach(() => { vi.useRealTimers(); vi.resetAllMocks(); - if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor); - else Reflect.deleteProperty(navigator, "clipboard"); }); describe("CopyButton", () => { @@ -28,10 +24,25 @@ describe("CopyButton", () => { expect(screen.queryByRole("button")).not.toBeInTheDocument(); }); - it("does not show the checkmark when copying fails", async () => { + it("shows the icon confirmation only after copying succeeds", async () => { + const user = userEvent.setup(); + const pendingCopy = Promise.withResolvers(); + vi.mocked(copyToClipboard).mockReturnValue(pendingCopy.promise); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: "Copy value" })); + + expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Value copied"); + expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument(); + expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument(); + + await act(async () => pendingCopy.resolve(true)); + expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument(); + }); + + it("keeps the icon unchanged when copying fails", async () => { const user = userEvent.setup(); vi.mocked(copyToClipboard).mockResolvedValue(false); - renderWithProviders(); await user.click(screen.getByRole("button", { name: "Copy value" })); @@ -41,19 +52,6 @@ describe("CopyButton", () => { expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument(); }); - it("still copies when the clipboard API is unavailable, as on a plain HTTP origin", async () => { - const user = userEvent.setup(); - Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); - vi.mocked(copyToClipboard).mockResolvedValue(true); - - renderWithProviders(); - - await user.click(screen.getByRole("button", { name: "Copy value" })); - - expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Copied"); - expect(await screen.findByRole("img", { name: "Copied icon" })).toBeInTheDocument(); - }); - it("keeps the default icon styling, custom classes, title and 1200ms confirmation", async () => { vi.useFakeTimers(); vi.mocked(copyToClipboard).mockResolvedValue(true); diff --git a/ui/litellm-dashboard/src/components/shared/CopyButton.tsx b/ui/litellm-dashboard/src/components/shared/CopyButton.tsx index 89f99f1367b..225154e4bcf 100644 --- a/ui/litellm-dashboard/src/components/shared/CopyButton.tsx +++ b/ui/litellm-dashboard/src/components/shared/CopyButton.tsx @@ -53,7 +53,9 @@ const CopyButton: React.FC = ({ if (value == null || (!value && !isAction)) return null; - const handleCopy = async () => setCopied(await copyToClipboard(value, copiedLabel)); + const handleCopy = async () => { + setCopied(await copyToClipboard(value, copiedLabel)); + }; const button = (