scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

Merge branch 'litellm_internal_staging' into litellm_mcp_server_metadata

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mubashir Osmani 2026-07-24 22:38:10 +00:00
commit 30f2eb307a
243 changed files with 16144 additions and 6116 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/_test-unit-base.yml vendored
View file

@ -80,7 +80,7 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'

14
.github/workflows/codspeed.yml vendored
View file

@ -5,10 +5,24 @@ on:
branches:
- main
- litellm_internal_staging
paths:
- "litellm/**"
- "tests/benchmarks/**"
- "pyproject.toml"
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
pull_request:
branches:
- main
- litellm_internal_staging
paths:
- "litellm/**"
- "tests/benchmarks/**"
- "pyproject.toml"
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
# Allow CodSpeed to trigger backtest performance analysis
# in order to generate initial data
workflow_dispatch:

2
.github/workflows/mutation-test.yml vendored
View file

@ -55,7 +55,7 @@ jobs:
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
- name: Generate Prisma client
env:

2
Dockerfile
View file

@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
# Copy full source tree
@ -84,6 +85,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

2
docker/Dockerfile.database
View file

@ -62,6 +62,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
# Copy full source tree
@ -82,6 +83,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

3
docker/Dockerfile.non_root
View file

@ -68,6 +68,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
# Copy full source tree
@ -94,6 +95,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3 \
--no-sources-package litellm-proxy-extras; \
else \
@ -102,6 +104,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3; \
fi

2
gateway/Dockerfile
View file

@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--python python3
# Stage 2 — copy source and install the project + workspace members.
@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--python python3
RUN mkdir -p /home/nonroot && \

3
litellm/__init__.py
View file

@ -211,6 +211,9 @@ filter_invalid_headers: Optional[bool] = False
add_user_information_to_llm_headers: Optional[bool] = (
None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers
)
overwrite_user_with_key_hash: bool = (
False # force the outgoing `user` param to the hashed api key, so providers see a stable, tamper-proof id
)
store_audit_logs = False # Enterprise feature, allow users to see audit logs
skip_system_message_in_guardrail: bool = False
skip_tool_message_in_guardrail: bool = False

1
litellm/constants.py
View file

@ -1145,6 +1145,7 @@ BEDROCK_CONVERSE_MODELS = [
"anthropic.claude-sonnet-4-5-20250929-v1:0",
"anthropic.claude-fable-5",
"anthropic.claude-sonnet-5",
"anthropic.claude-opus-5",
"anthropic.claude-opus-4-8",
"anthropic.claude-opus-4-7",
"anthropic.claude-opus-4-6-v1:0",

9
litellm/cost_calculator.py
View file

@ -2191,6 +2191,13 @@ def batch_cost_calculator(
return total_prompt_cost, total_completion_cost
def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str]:
field_names = list(type(prompt_tokens_details).model_fields)
if getattr(prompt_tokens_details, "cache_write_tokens", None) is None:
return field_names
return [attr for attr in field_names if attr != "cache_creation_tokens"]
class BaseTokenUsageProcessor:
@staticmethod
def combine_usage_objects(usage_objects: List[Usage]) -> Usage:
@ -2225,7 +2232,7 @@ class BaseTokenUsageProcessor:
# Check what keys exist in the model's prompt_tokens_details
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
for attr in type(usage.prompt_tokens_details).model_fields:
for attr in _summable_prompt_token_fields(usage.prompt_tokens_details):
if (
hasattr(usage.prompt_tokens_details, attr)
and not attr.startswith("_")

44
litellm/integrations/custom_guardrail.py
View file

@ -1,3 +1,4 @@
import contextvars
import hashlib
import os
import secrets
@ -64,6 +65,10 @@ from litellm.exceptions import (
# proxy's metadata sanitizer.
_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16)
_guardrail_self_recorded: contextvars.ContextVar[bool] = contextvars.ContextVar(
"litellm_guardrail_self_recorded", default=False
)
def _strict_guardrail_modes_enabled() -> bool:
"""Whether guardrail-mode validation raises (default) or logs a warning.
@ -117,6 +122,7 @@ class CustomGuardrail(CustomLogger):
on_sensitive_data: Optional[str] = None,
sensitive_data_route_to_model: Optional[str] = None,
sticky_session_routing: bool = True,
run_in_parallel: bool = False,
only_scan_new_messages: bool = False,
**kwargs,
):
@ -136,6 +142,9 @@ class CustomGuardrail(CustomLogger):
on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route'
sensitive_data_route_to_model: Model to route to when on_sensitive_data='route'
sticky_session_routing: When True, all subsequent requests in the session use the same model
run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with
other opted-in guardrails of the same hook. Only safe for block-only guardrails that
do not mutate the request or response.
"""
self.guardrail_name = guardrail_name
self.supported_event_hooks = supported_event_hooks
@ -150,6 +159,7 @@ class CustomGuardrail(CustomLogger):
self.on_sensitive_data: Optional[str] = on_sensitive_data
self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model
self.sticky_session_routing: bool = sticky_session_routing
self.run_in_parallel: bool = run_in_parallel
self.only_scan_new_messages: bool = only_scan_new_messages
if supported_event_hooks:
@ -956,6 +966,8 @@ class CustomGuardrail(CustomLogger):
request_data["metadata"] = {}
_append_guardrail_info(request_data["metadata"])
_guardrail_self_recorded.set(True)
# Emit the otel guardrail span here, where every guardrail execution lands,
# rather than relying on a post-call hook that does not fire on every path
# (e.g. a pass-through request that passes its guardrails).
@ -1238,8 +1250,12 @@ def log_guardrail_information(func):
(structured detections, tracing detail) than this decorator's
"allow"/"mask"/raw-response default. To avoid double-recording in that
case (which would emit two spans, two Datadog records, two spend-log
entries, etc.), snapshot the entry count before invocation: if the
wrapped function already appended its own entry, skip the auto-record.
entries, etc.), a context-local flag records whether the wrapped function
appended its own entry; if so, the auto-record is skipped. The flag is a
``ContextVar`` rather than a count of entries in the shared ``request_data``
so it stays correct when guardrails run concurrently (asyncio copies the
context into each gathered task): counting shared entries would let one
guardrail's append hide another guardrail's missing record.
"""
import functools
import inspect
@ -1259,16 +1275,6 @@ def log_guardrail_information(func):
return GuardrailEventHooks.post_call
return None
def _count_recorded_guardrail_entries(request_data: dict) -> int:
total = 0
for container_key in ("metadata", "litellm_metadata"):
container = request_data.get(container_key)
if isinstance(container, dict):
entries = container.get("standard_logging_guardrail_information")
if isinstance(entries, list):
total += len(entries)
return total
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = datetime.now() # Move start_time inside the wrapper
@ -1282,10 +1288,10 @@ def log_guardrail_information(func):
original_inputs = kwargs.get("inputs")
logging_obj = kwargs.get("logging_obj")
entries_before = _count_recorded_guardrail_entries(request_data)
self_recorded_token = _guardrail_self_recorded.set(False)
try:
response = await func(*args, **kwargs)
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
return response
return self._process_response(
response=response,
@ -1297,7 +1303,7 @@ def log_guardrail_information(func):
original_inputs=original_inputs,
)
except Exception as e:
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
raise
return self._process_error(
e=e,
@ -1308,6 +1314,7 @@ def log_guardrail_information(func):
event_type=event_type,
)
finally:
_guardrail_self_recorded.reset(self_recorded_token)
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
@functools.wraps(func)
@ -1323,10 +1330,10 @@ def log_guardrail_information(func):
original_inputs = kwargs.get("inputs")
logging_obj = kwargs.get("logging_obj")
entries_before = _count_recorded_guardrail_entries(request_data)
self_recorded_token = _guardrail_self_recorded.set(False)
try:
response = func(*args, **kwargs)
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
return response
return self._process_response(
response=response,
@ -1336,7 +1343,7 @@ def log_guardrail_information(func):
original_inputs=original_inputs,
)
except Exception as e:
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
raise
return self._process_error(
e=e,
@ -1345,6 +1352,7 @@ def log_guardrail_information(func):
event_type=event_type,
)
finally:
_guardrail_self_recorded.reset(self_recorded_token)
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
@functools.wraps(func)

27
litellm/integrations/otel/model/spans.py
View file

@ -18,12 +18,14 @@ before the LLM call even starts), so a guardrail is a sibling of the LLM call,
not a child of it. The emitter parents every span to the ambient OTel context
(the active server span), which matches this.
MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) are intentionally NOT in this
tree. Per the OTel GenAI MCP semconv, MCP and the HTTP transport are independent
contexts, so an MCP span parents to the trace context the client propagated in
``params._meta`` (or starts its own root when none is propagated) and records the
``PROXY_REQUEST`` transport span as a span *link*, never a parent. The registry
encodes this as ``parent=None, links=PROXY_REQUEST``.
MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) have two shapes, chosen at emit
time by :func:`resolve_mcp_span_context`. When the client propagates trace context
in ``params._meta`` MCP and the HTTP transport are independent contexts per the
OTel GenAI MCP semconv, so the span parents to that propagated context and records
the ``PROXY_REQUEST`` transport span as a span *link*, never a parent — the shape
this registry's ``parent=None, links=PROXY_REQUEST`` entry encodes. When nothing is
propagated (the common case) the span nests under the transport span of the request
carrying that message, so the tool call stays in one trace.
Not every service call becomes a span — :func:`span_role_for_service` decides:
@ -89,12 +91,13 @@ class SpanSpec:
SPAN_REGISTRY: dict[SpanRole, SpanSpec] = {
SpanRole.PROXY_REQUEST: SpanSpec(SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None),
SpanRole.LLM_CALL: SpanSpec(SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST),
# MCP and the HTTP transport are independent contexts (OTel GenAI MCP semconv),
# so an MCP span does not nest under the transport span. The proxy is an MCP
# client to the upstream server, so it's a CLIENT span; it parents to the trace
# context the client propagated in ``params._meta`` (or starts its own root when
# none is propagated) and records the PROXY_REQUEST transport span as a span
# *link*, never a parent — hence ``parent=None, links=PROXY_REQUEST``.
# The proxy is an MCP client to the upstream server, so MCP spans are CLIENT
# spans. With trace context propagated in ``params._meta``, MCP and the HTTP
# transport are independent contexts (OTel GenAI MCP semconv): the span parents
# to the propagated context and records the PROXY_REQUEST transport span as a
# span *link*, never a parent — the shape ``parent=None, links=PROXY_REQUEST``
# encodes. With nothing propagated, ``resolve_mcp_span_context`` nests the span
# under that message's transport span instead, keeping the call in one trace.
SpanRole.MCP_TOOL_CALL: SpanSpec(
SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=None, links=SpanRole.PROXY_REQUEST
),

110
litellm/integrations/otel/plumbing/context.py
View file

@ -5,7 +5,14 @@ from typing import Mapping
from opentelemetry import baggage
from opentelemetry.context import Context, get_current
from opentelemetry.trace import Link, Span, get_current_span, set_span_in_context
from opentelemetry.trace import (
Link,
NonRecordingSpan,
Span,
SpanContext,
get_current_span,
set_span_in_context,
)
from opentelemetry.trace.propagation.tracecontext import (
TraceContextTextMapPropagator,
)
@ -72,6 +79,62 @@ def reset_mcp_message_trace_carrier(token: "Token[Mapping[str, str] | None]") ->
_mcp_message_trace_carrier.reset(token)
# The transport span of the HTTP request carrying the CURRENT MCP message, as a
# plain ``SpanContext`` so it can cross a task boundary.
#
# ``_request_root_span`` above cannot be used for MCP: a *stateful* streamable-HTTP
# session runs every message on the single task spawned by that session's
# ``initialize`` POST, so the ContextVar the ASGI request task writes at auth time
# is frozen at ``initialize`` there and never sees the later ``tools/call`` POSTs.
# Reading it from the message handler would parent every tool call in the session
# to the first request's (already ended) server span. The gateway instead resolves
# the current message's transport span on the request task and hands it over the
# same way it hands over per-request auth, and the handler publishes it here for
# the span emitter to pick up.
_mcp_message_transport_span_context: "ContextVar[SpanContext | None]" = ContextVar(
"litellm_otel_mcp_message_transport_span_context", default=None
)
def set_mcp_message_transport_span_context(
span_context: "SpanContext | None",
) -> "Token[SpanContext | None]":
"""Publish the transport span of the request carrying the current MCP message.
Returns the reset token; the caller must reset it once the message is handled
so the transport never leaks to the next message on the same session task.
"""
return _mcp_message_transport_span_context.set(span_context)
def reset_mcp_message_transport_span_context(token: "Token[SpanContext | None]") -> None:
_mcp_message_transport_span_context.reset(token)
def request_root_span_context() -> "SpanContext | None":
"""The anchored request root span's context, safe to hand to another task.
A ``SpanContext`` is an immutable value, unlike the live ``Span``, so passing it
across the MCP session-task boundary cannot keep a finished span alive or invite
writes to it from the wrong request.
"""
span = request_root_span()
return span.get_span_context() if span is not None else None
def _mcp_transport_span_context() -> "SpanContext | None":
"""The transport span an MCP message span should attach to.
Prefers the transport the gateway published for this specific message; falls
back to the ambient request anchor for paths that emit an MCP span on the
request task itself (the REST MCP endpoints, the SDK).
"""
published = _mcp_message_transport_span_context.get()
if published is not None and published.is_valid:
return published
return request_root_span_context()
def set_request_baggage(values: Mapping[str, str], context: Context | None = None) -> Context:
"""Return a context with ``values`` written into Baggage."""
ctx = context
@ -132,33 +195,44 @@ def resolve_request_span_context() -> Context:
def resolve_mcp_span_context(
carrier: "Mapping[str, str] | None" = None,
) -> "tuple[Context, tuple[Link, ...]]":
"""Parent context + links for an MCP message span, per the OTel GenAI MCP semconv.
"""Parent context + links for an MCP message span.
MCP and the underlying transport (HTTP) are independent lifecycles — one
streamable-HTTP session multiplexes many messages, so nesting the message span
under the HTTP/session span is wrong (it renders the message at the session's
start, skewed by however long the session has been open). Instead:
When the client propagates W3C trace context in the request's ``params._meta``
(SEP-414), MCP and the underlying transport are independent lifecycles — one
streamable-HTTP session multiplexes many messages, and the client's own span is
the truthful parent. So, per the OTel GenAI MCP semconv:
* parent to the trace context the client propagated in the request's
``params._meta`` (a *remote* parent), and
* record the transport/session span as a *link*, never the parent.
* parent to the trace context the client propagated (a *remote* parent), and
* record the transport span as a *link*, never the parent.
Almost no client implements SEP-414 yet, so in practice nothing is propagated.
Rooting the span there splits a single tool call into two disconnected traces
joined only by a link, which is how it surfaces in APM: the ``POST`` transaction
and the ``tools/call`` span share no trace. With no remote parent to honor,
parent to the transport span of the request carrying this message instead, so
the call stays in one trace; no link is added since the transport is now the
real parent. The transport comes from :func:`_mcp_transport_span_context`, which
is the *current message's* POST rather than whatever request happened to open
the session, so a long-lived session does not glue every message under its
first request. With neither a remote parent nor a transport the returned context
carries no span and the span legitimately starts its own root trace.
Only trace context (``traceparent``/``tracestate``) is extracted, never the
client's W3C Baggage: ``params._meta`` is caller-controlled, and the otel
baggage processor stamps allowlisted baggage keys (``litellm.team.id``,
``litellm.metadata.*``, ...) onto the span as attributes, so honoring remote
baggage would let a client spoof a span's identity attribution.
With no propagated context the returned context carries no span, so the span
starts its own root trace (still linked to the transport). The base context is
explicitly empty so an absent ``traceparent`` can never fall through to the
ambient (stale session) span.
baggage would let a client spoof a span's identity attribution. The base context
for extraction is explicitly empty so an absent or malformed ``traceparent`` can
never fall through to the ambient (stale session) span.
"""
source = carrier if carrier is not None else _mcp_message_trace_carrier.get()
parent = _PROPAGATOR.extract(dict(source or {}), context=Context())
transport = request_root_span()
links = (Link(transport.get_span_context()),) if transport is not None else ()
return parent, links
transport = _mcp_transport_span_context()
if is_recordable_span(get_current_span(parent)):
return parent, (Link(transport),) if transport is not None else ()
if transport is not None:
return context_from_span(NonRecordingSpan(transport)), ()
return parent, ()
def is_recordable_span(obj: object) -> bool:

47
litellm/litellm_core_utils/litellm_logging.py
View file

@ -1612,6 +1612,35 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
async def dispatch_failure_handlers(
self,
exception: Exception,
traceback_exception: str,
prefer_async_handlers: bool = False,
) -> None:
"""Route failure logging to async and/or sync handlers for this request.
Mirrors ``dispatch_success_handlers``: the sync ``failure_handler`` never runs
concurrently with ``async_failure_handler`` on the shared logging object, so the
two paths cannot mutate it at the same time. ``prefer_async_handlers`` only
bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from
``completion()``); legacy string callbacks still run via
``executor.submit(failure_handler)`` when configured.
"""
litellm_params = self.model_call_details.get("litellm_params", {}) or {}
sync_sdk = self._is_sync_litellm_request(litellm_params)
passthrough = self.call_type == CallTypes.pass_through.value
if sync_sdk and not prefer_async_handlers and not passthrough:
self.failure_handler(exception, traceback_exception)
return
await self.async_failure_handler(exception, traceback_exception)
if not self._should_run_sync_failure_callbacks_for_async_calls():
return
executor.submit(self.failure_handler, exception, traceback_exception)
def should_run_logging(
self,
event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"],
@ -3076,6 +3105,24 @@ class Logging(LiteLLMLoggingBaseClass):
_filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks)
return len(_filtered_success_callbacks) > 0
def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool:
"""
Returns:
- bool: True if sync failure callbacks should be run for async calls. eg. `langfuse`, `s3`
Mirrors ``_should_run_sync_callbacks_for_async_calls`` but reads the failure
callback lists. Gating the legacy sync ``failure_handler`` on the success lists
would drop sync failure callbacks for any caller that configures only failure
callbacks, so streaming errors would be logged nowhere.
"""
_combined_sync_callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_failure_callbacks,
global_callbacks=litellm.failure_callback,
)
_filtered_failure_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks)
_filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks)
return len(_filtered_failure_callbacks) > 0
def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List:
if dynamic_success_callbacks is None:
return list(global_callbacks)

7
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -457,7 +457,8 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
cache_creation_tokens = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0),
getattr(usage.prompt_tokens_details, "cache_write_tokens", 0)
or getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0),
)
or 0
)
@ -906,10 +907,6 @@ def get_token_type_cost_breakdown(
cache_read_tokens = prompt_tokens_details["cache_hit_tokens"]
cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"]
cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"]
# Some OpenAI-compatible providers (e.g. kimi-k2) report cache-write tokens
# under `cache_write_tokens`; mirror the total-cost normalization path.
if not cache_creation_tokens:
cache_creation_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "cache_write_tokens", 0))
# Fall back to the private top-level counters the Usage constructor mirrors cache
# tokens onto, so providers/callers that bypass prompt_tokens_details are covered.
if not cache_read_tokens:

16
litellm/litellm_core_utils/streaming_handler.py
View file

@ -2008,12 +2008,9 @@ class CustomStreamWrapper:
if self.logging_obj is not None:
self._record_partial_usage_for_failure()
## LOGGING
threading.Thread(
target=self.logging_obj.failure_handler,
args=(e, traceback_exception),
).start() # log response
# Handle any exceptions that might occur during streaming
asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception))
asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
)
self._handle_stream_fallback_error(e)
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
if self.received_finish_reason is None:
@ -2122,13 +2119,8 @@ class CustomStreamWrapper:
if self.logging_obj is not None:
self._record_partial_usage_for_failure()
## LOGGING
threading.Thread(
target=self.logging_obj.failure_handler,
args=(e, traceback_exception),
).start() # log response
# Handle any exceptions that might occur during streaming
asyncio.create_task(
self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
)
self._handle_stream_fallback_error(e)

3
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -123,7 +123,8 @@ class VertexAIBatchTransformation:
Gets the output file id from the Vertex AI Batch response
"""
output_file_id: str = response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "")
output_info = response.get("outputInfo") or OutputInfo()
output_file_id: str = output_info.get("gcsOutputDirectory", "")
if output_file_id:
output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl"
if output_file_id and output_file_id != "/predictions.jsonl":

54
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -3333,27 +3333,37 @@ class ModelResponseIterator:
return self.chunk_parser(chunk=json_chunk)
def handle_accumulated_json_chunk(self, chunk: str) -> Optional["ModelResponseStream"]:
chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
message = chunk.replace("\n\n", "")
def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]:
message = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip()
self.accumulated_json += message
# json.loads on the whole buffer after every fragment is O(n^2) and
# holds the GIL, freezing the event loop for seconds on large responses
# (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini
# chunk is a JSON object/array, so only attempt the parse once the
# buffer's last non-whitespace byte can close one.
stripped = self.accumulated_json.rstrip()
if not stripped or stripped[-1] not in "}]":
# Mid-stream, defer parsing until the buffer's last byte can close a value:
# attempting a parse after every fragment of one large object is O(n^2) and
# holds the GIL, freezing the event loop. At end of stream (is_final) no more
# data is coming, so drain whatever complete values remain regardless of the
# trailing byte, otherwise a complete leading value sitting behind a truncated
# trailing one would be silently dropped.
if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"):
return None
try:
_data = json.loads(self.accumulated_json)
self.accumulated_json = "" # reset after successful parsing
return self.chunk_parser(chunk=_data)
except json.JSONDecodeError:
return None
# Peel one complete JSON value from the front of the buffer and keep the
# unconsumed tail. Running json.loads over the whole buffer would fail
# forever once it held more than one concatenated value ("Extra data") while
# never resetting the buffer, so the buffer grew without bound and pinned the
# core. raw_decode reports where the value ended, so concatenated values drain
# one call at a time. A leading non-dict value (never emitted by Gemini in
# practice) is consumed and skipped so it cannot block the dict values behind it.
decoder = json.JSONDecoder()
while self.accumulated_json:
try:
raw_value = decoder.raw_decode(self.accumulated_json)
except json.JSONDecodeError:
return None
decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int]
self.accumulated_json = self.accumulated_json[end_index:].strip()
if isinstance(decoded, dict):
return self.chunk_parser(chunk=decoded)
return None
def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]:
try:
@ -3378,7 +3388,9 @@ class ModelResponseIterator:
chunk = self.response_iterator.__next__()
except StopIteration:
if self.chunk_type == "accumulated_json" and self.accumulated_json:
return self.handle_accumulated_json_chunk(chunk="")
result = self.handle_accumulated_json_chunk(chunk="", is_final=True)
if result is not None:
return result
raise StopIteration
except ValueError as e:
raise RuntimeError(f"Error receiving chunk from stream: {e}")
@ -3400,7 +3412,9 @@ class ModelResponseIterator:
chunk = await self.async_response_iterator.__anext__()
except StopAsyncIteration:
if self.chunk_type == "accumulated_json" and self.accumulated_json:
return self.handle_accumulated_json_chunk(chunk="")
result = self.handle_accumulated_json_chunk(chunk="", is_final=True)
if result is not None:
return result
raise StopAsyncIteration
except ValueError as e:
raise RuntimeError(f"Error receiving chunk from stream: {e}")

350
litellm/model_prices_and_context_window_backup.json
View file

@ -1502,6 +1502,222 @@
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 1024
},
"anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"global.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"us.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"eu.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"au.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"jp.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"anthropic.claude-opus-4-8": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
@ -2756,6 +2972,38 @@
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true
},
"azure_ai/claude-opus-5": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2.5e-05,
"litellm_provider": "azure_ai",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"prompt_cache_min_tokens": 512
},
"azure_ai/claude-opus-4-8": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
@ -11846,6 +12094,44 @@
"supports_output_config": true,
"prompt_cache_min_tokens": 512
},
"claude-opus-5": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "anthropic",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"provider_specific_entry": {
"us": 1.1,
"fast": 2.0
},
"supports_output_config": true,
"supports_speed": true,
"prompt_cache_min_tokens": 512
},
"claude-opus-4-8": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
@ -36896,6 +37182,70 @@
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true
},
"vertex_ai/claude-opus-5": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"prompt_cache_min_tokens": 512
},
"vertex_ai/claude-opus-5@default": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"prompt_cache_min_tokens": 512
},
"vertex_ai/claude-opus-4-8": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,

9
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
View file

@ -1,9 +1,12 @@
from typing import Dict, List, Optional
from typing import TYPE_CHECKING, Dict, List, Optional
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from litellm.proxy._types import UserAPIKeyAuth
if TYPE_CHECKING:
from opentelemetry.trace import SpanContext
class MCPAuthenticatedUser(AuthenticatedUser):
"""
@ -16,6 +19,8 @@ class MCPAuthenticatedUser(AuthenticatedUser):
4. Server-specific authentication headers
5. OAuth2 headers
6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin.
7. Transport span context - the tracing span of the HTTP request carrying the current
message, which a stateful session's message handler cannot read from its own task.
"""
def __init__(
@ -28,6 +33,7 @@ class MCPAuthenticatedUser(AuthenticatedUser):
mcp_protocol_version: Optional[str] = None,
raw_headers: Optional[Dict[str, str]] = None,
client_ip: Optional[str] = None,
transport_span_context: Optional["SpanContext"] = None,
):
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header
@ -37,3 +43,4 @@ class MCPAuthenticatedUser(AuthenticatedUser):
self.oauth2_headers = oauth2_headers
self.raw_headers = raw_headers
self.client_ip = client_ip
self.transport_span_context = transport_span_context

192
litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
View file

@ -1,192 +0,0 @@
"""
OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers.
Exchanges a user's incoming JWT (subject_token) for a scoped access token
at an IDP's token exchange endpoint. The exchanged token is then used to
authenticate requests to the upstream MCP server.
See: https://datatracker.ietf.org/doc/html/rfc8693
"""
import asyncio
import hashlib
import weakref
from typing import TYPE_CHECKING, Dict, Tuple
import httpx
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
build_token_endpoint_client_auth,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
# RFC 8693 grant type constant
TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
class TokenExchangeHandler:
"""Handles OAuth 2.0 Token Exchange (RFC 8693) for MCP servers.
Caches exchanged tokens keyed by ``hash(subject_token + server_id)`` so
repeated calls with the same user token skip the IDP round-trip.
"""
def __init__(self) -> None:
self._cache = InMemoryCache(
max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
)
# WeakValueDictionary so locks are GC'd once no coroutine holds a reference,
# preventing unbounded growth with many rotating user tokens.
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
def _get_lock(self, cache_key: str) -> asyncio.Lock:
lock = self._locks.get(cache_key)
if lock is None:
lock = asyncio.Lock()
self._locks[cache_key] = lock
return lock
@staticmethod
def _cache_key(subject_token: str, server_id: str) -> str:
raw = f"{subject_token}:{server_id}"
return hashlib.sha256(raw.encode()).hexdigest()
async def exchange_token(
self,
subject_token: str,
server: "MCPServer",
) -> str:
"""Exchange *subject_token* for a scoped access token.
Returns the exchanged ``access_token`` string (suitable for a
``Bearer`` header).
Raises ``ValueError`` on configuration or IDP errors.
"""
cache_key = self._cache_key(subject_token, server.server_id)
# Fast path
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
# Slow path — one exchange at a time per (user, server) pair
async with self._get_lock(cache_key):
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
token, ttl = await self._do_exchange(subject_token, server)
self._cache.set_cache(cache_key, token, ttl=ttl)
return token
async def _do_exchange(
self,
subject_token: str,
server: "MCPServer",
) -> Tuple[str, int]:
"""POST to the token exchange endpoint with RFC 8693 parameters.
Returns ``(access_token, ttl_seconds)``.
"""
endpoint = server.token_exchange_endpoint or server.token_url
if not endpoint:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but no token_exchange_endpoint or token_url configured"
)
if not server.client_id or not server.client_secret:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but missing client_id or client_secret"
)
client_auth = build_token_endpoint_client_auth(
auth_method=server.token_endpoint_auth_method,
client_id=server.client_id,
client_secret=server.client_secret,
)
data: Dict[str, str] = {
"grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": subject_token,
"subject_token_type": server.subject_token_type or DEFAULT_SUBJECT_TOKEN_TYPE,
**client_auth.body,
}
if server.audience:
data["audience"] = server.audience
if server.scopes:
data["scope"] = " ".join(server.scopes)
verbose_logger.debug(
"Exchanging token for MCP server %s at %s (audience=%s)",
server.server_id,
endpoint,
server.audience,
)
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
post_kwargs = {"data": data, **({"headers": client_auth.headers} if client_auth.headers else {})}
try:
response = await client.post(endpoint, **post_kwargs)
response.raise_for_status()
except httpx.HTTPStatusError as exc:
verbose_logger.debug(
"Token exchange IDP error for MCP server %s (status %d)",
server.server_id,
exc.response.status_code,
)
raise ValueError(
f"Token exchange for MCP server '{server.server_id}' failed with status {exc.response.status_code}"
) from exc
body = response.json()
if not isinstance(body, dict):
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"returned non-object JSON (got {type(body).__name__})"
)
access_token = body.get("access_token")
if not access_token:
raise ValueError(f"Token exchange response for MCP server '{server.server_id}' missing 'access_token'")
raw_expires_in = body.get("expires_in")
try:
expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
except (TypeError, ValueError):
expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
ttl = max(
expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
)
verbose_logger.info(
"Token exchange succeeded for MCP server %s (expires in %ds)",
server.server_id,
expires_in,
)
return access_token, ttl
def invalidate(self, subject_token: str, server_id: str) -> None:
"""Remove a cached exchanged token (e.g. after a 401)."""
cache_key = self._cache_key(subject_token, server_id)
self._cache.delete_cache(cache_key)
# Module-level singleton
mcp_token_exchange_handler = TokenExchangeHandler()

4
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -3086,9 +3086,7 @@ class MCPServerManager:
)
):
spec = None
auth_value = (
await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) if spec is None else None
)
auth_value = await resolve_mcp_auth(server, mcp_auth_header) if spec is None else None
# Create sampling and elicitation callbacks for this client
sampling_cb = _create_sampling_callback(user_api_key_auth=user_api_key_auth) if server.allow_sampling else None

34
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
View file

@ -26,7 +26,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy._experimental.mcp_server.auth import token_exchange
from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
build_token_endpoint_client_auth,
)
@ -58,17 +57,12 @@ class MCPOAuth2TokenCache(InMemoryCache):
def _has_client_credentials_config(server: "MCPServer") -> bool:
return bool(server.client_id and server.client_secret and server.token_url)
async def async_get_token(
self,
server: "MCPServer",
*,
require_client_credentials_flow: bool = True,
) -> Optional[str]:
async def async_get_token(self, server: "MCPServer") -> Optional[str]:
"""Return a valid access token, fetching or refreshing as needed.
Returns ``None`` when the server lacks client credentials config.
"""
if require_client_credentials_flow and not server.has_client_credentials:
if not server.has_client_credentials:
return None
if not self._has_client_credentials_config(server):
return None
@ -278,36 +272,16 @@ mcp_per_user_token_cache = MCPPerUserTokenCache()
async def resolve_mcp_auth(
server: "MCPServer",
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
subject_token: Optional[str] = None,
) -> Optional[Union[str, Dict[str, str]]]:
"""Resolve the auth value for an MCP server.
Priority:
1. ``mcp_auth_header`` — per-request/per-user override
2. OAuth2 Token Exchange (OBO / RFC 8693) — exchange user token for scoped token
3. OAuth2 client_credentials token — auto-fetched and cached
4. ``server.authentication_token`` — static token from config/DB
2. OAuth2 client_credentials token — auto-fetched and cached
3. ``server.authentication_token`` — static token from config/DB
"""
if mcp_auth_header:
return mcp_auth_header
if server.has_token_exchange_config:
if subject_token:
return await token_exchange.mcp_token_exchange_handler.exchange_token(subject_token, server)
# No subject_token — fall back to client_credentials using the same client
# credentials and token_url so M2M scenarios still work.
if server.client_id and server.client_secret and server.token_url:
return await mcp_oauth2_token_cache.async_get_token(
server,
require_client_credentials_flow=False,
)
# OBO configured but no subject_token and missing client credentials — warn
# rather than silently proceeding unauthenticated.
verbose_logger.warning(
"MCP server '%s' is configured for token exchange (OBO) but no subject_token "
"was provided and client credentials (client_id/client_secret/token_url) are "
"incomplete. The request will proceed without authentication.",
server.server_id,
)
if server.has_client_credentials:
return await mcp_oauth2_token_cache.async_get_token(server)
return server.authentication_token

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

@ -230,16 +230,33 @@ if MCP_AVAILABLE:
return server_auth
return mcp_auth_header
def _get_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]:
"""Return the subset of *allowed_server_ids* whose servers use OAuth2 auth.
def _is_v1_resolved_oauth2_server(server: Optional[MCPServer]) -> bool:
"""Whether this server's per-user OAuth2 token is still resolved by v1.
Used as a cheap pre-flight check to skip bulk credential fetching when no
OAuth2 servers are involved in the current request.
A server the v2 resolver owns reads its stored token from the resolver at connect
time and drops any Authorization built for it here, so the v1 lookup would be a DB
round-trip whose result is discarded. Mirrors the same guard on the protocol listing
path and in ``_resolve_oauth2_headers_for_tool_call``.
"""
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
to_server_spec,
)
if getattr(server, "auth_type", None) != MCPAuth.oauth2:
return False
return to_server_spec(server) is None
def _v1_resolved_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]:
"""Return the subset of *allowed_server_ids* whose per-user OAuth2 token is still
resolved by v1.
Used as a cheap pre-flight check to skip bulk credential fetching when no such
server is involved in the current request.
"""
return {
sid
for sid in allowed_server_ids
if getattr(global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None) == MCPAuth.oauth2
if _is_v1_resolved_oauth2_server(global_mcp_server_manager.get_mcp_server_by_id(sid))
}
async def _get_user_oauth_extra_headers(
@ -253,11 +270,13 @@ if MCP_AVAILABLE:
the MCP server the same way the admin "Add MCP / Authorize and Fetch" flow does.
Returns None for non-OAuth2 servers or when no credential is stored.
A server the v2 resolver owns is skipped; see ``_is_v1_resolved_oauth2_server``.
Args:
prefetched_creds: Optional dict keyed by server_id with credential payloads.
When provided, avoids a per-server DB round-trip.
"""
if getattr(server, "auth_type", None) != MCPAuth.oauth2:
if not _is_v1_resolved_oauth2_server(server):
return None
user_id = getattr(user_api_key_dict, "user_id", None)
server_id = getattr(server, "server_id", None)
@ -320,38 +339,6 @@ if MCP_AVAILABLE:
verbose_logger.warning(f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}")
return {}
async def _get_bulk_user_oauth_headers(
user_api_key_dict: UserAPIKeyAuth,
) -> Dict[str, Dict[str, str]]:
"""
Fetch ALL OAuth2 credentials for the current user in a single DB query and
return a mapping of server_id → {"Authorization": "Bearer <token>"}.
This is the batch alternative to calling _get_user_oauth_extra_headers
per-server inside a loop (N+1 DB queries).
"""
user_id = getattr(user_api_key_dict, "user_id", None)
if not user_id:
return {}
try:
from litellm.proxy._experimental.mcp_server.db import (
list_user_oauth_credentials,
)
from litellm.proxy.utils import get_prisma_client_or_throw
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
creds = await list_user_oauth_credentials(prisma_client, user_id)
return {
c["server_id"]: {"Authorization": f"Bearer {c['access_token']}"}
for c in creds
if c.get("access_token") and c.get("server_id")
}
except Exception:
verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True)
return {}
def _create_tool_response_objects(tools, server: MCPServer):
"""Helper function to create tool response objects.
@ -825,7 +812,7 @@ if MCP_AVAILABLE:
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
prefetched_oauth_creds = (
await _prefetch_user_oauth_creds(user_api_key_dict)
if _get_oauth2_server_ids(allowed_server_ids)
if _v1_resolved_oauth2_server_ids(allowed_server_ids)
else {}
)

92
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -15,6 +15,7 @@ import types
import uuid
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Callable,
@ -107,6 +108,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100
# arbitrarily large body just to make a routing decision.
_MCP_ROUTING_PEEK_MAX_BYTES = 4096
if TYPE_CHECKING:
from opentelemetry.trace import SpanContext
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
"""Remove a (user_id, server_id) entry from the BYOK credential cache.
@ -242,10 +246,12 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> Optional[dict[str, str]]:
"""The W3C trace context (``traceparent``/``tracestate``) the MCP client
propagated in the request's ``params._meta`` (SEP-414), or ``None``.
Per the OTel MCP semconv the MCP span parents to this propagated context rather
than to the HTTP/session transport (which is recorded as a link instead), so a
streamable-HTTP session that multiplexes many messages does not glue every
message under the session's first request. The client's W3C Baggage is
When present, per the OTel MCP semconv the MCP span parents to this propagated
context rather than to the HTTP transport (which is recorded as a link instead).
When absent, the span nests under the transport span of the request carrying
this specific message, so a streamable-HTTP session that multiplexes many
messages still does not glue every message under the session's first request;
see ``resolve_mcp_span_context``. The client's W3C Baggage is
deliberately excluded: it is caller-controlled, and the otel baggage processor
stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``,
...) onto the span, so honoring remote baggage would let a client spoof a
@ -288,6 +294,56 @@ def _otel_reset_mcp_trace_carrier(token: object) -> None:
return
def _otel_request_transport_span_context() -> Optional["SpanContext"]:
"""The tracing span of the HTTP request being handled, as a portable value.
Resolved on the ASGI request task, where the proxy's server span is anchored,
and carried to the MCP message handler on the authenticated-user object. A
stateful streamable-HTTP session handles every message on the task spawned by
its ``initialize`` POST, so the handler's own task cannot see later requests'
spans; this is the same reason per-request auth is carried across rather than
read from a ContextVar. Lazily imported so opentelemetry stays an optional
dependency; returns ``None`` when otel_v2 is unavailable or no request span is
anchored."""
try:
from litellm.integrations.otel.plumbing.context import (
request_root_span_context,
)
return request_root_span_context()
except ImportError:
return None
def _otel_set_mcp_transport_span_context(span_context: Optional["SpanContext"]) -> object:
"""Publish the current message's transport span for the otel_v2 MCP span and
return a reset token, or ``None`` when otel_v2 is unavailable."""
if span_context is None:
return None
try:
from litellm.integrations.otel.plumbing.context import (
set_mcp_message_transport_span_context,
)
return set_mcp_message_transport_span_context(span_context)
except ImportError:
return None
def _otel_reset_mcp_transport_span_context(token: object) -> None:
"""Paired with ``_otel_set_mcp_transport_span_context``."""
if token is None:
return
try:
from litellm.integrations.otel.plumbing.context import (
reset_mcp_message_transport_span_context,
)
reset_mcp_message_transport_span_context(token)
except ImportError:
return
def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException:
"""Map a ``ProxyException`` to an ``HTTPException`` that preserves its real
status code and headers.
@ -654,6 +710,18 @@ if MCP_AVAILABLE:
############### MCP Server Routes #######################
########################################################
def _current_transport_span_context() -> Optional["SpanContext"]:
"""The transport span of the HTTP request carrying the message being handled.
Published by the ASGI request task onto the authenticated-user object, because
a stateful session's message handler runs on the task spawned by that session's
``initialize`` POST and so cannot read later requests' spans from its own task.
"""
auth_user = auth_context_var.get()
if not isinstance(auth_user, MCPAuthenticatedUser):
auth_user = _recover_auth_from_session()
return auth_user.transport_span_context if auth_user is not None else None
@server.list_tools()
async def handle_list_tools() -> "ListToolsResult | List[Tool]":
"""
@ -670,9 +738,11 @@ if MCP_AVAILABLE:
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_trace_token = None
_transport_token = None
try:
_trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx))
_transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context())
# Get user authentication from context variable
(
user_api_key_auth,
@ -728,6 +798,7 @@ if MCP_AVAILABLE:
# This prevents the HTTP stream from failing and allows the client to get a response
return []
finally:
_otel_reset_mcp_transport_span_context(_transport_token)
_otel_reset_mcp_trace_carrier(_trace_token)
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
@ -901,9 +972,11 @@ if MCP_AVAILABLE:
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_trace_token = None
_transport_token = None
try:
_trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx))
_transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context())
# Validate arguments
(
user_api_key_auth,
@ -1042,6 +1115,7 @@ if MCP_AVAILABLE:
return response
finally:
_otel_reset_mcp_transport_span_context(_transport_token)
_otel_reset_mcp_trace_carrier(_trace_token)
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
@ -4197,6 +4271,7 @@ if MCP_AVAILABLE:
session_id=session_id if use_stateful else None,
touch_last_seen=(scope.get("method") or "").upper() != "DELETE",
copy_existing_session_auth_context=is_initialize,
transport_span_context=_otel_request_transport_span_context(),
)
local_send = send
if use_stateful and is_initialize:
@ -4421,6 +4496,7 @@ if MCP_AVAILABLE:
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
client_ip: Optional[str] = None,
transport_span_context: Optional["SpanContext"] = None,
) -> None:
auth_user.user_api_key_auth = user_api_key_auth
auth_user.mcp_auth_header = mcp_auth_header
@ -4429,6 +4505,7 @@ if MCP_AVAILABLE:
auth_user.oauth2_headers = oauth2_headers
auth_user.raw_headers = raw_headers
auth_user.client_ip = client_ip
auth_user.transport_span_context = transport_span_context
def set_auth_context(
user_api_key_auth: Optional[UserAPIKeyAuth],
@ -4438,6 +4515,7 @@ if MCP_AVAILABLE:
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
client_ip: Optional[str] = None,
transport_span_context: Optional["SpanContext"] = None,
) -> MCPAuthenticatedUser:
"""
Set the UserAPIKeyAuth in the auth context variable.
@ -4448,6 +4526,7 @@ if MCP_AVAILABLE:
mcp_servers: Optional list of server names and access groups to filter by
mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value}
client_ip: Client IP address for MCP access control
transport_span_context: Tracing span of the HTTP request carrying this message
"""
auth_user = MCPAuthenticatedUser(
user_api_key_auth=user_api_key_auth,
@ -4457,6 +4536,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=client_ip,
transport_span_context=transport_span_context,
)
auth_context_var.set(auth_user)
return auth_user
@ -4472,6 +4552,7 @@ if MCP_AVAILABLE:
session_id: Optional[str] = None,
touch_last_seen: bool = True,
copy_existing_session_auth_context: bool = False,
transport_span_context: Optional["SpanContext"] = None,
) -> MCPAuthenticatedUser:
auth_user = _stateful_session_auth_contexts.get(session_id) if session_id else None
if auth_user is not None and session_id is not None:
@ -4486,6 +4567,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=client_ip,
transport_span_context=transport_span_context,
)
_update_auth_context(
auth_user=auth_user,
@ -4496,6 +4578,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=client_ip,
transport_span_context=transport_span_context,
)
auth_context_var.set(auth_user)
return auth_user
@ -4507,6 +4590,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=client_ip,
transport_span_context=transport_span_context,
)
def _wrap_send_with_stateful_session_auth_context(

189
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -26762,6 +26762,113 @@
"title": "ToolPolicyUpdateResponse",
"type": "object"
},
"ToolSpendDailyEntry": {
"description": "Spend attributed to one tool on one UTC day.",
"properties": {
"call_count": {
"default": 0,
"title": "Call Count",
"type": "integer"
},
"date": {
"title": "Date",
"type": "string"
},
"spend": {
"default": 0.0,
"title": "Spend",
"type": "number"
},
"tool_name": {
"title": "Tool Name",
"type": "string"
}
},
"required": [
"date",
"tool_name"
],
"title": "ToolSpendDailyEntry",
"type": "object"
},
"ToolSpendEntry": {
"description": "Total spend attributed to one tool over the requested window.",
"properties": {
"call_count": {
"default": 0,
"title": "Call Count",
"type": "integer"
},
"spend": {
"default": 0.0,
"description": "Attributed spend: a request that used several tools counts its full spend toward each of them",
"title": "Spend",
"type": "number"
},
"tool_name": {
"title": "Tool Name",
"type": "string"
},
"total_tokens": {
"default": 0,
"title": "Total Tokens",
"type": "integer"
}
},
"required": [
"tool_name"
],
"title": "ToolSpendEntry",
"type": "object"
},
"ToolSpendResponse": {
"properties": {
"by_tool": {
"items": {
"$ref": "#/components/schemas/ToolSpendEntry"
},
"title": "By Tool",
"type": "array"
},
"daily": {
"items": {
"$ref": "#/components/schemas/ToolSpendDailyEntry"
},
"title": "Daily",
"type": "array"
},
"end_date": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "End Date"
},
"start_date": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Start Date"
},
"total_spend": {
"default": 0.0,
"description": "Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist",
"title": "Total Spend",
"type": "number"
}
},
"title": "ToolSpendResponse",
"type": "object"
},
"ToolUsageLogEntry": {
"description": "One spend log row for a tool call (for UI \"recent logs\" table).",
"properties": {
@ -26858,6 +26965,13 @@
},
"ValidationError": {
"properties": {
"ctx": {
"title": "Context",
"type": "object"
},
"input": {
"title": "Input"
},
"loc": {
"items": {
"anyOf": [
@ -27301,6 +27415,81 @@
]
}
},
"/v1/tool/spend": {
"get": {
"description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nJoins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to\n``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools\ncounts its full spend toward each of those tools, so per-tool numbers are\nattributions. ``total_spend`` is the deduplicated spend of every request that\ncalled at least one tool in the window, so it never double counts.",
"operationId": "get_tool_spend_v1_tool_spend_get",
"parameters": [
{
"description": "YYYY-MM-DD (defaults to 30 days ago)",
"in": "query",
"name": "start_date",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "YYYY-MM-DD (defaults to 30 days ago)",
"title": "Start Date"
}
},
{
"description": "YYYY-MM-DD (defaults to today)",
"in": "query",
"name": "end_date",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "YYYY-MM-DD (defaults to today)",
"title": "End Date"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ToolSpendResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Get Tool Spend",
"tags": [
"tools"
]
}
},
"/v1/tool/{tool_name}": {
"get": {
"description": "Get details for a single tool.",

12
litellm/proxy/_types.py
View file

@ -2616,6 +2616,17 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
# key off. Server-only and stripped from validated input for the same reason as the marker
# above: a forged entry would let a caller pick which team's rpm bucket it is charged against.
mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True)
via_virtual_key: bool = Field(
default=False,
exclude=True,
description=(
"Server-only marker set exclusively by the DB virtual-key and master-key auth paths via "
"post-construction assignment. Stripped from validated input so custom auth handlers, JWT "
"claims, or key metadata cannot forge it. Gates overwrite_user_with_key_hash stamping: only "
"a credential the proxy itself validated as a key may be forwarded as the provider-facing "
"user id."
),
)
budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True)
budget_throttle_pct: Optional[float] = Field(default=None, exclude=True)
user: Optional[Any] = None # Expanded user object when expand=user is used
@ -2641,6 +2652,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
# kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data.
values.pop("mcp_admitted_user_subject", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("via_virtual_key", None)
if values.get("api_key") is not None:
values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))})
if isinstance(values.get("api_key"), str):

13
litellm/proxy/auth/user_api_key_auth.py
View file

@ -1497,6 +1497,13 @@ async def _user_api_key_auth_builder(
check_cache_only=True,
).resolve(hashed_token=hash_token(api_key))
)
# Key-cache entries are written only after the proxy validated a
# virtual key or the master key, but via_virtual_key is exclude=True
# so serialization drops it; restore it at this trusted boundary.
# The UI-login JWT fallback below constructs its token from a
# decrypted blob, not this cache, and stays unmarked.
if isinstance(valid_token, UserAPIKeyAuth):
valid_token.via_virtual_key = True
except Exception:
verbose_logger.debug("api key not found in cache.")
valid_token = None
@ -1614,6 +1621,7 @@ async def _user_api_key_auth_builder(
_user_api_key_obj = update_valid_token_with_end_user_params(
valid_token=_user_api_key_obj, end_user_params=end_user_params
)
_user_api_key_obj.via_virtual_key = True
return _user_api_key_obj
@ -2021,7 +2029,7 @@ async def _user_api_key_auth_builder(
# No token was found when looking up in the DB
raise Exception("Invalid proxy server token passed")
if valid_token_dict is not None:
return await _return_user_api_key_auth_obj(
virtual_key_auth_obj = await _return_user_api_key_auth_obj(
user_obj=user_obj,
api_key=api_key,
parent_otel_span=parent_otel_span,
@ -2029,6 +2037,8 @@ async def _user_api_key_auth_builder(
route=route,
start_time=start_time,
)
virtual_key_auth_obj.via_virtual_key = True
return virtual_key_auth_obj
except Exception as e:
return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
e=e,
@ -2442,6 +2452,7 @@ async def _reserve_budget_after_common_checks(
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
)

4
litellm/proxy/config_resolvers/sso.py
View file

@ -36,6 +36,10 @@ SSO_DESCRIPTORS: tuple[FieldDescriptor, ...] = (
FieldDescriptor("generic_token_endpoint", "generic_token_endpoint", "GENERIC_TOKEN_ENDPOINT"),
FieldDescriptor("generic_userinfo_endpoint", "generic_userinfo_endpoint", "GENERIC_USERINFO_ENDPOINT"),
FieldDescriptor("generic_scope", "generic_scope", "GENERIC_SCOPE", default="openid email profile"),
FieldDescriptor("saml_idp_metadata_url", "saml_idp_metadata_url", "SAML_IDP_METADATA_URL"),
FieldDescriptor("saml_idp_metadata_xml", "saml_idp_metadata_xml", "SAML_IDP_METADATA_XML"),
FieldDescriptor("saml_sp_entity_id", "saml_sp_entity_id", "SAML_SP_ENTITY_ID"),
FieldDescriptor("saml_allow_unsolicited", "saml_allow_unsolicited", "SAML_ALLOW_UNSOLICITED"),
FieldDescriptor("proxy_base_url", "proxy_base_url", "PROXY_BASE_URL"),
)

2
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -432,7 +432,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
Override to store only the Model Armor API response, not the entire data dict.
This prevents circular references in logging.
"""
metadata = request_data.get("metadata", {}) if isinstance(request_data, dict) else {}
metadata = (request_data.get("metadata") or {}) if isinstance(request_data, dict) else {}
guardrail_response = metadata.get("_model_armor_response", {})
# Determine status – default to "success" but prefer the explicit value if present.

3
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -489,6 +489,9 @@ class InMemoryGuardrailHandler:
"skip_tool_message_in_guardrail",
getattr(litellm_params, "skip_tool_message_in_guardrail", None),
)
configured_run_in_parallel = getattr(litellm_params, "run_in_parallel", None)
if configured_run_in_parallel is not None:
custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel)
parsed_guardrail = Guardrail(
guardrail_id=guardrail.get("guardrail_id"),

25
litellm/proxy/litellm_pre_call_utils.py
View file

@ -13,7 +13,7 @@ from starlette.datastructures import Headers
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
iter_client_callback_metadata_dicts,
@ -48,6 +48,24 @@ _EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session-
# Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores
# (covers UUIDs and most common session-id formats).
_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
_SHA256_HEX_RE = re.compile(r"^[0-9a-f]{64}$")
def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None:
"""Only proxy-validated keys are stamped, proven by the unforgeable
via_virtual_key marker AND a known non-secret shape: the sha256 hex digest
UserAPIKeyAuth stores virtual keys in, or the master key's stable alias.
Custom-auth credentials arrive raw (never forward auth material) and hashed
JWTs rotate on re-issue (useless as a stable ban id), so both are skipped."""
api_key = user_api_key_dict.api_key
if not user_api_key_dict.via_virtual_key or api_key is None:
return None
if api_key == LITELLM_PROXY_MASTER_KEY_ALIAS or _SHA256_HEX_RE.fullmatch(api_key):
return api_key
return None
_ANTHROPIC_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]+$")
@ -1447,6 +1465,11 @@ async def add_litellm_data_to_request(
if "user" not in data:
data["user"] = user
if litellm.overwrite_user_with_key_hash is True:
stampable_hash = _stampable_key_hash(user_api_key_dict)
if stampable_hash is not None:
data["user"] = stampable_hash
data["secret_fields"] = SecretFields(raw_headers=_raw_headers)
## Dynamic api version (Azure OpenAI endpoints) ##

493
litellm/proxy/management_endpoints/sso/saml_sso.py Normal file
View file

@ -0,0 +1,493 @@
"""
SAML 2.0 SSO for the LiteLLM proxy admin UI.
Supports both SP-initiated and IdP-initiated login via the HTTP-POST binding,
using the OneLogin python3-saml toolkit for signature, audience and time
validation. The IdP is configured from its metadata (``SAML_IDP_METADATA_URL``
or inline ``SAML_IDP_METADATA_XML``); a successful login is mapped to a
``CustomOpenID`` and handed to the shared post-login path used by every other
SSO provider.
python3-saml pulls in the native ``xmlsec``/``libxml2`` libraries, so it is an
optional dependency. When it is not installed the SAML routes return a clear
error instead of breaking proxy startup.
"""
# python3-saml ships no type stubs, so the type checker sees every onelogin call
# as Unknown and the guarded optional import as possibly-unbound. Values crossing
# that boundary are cast() to concrete types at each use site; these directives
# silence only the unavoidable noise from the untyped dependency in this module.
# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
# pyright: reportUnknownArgumentType=false, reportUnknownParameterType=false
# pyright: reportMissingTypeStubs=false, reportPossiblyUnboundVariable=false
# pyright: reportConstantRedefinition=false
import asyncio
import hashlib
import os
import secrets
import time
from typing import cast
from urllib.parse import parse_qsl
from fastapi import HTTPException, Request, status
from fastapi.responses import RedirectResponse
from pydantic import ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
from litellm.proxy.utils import get_custom_url
try:
from onelogin.saml2.auth import OneLogin_Saml2_Auth
from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser
from onelogin.saml2.settings import OneLogin_Saml2_Settings
from onelogin.saml2.xml_utils import OneLogin_Saml2_XML
SAML_AVAILABLE = True
except ImportError:
SAML_AVAILABLE = False
SAML_LOGIN_ROUTE = "sso/saml/login"
SAML_CALLBACK_ROUTE = "sso/saml/callback"
SAML_METADATA_ROUTE = "sso/saml/metadata"
_SAML_AUTHN_STATE_COOKIE = "litellm_saml_authn"
_SAML_IDP_SETTINGS_CACHE_PREFIX = "saml_idp_settings"
_SAML_AUTHN_REQUEST_CACHE_PREFIX = "saml_authn_request"
_SAML_CONSUMED_ASSERTION_CACHE_PREFIX = "saml_consumed_assertion"
_SAML_AUTHN_REQUEST_TTL_SECONDS = 600
_SAML_IDP_METADATA_TTL_SECONDS = 3600
_SAML_METADATA_FETCH_TIMEOUT_SECONDS = 10
_SAML_MAX_POST_BYTES = 5 * 1024 * 1024
# The replay guard tracks each assertion's NotOnOrAfter so it spans the full
# validity window; the floor covers IdPs that issue hour-long assertions or omit
# the timestamp, and the cap bounds cache growth.
_SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS = 3600
_SAML_REPLAY_GUARD_MAX_TTL_SECONDS = 86400
_EMAIL_ATTRIBUTE_CANDIDATES = (
"urn:oid:0.9.2342.19200300.100.1.3",
"http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress",
"email",
"emailAddress",
"mail",
"Email",
)
_FIRST_NAME_ATTRIBUTE_CANDIDATES = (
"urn:oid:2.5.4.42",
"http://schemas.xmlsoap.org/ws/2005/05/identity/claims/givenname",
"givenName",
"first_name",
"firstName",
)
_LAST_NAME_ATTRIBUTE_CANDIDATES = (
"urn:oid:2.5.4.4",
"http://schemas.xmlsoap.org/ws/2005/05/identity/claims/surname",
"sn",
"surname",
"last_name",
"lastName",
)
_ROLE_ATTRIBUTE_CANDIDATES = ("role", "roles", "litellm_role")
_TEAM_IDS_ATTRIBUTE_CANDIDATES = ("teams", "team_ids", "groups")
def _saml_unavailable_error() -> HTTPException:
return HTTPException(
status_code=status.HTTP_501_NOT_IMPLEMENTED,
detail=(
"SAML SSO requires the optional 'python3-saml' dependency, which is "
"not installed. Re-install litellm with the saml extra: "
"'pip install litellm[saml]'. The saml extra bundles the native "
"xmlsec/libxml2 libraries, so no system packages are required."
),
)
class SAMLAuthHandler:
"""SP- and IdP-initiated SAML 2.0 login for the admin UI."""
@staticmethod
def _env(name: str, default: str | None = None) -> str | None:
return os.getenv(name, default)
@staticmethod
def is_saml_configured() -> bool:
return bool(SAMLAuthHandler._env("SAML_IDP_METADATA_URL") or SAMLAuthHandler._env("SAML_IDP_METADATA_XML"))
@staticmethod
def _bool_env(name: str, default: bool) -> bool:
raw = SAMLAuthHandler._env(name)
if raw is None:
return default
return raw.strip().lower() in ("true", "1", "yes", "on")
@staticmethod
def _base_url(request: Request) -> str:
base = get_custom_url(request_base_url=str(request.base_url))
return base if base.endswith("/") else base + "/"
@staticmethod
def _is_https(request: Request) -> bool:
return SAMLAuthHandler._base_url(request).startswith("https")
@staticmethod
def _acs_url(request: Request) -> str:
return SAMLAuthHandler._base_url(request) + SAML_CALLBACK_ROUTE
@staticmethod
def _metadata_url(request: Request) -> str:
return SAMLAuthHandler._base_url(request) + SAML_METADATA_ROUTE
@staticmethod
def _sp_entity_id(request: Request) -> str:
return SAMLAuthHandler._env("SAML_SP_ENTITY_ID") or SAMLAuthHandler._metadata_url(request)
@staticmethod
async def _load_idp_settings(cache: DualCache) -> dict[str, object]:
metadata_url = SAMLAuthHandler._env("SAML_IDP_METADATA_URL")
metadata_xml = SAMLAuthHandler._env("SAML_IDP_METADATA_XML")
source = metadata_url or metadata_xml
if source is None:
raise HTTPException(
status_code=status.HTTP_501_NOT_IMPLEMENTED,
detail="SAML SSO is not configured. Set SAML_IDP_METADATA_URL or SAML_IDP_METADATA_XML.",
)
cache_key = f"{_SAML_IDP_SETTINGS_CACHE_PREFIX}:{hashlib.sha256(source.encode()).hexdigest()}"
cached = cache.get_cache(key=cache_key)
if isinstance(cached, dict):
return cast(dict[str, object], cached) # cast-ok: untyped python3-saml
if metadata_url is not None:
parsed = await asyncio.to_thread(
OneLogin_Saml2_IdPMetadataParser.parse_remote,
metadata_url,
validate_cert=SAMLAuthHandler._bool_env("SAML_IDP_METADATA_VALIDATE_CERT", True),
timeout=_SAML_METADATA_FETCH_TIMEOUT_SECONDS,
)
else:
parsed = OneLogin_Saml2_IdPMetadataParser.parse(cast(str, metadata_xml)) # cast-ok: untyped python3-saml
idp_settings = cast(dict[str, object], parsed) # cast-ok: untyped python3-saml
if not idp_settings.get("idp"):
raise HTTPException(
status_code=status.HTTP_502_BAD_GATEWAY,
detail="Could not parse an IdP entityID/SSO URL/certificate from the SAML metadata.",
)
cache.set_cache(key=cache_key, value=idp_settings, ttl=_SAML_IDP_METADATA_TTL_SECONDS)
return idp_settings
@staticmethod
def _build_settings(request: Request, idp_settings: dict[str, object]) -> dict[str, object]:
sp_settings: dict[str, object] = {
"strict": SAMLAuthHandler._bool_env("SAML_STRICT", True),
"debug": False,
"sp": {
"entityId": SAMLAuthHandler._sp_entity_id(request),
"assertionConsumerService": {
"url": SAMLAuthHandler._acs_url(request),
"binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
},
"NameIDFormat": SAMLAuthHandler._env(
"SAML_SP_NAME_ID_FORMAT",
"urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress",
),
},
"security": {
"wantAssertionsSigned": SAMLAuthHandler._bool_env("SAML_WANT_ASSERTIONS_SIGNED", True),
"wantMessagesSigned": SAMLAuthHandler._bool_env("SAML_WANT_MESSAGES_SIGNED", False),
"authnRequestsSigned": SAMLAuthHandler._bool_env("SAML_AUTHN_REQUESTS_SIGNED", False),
"wantNameId": True,
"requestedAuthnContext": False,
"rejectUnsolicitedResponsesWithInResponseTo": False,
},
}
return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings)
@staticmethod
def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]:
base = SAMLAuthHandler._base_url(request)
scheme, _, host_part = base.partition("://")
host = host_part.split("/", 1)[0]
return {
"https": "on" if scheme == "https" else "off",
"http_host": host,
"script_name": "/" + SAML_CALLBACK_ROUTE,
"get_data": dict(request.query_params),
"post_data": post_data or {},
}
@staticmethod
async def _build_auth(
request: Request,
cache: DualCache,
post_data: dict[str, str] | None = None,
) -> "OneLogin_Saml2_Auth":
if not SAML_AVAILABLE:
raise _saml_unavailable_error()
idp_settings = await SAMLAuthHandler._load_idp_settings(cache)
settings = SAMLAuthHandler._build_settings(request, idp_settings)
request_data = SAMLAuthHandler._prepare_request_data(request, post_data)
try:
return OneLogin_Saml2_Auth(request_data, old_settings=settings)
except Exception as e: # noqa: BLE001 - toolkit exposes no common exception base; fail closed
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Invalid SAML configuration: {e}",
)
@staticmethod
async def build_login_redirect(
request: Request, cache: DualCache, relay_state: str | None = None
) -> RedirectResponse:
auth = await SAMLAuthHandler._build_auth(request, cache)
redirect_url = cast(str, auth.login(return_to=relay_state)) # cast-ok: untyped python3-saml
response = RedirectResponse(url=redirect_url, status_code=303)
request_id = cast(str | None, auth.get_last_request_id()) # cast-ok: untyped python3-saml
if request_id is not None:
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}",
value="1",
ttl=_SAML_AUTHN_REQUEST_TTL_SECONDS,
)
secure = SAMLAuthHandler._is_https(request)
response.set_cookie(
key=_SAML_AUTHN_STATE_COOKIE,
value=request_id,
max_age=_SAML_AUTHN_REQUEST_TTL_SECONDS,
httponly=True,
secure=secure,
samesite="none" if secure else "lax",
)
return response
@staticmethod
async def build_sp_metadata(request: Request, cache: DualCache) -> str:
if not SAML_AVAILABLE:
raise _saml_unavailable_error()
idp_settings = await SAMLAuthHandler._load_idp_settings(cache)
settings = SAMLAuthHandler._build_settings(request, idp_settings)
saml_settings = OneLogin_Saml2_Settings(settings, sp_validation_only=True)
metadata = cast(str, saml_settings.get_sp_metadata()) # cast-ok: untyped python3-saml
errors = cast(list[str], saml_settings.validate_metadata(metadata)) # cast-ok: untyped python3-saml
if errors:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Invalid SP metadata: {', '.join(errors)}",
)
return metadata
@staticmethod
async def read_acs_post_data(request: Request) -> dict[str, str]:
"""Read the ACS POST form under a hard size cap before any base64/XML decoding.
Bounds both Content-Length-declared and chunked requests so an unauthenticated
caller cannot force unbounded buffering while decoding the SAMLResponse."""
declared = request.headers.get("content-length")
if declared is not None and declared.isdigit() and int(declared) > _SAML_MAX_POST_BYTES:
raise HTTPException(
status_code=status.HTTP_413_CONTENT_TOO_LARGE,
detail="SAML response exceeds the maximum allowed size.",
)
body = bytearray()
async for chunk in request.stream():
body += chunk
if len(body) > _SAML_MAX_POST_BYTES:
raise HTTPException(
status_code=status.HTTP_413_CONTENT_TOO_LARGE,
detail="SAML response exceeds the maximum allowed size.",
)
return dict(parse_qsl(body.decode("utf-8", "replace")))
@staticmethod
async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID:
auth = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data)
browser_request_id = request.cookies.get(_SAML_AUTHN_STATE_COOKIE)
try:
auth.process_response(request_id=browser_request_id)
except Exception as e: # noqa: BLE001 - toolkit exposes no common exception base; fail closed
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=f"Could not process SAML response: {e}",
)
errors = cast(list[str], auth.get_errors()) # cast-ok: untyped python3-saml
if errors or not auth.is_authenticated():
reason = auth.get_last_error_reason()
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=f"SAML authentication failed: {reason or ', '.join(errors)}",
)
await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id)
return SAMLAuthHandler._result_from_auth(auth)
@staticmethod
def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int:
not_on_or_after = auth.get_last_assertion_not_on_or_after()
if not isinstance(not_on_or_after, int):
return _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS
remaining = not_on_or_after - int(time.time())
return min(
max(remaining, _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS),
_SAML_REPLAY_GUARD_MAX_TTL_SECONDS,
)
@staticmethod
def _response_in_response_to(auth: "OneLogin_Saml2_Auth") -> str | None:
"""The request id this response answers, read from the Response element or, when the
IdP only stamps it on the bearer SubjectConfirmationData, from there. A non-None value
marks the response as solicited (SP-initiated) and so requiring browser binding."""
value = cast(str | None, auth.get_last_response_in_response_to()) # cast-ok: untyped python3-saml
if value:
return value
xml = cast(bytes | None, auth.get_last_response_xml()) # cast-ok: untyped python3-saml
if not xml:
return None
root = OneLogin_Saml2_XML.to_etree(xml)
for node in OneLogin_Saml2_XML.query(root, "//saml:SubjectConfirmationData[@InResponseTo]"):
irt = cast(str | None, node.get("InResponseTo")) # cast-ok: untyped python3-saml
if irt:
return irt
return None
@staticmethod
async def _enforce_response_binding(
auth: "OneLogin_Saml2_Auth",
cache: DualCache,
browser_request_id: str | None,
) -> None:
in_response_to = SAMLAuthHandler._response_in_response_to(auth)
if in_response_to is not None:
authn_key = f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{in_response_to}"
if cache.get_cache(key=authn_key) is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="SAML response references an unknown or already-used login request.",
)
if browser_request_id is None or not secrets.compare_digest(browser_request_id, in_response_to):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="SAML response is not bound to this browser's login request.",
)
elif browser_request_id is not None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="SAML response is not bound to this browser's login request.",
)
elif not SAMLAuthHandler._bool_env("SAML_ALLOW_UNSOLICITED", False):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Unsolicited (IdP-initiated) SAML responses are disabled.",
)
elif cache.redis_cache is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=(
"Unsolicited (IdP-initiated) SAML responses require a shared Redis cache "
"so the replay guard is enforced across every worker."
),
)
assertion_id = cast(str | None, auth.get_last_assertion_id()) # cast-ok: untyped python3-saml
if assertion_id is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="SAML assertion is missing the required ID attribute.",
)
consumed_key = f"{_SAML_CONSUMED_ASSERTION_CACHE_PREFIX}:{assertion_id}"
consumed_count = await cache.async_increment_cache(
key=consumed_key, value=1, ttl=SAMLAuthHandler._replay_guard_ttl(auth)
)
if consumed_count is not None and consumed_count > 1:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="SAML assertion has already been used (replay detected).",
)
@staticmethod
def _result_from_auth(auth: "OneLogin_Saml2_Auth") -> CustomOpenID:
attributes = cast(dict[str, list[str]], auth.get_attributes()) # cast-ok: untyped python3-saml
name_id = cast(str | None, auth.get_nameid()) # cast-ok: untyped python3-saml
email = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_EMAIL", _EMAIL_ATTRIBUTE_CANDIDATES)
if email is None and name_id is not None and "@" in name_id:
email = name_id
if email is None and SAMLAuthHandler._env("ALLOWED_EMAIL_DOMAINS") is not None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=(
"SAML assertion did not contain an email address, but ALLOWED_EMAIL_DOMAINS "
"restricts sign-in by email domain."
),
)
user_id = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_USER_ID", ()) or name_id or email
if user_id is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="SAML assertion did not contain a usable subject (NameID) or email.",
)
first_name = SAMLAuthHandler._attribute_value(
attributes, "SAML_ATTRIBUTE_FIRST_NAME", _FIRST_NAME_ATTRIBUTE_CANDIDATES
)
last_name = SAMLAuthHandler._attribute_value(
attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES
)
role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES)
team_ids = SAMLAuthHandler._attribute_values(
attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES
)
display_name = " ".join(part for part in (first_name, last_name) if part) or email
verbose_proxy_logger.info(f"SAML login: subject={user_id}, email={email}, attributes={list(attributes.keys())}")
try:
return CustomOpenID(
id=user_id,
email=email,
first_name=first_name,
last_name=last_name,
display_name=display_name,
picture=None,
provider="saml",
team_ids=team_ids,
user_role=get_litellm_user_role(role_value) if role_value else None,
)
except ValidationError as e:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=f"SAML assertion contained an invalid subject or email: {e}",
)
@staticmethod
def _attribute_value(
attributes: dict[str, list[str]],
env_override: str,
candidates: tuple[str, ...],
) -> str | None:
values = SAMLAuthHandler._attribute_values(attributes, env_override, candidates)
return values[0] if values else None
@staticmethod
def _attribute_values(
attributes: dict[str, list[str]],
env_override: str,
candidates: tuple[str, ...],
) -> list[str]:
override = SAMLAuthHandler._env(env_override)
keys = (override, *candidates) if override else candidates
for key in keys:
values = attributes.get(key)
if values:
return [v for v in values if v]
return []

152
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

@ -10,16 +10,18 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a
"""
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, List, Optional
from datetime import datetime, timedelta, timezone
from itertools import groupby
from typing import TYPE_CHECKING, Annotated, Any, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, TypeAdapter
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.table_repositories import (
@ -39,6 +41,9 @@ from litellm.types.tool_management import (
ToolPolicyOptionsResponse,
ToolPolicyUpdateRequest,
ToolPolicyUpdateResponse,
ToolSpendDailyEntry,
ToolSpendEntry,
ToolSpendResponse,
ToolUsageLogEntry,
ToolUsageLogsResponse,
)
@ -124,6 +129,147 @@ async def list_tools(
raise HTTPException(status_code=500, detail=str(e))
def _parse_day_start(value: str | None) -> datetime | None:
if not value:
return None
try:
return datetime.strptime(value.strip(), "%Y-%m-%d").replace(tzinfo=timezone.utc)
except ValueError:
raise HTTPException(
status_code=400,
detail=f"Invalid date format: {value}. Expected: 'YYYY-MM-DD'",
)
class _ToolSpendRow(BaseModel):
date: str
tool_name: str
call_count: int
spend: float
total_tokens: int
class _RequestTotalRow(BaseModel):
total_spend: float
_TOOL_SPEND_ROWS = TypeAdapter(list[_ToolSpendRow])
_REQUEST_TOTAL_ROWS = TypeAdapter(list[_RequestTotalRow])
def _summarize_tool(name: str, grp: tuple[_ToolSpendRow, ...]) -> ToolSpendEntry:
return ToolSpendEntry(
tool_name=name,
spend=sum(r.spend for r in grp),
call_count=sum(r.call_count for r in grp),
total_tokens=sum(r.total_tokens for r in grp),
)
def _build_tool_spend_response(
rows: list[_ToolSpendRow],
total_spend: float,
start_date: str,
end_date: str,
) -> ToolSpendResponse:
daily = [
ToolSpendDailyEntry(date=r.date, tool_name=r.tool_name, spend=r.spend, call_count=r.call_count) for r in rows
]
grouped = groupby(sorted(rows, key=lambda r: r.tool_name), key=lambda r: r.tool_name)
by_tool = sorted(
(_summarize_tool(name, tuple(grp)) for name, grp in grouped),
key=lambda e: e.spend,
reverse=True,
)
return ToolSpendResponse(
by_tool=by_tool,
daily=daily,
total_spend=total_spend,
start_date=start_date,
end_date=end_date,
)
@router.get(
"/v1/tool/spend",
tags=["tool management"],
dependencies=[Depends(user_api_key_auth)],
response_model=ToolSpendResponse,
)
async def get_tool_spend(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
start_date: Annotated[str | None, Query(description="YYYY-MM-DD (defaults to 30 days ago)")] = None,
end_date: Annotated[str | None, Query(description="YYYY-MM-DD (defaults to today)")] = None,
):
"""
Spend attributed to each tool over a date range, for the Cost Optimization dashboard.
Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to
``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools
counts its full spend toward each of those tools, so per-tool numbers are
attributions. ``total_spend`` is the deduplicated spend of every request that
called at least one tool in the window, so it never double counts.
"""
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
):
raise HTTPException(
status_code=403,
detail="Only proxy admin roles can view tool spend across the deployment",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
now = datetime.now(timezone.utc)
end_day = _parse_day_start(end_date)
start_dt = _parse_day_start(start_date) or ((end_day or now) - timedelta(days=30))
end_exclusive = (end_day + timedelta(days=1)) if end_day else now
rows = await prisma_client.db.query_raw(
"""
SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date,
ti.tool_name AS tool_name,
COUNT(*)::int AS call_count,
COALESCE(SUM(sl.spend), 0)::double precision AS spend,
COALESCE(SUM(sl.total_tokens), 0)::bigint AS total_tokens
FROM "LiteLLM_SpendLogToolIndex" ti
JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id
WHERE ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC')
AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC')
GROUP BY date, ti.tool_name
ORDER BY date ASC, spend DESC
""",
start_dt.isoformat(),
end_exclusive.isoformat(),
)
totals = await prisma_client.db.query_raw(
"""
SELECT COALESCE(SUM(sl.spend), 0)::double precision AS total_spend
FROM "LiteLLM_SpendLogs" sl
WHERE EXISTS (
SELECT 1
FROM "LiteLLM_SpendLogToolIndex" ti
WHERE ti.request_id = sl.request_id
AND ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC')
AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC')
)
""",
start_dt.isoformat(),
end_exclusive.isoformat(),
)
total_rows = _REQUEST_TOTAL_ROWS.validate_python(totals or [])
return _build_tool_spend_response(
rows=_TOOL_SPEND_ROWS.validate_python(rows or []),
total_spend=total_rows[0].total_spend if total_rows else 0.0,
start_date=start_dt.strftime("%Y-%m-%d"),
end_date=(end_day or now).strftime("%Y-%m-%d"),
)
@router.get(
"/v1/tool/{tool_name:path}/detail",
tags=["tool management"],

137
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -100,6 +100,7 @@ from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO
from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler
from litellm.proxy.management_endpoints.sso_helper_utils import (
check_is_admin_only_access,
has_admin_ui_access,
@ -857,6 +858,27 @@ def process_sso_jwt_access_token(
return None
async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None:
"""Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license."""
if premium_user is True:
return
if prisma_client is None:
raise ProxyException(
message=CommonProxyErrors.db_not_connected_error.value,
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
billable_users = await UserRepository(prisma_client).count_billable_users()
if billable_users and billable_users > 5:
raise ProxyException(
message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You configured SSO (one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, `GENERIC_CLIENT_ID`, or SAML) in your env. Please unset it",
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
@router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False)
async def google_login(
request: Request,
@ -876,6 +898,7 @@ async def google_login(
general_settings,
premium_user,
prisma_client,
user_api_key_cache,
user_custom_ui_sso_sign_in_handler,
)
@ -891,25 +914,13 @@ async def google_login(
return admin_ui_disabled()
####### Check if user is a Enterprise / Premium User #######
if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None:
if premium_user is not True:
# Check if under 'free SSO user' limit
if prisma_client is not None:
billable_users = await UserRepository(prisma_client).count_billable_users()
if billable_users and billable_users > 5:
raise ProxyException(
message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this",
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
else:
raise ProxyException(
message=CommonProxyErrors.db_not_connected_error.value,
type=ProxyErrorTypes.auth_error,
param="premium_user",
code=status.HTTP_403_FORBIDDEN,
)
if (
microsoft_client_id is not None
or google_client_id is not None
or generic_client_id is not None
or SAMLAuthHandler.is_saml_configured()
):
await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client)
####### Detect DB + MASTER KEY in .env #######
missing_env_vars = show_missing_vars_in_env()
@ -947,6 +958,19 @@ async def google_login(
"Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise."
)
if (
microsoft_client_id is None
and google_client_id is None
and generic_client_id is None
and SAMLAuthHandler.is_saml_configured()
):
verbose_proxy_logger.info("Redirecting to SAML SSO login")
return await SAMLAuthHandler.build_login_redirect(
request=request,
cache=user_api_key_cache,
relay_state=return_to,
)
# Check if we should use SSO handler
if (
SSOAuthenticationHandler.should_use_sso_handler(
@ -1913,6 +1937,81 @@ async def auth_callback(request: Request, state: Optional[str] = None):
)
@router.get("/sso/saml/login", tags=["experimental"], include_in_schema=False)
async def saml_login(request: Request, return_to: str | None = None):
"""SP-initiated SAML login. Redirects the user to the configured IdP."""
from litellm.proxy.proxy_server import user_api_key_cache
_disable_ui_flag = os.getenv("DISABLE_ADMIN_UI")
if _disable_ui_flag is not None and str_to_bool(value=_disable_ui_flag):
return admin_ui_disabled()
return await SAMLAuthHandler.build_login_redirect(request=request, cache=user_api_key_cache, relay_state=return_to)
@router.get("/sso/saml/metadata", tags=["experimental"], include_in_schema=False)
async def saml_metadata(request: Request):
"""Service Provider metadata XML, for registering this proxy at the IdP."""
from litellm.proxy.proxy_server import user_api_key_cache
metadata = await SAMLAuthHandler.build_sp_metadata(request=request, cache=user_api_key_cache)
return Response(content=metadata, media_type="application/xml")
@router.post("/sso/saml/callback", tags=["experimental"], include_in_schema=False)
async def saml_callback(request: Request):
"""Assertion Consumer Service. Validates the IdP assertion and issues a UI session."""
from litellm.proxy.proxy_server import (
general_settings,
jwt_handler,
master_key,
premium_user,
prisma_client,
user_api_key_cache,
)
_disable_ui_flag = os.getenv("DISABLE_ADMIN_UI")
if _disable_ui_flag is not None and str_to_bool(value=_disable_ui_flag):
return admin_ui_disabled()
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
if master_key is None:
raise ProxyException(
message="Master Key not set for Proxy. Set `LITELLM_MASTER_KEY` in .env or general_settings:master_key in config.yaml.",
type=ProxyErrorTypes.auth_error,
param="master_key",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
post_data = await SAMLAuthHandler.read_acs_post_data(request)
if "SAMLResponse" not in post_data:
raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.")
result = await SAMLAuthHandler.handle_acs(request=request, cache=user_api_key_cache, post_data=post_data)
await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client)
ui_access_mode = general_settings.get("ui_access_mode", None)
relay_state = post_data.get("RelayState")
cp_return_to: str | None = (
relay_state
if isinstance(relay_state, str) and SSOAuthenticationHandler._validate_return_to(relay_state)
else None
)
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=result,
request=request,
received_response=None,
generic_client_id=None,
ui_access_mode=ui_access_mode,
access_token_payload=None,
jwt_handler=jwt_handler,
return_to=cp_return_to,
)
async def _build_cli_sso_user_defined_values(
result: Union[OpenID, dict],
parsed_openid_result: ParsedOpenIDResult,

10
litellm/proxy/management_helpers/utils.py
View file

@ -295,10 +295,16 @@ async def add_new_member(
# same new user is provisioned concurrently), seeding teams on create.
# The teams append lives in the filtered update below rather than the
# upsert's update branch so an already-existing user does not get a
# duplicate team id.
# duplicate team id. The update branch still has to write something:
# Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
# is non-empty, and falls back to a racy SELECT-then-INSERT when it is
# not, so this re-states user_id as a no-op rather than being empty.
_returned_user = await UserRepository(prisma_client).table.upsert(
where={"user_id": new_member.user_id},
data={"create": {"teams": [team_id], **new_user_defaults}, "update": {}},
data={
"create": {"teams": [team_id], **new_user_defaults},
"update": {"user_id": new_member.user_id},
},
)
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id)
if _returned_user is not None:

23
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -4,7 +4,9 @@ import asyncio
import json
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Mapping, Optional, Sequence, cast
from typing import Any, Dict, List, Mapping, NoReturn, Optional, Sequence, cast
from fastapi import HTTPException, status
import litellm
from litellm._logging import verbose_proxy_logger
@ -59,6 +61,22 @@ class _CounterReservationUnavailable(Exception):
super().__init__("Counter reservation unavailable")
def _raise_reservation_unavailable(counter_key: str) -> NoReturn:
verbose_proxy_logger.warning(
"fail_closed_budget_enforcement: rejecting request — budget reservation for %s could not be written",
counter_key,
)
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail=(
"Budget enforcement unavailable: the budget reservation could not "
"be written to the spend counter backend, and "
"fail_closed_budget_enforcement is enabled, so the request was "
"rejected to avoid exceeding the configured budget. Retry shortly."
),
)
def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set:
if not budget_reservation:
return set()
@ -138,6 +156,7 @@ async def reserve_budget_for_request(
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
fail_closed_budget_enforcement: bool = False,
) -> Optional[dict]:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
return None
@ -193,6 +212,8 @@ async def reserve_budget_for_request(
default_reserved_cost=reservation_cost,
)
applied_entries.remove(entry)
if fail_closed_budget_enforcement:
_raise_reservation_unavailable(counter_key=counter.counter_key)
continue
if reserved_value is not None:

108
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -1695,13 +1695,10 @@ async def ui_view_spend_logs(
code=status.HTTP_401_UNAUTHORIZED,
)
if start_date is None or end_date is None:
raise ProxyException(
message="Start date and end date are required",
type="bad_request",
param="None",
code=status.HTTP_400_BAD_REQUEST,
)
# Inline import — auth_utils participates in a proxy import cycle.
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
is_v2 = "/spend/logs/v2" in get_request_route(request)
# Validate sort_by and sort_order
valid_sort_fields = {
@ -1729,36 +1726,50 @@ async def ui_view_spend_logs(
)
try:
# Inline import — auth_utils participates in a proxy import cycle.
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
is_request_id_lookup = request_id is not None and not is_v2
is_v2 = "/spend/logs/v2" in get_request_route(request)
formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"]
if is_request_id_lookup:
# request_id is the @id primary key: it identifies a single row, so a
# time window is meaningless. The dashboard always sends a default 24h
# window, which hid ids copied from an older page (LIT-3981). Drop the
# window for the id lookup so it resolves across all time; every other
# query, including the public v2 route, still requires one (below).
start_date_obj: datetime | None = None
end_date_obj: datetime | None = None
else:
if start_date is None or end_date is None:
raise ProxyException(
message="Start date and end date are required",
type="bad_request",
param="None",
code=status.HTTP_400_BAD_REQUEST,
)
formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"]
def parse_date(date_str: str) -> datetime:
date_str = date_str.strip()
for fmt in formats:
try:
return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc)
except ValueError:
continue
expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'"
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Invalid date format: {date_str}. Expected: {expected}",
)
def parse_date(date_str: str) -> datetime:
date_str = date_str.strip()
for fmt in formats:
try:
return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc)
except ValueError:
continue
expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'"
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Invalid date format: {date_str}. Expected: {expected}",
)
start_date_obj = parse_date(start_date)
end_date_obj = parse_date(end_date)
# Convert to ISO format strings for Prisma
start_date_iso = start_date_obj.isoformat() # Already in UTC, no need to add Z
end_date_iso = end_date_obj.isoformat() # Already in UTC, no need to add Z
start_date_obj = parse_date(start_date)
end_date_obj = parse_date(end_date)
# Build where conditions
where_conditions: dict[str, Any] = {
"startTime": {"gte": start_date_iso, "lte": end_date_iso},
}
where_conditions: dict[str, Any] = {}
if start_date_obj is not None and end_date_obj is not None:
where_conditions["startTime"] = {
"gte": start_date_obj.isoformat(), # Already in UTC, no need to add Z
"lte": end_date_obj.isoformat(),
}
if team_id is not None:
where_conditions["team_id"] = team_id
@ -1827,9 +1838,19 @@ async def ui_view_spend_logs(
where_conditions["spend"]["gte"] = min_spend
if max_spend is not None:
where_conditions["spend"]["lte"] = max_spend
is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
# A request_id lookup drops the date window, so a non-admin could otherwise
# reach any single row by id; require they own it, mirroring the detail
# endpoint. That ownership check fully authorizes the one row, so the
# general scoping below is skipped for id lookups. Scoped to the UI route
# so the public v2 contract is unchanged.
if request_id is not None and not is_v2 and not is_admin_view:
await _assert_user_can_view_request_id(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
request_id=request_id,
)
permitted_team_ids: List[str] | None = None
if not is_admin_view:
if not is_request_id_lookup and not is_admin_view:
if team_id is not None:
can_view_team = await _can_team_member_view_log(
prisma_client=prisma_client,
@ -1875,15 +1896,16 @@ async def ui_view_spend_logs(
sql_params: List[Any] = []
p = 1 # parameter index counter
# Date range (always present). Wrap the param side with
# `AT TIME ZONE 'UTC'` so comparison against the plain `timestamp`
# column does not depend on the DB session timezone (see #22529).
sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')")
sql_params.append(start_date_obj)
p += 1
sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')")
sql_params.append(end_date_obj)
p += 1
# Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison
# against the plain `timestamp` column does not depend on the DB session
# timezone (see #22529). Absent for a request_id-only lookup (see above).
if start_date_obj is not None and end_date_obj is not None:
sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')")
sql_params.append(start_date_obj)
p += 1
sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')")
sql_params.append(end_date_obj)
p += 1
# Equality filters - read effective values from where_conditions (post-authorization)
for sql_col, wc_key in [

16
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -374,12 +374,22 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
if isinstance(v, BaseModel):
v = v.model_dump()
additional_usage_values.update({k: v})
if "cache_read_input_tokens" not in additional_usage_values:
prompt_tokens_details = additional_usage_values.get("prompt_tokens_details")
if isinstance(prompt_tokens_details, dict):
prompt_tokens_details = additional_usage_values.get("prompt_tokens_details")
if not isinstance(prompt_tokens_details, dict):
usage_object = clean_metadata.get("usage_object")
if isinstance(usage_object, dict):
prompt_tokens_details = usage_object.get("prompt_tokens_details")
if isinstance(prompt_tokens_details, dict):
if "cache_read_input_tokens" not in additional_usage_values:
cached_tokens = prompt_tokens_details.get("cached_tokens")
if isinstance(cached_tokens, int) and cached_tokens > 0:
additional_usage_values["cache_read_input_tokens"] = cached_tokens
if "cache_creation_input_tokens" not in additional_usage_values:
cache_write_tokens = prompt_tokens_details.get("cache_write_tokens") or prompt_tokens_details.get(
"cache_creation_tokens"
)
if isinstance(cache_write_tokens, int) and cache_write_tokens > 0:
additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens
clean_metadata["additional_usage_values"] = additional_usage_values
if litellm.cache is not None:

135
litellm/proxy/utils.py
View file

@ -1400,6 +1400,14 @@ class ProxyLogging:
self._process_guardrail_metadata(data)
return data
parallel_guardrails: tuple[CustomGuardrail, ...] = tuple(
cb
for cb in caps.resolved_callbacks
if isinstance(cb, CustomGuardrail)
and getattr(cb, "run_in_parallel", False)
and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed)
)
deferred_route_exc: Optional[SensitiveDataRouteException] = None
for _callback in caps.resolved_callbacks:
start_time = time.time()
@ -1409,6 +1417,9 @@ class ProxyLogging:
if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed:
continue
if getattr(_callback, "run_in_parallel", False):
continue
result = await self._process_guardrail_callback(
callback=_callback,
data=data, # type: ignore
@ -1465,6 +1476,14 @@ class ProxyLogging:
if deferred_route_exc is not None and data is not None:
data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict)
if parallel_guardrails and data is not None:
await self._run_parallel_pre_call_guardrails(
guardrails=parallel_guardrails,
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
if data is not None:
self._process_guardrail_metadata(data)
@ -1477,6 +1496,47 @@ class ProxyLogging:
except Exception as e:
raise e
async def _run_parallel_pre_call_guardrails(
self,
guardrails: tuple[CustomGuardrail, ...],
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
) -> None:
"""
Run opted-in pre_call guardrails concurrently against one shared payload
snapshot. These guardrails are declared block-only, so any modified data
they return is discarded; they run for their blocking side effect (raising
to reject the request before it reaches the LLM). Every guardrail is
awaited to completion (``return_exceptions=True``) so a raise by one never
leaves the others running as unobserved background tasks. A guardrail that
blocks (any exception other than a reroute or passthrough) takes precedence
over one that only changes the request flow, so a fast reroute can never
let a slower block be bypassed; the request is rejected before it reaches
the LLM, preserving the pre-call barrier that ``during_call`` guardrails
cannot provide. Per-guardrail latency is recorded by
``_process_guardrail_callback``'s own metrics.
"""
results = await asyncio.gather(
*(
self._process_guardrail_callback(
callback=callback,
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
event_type=GuardrailEventHooks.pre_call,
)
for callback in guardrails
),
return_exceptions=True,
)
raised = tuple(result for result in results if isinstance(result, BaseException))
blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
if blocking is not None:
raise blocking
if raised:
raise raised[0]
async def _handle_sensitive_data_route_exception(
self,
exc: SensitiveDataRouteException,
@ -2277,9 +2337,16 @@ class ProxyLogging:
# Merge model-level guardrails before checking which guardrails to run
guardrail_data = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
parallel_guardrails: tuple[CustomGuardrail, ...] = tuple(
callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False)
)
for callback in guardrail_callbacks:
# Main - V2 Guardrails implementation
if getattr(callback, "run_in_parallel", False):
continue
if (
callback.should_run_guardrail(
data=guardrail_data,
@ -2316,6 +2383,15 @@ class ProxyLogging:
if guardrail_response is not None:
response = guardrail_response
if parallel_guardrails:
await self._run_parallel_post_call_guardrails(
guardrails=parallel_guardrails,
data=data,
guardrail_data=guardrail_data,
response=response,
user_api_key_dict=user_api_key_dict,
)
############ Handle CustomLogger ###############################
#################################################################
@ -2329,6 +2405,65 @@ class ProxyLogging:
raise e
return response
async def _run_parallel_post_call_guardrails(
self,
guardrails: tuple[CustomGuardrail, ...],
data: dict,
guardrail_data: dict,
response: LLMResponseTypes,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Run opted-in post_call guardrails concurrently against the response
produced by the sequential guardrails. These guardrails are declared
block-only, so any modified response they return is discarded; they run
for their blocking side effect (raising to reject the response before it
reaches the client). Every guardrail is awaited to completion
(``return_exceptions=True``) so a raise by one never leaves the others
running as unobserved background tasks. A guardrail that blocks (any
exception other than a passthrough) takes precedence over one that only
changes the response flow, so a fast passthrough can never let a slower
block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply``
immediately before awaiting, and the unified hook pops it before its first
suspension point, so concurrent guardrails never race on that key.
"""
async def _run_one(callback: CustomGuardrail) -> None:
if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True:
return
if "apply_guardrail" in type(callback).__dict__:
data["guardrail_to_apply"] = callback
await self._run_guardrail_with_metrics(
callback,
unified_guardrail.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
),
"post_call",
)
else:
await self._run_guardrail_with_metrics(
callback,
callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
),
"post_call",
)
results = await asyncio.gather(
*(_run_one(callback) for callback in guardrails),
return_exceptions=True,
)
raised = tuple(result for result in results if isinstance(result, BaseException))
blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
if blocking is not None:
raise blocking
if raised:
raise raised[0]
async def post_call_response_headers_hook(
self,
data: dict,

1
litellm/responses/utils.py
View file

@ -1049,6 +1049,7 @@ class ResponseAPILoggingUtils:
audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None),
text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None),
image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None),
cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None),
)
completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None
output_tokens_details = getattr(response_api_usage, "output_tokens_details", None)

3
litellm/setup_wizard.py
View file

@ -52,12 +52,13 @@ PROVIDERS: List[Dict] = [
{
"id": "anthropic",
"name": "Anthropic",
"description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5",
"description": "Claude Fable 5, Opus 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 5, Sonnet 4.6, Haiku 4.5",
"env_key": "ANTHROPIC_API_KEY",
"key_hint": "sk-ant-...",
"test_model": "claude-haiku-4-5-20251001",
"models": [
"claude-fable-5",
"claude-opus-5",
"claude-sonnet-5",
"claude-opus-4-8",
"claude-opus-4-7",

11
litellm/types/guardrails.py
View file

@ -910,6 +910,17 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
),
)
run_in_parallel: Optional[bool] = Field(
default=None,
description=(
"When True, this pre_call or post_call guardrail runs concurrently with other opted-in "
"guardrails of the same hook, after the sequential guardrails have run. Use only for "
"block-only guardrails that inspect and reject; do not enable it for guardrails that "
"modify the request or response (e.g. PII masking or sensitive-data routing), since "
"parallel runs share one snapshot and their mutations would race."
),
)
@field_validator(
"mode",
"default_action",

9
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -261,12 +261,3 @@ class MCPServer(BaseModel):
if self.oauth_passthrough is not True:
return False
return any(h.lower() == "authorization" for h in self.extra_headers)
@property
def has_token_exchange_config(self) -> bool:
"""True if this server is configured for OAuth2 token exchange (OBO / RFC 8693)."""
return (
self.auth_type == MCPAuth.oauth2_token_exchange
and bool(self.client_id and self.client_secret)
and bool(self.token_exchange_endpoint or self.token_url)
)

18
litellm/types/proxy/management_endpoints/ui_sso.py
View file

@ -153,6 +153,24 @@ class SSOConfig(LiteLLMPydanticObjectBase):
description="Space-separated OAuth scopes requested from the generic provider, e.g. 'openid email profile'",
)
# SAML SSO
saml_idp_metadata_url: Optional[str] = Field(
default=None,
description="URL of the SAML IdP metadata to fetch and parse for SSO authentication",
)
saml_idp_metadata_xml: Optional[str] = Field(
default=None,
description="Inline SAML IdP metadata XML, used when a metadata URL is not available",
)
saml_sp_entity_id: Optional[str] = Field(
default=None,
description="SAML Service Provider entityID; defaults to the proxy's /sso/saml/metadata URL",
)
saml_allow_unsolicited: Optional[str] = Field(
default=None,
description="'true' to accept IdP-initiated (unsolicited) SAML responses, which cannot be browser-bound against login CSRF",
)
# Common settings
proxy_base_url: Optional[str] = Field(
default=None,

35
litellm/types/tool_management.py
View file

@ -98,3 +98,38 @@ class ToolUsageLogsResponse(BaseModel):
total: int
page: int
page_size: int
class ToolSpendEntry(BaseModel):
"""Total spend attributed to one tool over the requested window."""
tool_name: str
spend: float = Field(
0.0,
description="Attributed spend: a request that used several tools counts its full spend toward each of them",
)
call_count: int = 0
total_tokens: int = 0
class ToolSpendDailyEntry(BaseModel):
"""Spend attributed to one tool on one UTC day."""
date: str
tool_name: str
spend: float = 0.0
call_count: int = 0
class ToolSpendResponse(BaseModel):
by_tool: List[ToolSpendEntry] = Field(default_factory=list)
daily: List[ToolSpendDailyEntry] = Field(default_factory=list)
total_spend: float = Field(
0.0,
description=(
"Deduplicated spend of every request that called at least one tool in the window; "
"less than the sum of per-tool attributed spend whenever multi-tool requests exist"
),
)
start_date: str | None = None
end_date: str | None = None

21
litellm/types/utils.py
View file

@ -1534,14 +1534,27 @@ class PromptTokensDetailsWrapper(
audio_length_seconds: Optional[float] = None
"""Length of audio sent to the model. Used for multimodal embeddings priced per audio-second."""
cache_write_tokens: Optional[int] = None
"""Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field."""
cache_creation_tokens: Optional[int] = None
"""Number of cache creation tokens sent to the model. Used for Anthropic prompt caching."""
"""Number of cache creation tokens sent to the model. Anthropic/Bedrock naming; kept in sync with cache_write_tokens (assigning either mirrors to the other)."""
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
"""Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching."""
def __setattr__(self, name: str, value: object) -> None:
super().__setattr__(name, value)
if name == "cache_write_tokens":
super().__setattr__("cache_creation_tokens", value)
elif name == "cache_creation_tokens":
super().__setattr__("cache_write_tokens", value)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.cache_write_tokens = (
self.cache_write_tokens if self.cache_write_tokens is not None else self.cache_creation_tokens
)
if self.character_count is None:
del self.character_count
if self.image_count is None:
@ -1554,6 +1567,8 @@ class PromptTokensDetailsWrapper(
del self.web_search_requests
if self.tool_use_tokens is None:
del self.tool_use_tokens
if self.cache_write_tokens is None:
del self.cache_write_tokens
if self.cache_creation_tokens is None:
del self.cache_creation_tokens
if self.cache_creation_token_details is None:
@ -1662,10 +1677,10 @@ class Usage(SafeAttributeModel, CompletionUsage):
if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int):
if _prompt_tokens_details is None:
_prompt_tokens_details = PromptTokensDetailsWrapper(
cache_creation_tokens=params["cache_creation_input_tokens"]
cache_write_tokens=params["cache_creation_input_tokens"]
)
else:
_prompt_tokens_details.cache_creation_tokens = params["cache_creation_input_tokens"]
_prompt_tokens_details.cache_write_tokens = params["cache_creation_input_tokens"]
super().__init__(
prompt_tokens=prompt_tokens or 0,

350
model_prices_and_context_window.json
View file

@ -1502,6 +1502,222 @@
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 1024
},
"anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"global.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"us.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"eu.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"au.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"jp.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.75e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_output_config": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 512
},
"anthropic.claude-opus-4-8": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
@ -2756,6 +2972,38 @@
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true
},
"azure_ai/claude-opus-5": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2.5e-05,
"litellm_provider": "azure_ai",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"prompt_cache_min_tokens": 512
},
"azure_ai/claude-opus-4-8": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
@ -11846,6 +12094,44 @@
"supports_output_config": true,
"prompt_cache_min_tokens": 512
},
"claude-opus-5": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "anthropic",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"provider_specific_entry": {
"us": 1.1,
"fast": 2.0
},
"supports_output_config": true,
"supports_speed": true,
"prompt_cache_min_tokens": 512
},
"claude-opus-4-8": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
@ -36987,6 +37273,70 @@
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true
},
"vertex_ai/claude-opus-5": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"prompt_cache_min_tokens": 512
},
"vertex_ai/claude-opus-5@default": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": true,
"prompt_cache_min_tokens": 512
},
"vertex_ai/claude-opus-4-8": {
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,

4
pyproject.toml
View file

@ -99,6 +99,10 @@ utils = [
"numpydoc>=1.8.0,<2.0",
]
caching = ["diskcache>=5.6.3,<6.0"]
# SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels
# bundle the native libxmlsec1/libxml2 libraries, so no system packages are
# required. Kept out of the base `proxy` extra so it stays optional.
saml = ["python3-saml>=1.16.0,<2.0"]
semantic-router = [
"semantic-router>=0.1.15,<1.0; python_version < '3.14'",
"aurelio-sdk>=0.0.19,<1.0; python_version < '3.14'",

4
ruff-strict-budget.json
View file

@ -222,7 +222,7 @@
"limit": 38
},
"RET504": {
"limit": 721
"limit": 719
},
"RUF010": {
"limit": 874
@ -324,7 +324,7 @@
"limit": 883
},
"UP006": {
"limit": 12792
"limit": 12789
},
"UP007": {
"limit": 2570

55
tests/e2e/a2a/a2a_client.py
View file

@ -85,6 +85,7 @@ class A2ABridgeParams(BaseModel):
custom_llm_provider: str
model: str
api_key: str | None = None
class AgentRegisterBody(BaseModel):
@ -190,11 +191,52 @@ class A2ATaskStatus(BaseModel):
message: A2AResponseMessage | None = None
class A2AListingLocation(BaseModel):
"""Only the location fields a test reads back off a returned listing."""
un_locode: str | None = None
class A2AListing(BaseModel):
"""A single property card from the agent's `search_results` artifact; only the
identity/location fields a test asserts on are modelled."""
raia_id: str
property_type: str | None = None
service_type: str | None = None
location: A2AListingLocation = A2AListingLocation()
class A2ASearchResults(BaseModel):
"""The DataPart payload the property agent's `search_properties` skill returns:
the run count plus the listing cards themselves. Proof the tool actually ran and
matched, not just that the task completed with some text."""
total: int
count: int
listings: list[A2AListing] = []
class A2AArtifactPart(BaseModel):
kind: str | None = None
data: A2ASearchResults | None = None
class A2AArtifact(BaseModel):
model_config = ConfigDict(populate_by_name=True)
artifact_id: str | None = Field(default=None, alias="artifactId")
name: str | None = None
parts: list[A2AArtifactPart] = []
class A2AResult(BaseModel):
"""A message/send result. In 0.3 the message fields sit directly on the result
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that
runs a task replies with a `task` whose agent text lives on `status.message`.
`text` reads the agent's reply from whichever shape the served version produced."""
runs a task replies with a `task` whose agent text lives on `status.message` and
whose tool output lives on `artifacts`. `text` reads the agent's reply from
whichever shape the served version produced; `search_results` reads the tool's
structured output when the agent ran a skill."""
model_config = ConfigDict(populate_by_name=True)
@ -204,6 +246,7 @@ class A2AResult(BaseModel):
parts: list[A2AResponsePart] = []
message: A2AResponseMessage | None = None
status: A2ATaskStatus | None = None
artifacts: list[A2AArtifact] = []
@property
def text(self) -> str:
@ -221,6 +264,14 @@ class A2AResult(BaseModel):
def is_nested_v1_shape(self) -> bool:
return self.message is not None
@property
def search_results(self) -> A2ASearchResults | None:
for artifact in self.artifacts:
for part in artifact.parts:
if part.data is not None:
return part.data
return None
class A2AError(BaseModel):
code: int

16
tests/e2e/a2a/test_a2a_agent_e2e.py
View file

@ -32,7 +32,11 @@ from e2e_config import unique_marker
from e2e_http import Result, UnknownApiError, unwrap
from lifecycle import ResourceManager
BRIDGE = A2ABridgeParams(custom_llm_provider="anthropic", model="claude-haiku-4-5")
BRIDGE = A2ABridgeParams(
custom_llm_provider="anthropic",
model="claude-haiku-4-5",
api_key="os.environ/ANTHROPIC_API_KEY",
)
MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json"
MOVEHOME_ORIGIN = "https://movehome.org"
@ -113,6 +117,7 @@ class TestA2AAgentLifecycle:
agent = unwrap(client.register_agent(body))
resources.defer(lambda: client.delete_agent(agent.agent_id))
assert agent.agent_card_params.protocol_version == "0.3"
location = "GBLON"
request = A2AJsonRpcRequest(
id=f"e2e-{unique_marker()}",
params=A2AMessageSendParams(
@ -121,7 +126,7 @@ class TestA2AAgentLifecycle:
A2ADataPart(
data=A2ASkillInvocation(
skill="search_properties",
params=A2ASearchPropertiesParams(un_locode="USSFO", service_type="sale", asking_price_max=2_000_000, limit=3),
params=A2ASearchPropertiesParams(un_locode=location, service_type="long_term", limit=3),
)
)
],
@ -132,7 +137,12 @@ class TestA2AAgentLifecycle:
response = unwrap(client.send_message(agent.agent_id, scoped_key, request))
assert response.error is None
assert response.result is not None
assert response.result.text.strip() != ""
results = response.result.search_results
assert results is not None, "agent returned no search_results artifact; skill did not run"
assert results.total > 0
assert results.listings, "search_properties matched nothing; agent returned no property cards"
assert all(listing.raia_id for listing in results.listings)
assert all(listing.location.un_locode == location for listing in results.listings)
@pytest.mark.covers("other.a2a.discovery.proxy_fronted_card")
def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:

13
tests/e2e/batches/test_batches_e2e.py
View file

@ -23,7 +23,7 @@ from typing import Callable
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from batch_client import (
UPLOAD_FILENAME,
@ -702,14 +702,7 @@ class TestBedrockBatchAssumeRole:
def test_unified_batch_create_with_assume_role(
self, client: BatchClient, resources: ResourceManager
) -> None:
(role_arn,) = require_env("AWS_ROLE_NAME")
require_env(
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_REGION",
"AWS_BATCH_S3_BUCKET",
"AWS_BATCH_ROLE_ARN",
)
role_arn = os.environ["AWS_ROLE_NAME"]
session_name = f"e2e-batch-sts-{unique_marker()}"[:64]
model_name = batch_model_name("bedrock-sts-batch")
@ -819,7 +812,7 @@ class TestHostedVllmBatch:
def test_unified_file_and_batch_create(
self, client: BatchClient, resources: ResourceManager
) -> None:
(api_base,) = require_env("HOSTED_VLLM_API_BASE")
api_base = os.environ["HOSTED_VLLM_API_BASE"]
api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None
model_id = (
os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct"

16
tests/e2e/e2e_config.py
View file

@ -102,22 +102,6 @@ ANOMALY_SPEND_SETTLE_SECONDS = float(
)
def require_env(*names: str) -> tuple[str, ...]:
"""Return the non-empty values for each env name, or hard-fail naming which are missing.
Live e2e never skips for missing credentials: a missing key is a red run so
ops knows the suite cannot prove the product path.
"""
missing = tuple(name for name in names if not (os.environ.get(name) or "").strip())
if missing:
joined = ", ".join(missing)
raise AssertionError(
f"missing required env for e2e: {joined}. "
"Add them to tests/e2e/.env locally and to litellm ops for stage/CI."
)
return tuple((os.environ.get(name) or "").strip() for name in names)
def datadog_mcp_url(*, toolsets: str = "core") -> str:
"""Regional Datadog remote MCP endpoint for this process's DD_SITE.

11
tests/e2e/guardrails/test_bedrock_guardrail_e2e.py
View file

@ -8,9 +8,11 @@ a 200 means the guardrail never ran.
from __future__ import annotations
import os
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import UnknownApiError
from guardrails_client import GuardrailsClient
from lifecycle import ResourceManager
@ -33,11 +35,8 @@ class TestBedrockGuardrail:
def test_bedrock_pre_call_blocks_harmful_prompt(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
(identifier, version) = require_env(
"BEDROCK_GUARDRAIL_IDENTIFIER",
"BEDROCK_GUARDRAIL_VERSION",
)
require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION")
identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
name = f"e2e-bedrock-guard-{unique_marker()}"
guardrail_id = client.create_bedrock_guardrail(

3
tests/e2e/guardrails/test_block_code_execution_guardrail_e2e.py
View file

@ -16,7 +16,7 @@ from __future__ import annotations
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import unwrap
from guardrails_client import BlockCodeExecutionParamsBody, GuardrailsClient
from lifecycle import ResourceManager
@ -46,7 +46,6 @@ class TestBlockCodeExecutionGuardrail:
def test_blocks_execution_request_but_allows_explanation(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
require_env("GEMINI_API_KEY")
model = client.create_backend_model(resources, prefix="e2e-blockcode-backend")
name = f"e2e-block-code-{unique_marker()}"

3
tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py
View file

@ -14,7 +14,7 @@ from __future__ import annotations
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import UnknownApiError, unwrap
from guardrails_client import GuardrailsClient, OpenAIModerationParamsBody
from lifecycle import ResourceManager
@ -34,7 +34,6 @@ class TestOpenAIModerationGuardrail:
def test_moderation_blocks_flagged_input(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
require_env("OPENAI_API_KEY", "GEMINI_API_KEY")
model = client.create_backend_model(resources, prefix="e2e-moderation-backend")
name = f"e2e-openai-moderation-{unique_marker()}"

11
tests/e2e/guardrails/test_presidio_guardrail_e2e.py
View file

@ -25,11 +25,12 @@ The chat backend is a gemini deployment created for the test.
from __future__ import annotations
import os
import time
import pytest
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, require_env, unique_marker
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker
from e2e_http import NoBody, require_successful_call, unwrap
from guardrails_client import GuardrailMode, GuardrailsClient, PresidioParamsBody
from lifecycle import ResourceManager
@ -88,9 +89,8 @@ def _poll_logged_prompt(reader: OtelReader, *, call_id: str, genai_span: str) ->
def _presidio_params(
mode: GuardrailMode, *, apply_to_output: bool = False, logging_only: bool = False
) -> PresidioParamsBody:
analyzer, anonymizer = require_env(
"PRESIDIO_ANALYZER_API_BASE", "PRESIDIO_ANONYMIZER_API_BASE"
)
analyzer = os.environ["PRESIDIO_ANALYZER_API_BASE"]
anonymizer = os.environ["PRESIDIO_ANONYMIZER_API_BASE"]
return PresidioParamsBody(
mode=mode,
default_on=False,
@ -124,7 +124,6 @@ class TestPresidioGuardrail:
def test_pre_call_masks_pii_before_the_model_sees_it(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
require_env("GEMINI_API_KEY")
model = client.create_backend_model(resources, prefix="e2e-presidio-pre")
name = f"e2e-presidio-pre-{unique_marker()}"
guardrail_id = client.register(name, _presidio_params("pre_call"))
@ -149,7 +148,6 @@ class TestPresidioGuardrail:
def test_post_call_masks_pii_in_model_output(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
require_env("GEMINI_API_KEY")
model = client.create_backend_model(resources, prefix="e2e-presidio-post")
name = f"e2e-presidio-post-{unique_marker()}"
guardrail_id = client.register(name, _presidio_params("post_call", apply_to_output=True))
@ -173,7 +171,6 @@ class TestPresidioGuardrail:
def test_logging_only_masks_the_logged_prompt(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
require_env("GEMINI_API_KEY")
_require_otel_v2_active(client)
reader = build_otel_reader()

15
tests/e2e/llm_translation/test_chat_completions_regression_e2e.py
View file

@ -21,7 +21,7 @@ import os
import pytest
from pydantic import BaseModel
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import StreamingResponse, unwrap
from lifecycle import ResourceManager
from models import (
@ -250,7 +250,7 @@ class TestCohereChat:
def test_cohere_chat_returns_content(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
(cohere_key,) = require_env("COHERE_API_KEY")
cohere_key = os.environ["COHERE_API_KEY"]
model = f"e2e-cohere-chat-{unique_marker()}"
model_id = client.proxy.create_model(
model,
@ -343,7 +343,7 @@ class TestHostedVllmChat:
def test_hosted_vllm_chat_returns_content(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
(api_base,) = require_env("HOSTED_VLLM_API_BASE")
api_base = os.environ["HOSTED_VLLM_API_BASE"]
api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None
backend = (
os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct"
@ -395,7 +395,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_streams_real_content(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-chat-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -423,7 +422,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_logs_cost(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-cost-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -457,7 +455,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_returns_tool_call(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-tool-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -488,7 +485,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_structured_output_conforms_to_schema(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-schema-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -522,7 +518,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_reasoning_reports_reasoning_tokens(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-reasoning-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -561,7 +556,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_vision_describes_image(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-vision-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_VISION_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -579,7 +573,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_prompt_cache_hits_on_repeat(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-cache-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -610,7 +603,6 @@ class TestOpenAIChatCompletions:
def test_openai_chat_streams_tool_call(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
require_env("OPENAI_API_KEY")
model = f"e2e-openai-tool-stream-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY")
@ -645,7 +637,6 @@ class TestBedrockConverseChatCompletions:
"""
def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str:
require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION")
model = f"{prefix}-{unique_marker()}"
model_id = client.proxy.create_model(model, _bedrock_params())
resources.defer(lambda: client.proxy.delete_model(model_id))

3
tests/e2e/llm_translation/test_image_generation_e2e.py
View file

@ -9,7 +9,7 @@ from __future__ import annotations
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call
from endpoints_client import EndpointsClient, ImagesResult
from lifecycle import ResourceManager
@ -50,7 +50,6 @@ class TestImageGeneration:
def test_bedrock_image_generation_returns_image(
self, endpoints_client: EndpointsClient, resources: ResourceManager
) -> None:
require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION")
model = f"e2e-bedrock-image-{unique_marker()}"
model_id = endpoints_client.create_model(
model,

3
tests/e2e/llm_translation/test_messages_e2e.py
View file

@ -10,7 +10,7 @@ from __future__ import annotations
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call, unwrap
from endpoints_client import EndpointsClient, MessagesResult
from lifecycle import ResourceManager
@ -73,7 +73,6 @@ class TestAnthropicMessages:
def test_messages_logs_cost_matching_the_response_header(
self, endpoints_client: EndpointsClient, resources: ResourceManager
) -> None:
require_env("ANTHROPIC_API_KEY")
model = f"e2e-messages-cost-{unique_marker()}"
model_id = endpoints_client.create_model(
model,

3
tests/e2e/llm_translation/test_rerank_e2e.py
View file

@ -9,7 +9,7 @@ from __future__ import annotations
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call
from endpoints_client import EndpointsClient, RerankResult
from lifecycle import ResourceManager
@ -56,7 +56,6 @@ class TestRerank:
def test_bedrock_rerank_scores_top_n(
self, endpoints_client: EndpointsClient, resources: ResourceManager
) -> None:
require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION")
model = f"e2e-bedrock-rerank-{unique_marker()}"
model_id = endpoints_client.create_model(
model,

4
tests/e2e/llm_translation/test_responses_e2e.py
View file

@ -13,7 +13,7 @@ from typing import cast
import pytest
from pydantic import BaseModel, ValidationError
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call
from endpoints_client import (
EndpointsClient,
@ -255,7 +255,6 @@ class TestResponses:
def test_responses_bedrock_returns_completion(
self, endpoints_client: EndpointsClient, resources: ResourceManager
) -> None:
require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION")
model = f"e2e-responses-{unique_marker()}"
model_id = endpoints_client.create_model(model, _bedrock_params())
resources.defer(lambda: endpoints_client.delete_model(model_id))
@ -270,7 +269,6 @@ class TestResponses:
def test_responses_bedrock_returns_function_call(
self, endpoints_client: EndpointsClient, resources: ResourceManager
) -> None:
require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION")
model = f"e2e-responses-{unique_marker()}"
model_id = endpoints_client.create_model(model, _bedrock_params())
resources.defer(lambda: endpoints_client.delete_model(model_id))

4
tests/e2e/llm_translation/test_responses_metadata_e2e.py
View file

@ -14,7 +14,7 @@ import time
import pytest
from pydantic import BaseModel, ConfigDict
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call
from endpoints_client import EndpointsClient, ResponsesResult
from lifecycle import ResourceManager
@ -42,7 +42,7 @@ class RedisKeyInfo(BaseModel):
def _redis_scan(marker: str) -> tuple[RedisKeyInfo, ...]:
import redis
(host,) = require_env("REDIS_HOST")
host = os.environ["REDIS_HOST"]
port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379")
try:
with socket.create_connection((host, port), timeout=3):

4
tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py
View file

@ -11,7 +11,7 @@ import socket
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call
from lifecycle import ResourceManager
from models import KeyGenerateBody, LiteLLMParamsBody
@ -23,7 +23,7 @@ BACKEND = "anthropic/claude-haiku-4-5-20251001"
def _require_redis_reachable() -> None:
(host,) = require_env("REDIS_HOST")
host = os.environ["REDIS_HOST"]
port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379")
try:
with socket.create_connection((host, port), timeout=3):

4
tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py
View file

@ -13,7 +13,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from e2e_config import require_env, unique_marker
from e2e_config import unique_marker
from e2e_http import require_successful_call
from lifecycle import ResourceManager
from models import KeyGenerateBody, LiteLLMParamsBody
@ -28,7 +28,7 @@ RECOVERY_TIMEOUT = float(
def _require_redis() -> None:
(host,) = require_env("REDIS_HOST")
host = os.environ["REDIS_HOST"]
port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379")
try:
with socket.create_connection((host, port), timeout=3):

34
tests/e2e/ui/tests/modelsPage/addModel.spec.ts
View file

@ -1,5 +1,5 @@
import { test, expect } from "@playwright/test";
import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID } from "../../constants";
import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants";
import { Role, users } from "../../fixtures/users";
import { navigateToPage } from "../../helpers/navigation";
import { Page } from "../../fixtures/pages";
@ -56,17 +56,17 @@ test.describe("Add Model", () => {
},
});
expect(createResponse.ok()).toBe(true);
const createdModelId = (await createResponse.json()).model_info?.id;
expect(createdModelId, "model id from /model/new").toBeTruthy();
// Navigate to Models + Endpoints
await page.goto("/ui");
await page.getByText("Models + Endpoints").click();
// Click the new model row to open its detail view. The table renders
// a clickable outer row plus a nested detail row for the same model,
// so we target the first match (outer row) explicitly.
const modelRow = page.locator("tr", { hasText: modelName }).first();
await expect(modelRow).toBeVisible({ timeout: 10_000 });
await modelRow.click();
// The Model ID cell is the drill-in control; the row itself is not clickable.
const modelIdCell = page.getByTestId(`model-id-${createdModelId}`);
await expect(modelIdCell).toBeVisible({ timeout: 10_000 });
await modelIdCell.click();
await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000 });
@ -137,11 +137,11 @@ test.describe("Add Model", () => {
await page.waitForTimeout(2000);
// Search for the model we just added
await page.locator('input[placeholder="Search model names..."]').fill("claude-haiku-4-5");
await page.getByPlaceholder("Search model names").fill("claude-haiku-4-5");
await page.waitForTimeout(1000);
// Verify the model appears in the results count (not "Showing 0 results")
await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, {
await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, {
timeout: 15_000,
});
@ -228,24 +228,24 @@ test.describe("Add Model", () => {
// searching.
await page.waitForTimeout(2000);
await page.locator('input[placeholder="Search model names..."]').fill("cohere");
await page.getByPlaceholder("Search model names").fill("cohere");
await page.waitForTimeout(1000);
// Confirm the search returned at least one result — gives a clear
// failure message when the table is empty instead of timing out on a
// row assertion.
await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, {
await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, {
timeout: 15_000,
});
// Stronger than "alias appears somewhere in tbody" — pin the assertion
// Stronger than "the team appears somewhere in tbody" — pin the assertion
// to a single row that has BOTH the cohere model_name AND the seeded
// team alias, so a stale cohere row from "Add wildcard route" (no team)
// can't satisfy the check.
// team, so a stale cohere row from "Add wildcard route" (no team) can't
// satisfy the check. The Team ID column renders the id, not the alias.
const teamCohereRow = page
.locator("table tbody tr")
.filter({ hasText: "cohere/" })
.filter({ hasText: E2E_TEAM_CRUD_ALIAS });
.filter({ hasText: E2E_TEAM_CRUD_ID });
await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 });
} finally {
await deleteTeamScopedCohereModels();
@ -281,11 +281,11 @@ test.describe("Add Model", () => {
await page.waitForTimeout(2000);
// Search for the wildcard model
await page.locator('input[placeholder="Search model names..."]').fill("cohere");
await page.getByPlaceholder("Search model names").fill("cohere");
await page.waitForTimeout(1000);
// Verify the model appears in the results count (not "Showing 0 results")
await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, {
await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, {
timeout: 15_000,
});

7
tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts
View file

@ -67,9 +67,10 @@ test.describe("Clear custom pricing on a deployment", () => {
await page.goto("/ui");
await page.getByText("Models + Endpoints").click();
const modelRow = page.locator("tr", { hasText: modelName }).first();
await expect(modelRow).toBeVisible({ timeout: 15_000 });
await modelRow.click();
// The Model ID cell is the drill-in control; the row itself is not clickable.
const modelIdCell = page.getByTestId(`model-id-${createdModelId}`);
await expect(modelIdCell).toBeVisible({ timeout: 15_000 });
await modelIdCell.click();
await expect(page.getByText("Back to Models").first()).toBeVisible({
timeout: 10_000,
});

7
tests/llm_translation/reasoning_effort_grid/grid_spec.py
View file

@ -167,6 +167,13 @@ ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = (
"once the model is available."
),
),
ModelEntry(
alias="claude-opus-5",
model="anthropic/claude-opus-5",
mode="adaptive",
required_env=_ANTHROPIC_REQ,
caps=_CAPS_XHIGH_MAX,
),
ModelEntry(
alias="claude-opus-4-8",
model="anthropic/claude-opus-4-8",

6
tests/llm_translation/test_openai.py
View file

@ -518,7 +518,7 @@ async def test_openai_codex_stream(sync_mode):
from litellm.main import stream_chunk_builder
kwargs = {
"model": "openai/gpt-5.2-codex",
"model": "openai/gpt-5.3-codex",
"messages": [{"role": "user", "content": "Hey!"}],
"stream": True,
}
@ -550,7 +550,7 @@ async def test_openai_codex(sync_mode):
{
"model_name": "openai-codex-mini-latest",
"litellm_params": {
"model": "openai/gpt-5.2-codex",
"model": "openai/gpt-5.3-codex",
},
}
]
@ -838,7 +838,7 @@ def test_gpt_5_reasoning_streaming():
def test_openai_gpt_5_codex_reasoning():
litellm._turn_on_debug()
completion_kwargs = {
"model": "gpt-5-codex",
"model": "gpt-5.3-codex",
"messages": [
{
"role": "system",

461
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Union
from unittest.mock import Mock
import pytest
from fastapi import Request
from fastapi import HTTPException, Request
from starlette.datastructures import State
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.utils import _get_docs_url, _get_openapi_url, _get_redoc_url
from litellm.types.guardrails import GuardrailEventHooks
sys.path.insert(
0, os.path.abspath("../..")
@ -2638,6 +2640,463 @@ async def test_during_call_hook_parallel_execution_with_error():
litellm.callbacks = original_callbacks
class _PreCallGuardrail(CustomGuardrail):
"""Test double for pre_call guardrails; records timing and observed payload."""
def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True):
super().__init__(
guardrail_name=name,
event_hook=GuardrailEventHooks.pre_call,
default_on=default_on,
run_in_parallel=run_in_parallel,
)
self.name = name
self.sleep = sleep
self.execution_order = execution_order
self.observed_content = None
self.was_called = False
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.was_called = True
self.observed_content = data["messages"][0]["content"]
self.execution_order.append(f"{self.name}_start")
await asyncio.sleep(self.sleep)
self.execution_order.append(f"{self.name}_end")
return None
@pytest.mark.asyncio
async def test_pre_call_hook_runs_opted_in_guardrails_in_parallel():
"""run_in_parallel pre_call guardrails execute concurrently (all start before any ends)."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PreCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3)
]
result = await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item)
starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item)
assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}"
assert result["model"] == "gpt-4"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_runs_default_guardrails_sequentially():
"""Guardrails without run_in_parallel keep the sequential, one-at-a-time behavior."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PreCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2)
]
start = asyncio.get_event_loop().time()
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
elapsed = asyncio.get_event_loop().time() - start
assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"]
assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_sequential_mutations_precede_parallel_batch():
"""Sequential (mutating) guardrails run before the parallel batch, which sees their changes."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class MaskingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="masker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=False,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
data["messages"][0]["content"] = "MASKED"
return data
parallel_observer = _PreCallGuardrail("observer", run_in_parallel=True, execution_order=execution_order)
try:
litellm.callbacks = [parallel_observer, MaskingGuardrail()]
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "secret"}]},
call_type="completion",
)
assert parallel_observer.observed_content == "MASKED"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_guardrail_blocks_request():
"""A raising parallel guardrail blocks the request before it reaches the LLM."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class BlockingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
raise HTTPException(status_code=400, detail="blocked by guardrail")
try:
litellm.callbacks = [BlockingGuardrail()]
with pytest.raises(HTTPException) as exc_info:
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert exc_info.value.status_code == 400
assert "blocked by guardrail" in str(exc_info.value.detail)
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_guardrail_skipped_when_should_not_run():
"""A parallel guardrail that should_run_guardrail rejects is never invoked."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
guardrail = _PreCallGuardrail(
"off_by_default", run_in_parallel=True, execution_order=execution_order, default_on=False
)
litellm.callbacks = [guardrail]
result = await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert guardrail.was_called is False
assert result["model"] == "gpt-4"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_block_wins_over_reroute():
"""A slower block must win over a faster reroute so crafted input cannot bypass a block."""
from litellm.caching.caching import DualCache
from litellm.exceptions import SensitiveDataRouteException
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class FastRerouteGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="rerouter",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
raise SensitiveDataRouteException(route_to_model="on-prem", session_id="s1", guardrail_name="rerouter")
class SlowBlockingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
await asyncio.sleep(0.1)
raise HTTPException(status_code=400, detail="blocked by guardrail")
try:
litellm.callbacks = [FastRerouteGuardrail(), SlowBlockingGuardrail()]
with pytest.raises(HTTPException) as exc_info:
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert exc_info.value.status_code == 400
assert "blocked by guardrail" in str(exc_info.value.detail)
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_awaits_all_when_one_blocks():
"""A block must not orphan sibling guardrails; every parallel guardrail runs to completion."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
completed = []
class FastBlockingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="fast_blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
raise HTTPException(status_code=400, detail="blocked")
class SlowGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="slow",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
await asyncio.sleep(0.1)
completed.append("slow")
return None
try:
litellm.callbacks = [FastBlockingGuardrail(), SlowGuardrail()]
with pytest.raises(HTTPException):
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert completed == ["slow"], "slow guardrail was orphaned instead of awaited to completion"
finally:
litellm.callbacks = original_callbacks
class _PostCallGuardrail(CustomGuardrail):
"""Test double for post_call guardrails; records timing and invocation."""
def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True):
super().__init__(
guardrail_name=name,
event_hook=GuardrailEventHooks.post_call,
default_on=default_on,
run_in_parallel=run_in_parallel,
)
self.name = name
self.sleep = sleep
self.execution_order = execution_order
self.was_called = False
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
self.was_called = True
self.execution_order.append(f"{self.name}_start")
await asyncio.sleep(self.sleep)
self.execution_order.append(f"{self.name}_end")
return None
@pytest.mark.asyncio
async def test_post_call_hook_runs_opted_in_guardrails_in_parallel():
"""run_in_parallel post_call guardrails execute concurrently (all start before any ends)."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PostCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3)
]
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item)
starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item)
assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_post_call_hook_runs_default_guardrails_sequentially():
"""post_call guardrails without run_in_parallel keep the sequential behavior."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PostCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2)
]
start = asyncio.get_event_loop().time()
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
elapsed = asyncio.get_event_loop().time() - start
assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"]
assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_post_call_hook_parallel_guardrail_blocks_response():
"""A raising parallel post_call guardrail blocks the response before it reaches the client."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class BlockingPostCallGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="post_blocker",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
raise HTTPException(status_code=400, detail="blocked response by guardrail")
try:
litellm.callbacks = [BlockingPostCallGuardrail()]
with pytest.raises(HTTPException) as exc_info:
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
assert exc_info.value.status_code == 400
assert "blocked response by guardrail" in str(exc_info.value.detail)
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_post_call_hook_parallel_awaits_all_when_one_blocks():
"""A blocking post_call guardrail must not orphan its siblings; all run to completion."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
completed = []
class FastBlockingPostCall(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="fast_post_blocker",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
raise HTTPException(status_code=400, detail="blocked")
class SlowPostCall(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="slow_post",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
await asyncio.sleep(0.1)
completed.append("slow")
return None
try:
litellm.callbacks = [FastBlockingPostCall(), SlowPostCall()]
with pytest.raises(HTTPException):
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
assert completed == ["slow"], "slow post_call guardrail was orphaned instead of awaited to completion"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_handle_logging_proxy_only_error_preserves_pass_through_call_type():
"""Ensure _handle_logging_proxy_only_error does not overwrite call_type

136
tests/test_litellm/integrations/otel/test_otel_v2_logger.py
View file

@ -29,7 +29,9 @@ from litellm.integrations.otel import ( # noqa: E402
from litellm.integrations.otel.plumbing import providers # noqa: E402
from litellm.integrations.otel.plumbing.context import ( # noqa: E402
reset_mcp_message_trace_carrier,
reset_mcp_message_transport_span_context,
set_mcp_message_trace_carrier,
set_mcp_message_transport_span_context,
set_request_root_span,
)
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
@ -56,9 +58,11 @@ def _reset_request_root_span():
_otel_context._request_root_span.set(None)
_otel_context._mcp_message_trace_carrier.set(None)
_otel_context._mcp_message_transport_span_context.set(None)
yield
_otel_context._request_root_span.set(None)
_otel_context._mcp_message_trace_carrier.set(None)
_otel_context._mcp_message_transport_span_context.set(None)
def _payload(**overrides):
@ -528,15 +532,15 @@ _MCP_SPAN_CASES = [
@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES)
def test_mcp_span_roots_and_links_transport_without_propagated_context(
def test_mcp_span_nests_under_transport_without_propagated_context(
make_payload, span_name
):
"""MCP and the HTTP transport are independent lifecycles (one streamable-HTTP
session multiplexes many messages), so per the MCP semconv the message span
must NOT nest under the session/transport span — that is what made it render
skewed at the session's start. With no propagated ``params._meta`` context it
starts its own root trace and records the transport span as a *link*, never
the parent."""
"""Almost no MCP client implements SEP-414, so ``params._meta`` normally carries
no trace context. Rooting the span there split one tool call into two traces
joined only by a link, which is how it surfaced in APM: the ``POST`` transaction
and the ``tools/call`` span shared no ``trace_id``. With no remote parent to
honor the span nests under the transport span instead, and records no link since
the transport is now the real parent."""
logger, exporter = _logger()
transport = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
@ -549,11 +553,75 @@ def test_mcp_span_roots_and_links_transport_without_propagated_context(
)
transport.end()
span = next(s for s in exporter.get_finished_spans() if s.name == span_name)
assert span.parent is not None
assert span.parent.span_id == transport.get_span_context().span_id
assert span.context.trace_id == transport.get_span_context().trace_id
assert span.links == ()
@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES)
def test_mcp_span_nests_under_this_messages_transport_not_the_session_opener(
make_payload, span_name
):
"""A *stateful* streamable-HTTP session runs every message on the single task
spawned by that session's ``initialize`` POST, so the ``_request_root_span``
ContextVar the ASGI request task writes is frozen at ``initialize`` inside the
handler and never sees the later ``tools/call`` POST. Nesting on that anchor
would hang every tool call of the session off the first request's (already
ended) span, rendering skewed at the session's start. The gateway resolves the
current message's transport on the request task and publishes it, so the span
parents to the POST that actually carried this message."""
logger, exporter = _logger()
session_opener = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
)
this_message = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
)
async def session_task():
token = set_mcp_message_transport_span_context(
this_message.get_span_context()
)
try:
await logger.async_log_success_event(
{"standard_logging_object": make_payload()}, None, None, None
)
finally:
reset_mcp_message_transport_span_context(token)
async def initialize_request():
# The anchor the session task inherits is the one ``initialize`` left behind;
# spawning here reproduces the SDK's session task, which outlives this request.
set_request_root_span(session_opener)
await asyncio.create_task(session_task())
asyncio.run(initialize_request())
session_opener.end()
this_message.end()
span = next(s for s in exporter.get_finished_spans() if s.name == span_name)
assert span.parent is not None
assert span.parent.span_id == this_message.get_span_context().span_id
assert span.context.trace_id == this_message.get_span_context().trace_id
assert span.parent.span_id != session_opener.get_span_context().span_id
assert span.context.trace_id != session_opener.get_span_context().trace_id
@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES)
def test_mcp_span_roots_without_transport_or_propagated_context(
make_payload, span_name
):
"""With neither a remote parent nor a transport span there is nothing to nest
under, so the span legitimately starts its own root trace with no links."""
logger, exporter = _logger()
asyncio.run(
logger.async_log_success_event(
{"standard_logging_object": make_payload()}, None, None, None
)
)
span = next(s for s in exporter.get_finished_spans() if s.name == span_name)
assert span.parent is None
assert span.context.trace_id != transport.get_span_context().trace_id
assert [link.context.span_id for link in span.links] == [
transport.get_span_context().span_id
]
assert span.links == ()
@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES)
@ -641,10 +709,11 @@ def test_mcp_span_carries_authenticated_identity(make_payload, span_name):
assert span.attributes[LiteLLM.TEAM_ID] == "t1"
def test_mcp_span_malformed_traceparent_starts_root():
def test_mcp_span_malformed_traceparent_nests_under_transport():
"""A malformed traceparent in ``params._meta`` must not crash or parent to a
bogus span: the propagator ignores it, so the span starts its own root trace and
still links the transport span."""
bogus span: the propagator ignores it, leaving no remote parent, so the span
falls back to nesting under the transport span rather than starting a
disconnected root trace."""
logger, exporter = _logger()
transport = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
@ -661,9 +730,44 @@ def test_mcp_span_malformed_traceparent_starts_root():
reset_mcp_message_trace_carrier(token)
transport.end()
span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list")
assert span.parent is None
assert span.parent is not None
assert span.parent.span_id == transport.get_span_context().span_id
assert span.links == ()
def test_mcp_span_links_this_messages_transport_when_context_is_propagated():
"""On the semconv path the transport is recorded as a link, and that link must
point at the POST carrying this message too. Reading the stale session anchor
would attribute the tool call to whichever request opened the session."""
logger, exporter = _logger()
session_opener = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
)
this_message = logger._emitter.start_span(
SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME
)
set_request_root_span(session_opener)
trace_token = set_mcp_message_trace_carrier(
{"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"}
)
transport_token = set_mcp_message_transport_span_context(
this_message.get_span_context()
)
try:
asyncio.run(
logger.async_log_success_event(
{"standard_logging_object": _mcp_list_payload()}, None, None, None
)
)
finally:
reset_mcp_message_transport_span_context(transport_token)
reset_mcp_message_trace_carrier(trace_token)
session_opener.end()
this_message.end()
span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list")
assert span.parent is not None and span.parent.span_id == 0x2222222222222222
assert [link.context.span_id for link in span.links] == [
transport.get_span_context().span_id
this_message.get_span_context().span_id
]

33
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -1,3 +1,4 @@
import asyncio
from unittest.mock import AsyncMock
import pytest
@ -1394,6 +1395,38 @@ class TestEventTypeLogging:
assert len(logged_info) == 1
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
@pytest.mark.asyncio
async def test_log_guardrail_information_records_every_concurrent_guardrail(self):
"""Guardrails run concurrently (parallel pre_call/post_call, during_call) share one
request_data dict. Each must still record its own entry. The previous guard counted
entries in that shared dict, so a sibling's append made a guardrail think it had already
recorded and skip its own auto-record — silently dropping lifecycle logs the UI shows."""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class SleeperGuardrail(CustomGuardrail):
def __init__(self, name, sleep):
super().__init__(guardrail_name=name, event_hook=GuardrailEventHooks.pre_call)
self._sleep = sleep
@log_guardrail_information
async def async_pre_call_hook(self, data: dict, **kwargs):
await asyncio.sleep(self._sleep)
return data
request_data = {"metadata": {}}
# Different sleeps guarantee overlapping execution windows: the faster guardrail
# records while the slower one is still awaiting, which is exactly what tripped the
# old shared-count guard.
await asyncio.gather(
SleeperGuardrail("guardrail-a", 0.05).async_pre_call_hook(data=request_data),
SleeperGuardrail("guardrail-b", 0.15).async_pre_call_hook(data=request_data),
)
logged = request_data["metadata"]["standard_logging_guardrail_information"]
assert {entry["guardrail_name"] for entry in logged} == {"guardrail-a", "guardrail-b"}
assert len(logged) == 2
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
self,
):

90
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -2111,6 +2111,37 @@ def test_token_type_cost_breakdown_reads_cache_write_tokens():
)
def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6():
"""
Regression: OpenAI gpt-5.6 reports cache-write tokens under
prompt_tokens_details.cache_write_tokens (not the Anthropic cache_creation_tokens
name). Those tokens must be billed at the cache-write rate rather than the plain
input rate. Customer report: cache creation tokens were never counted for the
GPT-5.6 series, so cost was undercounted on cache-write requests.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gpt-5.6"
usage = Usage(
prompt_tokens=1000,
completion_tokens=10,
total_tokens=1010,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_write_tokens=800),
)
assert usage.prompt_tokens_details.cache_write_tokens == 800
assert usage.prompt_tokens_details.cache_creation_tokens == 800
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
info = litellm.get_model_info(model=model, custom_llm_provider="openai")
expected_prompt = (1000 - 800) * info["input_cost_per_token"] + 800 * info["cache_creation_input_token_cost"]
assert prompt_cost == pytest.approx(expected_prompt)
assert info["cache_creation_input_token_cost"] > info["input_cost_per_token"]
assert prompt_cost > 1000 * info["input_cost_per_token"]
def test_token_type_cost_breakdown_reconciles_with_generic_total():
"""
Both-ways check: the reasoning subset must sum with the remaining (text) output
@ -2166,6 +2197,65 @@ def test_token_type_cost_breakdown_zero_without_special_tokens():
)
@pytest.mark.parametrize(
"raw_usage, expect_read, expect_write",
[
(
{
"input_tokens": 5000,
"output_tokens": 10,
"total_tokens": 5010,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 4012},
},
False,
True,
),
(
{
"input_tokens": 5000,
"output_tokens": 10,
"total_tokens": 5010,
"input_tokens_details": {"cached_tokens": 4012, "cache_write_tokens": 0},
},
True,
False,
),
],
)
def test_token_type_cost_breakdown_openai_responses_api_cache_write_read(
raw_usage, expect_read, expect_write
):
"""Regression for #34309: OpenAI Responses API reports cache tokens under
input_tokens_details.{cached_tokens, cache_write_tokens}, not the Anthropic-style
top-level cache_creation_input_tokens. The itemized breakdown must still populate
cache_read_cost / cache_creation_cost from the transformed usage."""
from litellm.responses.utils import ResponseAPILoggingUtils
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gpt-5.6"
usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage)
breakdown = get_token_type_cost_breakdown(
model=model, custom_llm_provider="openai", usage=usage
)
info = litellm.get_model_info(model=model, custom_llm_provider="openai")
if expect_write:
assert breakdown.cache_creation_cost == pytest.approx(
4012 * info["cache_creation_input_token_cost"]
)
assert breakdown.cache_creation_cost > 0
assert breakdown.cache_read_cost == 0.0
if expect_read:
assert breakdown.cache_read_cost == pytest.approx(
4012 * info["cache_read_input_token_cost"]
)
assert breakdown.cache_read_cost > 0
assert breakdown.cache_creation_cost == 0.0
def test_token_type_cost_breakdown_handles_unknown_model_gracefully():
"""A model with no pricing must yield zeros, never raise."""
breakdown = get_token_type_cost_breakdown(

165
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -1111,6 +1111,171 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through
litellm._async_success_callback = original_async_callbacks
@pytest.mark.asyncio
async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handler(
logging_obj,
):
"""prefer_async_handlers must await async_failure_handler and never submit the sync failure_handler.
Submitting the sync ``failure_handler`` while awaiting ``async_failure_handler``
lets both mutate the shared logging_obj at once, which is the concurrent-mutation
crash this dispatch guard exists to prevent.
"""
exception = ValueError("boom")
traceback_exception = "traceback"
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(
logging_obj, "async_failure_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "failure_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_failure_callbacks_for_async_calls",
return_value=False,
),
patch(
"litellm.litellm_core_utils.litellm_logging.executor.submit"
) as mock_submit,
):
await logging_obj.dispatch_failure_handlers(
exception,
traceback_exception,
prefer_async_handlers=True,
)
mock_async.assert_awaited_once_with(exception, traceback_exception)
mock_sync.assert_not_called()
mock_submit.assert_not_called()
@pytest.mark.asyncio
async def test_dispatch_failure_handlers_async_completes_before_sync_submit(
logging_obj,
):
"""The async failure handler must fully finish before the legacy sync handler is scheduled.
Ordering proves there is no window where both handlers touch the shared
logging_obj concurrently: the sync submit only happens after the await returns.
"""
exception = ValueError("boom")
traceback_exception = "traceback"
events: list[str] = []
async def _async_failure(exc, tb, **kwargs):
events.append("async_start")
await asyncio.sleep(0)
events.append("async_end")
def _submit(*args, **kwargs):
events.append("sync_submit")
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure),
patch.object(logging_obj, "failure_handler", new_callable=MagicMock),
patch.object(
logging_obj,
"_should_run_sync_failure_callbacks_for_async_calls",
return_value=True,
),
patch(
"litellm.litellm_core_utils.litellm_logging.executor.submit",
side_effect=_submit,
),
):
await logging_obj.dispatch_failure_handlers(
exception,
traceback_exception,
prefer_async_handlers=True,
)
assert events == ["async_start", "async_end", "sync_submit"]
@pytest.mark.asyncio
async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks(
logging_obj,
):
"""A sync failure callback must still run when only failure callbacks are configured.
The legacy thread-based path always submitted the sync failure_handler, so gating it on
the success callback list would silently drop failure logging for any deployment that
registers only failure callbacks and no success callbacks. This drives the real predicate
(unmocked), so gating the sync failure handler on the success list fails this test.
"""
exception = ValueError("boom")
traceback_exception = "traceback"
def _sync_failure_callback(*args, **kwargs):
return None
logging_obj.model_call_details["litellm_params"] = {}
logging_obj.dynamic_success_callbacks = None
logging_obj.dynamic_failure_callbacks = None
with (
patch.object(litellm, "success_callback", []),
patch.object(litellm, "failure_callback", [_sync_failure_callback]),
patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock),
patch.object(
logging_obj, "failure_handler", new_callable=MagicMock
) as mock_sync,
patch(
"litellm.litellm_core_utils.litellm_logging.executor.submit"
) as mock_submit,
):
await logging_obj.dispatch_failure_handlers(
exception,
traceback_exception,
prefer_async_handlers=True,
)
mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception)
@pytest.mark.asyncio
async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inline(
logging_obj,
):
"""A sync-SDK request (prefer_async_handlers=False) runs failure_handler inline.
``async for`` over a stream from ``completion()`` passes prefer_async_handlers=True; a
plain sync request leaves it False, so the legacy sync handler runs directly and the
async handler is never awaited, matching dispatch_success_handlers.
"""
exception = ValueError("boom")
traceback_exception = "traceback"
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(
logging_obj, "async_failure_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "failure_handler", new_callable=MagicMock
) as mock_sync,
patch(
"litellm.litellm_core_utils.litellm_logging.executor.submit"
) as mock_submit,
):
await logging_obj.dispatch_failure_handlers(
exception,
traceback_exception,
prefer_async_handlers=False,
)
mock_sync.assert_called_once_with(exception, traceback_exception)
mock_async.assert_not_awaited()
mock_submit.assert_not_called()
def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj):
"""Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False."""
import datetime

12
tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
View file

@ -226,6 +226,18 @@ def test_get_output_file_id_empty_output_info_falls_through_to_output_config():
assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl"
def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config():
resp = {
"outputInfo": None,
"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}},
}
assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl"
def test_get_output_file_id_output_info_explicit_none_and_no_output_config():
assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == ""
def test_get_output_file_id_no_output_info_and_no_output_config():
assert T._get_output_file_id_from_vertex_ai_batch_response({}) == ""

116
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -2087,6 +2087,9 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped():
async def async_failure_handler(self, *args, **kwargs):
return None
async def dispatch_failure_handlers(self, *args, **kwargs):
return None
async def failing_make_call(client=None, **kwargs):
raise VertexAIError(status_code=400, message="bad input", headers={})
@ -5404,3 +5407,116 @@ def test_process_candidates_merges_thought_signatures_and_server_side_tools():
fields = model_response.choices[-1].message.provider_specific_fields
assert fields["thought_signatures"] == ["sig-text"]
assert fields["server_side_tool_invocations"][0]["id"] == "tool-1"
def _accumulating_gemini_iterator():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
iterator = ModelResponseIterator(
streaming_response=[], sync_stream=True, logging_obj=MagicMock()
)
iterator.chunk_type = "accumulated_json"
return iterator
def test_accumulated_json_chunk_multi_value_buffer_does_not_wedge():
"""Two complete Gemini objects buffered together must both surface.
A whole-buffer json.loads raises "Extra data" on concatenated values and, since
the buffer was never reset on failure, returned None forever while growing without
bound. Peeling one value from the front keeps the remainder for the next call.
"""
obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}'
iterator = _accumulating_gemini_iterator()
first = iterator.handle_accumulated_json_chunk(chunk=obj + obj)
assert first is not None
assert first.choices[0].delta.content == "a"
second = iterator.handle_accumulated_json_chunk(chunk="")
assert second is not None
assert second.choices[0].delta.content == "a"
assert iterator.accumulated_json.strip() == ""
def test_accumulated_json_end_of_stream_drains_all_buffered_values():
"""End of stream must drain every buffered value and then terminate.
With concatenated values a whole-buffer parse never succeeds, so __next__ kept
returning None without shrinking the buffer - an unrecoverable per-request spin.
The bounded loop asserts the iterator both surfaces all values and terminates.
"""
obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}'
iterator = _accumulating_gemini_iterator()
iterator.response_iterator = iter([])
iterator.accumulated_json = obj + obj + obj
out = []
terminated = False
for _ in range(100):
try:
chunk = iterator.__next__()
except StopIteration:
terminated = True
break
if chunk is not None:
out.append(chunk)
assert terminated, "iterator did not terminate - accumulated buffer wedged"
assert len(out) == 3
assert iterator.accumulated_json.strip() == ""
def test_accumulated_json_end_of_stream_surfaces_leading_value_before_truncated_tail():
"""A complete leading value must survive a truncated trailing value at end of stream.
The mid-stream perf guard only inspects the buffer's last byte, so a complete leading
object followed by a truncated one (a server that cut the stream mid-object, last byte
not a closer) would keep the guard from ever parsing and drop the complete value. At end
of stream the drain ignores that guard, surfaces the complete value, and discards only
the truncated tail.
"""
obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}'
iterator = _accumulating_gemini_iterator()
iterator.response_iterator = iter([])
iterator.accumulated_json = obj + '{"candidates":'
out = []
for _ in range(100):
try:
chunk = iterator.__next__()
except StopIteration:
break
if chunk is not None:
out.append(chunk)
assert len(out) == 1
assert out[0].choices[0].delta.content == "a"
def test_accumulated_json_skips_non_dict_leading_value():
"""A non-dict value at the front must not block the dict values behind it.
raw_decode advances past a decoded value, so a leading non-dict (a JSON array or scalar,
which Gemini never emits but a malformed stream could) must be consumed and skipped. If
the drain stopped on it, the trailing objects would be lost at end of stream.
"""
obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}'
iterator = _accumulating_gemini_iterator()
iterator.response_iterator = iter([])
iterator.accumulated_json = "[1, 2]" + obj
out = []
for _ in range(100):
try:
chunk = iterator.__next__()
except StopIteration:
break
if chunk is not None:
out.append(chunk)
assert len(out) == 1
assert out[0].choices[0].delta.content == "a"

539
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py
View file

@ -1,539 +0,0 @@
"""
Tests for OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers.
Covers: exchange flow, caching, error handling, resolve_mcp_auth integration,
bearer token extraction, and config loading.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from litellm.proxy._experimental.mcp_server.auth.token_exchange import (
TOKEN_EXCHANGE_GRANT_TYPE,
TokenExchangeHandler,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
resolve_mcp_auth,
)
from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def _obo_server(**overrides) -> MCPServer:
defaults = dict(
server_id="srv-obo-1",
name="test-obo",
url="https://mcp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
client_id="litellm-client-id",
client_secret="litellm-client-secret",
token_exchange_endpoint="https://idp.example.com/oauth2/token",
audience="api://mcp-server",
scopes=["mcp.tools.read", "mcp.tools.execute"],
)
defaults.update(overrides)
return MCPServer(**defaults)
def _exchange_response(token="exchanged-tok-abc", expires_in=3600):
resp = MagicMock()
resp.json.return_value = {
"access_token": token,
"token_type": "Bearer",
"expires_in": expires_in,
}
resp.raise_for_status = MagicMock()
resp.text = ""
return resp
# ── Exchange Flow ──
@pytest.mark.asyncio
async def test_exchange_token_success():
"""Token exchange sends correct RFC 8693 parameters and returns access_token."""
handler = TokenExchangeHandler()
server = _obo_server()
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("scoped-token-1")
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
result = await handler.exchange_token("user-jwt-xyz", server)
assert result == "scoped-token-1"
mock_client.post.assert_called_once()
_, kwargs = mock_client.post.call_args
data = kwargs["data"]
assert data["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE
assert data["subject_token"] == "user-jwt-xyz"
assert data["subject_token_type"] == "urn:ietf:params:oauth:token-type:access_token"
assert data["audience"] == "api://mcp-server"
assert data["scope"] == "mcp.tools.read mcp.tools.execute"
assert data["client_id"] == "litellm-client-id"
assert data["client_secret"] == "litellm-client-secret"
@pytest.mark.asyncio
async def test_exchange_token_no_audience():
"""When audience is None, it is omitted from the request."""
handler = TokenExchangeHandler()
server = _obo_server(audience=None)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response()
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
await handler.exchange_token("user-jwt", server)
_, kwargs = mock_client.post.call_args
assert "audience" not in kwargs["data"]
@pytest.mark.asyncio
async def test_exchange_token_no_scopes():
"""When scopes is None, scope param is omitted from the request."""
handler = TokenExchangeHandler()
server = _obo_server(scopes=None)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response()
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
await handler.exchange_token("user-jwt", server)
_, kwargs = mock_client.post.call_args
assert "scope" not in kwargs["data"]
# ── Caching ──
@pytest.mark.asyncio
async def test_exchange_token_cached():
"""Second call with same user token uses cache — only 1 HTTP POST."""
handler = TokenExchangeHandler()
server = _obo_server()
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("cached-exchange-tok")
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
t1 = await handler.exchange_token("same-jwt", server)
t2 = await handler.exchange_token("same-jwt", server)
assert t1 == t2 == "cached-exchange-tok"
assert mock_client.post.call_count == 1
@pytest.mark.asyncio
async def test_different_user_tokens_not_shared():
"""Different user JWTs get different exchanged tokens."""
handler = TokenExchangeHandler()
server = _obo_server()
call_count = 0
async def mock_post(url, data=None):
nonlocal call_count
call_count += 1
resp = MagicMock()
resp.json.return_value = {
"access_token": f"exchanged-{call_count}",
"expires_in": 3600,
}
resp.raise_for_status = MagicMock()
return resp
mock_client = AsyncMock()
mock_client.post = mock_post
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
t1 = await handler.exchange_token("user-a-jwt", server)
t2 = await handler.exchange_token("user-b-jwt", server)
assert t1 == "exchanged-1"
assert t2 == "exchanged-2"
assert call_count == 2
# ── Error Handling ──
@pytest.mark.asyncio
async def test_exchange_token_http_error():
"""HTTP errors from the IDP are wrapped in a ValueError."""
handler = TokenExchangeHandler()
server = _obo_server()
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = "invalid_grant"
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
"Bad Request",
request=MagicMock(),
response=mock_response,
)
mock_client = AsyncMock()
mock_client.post.return_value = mock_response
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
),
pytest.raises(ValueError, match="failed with status 400"),
):
await handler.exchange_token("bad-jwt", server)
@pytest.mark.asyncio
async def test_exchange_token_http_error_does_not_log_response_body():
"""Raw IDP error bodies are not logged because they can contain credentials."""
handler = TokenExchangeHandler()
server = _obo_server()
raw_response_body = "client_secret=do-not-log"
mock_response = MagicMock()
mock_response.status_code = 401
mock_response.text = raw_response_body
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
"Unauthorized",
request=MagicMock(),
response=mock_response,
)
mock_client = AsyncMock()
mock_client.post.return_value = mock_response
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
),
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.verbose_logger.debug"
) as mock_debug,
pytest.raises(ValueError, match="failed with status 401"),
):
await handler.exchange_token("bad-jwt", server)
logged_values = " ".join(
str(value)
for call in mock_debug.call_args_list
for value in [*call.args, *call.kwargs.values()]
)
assert raw_response_body not in logged_values
@pytest.mark.asyncio
async def test_exchange_token_missing_access_token():
"""Response without access_token raises ValueError."""
handler = TokenExchangeHandler()
server = _obo_server()
resp = MagicMock()
resp.json.return_value = {"token_type": "Bearer"}
resp.raise_for_status = MagicMock()
mock_client = AsyncMock()
mock_client.post.return_value = resp
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
),
pytest.raises(ValueError, match="missing 'access_token'"),
):
await handler.exchange_token("jwt", server)
@pytest.mark.asyncio
async def test_exchange_token_missing_endpoint():
"""Missing token_exchange_endpoint and token_url raises ValueError."""
handler = TokenExchangeHandler()
server = _obo_server(token_exchange_endpoint=None, token_url=None)
with pytest.raises(ValueError, match="no token_exchange_endpoint or token_url"):
await handler.exchange_token("jwt", server)
@pytest.mark.asyncio
async def test_exchange_token_missing_credentials():
"""Missing client_id or client_secret raises ValueError."""
handler = TokenExchangeHandler()
server = _obo_server(client_id=None, client_secret=None)
# has_token_exchange_config will be False, so we call _do_exchange directly
with pytest.raises(ValueError, match="missing client_id or client_secret"):
await handler._do_exchange("jwt", server)
# ── resolve_mcp_auth Integration ──
@pytest.mark.asyncio
async def test_resolve_mcp_auth_with_token_exchange():
"""resolve_mcp_auth delegates to token exchange when server has OBO config and subject_token provided."""
server = _obo_server()
mock_handler = AsyncMock()
mock_handler.exchange_token.return_value = "obo-scoped-token"
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.mcp_token_exchange_handler",
mock_handler,
):
result = await resolve_mcp_auth(server, subject_token="user-jwt")
assert result == "obo-scoped-token"
mock_handler.exchange_token.assert_called_once_with("user-jwt", server)
@pytest.mark.asyncio
async def test_resolve_mcp_auth_obo_without_subject_token_falls_through():
"""Without a subject_token, resolve_mcp_auth falls through to client_credentials."""
server = _obo_server(
token_url="https://auth.example.com/token",
)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("cc-token")
with patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
return_value=mock_client,
):
result = await resolve_mcp_auth(server, subject_token=None)
# Falls through to client_credentials since subject_token is None
# The server has client_id/client_secret/token_url so has_client_credentials is True
assert result == "cc-token"
@pytest.mark.asyncio
async def test_resolve_mcp_auth_obo_without_subject_token_uses_cached_client_credentials():
"""The M2M fallback for OBO servers reuses the client_credentials cache."""
server = _obo_server(
server_id="srv-obo-m2m-cache",
token_url="https://auth.example.com/token",
)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("cached-cc-token")
with patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
return_value=mock_client,
):
first = await resolve_mcp_auth(server, subject_token=None)
second = await resolve_mcp_auth(server, subject_token=None)
assert first == second == "cached-cc-token"
mock_client.post.assert_called_once()
@pytest.mark.asyncio
async def test_resolve_mcp_auth_header_beats_obo():
"""An explicit mcp_auth_header takes priority over OBO token exchange."""
server = _obo_server()
result = await resolve_mcp_auth(
server, mcp_auth_header="Bearer override", subject_token="user-jwt"
)
assert result == "Bearer override"
# ── Bearer Token Extraction ──
def test_extract_bearer_token_from_oauth2_headers():
"""Extracts token from oauth2_headers Authorization header."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers={"Authorization": "Bearer my-jwt-token"},
raw_headers=None,
)
assert result == "my-jwt-token"
def test_extract_bearer_token_from_raw_headers():
"""Falls back to raw_headers when oauth2_headers missing."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers=None,
raw_headers={"authorization": "Bearer raw-jwt"},
)
assert result == "raw-jwt"
def test_extract_bearer_token_no_bearer_prefix():
"""Returns token as-is when no Bearer prefix."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers={"Authorization": "some-opaque-token"},
raw_headers=None,
)
assert result == "some-opaque-token"
def test_extract_bearer_token_none():
"""Returns None when no auth headers present."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers=None,
raw_headers=None,
)
assert result is None
# ── MCPServer Properties ──
def test_has_token_exchange_config_true():
"""has_token_exchange_config is True for a fully configured OBO server."""
server = _obo_server()
assert server.has_token_exchange_config is True
def test_has_token_exchange_config_false_wrong_auth_type():
"""has_token_exchange_config is False when auth_type is not oauth2_token_exchange."""
server = _obo_server(auth_type=MCPAuth.oauth2)
assert server.has_token_exchange_config is False
def test_has_token_exchange_config_false_missing_creds():
"""has_token_exchange_config is False when client_id/client_secret missing."""
server = _obo_server(client_id=None)
assert server.has_token_exchange_config is False
def test_has_token_exchange_config_uses_token_url_fallback():
"""has_token_exchange_config is True when token_url is set instead of token_exchange_endpoint."""
server = _obo_server(
token_exchange_endpoint=None,
token_url="https://idp.example.com/token",
)
assert server.has_token_exchange_config is True
# ── Config Loading ──
@pytest.mark.asyncio
async def test_config_loading_token_exchange_fields():
"""load_servers_from_config correctly maps OBO config fields to MCPServer."""
manager = MCPServerManager()
config = {
"my_obo_server": {
"url": "https://mcp.example.com/mcp",
"transport": "http",
"auth_type": "oauth2_token_exchange",
"client_id": "my-client",
"client_secret": "my-secret",
"token_exchange_endpoint": "https://idp.example.com/oauth2/token",
"audience": "api://my-mcp",
"scopes": ["read", "write"],
"subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
}
}
await manager.load_servers_from_config(config)
servers = list(manager.config_mcp_servers.values())
assert len(servers) == 1
server = servers[0]
assert server.auth_type == MCPAuth.oauth2_token_exchange
assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token"
assert server.audience == "api://my-mcp"
assert server.subject_token_type == "urn:ietf:params:oauth:token-type:jwt"
assert server.client_id == "my-client"
assert server.client_secret == "my-secret"
assert server.scopes == ["read", "write"]
assert server.has_token_exchange_config is True
@pytest.mark.asyncio
async def test_config_loading_default_subject_token_type():
"""subject_token_type defaults to access_token when not specified in config."""
manager = MCPServerManager()
config = {
"obo_defaults": {
"url": "https://mcp.example.com/mcp",
"transport": "http",
"auth_type": "oauth2_token_exchange",
"client_id": "cid",
"client_secret": "csec",
"token_exchange_endpoint": "https://idp.example.com/token",
}
}
await manager.load_servers_from_config(config)
server = list(manager.config_mcp_servers.values())[0]
assert server.subject_token_type == "urn:ietf:params:oauth:token-type:access_token"
@pytest.mark.asyncio
async def test_database_loading_token_exchange_scopes_from_credentials():
"""DB-loaded OBO server credentials retain configured scopes."""
manager = MCPServerManager()
db_server = LiteLLM_MCPServerTable(
server_id="srv-obo-db",
server_name="obo_db_server",
url="https://mcp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
credentials={
"client_id": "db-client",
"client_secret": "db-secret",
"token_exchange_endpoint": "https://idp.example.com/oauth2/token",
"audience": "api://db-mcp",
"scopes": ["db.read", "db.write"],
},
)
server = await manager.build_mcp_server_from_table(
db_server,
credentials_are_encrypted=False,
)
assert server.auth_type == MCPAuth.oauth2_token_exchange
assert server.client_id == "db-client"
assert server.client_secret == "db-secret"
assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token"
assert server.audience == "api://db-mcp"
assert server.scopes == ["db.read", "db.write"]
@pytest.mark.asyncio
async def test_exchange_token_uses_client_secret_basic_when_configured():
"""LIT-4091: token exchange with token_endpoint_auth_method=client_secret_basic sends the
client credentials as HTTP Basic and omits client_secret from the body."""
import base64
handler = TokenExchangeHandler()
server = _obo_server(
server_id="srv-obo-basic", token_endpoint_auth_method="client_secret_basic"
)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("scoped-basic")
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
result = await handler.exchange_token("user-jwt-basic", server)
assert result == "scoped-basic"
_, kwargs = mock_client.post.call_args
expected = "Basic " + base64.b64encode(b"litellm-client-id:litellm-client-secret").decode()
assert kwargs["headers"]["Authorization"] == expected
assert "client_secret" not in kwargs["data"]
assert "client_id" not in kwargs["data"]
assert kwargs["data"]["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE

53
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -2333,6 +2333,59 @@ class TestMCPServerManager:
assert emitted.headers["Authorization"] == "Bearer upstream-token"
assert not kwargs["extra_headers"] or "authorization" not in {k.lower() for k in kwargs["extra_headers"]}
@pytest.mark.asyncio
async def test_create_mcp_client_token_exchange_never_falls_back_to_v1(self):
"""A configured OBO server is owned end to end by the v2 token_exchange arm, even when the
caller supplies an x-mcp-* override. This is what makes the v1 OBO handler unreachable, so if
it ever defers to v1 again the deleted handler is silently needed back."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
OAuthToken,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import (
UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
class _StubExchanger:
def __init__(self):
self.subject_tokens = []
async def exchange(self, subject_token, server, config, *, tenant_id=""):
self.subject_tokens.append(subject_token)
return Ok(OAuthToken(access_token="exchanged-token"))
async def invalidate(self, subject_token, server, config, *, tenant_id=""):
return None
exchanger = _StubExchanger()
manager = MCPServerManager()
server = MCPServer(
server_id="obo-egress",
name="obo",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
client_id="gateway-client",
client_secret="gateway-secret",
token_exchange_endpoint="https://idp.example.com/oauth2/token",
)
with (
patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.resolve_mcp_auth",
new_callable=AsyncMock,
) as mock_resolve,
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as mock_client_cls,
):
await manager._create_mcp_client(
server=server,
mcp_auth_header="Bearer caller-override",
subject_token="eyJ-subject-token",
cred_provider=UpstreamCredentialProvider(token_exchanger=exchanger),
)
mock_resolve.assert_not_awaited()
assert exchanger.subject_tokens == ["eyJ-subject-token"]
assert self._emitted_authorization(mock_client_cls) == "Bearer exchanged-token"
@staticmethod
def _emitted_authorization(mock_client_cls) -> str:
kwargs = mock_client_cls.call_args.kwargs

4
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
View file

@ -237,7 +237,7 @@ class TestListToolRestApiWithToolSearch:
return_value={},
),
patch(
"litellm.proxy._experimental.mcp_server.rest_endpoints._get_oauth2_server_ids",
"litellm.proxy._experimental.mcp_server.rest_endpoints._v1_resolved_oauth2_server_ids",
return_value=[],
),
patch(
@ -316,7 +316,7 @@ class TestListToolRestApiWithToolSearch:
return_value={},
),
patch(
"litellm.proxy._experimental.mcp_server.rest_endpoints._get_oauth2_server_ids",
"litellm.proxy._experimental.mcp_server.rest_endpoints._v1_resolved_oauth2_server_ids",
return_value=[],
),
patch(

65
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -2783,3 +2783,68 @@ class TestRestListToolsetFiltering:
)
assert [tool.name for tool in result] == ["lookup_status"]
class TestV1ResolvedOauth2Gate:
"""The REST surface must stop resolving per-user OAuth2 tokens for servers the v2 resolver owns.
``_resolve_v2_auth`` drops any Authorization built here for an ``authorization_code`` server and
injects the resolver's own token, so the v1 lookup was a DB round-trip whose result was discarded.
A server that still defers to v1 (upstream-delegated oauth2) must keep resolving, which is what
makes these assertions non-vacuous.
"""
@staticmethod
def _oauth2_server(*, delegate_auth_to_upstream: bool) -> Any:
from litellm.proxy._experimental.mcp_server.server import MCPServer
from litellm.types.mcp import MCPTransport
return MCPServer(
server_id="oauth2-srv",
name="oauth2-srv",
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=delegate_auth_to_upstream,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"delegate_auth_to_upstream, expected_headers, expected_lookups",
[
(False, None, 0),
(True, {"Authorization": "Bearer stored-token"}, 1),
],
)
async def test_user_oauth_headers_skip_v2_owned_servers(
self, delegate_auth_to_upstream, expected_headers, expected_lookups, monkeypatch
):
from litellm.proxy._experimental.mcp_server import db as mcp_db
server = self._oauth2_server(delegate_auth_to_upstream=delegate_auth_to_upstream)
resolve_token = AsyncMock(return_value={"access_token": "stored-token"})
monkeypatch.setattr(mcp_db, "resolve_valid_user_oauth_token", resolve_token)
headers = await rest_endpoints._get_user_oauth_extra_headers(
server,
UserAPIKeyAuth(user_id="alice", api_key="sk-1234"),
prefetched_creds={"oauth2-srv": {"access_token": "stored-token"}},
)
assert headers == expected_headers
assert resolve_token.await_count == expected_lookups
def test_prefetch_preflight_only_counts_v1_resolved_servers(self, monkeypatch):
v2_owned = self._oauth2_server(delegate_auth_to_upstream=False)
v1_resolved = self._oauth2_server(delegate_auth_to_upstream=True)
v1_resolved.server_id = "delegate-srv"
registry = {"oauth2-srv": v2_owned, "delegate-srv": v1_resolved}
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_mcp_server_by_id",
lambda server_id: registry.get(server_id),
)
assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv"]) == set()
assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv", "delegate-srv"]) == {"delegate-srv"}

283
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -179,6 +179,45 @@ async def test_budget_reservation_runs_when_not_disabled():
assert user_api_key_auth_obj.budget_reservation == reservation
@pytest.mark.asyncio
@pytest.mark.parametrize(
"general_settings,expected_flag",
[
({"fail_closed_budget_enforcement": True}, True),
({}, False),
],
)
async def test_fail_closed_budget_enforcement_reaches_reservation(
general_settings, expected_flag
):
"""#33923: the strict flag must be threaded into reserve_budget_for_request so a
failed reservation write can reject instead of failing open."""
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
with patch(
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
new=AsyncMock(return_value=None),
) as mock_reserve:
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request_data={"model": "gpt-4o"},
route="/v1/chat/completions",
llm_router=None,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
skip_budget_checks=False,
general_settings=general_settings,
)
assert (
mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"]
is expected_flag
)
@pytest.mark.asyncio
async def test_should_not_reuse_cached_key_object_for_request_state():
key_cache = DualCache()
@ -1290,6 +1329,250 @@ async def test_scim_deactivated_user_key_is_rejected():
setattr(_proxy_server_mod, attr, val)
@pytest.mark.asyncio
async def test_cached_proxy_admin_key_sets_via_virtual_key_marker():
"""Cached PROXY_ADMIN auth objects early-return before the marked DB and
master-key returns, and cache serialization drops the exclude=True marker;
the cache-hit boundary must restore it or cached admin traffic silently
bypasses overwrite_user_with_key_hash stamping."""
from fastapi import Request
from starlette.datastructures import URL
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
from litellm.proxy.proxy_server import hash_token
api_key = "sk-cached-admin-marker-test"
hashed_key = hash_token(api_key)
cached_token = UserAPIKeyAuth(
api_key=api_key,
token=hashed_key,
user_id="cached-admin-user",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
assert cached_token.via_virtual_key is False
mock_cache = AsyncMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.delete_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = (
AsyncMock()
)
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
import litellm.proxy.proxy_server as _proxy_server_mod
_attrs_to_set = {
"prisma_client": MagicMock(),
"user_api_key_cache": mock_cache,
"proxy_logging_obj": mock_proxy_logging_obj,
"master_key": "sk-master-key",
"general_settings": {},
"llm_model_list": [],
"llm_router": None,
"open_telemetry_logger": None,
"model_max_budget_limiter": MagicMock(),
"user_custom_auth": None,
"jwt_handler": None,
"litellm_proxy_admin_name": "admin",
}
_original_values = {
attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set
}
try:
for attr, val in _attrs_to_set.items():
setattr(_proxy_server_mod, attr, val)
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
with patch(
"litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
new_callable=AsyncMock,
return_value=cached_token,
):
result = await _user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
assert isinstance(result, UserAPIKeyAuth)
assert result.user_role == LitellmUserRoles.PROXY_ADMIN
assert result.via_virtual_key is True
assert result.api_key == hashed_key
finally:
for attr, val in _original_values.items():
setattr(_proxy_server_mod, attr, val)
@pytest.mark.asyncio
async def test_master_key_auth_sets_via_virtual_key_marker():
"""Master-key requests must also be stamped by overwrite_user_with_key_hash;
the auth path substitutes the stable alias for api_key and must mark the
result as proxy-validated."""
from fastapi import Request
from starlette.datastructures import URL
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
master_key = "sk-master-key"
mock_cache = AsyncMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.delete_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = (
AsyncMock()
)
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
import litellm.proxy.proxy_server as _proxy_server_mod
_attrs_to_set = {
"prisma_client": MagicMock(),
"user_api_key_cache": mock_cache,
"proxy_logging_obj": mock_proxy_logging_obj,
"master_key": master_key,
"general_settings": {},
"llm_model_list": [],
"llm_router": None,
"open_telemetry_logger": None,
"model_max_budget_limiter": MagicMock(),
"user_custom_auth": None,
"jwt_handler": None,
"litellm_proxy_admin_name": "admin",
}
_original_values = {
attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set
}
try:
for attr, val in _attrs_to_set.items():
setattr(_proxy_server_mod, attr, val)
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
result = await _user_api_key_auth_builder(
request=request,
api_key=f"Bearer {master_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
assert isinstance(result, UserAPIKeyAuth)
assert result.via_virtual_key is True
assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS
finally:
for attr, val in _original_values.items():
setattr(_proxy_server_mod, attr, val)
@pytest.mark.asyncio
async def test_db_virtual_key_auth_sets_via_virtual_key_marker():
"""via_virtual_key gates overwrite_user_with_key_hash stamping and is
forge-stripped from validated input, so the DB auth path setting it by
post-construction assignment is the only thing that turns stamping on."""
from fastapi import Request
from starlette.datastructures import URL
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
from litellm.proxy.proxy_server import hash_token
api_key = "sk-via-virtual-key-marker-test"
hashed_key = hash_token(api_key)
valid_token = UserAPIKeyAuth(
api_key=api_key,
token=hashed_key,
user_id="marker-test-user",
)
mock_cache = AsyncMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.delete_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = (
AsyncMock()
)
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
mock_prisma_client = MagicMock()
import litellm.proxy.proxy_server as _proxy_server_mod
_attrs_to_set = {
"prisma_client": mock_prisma_client,
"user_api_key_cache": mock_cache,
"proxy_logging_obj": mock_proxy_logging_obj,
"master_key": "sk-master-key",
"general_settings": {},
"llm_model_list": [],
"llm_router": None,
"open_telemetry_logger": None,
"model_max_budget_limiter": MagicMock(),
"user_custom_auth": None,
"jwt_handler": None,
"litellm_proxy_admin_name": "admin",
}
_original_values = {
attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set
}
try:
for attr, val in _attrs_to_set.items():
setattr(_proxy_server_mod, attr, val)
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
with (
patch(
"litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
new_callable=AsyncMock,
return_value=valid_token,
),
patch(
"litellm.proxy.auth.user_api_key_auth.get_user_object",
new_callable=AsyncMock,
return_value=None,
),
):
result = await _user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
assert isinstance(result, UserAPIKeyAuth)
assert result.via_virtual_key is True
assert result.api_key == hashed_key
finally:
for attr, val in _original_values.items():
setattr(_proxy_server_mod, attr, val)
@pytest.mark.asyncio
async def test_return_user_api_key_auth_obj_user_spend_and_budget():
"""

21
tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py
View file

@ -63,6 +63,27 @@ def test_sso_descriptor_mapping_is_single_sourced():
)
def test_sso_descriptor_mapping_covers_saml_fields():
# SAML config is stored and read through the same descriptor table as the
# OAuth providers; the login path reads these env vars, so the save path must
# map every SAML field to its uppercase env var.
assert SSO_FIELD_ENV_VARS["saml_idp_metadata_url"] == "SAML_IDP_METADATA_URL"
assert SSO_FIELD_ENV_VARS["saml_idp_metadata_xml"] == "SAML_IDP_METADATA_XML"
assert SSO_FIELD_ENV_VARS["saml_sp_entity_id"] == "SAML_SP_ENTITY_ID"
assert SSO_FIELD_ENV_VARS["saml_allow_unsolicited"] == "SAML_ALLOW_UNSOLICITED"
def test_resolve_sso_config_resolves_saml_fields():
resolved = resolve_sso_config(
{"saml_idp_metadata_url": "https://idp.example.com/metadata"},
{"SAML_ALLOW_UNSOLICITED": "true"},
)
assert resolved.config.saml_idp_metadata_url == "https://idp.example.com/metadata"
assert resolved.provenance["saml_idp_metadata_url"] == "db"
assert resolved.config.saml_allow_unsolicited == "true"
assert resolved.provenance["saml_allow_unsolicited"] == "env"
def test_resolve_sso_config_returns_unmasked_secret_and_provenance():
# The resolver hands back plaintext; masking is the endpoint's job. If the
# resolver masked, the login path would consume a masked secret and fail.

16
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

@ -3680,6 +3680,22 @@ async def test_pre_call_hook_skips_chat_traffic_when_configured_for_pre_mcp_call
mock_post.assert_not_called()
def test_process_response_with_none_metadata_does_not_crash():
guardrail = _make_guardrail()
response = {"id": "batch_123", "status": "validating"}
request_data = {"model": "gemini-2.5-flash", "metadata": None}
result = guardrail._process_response(
response=response,
request_data=request_data,
event_type=GuardrailEventHooks.post_call,
)
assert result is response
assert isinstance(request_data["metadata"], dict)
assert "standard_logging_guardrail_information" in request_data["metadata"]
@pytest.mark.asyncio
async def test_moderation_hook_scans_mcp_tool_call_when_configured_for_during_mcp_call():
"""A guardrail configured with mode `during_mcp_call` must scan MCP tool calls.

40
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -1,3 +1,5 @@
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.guardrails.guardrail_registry import (
get_guardrail_initializer_from_hooks,
@ -32,6 +34,44 @@ def test_noma_registry_resolution():
assert "noma_v2" in guardrail_initializer_registry
@pytest.mark.parametrize(
"configured, expected",
[(None, True), (False, False), (True, True)],
)
def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected):
"""
A guardrail whose constructor sets run_in_parallel=True must keep that default when
the config omits the key; only an explicit config value may override it. The
previous code wrote bool(None)==False on every instance, silently disabling the
opt-in for such guardrails.
"""
from litellm.proxy.guardrails import guardrail_registry as registry_module
def _initializer(litellm_params, guardrail):
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer
try:
params = {"guardrail": "parallel_default_test", "mode": "pre_call"}
if configured is not None:
params["run_in_parallel"] = configured
handler = InMemoryGuardrailHandler()
result = handler.initialize_guardrail(
guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params},
)
stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
assert stored.run_in_parallel is expected
finally:
registry_module.guardrail_initializer_registry.pop("parallel_default_test", None)
def test_update_in_memory_guardrail():
handler = InMemoryGuardrailHandler()
handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(

24
tests/test_litellm/proxy/guardrails/test_init_guardrails.py
View file

@ -62,3 +62,27 @@ def test_initialize_guardrail_preserves_guardrail_info():
assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"}
stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]]
assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"}
@pytest.mark.parametrize(
"config_value, expected",
[(True, True), (False, False), (None, False)],
)
def test_initialize_guardrail_sets_run_in_parallel(config_value, expected):
"""run_in_parallel from litellm_params must reach the built guardrail instance."""
litellm_params = {
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
"mode": "pre_call",
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
}
if config_value is not None:
litellm_params["run_in_parallel"] = config_value
guardrail_handler = InMemoryGuardrailHandler()
result = guardrail_handler.initialize_guardrail(
guardrail={"guardrail_name": "test_parallel_flag", "litellm_params": litellm_params},
)
custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
assert custom_guardrail.run_in_parallel is expected

644
tests/test_litellm/proxy/management_endpoints/test_saml_sso.py Normal file
View file

@ -0,0 +1,644 @@
"""
Regression tests for SAML 2.0 SSO (SP- and IdP-initiated) on the admin UI.
These exercise the real OneLogin python3-saml validation by generating signed
SAML responses with a freshly minted IdP keypair, so a mutation that weakens
signature, signing-requirement, expiry, replay or attribute-mapping handling
makes a test fail.
"""
import base64
import datetime
import time
import pytest
from fastapi import HTTPException, Request
pytest.importorskip(
"onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests"
)
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
from onelogin.saml2.utils import OneLogin_Saml2_Utils
from starlette.datastructures import URL
from typing import cast
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.management_endpoints.sso.saml_sso import (
_SAML_AUTHN_REQUEST_CACHE_PREFIX,
_SAML_AUTHN_STATE_COOKIE,
_SAML_MAX_POST_BYTES,
_SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS,
_SAML_REPLAY_GUARD_MAX_TTL_SECONDS,
SAMLAuthHandler,
)
def _shared_cache(store=None):
"""A DualCache whose replay guard is backed by a shared, atomic store.
An InMemoryCache instance stands in for Redis; passing the same instance to
two DualCaches simulates two workers sharing one atomic backend."""
return DualCache(redis_cache=cast(RedisCache, store or InMemoryCache()))
IDP_ENTITY = "https://idp.example.com/metadata"
SP_ENTITY = "https://proxy.example.com/sso/saml/metadata"
ACS = "https://proxy.example.com/sso/saml/callback"
SSO_URL = "https://idp.example.com/sso"
PROXY_BASE_URL = "https://proxy.example.com"
def _make_idp_keypair():
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "idp.example.com")])
cert = (
x509.CertificateBuilder()
.subject_name(name)
.issuer_name(name)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow() - datetime.timedelta(days=1))
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
.sign(key, hashes.SHA256())
)
key_pem = key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.TraditionalOpenSSL,
serialization.NoEncryption(),
).decode()
cert_pem = cert.public_bytes(serialization.Encoding.PEM).decode()
return key_pem, cert_pem
def _idp_metadata_xml(cert_pem):
cert_body = "".join(
line for line in cert_pem.splitlines() if "CERTIFICATE" not in line
)
return (
'<?xml version="1.0"?>'
f'<EntityDescriptor xmlns="urn:oasis:names:tc:SAML:2.0:metadata" entityID="{IDP_ENTITY}">'
'<IDPSSODescriptor protocolSupportEnumeration="urn:oasis:names:tc:SAML:2.0:protocol">'
'<KeyDescriptor use="signing"><KeyInfo xmlns="http://www.w3.org/2000/09/xmldsig#">'
f"<X509Data><X509Certificate>{cert_body}</X509Certificate></X509Data>"
"</KeyInfo></KeyDescriptor>"
'<SingleSignOnService Binding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" '
f'Location="{SSO_URL}"/>'
"</IDPSSODescriptor></EntityDescriptor>"
)
def _saml_time(delta_seconds):
t = datetime.datetime.utcnow() + datetime.timedelta(seconds=delta_seconds)
return t.strftime("%Y-%m-%dT%H:%M:%SZ")
def _build_signed_response(
key_pem,
cert_pem,
*,
in_response_to=None,
response_level_in_response_to=True,
email="alice@example.com",
attributes=None,
not_before_delta=-60,
not_on_or_after_delta=300,
sign=True,
):
if attributes is None:
attributes = {
"email": [email],
"givenName": ["Alice"],
"sn": ["Smith"],
"role": ["internal_user"],
}
assertion_id = "_assertion_" + OneLogin_Saml2_Utils.generate_unique_id()
response_id = "_response_" + OneLogin_Saml2_Utils.generate_unique_id()
not_before = _saml_time(not_before_delta)
not_on_or_after = _saml_time(not_on_or_after_delta)
issue_instant = _saml_time(-1)
irt = f'InResponseTo="{in_response_to}"' if in_response_to else ""
response_irt = irt if response_level_in_response_to else ""
attr_xml = "".join(
f'<saml:Attribute Name="{name}">'
+ "".join(f"<saml:AttributeValue>{v}</saml:AttributeValue>" for v in values)
+ "</saml:Attribute>"
for name, values in attributes.items()
)
assertion = (
'<saml:Assertion xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" '
f'ID="{assertion_id}" Version="2.0" IssueInstant="{issue_instant}">'
f"<saml:Issuer>{IDP_ENTITY}</saml:Issuer>"
"<saml:Subject>"
'<saml:NameID Format="urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress">'
f"{email}</saml:NameID>"
'<saml:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">'
f'<saml:SubjectConfirmationData {irt} NotOnOrAfter="{not_on_or_after}" Recipient="{ACS}"/>'
"</saml:SubjectConfirmation></saml:Subject>"
f'<saml:Conditions NotBefore="{not_before}" NotOnOrAfter="{not_on_or_after}">'
f"<saml:AudienceRestriction><saml:Audience>{SP_ENTITY}</saml:Audience>"
"</saml:AudienceRestriction></saml:Conditions>"
f'<saml:AuthnStatement AuthnInstant="{issue_instant}" SessionIndex="_session">'
"<saml:AuthnContext><saml:AuthnContextClassRef>"
"urn:oasis:names:tc:SAML:2.0:ac:classes:Password"
"</saml:AuthnContextClassRef></saml:AuthnContext></saml:AuthnStatement>"
f"<saml:AttributeStatement>{attr_xml}</saml:AttributeStatement>"
"</saml:Assertion>"
)
if sign:
signed = OneLogin_Saml2_Utils.add_sign(assertion, key_pem, cert_pem)
assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace(
'<?xml version="1.0"?>', ""
)
return (
'<?xml version="1.0"?>'
'<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" '
'xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" '
f'ID="{response_id}" Version="2.0" IssueInstant="{issue_instant}" '
f'Destination="{ACS}" {response_irt}>'
f"<saml:Issuer>{IDP_ENTITY}</saml:Issuer>"
'<samlp:Status><samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>'
"</samlp:Status>"
f"{assertion}</samlp:Response>"
)
def _b64(xml):
return base64.b64encode(xml.encode()).decode()
def _fake_request(cookies=None):
return type(
"Req",
(),
{
"base_url": URL(PROXY_BASE_URL + "/"),
"query_params": {},
"cookies": cookies or {},
},
)()
async def _acs(b64, cache, cookies=None):
return await SAMLAuthHandler.handle_acs(
_fake_request(cookies), cache, {"SAMLResponse": b64}
)
@pytest.fixture
def saml_env(monkeypatch):
key_pem, cert_pem = _make_idp_keypair()
monkeypatch.setenv("SAML_IDP_METADATA_XML", _idp_metadata_xml(cert_pem))
monkeypatch.setenv("SAML_SP_ENTITY_ID", SP_ENTITY)
monkeypatch.setenv("PROXY_BASE_URL", PROXY_BASE_URL)
for var in (
"SAML_IDP_METADATA_URL",
"SAML_ATTRIBUTE_EMAIL",
"SAML_ATTRIBUTE_TEAM_IDS",
"SAML_ALLOW_UNSOLICITED",
"ALLOWED_EMAIL_DOMAINS",
):
monkeypatch.delenv(var, raising=False)
return key_pem, cert_pem
@pytest.fixture
def saml_env_idp_initiated(saml_env, monkeypatch):
monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "true")
return saml_env
@pytest.mark.asyncio
async def test_valid_idp_initiated_login_maps_assertion_to_user(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(key_pem, cert_pem)
result = await _acs(_b64(resp), _shared_cache())
assert result.email == "alice@example.com"
assert result.id == "alice@example.com"
assert result.first_name == "Alice"
assert result.last_name == "Smith"
assert result.user_role == LitellmUserRoles.INTERNAL_USER
assert result.provider == "saml"
@pytest.mark.asyncio
async def test_tampered_assertion_is_rejected(saml_env):
key_pem, cert_pem = saml_env
resp = _build_signed_response(key_pem, cert_pem)
tampered = resp.replace("alice@example.com", "attacker@example.com")
with pytest.raises(HTTPException) as exc:
await _acs(_b64(tampered), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_unsigned_assertion_is_rejected(saml_env):
key_pem, cert_pem = saml_env
resp = _build_signed_response(key_pem, cert_pem, sign=False)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_signature_from_untrusted_key_is_rejected(saml_env):
_, cert_pem = saml_env
attacker_key, attacker_cert = _make_idp_keypair()
resp = _build_signed_response(attacker_key, attacker_cert)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_expired_assertion_is_rejected(saml_env):
key_pem, cert_pem = saml_env
resp = _build_signed_response(
key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600
)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_sp_initiated_unknown_in_response_to_is_rejected(saml_env):
key_pem, cert_pem = saml_env
resp = _build_signed_response(key_pem, cert_pem, in_response_to="_never_issued")
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_sp_initiated_known_request_succeeds_once_then_replay_rejected(saml_env):
key_pem, cert_pem = saml_env
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id)
cookies = {_SAML_AUTHN_STATE_COOKIE: request_id}
result = await _acs(_b64(resp), cache, cookies=cookies)
assert result.email == "alice@example.com"
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), cache, cookies=cookies)
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env):
key_pem, cert_pem = saml_env
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), cache)
assert exc.value.status_code == 401
with pytest.raises(HTTPException) as exc:
await _acs(
_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"}
)
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_subjectconfirmation_only_in_response_to_without_cookie_is_rejected(
saml_env_idp_initiated,
):
"""An IdP that stamps InResponseTo only on the SubjectConfirmationData (not the
Response element) is still solicited and must be browser-bound: with unsolicited
explicitly allowed, a missing cookie must still 401 rather than slip through."""
key_pem, cert_pem = saml_env_idp_initiated
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
resp = _build_signed_response(
key_pem,
cert_pem,
in_response_to=request_id,
response_level_in_response_to=False,
)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), cache)
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml_env):
key_pem, cert_pem = saml_env
cache = DualCache()
request_id = "_authn_req_known"
cache.set_cache(
key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600
)
resp = _build_signed_response(
key_pem,
cert_pem,
in_response_to=request_id,
response_level_in_response_to=False,
)
result = await _acs(
_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id}
)
assert result.email == "alice@example.com"
@pytest.mark.asyncio
async def test_unsolicited_response_rejected_by_default(saml_env):
key_pem, cert_pem = saml_env
resp = _build_signed_response(key_pem, cert_pem)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_idp_initiated_assertion_replay_is_rejected(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
cache = _shared_cache()
resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com")
first = await _acs(_b64(resp), cache)
assert first.email == "bob@example.com"
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), cache)
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_assertion_without_id_is_rejected(saml_env_idp_initiated):
"""An assertion with no ID attribute has no stable replay key. On the unsolicited
path there is no browser binding, so the consumed-assertion guard is the only replay
defense; a missing ID must be rejected rather than silently skipping the guard."""
class _AuthNoAssertionId:
def get_last_response_in_response_to(self):
return None
def get_last_response_xml(self):
return None
def get_last_assertion_id(self):
return None
with pytest.raises(HTTPException) as exc:
await SAMLAuthHandler._enforce_response_binding(
_AuthNoAssertionId(), _shared_cache(), None
)
assert exc.value.status_code == 401
assert "ID" in exc.value.detail
@pytest.mark.asyncio
async def test_unsolicited_response_rejected_when_disabled(saml_env, monkeypatch):
key_pem, cert_pem = saml_env
monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "false")
resp = _build_signed_response(key_pem, cert_pem)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
@pytest.mark.asyncio
async def test_invalid_email_in_assertion_is_rejected_cleanly(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(
key_pem,
cert_pem,
email="not-an-email",
attributes={"email": ["not-an-email"], "givenName": ["X"]},
)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), _shared_cache())
assert exc.value.status_code == 401
assert "invalid subject or email" in exc.value.detail
@pytest.mark.asyncio
async def test_email_less_assertion_rejected_when_domain_restriction_configured(
saml_env_idp_initiated, monkeypatch
):
key_pem, cert_pem = saml_env_idp_initiated
monkeypatch.setenv("ALLOWED_EMAIL_DOMAINS", "example.com")
resp = _build_signed_response(
key_pem,
cert_pem,
email="opaque-persistent-id-123",
attributes={"givenName": ["Alice"]},
)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), _shared_cache())
assert exc.value.status_code == 401
assert "ALLOWED_EMAIL_DOMAINS" in exc.value.detail
@pytest.mark.asyncio
async def test_email_less_assertion_allowed_without_domain_restriction(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(
key_pem,
cert_pem,
email="opaque-persistent-id-123",
attributes={"givenName": ["Alice"]},
)
result = await _acs(_b64(resp), _shared_cache())
assert result.email is None
assert result.id == "opaque-persistent-id-123"
@pytest.mark.asyncio
async def test_custom_email_attribute_override(saml_env_idp_initiated, monkeypatch):
key_pem, cert_pem = saml_env_idp_initiated
monkeypatch.setenv("SAML_ATTRIBUTE_EMAIL", "corpMail")
resp = _build_signed_response(
key_pem,
cert_pem,
email="ignored@example.com",
attributes={
"corpMail": ["real@corp.example.com"],
"givenName": ["Real"],
},
)
result = await _acs(_b64(resp), _shared_cache())
assert result.email == "real@corp.example.com"
@pytest.mark.asyncio
async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(
key_pem,
cert_pem,
attributes={
"email": ["carol@example.com"],
"groups": ["team-a", "team-b"],
},
)
result = await _acs(_b64(resp), _shared_cache())
assert result.team_ids == ["team-a", "team-b"]
@pytest.mark.asyncio
async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env):
cache = DualCache()
redirect = await SAMLAuthHandler.build_login_redirect(_fake_request(), cache)
location = redirect.headers["location"]
assert location.startswith(SSO_URL)
assert "SAMLRequest=" in location
cached = [
k
for k in cache.in_memory_cache.cache_dict
if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX)
]
assert len(cached) == 1
request_id = cached[0].split(":", 1)[1]
set_cookie = redirect.headers["set-cookie"]
assert f"{_SAML_AUTHN_STATE_COOKIE}={request_id}" in set_cookie
assert "httponly" in set_cookie.lower()
@pytest.mark.asyncio
async def test_sp_metadata_contains_acs_and_entity_id(saml_env):
metadata = await SAMLAuthHandler.build_sp_metadata(_fake_request(), DualCache())
assert ACS in metadata
assert SP_ENTITY in metadata
assert "AssertionConsumerService" in metadata
def test_replay_guard_ttl_tracks_assertion_validity():
class _Auth:
def __init__(self, not_on_or_after):
self._not_on_or_after = not_on_or_after
def get_last_assertion_not_on_or_after(self):
return self._not_on_or_after
now = int(time.time())
long_lived = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 7200))
assert long_lived >= 7200
short_lived = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 60))
assert short_lived == _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS
missing = SAMLAuthHandler._replay_guard_ttl(_Auth(None))
assert missing == _SAML_REPLAY_GUARD_DEFAULT_TTL_SECONDS
capped = SAMLAuthHandler._replay_guard_ttl(_Auth(now + 10 * 86400))
assert capped == _SAML_REPLAY_GUARD_MAX_TTL_SECONDS
def test_is_saml_configured_reflects_env(monkeypatch):
monkeypatch.delenv("SAML_IDP_METADATA_URL", raising=False)
monkeypatch.delenv("SAML_IDP_METADATA_XML", raising=False)
assert SAMLAuthHandler.is_saml_configured() is False
monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml")
assert SAMLAuthHandler.is_saml_configured() is True
@pytest.mark.asyncio
async def test_idp_initiated_rejected_without_shared_cache(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
resp = _build_signed_response(key_pem, cert_pem)
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), DualCache())
assert exc.value.status_code == 401
assert "shared Redis cache" in exc.value.detail
@pytest.mark.asyncio
async def test_idp_initiated_replay_rejected_across_workers(saml_env_idp_initiated):
key_pem, cert_pem = saml_env_idp_initiated
shared_store = InMemoryCache()
worker_one = _shared_cache(shared_store)
worker_two = _shared_cache(shared_store)
resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com")
first = await _acs(_b64(resp), worker_one)
assert first.email == "bob@example.com"
with pytest.raises(HTTPException) as exc:
await _acs(_b64(resp), worker_two)
assert exc.value.status_code == 401
class _FakeChunkedRequest:
def __init__(self, chunks, content_length=None):
self._chunks = chunks
self.headers = {} if content_length is None else {"content-length": content_length}
async def stream(self):
for chunk in self._chunks:
yield chunk
@pytest.mark.asyncio
async def test_read_acs_post_data_parses_form():
body = b"SAMLResponse=abc123&RelayState=%2Fui%2F"
request = _FakeChunkedRequest([body], content_length=str(len(body)))
post_data = await SAMLAuthHandler.read_acs_post_data(cast(Request, request))
assert post_data == {"SAMLResponse": "abc123", "RelayState": "/ui/"}
@pytest.mark.asyncio
async def test_read_acs_post_data_rejects_oversized_content_length():
request = _FakeChunkedRequest([b""], content_length=str(_SAML_MAX_POST_BYTES + 1))
with pytest.raises(HTTPException) as exc:
await SAMLAuthHandler.read_acs_post_data(cast(Request, request))
assert exc.value.status_code == 413
@pytest.mark.asyncio
async def test_read_acs_post_data_rejects_oversized_stream_without_content_length():
chunk = b"a" * (1024 * 1024)
chunk_count = _SAML_MAX_POST_BYTES // len(chunk) + 2
request = _FakeChunkedRequest([chunk] * chunk_count)
with pytest.raises(HTTPException) as exc:
await SAMLAuthHandler.read_acs_post_data(cast(Request, request))
assert exc.value.status_code == 413

125
tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py
View file

@ -13,12 +13,17 @@ from datetime import datetime, timezone
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.management_endpoints.tool_management_endpoints import router
from litellm.proxy.management_endpoints.tool_management_endpoints import (
_build_tool_spend_response,
_ToolSpendRow,
router,
)
from litellm.types.tool_management import LiteLLM_ToolTableRow
# --- helpers ---
@ -50,9 +55,9 @@ def _make_app() -> FastAPI:
# Stub the auth dependency so we don't need a real proxy running.
def _override_auth():
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
return UserAPIKeyAuth(api_key="sk-test", user_id="admin")
return UserAPIKeyAuth(api_key="sk-test", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
# A real (non-None) prisma stub for truthiness checks.
@ -147,3 +152,117 @@ class TestToolManagementEndpoints:
json={"tool_name": "my_tool", "input_policy": "invalid_value"},
)
assert resp.status_code == 422
def test_tool_spend_route_not_shadowed_by_get_tool(self):
prisma = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=[])
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
resp = self.client.get("/v1/tool/spend")
assert resp.status_code == 200
assert resp.json()["by_tool"] == []
def test_tool_spend_aggregates_and_sorts(self):
rows = [
{"date": "2026-07-01", "tool_name": "search", "call_count": 2, "spend": 1.0, "total_tokens": 100},
{"date": "2026-07-02", "tool_name": "search", "call_count": 1, "spend": 4.0, "total_tokens": 50},
{"date": "2026-07-01", "tool_name": "read_file", "call_count": 3, "spend": 2.0, "total_tokens": 300},
]
prisma = MagicMock()
prisma.db.query_raw = AsyncMock(side_effect=[rows, [{"total_spend": 5.5}]])
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02")
assert resp.status_code == 200
body = resp.json()
assert [t["tool_name"] for t in body["by_tool"]] == ["search", "read_file"]
search = body["by_tool"][0]
assert search["spend"] == 5.0
assert search["call_count"] == 3
assert search["total_tokens"] == 150
assert len(body["daily"]) == 3
assert body["start_date"] == "2026-07-01"
assert body["end_date"] == "2026-07-02"
assert body["total_spend"] == 5.5
@patch("litellm.proxy.proxy_server.prisma_client", None)
def test_tool_spend_no_db_returns_500(self):
resp = self.client.get("/v1/tool/spend")
assert resp.status_code == 500
def test_tool_spend_end_date_is_inclusive_via_exclusive_next_day_bound(self):
prisma = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=[])
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02")
assert resp.status_code == 200
expected_binds = (
datetime(2026, 7, 1, tzinfo=timezone.utc).isoformat(),
datetime(2026, 7, 3, tzinfo=timezone.utc).isoformat(),
)
assert prisma.db.query_raw.await_count == 2
for call in prisma.db.query_raw.await_args_list:
assert tuple(call.args[1:]) == expected_binds
assert resp.json()["end_date"] == "2026-07-02"
@pytest.mark.parametrize(
"query",
[
"start_date=not-a-date",
"start_date=2026-02-30",
"start_date=07/01/2026",
"end_date=2026-13-01",
"end_date=20260701",
],
)
def test_tool_spend_malformed_date_returns_400(self, query: str):
prisma = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=[])
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
resp = self.client.get(f"/v1/tool/spend?{query}")
assert resp.status_code == 400
assert "Invalid date format" in resp.json()["detail"]
prisma.db.query_raw.assert_not_awaited()
def test_tool_spend_non_admin_returns_403(self):
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
app = _make_app()
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="sk-user", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER
)
client = TestClient(app, raise_server_exceptions=True)
prisma = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=[])
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
resp = client.get("/v1/tool/spend")
assert resp.status_code == 403
prisma.db.query_raw.assert_not_awaited()
def _spend_row(date: str, tool_name: str, spend: float, call_count: int = 1, total_tokens: int = 10) -> _ToolSpendRow:
return _ToolSpendRow(date=date, tool_name=tool_name, call_count=call_count, spend=spend, total_tokens=total_tokens)
class TestBuildToolSpendResponse:
def test_multi_tool_attribution_double_counts_per_tool_but_not_total(self):
rows = [
_spend_row("2026-07-01", "a", spend=3.0),
_spend_row("2026-07-01", "b", spend=3.0),
]
resp = _build_tool_spend_response(rows, total_spend=3.0, start_date="2026-07-01", end_date="2026-07-01")
by_tool = {t.tool_name: t.spend for t in resp.by_tool}
assert by_tool == {"a": 3.0, "b": 3.0}
assert resp.total_spend == 3.0
def test_groups_across_days_and_sorts_by_spend(self):
rows = [
_spend_row("2026-07-01", "b", spend=1.0, call_count=2, total_tokens=100),
_spend_row("2026-07-02", "b", spend=4.0, call_count=1, total_tokens=50),
_spend_row("2026-07-01", "a", spend=2.0, call_count=3, total_tokens=300),
]
resp = _build_tool_spend_response(rows, total_spend=7.0, start_date="2026-07-01", end_date="2026-07-02")
assert [(t.tool_name, t.spend, t.call_count, t.total_tokens) for t in resp.by_tool] == [
("b", 5.0, 3, 150),
("a", 2.0, 3, 300),
]
assert len(resp.daily) == 3

65
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -7391,6 +7391,71 @@ async def test_legacy_login_page_hides_credentials_hint_via_general_settings():
assert "MASTER_KEY" not in body
@pytest.mark.asyncio
async def test_saml_callback_blocked_when_admin_ui_disabled():
"""An IdP-initiated assertion must not mint a UI session when the admin UI is
disabled; the ACS enforces DISABLE_ADMIN_UI like the SP-initiated login route."""
from litellm.proxy.management_endpoints.ui_sso import saml_callback
with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "true"}):
response = await saml_callback(SimpleNamespace(cookies={}))
assert response.status_code == 200
assert "Admin UI is Disabled" in response.body.decode()
@pytest.mark.asyncio
async def test_saml_callback_enforces_free_sso_user_limit_after_validation():
"""An IdP-initiated assertion must not bypass the >5 free-SSO-user Enterprise gate
that /sso/key/generate enforces; the ACS re-checks it after validating the assertion,
so the entitlement DB query never runs on unvalidated input."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.ui_sso import saml_callback
from litellm.proxy.management_endpoints.types import CustomOpenID
call_order: list[str] = []
async def _fake_handle_acs(**kwargs):
call_order.append("validate")
return CustomOpenID(
id="dana@litellm.ai",
email="dana@litellm.ai",
first_name=None,
last_name=None,
display_name="dana",
picture=None,
provider="saml",
team_ids=[],
user_role=None,
)
async def _fake_count_billable_users():
call_order.append("count")
return 6
async def _stream():
yield b"SAMLResponse=signed-response"
request_double = SimpleNamespace(cookies={}, headers={}, stream=_stream)
with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), patch(
"litellm.proxy.proxy_server.premium_user", False
), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch(
"litellm.proxy.proxy_server.master_key", "sk-1234"
), patch(
"litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs",
new=_fake_handle_acs,
), patch(
"litellm.repositories.user_repository.UserRepository.count_billable_users",
new=AsyncMock(side_effect=_fake_count_billable_users),
):
with pytest.raises(ProxyException) as exc:
await saml_callback(request_double)
assert str(exc.value.code) == "403"
assert call_order == ["validate", "count"]
@pytest.mark.asyncio
async def test_cli_poll_key_tolerates_missing_user_row():
"""The CLI poll must still mint the JWT when the user lookup raises,

13
tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
View file

@ -1076,6 +1076,13 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert():
would race a check-then-create into a duplicate-key failure. The upsert seeds
teams on create, and the filtered append is a no-op because the team is
already present on the freshly created row.
Calling upsert is not on its own enough to be atomic: Prisma only compiles it
down to a single INSERT ... ON CONFLICT when the update branch is non-empty,
and otherwise emits SELECT-then-INSERT, which loses the race. That is how
parallel /team/new calls naming the same new member started returning 500
"Unique constraint failed on the fields: (user_id)", so the shape of both
branches is pinned here.
"""
from litellm.proxy._types import LitellmUserRoles
@ -1114,5 +1121,7 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert():
# non-atomic standalone create that could race under concurrent provisioning
mock_prisma_client.db.litellm_usertable.upsert.assert_called_once()
mock_prisma_client.db.litellm_usertable.create.assert_not_called()
create_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["create"]
assert create_data["teams"] == ["team-1"]
upsert_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]
assert upsert_data["create"]["teams"] == ["team-1"]
assert upsert_data["update"], "empty update branch degrades the upsert to a racy SELECT-then-INSERT"
assert "teams" not in upsert_data["update"]

220
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -1628,6 +1628,226 @@ async def test_ui_view_spend_logs_date_range_filter(client, monkeypatch):
assert data["data"][0]["id"] == "log2"
@pytest.mark.asyncio
async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window(
client, monkeypatch
):
"""
LIT-3981: a request_id lookup on the UI route resolves across all time even
when the caller sends a date window that excludes the log (the dashboard
always sends a window). The window is dropped and request_id alone scopes
the query. Pre-fix the window was always applied, so an id from an older
page returned nothing.
"""
today = datetime.datetime.now(timezone.utc)
mock_spend_logs = [
{
"id": "log_old",
"request_id": "req-old",
"api_key": "sk-test-key",
"user": "test_user_1",
"team_id": "team1",
"spend": 0.05,
"startTime": (today - datetime.timedelta(days=90)).isoformat(),
"model": "gpt-4",
},
]
captured: dict = {}
def filter_fn(where):
captured["where"] = where
rows = _filter_logs_by_date_range(mock_spend_logs, where)
if where.get("request_id"):
rows = [r for r in rows if r["request_id"] == where["request_id"]]
return rows
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn),
)
# A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends.
start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S")
end_date = today.strftime("%Y-%m-%d %H:%M:%S")
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
response = client.get(
"/spend/logs/ui",
params={
"request_id": "req-old",
"start_date": start_date,
"end_date": end_date,
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 1
assert data["data"][0]["request_id"] == "req-old"
# Query dropped the time window and scoped solely by the primary key.
assert "startTime" not in captured["where"]
assert captured["where"]["request_id"] == "req-old"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_requires_dates_without_request_id(
client, monkeypatch
):
"""The date window stays mandatory on the UI route when no request_id is set."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma([], lambda where: []),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
response = client.get(
"/spend/logs/ui", headers={"Authorization": "Bearer sk-test"}
)
assert response.status_code == 400
assert "date" in response.text.lower()
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_spend_logs_v2_still_requires_dates_with_request_id(client, monkeypatch):
"""The public /spend/logs/v2 contract is unchanged: dates remain required even
when request_id is supplied. Only the internal UI route relaxes the window."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
make_ui_spend_logs_mock_prisma([], lambda where: []),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
response = client.get(
"/spend/logs/v2",
params={"request_id": "req-old"},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 400
assert "date" in response.text.lower()
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatch):
"""A non-admin looking up a request_id they do not own is rejected (403), so
the relaxed date window cannot read another tenant's log by id."""
class _ForeignRow:
user = "other_user"
team_id = None
class _SpendLogs:
async def find_unique(self, where, include=None):
return _ForeignRow()
class _DB:
def __init__(self):
self.litellm_spendlogs = _SpendLogs()
class _Prisma:
def __init__(self):
self.db = _DB()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _Prisma())
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1"
)
try:
response = client.get(
"/spend/logs/ui",
params={"request_id": "foreign-req"},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 403
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
client, monkeypatch
):
"""A non-admin owner looking up their own request_id resolves across all time.
The ownership check authorizes the single row, so the query drops both the date
window and the general user/team scoping and filters by the primary key alone;
without that skip an internal user would have a `user`/`OR` clause added."""
today = datetime.datetime.now(timezone.utc)
mock_spend_logs = [
{
"id": "log_old",
"request_id": "req-old",
"api_key": "sk-test-key",
"user": "user_1",
"team_id": "team1",
"spend": 0.05,
"startTime": (today - datetime.timedelta(days=90)).isoformat(),
"model": "gpt-4",
},
]
captured: dict = {}
def filter_fn(where):
captured["where"] = where
rows = _filter_logs_by_date_range(mock_spend_logs, where)
if where.get("request_id"):
rows = [r for r in rows if r["request_id"] == where["request_id"]]
return rows
mock_prisma = make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn)
class _OwnedRow:
user = "user_1"
team_id = "team1"
async def _find_unique(where, include=None):
return _OwnedRow()
mock_prisma.db.find_unique = _find_unique
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends.
start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S")
end_date = today.strftime("%Y-%m-%d %H:%M:%S")
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1"
)
try:
response = client.get(
"/spend/logs/ui",
params={
"request_id": "req-old",
"start_date": start_date,
"end_date": end_date,
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 1
assert data["data"][0]["request_id"] == "req-old"
assert "startTime" not in captured["where"]
assert captured["where"]["request_id"] == "req-old"
assert "user" not in captured["where"]
assert "OR" not in captured["where"]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_unauthorized(client):
# Test without authorization header

137
tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -109,6 +109,143 @@ def test_get_logging_payload_does_not_map_missing_or_zero_cached_tokens(prompt_t
assert "cache_read_input_tokens" not in additional_usage_values
def test_get_logging_payload_maps_openai_cache_write_tokens_to_cache_creation_input_tokens():
additional_usage_values = _get_additional_usage_values_for_usage(
litellm.Usage(
prompt_tokens=1000,
completion_tokens=2,
total_tokens=1002,
prompt_tokens_details={"cached_tokens": 0, "cache_write_tokens": 800},
)
)
assert additional_usage_values["cache_creation_input_tokens"] == 800
assert additional_usage_values["prompt_tokens_details"]["cache_write_tokens"] == 800
def test_get_logging_payload_preserves_anthropic_cache_creation_input_tokens():
additional_usage_values = _get_additional_usage_values_for_usage(
litellm.Usage(
prompt_tokens=1000,
completion_tokens=2,
total_tokens=1002,
cache_creation_input_tokens=300,
)
)
assert additional_usage_values["cache_creation_input_tokens"] == 300
@pytest.mark.parametrize(
"prompt_tokens_details",
[None, {"cached_tokens": 100}, {"cached_tokens": 100, "cache_write_tokens": 0}],
)
def test_get_logging_payload_does_not_map_missing_or_zero_cache_write_tokens(prompt_tokens_details):
additional_usage_values = _get_additional_usage_values_for_usage(
litellm.Usage(
prompt_tokens=10,
completion_tokens=2,
total_tokens=12,
prompt_tokens_details=prompt_tokens_details,
)
)
assert "cache_creation_input_tokens" not in additional_usage_values
def _make_standard_logging_payload_with_usage_object(usage_object: dict) -> StandardLoggingPayload:
return StandardLoggingPayload(
id="test-id-responses",
call_type="responses",
stream=False,
response_cost=0.02,
status="success",
total_tokens=1010,
prompt_tokens=1000,
completion_tokens=10,
startTime=1234567890.0,
endTime=1234567891.0,
completionStartTime=None,
model_map_information=StandardLoggingModelInformation(model_map_key="gpt-5.6", model_map_value=None),
model="gpt-5.6",
model_id="model-123",
model_group="openai",
custom_llm_provider="openai",
api_base="https://api.openai.com",
metadata=StandardLoggingMetadata(
user_api_key_hash="test_hash",
user_api_key_alias=None,
user_api_key_team_id=None,
user_api_key_org_id=None,
user_api_key_user_id=None,
user_api_key_team_alias=None,
spend_logs_metadata=None,
requester_ip_address=None,
requester_metadata=None,
user_api_key_end_user_id=None,
usage_object=usage_object,
),
cache_hit=False,
cache_key=None,
saved_cache_cost=0.0,
request_tags=[],
end_user=None,
requester_ip_address=None,
messages=[],
response={},
error_str=None,
model_parameters={},
hidden_params=StandardLoggingHiddenParams(
model_id="model-123",
cache_key=None,
api_base="https://api.openai.com",
response_cost="0.02",
litellm_overhead_time_ms=None,
additional_headers=None,
batch_models=None,
litellm_model_name=None,
usage_object=None,
),
)
def test_get_logging_payload_maps_responses_api_cache_write_tokens_from_usage_object():
"""Responses API (/v1/responses) usage is not chat-Usage-shaped, so
additional_usage_values can't derive cache tokens from response_obj.usage.
The Admin UI Logs "Cache Creation Tokens" row reads
additional_usage_values.cache_creation_input_tokens, so it must be filled
from the normalized standard_logging usage_object (LIT-4633)."""
standard_logging_payload = _make_standard_logging_payload_with_usage_object(
usage_object={
"prompt_tokens": 1000,
"completion_tokens": 10,
"total_tokens": 1010,
"prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 800, "cache_creation_tokens": 800},
}
)
payload = get_logging_payload(
kwargs={
"model": "gpt-5.6",
"call_type": "responses",
"litellm_params": {"metadata": {"user_api_key": "test-key"}},
"standard_logging_object": standard_logging_payload,
},
response_obj={
"id": "resp-test",
"usage": {
"input_tokens": 1000,
"output_tokens": 10,
"total_tokens": 1010,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 800},
},
},
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
additional_usage_values = json.loads(payload["metadata"])["additional_usage_values"]
assert additional_usage_values["cache_creation_input_tokens"] == 800
def test_sanitize_request_body_for_spend_logs_payload_basic():
request_body = {
"messages": [{"role": "user", "content": "Hello, how are you?"}],

95
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
import litellm
from litellm.caching.dual_cache import DualCache
@ -1562,6 +1563,100 @@ async def test_should_skip_reservation_when_counter_increment_fails(
)
@pytest.mark.asyncio
async def test_should_raise_503_when_counter_increment_fails_and_fail_closed(
spend_counter_state,
monkeypatch,
):
"""#33923: with fail_closed_budget_enforcement on, a failed reservation write
must reject instead of silently degrading to read-time-only enforcement."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-budget-reserve-fail-closed",
spend=0.0,
max_budget=1.0,
)
async def fail_increment_cache(*args, **kwargs):
raise RuntimeError("counter unavailable")
monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache)
with patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=0.5,
):
with pytest.raises(HTTPException) as exc_info:
await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=valid_token,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
fail_closed_budget_enforcement=True,
)
assert exc_info.value.status_code == 503
assert (
counter_cache.in_memory_cache.get_cache(
key="spend:key:key-budget-reserve-fail-closed"
)
is None
)
@pytest.mark.asyncio
async def test_fail_closed_releases_earlier_counters_before_503(
spend_counter_state,
):
"""#33923: when a later counter's reservation write fails in strict mode, the
counters that already reserved must be released before the 503 propagates."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-budget-fail-closed-release",
spend=0.0,
max_budget=1.0,
budget_limits=[
{
"budget_duration": "1h",
"max_budget": 1.0,
}
],
)
with patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=0.5,
):
with pytest.raises(HTTPException) as exc_info:
await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=valid_token,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
fail_closed_budget_enforcement=True,
)
assert exc_info.value.status_code == 503
assert (
counter_cache.in_memory_cache.get_cache(
key="spend:key:key-budget-fail-closed-release"
)
== 0.0
)
@pytest.mark.asyncio
async def test_should_skip_reservation_when_counter_initialization_fails(
spend_counter_state,

195
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -5226,3 +5226,198 @@ async def test_add_litellm_data_to_request_unions_metadata_tags_with_header_tags
tags = updated["litellm_metadata"]["tags"]
assert "header-tag" in tags
assert "body-tag" in tags
def _make_chat_request_mock() -> MagicMock:
return _make_request_mock("/v1/chat/completions", {"Content-Type": "application/json"})
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_clobbers_caller_supplied_user(monkeypatch):
"""The flag exists so providers can ban by a tamper-proof id; a caller-chosen
`user` must never survive, and the raw sk- key must never be forwarded."""
from litellm.proxy._types import hash_token
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
raw_key = "sk-overwrite-user-test-1234"
user_api_key_dict = UserAPIKeyAuth(api_key=raw_key)
user_api_key_dict.via_virtual_key = True
data = {"model": "gpt-4o", "user": "attacker-chosen-id"}
updated_data = await add_litellm_data_to_request(
data=data,
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == hash_token(raw_key)
assert updated_data["user"] != "attacker-chosen-id"
assert raw_key not in updated_data["user"]
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_sets_user_when_absent(monkeypatch):
from litellm.proxy._types import hash_token
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
raw_key = "sk-overwrite-user-test-5678"
user_api_key_dict = UserAPIKeyAuth(api_key=raw_key)
user_api_key_dict.via_virtual_key = True
data = {"model": "gpt-4o"}
updated_data = await add_litellm_data_to_request(
data=data,
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == hash_token(raw_key)
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_disabled_preserves_caller_user():
assert litellm.overwrite_user_with_key_hash is False
user_api_key_dict = UserAPIKeyAuth(api_key="sk-overwrite-user-test-9999")
user_api_key_dict.via_virtual_key = True
data = {"model": "gpt-4o", "user": "caller-chosen-id"}
updated_data = await add_litellm_data_to_request(
data=data,
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == "caller-chosen-id"
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_skips_custom_auth_credential(monkeypatch):
"""Custom-auth credentials are not sk-prefixed or JWTs, so UserAPIKeyAuth stores
them raw; the stamp must skip them entirely so auth material never leaks."""
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
raw_credential = "my-custom-auth-credential-abc123"
user_api_key_dict = UserAPIKeyAuth(api_key=raw_credential)
assert user_api_key_dict.api_key == raw_credential
updated_data = await add_litellm_data_to_request(
data={"model": "gpt-4o", "user": "caller-chosen-id"},
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == "caller-chosen-id"
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_skips_jwt_auth(monkeypatch):
"""A hashed JWT rotates on every token re-issue, so it is useless as a stable
ban id; JWT-authenticated requests are not stamped."""
from litellm.proxy._types import hash_token
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
hashed_jwt = f"hashed-jwt-{hash_token('some-jwt-token')}"
user_api_key_dict = UserAPIKeyAuth(api_key=hashed_jwt)
updated_data = await add_litellm_data_to_request(
data={"model": "gpt-4o", "user": "caller-chosen-id"},
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == "caller-chosen-id"
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_skips_hex_shaped_custom_credential(monkeypatch):
"""A custom-auth credential that happens to be 64 hex chars is indistinguishable
from a key hash by shape alone; only the server-set via_virtual_key marker may
authorize stamping, so this raw credential must never be forwarded."""
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
hex_shaped_credential = "a" * 64
user_api_key_dict = UserAPIKeyAuth(api_key=hex_shaped_credential)
assert user_api_key_dict.api_key == hex_shaped_credential
assert user_api_key_dict.via_virtual_key is False
updated_data = await add_litellm_data_to_request(
data={"model": "gpt-4o", "user": "caller-chosen-id"},
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == "caller-chosen-id"
def test_via_virtual_key_cannot_be_forged_from_validated_input():
from_kwargs = UserAPIKeyAuth(api_key="b" * 64, via_virtual_key=True)
assert from_kwargs.via_virtual_key is False
from_dict = UserAPIKeyAuth.model_validate({"api_key": "b" * 64, "via_virtual_key": True})
assert from_dict.via_virtual_key is False
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_stamps_master_key_alias(monkeypatch):
"""Master-key requests carry the stable alias instead of a hash (so the master
key never propagates anywhere); the alias is the stampable id for them."""
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
user_api_key_dict = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS)
user_api_key_dict.via_virtual_key = True
updated_data = await add_litellm_data_to_request(
data={"model": "gpt-4o", "user": "attacker-chosen-id"},
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == LITELLM_PROXY_MASTER_KEY_ALIAS
@pytest.mark.asyncio
async def test_overwrite_user_with_key_hash_rejects_alias_without_marker(monkeypatch):
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
monkeypatch.setattr(litellm, "overwrite_user_with_key_hash", True)
user_api_key_dict = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS)
assert user_api_key_dict.via_virtual_key is False
updated_data = await add_litellm_data_to_request(
data={"model": "gpt-4o", "user": "caller-chosen-id"},
request=_make_chat_request_mock(),
user_api_key_dict=user_api_key_dict,
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated_data["user"] == "caller-chosen-id"

69
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
View file

@ -617,6 +617,75 @@ class TestProxySettingEndpoints:
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] == "new_google_client_id"
def test_update_sso_settings_maps_saml_fields_to_env_vars(
self, mock_proxy_config, mock_auth, monkeypatch
):
"""SAML settings entered in the admin UI must be applied as the SAML_* env
vars the SAML handler reads, and the allow-unsolicited toggle must map to
the 'true'/'false' string the handler expects."""
import json
import os
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
mock_prisma.db.litellm_config = MagicMock()
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_config.update = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(
proxy_config,
"_encrypt_env_variables",
lambda environment_variables: environment_variables,
)
for var in (
"SAML_IDP_METADATA_URL",
"SAML_IDP_METADATA_XML",
"SAML_SP_ENTITY_ID",
"SAML_ALLOW_UNSOLICITED",
):
monkeypatch.delenv(var, raising=False)
new_sso_settings = {
"saml_idp_metadata_url": "https://idp.example.com/metadata",
"saml_sp_entity_id": "https://proxy.example.com/sso/saml/metadata",
"saml_allow_unsolicited": "true",
"proxy_base_url": "https://proxy.example.com",
"user_email": "admin@example.com",
}
try:
response = client.patch("/update/sso_settings", json=new_sso_settings)
assert response.status_code == 200
assert os.environ.get("SAML_IDP_METADATA_URL") == "https://idp.example.com/metadata"
assert os.environ.get("SAML_SP_ENTITY_ID") == "https://proxy.example.com/sso/saml/metadata"
assert os.environ.get("SAML_ALLOW_UNSOLICITED") == "true"
assert "SAML_IDP_METADATA_XML" not in os.environ
stored = json.loads(
mock_prisma.db.litellm_ssoconfig.upsert.call_args.kwargs["data"]["create"]["sso_settings"]
)
assert stored["saml_idp_metadata_url"] == "https://idp.example.com/metadata"
assert stored["saml_allow_unsolicited"] == "true"
finally:
for var in (
"SAML_IDP_METADATA_URL",
"SAML_IDP_METADATA_XML",
"SAML_SP_ENTITY_ID",
"SAML_ALLOW_UNSOLICITED",
):
os.environ.pop(var, None)
def test_update_sso_settings_audits_when_env_cleanup_fails(
self, mock_proxy_config, mock_auth, monkeypatch
):

1
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -626,6 +626,7 @@ def _moderation_guardrail() -> MagicMock:
cb.should_run_guardrail = MagicMock(return_value=True)
cb.async_moderation_hook = AsyncMock(return_value=None)
cb.async_post_call_success_hook = AsyncMock(return_value=None)
cb.run_in_parallel = False
return cb

1
tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py
View file

@ -28,6 +28,7 @@ def _make_guardrail(name="g", should_run=True, override=None):
cb.event_hook = GuardrailEventHooks.post_call
cb.should_run_guardrail = MagicMock(return_value=should_run)
cb.async_post_call_success_hook = AsyncMock(return_value=override)
cb.run_in_parallel = False
return cb

26
tests/test_litellm/responses/test_responses_utils.py
View file

@ -369,6 +369,32 @@ class TestResponseAPILoggingUtils:
assert result.completion_tokens_details.image_tokens == 272
assert result.completion_tokens_details.text_tokens == 100
def test_transform_response_api_usage_maps_cache_write_tokens(self):
"""Responses API (/v1/responses) cache-write tokens must survive the usage transform.
gpt-5.6 returns usage.input_tokens_details.cache_write_tokens (an extra field
not typed on InputTokensDetails). Before the fix the transform rebuilt the token
details and dropped it, leaving the cache-creation metric empty (LIT-4633).
"""
usage = {
"input_tokens": 10062,
"output_tokens": 16,
"total_tokens": 10078,
"input_tokens_details": {
"cached_tokens": 0,
"cache_write_tokens": 10059,
},
}
result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
usage
)
assert result.prompt_tokens_details is not None
assert result.prompt_tokens_details.cache_write_tokens == 10059
assert result.prompt_tokens_details.cache_creation_tokens == 10059
assert result.prompt_tokens_details.cached_tokens == 0
def test_transform_response_api_usage_mixed_details(self):
"""Test transformation handles mixed token details (cached + image + audio)."""
# Setup - hypothetical usage with mixed token types

277
tests/test_litellm/test_claude_opus_5_config.py Normal file
View file

@ -0,0 +1,277 @@
"""
Validate Claude Opus 5 model configuration entries.
Opus 5 carries Opus 4.8's pricing ($5 / $25 per MTok) and the gen-5 adaptive
thinking profile, but differs from 4.8 in two ways that are behavior-bearing in
LiteLLM: the cacheable-prefix minimum drops to 512 tokens, and Bedrock's Opus 5
validator accepts the full effort ladder, so the entries must not carry the
``bedrock_output_config_effort_ceiling`` that silently clamps ``max`` to
``xhigh`` on 4.8. The cost-map entries are also what populate
``litellm.anthropic_models`` at import, which is what lets a bare
``claude-opus-5`` name resolve to the ``anthropic`` provider (and match an
``anthropic/*`` wildcard deployment).
"""
import json
import os
import pytest
import litellm
from litellm.constants import BEDROCK_CONVERSE_MODELS
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..")
ALL_OPUS_5_VARIANTS = (
"claude-opus-5",
"anthropic.claude-opus-5",
"global.anthropic.claude-opus-5",
"us.anthropic.claude-opus-5",
"eu.anthropic.claude-opus-5",
"au.anthropic.claude-opus-5",
"jp.anthropic.claude-opus-5",
"vertex_ai/claude-opus-5",
"vertex_ai/claude-opus-5@default",
"azure_ai/claude-opus-5",
)
BEDROCK_OPUS_5_VARIANTS = (
"anthropic.claude-opus-5",
"global.anthropic.claude-opus-5",
"us.anthropic.claude-opus-5",
"eu.anthropic.claude-opus-5",
"au.anthropic.claude-opus-5",
"jp.anthropic.claude-opus-5",
)
def _load_root_cost_map() -> dict:
json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json")
with open(json_path) as f:
return json.load(f)
@pytest.fixture
def local_model_cost_map(monkeypatch):
"""Force the bundled backup cost map so assertions don't depend on the
network-fetched ``main`` copy (which lags this branch until merge)."""
original_model_cost = litellm.model_cost
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.get_model_info.cache_clear()
try:
yield
finally:
litellm.model_cost = original_model_cost
litellm.get_model_info.cache_clear()
def test_opus_5_pricing_and_capabilities():
model_data = _load_root_cost_map()
expected_providers = {
"claude-opus-5": "anthropic",
"anthropic.claude-opus-5": "bedrock_converse",
"vertex_ai/claude-opus-5": "vertex_ai-anthropic_models",
"azure_ai/claude-opus-5": "azure_ai",
}
for model_name, provider in expected_providers.items():
assert model_name in model_data, f"Missing model entry: {model_name}"
info = model_data[model_name]
assert info["litellm_provider"] == provider
assert info["mode"] == "chat"
assert info["max_input_tokens"] == 1000000
assert info["max_output_tokens"] == 128000
assert info["max_tokens"] == 128000
# Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard
# 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers.
assert info["input_cost_per_token"] == 5e-06
assert info["output_cost_per_token"] == 2.5e-05
assert info["cache_creation_input_token_cost"] == 6.25e-06
assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05
assert info["cache_read_input_token_cost"] == 5e-07
# Flat rate across the full 1M window, no long-context premium.
assert "input_cost_per_token_above_200k_tokens" not in info
assert "output_cost_per_token_above_200k_tokens" not in info
# gen-5 adaptive-thinking profile: effort-driven, no sampling params, no
# assistant prefill.
assert info["supports_adaptive_thinking"] is True
assert info["supports_reasoning"] is True
assert info["supports_sampling_params"] is False
assert info["supports_assistant_prefill"] is False
assert info["supports_xhigh_reasoning_effort"] is True
assert info["supports_max_reasoning_effort"] is True
assert info["supports_function_calling"] is True
assert info["supports_prompt_caching"] is True
assert info["supports_tool_choice"] is True
assert info["supports_vision"] is True
def test_opus_5_bedrock_regional_pricing():
"""Global/base endpoints use base pricing; the us./eu./au./jp. regional
cross-region inference profiles carry a 10% premium."""
model_data = _load_root_cost_map()
base_pricing = {
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2.5e-05,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
}
regional_pricing = {
"input_cost_per_token": 5.5e-06,
"output_cost_per_token": 2.75e-05,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
}
expected = {
"anthropic.claude-opus-5": base_pricing,
"global.anthropic.claude-opus-5": base_pricing,
"us.anthropic.claude-opus-5": regional_pricing,
"eu.anthropic.claude-opus-5": regional_pricing,
"au.anthropic.claude-opus-5": regional_pricing,
"jp.anthropic.claude-opus-5": regional_pricing,
}
for model_name, pricing in expected.items():
assert model_name in model_data, f"Missing model entry: {model_name}"
info = model_data[model_name]
assert info["litellm_provider"] == "bedrock_converse"
for key, value in pricing.items():
assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}"
@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS)
def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name):
"""Bedrock accepts every effort level for Opus 5, so no clamp belongs here.
Opus 4.7/4.8 carry ``bedrock_output_config_effort_ceiling: "xhigh"``, which
is what ``normalize_bedrock_opus_output_config_effort`` reads to rewrite a
caller's effort down. Verified against Bedrock on 2026-07-24 that
``output_config.effort="max"`` returns 200 for the Opus 5 profiles, so the
ceiling is deliberately absent; adding one back would silently downgrade
requests.
This asserts the cost-map entry rather than calling the normalizer because
``_BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER`` currently ranks ``max`` (3) below
``xhigh`` (4), so an ``xhigh`` ceiling never clamps ``max`` and a behavioral
assertion would pass either way. Keeping the entry clean means Opus 5 stays
correct once that ordering is fixed."""
info = _load_root_cost_map()[model_name]
assert "bedrock_output_config_effort_ceiling" not in info
@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS)
def test_opus_5_bedrock_rejects_strict_tools(model_name, local_model_cost_map):
"""Bedrock Converse routes Opus through a validator that rejects
``toolSpec.strict`` (``tools.0.custom.strict: Extra inputs are not
permitted``), same as Opus 4.7/4.8; verified against Bedrock on 2026-07-24.
Without the flag LiteLLM forwards ``strict`` and every tool call 400s."""
from litellm.llms.bedrock.common_utils import bedrock_converse_supports_strict_tools
assert bedrock_converse_supports_strict_tools(model_name) is False
def test_opus_5_prompt_cache_minimum_is_512(local_model_cost_map):
"""Opus 5 halves the cacheable-prefix minimum (Opus 4.8 is 1024).
The router's prompt-caching deployment check reads this value, so a stale
1024 would route prompts of 512-1023 tokens away from a warm Opus 5
deployment even though they cache fine."""
from litellm.utils import get_prompt_cache_min_tokens
assert get_prompt_cache_min_tokens(model="claude-opus-5") == 512
assert get_prompt_cache_min_tokens(model="us.anthropic.claude-opus-5") == 512
def test_opus_5_supports_fast_mode(local_model_cost_map):
"""Fast mode is Opus 5 on the first-party API at $10 / $50 per MTok, i.e. 2x
base. ``supports_speed`` gates whether ``speed="fast"`` is forwarded at all,
and ``provider_specific_entry.fast`` is what prices the response."""
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.cost_calculation import (
cost_per_token as anthropic_cost_per_token,
)
from litellm.types.utils import Usage
assert (
AnthropicConfig._model_supports_speed_param("claude-opus-5", "anthropic") is True
)
usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
usage.speed = "fast"
prompt_cost, completion_cost = anthropic_cost_per_token(
model="claude-opus-5", usage=usage
)
assert prompt_cost == pytest.approx(1000 * 5e-06 * 2.0)
assert completion_cost == pytest.approx(500 * 2.5e-05 * 2.0)
def test_opus_5_present_in_bundled_backup():
"""The bundled backup is the runtime fallback (and what tests load with
``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the
root cost map, otherwise the model resolves on one path but not the other."""
backup = GetModelCostMap.load_local_model_cost_map()
for model_name in ALL_OPUS_5_VARIANTS:
assert model_name in backup, f"Missing from backup cost map: {model_name}"
def test_opus_5_registered_for_bedrock_converse():
assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS
def test_opus_5_provider_resolves_via_model_info(local_model_cost_map):
"""Regression: ``claude-opus-5`` must resolve to provider ``anthropic``.
Without the cost-map entry the model is unknown to LiteLLM, so it cannot be
tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment
would not match it."""
info = litellm.get_model_info(model="claude-opus-5")
assert info["litellm_provider"] == "anthropic"
assert info["max_input_tokens"] == 1000000
assert info["max_output_tokens"] == 128000
@pytest.mark.parametrize(
"cost_map",
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
ids=["root", "bundled_backup"],
)
def test_opus_5_all_variants_carry_adaptive_thinking_flag(cost_map):
"""Every Opus 5 entry must advertise ``supports_adaptive_thinking``.
Adaptive-thinking detection is cost-map driven, so a single variant missing
the flag silently sends the legacy ``thinking.type='enabled'`` shape, which
Opus 5 rejects with a 400."""
variants = [k for k in cost_map if "claude-opus-5" in k]
assert variants, "no claude-opus-5 entries found in cost map"
missing = [
k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True
]
assert not missing, f"missing supports_adaptive_thinking: {missing}"
@pytest.mark.parametrize(
"cost_map",
[_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()],
ids=["root", "bundled_backup"],
)
def test_opus_5_all_variants_carry_512_token_cache_minimum(cost_map):
variants = [k for k in cost_map if "claude-opus-5" in k]
assert variants, "no claude-opus-5 entries found in cost map"
wrong = {
k: cost_map[k].get("prompt_cache_min_tokens")
for k in variants
if cost_map[k].get("prompt_cache_min_tokens") != 512
}
assert not wrong, f"prompt_cache_min_tokens must be 512: {wrong}"

30
tests/test_litellm/test_cost_calculator.py
View file

@ -12,6 +12,7 @@ from pydantic import BaseModel
import litellm
from litellm.cost_calculator import (
BaseTokenUsageProcessor,
RealtimeAPITokenUsageProcessor,
completion_cost,
cost_per_token,
@ -3479,3 +3480,32 @@ def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate():
)
assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2)
def test_combine_usage_objects_sums_mirrored_cache_write_fields_once():
"""
cache_write_tokens and cache_creation_tokens mirror each other on
PromptTokensDetailsWrapper, so field-iterating aggregation must sum the pair
once: a single 50-token usage stays 50 and two combine to 100, not double.
"""
single = Usage(
prompt_tokens=100,
completion_tokens=10,
total_tokens=110,
prompt_tokens_details=PromptTokensDetailsWrapper(cache_write_tokens=50),
)
combined = BaseTokenUsageProcessor.combine_usage_objects([single])
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.cache_write_tokens == 50
assert combined.prompt_tokens_details.cache_creation_tokens == 50
anthropic_style = Usage(
prompt_tokens=100,
completion_tokens=10,
total_tokens=110,
cache_creation_input_tokens=50,
)
combined_pair = BaseTokenUsageProcessor.combine_usage_objects([anthropic_style, anthropic_style])
assert combined_pair.prompt_tokens_details is not None
assert combined_pair.prompt_tokens_details.cache_write_tokens == 100
assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100

Some files were not shown because too many files have changed in this diff Show more