scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

Merge remote-tracking branch 'upstream/main' into fix-ui-copy-button-fallback

Browse source 
# Conflicts:
#	ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx
#	ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx
#	ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx
#	ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx
#	ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx
#	ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx
#	ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx
#	ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx
#	ui/litellm-dashboard/src/components/shared/CopyButton.tsx
This commit is contained in:
Vivek Raj 2026-10-27 00:49:44 +05:30
commit 5801a6e284
39 changed files with 2379 additions and 445 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/integrations/otel/model/destination.py
View file

@ -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."
),
)

160
litellm/integrations/otel/plumbing/providers.py
View file

@ -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:

2
litellm/llms/pass_through/guardrail_translation/handler.py
View file

@ -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)

123
litellm/model_prices_and_context_window_backup.json
View file

@ -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
}
}

46
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py
View file

@ -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)

2
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -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:

7
litellm/utils.py
View file

@ -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:

123
model_prices_and_context_window.json
View file

@ -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
}
}

723
tests/integration/observability/test_otel_team_sampling_across_pods.py Normal file
View file

@ -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

153
tests/integration/observability/test_passthrough_guardrail_scan_text.py Normal file
View file

@ -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

33
tests/integration/sdk/test_sync_stream_connection_release_wire.py
View file

@ -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"
)

56
tests/unit/experimental_mcp_client/test_mcp_client.py
View file

@ -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 (

403
tests/unit/integrations/otel/test_otel_v2_destinations.py
View file

@ -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),))

118
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -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

16
tests/unit/llms/pass_through/guardrail_translation/test_handler.py
View file

@ -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",
[

86
tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
View file

@ -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)

65
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
View file

@ -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}

33
tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -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"]

5
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -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

36
ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx
View file

@ -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(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
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(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);

9
ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx
View file

@ -68,6 +68,15 @@ const SettingValueEditor: React.FC<{
/>
);
}
if (setting.field_type === "String") {
return (
<Input
aria-label={setting.field_name}
value={setting.field_value ?? ""}
onChange={(event) => onChange(setting.field_name, event.target.value)}
/>
);
}
if (setting.field_type === "Integer") {
return (
<Input

23
ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.test.tsx
View file

@ -1,10 +1,16 @@
import { copyToClipboard } from "@/utils/dataUtils";
import { describe, it, expect, vi, beforeEach } from "vitest";
import userEvent from "@testing-library/user-event";
import { render, screen, waitFor, within } from "@testing-library/react";
import { render, screen, within } from "@testing-library/react";
import moment from "moment";
import { AuditLogDrawer } from "./AuditLogDrawer";
import { AuditLogEntry } from "../AuditLogsTableColumns";
vi.mock("@/utils/dataUtils", async (importOriginal) => ({
...(await importOriginal<typeof import("@/utils/dataUtils")>()),
copyToClipboard: vi.fn(),
}));
vi.mock("../../../common_components/DefaultProxyAdminTag", () => ({
default: ({ userId }: { userId: string }) => <span>{userId}</span>,
}));
@ -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(<AuditLogDrawer {...defaultProps} />);
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", () => {

19
ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx
View file

@ -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<stri
const [copied, setCopied] = useState(false);
const handleCopy = useCallback(async () => {
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]);

34
ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx
View file

@ -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<typeof import("@/utils/dataUtils")>()),
copyToClipboard: vi.fn(),
}));
const logEntry = (overrides: Partial<LogEntry>): 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");
});
});

9
ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx
View file

@ -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 (

42
ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx
View file

@ -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(<InputCard messages={[{ role: "user", content: "Copy this text" }]} />);
render(<InputCard messages={messages} />);
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 () => {

5
ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx
View file

@ -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 (
<div

41
ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.integration.test.tsx
View file

@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event";
import { describe, expect, it, vi } from "vitest";
import { GuardrailJumpLink, LogDetailContent } from "./LogDetailContent";
import type { LogEntry } from "../types";
import { toast } from "@/lib/toast";
vi.mock("./sections/GuardrailViewer/GuardrailViewer", () => ({
default: ({ data }: { data: unknown }) => <div data-testid="guardrail-viewer">{JSON.stringify(data)}</div>,
@ -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(<LogDetailContent logEntry={log} />);
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) => {

10
ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx
View file

@ -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 (

22
ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx
View file

@ -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<typeof import("@/utils/dataUtils")>()),
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);
});
});

11
ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx
View file

@ -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;

35
ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx
View file

@ -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(<OutputCard message={mockMessage} />);
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();
});
});

6
ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx
View file

@ -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 (
<div className="overflow-hidden rounded-md" style={{ border: `1px solid ${COLOR_BORDER}` }}>

22
ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx
View file

@ -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(<RealtimePrettyView response={sampleRealtimeResponse} />);
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", () => {

4
ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx
View file

@ -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 (

36
ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx
View file

@ -13,13 +13,9 @@ vi.mock("lucide-react", async (importOriginal) => ({
Check: (props: ComponentProps<"svg">) => <svg {...props} role="img" aria-label="Copied icon" />,
}));
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<boolean>();
vi.mocked(copyToClipboard).mockReturnValue(pendingCopy.promise);
renderWithProviders(<CopyButton value="secret-value" label="Copy value" copiedLabel="Value copied" />);
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(<CopyButton value="secret-value" label="Copy value" />);
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(<CopyButton value="secret-value" label="Copy value" />);
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);

4
ui/litellm-dashboard/src/components/shared/CopyButton.tsx
View file

@ -53,7 +53,9 @@ const CopyButton: React.FC<CopyButtonProps> = ({
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 = (
<Button

277
ui/litellm-dashboard/src/utils/dataUtils.test.ts
View file

@ -1,5 +1,6 @@
// @vitest-environment jsdom
import "@testing-library/jest-dom/vitest";
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
import {
copyToClipboard,
@ -16,8 +17,6 @@ const mockToast = vi.mocked(toast);
describe("dataUtils", () => {
beforeEach(() => {
vi.clearAllMocks();
// Reset document.execCommand mock
delete (document as any).execCommand;
});
afterEach(() => {
@ -134,194 +133,114 @@ describe("dataUtils", () => {
});
describe("copyToClipboard", () => {
describe("when Clipboard API is available (HTTPS scenario)", () => {
beforeEach(() => {
// Mock modern Clipboard API
Object.assign(navigator, {
clipboard: {
writeText: vi.fn().mockResolvedValue(undefined),
},
});
});
it("should use navigator.clipboard.writeText when available", async () => {
const result = await copyToClipboard("test text");
expect(navigator.clipboard.writeText).toHaveBeenCalledWith("test text");
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
expect(result).toBe(true);
});
it("should use custom message when provided", async () => {
await copyToClipboard("test text", "Custom message");
expect(mockToast.success).toHaveBeenCalledWith("Custom message");
});
it("should return false for null/undefined text", async () => {
expect(await copyToClipboard(null)).toBe(false);
expect(await copyToClipboard(undefined)).toBe(false);
expect(navigator.clipboard.writeText).not.toHaveBeenCalled();
});
it("should fall back to execCommand when clipboard API fails", async () => {
// Make clipboard API fail
navigator.clipboard.writeText = vi.fn().mockRejectedValue(new Error("Permission denied"));
// Mock successful execCommand
document.execCommand = vi.fn().mockReturnValue(true);
// Mock DOM methods
const mockTextArea = {
value: "",
style: {} as Record<string, string>,
setAttribute: vi.fn(),
focus: vi.fn(),
select: vi.fn(),
};
document.createElement = vi.fn().mockReturnValue(mockTextArea);
document.body.appendChild = vi.fn();
document.body.removeChild = vi.fn();
const result = await copyToClipboard("test text");
expect(navigator.clipboard.writeText).toHaveBeenCalledWith("test text");
expect(document.execCommand).toHaveBeenCalledWith("copy");
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
expect(result).toBe(true);
const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard");
const execCommandDescriptor = Object.getOwnPropertyDescriptor(document, "execCommand");
const writeText = vi.fn<(text: string) => Promise<void>>();
const execCommand = vi.fn<(command: string) => boolean>();
const selectedText = vi.fn(() => {
const element = document.activeElement;
return element instanceof HTMLTextAreaElement
? element.value.slice(element.selectionStart, element.selectionEnd)
: undefined;
});
const trigger = document.createElement("button");
const text = "Request α\nsecond line";
beforeEach(() => {
writeText.mockResolvedValue(undefined);
execCommand.mockImplementation(() => {
selectedText();
return true;
});
Object.defineProperty(navigator, "clipboard", { configurable: true, value: { writeText } });
Object.defineProperty(document, "execCommand", { configurable: true, value: execCommand });
vi.spyOn(console, "error").mockImplementation(() => {});
document.body.appendChild(trigger);
trigger.focus();
});
describe("when Clipboard API is not available (HTTP scenario)", () => {
beforeEach(() => {
// Mock HTTP scenario - no clipboard API
Object.assign(navigator, {
clipboard: undefined,
});
// Mock successful execCommand
document.execCommand = vi.fn().mockReturnValue(true);
// Mock DOM methods
const mockTextArea = {
value: "",
style: {} as Record<string, string>,
setAttribute: vi.fn(),
focus: vi.fn(),
select: vi.fn(),
};
document.createElement = vi.fn().mockReturnValue(mockTextArea);
document.body.appendChild = vi.fn();
document.body.removeChild = vi.fn();
});
it("should fall back to execCommand when clipboard API is not available", async () => {
const result = await copyToClipboard("test text");
expect(document.createElement).toHaveBeenCalledWith("textarea");
expect(document.execCommand).toHaveBeenCalledWith("copy");
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
expect(result).toBe(true);
});
it("should set textarea properties correctly", async () => {
const mockTextArea = {
value: "",
style: {} as Record<string, string>,
setAttribute: vi.fn(),
focus: vi.fn(),
select: vi.fn(),
};
document.createElement = vi.fn().mockReturnValue(mockTextArea);
await copyToClipboard("test text");
expect(mockTextArea.value).toBe("test text");
/* eslint-disable jest-dom/prefer-to-have-style -- the subject is a plain mock object, not a DOM node, so toHaveStyle cannot read it */
expect(mockTextArea.style.position).toBe("fixed");
expect(mockTextArea.style.left).toBe("-999999px");
expect(mockTextArea.style.top).toBe("-999999px");
/* eslint-enable jest-dom/prefer-to-have-style */
expect(mockTextArea.setAttribute).toHaveBeenCalledWith("readonly", "");
expect(mockTextArea.focus).toHaveBeenCalled();
expect(mockTextArea.select).toHaveBeenCalled();
});
it("should handle execCommand failure", async () => {
document.execCommand = vi.fn().mockReturnValue(false);
const result = await copyToClipboard("test text");
expect(document.execCommand).toHaveBeenCalledWith("copy");
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
expect(result).toBe(false);
});
it("should handle DOM manipulation errors", async () => {
document.createElement = vi.fn().mockImplementation(() => {
throw new Error("DOM error");
});
const result = await copyToClipboard("test text");
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
expect(result).toBe(false);
});
it("should clean up textarea element after successful copy", async () => {
const mockTextArea = {
value: "",
style: {} as Record<string, string>,
setAttribute: vi.fn(),
focus: vi.fn(),
select: vi.fn(),
};
document.createElement = vi.fn().mockReturnValue(mockTextArea);
await copyToClipboard("test text");
expect(document.body.appendChild).toHaveBeenCalledWith(mockTextArea);
expect(document.body.removeChild).toHaveBeenCalledWith(mockTextArea);
});
afterEach(() => {
document.body.replaceChildren();
if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor);
else Reflect.deleteProperty(navigator, "clipboard");
if (execCommandDescriptor) Object.defineProperty(document, "execCommand", execCommandDescriptor);
else Reflect.deleteProperty(document, "execCommand");
});
describe("edge cases", () => {
beforeEach(() => {
// Mock scenario where navigator exists but clipboard is null
Object.assign(navigator, {
clipboard: null,
});
it("confirms a modern clipboard write only after it succeeds", async () => {
const pendingWrite = Promise.withResolvers<void>();
writeText.mockReturnValue(pendingWrite.promise);
const result = copyToClipboard(text, "Request copied");
document.execCommand = vi.fn().mockReturnValue(true);
const mockTextArea = {
value: "",
style: {} as Record<string, string>,
setAttribute: vi.fn(),
focus: vi.fn(),
select: vi.fn(),
};
document.createElement = vi.fn().mockReturnValue(mockTextArea);
document.body.appendChild = vi.fn();
document.body.removeChild = vi.fn();
expect(writeText).toHaveBeenCalledWith(text);
expect(mockToast.success).not.toHaveBeenCalled();
pendingWrite.resolve();
expect(await result).toBe(true);
expect(mockToast.success).toHaveBeenCalledWith("Request copied");
expect(execCommand).not.toHaveBeenCalled();
expect(trigger).toHaveFocus();
});
it.each([null, undefined, ""])("does not copy absent content %s", async (value) => {
expect(await copyToClipboard(value)).toBe(false);
expect(writeText).not.toHaveBeenCalled();
expect(execCommand).not.toHaveBeenCalled();
expect(mockToast.success).not.toHaveBeenCalled();
});
it.each([undefined, null, {}])("copies the full selection when the Clipboard API is %s", async (clipboard) => {
Object.defineProperty(navigator, "clipboard", { configurable: true, value: clipboard });
expect(await copyToClipboard(text)).toBe(true);
expect(execCommand).toHaveBeenCalledWith("copy");
expect(selectedText).toHaveReturnedWith(text);
expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard");
expect(document.querySelector("textarea")).not.toBeInTheDocument();
expect(trigger).toHaveFocus();
});
it("falls back after a rejected modern write", async () => {
writeText.mockRejectedValue(new Error("Permission denied"));
expect(await copyToClipboard(text, "Request copied")).toBe(true);
expect(writeText).toHaveBeenCalledWith(text);
expect(execCommand).toHaveBeenCalledWith("copy");
expect(selectedText).toHaveReturnedWith(text);
expect(mockToast.success).toHaveBeenCalledExactlyOnceWith("Request copied");
expect(document.querySelector("textarea")).not.toBeInTheDocument();
expect(trigger).toHaveFocus();
});
it.each(["false", "throw"])("cleans up and reports failure when the fallback returns %s", async (outcome) => {
Object.defineProperty(navigator, "clipboard", { configurable: true, value: undefined });
execCommand.mockImplementation(() => {
selectedText();
if (outcome === "throw") throw new Error("Copy denied");
return false;
});
it("should handle navigator.clipboard being null", async () => {
const result = await copyToClipboard("test text");
expect(await copyToClipboard(text)).toBe(false);
expect(document.execCommand).toHaveBeenCalledWith("copy");
expect(result).toBe(true);
expect(selectedText).toHaveReturnedWith(text);
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
expect(mockToast.success).not.toHaveBeenCalled();
expect(document.querySelector("textarea")).not.toBeInTheDocument();
expect(trigger).toHaveFocus();
});
it("reports failure when the fallback cannot create its selection", async () => {
Object.defineProperty(navigator, "clipboard", { configurable: true, value: undefined });
vi.spyOn(document, "createElement").mockImplementationOnce(() => {
throw new Error("DOM error");
});
it("should handle navigator.clipboard.writeText being undefined", async () => {
Object.assign(navigator, {
clipboard: {}, // clipboard exists but writeText doesn't
});
const result = await copyToClipboard("test text");
expect(document.execCommand).toHaveBeenCalledWith("copy");
expect(result).toBe(true);
});
expect(await copyToClipboard(text)).toBe(false);
expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard");
expect(mockToast.success).not.toHaveBeenCalled();
expect(trigger).toHaveFocus();
});
});
});

23
ui/litellm-dashboard/src/utils/dataUtils.ts
View file

@ -103,19 +103,20 @@ const fallbackCopyToClipboard = (text: string, messageText: string): boolean =>
textArea.setAttribute("readonly", "");
document.body.appendChild(textArea);
textArea.focus();
textArea.select();
try {
textArea.focus();
textArea.select();
const successful = document.execCommand("copy");
document.body.removeChild(textArea);
if (previouslyFocused instanceof HTMLElement) previouslyFocused.focus();
if (successful) {
toast.success(messageText);
return true;
} else {
throw new Error("execCommand failed");
if (document.execCommand("copy")) {
toast.success(messageText);
return true;
}
} finally {
textArea.remove();
if (previouslyFocused instanceof HTMLElement) previouslyFocused.focus();
}
toast.fromError("Failed to copy to clipboard");
return false;
} catch (err) {
toast.fromError("Failed to copy to clipboard");
console.error("Failed to copy: ", err);