Merge branch 'main' into fix/sandbox-async-for-guards

This commit is contained in:
LeizhenZhang 2026-10-26 16:24:52 -05:00
commit 11a7320381
92 changed files with 5247 additions and 1105 deletions

View file

@ -45,7 +45,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@ -132,7 +132,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:

View file

@ -33,7 +33,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
steps:
@ -76,6 +76,12 @@ jobs:
- name: Import and CLI checks
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli
- name: Routing and UUID compatibility
run: >-
$IN_NETNS .venv/bin/python -m pytest
tests/unit/router_strategy/test_complexity_router.py::TestCustomDimensions
tests/unit/test_uuid_helper.py -q
- name: Proxy startup check
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics"

View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@ -124,7 +124,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \

View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

View file

@ -1,8 +1,8 @@
# syntax=docker/dockerfile:1.7
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.26.0
@ -140,7 +140,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer

View file

@ -8202,18 +8202,18 @@ dependencies = [
[[package]]
name = "webpki-root-certs"
version = "1.0.10"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb"
checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "webpki-roots"
version = "1.0.10"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9"
checksum = "f5f280bbb2b287bc8e7de7e3852945a0ff50d1bbcb5ed42b6c7a568c546df171"
dependencies = [
"rustls-pki-types",
]

View file

@ -1,10 +1,14 @@
"""
Internal unified UUID helper.
Always uses fastuuid for performance.
Uses fastuuid where supported and the standard library on Python 3.15+.
"""
import fastuuid as _uuid
import sys
import uuid as _uuid
if sys.version_info < (3, 15):
import fastuuid as _uuid
# Expose a module-like alias so callers can use: uuid.uuid4()
uuid = _uuid

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

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:

View file

@ -17,6 +17,8 @@ from litellm.utils import supports_reasoning, supports_vision
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "xhigh", "max", "ultra")
class DeepSeekChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list:
@ -33,7 +35,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
) -> dict[str, object]:
"""
Map OpenAI params to DeepSeek params.
@ -60,6 +62,13 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
elif reasoning_effort is not None:
optional_params["thinking"] = {"type": "disabled" if reasoning_effort == "none" else "enabled"}
if (
isinstance(reasoning_effort, str)
and reasoning_effort in SUPPORTED_REASONING_EFFORTS
and optional_params["thinking"] == {"type": "enabled"}
):
return {**optional_params, "reasoning_effort": reasoning_effort}
return optional_params
def _fill_reasoning_content(self, messages: list[AllMessageValues]) -> list[AllMessageValues]:

View file

@ -56,9 +56,9 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig):
# Response format mappings from OpenAI to ElevenLabs
FORMAT_MAPPINGS = {
"mp3": "mp3_44100_128",
"pcm": "pcm_44100",
"pcm": "pcm_24000",
"opus": "opus_48000_128",
# ElevenLabs does not support WAV, AAC, or FLAC formats.
"wav": "wav_24000",
}
ELEVENLABS_QUERY_PARAMS_KEY = "__elevenlabs_query_params__"

View file

@ -11,7 +11,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields
from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_internal_message_fields
from litellm.llms.openai.chat.gpt_transformation import (
OpenAIChatCompletionStreamingHandler,
)
@ -154,7 +154,7 @@ class GroqChatConfig(OpenAILikeChatConfig):
for k, v in _message.items():
if v is not None:
new_message[k] = v
messages[idx] = strip_litellm_response_only_fields(new_message)
messages[idx] = strip_litellm_internal_message_fields(new_message)
if is_async:
return super()._transform_messages(messages=messages, model=model, is_async=True)

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)

View file

@ -7,12 +7,14 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
import asyncio
import json
import threading
from collections import OrderedDict
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
from urllib.parse import urlparse
import litellm
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@ -70,16 +72,20 @@ class _VertexCredentialsObject(Protocol):
def refresh(self, request: object) -> None: ...
_VertexCredentialCacheKey = tuple[str | None, str | None]
class VertexBase:
def __init__(self) -> None:
super().__init__()
self.access_token: str | None = None
self.refresh_token: str | None = None
self._credentials: GoogleCredentialsObject | None = None
self._credentials_project_mapping: dict[
tuple[VERTEX_CREDENTIALS_TYPES | None, str | None],
self._credentials_project_mapping: OrderedDict[
_VertexCredentialCacheKey,
tuple[_VertexCredentialsObject, str | None],
] = {}
] = OrderedDict()
self._credentials_cache_lock = threading.RLock()
self.project_id: str | None = None
self.async_handler: AsyncHTTPHandler | None = None
# Per-credential-key asyncio.Lock for single-flight async refresh.
@ -91,10 +97,10 @@ class VertexBase:
# each lock; the entry is pruned when the count reaches zero, so the
# dict stays bounded even in long-running high-cardinality deployments
# without depending on any private asyncio internals.
self._async_refresh_locks: dict[tuple, asyncio.Lock] = {}
self._async_refresh_lock_refcounts: dict[tuple, int] = {}
self._async_refresh_locks: dict[_VertexCredentialCacheKey, asyncio.Lock] = {}
self._async_refresh_lock_refcounts: dict[_VertexCredentialCacheKey, int] = {}
# Tracks in-flight background refresh tasks to avoid duplicate refreshes.
self._background_refresh_tasks: dict[tuple, asyncio.Task] = {}
self._background_refresh_tasks: dict[_VertexCredentialCacheKey, asyncio.Task] = {}
# Protects the sync get_access_token refresh path.
# Use RLock so that the reauthentication retry path (which calls
# back into get_access_token while still holding the lock) can
@ -378,7 +384,7 @@ class VertexBase:
with self._sync_refresh_lock:
credentials.refresh(Request())
def _acquire_async_refresh_lock(self, credential_cache_key: tuple) -> asyncio.Lock:
def _acquire_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey) -> asyncio.Lock:
"""Increment the refcount and return the lock for ``credential_cache_key``.
Every call must be paired with ``_release_async_refresh_lock`` once the
@ -391,7 +397,7 @@ class VertexBase:
)
return lock
def _release_async_refresh_lock(self, credential_cache_key: tuple, lock: asyncio.Lock) -> None:
def _release_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey, lock: asyncio.Lock) -> None:
"""Decrement the refcount and drop the lock entry when it reaches zero.
Must be called only after the caller has released ``lock`` (i.e. once
@ -409,7 +415,7 @@ class VertexBase:
def _try_get_cached_token(
self,
credential_cache_key: tuple,
credential_cache_key: _VertexCredentialCacheKey,
project_id: str | None,
) -> tuple[str, str] | None:
"""
@ -432,7 +438,7 @@ class VertexBase:
def _try_get_usable_cached_token(
self,
credential_cache_key: tuple,
credential_cache_key: _VertexCredentialCacheKey,
project_id: str | None,
) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None:
"""
@ -458,19 +464,37 @@ class VertexBase:
return creds.token, resolved_project, token_state, creds, cached_project_id
def _unpack_cached_credentials(
self, credential_cache_key: tuple
self, credential_cache_key: _VertexCredentialCacheKey
) -> tuple[_VertexCredentialsObject | None, str | None]:
"""
Return (credentials, project_id) from the cache, or (None, None) if
not cached. Handles both tuple and legacy cache formats.
"""
if credential_cache_key not in self._credentials_project_mapping:
return None, None
cached_entry: Final = self._credentials_project_mapping[credential_cache_key]
with self._credentials_cache_lock:
cached_entry: Final = self._credentials_project_mapping.get(credential_cache_key)
if cached_entry is None:
return None, None
self._credentials_project_mapping.move_to_end(credential_cache_key)
if isinstance(cached_entry, tuple):
return cached_entry
return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None)
def _cache_credentials(
self,
credential_cache_key: _VertexCredentialCacheKey,
credentials: _VertexCredentialsObject,
project_id: str | None,
*,
replace: bool = True,
) -> None:
with self._credentials_cache_lock:
if not replace and credential_cache_key in self._credentials_project_mapping:
return
self._credentials_project_mapping[credential_cache_key] = (credentials, project_id)
self._credentials_project_mapping.move_to_end(credential_cache_key)
while len(self._credentials_project_mapping) > max(0, DEFAULT_MAX_LRU_CACHE_SIZE):
self._credentials_project_mapping.popitem(last=False)
def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState":
"""
Return the token state using google-auth's TokenState enum.
@ -494,7 +518,7 @@ class VertexBase:
self,
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
credential_cache_key: tuple,
credential_cache_key: _VertexCredentialCacheKey,
) -> tuple[_VertexCredentialsObject, str | None]:
"""Load credentials via load_auth (in thread) and cache the result."""
try:
@ -507,17 +531,12 @@ class VertexBase:
raise
if _credentials is None:
raise ValueError("Could not resolve credentials")
self._credentials_project_mapping[credential_cache_key] = (
_credentials,
credential_project_id,
)
self._cache_credentials(credential_cache_key, _credentials, credential_project_id)
return _credentials, credential_project_id
async def _background_refresh_credentials(
self,
credentials: _VertexCredentialsObject,
credential_cache_key: tuple,
credential_project_id: str | None,
) -> None:
"""
Refresh credentials in the background without blocking the calling request.
@ -528,25 +547,13 @@ class VertexBase:
try:
verbose_logger.debug("Background proactive credential refresh")
await asyncify(self.refresh_auth)(credentials)
# Only update the cache if it still points at the credentials
# object we just refreshed. The per-key async lock is not held
# here, so a concurrent INVALID path may have already replaced
# this entry (e.g. via _handle_reauthentication_async, which
# creates a fresh credentials object). In that case our write
# would clobber the newer entry with a stale reference.
cached_creds, _ = self._unpack_cached_credentials(credential_cache_key)
if cached_creds is credentials:
self._credentials_project_mapping[credential_cache_key] = (
credentials,
credential_project_id,
)
except Exception:
verbose_logger.debug(
"Background credential refresh failed, will retry on next request",
exc_info=True,
)
async def _await_in_flight_background_refresh(self, credential_cache_key: tuple) -> None:
async def _await_in_flight_background_refresh(self, credential_cache_key: _VertexCredentialCacheKey) -> None:
"""Wait for an in-flight background refresh to finish, if any.
google-auth's ``Credentials.refresh()`` is not safe to invoke
@ -568,8 +575,7 @@ class VertexBase:
def _schedule_background_refresh(
self,
credentials: _VertexCredentialsObject,
credential_cache_key: tuple,
credential_project_id: str | None,
credential_cache_key: _VertexCredentialCacheKey,
) -> None:
"""Kick off a single background refresh for ``credential_cache_key``.
@ -581,9 +587,7 @@ class VertexBase:
if existing is not None and not existing.done():
return
self._background_refresh_tasks.pop(credential_cache_key, None)
task: Final = asyncio.create_task(
self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id)
)
task: Final = asyncio.create_task(self._background_refresh_credentials(credentials))
def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None:
if self._background_refresh_tasks.get(credential_cache_key) is _fut:
@ -796,7 +800,7 @@ class VertexBase:
self,
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
credential_cache_key: tuple,
credential_cache_key: _VertexCredentialCacheKey,
error: Exception,
) -> tuple[str, str]:
"""
@ -822,8 +826,8 @@ class VertexBase:
)
# Clear the cached credentials
if credential_cache_key in self._credentials_project_mapping:
del self._credentials_project_mapping[credential_cache_key]
with self._credentials_cache_lock:
self._credentials_project_mapping.pop(credential_cache_key, None)
# Retry once with _retry_reauth=True to prevent infinite recursion
try:
@ -846,7 +850,7 @@ class VertexBase:
self,
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
credential_cache_key: tuple,
credential_cache_key: _VertexCredentialCacheKey,
error: Exception,
) -> tuple[str, str]:
"""
@ -856,7 +860,8 @@ class VertexBase:
"Handling async reauthentication for project_id: %s. Clearing cache and retrying once.", project_id
)
self._credentials_project_mapping.pop(credential_cache_key, None)
with self._credentials_cache_lock:
self._credentials_project_mapping.pop(credential_cache_key, None)
try:
(
@ -875,10 +880,7 @@ class VertexBase:
# references the OLD credentials object we just replaced, and
# leaving it would force the next request to do a redundant
# refresh/reauth before realizing the cached creds are stale.
self._credentials_project_mapping[resolved_cache_key] = (
_credentials,
credential_project_id,
)
self._cache_credentials(resolved_cache_key, _credentials, credential_project_id)
if _credentials.token is None or not isinstance(_credentials.token, str):
raise ValueError(
@ -924,21 +926,11 @@ class VertexBase:
# Convert dict credentials to string for caching
cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials
credential_cache_key: Final = (cache_credentials, project_id)
_credentials: _VertexCredentialsObject | None = None
verbose_logger.debug("Checking cached credentials for project_id: %s", project_id)
if credential_cache_key in self._credentials_project_mapping:
_credentials, credential_project_id = self._unpack_cached_credentials(credential_cache_key)
if _credentials is not None:
verbose_logger.debug("Cached credentials found for project_id: %s.", project_id)
# Retrieve both credentials and cached project_id
cached_entry: Final = self._credentials_project_mapping[credential_cache_key]
verbose_logger.debug("cached_entry: %s", cached_entry)
if isinstance(cached_entry, tuple):
_credentials, credential_project_id = cached_entry
else:
# Backward compatibility with old cache format
_credentials = cached_entry
credential_project_id = _credentials.quota_project_id or getattr(_credentials, "project_id", None)
verbose_logger.debug(
"Using cached credentials for project_id: %s",
credential_project_id,
@ -963,10 +955,7 @@ class VertexBase:
f"Could not resolve credentials - either dynamically or from environment, for project_id: {project_id}"
)
# Cache the project_id and credentials from load_auth result (resolved project_id)
self._credentials_project_mapping[credential_cache_key] = (
_credentials,
credential_project_id,
)
self._cache_credentials(credential_cache_key, _credentials, credential_project_id)
## VALIDATE CREDENTIALS
verbose_logger.debug("Validating credentials")
@ -974,11 +963,7 @@ class VertexBase:
project_id = credential_project_id
# Update cache with resolved project_id for future lookups
resolved_cache_key: Final = (cache_credentials, project_id)
if resolved_cache_key not in self._credentials_project_mapping:
self._credentials_project_mapping[resolved_cache_key] = (
_credentials,
credential_project_id,
)
self._cache_credentials(resolved_cache_key, _credentials, credential_project_id, replace=False)
# Check if credentials are None before accessing attributes
if _credentials is None:
@ -991,10 +976,6 @@ class VertexBase:
try:
verbose_logger.debug("Credentials expired, refreshing")
self.refresh_auth(_credentials)
self._credentials_project_mapping[credential_cache_key] = (
_credentials,
credential_project_id,
)
except Exception as e:
# if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login`
# in this case, we should try to reload the credentials by clearing the cache and retrying
@ -1043,9 +1024,9 @@ class VertexBase:
# callers on the lock just to schedule that refresh.
usable: Final = self._try_get_usable_cached_token(credential_cache_key, project_id)
if usable is not None:
cached_token, resolved_project, token_state, creds, cached_project_id = usable
cached_token, resolved_project, token_state, creds, _cached_project_id = usable
if token_state == TokenState.STALE:
self._schedule_background_refresh(creds, credential_cache_key, cached_project_id)
self._schedule_background_refresh(creds, credential_cache_key)
return cached_token, resolved_project
# === SLOW PATH (per-key lock) ===
@ -1075,10 +1056,7 @@ class VertexBase:
# reauth that only repopulated the unresolved key), which
# would force the next request through an unnecessary
# refresh/reauth cycle.
self._credentials_project_mapping[resolved_cache_key] = (
_credentials,
credential_project_id,
)
self._cache_credentials(resolved_cache_key, _credentials, credential_project_id)
# Use google-auth's token_state to decide refresh strategy:
# - STALE: token is usable but within REFRESH_THRESHOLD (3:45) of
@ -1098,7 +1076,6 @@ class VertexBase:
self._schedule_background_refresh(
_credentials,
credential_cache_key,
credential_project_id,
)
return current_token, project_id
@ -1117,10 +1094,6 @@ class VertexBase:
try:
verbose_logger.debug("Credentials expired, refreshing")
await asyncify(self.refresh_auth)(_credentials)
self._credentials_project_mapping[credential_cache_key] = (
_credentials,
credential_project_id,
)
except Exception as e:
if "Reauthentication is needed" in str(e):
verbose_logger.debug("Reauthentication needed, clearing cache and retrying")

View file

@ -33698,7 +33698,7 @@
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33722,7 +33722,7 @@
"input_cost_per_token": 1e-07,
"input_cost_per_token_above_272k_tokens": 2e-07,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33746,7 +33746,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33770,7 +33770,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -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,
@ -42726,6 +42740,29 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"mistral/zai-glm-5-3-flash": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "mistral",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 5e-07,
"reasoning_effort_levels": [
"low",
"high",
"max"
],
"source": "https://docs.mistral.ai/models/zai-glm-5-3-flash",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"mistral/zai-glm-5": {
"cache_read_input_token_cost": 1.4e-07,
"input_cost_per_token": 1.4e-06,
@ -75023,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,
@ -76416,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,
@ -78629,5 +78671,227 @@
"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
},
"global.openai.gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "bedrock_converse",
"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,
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_bedrock_converse_tool_result_images": false,
"supports_bedrock_runtime_chat_completions_response_format": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": false,
"supports_none_reasoning_effort": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"us.openai.gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.375e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.75e-06,
"cache_read_input_token_cost": 5.5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1.1e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_above_272k_tokens": 2.2e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5.5e-06,
"output_cost_per_token_above_272k_tokens": 8.25e-06,
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_bedrock_converse_tool_result_images": false,
"supports_bedrock_runtime_chat_completions_response_format": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": false,
"supports_none_reasoning_effort": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"baseten/decision-models/baseten-elephant-9b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"baseten/decision-models/baseten-mammoth-27b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"github_copilot/gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot",
"supported_endpoints": [
"/v1/responses"
],
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true
}
}

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)

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:

View file

@ -1766,8 +1766,10 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase):
max_parallel_requests: int | None = Field(
default=None, description="Max concurrent requests allowed for this budget id."
)
tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.")
rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.")
tpm_limit: int | None = Field(default=None, ge=0, description="Max tokens per minute, allowed for this budget id.")
rpm_limit: int | None = Field(
default=None, ge=0, description="Max requests per minute, allowed for this budget id."
)
tpd_limit: int | None = Field(
default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id."
)
@ -1878,6 +1880,8 @@ RouterSettingsDict = Annotated[
class NewTeamRequest(TeamBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
router_settings: RouterSettingsDict | None = None
model_aliases: dict | None = None
model_max_budget: GenericBudgetConfigType | None = Field(
@ -1957,8 +1961,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
team_alias: str | None = None
organization_id: str | None = None
metadata: dict | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
tpd_limit: int | None = None
max_budget: float | None = None
soft_budget: float | None = None

View file

@ -119,6 +119,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
)
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
from litellm.proxy.native_compaction import with_proxy_compaction_executor
from litellm.proxy.route_llm_request import (
route_request,
@ -1416,6 +1417,8 @@ async def open_sse_before_first_byte(
if interval is None:
return await produce_response
# Share admission state with stream cleanup before the producer copies this context.
get_or_create_request_stash()
produce_task: Final = asyncio.ensure_future(produce_response)
await asyncio.wait((produce_task,), timeout=interval)
if produce_task.done():

View file

@ -31,6 +31,7 @@ _BINARY_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobu
_ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required})
_DECODED_JSON: Final = TypeAdapter(object)
_JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object])
def resolve_inference_model(
@ -272,7 +273,7 @@ async def read_request_body(request: Request | None) -> dict:
parsed_body = {}
else:
try:
parsed_body = orjson.loads(body)
parsed_body = dict(_JSON_REQUEST_BODY.validate_python(orjson.loads(body)))
except orjson.JSONDecodeError as e:
# The surrogate-repair fallback below runs two full-body re.sub
# passes, which block the event loop on multi-MB malformed bodies.
@ -415,7 +416,9 @@ def rewrite_request_model(
if request is None:
return
cached_body: Final = _safe_get_request_parsed_body(request=request)
body: Final = {**cached_body, "model": model} if cached_body is not None else request_data
body: Final = (
{**_JSON_REQUEST_BODY.validate_python(cached_body), "model": model} if cached_body is not None else request_data
)
safe_set_request_parsed_body(request=request, parsed_body=body)
request._json = body # pyright: ignore[reportPrivateUsage] # Starlette JSON cache
request._body = orjson.dumps(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache

View file

@ -65,6 +65,7 @@ from litellm.repositories.table_repositories import (
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.unit_of_work import (
LinkedSpendResetWrites,
SpendResetWrites,
budget_cascade_unit_of_work,
enduser_spend_reset_unit_of_work,
spend_reset_unit_of_work,
@ -235,7 +236,7 @@ def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]:
def _queue_budget_linked_resets(
writes: LinkedSpendResetWrites,
writes: SpendResetWrites,
cascade: "_BudgetCascade",
extra: Mapping[str, object] = MappingProxyType({}),
) -> None:
@ -908,6 +909,7 @@ class ResetBudgetJob:
db_span("reset_budget_cascade", "LiteLLM_BudgetTable"),
budget_cascade_unit_of_work(self._new_batch) as uow,
):
uow.team_memberships.queue_total_spend_repair(cascade.budget_ids)
_queue_budget_linked_resets(uow.team_memberships, cascade)
_queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)

View file

@ -472,6 +472,8 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS
# recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma
# types a raw array parameter from the first batch a connection sees, so after an all-$0 batch
# (integers) every later fractional batch on that connection failed with "improper binary format".
# total_spend is floored at spend before adding: rows written before total_spend existed (migration
# 20260421135425, default 0) carry their history only in spend, so a plain increment keeps them behind forever.
_TEAM_MEMBER_SPEND_SQL: Final = """
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
SELECT member.user_id, member.team_id, member.cost, member.cost
@ -487,7 +489,7 @@ WHERE EXISTS (
)
ON CONFLICT (user_id, team_id) DO UPDATE
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend
total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend) + EXCLUDED.total_spend
"""

View file

@ -1013,7 +1013,7 @@ async def get_all_team_memberships(
returned_tm: Final[list[LiteLLM_TeamMembership]] = []
for tm in team_memberships:
returned_tm.append(LiteLLM_TeamMembership.model_validate(tm.model_dump()))
returned_tm.append(with_total_spend_floor(LiteLLM_TeamMembership.model_validate(tm.model_dump())))
return returned_tm
@ -4202,6 +4202,23 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu
_raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.")
# One statement so the validated spend cannot go stale between the read and the write: the WHERE re-checks
# reset_to <= spend, and total_spend is floored at spend before spend drops so the history survives the reset.
_RESET_TEAM_MEMBER_SPEND_SQL: Final = """
UPDATE "LiteLLM_TeamMembership"
SET total_spend = GREATEST(total_spend, spend), spend = $3::float8
WHERE user_id = $1 AND team_id = $2 AND spend >= $3::float8
"""
def with_total_spend_floor(membership: LiteLLM_TeamMembership) -> LiteLLM_TeamMembership:
"""Rows written before total_spend existed (migration 20260421135425) hold their history only in spend."""
spend: Final = membership.spend or 0.0
if (membership.total_spend or 0.0) >= spend:
return membership
return membership.model_copy(update={"total_spend": spend})
def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
detail: Final = {"error": message}
raise HTTPException(status_code=status_code, detail=detail)
@ -4286,10 +4303,12 @@ async def reset_team_member_spend_fn(
current_spend: Final = membership.spend or 0.0
reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership)
await _team_membership_db(prisma_client).update(
where=membership_where,
data={"spend": reset_to},
)
reset_rows: Final = await _execute_raw(prisma_client)(_RESET_TEAM_MEMBER_SPEND_SQL, user_id, team_id, reset_to)
if reset_rows == 0:
_raise_reset_spend_error(
status.HTTP_409_CONFLICT,
f"Member spend changed since it was read; reset_to ({reset_to}) must be <= current spend. Retry.",
)
await invalidate_team_member_spend_state(
user_id=user_id,
@ -5548,16 +5567,17 @@ async def team_member_me(
litellm_budget_table=None,
)
floored: Final = with_total_spend_floor(membership)
return TeamMemberInfoResponse(
user_id=caller_user_id,
team_id=team_id,
team_alias=team_table.team_alias,
role=member_role,
user_email=user_email,
spend=membership.spend,
total_spend=membership.total_spend,
budget_id=membership.budget_id,
litellm_budget_table=membership.litellm_budget_table,
spend=floored.spend,
total_spend=floored.total_spend,
budget_id=floored.budget_id,
litellm_budget_table=floored.litellm_budget_table,
)

View file

@ -171,4 +171,6 @@ class PrismaBatch(Protocol):
@property
def litellm_projecttable(self) -> BatchTable: ...
def execute_raw(self, query: str, *args: object) -> None: ...
async def commit(self) -> None: ...

View file

@ -15,11 +15,11 @@ instead of full-model writes, which trip ``prisma.errors.DataError`` on rows
carrying fields the update input type rejects (see #27730).
"""
from collections.abc import AsyncGenerator, Callable, Mapping
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import datetime
from typing import Final
from typing import Final, Protocol, runtime_checkable
from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch
@ -62,6 +62,13 @@ class TeamSpendResetWrites:
)
@runtime_checkable
class SpendResetWrites(Protocol):
def queue_spend_zero(self, where: Mapping[str, object]) -> None: ...
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: ...
@dataclass(frozen=True, slots=True)
class LinkedSpendResetWrites:
table: BatchTable
@ -78,6 +85,30 @@ class LinkedSpendResetWrites:
)
# Membership rows written before total_spend existed (migration 20260421135425, default 0) hold their
# history only in spend. Zeroing spend would erase it, so the floor lands in the same batch, ahead of the reset.
_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL: Final = """
UPDATE "LiteLLM_TeamMembership"
SET total_spend = spend
WHERE budget_id = ANY($1::text[]) AND spend > total_spend
"""
@dataclass(frozen=True, slots=True)
class TeamMembershipSpendResetWrites:
resets: LinkedSpendResetWrites
batch: PrismaBatch
def queue_spend_zero(self, where: Mapping[str, object]) -> None:
self.resets.queue_spend_zero(where)
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None:
self.resets.queue_spend_decrement(where, amount)
def queue_total_spend_repair(self, budget_ids: Sequence[str]) -> None:
self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids))
@dataclass(frozen=True, slots=True)
class BudgetWindowWrites:
table: BatchTable
@ -106,7 +137,7 @@ class BudgetCascadeUnitOfWork:
``enduser_spend_reset_unit_of_work`` because one tier can gate millions.
"""
team_memberships: LinkedSpendResetWrites
team_memberships: TeamMembershipSpendResetWrites
keys: LinkedSpendResetWrites
organizations: LinkedSpendResetWrites
tags: LinkedSpendResetWrites
@ -132,7 +163,9 @@ async def budget_cascade_unit_of_work(
) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]:
batch = new_batch()
yield BudgetCascadeUnitOfWork(
team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership),
team_memberships=TeamMembershipSpendResetWrites(
resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch
),
keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken),
organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable),
tags=LinkedSpendResetWrites(table=batch.litellm_tagtable),

View file

@ -7,12 +7,12 @@ All values are configurable via proxy config.yaml.
import math
import re
import warnings
import sys
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from enum import Enum
from types import MappingProxyType
from typing import Annotated, Final, Literal, NamedTuple
from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple
from pydantic import (
ConfigDict,
@ -27,8 +27,13 @@ from pydantic import (
from litellm.types.llms.base import LiteLLMBaseModel
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
if TYPE_CHECKING:
import sre_constants
import sre_parse
elif sys.version_info >= (3, 11):
from re import _constants as sre_constants
from re import _parser as sre_parse
else:
import sre_constants
import sre_parse

View file

@ -55,6 +55,8 @@ class LiteLLMKeyType(str, enum.Enum):
class GenerateKeyRequest(KeyRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
soft_budget: float | None = None
send_invite_email: bool | None = None
key_type: LiteLLMKeyType | None = Field(
@ -82,6 +84,8 @@ def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok:
class UpdateKeyRequest(KeyRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
duration: str | None = None

View file

@ -16,6 +16,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
class NewUserRequest(GenerateRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
max_budget: float | None = None
user_email: str | None = None
user_alias: str | None = None
@ -47,6 +49,8 @@ class NewUserRequest(GenerateRequestBase):
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
password: str | None = Field(default=None, repr=False)
spend: float | None = None

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:

View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@ -137,7 +137,9 @@ ENV HOME=/home/nonroot \
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
PRISMA_OFFLINE_MODE=true
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/run.py /app/run.py
COPY --from=builder /opt/prisma /opt/prisma
RUN chmod -R a+rX /opt/prisma && \

View file

@ -33698,7 +33698,7 @@
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33722,7 +33722,7 @@
"input_cost_per_token": 1e-07,
"input_cost_per_token_above_272k_tokens": 2e-07,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33746,7 +33746,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33770,7 +33770,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -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,
@ -42726,6 +42740,29 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"mistral/zai-glm-5-3-flash": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "mistral",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 5e-07,
"reasoning_effort_levels": [
"low",
"high",
"max"
],
"source": "https://docs.mistral.ai/models/zai-glm-5-3-flash",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"mistral/zai-glm-5": {
"cache_read_input_token_cost": 1.4e-07,
"input_cost_per_token": 1.4e-06,
@ -75023,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,
@ -76416,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,
@ -78629,5 +78671,227 @@
"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
},
"global.openai.gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "bedrock_converse",
"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,
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_bedrock_converse_tool_result_images": false,
"supports_bedrock_runtime_chat_completions_response_format": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": false,
"supports_none_reasoning_effort": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"us.openai.gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.375e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.75e-06,
"cache_read_input_token_cost": 5.5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1.1e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_above_272k_tokens": 2.2e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5.5e-06,
"output_cost_per_token_above_272k_tokens": 8.25e-06,
"source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_bedrock_converse_tool_result_images": false,
"supports_bedrock_runtime_chat_completions_response_format": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": false,
"supports_none_reasoning_effort": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"baseten/decision-models/baseten-elephant-9b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"baseten/decision-models/baseten-mammoth-27b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"github_copilot/gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot",
"supported_endpoints": [
"/v1/responses"
],
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true
}
}

View file

@ -3,14 +3,14 @@ name = "litellm-core"
version = "1.105.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
requires-python = ">=3.10, <3.16"
license = "MIT"
license-files = ["LICENSE"]
authors = [
{ name = "BerriAI" },
]
dependencies = [
"fastuuid>=0.14.0,<1.0",
"fastuuid>=0.14.0,<1.0; python_version < '3.15'",
"filelock>=3.16.1,<4.0",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<4.0.0",
@ -20,14 +20,16 @@ dependencies = [
"packaging>=24.0",
"importlib-metadata>=8.0.0,<9.0",
"tiktoken>=0.8.0,<1.0; python_version < '3.14'",
"tiktoken>=0.12.0,<1.0; python_version >= '3.14'",
"tiktoken>=0.12.0,<1.0; python_version == '3.14'",
"tiktoken>=0.14.0,<1.0; python_version >= '3.15'",
"click>=8.0.0,<9.0",
"jinja2>=3.1.6,<4.0",
"aiohttp>=3.14.2,<4.0",
"anyio>=4.14.0,<5.0",
"async-timeout>=4.0.3,<6.0; python_version < '3.11'",
"pydantic>=2.11.0,<3.0.0; python_version < '3.14'",
"pydantic>=2.12.0,<3.0.0; python_version >= '3.14'",
"pydantic>=2.12.0,<3.0.0; python_version == '3.14'",
"pydantic>=2.14.0,<3.0.0; python_version >= '3.15'",
"pydantic-settings>=2.14.1,<3.0",
"jsonschema>=4.0.0,<5.0",
"certifi>=2026.10.19",

View file

@ -3,7 +3,7 @@ name = "litellm"
version = "1.108.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
requires-python = ">=3.10, <3.16"
license = "MIT"
license-files = ["LICENSE"]
authors = [
@ -14,7 +14,7 @@ dependencies = [
# deps. Reproducibility for our Docker/CI comes from `uv.lock`.
# When changing a floor, verify it installs + imports on every supported
# Python with: `uv pip install --resolution=lowest-direct .`
"fastuuid>=0.14.0,<1.0",
"fastuuid>=0.14.0,<1.0; python_version < '3.15'",
"filelock>=3.16.1,<4.0",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<4.0.0",
@ -23,7 +23,8 @@ dependencies = [
"packaging>=24.0",
"importlib-metadata>=8.0.0,<9.0",
"tiktoken>=0.8.0,<1.0; python_version < '3.14'",
"tiktoken>=0.12.0,<1.0; python_version >= '3.14'",
"tiktoken>=0.12.0,<1.0; python_version == '3.14'",
"tiktoken>=0.14.0,<1.0; python_version >= '3.15'",
"tokenizers>=0.21.0,<1.0",
"huggingface-hub>=0.34.0,<2.0",
"click>=8.0.0,<9.0",
@ -32,7 +33,8 @@ dependencies = [
"anyio>=4.14.0,<5.0",
"async-timeout>=4.0.3,<6.0; python_version < '3.11'",
"pydantic>=2.11.0,<3.0.0; python_version < '3.14'",
"pydantic>=2.12.0,<3.0.0; python_version >= '3.14'",
"pydantic>=2.12.0,<3.0.0; python_version == '3.14'",
"pydantic>=2.14.0,<3.0.0; python_version >= '3.15'",
"pydantic-settings>=2.14.1,<3.0",
"jsonschema>=4.0.0,<5.0",
"boto3>=1.43.117,<2.0",
@ -55,17 +57,20 @@ Documentation = "https://docs.litellm.ai"
proxy = [
"gunicorn>=23.0.0,<24.0",
"uvicorn>=0.33.0,<1.0",
"granian>=2.7.4,<3.0",
"granian>=2.7.4,<3.0; python_version < '3.15'",
"granian>=2.8.0,<3.0; python_version >= '3.15'",
"uvloop>=0.22.1,<1.0; sys_platform != 'win32'",
"fastapi>=0.136.3,<1.0",
"starlette>=1.0.1,<2.0",
"backoff>=2.2.1,<3.0",
"pyyaml>=6.0.3,<7.0",
"rq>=2.7.0,<3.0",
"orjson>=3.11.6,<4.0",
"orjson>=3.11.6,<4.0; python_version < '3.15'",
"orjson>=3.12.0,<4.0; python_version >= '3.15'",
# redis-py's C response parser. It arrives with redis (via rq) either way; naming
# it here is what makes redis-py select _HiredisParser instead of the Python one.
"hiredis>=3.0.0,<4.0",
"hiredis>=3.0.0,<4.0; python_version < '3.15'",
"hiredis>=3.5.0,<4.0; python_version >= '3.15'",
"apscheduler>=3.11.2,<4.0",
"fastapi-sso>=0.19.0,<1.0",
"PyJWT>=2.13.0,<3.0",
@ -210,8 +215,8 @@ dev = [
"tomli==2.4.1; python_version < '3.11'",
"pytest-mock==3.15.1",
"pytest-asyncio==1.3.0",
"psycopg==3.3.3",
"psycopg-binary==3.3.3",
"psycopg==3.3.6",
"psycopg-binary==3.3.6",
"pytest-xdist==3.8.0",
"requests-mock==1.12.1",
"responses==0.26.0",
@ -374,11 +379,6 @@ exclude-newer = "3 days"
[tool.uv.exclude-newer-package]
# TODO: ryan, remove after 2026-10-23, once certifi 2026.10.19 (drops the compromised eMudhra root) clears the 3-day window
certifi = "2026-10-20T03:00:00Z"
# TODO(shrey): remove after 2026-10-25, once grpcio 1.84.1 (drops the compromised eMudhra roots from roots.pem) clears the 3-day window
grpcio = "2026-10-22T15:00:00Z"
# TODO(shrey): remove after 2026-10-24, once boto3/botocore 1.43.117 (drops the compromised eMudhra roots from cacert.pem) clear the 3-day window
boto3 = "2026-10-21T19:31:00Z"
botocore = "2026-10-21T19:31:00Z"
# TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window
brotli = "2026-10-22T14:01:00Z"
# TODO: ryan, remove after 2026-10-24, once mlflow 3.17.2 (GHSA-wqmh-qccw-pqwf) clears the 3-day window

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

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

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

View file

@ -0,0 +1,273 @@
from __future__ import annotations
import asyncio
import json
import os
import uuid
from concurrent.futures import ThreadPoolExecutor
from typing import Final
import httpx
import pytest
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
from litellm.proxy.utils import PrismaClient, ProxyLogging
from tests.integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
from tests.integration._support.database import read_rows, write_rows
from tests.integration._support.provider import PROVIDER_URL, SharedProvider
from tests.integration._support.wire import Reply
_CALL_COST: Final = 0.02
JSON_VALUE: Final = TypeAdapter(JsonValue)
_LEGACY_SPEND: Final = 612.4
_LEGACY_TOTAL: Final = 148.1
_HEALTHY_SPEND: Final = 12.25
_HEALTHY_TOTAL: Final = 97.8
def _completion() -> Reply:
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
).encode()
)
def _priced_model(scenario: Scenario) -> str:
return scenario.model(api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
def _member_with_budget(gateway: Gateway, scenario: Scenario, team_id: str, *, max_budget: float) -> str:
user_id: Final = scenario.user(user_role="internal_user")
gateway.post(
"/team/member_add",
{"team_id": team_id, "member": {"role": "user", "user_id": user_id}, "max_budget_in_team": max_budget},
)
return user_id
def _seed_membership(user_id: str, team_id: str, *, spend: float, total_spend: float) -> None:
write_rows(
'UPDATE "LiteLLM_TeamMembership" SET spend = %s::float8, total_spend = %s::float8 '
"WHERE user_id = %s AND team_id = %s",
(str(spend), str(total_spend), user_id, team_id),
)
def _membership(user_id: str, team_id: str) -> tuple[float, float]:
rows: Final = read_rows(
'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s',
(user_id, team_id),
)
assert len(rows) == 1, rows
return float(str(rows[0]["spend"])), float(str(rows[0]["total_spend"]))
def _membership_reaches(user_id: str, team_id: str, *, total_spend: float) -> tuple[float, float]:
return eventually(
lambda: _membership(user_id, team_id),
lambda row: row[1] == pytest.approx(total_spend),
seconds=70,
)
def _served(gateway: Gateway, provider: SharedProvider, model: str, key: str) -> None:
provider.expect(_completion())
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"reconcile {uuid.uuid4().hex}"}]},
key=key,
)
assert response.status_code == 200, response.text
assert len(provider.received()) == 1
def _reset_spend(gateway: Gateway, team_id: str, user_id: str, reset_to: float) -> httpx.Response:
return gateway.request("POST", f"/team/{team_id}/member/{user_id}/reset_spend", {"reset_to": reset_to})
def _membership_figures(row: JsonValue) -> tuple[str, tuple[float, float]]:
membership: Final = object_value(row)
return string_value(membership["user_id"]), (
float(str(membership["spend"])),
float(str(membership["total_spend"])),
)
def _reported_memberships(gateway: Gateway, path: str, params: dict[str, str]) -> dict[str, tuple[float, float]]:
"""/team/info answers with one team object, /team/list with a list of them."""
response: Final = gateway.request("GET", path, params=params)
assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}"
body: Final = JSON_VALUE.validate_json(response.content)
teams: Final = [object_value(body)] if isinstance(body, dict) else [object_value(team) for team in list_value(body)]
rows: Final = [
row for team in teams for row in list_value(team["team_memberships"])
] # comprehension-ok: flatten memberships across teams
return dict(_membership_figures(row) for row in rows)
def test_a_flush_lifts_a_legacy_total_to_the_cycle_spend_and_leaves_a_healthy_row_on_its_baseline(
gateway: Gateway, provider: SharedProvider
) -> None:
with gateway.scenario() as scenario:
model: Final = _priced_model(scenario)
team_id: Final = scenario.team(models=[model])
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
_seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL)
legacy_key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
healthy_key: Final = scenario.key(user_id=healthy, team_id=team_id, models=[model])
_served(gateway, provider, model, legacy_key)
_served(gateway, provider, model, healthy_key)
legacy_row: Final = _membership_reaches(
legacy, team_id, total_spend=max(_LEGACY_SPEND, _LEGACY_TOTAL) + _CALL_COST
)
assert legacy_row[0] == pytest.approx(_LEGACY_SPEND + _CALL_COST)
healthy_row: Final = _membership_reaches(healthy, team_id, total_spend=_HEALTHY_TOTAL + _CALL_COST)
assert healthy_row[0] == pytest.approx(_HEALTHY_SPEND + _CALL_COST)
def test_team_info_and_team_list_report_a_legacy_total_no_lower_than_the_cycle_spend(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
team_id: Final = scenario.team()
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
_seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL)
info: Final = _reported_memberships(gateway, "/team/info", {"team_id": team_id})
listed: Final = _reported_memberships(gateway, "/team/list", {"team_id": team_id})
for reported in (info, listed):
assert reported[legacy] == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_SPEND))
assert reported[healthy] == (pytest.approx(_HEALTHY_SPEND), pytest.approx(_HEALTHY_TOTAL))
assert _membership(legacy, team_id) == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_TOTAL))
@pytest.mark.parametrize("reset_to", [0.0, 100.0], ids=["to_zero", "to_nonzero"])
def test_an_admin_reset_keeps_the_legacy_cycle_spend_as_the_lifetime_floor(gateway: Gateway, reset_to: float) -> None:
with gateway.scenario() as scenario:
team_id: Final = scenario.team()
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
response: Final = _reset_spend(gateway, team_id, legacy, reset_to)
assert response.status_code == 200, response.text
assert _membership(legacy, team_id) == (pytest.approx(reset_to), pytest.approx(_LEGACY_SPEND))
def test_concurrent_admin_resets_leave_one_winner_and_a_lifetime_floor_at_the_old_cycle_spend(
gateway: Gateway,
) -> None:
"""Two admins reset the same member at once. Each UPDATE re-checks reset_to <= spend, so
the higher reset either lands first or is refused with 409 once the lower one has landed.
Whatever the order, the row ends at the lower value with the old cycle spend preserved."""
with gateway.scenario() as scenario:
team_id: Final = scenario.team()
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
with ThreadPoolExecutor(max_workers=2) as pool:
higher, lower = pool.map(lambda value: _reset_spend(gateway, team_id, legacy, value), (600.0, 500.0))
assert lower.status_code == 200, lower.text
assert higher.status_code in (200, 409), higher.text
assert _membership(legacy, team_id) == (pytest.approx(500.0), pytest.approx(_LEGACY_SPEND))
def test_an_admin_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost(
gateway: Gateway, provider: SharedProvider
) -> None:
with gateway.scenario() as scenario:
model: Final = _priced_model(scenario)
team_id: Final = scenario.team(models=[model])
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
_served(gateway, provider, model, key)
reset: Final = _reset_spend(gateway, team_id, legacy, 0.0)
assert reset.status_code == 200, reset.text
spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST)
assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total)
assert total >= spend
def _expire_member_budget(user_id: str, team_id: str, *, budget_duration: str) -> None:
write_rows(
"UPDATE \"LiteLLM_BudgetTable\" SET budget_duration = %s, budget_reset_at = NOW() - INTERVAL '1 hour' "
'WHERE budget_id = (SELECT budget_id FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s)',
(budget_duration, user_id, team_id),
)
async def _run_scheduled_reset() -> None:
client: Final = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache())
)
await client.connect()
try:
await ResetBudgetJob(
ProxyLogging(user_api_key_cache=DualCache()), client
).reset_budget_for_litellm_budget_table()
finally:
await client.db.disconnect()
@pytest.mark.parametrize(
"spend, max_budget, carried",
[(8.0, 10.0, 0.0), (10.0, 10.0, 0.0), (15.0, 10.0, 5.0)],
ids=["below_cap", "at_cap", "above_cap_carries_overage"],
)
def test_a_scheduled_reset_lifts_the_legacy_total_before_zeroing_or_carrying_the_cycle_spend(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch, spend: float, max_budget: float, carried: float
) -> None:
monkeypatch.setattr(litellm, "budget_rollover", True)
with gateway.scenario() as scenario:
team_id: Final = scenario.team()
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=max_budget)
_seed_membership(legacy, team_id, spend=spend, total_spend=3.0)
_expire_member_budget(legacy, team_id, budget_duration="1d")
asyncio.run(_run_scheduled_reset())
assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend))
asyncio.run(_run_scheduled_reset())
assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend))
def test_a_scheduled_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost(
gateway: Gateway, provider: SharedProvider
) -> None:
with gateway.scenario() as scenario:
model: Final = _priced_model(scenario)
team_id: Final = scenario.team(models=[model])
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
_expire_member_budget(legacy, team_id, budget_duration="1d")
key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
_served(gateway, provider, model, key)
asyncio.run(_run_scheduled_reset())
spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST)
assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total)
assert total >= spend

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 (

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

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

View file

@ -511,6 +511,54 @@ class TestDeepSeekVisionMultimodalContent:
assert any(block.get("type") == "image_url" for block in content)
@pytest.mark.usefixtures("_deepseek_httpx_transport")
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize("model", ("deepseek-flash", "deepseek-v4-pro", "deepseek-chat", "deepseek-reasoner"))
@pytest.mark.parametrize("effort", ("minimal", "low", "medium", "high", "xhigh", "max", "ultra"))
@pytest.mark.parametrize("thinking", (None, "enabled"))
@pytest.mark.respx(assert_all_called=True)
async def test_deepseek_forwards_native_effort(
respx_mock: respx.MockRouter, model: str, effort: str, thinking: str | None, stream: bool
) -> None:
route: Final = respx_mock.post(DEEPSEEK_CHAT_COMPLETIONS_URL).mock(return_value=_deepseek_response(stream))
response: Final = await litellm.acompletion(
model=f"deepseek/{model}",
messages=[{"role": "user", "content": "Hello"}],
api_key=DEEPSEEK_API_KEY,
reasoning_effort=effort,
stream=stream,
**({"thinking": {"type": thinking}} if thinking is not None else {}),
)
if stream:
async for _ in response:
pass
body: Final = json.loads(route.calls[0].request.content)
assert body["reasoning_effort"] == effort
assert body["thinking"] == {"type": "enabled"}
@pytest.mark.parametrize(
"thinking,effort,expected_thinking",
(
(None, None, None),
(None, "none", "disabled"),
("enabled", "none", "enabled"),
("disabled", "max", "disabled"),
("enabled", None, "enabled"),
(None, "default", "enabled"),
(None, "bogus", "enabled"),
(None, {"effort": "high"}, "enabled"),
),
)
def test_deepseek_preserves_toggle_without_forwarding_unsupported_effort(
thinking: str | None, effort: str | dict[str, str] | None, expected_thinking: str | None
) -> None:
params: Final = {"reasoning_effort": effort, **({"thinking": {"type": thinking}} if thinking is not None else {})}
result: Final = DeepSeekChatConfig().map_openai_params(params, {}, "deepseek-flash", False)
assert "reasoning_effort" not in result
assert result.get("thinking") == ({"type": expected_thinking} if expected_thinking else None)
class TestDeepSeekThinkingParams:
"""Test thinking and reasoning_effort parameter handling for DeepSeek."""
@ -735,9 +783,7 @@ async def test_deepseek_provider_async_completion(stream):
messages = [{"role": "user", "content": "Hello, world!"}]
# Mock AsyncHTTPHandler.post method for async test
with patch(
"litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post"
) as mock_post:
with patch("litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post") as mock_post:
mock_response_data = litellm.ModelResponse(
choices=[
litellm.Choices(
@ -775,9 +821,7 @@ async def test_deepseek_provider_async_completion(stream):
# Check request body
request_body = json.loads(call_args.kwargs["data"])
assert call_args.kwargs["url"] == "https://api.deepseek.com/beta/chat/completions"
assert (
request_body["model"] == "deepseek-reasoner"
) # Model name should be stripped of provider prefix
assert request_body["model"] == "deepseek-reasoner" # Model name should be stripped of provider prefix
assert request_body["messages"] == messages
assert request_body["stream"] == stream

View file

@ -314,4 +314,4 @@ class TestElevenLabsTextToSpeechTransformation:
)
assert voice_id in url
assert "output_format=pcm_44100" in url
assert "output_format=pcm_24000" in url

View file

@ -55,3 +55,31 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r
sent: Final = json.loads(respx_mock.calls[0].request.content)
assert "_litellm_undeclared_sentinel" not in sent, sent
assert sent["text"] == "hi"
@pytest.mark.parametrize(
("response_format", "expected_output_format"),
[("wav", "wav_24000"), ("pcm", "pcm_24000")],
)
def test_speech_maps_wav_and_pcm_to_elevenlabs_output_formats(
respx_mock: respx.MockRouter,
response_format: str,
expected_output_format: str,
) -> None:
api_base: Final = "http://localhost:12346"
mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock(
return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"})
)
litellm.speech(
model="elevenlabs/eleven_multilingual_v2",
input="hi",
voice="21m00Tcm4TlvDq8ikWAM",
response_format=response_format,
api_base=api_base,
api_key="fake_elevenlabs_api_key",
)
assert mock_route.called
request_url: Final = httpx.URL(str(respx_mock.calls[0].request.url))
assert request_url.params["output_format"] == expected_output_format

View file

@ -1,4 +1,5 @@
import logging
from typing import Final
from unittest.mock import patch
import httpx
@ -523,3 +524,68 @@ def test_tool_search_reply_sent_to_groq_drops_litellm_response_fields():
}
],
}
@pytest.mark.asyncio
@pytest.mark.parametrize("is_async", [False, True])
@pytest.mark.parametrize("with_tool_call", [False, True])
async def test_streamed_reasoning_reply_can_be_replayed(is_async: bool, with_tool_call: bool) -> None:
tool_calls: Final = [
{
"index": 0,
"id": "call_weather",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
}
]
handler: Final = GroqChatCompletionStreamingHandler(streaming_response=iter(()), sync_stream=True)
chunk: Final = handler.chunk_parser(
{
"id": "chatcmpl-reasoning",
"object": "chat.completion.chunk",
"created": 1769511767,
"model": "openai/gpt-oss-20b",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"reasoning": "I should check the weather.",
"content": "" if with_tool_call else "Hello there!",
**({"tool_calls": tool_calls} if with_tool_call else {}),
},
"finish_reason": "tool_calls" if with_tool_call else "stop",
}
],
}
)
reply: Final = litellm.stream_chunk_builder([chunk]).choices[0].message
assert reply.reasoning_content == "I should check the weather."
history: Final = [reply.model_dump(exclude_none=True)]
config: Final = GroqChatConfig()
if is_async:
replayed = await config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=True)
else:
replayed = config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=False)
assert replayed == [
{
"role": "assistant",
"content": "" if with_tool_call else "Hello there!",
**(
{
"tool_calls": [
{
"id": "call_weather",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
}
]
}
if with_tool_call
else {}
),
}
]
assert reply.reasoning_content == "I should check the weather."

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

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)

View file

@ -6,13 +6,17 @@ import os
import pathlib
import re
import traceback
from typing import Final
from datetime import datetime, timedelta
from threading import Event
from typing import Final, Literal
from unittest.mock import MagicMock, call, patch
import google.auth.exceptions
import pytest
import requests
import responses
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MAX_RECURSE_DEPTH
import litellm
@ -2589,3 +2593,120 @@ def test_get_access_token_reraises_malformed_service_account_credentials_without
assert [record.getMessage().split(". Error: ")[0] for record in caplog.records] == [
"Failed to load vertex credentials. Check to see if credentials containing partial/invalid information"
]
def _authorized_user_cache_credentials(identity: str) -> str:
return json.dumps(
{
"type": "authorized_user",
"client_id": identity,
"client_secret": "synthetic-secret",
"refresh_token": identity,
"quota_project_id": identity,
}
)
@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("inferred_project", [False, True], ids=["explicit", "inferred"])
@pytest.mark.asyncio
async def test_credential_cache_bounds_retention_and_reloads_least_recently_used(
is_async: bool, inferred_project: bool
) -> None:
vertex: Final = VertexBase()
limit: Final = max(0, DEFAULT_MAX_LRU_CACHE_SIZE)
with (
patch("google.auth._helpers.utcnow", return_value=datetime(2025, 1, 1)),
responses.RequestsMock() as oauth,
):
def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600})
oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
async def access(index: int, infer: bool = False) -> tuple[str, str]:
identity: Final = f"project-{index}"
credentials: Final = _authorized_user_cache_credentials(identity)
project: Final = None if infer else identity
result: Final = (
await vertex.get_access_token_async(credentials, project)
if is_async
else vertex.get_access_token(credentials, project)
)
assert result[1] == identity
assert len(vertex._credentials_project_mapping) <= limit
return result
if limit == 0:
assert await access(0, inferred_project) == ("token-1", "project-0")
assert await access(0, inferred_project) == ("token-2", "project-0")
return
for index in range(limit):
assert await access(index) == (f"token-{index + 1}", f"project-{index}")
assert await access(0) == ("token-1", "project-0")
assert len(oauth.calls) == limit
assert await access(limit, inferred_project) == (f"token-{limit + 1}", f"project-{limit}")
assert len(oauth.calls) == limit + 1
if limit > (2 if inferred_project else 1):
assert await access(0) == ("token-1", "project-0")
assert len(oauth.calls) == limit + 1
evicted: Final = 1 if limit > 1 else 0
assert await access(evicted) == (f"token-{limit + 2}", f"project-{evicted}")
assert len(oauth.calls) == limit + 2
@pytest.mark.parametrize("mode", ["sync", "async", "background"])
@pytest.mark.skipif(DEFAULT_MAX_LRU_CACHE_SIZE < 2, reason="Eviction during refresh requires two retained entries")
@pytest.mark.asyncio
async def test_evicted_credentials_finish_refresh_without_being_readmitted(
mode: Literal["sync", "async", "background"],
) -> None:
vertex: Final = VertexBase()
limit: Final = DEFAULT_MAX_LRU_CACHE_SIZE
now: Final = datetime(2025, 1, 1)
refresh_started: Final = Event()
release_refresh: Final = Event()
donor: Final = _authorized_user_cache_credentials("donor")
target: Final = _authorized_user_cache_credentials("target")
with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth:
def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
body: Final = request.body.decode() if isinstance(request.body, bytes) else request.body
assert isinstance(body, str)
is_target: Final = "refresh_token=target" in body
if is_target and len(oauth.calls) == limit:
refresh_started.set()
assert release_refresh.wait(5), "OAuth refresh was not released"
return 200, {}, json.dumps(
{"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600 if is_target else 86400}
)
oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
assert vertex.get_access_token(donor, None) == ("token-1", "donor")
assert vertex.get_access_token(donor, None) == ("token-1", "donor")
for index in range(limit - 2):
assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}")
assert vertex.get_access_token(target, "target") == (f"token-{limit}", "target")
clock.return_value = now + timedelta(seconds=3550 if mode == "background" else 7200)
refresh: Final = asyncio.create_task(
asyncio.to_thread(vertex.get_access_token, target, "target")
if mode == "sync"
else vertex.get_access_token_async(target, "target")
)
try:
assert await asyncio.to_thread(refresh_started.wait, 5), "OAuth refresh did not start"
for index in range(limit - 2):
assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}")
assert vertex.get_access_token(donor, None) == ("token-1", "donor")
assert len(vertex._credentials_project_mapping) == limit
assert (target, "target") not in vertex._credentials_project_mapping
finally:
release_refresh.set()
completed: Final = await refresh
await asyncio.gather(*tuple(vertex._background_refresh_tasks.values()))
expected_token: Final = f"token-{limit if mode == 'background' else limit + 1}"
assert completed == (expected_token, "target")
assert len(oauth.calls) == limit + 1
assert (target, "target") not in vertex._credentials_project_mapping
assert vertex.get_access_token(target, "target") == (f"token-{limit + 2}", "target")
assert len(vertex._credentials_project_mapping) == limit

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}

View file

@ -113,6 +113,9 @@ class MockBatcher:
self.litellm_projecttable = _Table("project", self)
self.litellm_endusertable = _Table("enduser", self)
def execute_raw(self, query: str, *args: object) -> None:
self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args})
async def commit(self):
self.committed = True
return self.calls
@ -2112,6 +2115,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo
assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")}
assert batcher.committed is True
assert {(call["table"], call["op"]) for call in batcher.calls} == {
("raw", "execute_raw"),
("team_membership", "update_many"),
("key", "update_many"),
("org", "update_many"),
@ -3393,6 +3397,35 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled(
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1")
def test_budget_cascade_floors_membership_total_spend_before_zeroing_spend(
rollover_enabled: None,
reset_budget_job: ResetBudgetJob,
mock_prisma_client: MockPrismaClient,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Membership rows written before total_spend existed hold their history only in
spend. The cascade lifts total_spend to spend for the expiring tier's rows, in the
same batch and ahead of the zero/decrement, so the reset cannot erase it."""
_make_counter_invalidation_job(monkeypatch)
budget: Final = _budget_row(budget_id="budget-legacy", budget_duration="7d", max_budget=10.0)
mock_prisma_client.data["budget"] = [budget]
membership: Final = type(
"Membership",
(),
{"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-legacy"},
)
mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership])
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
writes: Final = mock_prisma_client.db.batch_calls
repair_index: Final = next(i for i, call in enumerate(writes) if call["op"] == "execute_raw")
first_membership_reset: Final = next(i for i, call in enumerate(writes) if call["table"] == "team_membership")
assert repair_index < first_membership_reset, writes
assert writes[repair_index]["args"] == (["budget-legacy"],)
assert mock_prisma_client.db.batchers[-1].committed
def test_budget_cascade_carries_enduser_overage_when_rollover_enabled(
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
):

View file

@ -1128,7 +1128,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster
assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement
assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement
assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement
assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement
assert (
'total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend)'
" + EXCLUDED.total_spend"
) in statement
@pytest.mark.asyncio

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

View file

@ -7,7 +7,7 @@ import contextvars
import logging
import sys
import time
from collections.abc import Iterator, Sequence
from collections.abc import AsyncGenerator, Iterator, Sequence
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Final, List, Optional
@ -4644,6 +4644,140 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch):
)
_SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1)
@pytest.mark.asyncio
@pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"])
@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"])
@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"])
async def test_sse_keepalive_releases_parallel_capacity(
after_ping: bool, fails: bool, preadmitted: bool
) -> None:
from fastapi.responses import StreamingResponse
from litellm.proxy.common_request_processing import open_sse_before_first_byte
from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
api_key: Final = hash_token("sk-keepalive-release")
auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
async def request(call_id: str) -> None:
data: Final = {"model": "test-model", "litellm_call_id": call_id}
upstream_answers: Final = asyncio.Event()
if not after_ping:
upstream_answers.set()
async def admit() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion"
)
async def failure(exc: Exception) -> None:
await handler.async_post_call_failure_hook(
request_data=data, original_exception=exc, user_api_key_dict=auth
)
async def stream() -> AsyncGenerator[bytes, None]:
yield b"data: [DONE]\n\n"
await handler.async_log_success_event(
kwargs={
"litellm_call_id": call_id,
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)),
start_time=_SSE_KEEPALIVE_CALLBACK_TIME,
end_time=_SSE_KEEPALIVE_CALLBACK_TIME,
)
async def produce() -> StreamingResponse:
if not preadmitted:
await admit()
await upstream_answers.wait()
if fails:
raise HTTPException(status_code=502, detail="upstream failed")
return StreamingResponse(stream())
if preadmitted:
await admit()
response_call: Final = open_sse_before_first_byte(
produce(), 0.001 if after_ping else 60.0, on_late_failure=failure
)
if fails and not after_ping:
with pytest.raises(HTTPException) as rejected:
await response_call
assert rejected.value.status_code == 502
await failure(rejected.value)
return
response: Final = await response_call
first: Final = await anext(response.body_iterator)
if after_ping:
upstream_answers.set()
rest: Final = tuple([chunk async for chunk in response.body_iterator])
assert first == (SSE_COMMENT_PING_BYTES if after_ping else b"data: [DONE]\n\n")
if fails:
error_frame: Final = rest[-2]
assert isinstance(error_frame, bytes) and b"upstream failed" in error_frame
assert rest[-1] == b"data: [DONE]\n\n"
elif after_ping:
assert rest[-1] == b"data: [DONE]\n\n"
else:
assert rest == ()
for index in range(3):
await contextvars.Context().run(asyncio.create_task, request(f"keepalive-{index}"))
assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0
@pytest.mark.asyncio
async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None:
from fastapi.responses import StreamingResponse
from litellm.proxy.common_request_processing import open_sse_before_first_byte
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
api_key: Final = hash_token("sk-keepalive-overlap")
auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
admitted: Final = asyncio.Event()
finish: Final = asyncio.Event()
async def request(call_id: str) -> None:
async def stream() -> AsyncGenerator[bytes, None]:
await finish.wait()
yield b"data: [DONE]\n\n"
await handler.async_post_call_success_hook(
data={"litellm_call_id": call_id}, user_api_key_dict=auth, response=ModelResponse()
)
async def produce() -> StreamingResponse:
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={"model": "test-model", "litellm_call_id": call_id},
call_type="acompletion",
)
admitted.set()
return StreamingResponse(stream())
response: Final = await open_sse_before_first_byte(produce(), 60.0)
assert tuple([chunk async for chunk in response.body_iterator]) == (b"data: [DONE]\n\n",)
first: Final = contextvars.Context().run(asyncio.create_task, request("first"))
await admitted.wait()
try:
with pytest.raises(HTTPException) as rejected:
await contextvars.Context().run(asyncio.create_task, request("overlapping"))
assert rejected.value.status_code == 429
finally:
finish.set()
await first
await contextvars.Context().run(asyncio.create_task, request("after-completion"))
@pytest.mark.asyncio
async def test_read_only_gauge_check_counts_without_acquiring_v3():
"""

View file

@ -12561,6 +12561,36 @@ async def test_team_member_me_returns_caller_membership(mock_db_client):
assert call_kwargs["team_id"] == team_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"spend, stored_total, reported_total",
[(612.4, 148.1, 612.4), (12.25, 97.8, 97.8)],
ids=["legacy_row_floors_to_cycle_spend", "healthy_row_keeps_its_total"],
)
async def test_team_member_me_reports_total_spend_no_lower_than_cycle_spend(
mock_db_client, spend, stored_total, reported_total
):
from fastapi import Request
from litellm.proxy.management_endpoints.team_endpoints import team_member_me
team_id = "team-me-floor"
caller_id = "alice@example.com"
team = _build_team_for_me(team_id, [{"user_id": caller_id, "user_email": None, "role": "user"}])
membership = _build_membership_for_me(caller_id, team_id, spend=spend).model_copy(
update={"total_spend": stored_total}
)
p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=None)
with p_team, p_membership, p_user:
response = await team_member_me(
http_request=MagicMock(spec=Request),
team_id=team_id,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id),
)
assert (response.spend, response.total_spend) == (spend, reported_total)
@pytest.mark.asyncio
async def test_team_member_me_matches_email_only_member(mock_db_client):
"""
@ -16608,7 +16638,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0),
)
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
mock_prisma_client.db.execute_raw = AsyncMock(return_value=1)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
@ -16631,15 +16661,76 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
assert response["spend"] == 0.0
assert response["previous_spend"] == 10.0
assert response["max_budget"] == 50.0
mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
data={"spend": 0.0},
)
statement, reset_user_id, reset_team_id, reset_to = mock_prisma_client.db.execute_raw.await_args.args
assert (reset_user_id, reset_team_id, reset_to) == ("member-1", "team-1", 0.0)
assert "total_spend = GREATEST(total_spend, spend)" in statement
assert "spend >= $3" in statement
assert await real_cache.async_get_cache(key="team-1_member-1") is None
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_caches(monkeypatch):
"""The reset validates reset_to against a spend it read a moment ago. If the row moved
underneath it (a flush landed, another admin reset first) the guarded UPDATE matches no
row, and the endpoint must say so instead of reporting a reset that never happened."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
mock_prisma_client = MagicMock()
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="cached-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0)
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
), pytest.raises(HTTPException) as exc_info:
await reset_team_member_spend_fn(
team_id="team-1",
user_id="member-1",
data=ResetSpendRequest(reset_to=5.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert exc_info.value.status_code == 409
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") == "cached-membership"
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
@pytest.mark.asyncio
async def test_get_all_team_memberships_floors_total_spend_at_current_cycle_spend():
"""Rows written before total_spend existed carry their history only in spend. /team/info
and /team/list both read through this helper, so a legacy row reports its lifetime spend
as at least the current cycle, while a healthy row keeps its own total."""
from litellm.proxy.management_endpoints.team_endpoints import get_all_team_memberships
legacy_row = MagicMock()
legacy_row.model_dump = lambda: {"user_id": "legacy", "team_id": "team-1", "spend": 612.4, "total_spend": 148.1}
healthy_row = MagicMock()
healthy_row.model_dump = lambda: {"user_id": "healthy", "team_id": "team-1", "spend": 12.25, "total_spend": 97.8}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[legacy_row, healthy_row])
legacy, healthy = await get_all_team_memberships(mock_prisma_client, ["team-1"])
assert (legacy.spend, legacy.total_spend) == (612.4, 612.4)
assert (healthy.spend, healthy.total_spend) == (12.25, 97.8)
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch):
mock_prisma_client = MagicMock()
@ -16755,7 +16846,7 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey
membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0)
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
mock_prisma_client.db.execute_raw = AsyncMock(return_value=1)
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
@ -17994,6 +18085,7 @@ async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an
disable_audit_logging_for_mocked_team,
):
import contextlib
from litellm.models.budget import LiteLLM_BudgetTable
with contextlib.ExitStack() as stack:

View file

@ -2,23 +2,35 @@ import json
from typing import Final
import pytest
from pydantic import ValidationError
from pydantic import BaseModel, ValidationError
from litellm.proxy._types import (
ROLES_WITHIN_ORG,
BudgetNewRequest,
ChangePasswordRequest,
GenerateKeyRequest,
GenerateKeyResponse,
KeyRequest,
LiteLLM_AuditLogs,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
NewCustomerRequest,
NewMCPServerRequest,
NewTeamRequest,
NewUserRequest,
NewUserResponse,
OrganizationMemberUpdateRequest,
RegenerateKeyRequest,
ResetSpendRequest,
UpdateKeyRequest,
UpdateMCPServerRequest,
UpdateTeamRequest,
UpdateUserRequest,
UpdateUserRequestNoUserIDorEmail,
UserAPIKeyAuth,
)
@ -213,6 +225,75 @@ def test_a_max_budget_that_was_actually_sent_is_kept(sent):
assert GenerateKeyRequest(max_budget=sent).max_budget == sent
RATE_LIMIT_REQUESTS: Final = (
(GenerateKeyRequest, {}),
(RegenerateKeyRequest, {}),
(UpdateKeyRequest, {"key": "sk-test"}),
(NewUserRequest, {}),
(UpdateUserRequestNoUserIDorEmail, {}),
(UpdateUserRequest, {"user_id": "user-test"}),
(NewTeamRequest, {}),
(UpdateTeamRequest, {"team_id": "team-test"}),
(BudgetNewRequest, {}),
(NewCustomerRequest, {"user_id": "customer-test"}),
)
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
@pytest.mark.parametrize("field", ("rpm_limit", "tpm_limit"))
@pytest.mark.parametrize("value", (-1, -100))
def test_management_requests_reject_negative_rate_limits(
request_model: type[BaseModel], identifiers: dict[str, str], field: str, value: int
) -> None:
with pytest.raises(ValidationError) as error:
request_model.model_validate({**identifiers, field: value})
assert error.value.errors()[0]["loc"] == (field,)
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
@pytest.mark.parametrize("value", (None, 0, 100))
def test_management_requests_preserve_explicit_rate_limits(
request_model: type[BaseModel], identifiers: dict[str, str], value: int | None
) -> None:
limits: Final = {"rpm_limit": value, "tpm_limit": value}
request: Final = request_model.model_validate({**identifiers, **limits})
assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == limits
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
def test_management_requests_leave_omitted_rate_limits_unset(
request_model: type[BaseModel], identifiers: dict[str, str]
) -> None:
request: Final = request_model.model_validate({**identifiers})
assert request.model_dump(include={"rpm_limit", "tpm_limit"}) == {"rpm_limit": None, "tpm_limit": None}
assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == {}
@pytest.mark.parametrize(("response_model", "identifiers"), (
(GenerateKeyResponse, {"key": "sk-test"}),
(NewUserResponse, {"key": "sk-test"}),
(LiteLLM_TeamTable, {"team_id": "team-test"}),
(LiteLLM_UserTable, {"user_id": "user-test"}),
(LiteLLM_BudgetTable, {"budget_id": "budget-test"}),
))
def test_existing_negative_rate_limits_remain_readable(
response_model: type[BaseModel], identifiers: dict[str, str]
) -> None:
limits: Final = {"rpm_limit": -1, "tpm_limit": -100}
response: Final = response_model.model_validate({**identifiers, **limits})
assert response.model_dump(include={"rpm_limit", "tpm_limit"}) == limits
def test_existing_customer_negative_rate_limits_remain_readable() -> None:
customer: Final = LiteLLM_EndUserTable.model_validate({
"user_id": "customer-test", "blocked": False,
"litellm_budget_table": {"rpm_limit": -1, "tpm_limit": -100},
})
assert customer.litellm_budget_table is not None
assert customer.litellm_budget_table.rpm_limit == -1
assert customer.litellm_budget_table.tpm_limit == -100
USER_IDENTIFIER_CASES = (
({"user_id": "user-1"}, True),
({"user_email": "user@example.com"}, True),

View file

@ -15207,21 +15207,23 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun
assert_loop_stayed_free(took, lags)
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch):
from litellm.rust_bridge._native import Tokenizer
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None:
from queue import SimpleQueue
from threading import Thread, current_thread
from tokenizers import Tokenizer
from litellm import Router
from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
claude_tokenizer: Final = litellm.utils.select_tokenizer("claude-fable-5")["tokenizer"]
claude_tokenizer: Final = Tokenizer.from_str(litellm.utils.claude_json_str)
loop_thread: Final = current_thread()
loader_threads: Final[SimpleQueue[Thread]] = SimpleQueue()
class SlowHubTokenizer:
@staticmethod
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
time.sleep(0.3)
return claude_tokenizer
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
loader_threads.put(current_thread())
return claude_tokenizer
monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", SlowHubTokenizer.from_pretrained)
monkeypatch.setattr("tokenizers.Tokenizer", types.SimpleNamespace(from_pretrained=from_pretrained))
monkeypatch.setattr(
"litellm.proxy.proxy_server.llm_router",
Router(
@ -15237,15 +15239,16 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
),
)
response, took, lags = await timed_with_loop_lags(
lambda: proxy_server_module.token_counter(
litellm.utils._select_custom_tokenizer_helper.cache_clear()
try:
response: Final = await proxy_server_module.token_counter(
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
)
)
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.total_tokens > 0
assert_loop_stayed_free(took, lags)
assert loader_threads.get_nowait() is not loop_thread
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.total_tokens > 0
finally:
litellm.utils._select_custom_tokenizer_helper.cache_clear()
async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch):

View file

@ -6,6 +6,7 @@ import pytest
from litellm.repositories.unit_of_work import (
LinkedSpendResetWrites,
SpendResetWrites,
budget_cascade_unit_of_work,
enduser_spend_reset_unit_of_work,
spend_reset_unit_of_work,
@ -39,6 +40,9 @@ class FakeBatch:
self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls)
self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls)
def execute_raw(self, query: str, *args: object) -> None:
self.calls.append(("execute_raw", {"query": query}, {"args": args}))
async def commit(self) -> None:
self.commit_count += 1
@ -166,7 +170,7 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(
async with budget_cascade_unit_of_work(_new_batch) as uow:
writes = [getattr(uow, field.name) for field in fields(uow)]
for write in writes:
if isinstance(write, LinkedSpendResetWrites):
if isinstance(write, SpendResetWrites):
write.queue_spend_zero(where={"budget_id": "budget-1"})
else:
write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at)
@ -192,3 +196,18 @@ async def test_budget_cascade_raising_inside_block_skips_commit():
await _blow_up_mid_transaction()
assert batch.commit_count == 0
async def test_membership_total_spend_repair_is_queued_on_the_cascade_batch_ahead_of_the_reset():
batch = FakeBatch()
async with budget_cascade_unit_of_work(lambda: batch) as uow:
uow.team_memberships.queue_total_spend_repair(["budget-1", "budget-2"])
uow.team_memberships.queue_spend_zero(where={"budget_id": "budget-1"})
assert batch.commit_count == 0
assert batch.commit_count == 1
repair, reset = batch.calls
assert repair[0] == "execute_raw"
assert repair[2] == {"args": (["budget-1", "budget-2"],)}
assert reset == ("litellm_teammembership.update_many", {"budget_id": "budget-1"}, {"spend": 0})

View file

@ -1065,6 +1065,11 @@ class TestCustomDimensions:
"create table widgets (id integer); ALTER TABLE widgets ADD label text;",
id="regex",
),
pytest.param(
{"patterns": [r"\b(?P<verb>create|alter)\s{1,4}?table\b"]},
"CREATE TABLE widgets (id integer)",
id="named-group-alternation-lazy-repeat",
),
],
)
def test_custom_dimension_changes_only_matching_requests(

View file

@ -1,6 +1,7 @@
import base64
import binascii
import json
from collections import OrderedDict
from dataclasses import dataclass
from typing import Final
@ -120,7 +121,7 @@ def manager(monkeypatch: pytest.MonkeyPatch) -> GoogleSecretManager:
monkeypatch.setattr(litellm, "secret_manager_client", None)
monkeypatch.setattr(litellm, "_key_management_system", None)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings())
monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", {})
monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", OrderedDict())
monkeypatch.setattr(google.auth, "default", _default_credentials)
monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", PROJECT_ID)
monkeypatch.delenv("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", raising=False)

View file

@ -455,30 +455,23 @@ def is_active_priced_mistral_chat_row(name: str, entry: Mapping[str, object]) ->
)
def cache_read_is_tenth_of_input(entry: Mapping[str, object]) -> bool:
def cache_read_is_below_input(entry: Mapping[str, object]) -> bool:
cache_read: Final = entry.get("cache_read_input_token_cost")
input_cost: Final = entry.get("input_cost_per_token")
return (
isinstance(cache_read, float)
and isinstance(input_cost, (int, float))
and 0 < cache_read < input_cost
and cache_read == pytest.approx(input_cost / 10)
)
return isinstance(cache_read, float) and isinstance(input_cost, (int, float)) and 0 < cache_read < input_cost
@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup"))
def test_active_mistral_chat_rows_price_cache_reads_below_input(path: Path):
"""A Mistral chat row without a cache-read rate bills cached prompt tokens at zero, so every
active priced row must carry one, and it must be cheaper than a fresh input token. Mistral
bills cached tokens at 10% of the input price for every model (docs.mistral.ai/studio/
conversations/advanced/prompt-caching, read 2026-09-18), so the ratio is checked as well."""
active priced row must carry one, and it must be cheaper than a fresh input token."""
rows: Mapping[str, object] = json.loads(path.read_text())
drifted: Final = [
f"{name}: cache_read={entry.get('cache_read_input_token_cost')} input={entry.get('input_cost_per_token')}"
for name, entry in rows.items()
if isinstance(entry, dict)
and is_active_priced_mistral_chat_row(name, entry)
and not cache_read_is_tenth_of_input(entry)
and not cache_read_is_below_input(entry)
]
assert drifted == []

View file

@ -1086,6 +1086,7 @@ _COST_MAP_ROW_SCHEMA: Final = {
"chat",
"completion",
"container",
"decisions",
"image_edit",
"embedding",
"evaluation",

View file

@ -1,14 +1,20 @@
import importlib
from typing import Final
from uuid import NAMESPACE_DNS, UUID, uuid5 as stdlib_uuid5
from litellm._uuid import uuid, uuid4
def test_uses_fastuuid_and_uuid4_works():
mod = importlib.import_module("litellm._uuid")
fastuuid_mod = importlib.import_module("fastuuid")
assert hasattr(mod, "uuid4")
assert hasattr(mod, "uuid")
assert mod.uuid is fastuuid_mod
def test_uuid4_returns_canonical_random_uuid() -> None:
for value in (uuid4(), uuid.uuid4()):
encoded: Final = str(value)
parsed: Final = UUID(encoded)
assert parsed.version == 4
assert encoded == str(parsed)
assert value.bytes == parsed.bytes
# Ensure uuid4 returns something that looks like a UUID string
val = str(mod.uuid4())
assert isinstance(val, str)
assert len(val) >= 8
def test_uuid5_accepts_byte_names_for_stable_ids() -> None:
namespace: Final = uuid.UUID(str(NAMESPACE_DNS))
name: Final = "guardrail-\u00e9"
generated: Final = uuid.uuid5(namespace, name.encode("utf-8"))
assert str(generated) == str(stdlib_uuid5(NAMESPACE_DNS, name))

View file

@ -3,7 +3,7 @@
# UI container — Next.js static export served by nginx.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
ARG NGINX_VERSION=1.31.5-alpine3.24@sha256:34f40471dea485273c5e2a04dd5e97a682332ceb4a9adecd67de450dcb2fb390
ARG NGINX_VERSION=1.31.8-alpine3.24@sha256:1ef7ca94d7879df82421428ce13bd1eaa412404953859cf7a5b6512c25e67a60
# ---------- builder ----------
FROM ${UI_BUILD_IMAGE} AS builder
@ -27,6 +27,8 @@ RUN npm run build
# ---------- runtime ----------
FROM nginx:${NGINX_VERSION} AS runtime
RUN apk upgrade --no-cache brotli-libs
# Drop the upstream default :80 server; we own the config.
RUN rm -f /etc/nginx/conf.d/default.conf

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

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" />);

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

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", () => {

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]);

View file

@ -6,6 +6,12 @@ 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 =>
({
@ -118,3 +124,22 @@ describe("DrawerHeader back link", () => {
expect(screen.queryByRole("button", { name: /^Back to/ })).not.toBeInTheDocument();
});
});
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 });
vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success);
renderHeader(logEntry({ request_id: "request-copy", litellm_call_id: "call-copy" }), false);
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(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");
});
});

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";
@ -183,12 +184,9 @@ function CopyableId({
const [copied, setCopied] = useState(false);
const handleCopy = async () => {
try {
await navigator.clipboard.writeText(value);
if (await copyToClipboard(value)) {
setCopied(true);
setTimeout(() => setCopied(false), 1200);
} catch {
/* clipboard unavailable in non-secure contexts */
}
};

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 () => {

View file

@ -4,7 +4,7 @@
*/
import { useState } from "react";
import { toast } from "@/lib/toast";
import { copyToClipboard } from "@/utils/dataUtils";
import { ParsedMessage } from "./prettyMessagesTypes";
import { SectionHeader } from "./SectionHeader";
import { CollapsibleMessage } from "./CollapsibleMessage";
@ -32,8 +32,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps)
const handleCopy = () => {
const content = lastMessage?.content || "";
navigator.clipboard.writeText(content);
toast.success("Input copied");
void copyToClipboard(content, "Input copied");
};
return (

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) => {

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";
@ -304,12 +305,9 @@ function CopyButton({
const [copied, setCopied] = useState(false);
const handleCopy = async () => {
try {
await navigator.clipboard.writeText(getText());
if (await copyToClipboard(getText())) {
setCopied(true);
setTimeout(() => setCopied(false), 1200);
} catch {
/* clipboard unavailable in non-secure contexts */
}
};

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);
});
});

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";
@ -289,12 +290,9 @@ export function LogDetailsDrawer({
const handleCopyLeftPanelId = async () => {
if (!leftPanelId) return;
try {
await navigator.clipboard.writeText(leftPanelId);
if (await copyToClipboard(leftPanelId)) {
setCopiedLeftPanelId(true);
setTimeout(() => setCopiedLeftPanelId(false), 1200);
} catch {
/* clipboard unavailable in non-secure contexts */
}
};

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();
});
});

View file

@ -4,7 +4,7 @@
*/
import { useState } from "react";
import { toast } from "@/lib/toast";
import { copyToClipboard } from "@/utils/dataUtils";
import { COLOR_BORDER } from "./constants";
import { ParsedMessage } from "./prettyMessagesTypes";
import { SectionHeader } from "./SectionHeader";
@ -22,8 +22,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard
const handleCopy = () => {
if (!message) return;
navigator.clipboard.writeText(message.content || "");
toast.success("Output copied");
void copyToClipboard(message.content || "", "Output copied");
};
return (

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", () => {

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
@ -276,7 +277,7 @@ function ConversationCard({
(r.output || []).flatMap((o) => (o.content || []).map((c) => `${o.role}: ${c.transcript || c.text || ""}`)),
)
.join("\n");
navigator.clipboard.writeText(transcripts);
void copyToClipboard(transcripts);
};
return (

View file

@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event";
import { afterEach, describe, expect, it, vi } from "vitest";
import type { ComponentProps } from "react";
import { copyToClipboard } from "@/utils/dataUtils";
import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
import { renderWithProviders, screen } from "../../../tests/test-utils";
import CopyButton from "./CopyButton";
vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() }));
@ -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,55 +24,37 @@ describe("CopyButton", () => {
expect(screen.queryByRole("button")).not.toBeInTheDocument();
});
it("shows the confirmation checkmark only after a successful write", async () => {
it("shows the icon confirmation only after copying succeeds", async () => {
const user = userEvent.setup();
const writeText = vi.fn().mockResolvedValue(undefined);
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
const pendingCopy = Promise.withResolvers<boolean>();
vi.mocked(copyToClipboard).mockReturnValue(pendingCopy.promise);
renderWithProviders(<CopyButton value="secret-value" label="Copy value" copiedLabel="Value copied" />);
renderWithProviders(<CopyButton value="secret-value" label="Copy value" />);
await user.click(screen.getByRole("button", { name: "Copy value" }));
const button = 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 user.click(button);
expect(writeText).toHaveBeenCalledWith("secret-value");
expect(await screen.findByRole("img", { name: "Copied icon" })).toBeInTheDocument();
await act(async () => pendingCopy.resolve(true));
expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument();
});
it("does not show the checkmark when the clipboard write is rejected", async () => {
it("keeps the icon unchanged when copying fails", async () => {
const user = userEvent.setup();
const writeText = vi.fn().mockRejectedValue(new Error("permission denied"));
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
vi.mocked(copyToClipboard).mockResolvedValue(false);
renderWithProviders(<CopyButton value="secret-value" label="Copy value" />);
const button = screen.getByRole("button", { name: "Copy value" });
await user.click(button);
await user.click(screen.getByRole("button", { name: "Copy value" }));
await waitFor(() => expect(writeText).toHaveBeenCalledWith("secret-value"));
expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Copied");
expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument();
expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument();
});
it("does not show the checkmark when the clipboard API is unavailable", async () => {
const user = userEvent.setup();
Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true });
renderWithProviders(<CopyButton value="secret-value" label="Copy value" />);
const button = screen.getByRole("button", { name: "Copy value" });
await user.click(button);
expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument();
expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument();
expect(copyToClipboard).not.toHaveBeenCalled();
});
it("keeps the default icon styling, custom classes, title and 1200ms confirmation", async () => {
vi.useFakeTimers();
const writeText = vi.fn().mockResolvedValue(undefined);
Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true });
vi.mocked(copyToClipboard).mockResolvedValue(true);
renderWithProviders(
<CopyButton value="request-id" label="Copy request ID" className="shrink-0" iconClassName="size-3" />,
);
@ -87,8 +65,7 @@ describe("CopyButton", () => {
expect(screen.getByRole("img", { name: "Copy icon" })).toHaveClass("size-3");
await act(async () => fireEvent.click(button));
expect(writeText).toHaveBeenCalledWith("request-id");
expect(copyToClipboard).not.toHaveBeenCalled();
expect(copyToClipboard).toHaveBeenCalledWith("request-id", "Copied");
act(() => vi.advanceTimersByTime(1199));
expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument();
act(() => vi.advanceTimersByTime(1));

View file

@ -54,17 +54,7 @@ const CopyButton: React.FC<CopyButtonProps> = ({
if (value == null || (!value && !isAction)) return null;
const handleCopy = async () => {
if (isAction) {
setCopied(await copyToClipboard(value, copiedLabel));
return;
}
if (!navigator.clipboard) return;
try {
await navigator.clipboard.writeText(value);
setCopied(true);
} catch {
setCopied(false);
}
setCopied(await copyToClipboard(value, copiedLabel));
};
const button = (

View file

@ -45,5 +45,30 @@ describe("MyUserTab", () => {
expect(screen.getByText("member@example.com")).toBeInTheDocument();
expect(screen.getByText("model-one")).toBeInTheDocument();
expect(screen.getByText("TPM: 1,000")).toBeInTheDocument();
expect(screen.getByText("Never resets")).toBeInTheDocument();
});
it("shows the next reset time when the member budget has a reset window", () => {
vi.mocked(useMyTeamMember).mockReturnValue({
data: {
user_id: "user-1",
user_email: "member@example.com",
team_id: "team-1",
role: "user",
spend: 12.5,
total_spend: 30,
litellm_budget_table: {
max_budget: 100,
budget_duration: "30d",
budget_reset_at: new Date(2026, 6, 15, 12, 0, 0).toISOString(),
},
},
isLoading: false,
error: null,
} as ReturnType<typeof useMyTeamMember>);
renderWithProviders(<MyUserTab teamId="team-1" />);
expect(screen.getByText("Resets Jul 15, 2026 12:00")).toBeInTheDocument();
});
});

View file

@ -1,4 +1,4 @@
import { formatBudgetReset } from "@/utils/budgetUtils";
import { describeBudgetResetSchedule } from "@/utils/budgetUtils";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { SimpleTooltip } from "@/components/ui/tooltip";
import { Badge } from "@/components/ui/badge";
@ -67,7 +67,7 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
const totalSpend = data.total_spend ?? 0;
const tpmLimit = budgetTable?.tpm_limit ?? null;
const rpmLimit = budgetTable?.rpm_limit ?? null;
const budgetReset = formatBudgetReset(budgetTable?.budget_reset_at);
const budgetReset = describeBudgetResetSchedule(budgetTable?.budget_duration, budgetTable?.budget_reset_at);
const allowedModels = budgetTable?.allowed_models ?? null;
return (
@ -93,17 +93,16 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
<div className="grid grid-cols-1 gap-4 md:grid-cols-2">
<Card>
<CardContent>
{labelWithTooltip(
"Current Cycle Spend (USD)",
"Spend for the current budget cycle. Resets to $0 when the budget window rolls over.",
)}
{labelWithTooltip("Current Cycle Spend (USD)", budgetReset.tooltip)}
<div className="mt-2">
<h3 className="text-2xl font-semibold">${formatNumber(spend, 4)}</h3>
<span className="text-muted-foreground">
of {maxBudget === null ? "Unlimited" : `$${formatNumber(maxBudget, 4)}`}
</span>
</div>
{budgetReset && <div className="mt-1 text-muted-foreground">Resets {budgetReset}</div>}
<div className="mt-1 text-muted-foreground">
{budgetReset.resetAt ? `Resets ${budgetReset.label}` : budgetReset.label}
</div>
</CardContent>
</Card>
@ -120,7 +119,10 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
<Card>
<CardContent>
{labelWithTooltip("Total Spend (USD)", "Cumulative spend across all budget cycles within this team.")}
{labelWithTooltip(
"Total Spend (USD)",
"Cumulative spend across budget cycles within this team. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.",
)}
<h4 className="mt-2 text-xl font-semibold">${formatNumber(totalSpend, 4)}</h4>
</CardContent>
</Card>

View file

@ -313,7 +313,47 @@ describe("TeamMembersComponent", () => {
/>,
);
expect(screen.getByText("Jul 15, 2026")).toBeInTheDocument();
expect(screen.getByText("Jul 15, 2026 12:00")).toBeInTheDocument();
});
it("shows Never resets for a member whose budget has no reset window", () => {
const teamData = createMockTeamData();
teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null;
renderWithProviders(
<TeamMembersComponent
teamData={teamData}
canEditTeam={false}
handleMemberDelete={mockHandleMemberDelete}
onMemberSpendReset={mockOnMemberSpendReset}
onMemberBudgetReset={mockOnMemberBudgetReset}
setSelectedEditMember={mockSetSelectedEditMember}
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
/>,
);
expect(screen.queryByText("Jul 15, 2026 12:00")).not.toBeInTheDocument();
expect(screen.getAllByText("Never resets").length).toBeGreaterThan(0);
});
it("shows Unknown for a member whose budget has a duration but no scheduled reset yet", () => {
const teamData = createMockTeamData();
teamData.team_memberships[0].litellm_budget_table!.budget_duration = "30d";
teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null;
renderWithProviders(
<TeamMembersComponent
teamData={teamData}
canEditTeam={false}
handleMemberDelete={mockHandleMemberDelete}
onMemberSpendReset={mockOnMemberSpendReset}
onMemberBudgetReset={mockOnMemberBudgetReset}
setSelectedEditMember={mockSetSelectedEditMember}
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
/>,
);
expect(screen.getByText("Unknown")).toBeInTheDocument();
});
it("shows each member the values from its own membership row", () => {

View file

@ -10,8 +10,9 @@ import MemberTable from "@/components/common_components/MemberTable";
import type { StoredModelMaxBudget } from "@/components/key_team_helpers/modelMaxBudgetPayload";
import { Member } from "@/components/networking";
import { parseErrorMessage } from "@/components/shared/errorUtils";
import { DateCell, MoneyCell } from "@/components/shared/table_cells";
import { MoneyCell } from "@/components/shared/table_cells";
import { toast } from "@/lib/toast";
import { describeBudgetResetSchedule } from "@/utils/budgetUtils";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles";
import { CircleHelp } from "lucide-react";
@ -152,6 +153,16 @@ export default function TeamMemberTab({
return findMembership(userId)?.litellm_budget_table?.budget_reset_at ?? null;
};
const renderBudgetReset = (userId: string | null) => {
const budget = findMembership(userId)?.litellm_budget_table;
const schedule = describeBudgetResetSchedule(budget?.budget_duration ?? null, budget?.budget_reset_at ?? null);
return (
<SimpleTooltip content={schedule.tooltip}>
<span>{schedule.label}</span>
</SimpleTooltip>
);
};
const extraColumns: NonNullable<ComponentProps<typeof MemberTable>["extraColumns"]> = [
{
title: (
@ -190,7 +201,7 @@ export default function TeamMemberTab({
title: (
<span className="flex items-center gap-1">
Current Cycle Spend (USD)
<SimpleTooltip content="Spend for the current budget cycle. Resets to $0 when the member's budget window rolls over. This is the value checked against the member's budget.">
<SimpleTooltip content="Spend in this member's current budget cycle, the value checked against their budget. The Budget Reset column shows when this cycle ends, or Never resets. Spend above a rollover cap at reset carries into the next cycle.">
<CircleHelp className="size-4" aria-label="Current cycle spend information" />
</SimpleTooltip>
</span>
@ -204,7 +215,7 @@ export default function TeamMemberTab({
title: (
<span className="flex items-center gap-1">
Total Spend (USD)
<SimpleTooltip content="Cumulative spend by this member within this team, across all budget cycles. Tracking began 2026-04-21; spend from before that date is not included.">
<SimpleTooltip content="Cumulative spend by this member within this team, across budget cycles. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.">
<CircleHelp className="size-4" aria-label="Total spend information" />
</SimpleTooltip>
</span>
@ -255,7 +266,7 @@ export default function TeamMemberTab({
title: "Budget Reset",
key: "budget_reset",
sortValue: (record: Member) => getUserBudgetReset(record.user_id),
render: (record: Member) => <DateCell value={getUserBudgetReset(record.user_id)} precision="date" />,
render: (record: Member) => renderBudgetReset(record.user_id),
},
{
title: (

View file

@ -1,6 +1,6 @@
import { afterEach, describe, expect, it } from "vitest";
import { formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils";
import { describeBudgetResetSchedule, formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils";
const originalTimezone: string | undefined = process.env.TZ;
@ -21,3 +21,26 @@ describe("formatBudgetResetUtc", () => {
expect(formatBudgetReset(resetAt)).toBe("Oct 19, 2026");
});
});
describe("describeBudgetResetSchedule", () => {
it("reports Never resets when the budget has no duration and no scheduled reset", () => {
const schedule = describeBudgetResetSchedule(null, null);
expect(schedule.label).toBe("Never resets");
expect(schedule.resetAt).toBeNull();
});
it("reports Unknown when a duration is set but no reset has been scheduled yet", () => {
const schedule = describeBudgetResetSchedule("30d", null);
expect(schedule.label).toBe("Unknown");
expect(schedule.tooltip).toContain("every 30d");
});
it("shows the next reset with its time zone and the rollover carryover note", () => {
const resetAt = new Date(2026, 6, 15, 12, 0, 0);
const schedule = describeBudgetResetSchedule("30d", resetAt.toISOString());
expect(schedule.label).toBe("Jul 15, 2026 12:00");
expect(schedule.resetAt).toBe(resetAt.toISOString());
expect(schedule.tooltip).toContain(Intl.DateTimeFormat().resolvedOptions().timeZone);
expect(schedule.tooltip).toContain("carries into the next cycle");
});
});

View file

@ -16,3 +16,41 @@ export function formatBudgetResetUtc(iso: string | null | undefined): string | n
if (!resetDate.isValid()) return null;
return resetDate.format("MMM D, YYYY");
}
export interface BudgetResetSchedule {
label: string;
resetAt: string | null;
tooltip: string;
}
const ROLLOVER_NOTE = "Spend above a rollover cap at reset carries into the next cycle.";
export function describeBudgetResetSchedule(
budgetDuration: string | null | undefined,
budgetResetAt: string | null | undefined,
): BudgetResetSchedule {
const resetDate = budgetResetAt ? dayjs(budgetResetAt) : null;
if (resetDate?.isValid()) {
const cadence = budgetDuration ? `every ${budgetDuration}` : "on a schedule";
const timeZone = Intl.DateTimeFormat().resolvedOptions().timeZone;
const label = resetDate.format("MMM D, YYYY HH:mm");
return {
label,
resetAt: resetDate.toISOString(),
tooltip: `Current Cycle Spend resets ${cadence}. Next reset ${label} (${timeZone}). ${ROLLOVER_NOTE}`,
};
}
if (!budgetDuration) {
return {
label: "Never resets",
resetAt: null,
tooltip:
"This budget has no reset window, so Current Cycle Spend accumulates for as long as the membership exists.",
};
}
return {
label: "Unknown",
resetAt: null,
tooltip: `Resets every ${budgetDuration}, but the next reset time has not been scheduled yet. ${ROLLOVER_NOTE}`,
};
}

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();
});
});
});

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

1630
uv.lock generated

File diff suppressed because it is too large Load diff