scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge pull request #40942 from BerriAI/litellm_internal_staging

Browse source 
chore(ci): promote internal staging to main
This commit is contained in:
yuneng-jiang 2026-09-12 21:11:12 -07:00 • committed by GitHub
commit 15789ae39e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
28 changed files with 1174 additions and 790 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
.github/workflows/guard-main-branch.yml vendored
View file

@ -1,42 +0,0 @@
name: Guard main branch
on:
pull_request:
branches:
- main
merge_group:
permissions: {}
# DO NOT RENAME the job's `name:` — it is referenced by GitHub branch
# protection as a required status check on `main`. Renaming silently
# breaks the gate.
jobs:
guard:
name: Verify PR source branch
runs-on: ubuntu-latest
timeout-minutes: 2
steps:
- name: Reject merge_group events
if: github.event_name == 'merge_group'
run: |
echo "::error::Merge queue is not supported for main. Disable merge queue or update this guard."
exit 1
- name: Check head branch name
env:
HEAD_REF: ${{ github.head_ref }}
HEAD_REPO: ${{ github.event.pull_request.head.repo.full_name }}
BASE_REPO: ${{ github.repository }}
run: |
echo "PR head repo: $HEAD_REPO"
echo "PR head branch: $HEAD_REF"
if [ "$HEAD_REPO" != "$BASE_REPO" ]; then
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against 'litellm_internal_staging' instead."
exit 1
fi
if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then
echo "Allowed source branch."
exit 0
fi
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_internal_staging' instead."
exit 1

1
litellm/constants.py
View file

@ -571,6 +571,7 @@ ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS: Final = 5.0
LOGGING_WORKER_CLEAR_PERCENTAGE: Final = int(
os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50)
) # Percentage of queue to clear (default: 50%)

60
litellm/litellm_core_utils/logging_worker.py
View file

@ -18,11 +18,16 @@ from litellm.constants import (
LOGGING_WORKER_CONCURRENCY,
LOGGING_WORKER_MAX_QUEUE_SIZE,
LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS,
MAX_ITERATIONS_TO_CLEAR_QUEUE,
MAX_TIME_TO_CLEAR_QUEUE,
)
def _coroutine_name(coroutine: Coroutine) -> str:
return getattr(coroutine, "__qualname__", None) or getattr(coroutine, "__name__", None) or type(coroutine).__name__
class LoggingTask(TypedDict):
"""
A logging task with its associated context to ensure logging is executed in
@ -47,10 +52,12 @@ class LoggingWorker:
timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
concurrency: int = LOGGING_WORKER_CONCURRENCY,
timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS,
):
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self.timeout_summary_window = timeout_summary_window
self._queue: asyncio.Queue[LoggingTask] | None = None
self._worker_task: asyncio.Task | None = None
self._running_tasks: set[asyncio.Task] = set()
@ -59,6 +66,10 @@ class LoggingWorker:
self._bound_loop: asyncio.AbstractEventLoop | None = None
self._last_aggressive_clear_time: float = 0.0
self._aggressive_clear_in_progress: bool = False
self._timeout_total: int = 0
self._timeout_burst_count: int = 0
self._timeout_last_callback: str | None = None
self._timeout_summary_task: asyncio.Task | None = None
# Register cleanup handler to flush remaining events on exit
atexit.register(self._flush_on_exit)
@ -136,6 +147,8 @@ class LoggingWorker:
self._sem = None
self._worker_task = None
self._running_tasks.clear()
self._timeout_summary_task = None
self._timeout_burst_count = 0
self._queue = new_queue
self._bound_loop = current_loop
return
@ -156,12 +169,15 @@ class LoggingWorker:
"""Runs the logging task and handles cleanup. Releases semaphore when done."""
try:
if self._queue is not None:
# Run the coroutine in its original context
callback_task: Final = task["context"].run(asyncio.create_task, task["coroutine"])
try:
# Run the coroutine in its original context
await asyncio.wait_for(
task["context"].run(asyncio.create_task, task["coroutine"]),
timeout=self.timeout,
)
await asyncio.wait_for(callback_task, timeout=self.timeout)
except asyncio.TimeoutError as e:
if callback_task.cancelled():
self._record_callback_timeout(task["coroutine"])
else:
verbose_logger.exception("LoggingWorker error: %s", e)
except Exception as e:
verbose_logger.exception("LoggingWorker error: %s", e)
finally:
@ -171,6 +187,35 @@ class LoggingWorker:
# Always release semaphore, even if queue is None
sem.release()
def _record_callback_timeout(self, coroutine: Coroutine) -> None:
"""Count a callback timeout and arm a debounced summary, so a burst of timeouts
(e.g. a slow Redis timing out many callbacks at once) logs one bounded line rather
than a full ERROR stacktrace per callback."""
self._timeout_total += 1
self._timeout_burst_count += 1
self._timeout_last_callback = _coroutine_name(coroutine)
if self._timeout_summary_task is None or self._timeout_summary_task.done():
self._timeout_summary_task = asyncio.create_task(self._flush_timeout_summary())
async def _flush_timeout_summary(self) -> None:
"""After the burst settles, log one bounded summary covering every timeout in it."""
await asyncio.sleep(self.timeout_summary_window)
self._emit_timeout_summary()
def _emit_timeout_summary(self) -> None:
"""Log one bounded summary for the current burst and reset the burst counter."""
burst_count: Final = self._timeout_burst_count
self._timeout_burst_count = 0
if burst_count <= 0:
return
verbose_logger.warning(
"LoggingWorker: %d callback(s) timed out after %ss (callback: %s); %d timed out since start",
burst_count,
self.timeout,
self._timeout_last_callback,
self._timeout_total,
)
async def _worker_loop(self) -> None:
"""Main worker loop that gets tasks and schedules them to run concurrently."""
try:
@ -406,6 +451,11 @@ class LoggingWorker:
async def stop(self) -> None:
"""Stop the logging worker and clean up resources."""
if self._timeout_summary_task is not None:
self._timeout_summary_task.cancel()
self._timeout_summary_task = None
self._emit_timeout_summary()
if self._worker_task is None and not self._running_tasks:
# No worker launched and no in-flight tasks to drain.
return

8
litellm/llms/base_llm/base_utils.py
View file

@ -50,6 +50,14 @@ class BaseLLMModelInfo(ABC):
"""
return None
def get_model_cost_key(self, model: str) -> str | None:
"""
Maps the model name a user sends to the key `litellm.model_cost` stores it under, when the two differ.
`get_model_info` tries this key once the exact `model` and `provider/model` keys miss. The default None means
the provider's user-facing names already match the cost map, so there is nothing extra to try.
"""
return None
@abstractmethod
def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
"""

3
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -602,6 +602,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
return None
return max(matches, key=lambda match: len(match[0]))[1]
def get_model_cost_key(self, model: str) -> str:
return f"fireworks_ai/{resolve_fireworks_resource_name(model)}"
def get_provider_info(self, model: str) -> ProviderSpecificModelInfo:
supports_function_calling_value: Final = self._get_model_cost_capability(
model=model, capability="supports_function_calling"

5
litellm/proxy/_experimental/mcp_server/AGENTS.md
View file

@ -67,9 +67,8 @@ module materially harder to understand.
auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them
behind a single generic branch unless tests prove every mode still behaves
correctly.
- Be especially careful with `available_on_public_internet: false` combined with
`delegate_auth_to_upstream: true`. The local `CLAUDE.md` explains the anonymous
upstream PKCE path that must remain intentional.
- Be especially careful with legacy `delegate_auth_to_upstream: true`. The local
`CLAUDE.md` explains its admitted replacement and public discovery contract.
- Keep database-backed fields in sync across migrations, typed models under
`litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this
package, and dashboard state when the field is user-visible.

2
litellm/proxy/_experimental/mcp_server/CLAUDE.md
View file

@ -1 +1 @@
MCP note: **`available_on_public_internet: false` with `delegate_auth_to_upstream: true` (oauth2, interactive - not `client_credentials`)** - LiteLLM still allows the anonymous upstream PKCE path (no proxy API key for `/authorize` and matching MCP routes). The internal-only flag mainly affects other surfaces (e.g. IP-based discovery). Rely on the upstream IdP and network policy; the dashboard shows a warning when both are set, and the proxy logs a warning when the server is loaded from config or the database
MCP note: **`auth_type: oauth2` with `delegate_auth_to_upstream: true` is deprecated** - LiteLLM admission is required for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. OAuth discovery endpoints stay public so clients can start the RFC 9728 flow

119
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -129,10 +129,9 @@ def _is_mcp_passthrough_cold_start(mcp_servers: list[str] | None, client_ip: str
spec-compliant WWW-Authenticate challenge instead of surfacing a generic
admission error.
Uses "all" semantics (mirrors
:meth:`MCPRequestHandler._target_servers_delegate_auth_to_upstream`): one
non-passthrough target in a co-targeted set must not flip the bypass open
for the others. Fails closed when any target cannot be resolved."""
Uses "all" semantics: one non-passthrough target in a co-targeted set must
not flip the bypass open for the others. Fails closed when any target
cannot be resolved."""
if not mcp_servers:
return False
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@ -146,6 +145,27 @@ def _is_mcp_passthrough_cold_start(mcp_servers: list[str] | None, client_ip: str
return True
def _is_legacy_delegate_cold_start(mcp_servers: list[str] | None, client_ip: str | None) -> bool:
"""Allow only credential-free legacy delegates to reach the route's OAuth challenge."""
if not mcp_servers:
return False
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
for name in mcp_servers:
server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)
if server is None or server.auth_type != MCPAuth.oauth2:
return False
if server.delegate_auth_to_upstream is not True:
return False
if MCPServerManager.effective_oauth2_flow(server) == "client_credentials":
return False
return True
def _is_litellm_auth_admission_error(exc: Exception) -> bool:
if isinstance(exc, HTTPException):
return exc.status_code == 401
@ -277,9 +297,18 @@ def _admission_failure_fallback(
mcp_servers_from_path is not None
and not _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers)
and _is_litellm_auth_admission_error(exc)
and _is_mcp_passthrough_cold_start(
mcp_servers_from_path,
client_ip=IPAddressUtils.get_mcp_client_ip(request),
and (
_is_mcp_passthrough_cold_start(
mcp_servers_from_path,
client_ip=IPAddressUtils.get_mcp_client_ip(request),
)
or (
not bearer_presented
and _is_legacy_delegate_cold_start(
mcp_servers_from_path,
client_ip=IPAddressUtils.get_mcp_client_ip(request),
)
)
)
):
verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter")
@ -434,22 +463,6 @@ class MCPRequestHandler:
api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}",
request=request,
)
elif MCPRequestHandler._target_servers_delegate_auth_to_upstream(
path=request_route,
mcp_servers=mcp_servers,
client_ip=IPAddressUtils.get_mcp_client_ip(request),
):
# Operator opted this oauth2 server into upstream-delegated auth: the
# client authenticates directly with the upstream MCP server, so any
# Authorization bearer is an upstream token, never a LiteLLM key. Skip
# LiteLLM validation entirely — covering both the no-credential
# discovery request and the authenticated call carrying the upstream
# bearer — so a tool call that succeeds never carries a phantom 401
# auth span; the bearer is forwarded upstream unchanged. Gated by
# _target_servers_delegate_auth_to_upstream, which returns True only
# when EVERY target is auth_type=oauth2 with delegate_auth_to_upstream
# set; fails closed otherwise.
validated_user_api_key_auth = UserAPIKeyAuth()
elif MCPRequestHandler._target_servers_are_true_passthrough(
path=request_route,
mcp_servers=mcp_servers,
@ -660,64 +673,6 @@ class MCPRequestHandler:
return [single_server_match.group(1)]
return [servers_and_path]
@staticmethod
def _target_servers_delegate_auth_to_upstream(
path: str, mcp_servers: list[str] | None, client_ip: str | None
) -> bool:
"""
True only when EVERY MCP server the request targets is configured for
``auth_type == oauth2`` AND has ``delegate_auth_to_upstream=True``.
Fails closed when any target does not opt in or cannot be resolved.
Used by :meth:`process_mcp_request` to skip LiteLLM API-key/SSO auth
entirely (PKCE passthrough) so the client authenticates directly with
the upstream MCP server. Mixed-target requests (e.g. one delegated +
one non-delegated server) fall back to normal LiteLLM auth.
"""
# Inline imports avoid a circular dependency: mcp_server_manager imports
# from this module.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
# Must mirror the downstream header-vs-path override
# (``extract_mcp_auth_context``) or an attacker could set
# ``x-mcp-servers`` to a delegate-enabled server while the URL path
# targets a non-delegate server, skipping LiteLLM auth for it.
target_names: Final = MCPRequestHandler._resolve_target_server_names(path=path, mcp_servers_header=mcp_servers)
if not target_names:
return False
for name in target_names:
server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip)
if server is None or server.auth_type != MCPAuth.oauth2:
return False
# `is True` is intentional: opt-in must be an explicit boolean
# True. A MagicMock attribute (in tests) or any other truthy
# non-bool must not silently enable the bypass.
if getattr(server, "delegate_auth_to_upstream", False) is not True:
return False
# Never delegate for M2M (client_credentials) servers: LiteLLM
# fetches the upstream token automatically using stored credentials,
# so allowing anonymous bypass would let any external caller invoke
# tools authenticated as LiteLLM's service account.
#
# Resolve the flow rather than reading has_client_credentials directly:
# this is a security gate, and a legacy row whose oauth2_flow was never
# stamped still carries the M2M credential shape (client_id/secret +
# token_url, no authorization_url). Treating an unstamped-but-M2M-shaped
# row as non-M2M here would reopen the anonymous bypass the explicit
# column no longer closes on its own. Shares the one resolution helper
# with the egress backstop and the anonymous-delegate allowlist; all fail
# closed on the ambiguous shape and are removed together once no null rows
# remain. A pure-PKCE delegate server (no stored credentials) resolves to a
# non-M2M flow and keeps its bypass.
if MCPServerManager.effective_oauth2_flow(server) == "client_credentials":
return False
return True
@staticmethod
def _target_servers_are_true_passthrough(path: str, mcp_servers: list[str] | None, client_ip: str | None) -> bool:
"""
@ -726,7 +681,7 @@ class MCPRequestHandler:
Used by :meth:`process_mcp_request` to skip LiteLLM admission auth entirely: the gateway is a
transparent proxy and the caller's ``Authorization`` is an upstream token, never a LiteLLM key.
Mirrors :meth:`_target_servers_delegate_auth_to_upstream`; a mixed-target request keeps normal auth.
A mixed-target request keeps normal auth.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,

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

@ -1055,7 +1055,7 @@ def _should_strip_caller_authorization(
pass-through cold-start case (RFC 9728) the bearer in
``Authorization`` is the upstream OAuth token and must be
forwarded, so we keep it.
- **oauth_delegate servers**: admission always runs and there is no
- **Delegated OAuth servers**: admission always runs and there is no
anonymous path, so the caller's separate ``Authorization`` is
forwarded only when a distinct ``x-litellm-api-key`` carried
admission. Without that header the ``Authorization`` *was* the
@ -1075,12 +1075,17 @@ def _should_strip_caller_authorization(
# upstream — it would override another user's stored credential. Delegate and
# pass-through return None from to_server_spec and keep forwarding the bearer.
return True
if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate):
is_delegated_oauth: Final = mcp_server.is_oauth_delegate or (
mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream
)
if not (mcp_server.is_oauth_passthrough or is_delegated_oauth):
return False
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
if mcp_server.is_oauth_delegate:
return not has_explicit_litellm_admission_header
if is_delegated_oauth:
return not has_explicit_litellm_admission_header or _authorization_is_litellm_admission_credential(
raw_headers, user_api_key_auth
)
return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or (
user_api_key_auth is None and not has_explicit_litellm_admission_header
)
@ -1107,15 +1112,11 @@ def _authorization_is_litellm_admission_credential(
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
same key in both headers.
"""
if user_api_key_auth is None or not user_api_key_auth.api_key:
return False
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
if not admission_header:
return True
authorization: Final = _raw_header_value(raw_headers, "authorization")
return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(
admission_header, "Bearer"
)
if admission_header and authorization:
return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer")
return bool(user_api_key_auth and user_api_key_auth.api_key and not admission_header)
def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
@ -1453,22 +1454,19 @@ def _warn_on_server_name_fields(
_warn("server_name", server_name)
def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str) -> None:
"""Surface internal + upstream PKCE delegate in logs for operators."""
def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None:
"""Direct legacy delegated OAuth configurations to the admitted replacement."""
if server.auth_type != MCPAuth.oauth2:
return
if getattr(server, "delegate_auth_to_upstream", False) is not True:
return
if getattr(server, "available_on_public_internet", True):
return
if server.has_client_credentials:
return
label: Final = get_server_prefix(server)
verbose_logger.warning(
"MCP server %r (id=%s, source=%s): internal-only (available_on_public_internet=false) "
"with delegate_auth_to_upstream=true. Anonymous callers can reach the upstream OAuth2 "
"/authorize flow and complete PKCE without a LiteLLM API key session; ensure the "
"upstream IdP and network enforce your access policy.",
"MCP server %r (id=%s, source=%s) uses deprecated auth_type=oauth2 with "
"delegate_auth_to_upstream=true. LiteLLM admission is now required; migrate to "
"auth_type=oauth_delegate for client-forwarded OAuth.",
label,
server.server_id,
source,
@ -2640,7 +2638,7 @@ class MCPServerManager:
oauth_identity_binding=server_config.get("oauth_identity_binding", None),
)
self._assign_unique_short_prefix(new_server)
_warn_internal_delegate_pkce_if_applicable(new_server, source="config")
_warn_legacy_delegate_auth_if_applicable(new_server, source="config")
_warn_config_id_jag_server_outruns_sso(new_server)
self._invalidate_discovery_lists(server_id)
self.config_mcp_servers[server_id] = new_server
@ -3185,7 +3183,7 @@ class MCPServerManager:
timeout=getattr(mcp_server, "timeout", None),
max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None),
)
_warn_internal_delegate_pkce_if_applicable(new_server, source="database")
_warn_legacy_delegate_auth_if_applicable(new_server, source="database")
self._set_oauth_discovery_deferred(
new_server.server_id,
_requires_oauth_discovery(server_url, use_issuer_anchor, new_server),
@ -3479,10 +3477,6 @@ class MCPServerManager:
)
)
# For anonymous callers (no user_id, no role), also surface any
# servers the operator has opted into upstream-delegated auth.
# These servers handle their own auth at the upstream level, so
# LiteLLM granting access here does not bypass any security gate.
is_anonymous: Final = not (
user_api_key_auth
and (
@ -3492,23 +3486,12 @@ class MCPServerManager:
)
)
if is_anonymous:
delegate_server_ids: Final = [
passthrough_server_ids: Final = [
server.server_id
for server in self.get_registry().values()
if (
getattr(server, "auth_type", None) == MCPAuth.oauth2
and getattr(server, "delegate_auth_to_upstream", False) is True
# M2M servers must not be exposed anonymously: an
# unauthenticated caller would get LiteLLM to proxy tool
# calls using its stored client_credentials. Resolve the flow
# rather than reading has_client_credentials so an unstamped
# M2M-shape row (null column, verbatim-read as non-M2M) still
# fails closed here, matching the anonymous-delegate auth gate.
and MCPServerManager.effective_oauth2_flow(server) != "client_credentials"
)
or getattr(server, "auth_type", None) == MCPAuth.true_passthrough
if getattr(server, "auth_type", None) == MCPAuth.true_passthrough
]
combined_servers.update(delegate_server_ids)
combined_servers.update(passthrough_server_ids)
restrict_allow_all: Final = (
resolved_general_settings.get("mcp_allow_all_keys_respects_mcp_scope", False)

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

@ -4257,20 +4257,6 @@ if MCP_AVAILABLE:
return None
return _get_authorization_header_from_scope(scope)
def _is_delegate_upstream_probe_target(server: MCPServer) -> bool:
"""Whether ``server`` is an interactive delegate-auth server whose client-supplied
token should be preflighted upstream.
Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is
resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed
(its stored client credentials drive egress; the caller's bearer is irrelevant).
"""
return (
server.auth_type == MCPAuth.oauth2
and server.delegate_auth_to_upstream is True
and MCPServerManager.effective_oauth2_flow(server) != "client_credentials"
)
async def _probe_upstream_auth(
url: str,
auth_header: str,
@ -4331,7 +4317,7 @@ if MCP_AVAILABLE:
mcp_servers: list[str] | None,
client_ip: str | None,
) -> None:
"""Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts.
"""Probe pass-through upstream servers in parallel before the MCP session starts.
Only servers the caller's key is already authorized to reach are probed —
the list is derived from _get_allowed_mcp_servers so that a user cannot
@ -4343,38 +4329,9 @@ if MCP_AVAILABLE:
if the upstream accepts it but forbids the caller.
Fails-open: network errors are logged and the request is allowed through.
Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``)
are probed with the caller's bare ``Authorization`` bearer. That bearer is only
an upstream token (never a LiteLLM key) when admission took the delegate bypass,
so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same
resolver admission used -- rather than the wider allowed-server prefix/access-group
matching. A name that only reaches a delegate server via server_id or an access
group would have been admitted as a real LiteLLM key, so probing it would leak that
key upstream; requiring the admission-resolver match closes that gap. Without the
probe a rejected token is absorbed by the tools/list handler and masked as an empty
tool list. Gated to single-server routes so one rejected token cannot 401 a
multi-server aggregate connect, matching the OBO preflight gating; the challenge
echoes the requested name so aliased routes get the same resource_metadata URL as
the tokenless preemptive challenge.
"""
forwarded_auth: Final = _get_forwarded_auth_from_scope(scope)
requested_single_target: Final = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None
# The bare Authorization header (no x-litellm-api-key) is a valid upstream token
# only when admission classified it as one, i.e. the single requested name resolves
# to a delegate server under admission's own resolver. Resolve it the same way here
# so a server_id- or access-group-named delegate (which admission would have treated
# as a LiteLLM key) is never probed with that key.
delegate_server: Final = (
global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip)
if requested_single_target
else None
)
delegate_auth: Final = (
_get_authorization_header_from_scope(scope)
if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server)
else None
)
if not forwarded_auth and not delegate_auth:
if not forwarded_auth:
return
# Use the authorized server set, not the raw user-supplied names, so that
@ -4384,35 +4341,20 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
client_ip=client_ip,
)
passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = (
tuple(
(srv, forwarded_auth, srv.name)
for srv in allowed_servers
# Restrict to genuine OAuth pass-through servers (auth_type none +
# Authorization in extra_headers). Gateway-managed OAuth2 servers
# must not receive the ``resource_metadata=`` challenge emitted
# below — they require ``authorization_uri=`` pointing at the
# gateway AS metadata. ``is_oauth_passthrough`` already requires
# ``auth_type in (None, MCPAuth.none)``, which is mutually
# exclusive with ``has_client_credentials`` (oauth2 + M2M flow),
# so M2M servers are implicitly excluded here.
if srv.is_oauth_passthrough
)
if forwarded_auth
else ()
passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple(
(srv, forwarded_auth, srv.name)
for srv in allowed_servers
# Restrict to genuine OAuth pass-through servers (auth_type none +
# Authorization in extra_headers). Gateway-managed OAuth2 servers
# must not receive the ``resource_metadata=`` challenge emitted
# below — they require ``authorization_uri=`` pointing at the
# gateway AS metadata. ``is_oauth_passthrough`` already requires
# ``auth_type in (None, MCPAuth.none)``, which is mutually
# exclusive with ``has_client_credentials`` (oauth2 + M2M flow),
# so M2M servers are implicitly excluded here.
if srv.is_oauth_passthrough
)
# Probe the admission-resolved delegate server only when the caller is actually
# authorized for it (present in the IP-filtered allowed set), keyed by server_id.
delegate_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = (
tuple(
(srv, delegate_auth, requested_single_target)
for srv in allowed_servers
if delegate_server is not None and srv.server_id == delegate_server.server_id
)
if delegate_auth and requested_single_target
else ()
)
probe_targets: Final = passthrough_targets + delegate_targets
probe_targets: Final = passthrough_targets
if not probe_targets:
return

36
litellm/proxy/db/spend_log_tool_index.py
View file

@ -5,8 +5,10 @@ At request time the spend writer builds one ToolUsageTransaction per request tha
invoked tools (MCP namespaced tool name plus response tool_calls; declared-but-not-
invoked tools are excluded) and queues it on the prisma client. The spend-log flush
job drains the queue into LiteLLM_SpendLogToolIndex (per-request drill-down) and
LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads) in a
single transaction, so a failed flush never leaves a partial rollup increment.
LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads). The
index rows are keyed on (request_id, tool_name) and written with skip_duplicates,
so they go out as bounded standalone statements; every rollup upsert stays in one
transaction, so a failed flush never leaves a partial rollup increment.
"""
from __future__ import annotations
@ -19,7 +21,10 @@ from datetime import datetime, timezone
from itertools import groupby
from typing import TYPE_CHECKING, Any, Final
from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.db.spend_log_batching import spend_log_write_batches
from litellm.repositories.table_repositories import SpendLogToolIndexRepository
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
@ -98,14 +103,19 @@ async def flush_tool_usage_transactions(
transactions: Sequence[ToolUsageTransaction],
n_retry_times: int = 3,
) -> None:
"""Write index rows and rollup upserts for a drained queue batch in one
transaction. Retries only ConnectError, the one failure that proves the
statements never reached the database. Post-send failures (Read timeouts
and errors) are ambiguous and are NOT retried: the engine can abandon the
transaction open on the pooled connection, so a retry's statements stack
into the same transaction and one commit applies both increment sets.
Ambiguous failures drop the batch; the caller logs it at error. Callers
must not add their own retry around this function."""
"""Write the index rows as bounded standalone statements, then every rollup
upsert for the drained queue batch in one transaction. One flush fans out to
transactions x tools index rows, so the index write is split by the spend-log
statement budgets; a split inside ``batch_()`` would not help, since the
batcher ships every queued statement to the query engine as one payload.
Retries only ConnectError, the one failure that proves the statements never
reached the database; replayed index rows are no-ops under skip_duplicates.
Post-send failures (Read timeouts and errors) are ambiguous and are NOT
retried: the engine can abandon the transaction open on the pooled
connection, so a retry's statements stack into the same transaction and one
commit applies both increment sets. Ambiguous failures drop the batch; the
caller logs it at error. Callers must not add their own retry around this
function."""
if not transactions:
return
@ -119,10 +129,14 @@ async def flush_tool_usage_transactions(
key=lambda entry: (entry[0], entry[1]),
)
index_table: Final = SpendLogToolIndexRepository(prisma_client).table
for attempt in range(n_retry_times + 1):
try:
for statement_rows in spend_log_write_batches(
index_rows, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS
):
await index_table.create_many(data=statement_rows, skip_duplicates=True)
async with prisma_client.db.batch_() as batcher:
batcher.litellm_spendlogtoolindex.create_many(data=index_rows, skip_duplicates=True)
for (date_key, tool_name), grouped in groupby(per_tool_day, key=lambda entry: (entry[0], entry[1])):
entries = tuple(grouped)
spend = sum(entry[2] for entry in entries)

13
litellm/proxy/guardrails/usage_tracking.py
View file

@ -17,8 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guardrail_cost_by_unit
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.db.spend_log_batching import spend_log_write_batches
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
@ -401,13 +403,12 @@ async def process_spend_logs_guardrail_usage(
return
try:
# Insert index rows (skip duplicates by request_id + guardrail_id)
if index_rows:
index_table: Final = SpendLogGuardrailIndexRepository(prisma_client).table
for statement_rows in spend_log_write_batches(
index_rows, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS
):
try:
await SpendLogGuardrailIndexRepository(prisma_client).table.create_many(
data=index_rows,
skip_duplicates=True,
)
await index_table.create_many(data=statement_rows, skip_duplicates=True)
except Exception as e:
verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e)

1
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -64,6 +64,7 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
CallTypes.pass_through.value,
CallTypes.llm_passthrough_route.value,
CallTypes.allm_passthrough_route.value,
CallTypes.call_mcp_tool.value,
# CheckBatchCost's synthetic logging_obj for a completed managed batch carries
# whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is
# None for a batch created before those columns were persisted, or by the master

109
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -117,6 +117,14 @@ class ResponsesToolChatForm:
web_search_options: OpenAIWebSearchOptions | None
@dataclass(frozen=True, slots=True)
class ResponsesReasoningChatForm:
"""The Responses ``reasoning`` object as the two params Chat Completions takes."""
effort: str | None
summary: str | None
if TYPE_CHECKING:
from openai.types.responses.response_apply_patch_tool_call import (
ResponseApplyPatchToolCall,
@ -309,6 +317,79 @@ class LiteLLMCompletionResponsesConfig:
)
return supported_params is not None and "web_search_options" not in supported_params
@staticmethod
def _completion_bridges_back_to_responses_api(
model: str,
custom_llm_provider: str | None,
tools: Sequence[ChatCompletionToolParam | OpenAIMcpServerTool] | None,
web_search_options: OpenAIWebSearchOptions | None,
reasoning_effort: str | None,
reasoning_summary: str | None,
api_base: str | None,
) -> bool:
"""
Whether ``litellm.completion`` will route this model back onto the Responses API.
Delegates to the same check ``litellm.completion`` itself runs, and is asked with the
params this transform is about to emit, so the two cannot reach different answers.
"""
from litellm.main import responses_api_bridge_check
try:
model_info, _ = responses_api_bridge_check(
model=model,
custom_llm_provider=custom_llm_provider or "",
web_search_options=web_search_options,
tools=tools,
reasoning_effort=reasoning_effort,
reasoning_summary=reasoning_summary,
api_base=api_base,
)
except Exception as e: # noqa: BLE001 # a capability probe must never fail the request it probes for
verbose_logger.debug("responses bridge: reasoning effort mode check failed: %s", e)
return False
return model_info.get("mode") == "responses"
@staticmethod
def _transform_reasoning_for_chat_completion(
reasoning_param: Reasoning | str | None,
model: str,
custom_llm_provider: str | None,
tools: Sequence[ChatCompletionToolParam | OpenAIMcpServerTool] | None = None,
web_search_options: OpenAIWebSearchOptions | None = None,
api_base: str | None = None,
) -> ResponsesReasoningChatForm:
"""
Split the Responses ``reasoning`` object into the params Chat Completions understands.
``reasoning_effort`` is a string enum there, so the object is never forwarded whole: a chat
provider either rejects it or silently drops it, and dropping it turns reasoning off while
still billing for the turn. ``summary`` has no chat equivalent, so it rides the
``reasoning_summary`` alias, which ``litellm.completion`` reassembles into ``{effort,
summary}`` when it bridges the model back onto the Responses API, and is sent to nothing
else.
"""
if not reasoning_param:
return ResponsesReasoningChatForm(effort=None, summary=None)
if isinstance(reasoning_param, str):
return ResponsesReasoningChatForm(effort=reasoning_param, summary=None)
effort: Final = reasoning_param.get("effort")
summary: Final = reasoning_param.get("summary")
if summary is None:
return ResponsesReasoningChatForm(effort=effort, summary=None)
bridges_back: Final = LiteLLMCompletionResponsesConfig._completion_bridges_back_to_responses_api(
model=model,
custom_llm_provider=custom_llm_provider,
tools=tools,
web_search_options=web_search_options,
reasoning_effort=effort,
reasoning_summary=summary,
api_base=api_base,
)
return ResponsesReasoningChatForm(effort=effort, summary=summary if bridges_back else None)
@staticmethod
def transform_responses_api_request_to_chat_completion_request(
model: str,
@ -339,23 +420,14 @@ class LiteLLMCompletionResponsesConfig:
if text_param:
response_format = LiteLLMCompletionResponsesConfig._transform_text_format_to_response_format(text_param)
# Extract reasoning_effort from reasoning parameter
reasoning_effort: Reasoning | str | None = None
reasoning_param: Final = responses_api_request.get("reasoning")
if reasoning_param:
if isinstance(reasoning_param, dict):
# reasoning can be {"effort": "low|medium|high", "summary": "detailed"}
# Keep the full dict when summary is set so the responses API bridge can
# forward it; otherwise use the effort string for chat completion (e.g. Gemini).
if "summary" in reasoning_param:
reasoning_effort = reasoning_param
elif "effort" in reasoning_param:
reasoning_effort = reasoning_param.get("effort")
else:
reasoning_effort = reasoning_param
elif isinstance(reasoning_param, str):
# reasoning could be a string directly
reasoning_effort = reasoning_param
reasoning: Final = LiteLLMCompletionResponsesConfig._transform_reasoning_for_chat_completion(
reasoning_param=responses_api_request.get("reasoning"),
model=model,
custom_llm_provider=custom_llm_provider,
tools=tools,
web_search_options=web_search_options,
api_base=kwargs.get("api_base"),
)
litellm_completion_request: dict = {
"messages": LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
@ -378,7 +450,8 @@ class LiteLLMCompletionResponsesConfig:
"service_tier": kwargs.get("service_tier"),
"web_search_options": web_search_options,
"response_format": response_format,
"reasoning_effort": reasoning_effort,
"reasoning_effort": reasoning.effort,
"reasoning_summary": reasoning.summary,
"context_management": responses_api_request.get("context_management"),
# litellm specific params
"custom_llm_provider": custom_llm_provider,

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

@ -155,11 +155,9 @@ class MCPServer(BaseModel):
access_groups: list[str] | None = None
allow_all_keys: bool = False
available_on_public_internet: bool = True
# Explicit opt-in to upstream-delegated authentication for ``oauth2``
# servers. When ``auth_type == oauth2`` and this is ``True``, MCP requests
# bypass LiteLLM API-key/SSO auth (and the pre-emptive 401) so the client
# completes PKCE directly with the upstream MCP server. See
# ``MCPRequestHandler._target_servers_delegate_auth_to_upstream``.
# Legacy opt-in to upstream-delegated authentication for ``oauth2``
# servers. LiteLLM admission still applies; use ``oauth_delegate`` for the
# supported client-forwarded OAuth flow.
#
# Honored only for ``auth_type == oauth2``; ignored for any other
# ``auth_type``. OAuth pass-through for non-oauth2 servers

11
litellm/utils.py
View file

@ -5563,12 +5563,21 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P
split_model = strip_bedrock_routing_prefix(split_model)
provider_model_info: Final = (
ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider))
if custom_llm_provider in LlmProvidersSet
else None
)
provider_cost_key: Final = (
provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None
)
return PotentialModelNamesAndCustomLLMProvider(
split_model=split_model,
combined_model_name=combined_model_name,
stripped_model_name=stripped_model_name,
combined_stripped_model_name=combined_stripped_model_name,
provider_prefixed_model_name=provider_prefixed_model_name,
provider_prefixed_model_name=provider_cost_key or provider_prefixed_model_name,
custom_llm_provider=cast(str, custom_llm_provider),
)

2
tests/e2e/coverage_registry/mcp.yaml
View file

@ -78,7 +78,7 @@
auth_family: none
assertions: [succeeds]
source: "mcp_server_manager.py:1485-1492"
rationale: Public/anonymous servers; delegate_auth_to_upstream
rationale: Explicitly anonymous true_passthrough servers
- id: mcp.call_tool.none.succeeds
module: mcp
tier: P1

241
tests/test_litellm/litellm_core_utils/test_logging_worker.py
View file

@ -14,6 +14,18 @@ from litellm.constants import LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS
from litellm.litellm_core_utils.logging_worker import LoggingWorker
class _RecordCollector(logging.Handler):
"""Captures emitted log records so a test can assert on real logging output
(level, message args, traceback) instead of patching the logger object."""
def __init__(self) -> None:
super().__init__()
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None:
self.records.append(record)
class TestLoggingWorker:
"""Test cases for LoggingWorker functionality."""
@ -26,9 +38,7 @@ class TestLoggingWorker:
async def test_graceful_shutdown_with_clear_queue(self, logging_worker):
"""Test that cancellation triggers clear_queue to prevent 'never awaited' warnings."""
# Mock the clear_queue method to verify it's called during cancellation
with patch.object(
logging_worker, "clear_queue", new_callable=AsyncMock
) as mock_clear_queue:
with patch.object(logging_worker, "clear_queue", new_callable=AsyncMock) as mock_clear_queue:
# Start the worker
logging_worker.start()
@ -195,9 +205,7 @@ class TestLoggingWorker:
async def test_worker_handles_cancellation_gracefully(self, logging_worker):
"""Test that the worker handles cancellation without throwing exceptions."""
# Mock verbose_logger to capture debug messages
with patch(
"litellm.litellm_core_utils.logging_worker.verbose_logger"
) as mock_logger:
with patch("litellm.litellm_core_utils.logging_worker.verbose_logger") as mock_logger:
# Start the worker
logging_worker.start()
@ -264,29 +272,21 @@ class TestLoggingWorker:
small_worker._ensure_queue()
# Mock verbose_logger to capture exception messages
with patch(
"litellm.litellm_core_utils.logging_worker.verbose_logger"
) as mock_logger:
with patch("litellm.litellm_core_utils.logging_worker.verbose_logger") as mock_logger:
# Fill the queue beyond capacity
mock_coro = AsyncMock()
for _ in range(5): # More than max_queue_size of 2
small_worker.enqueue(mock_coro())
# Should have logged queue full exceptions
exception_calls = [
call
for call in mock_logger.exception.call_args_list
if "queue is full" in str(call)
]
exception_calls = [call for call in mock_logger.exception.call_args_list if "queue is full" in str(call)]
assert len(exception_calls) > 0
@pytest.mark.asyncio
async def test_context_propagation(self, logging_worker):
"""Test that enqueued tasks execute in their original context."""
# Create a context variable for testing
test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar(
"test_context_var"
)
test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar("test_context_var")
# Track results from multiple tasks using asyncio.Event for synchronization
task_results = []
@ -364,36 +364,28 @@ class TestLoggingWorker:
task_results.sort(key=lambda x: x["task_id"])
# Verify that each task saw its own context
assert (
len(task_results) == 3
), f"Expected 3 results, got {len(task_results)}: {task_results}"
assert len(task_results) == 3, f"Expected 3 results, got {len(task_results)}: {task_results}"
# Task 1 should see "context_1"
task1_result = next((r for r in task_results if r["task_id"] == "task_1"), None)
assert task1_result is not None, "Task 1 result not found"
assert (
task1_result["context_accessible"] is True
), "Task 1 should have access to context variable"
assert (
task1_result["context_value"] == "context_1"
), f"Task 1 should see 'context_1', got: {task1_result['context_value']}"
assert task1_result["context_accessible"] is True, "Task 1 should have access to context variable"
assert task1_result["context_value"] == "context_1", (
f"Task 1 should see 'context_1', got: {task1_result['context_value']}"
)
# Task 2 should see "context_2"
task2_result = next((r for r in task_results if r["task_id"] == "task_2"), None)
assert task2_result is not None, "Task 2 result not found"
assert (
task2_result["context_accessible"] is True
), "Task 2 should have access to context variable"
assert (
task2_result["context_value"] == "context_2"
), f"Task 2 should see 'context_2', got: {task2_result['context_value']}"
assert task2_result["context_accessible"] is True, "Task 2 should have access to context variable"
assert task2_result["context_value"] == "context_2", (
f"Task 2 should see 'context_2', got: {task2_result['context_value']}"
)
# Task 3 should not have access to the context variable
task3_result = next((r for r in task_results if r["task_id"] == "task_3"), None)
assert task3_result is not None, "Task 3 result not found"
assert (
task3_result["context_accessible"] is False
), "Task 3 should not have access to context variable"
assert task3_result["context_accessible"] is False, "Task 3 should not have access to context variable"
@pytest.mark.asyncio
async def test_semaphore_concurrency_limit(self):
@ -525,3 +517,182 @@ class TestLoggingWorker:
asyncio.run(rebind_on_second_loop())
assert sorted(executed) == [0, 1, 2, 3, 4]
@pytest.mark.asyncio
async def test_timeout_burst_logs_one_bounded_summary(self):
"""Regression (LIT-7519): a burst of callback timeouts must log one bounded summary,
not a full ERROR traceback per timed-out callback.
Before the fix every timed-out callback hit ``verbose_logger.exception`` in
``_process_log_task``, so a slow Redis timing out many spend-tracking callbacks at once
produced one stacktrace each, clustered milliseconds apart. They must collapse into a
single WARNING that counts them, with no tracebacks.
"""
timeout_count = 25
worker = LoggingWorker(
timeout=0.05,
max_queue_size=100,
concurrency=100,
timeout_summary_window=1.0,
)
async def slow_callback() -> None:
await asyncio.sleep(10)
logger = logging.getLogger("LiteLLM")
collector = _RecordCollector()
previous_level = logger.level
logger.addHandler(collector)
logger.setLevel(logging.DEBUG)
try:
worker.start()
for _ in range(timeout_count):
worker.enqueue(slow_callback())
await worker.flush()
assert worker._timeout_summary_task is not None
await worker._timeout_summary_task
await worker.stop()
finally:
logger.removeHandler(collector)
logger.setLevel(previous_level)
errors = [r for r in collector.records if r.levelno >= logging.ERROR]
warnings = [r for r in collector.records if r.levelno == logging.WARNING]
assert errors == [], "a timeout burst must not emit any ERROR tracebacks"
assert len(warnings) == 1, "the whole burst must collapse into one summary line"
summary = warnings[0].getMessage()
assert f"{timeout_count} callback(s) timed out" in summary
assert f"{timeout_count} timed out since start" in summary
assert "slow_callback" in summary
@pytest.mark.asyncio
async def test_non_timeout_error_keeps_traceback(self):
"""A real programming error in a callback must still log a full traceback, so the
timeout aggregation never hides genuine failures.
"""
worker = LoggingWorker(timeout=5.0, max_queue_size=10)
async def failing_callback() -> None:
raise ValueError("boom")
logger = logging.getLogger("LiteLLM")
collector = _RecordCollector()
previous_level = logger.level
logger.addHandler(collector)
logger.setLevel(logging.DEBUG)
try:
worker.start()
worker.enqueue(failing_callback())
await worker.flush()
await worker.stop()
finally:
logger.removeHandler(collector)
logger.setLevel(previous_level)
errors = [r for r in collector.records if r.levelno >= logging.ERROR]
warnings = [r for r in collector.records if r.levelno == logging.WARNING]
assert len(errors) == 1, "a real error must still be logged once"
assert errors[0].exc_info is not None, "the traceback must be preserved"
assert warnings == [], "a single real error is not a timeout summary"
assert worker._timeout_summary_task is None
@pytest.mark.asyncio
async def test_callback_raised_timeout_keeps_traceback(self):
"""A callback that raises TimeoutError on its own did not hit the worker's deadline,
so it is a real failure and must keep its traceback rather than being folded into the
bounded burst summary.
"""
worker = LoggingWorker(timeout=5.0, max_queue_size=10, timeout_summary_window=1.0)
async def raises_own_timeout() -> None:
raise asyncio.TimeoutError("callback's own downstream timeout")
logger = logging.getLogger("LiteLLM")
collector = _RecordCollector()
previous_level = logger.level
logger.addHandler(collector)
logger.setLevel(logging.DEBUG)
try:
worker.start()
worker.enqueue(raises_own_timeout())
await worker.flush()
await worker.stop()
finally:
logger.removeHandler(collector)
logger.setLevel(previous_level)
errors = [r for r in collector.records if r.levelno >= logging.ERROR]
warnings = [r for r in collector.records if r.levelno == logging.WARNING]
assert len(errors) == 1, "a callback-raised TimeoutError must still be logged once"
assert errors[0].exc_info is not None, "the traceback must be preserved"
assert warnings == [], "a callback-raised TimeoutError is not a worker-deadline timeout"
assert worker._timeout_summary_task is None
@pytest.mark.asyncio
async def test_stop_flushes_pending_timeout_summary(self):
"""A burst still inside its summary window when the worker stops must still emit its one
summary, instead of losing it when the event loop tears down.
"""
worker = LoggingWorker(
timeout=0.05,
max_queue_size=10,
concurrency=10,
timeout_summary_window=30.0,
)
async def slow_callback() -> None:
await asyncio.sleep(10)
logger = logging.getLogger("LiteLLM")
collector = _RecordCollector()
previous_level = logger.level
logger.addHandler(collector)
logger.setLevel(logging.DEBUG)
try:
worker.start()
for _ in range(3):
worker.enqueue(slow_callback())
await worker.flush()
assert worker._timeout_summary_task is not None
assert not worker._timeout_summary_task.done(), "precondition: the summary window has not elapsed"
await worker.stop()
finally:
logger.removeHandler(collector)
logger.setLevel(previous_level)
warnings = [r for r in collector.records if r.levelno == logging.WARNING]
assert len(warnings) == 1, "stop must flush the pending summary exactly once"
assert "3 callback(s) timed out" in warnings[0].getMessage()
assert worker._timeout_summary_task is None
def test_loop_change_resets_timeout_summary_state(self):
"""On an event-loop change the summary task is bound to the dead loop; it and the pending
burst count must reset so timeouts on the new loop arm a fresh summary rather than a stale,
stuck one that silently drops later summaries.
"""
worker = LoggingWorker(timeout=1.0, max_queue_size=10, timeout_summary_window=30.0)
async def timed_out_callback() -> None:
await asyncio.sleep(10)
async def arm_on_first_loop() -> None:
worker._ensure_queue()
coro = timed_out_callback()
worker._record_callback_timeout(coro)
coro.close()
assert worker._timeout_summary_task is not None
assert worker._timeout_burst_count == 1
asyncio.run(arm_on_first_loop())
assert worker._timeout_summary_task is not None
async def rebind_on_second_loop() -> None:
worker._ensure_queue()
assert worker._timeout_summary_task is None, "stale summary task must drop on loop change"
assert worker._timeout_burst_count == 0, "stale burst count must reset on loop change"
asyncio.run(rebind_on_second_loop())

12
tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
View file

@ -1813,3 +1813,15 @@ def test_streaming_preserves_selected_model_for_private_accounting():
completion_response=assembled,
custom_llm_provider="fireworks_ai",
) == pytest.approx(expected_cost)
@pytest.mark.parametrize(
"model, expected",
[
("deepseek-r1", "fireworks_ai/accounts/fireworks/models/deepseek-r1"),
("glm-5p3-fast", "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast"),
("accounts/fireworks/models/deepseek-r1", "fireworks_ai/accounts/fireworks/models/deepseek-r1"),
],
)
def test_get_model_cost_key_resolves_short_names_to_long_keys(model: str, expected: str) -> None:
assert FireworksAIConfig().get_model_cost_key(model) == expected

206
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -1483,12 +1483,10 @@ class TestMCPOAuth2AuthFlow:
as LiteLLM API keys, causing auth failures and empty tool listings.
"""
async def test_oauth2_token_in_authorization_header_fallback(self):
async def test_oauth2_token_in_authorization_header_requires_litellm_admission(self):
"""
When only the Authorization header is present with a non-LiteLLM OAuth2
token AND the target server delegates auth to upstream, LiteLLM skips its
own validation entirely (so the upstream token is never mistaken for a
virtual key) and forwards the bearer upstream.
A bare Authorization token on the legacy delegated mode must establish
a LiteLLM principal rather than entering anonymously.
"""
from litellm.types.mcp import MCPAuth
@ -1510,6 +1508,7 @@ class TestMCPOAuth2AuthFlow:
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(user_id="admitted-user"),
) as mock_auth,
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
@ -1524,9 +1523,8 @@ class TestMCPOAuth2AuthFlow:
) = await MCPRequestHandler.process_mcp_request(scope)
assert isinstance(auth_result, UserAPIKeyAuth)
# The upstream token is never validated as a LiteLLM key ...
mock_auth.assert_not_called()
# ... and is preserved for upstream forwarding.
assert auth_result.user_id == "admitted-user"
mock_auth.assert_awaited_once()
assert oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-access-token-xyz"
async def test_explicit_litellm_key_with_oauth2_authorization(self):
@ -2367,11 +2365,9 @@ class TestMCPDelegateAuthToUpstream:
"""
Tests for the ``delegate_auth_to_upstream`` per-server flag.
When set on an ``auth_type=oauth2`` MCP server, LiteLLM must skip its own
API-key/SSO check entirely so the client completes PKCE directly with the
upstream MCP server. The gate must fail closed for any non-oauth2 server,
any mixed-target request, and any request where the target cannot be
resolved.
The legacy flag no longer bypasses LiteLLM admission. OAuth discovery may
still use the anonymous cold-start challenge, but a presented bearer must
authenticate to LiteLLM unless a separate admission credential is supplied.
"""
@staticmethod
@ -2386,6 +2382,13 @@ class TestMCPDelegateAuthToUpstream:
delegate_auth_to_upstream=delegate_auth_to_upstream,
)
def test_legacy_delegate_cold_start_fails_closed_without_targets(self):
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
_is_legacy_delegate_cold_start,
)
assert _is_legacy_delegate_cold_start(None, client_ip=None) is False
def test_build_mcp_server_table_preserves_delegate_auth_to_upstream(self):
"""Registry → API list rows must expose delegate_auth_to_upstream for the UI."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@ -2439,11 +2442,10 @@ class TestMCPDelegateAuthToUpstream:
not_passthrough = passthrough.model_copy(update={"oauth_passthrough": False})
assert manager._build_mcp_server_table(not_passthrough).oauth_passthrough is False
async def test_delegate_skips_litellm_auth_with_no_authorization(self):
async def test_delegate_without_authorization_attempts_litellm_auth_before_cold_start(self):
"""
oauth2 + delegate_auth_to_upstream=True, no Authorization header at
all → anonymous UserAPIKeyAuth and ``user_api_key_auth`` is never
called.
A credential-free discovery request attempts LiteLLM admission before
the route emits its RFC 9728 challenge.
"""
from litellm.types.mcp import MCPAuth
@ -2457,6 +2459,8 @@ class TestMCPDelegateAuthToUpstream:
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=401, detail="No key provided"),
) as mock_auth,
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
@ -2466,17 +2470,12 @@ class TestMCPDelegateAuthToUpstream:
)
auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope)
assert isinstance(auth_result, UserAPIKeyAuth)
mock_auth.assert_not_called()
mock_auth.assert_awaited_once()
async def test_delegate_with_upstream_token_in_authorization_skips_litellm_auth(
self,
):
async def test_delegate_with_only_upstream_token_requires_litellm_auth(self):
"""
oauth2 + delegate_auth_to_upstream=True with an upstream OAuth token in
``Authorization``: the delegate gate fires before any LiteLLM validation,
so ``user_api_key_auth`` is never called and the bearer is forwarded
upstream untouched. Skipping the doomed validation is what keeps a tool
call that actually succeeds from carrying a phantom 401 auth span.
An upstream token cannot establish a LiteLLM principal and must not
reopen anonymous admission.
"""
from litellm.types.mcp import MCPAuth
@ -2491,6 +2490,7 @@ class TestMCPDelegateAuthToUpstream:
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=401, detail="Invalid API key"),
) as mock_auth,
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
@ -2498,17 +2498,11 @@ class TestMCPDelegateAuthToUpstream:
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
(
auth_result,
_,
_,
_,
oauth2_headers,
_,
) = await MCPRequestHandler.process_mcp_request(scope)
assert isinstance(auth_result, UserAPIKeyAuth)
assert oauth2_headers.get("Authorization") == "Bearer upstream-pkce-token"
mock_auth.assert_not_called()
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(scope)
assert exc_info.value.status_code == 401
mock_auth.assert_awaited_once()
async def test_delegate_off_still_requires_litellm_auth(self):
"""
@ -2687,15 +2681,11 @@ class TestMCPDelegateAuthToUpstream:
assert auth_result.user_id == "real-user"
mock_auth.assert_called_once()
async def test_authorization_bearer_on_delegate_server_treated_as_upstream(self):
async def test_authorization_bearer_on_delegate_server_establishes_litellm_principal(self):
"""
On a delegate server the ``Authorization`` header is, by contract, an
upstream token rather than a LiteLLM key — even when it is sk-shaped. It
is forwarded upstream without LiteLLM validation, so ``user_api_key_auth``
is not called and no LiteLLM identity is resolved. Callers who need
LiteLLM identity / spend tracking on a delegate server must supply
``x-litellm-api-key`` (see
test_explicit_litellm_key_takes_precedence_over_delegate).
A bare Authorization bearer now follows normal LiteLLM admission. A
separate x-litellm-api-key is required when Authorization is intended
for the upstream server.
"""
from litellm.types.mcp import MCPAuth
@ -2727,9 +2717,9 @@ class TestMCPDelegateAuthToUpstream:
_,
) = await MCPRequestHandler.process_mcp_request(scope)
assert isinstance(auth_result, UserAPIKeyAuth)
assert auth_result.user_id is None
assert auth_result.user_id == "real-user"
assert oauth2_headers.get("Authorization") == "Bearer sk-1234"
mock_auth.assert_not_called()
mock_auth.assert_awaited_once()
async def test_delegate_ignored_for_client_credentials_server(self):
"""
@ -2828,13 +2818,10 @@ class TestMCPDelegateAuthToUpstream:
assert exc_info.value.status_code == 401
mock_auth.assert_called_once()
async def test_delegate_bypass_for_pure_pkce_server(self):
async def test_delegate_pkce_cold_start_attempts_litellm_auth(self):
"""
oauth2 + delegate + oauth2_flow=None and NO stored client credentials
(pure PKCE, the common delegate case) → bypass must still fire. The
shape resolves to a non-M2M flow, so the security gate leaves it alone;
the fail-closed rule targets the M2M shape specifically, not every
unstamped row.
A pure PKCE server may defer a credential-free request to the route's
challenge, but normal LiteLLM admission still runs first.
"""
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -2869,13 +2856,12 @@ class TestMCPDelegateAuthToUpstream:
):
mock_mgr.get_mcp_server_by_name.return_value = pkce_server
auth, *_rest = await MCPRequestHandler.process_mcp_request(scope)
mock_auth.assert_not_called()
mock_auth.assert_awaited_once()
assert auth.api_key is None
async def test_delegate_bypass_for_internal_server(self):
async def test_internal_delegate_cold_start_attempts_litellm_auth(self):
"""
Delegate + oauth2 interactive servers bypass LiteLLM auth even when
``available_on_public_internet`` is False (internal MCPs).
Internal delegated servers follow the same admission contract.
"""
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -2910,14 +2896,11 @@ class TestMCPDelegateAuthToUpstream:
):
mock_mgr.get_mcp_server_by_name.return_value = internal_server
auth, *_rest = await MCPRequestHandler.process_mcp_request(scope)
mock_auth.assert_not_called()
mock_auth.assert_awaited_once()
assert auth.api_key is None
async def test_get_allowed_servers_excludes_client_credentials_delegate(self):
"""
get_allowed_mcp_servers must not surface M2M (client_credentials) delegate
servers to anonymous callers even if delegate_auth_to_upstream=True.
"""
async def test_get_allowed_servers_excludes_legacy_delegates(self):
"""Legacy delegated servers are never added to anonymous access."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
@ -2955,102 +2938,7 @@ class TestMCPDelegateAuthToUpstream:
):
result = await manager.get_allowed_mcp_servers(None)
assert "pkce-server" in result
assert "m2m-server" not in result
async def test_get_allowed_servers_excludes_unstamped_m2m_shape_delegate(self):
"""
The anonymous allow-list must also exclude an M2M-shape delegate server whose
oauth2_flow was never stamped (null column, verbatim-read as non-M2M). Reading
the bare has_client_credentials here would surface it to anonymous callers; the
resolved-flow check fails closed on the shape, matching the auth gate.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
manager = MCPServerManager()
pkce_server = MCPServer(
server_id="pkce-server",
name="pkce_server",
transport="http",
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
available_on_public_internet=True,
)
unstamped_m2m = MCPServer(
server_id="unstamped-m2m",
name="unstamped_m2m",
transport="http",
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
oauth2_flow=None,
client_id="cid",
client_secret="csecret",
token_url="https://idp.example.com/token",
)
assert unstamped_m2m.has_client_credentials is False
manager.registry = {
pkce_server.server_id: pkce_server,
unstamped_m2m.server_id: unstamped_m2m,
}
with patch.object(
MCPRequestHandler,
"get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[],
):
result = await manager.get_allowed_mcp_servers(None)
assert "pkce-server" in result
assert "unstamped-m2m" not in result
async def test_get_allowed_servers_includes_internal_delegate(self):
"""
Internal-only (available_on_public_internet=False) delegate servers
appear in the anonymous allow-list like public delegate servers.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
manager = MCPServerManager()
public_server = MCPServer(
server_id="public-server",
name="public_server",
transport="http",
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
available_on_public_internet=True,
)
internal_server = MCPServer(
server_id="internal-server",
name="internal_server",
transport="http",
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
available_on_public_internet=False,
)
manager.registry = {
public_server.server_id: public_server,
internal_server.server_id: internal_server,
}
with patch.object(
MCPRequestHandler,
"get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[],
):
result = await manager.get_allowed_mcp_servers(None)
assert "public-server" in result
assert "internal-server" in result
assert result == []
async def test_true_passthrough_skips_litellm_auth_anonymously(self):
"""auth_type=true_passthrough performs no admission auth: the caller's Authorization is an

6
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
View file

@ -597,7 +597,11 @@ class TestHookHeaderMergePriority:
"Authorization": "Bearer oauth2-token",
"X-OAuth": "yes",
},
raw_headers=None,
raw_headers={
"x-litellm-api-key": "Bearer sk-litellm-key",
"authorization": "Bearer oauth2-token",
},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
proxy_logging_obj=None,
hook_extra_headers={
"Authorization": "Bearer hook-jwt",

386
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -521,6 +521,93 @@ def test_prepare_mcp_server_headers_oauth2_interactive_drops_caller_authorizatio
assert extra_headers is None
def test_prepare_mcp_server_headers_legacy_delegate_strips_admission_authorization():
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="legacy-delegate-admission",
name="legacy-delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer sk-litellm-key"},
raw_headers={"authorization": "Bearer sk-litellm-key"},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
)
assert server_auth_header is None
assert extra_headers is None
def test_prepare_mcp_server_headers_legacy_delegate_preserves_separate_upstream_authorization():
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="legacy-delegate-dual-credential",
name="legacy-delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer upstream-token"},
raw_headers={
"x-litellm-api-key": "Bearer sk-litellm-key",
"authorization": "Bearer upstream-token",
},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
)
assert server_auth_header is None
assert extra_headers == {"Authorization": "Bearer upstream-token"}
def test_prepare_mcp_server_headers_legacy_delegate_strips_repeated_admission_key():
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="legacy-delegate-repeated-key",
name="legacy-delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer sk-litellm-key"},
raw_headers={
"x-litellm-api-key": "Bearer sk-litellm-key",
"authorization": "Bearer sk-litellm-key",
},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
)
assert server_auth_header is None
assert extra_headers is None
def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers():
"""M2M must not merge caller Authorization from raw_headers when extra_headers lists it."""
try:
@ -6157,15 +6244,8 @@ def _patch_delegate_resolver(server: MCPServer, *resolvable_names: str):
@pytest.mark.asyncio
async def test_delegate_bad_token_gets_connect_time_401():
"""Regression (LIT-4194): a rejected upstream token on a delegate-auth server
must fail the connect with 401 + ``error="invalid_token"``, not be absorbed
into HTTP 200 + an empty tool list by the tools/list handler.
Delegate-mode clients send only ``Authorization`` (no ``x-litellm-api-key``),
so ``_get_forwarded_auth_from_scope`` returns None and, before the fix, the
preflight returned early without probing.
"""
async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract
"""A bare bearer is an admission credential and must never reach upstream."""
from litellm.proxy._experimental.mcp_server.server import (
_check_passthrough_upstream_auth,
)
@ -6184,48 +6264,6 @@ async def test_delegate_bad_token_gets_connect_time_401():
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')),
) as probe,
):
with pytest.raises(HTTPException) as exc_info:
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test"],
client_ip=None,
)
assert exc_info.value.status_code == 401
challenge = exc_info.value.headers["www-authenticate"]
assert 'error="invalid_token"' in challenge
assert (
'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge
)
probe.assert_awaited_once()
probe_url, probe_auth = probe.call_args.args
assert probe_url == "http://upstream:9401/mcp"
assert probe_auth == "Bearer bogus-token"
@pytest.mark.asyncio
async def test_delegate_valid_token_passes_preflight():
"""An upstream-accepted token must not be blocked by the delegate preflight."""
from litellm.proxy._experimental.mcp_server.server import (
_check_passthrough_upstream_auth,
)
from litellm.proxy._types import UserAPIKeyAuth
server = _delegate_auth_mcp_server()
scope = _delegate_scope([(b"authorization", b"Bearer good-token")])
with (
_patch_delegate_resolver(server, "delegate_test"),
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(200, None)),
) as probe,
):
await _check_passthrough_upstream_auth(
scope=scope,
@ -6234,43 +6272,103 @@ async def test_delegate_valid_token_passes_preflight():
client_ip=None,
)
probe.assert_awaited_once()
probe.assert_not_awaited()
@pytest.mark.asyncio
async def test_delegate_valid_token_forbidden_returns_403():
"""An upstream that accepts the token but forbids the caller (403) must surface
as a bare 403 with no ``WWW-Authenticate`` re-auth hint (a fresh token with the
same scopes would loop), not as an invalid_token challenge."""
async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract
"""A separate upstream bearer never triggers the removed legacy probe."""
from litellm.proxy._experimental.mcp_server.server import (
_check_passthrough_upstream_auth,
)
from litellm.proxy._types import UserAPIKeyAuth
server = _delegate_auth_mcp_server()
scope = _delegate_scope([(b"authorization", b"Bearer scoped-out-token")])
scope = _delegate_scope(
[
(b"x-litellm-api-key", b"sk-litellm-proxy-key"),
(b"authorization", b"Bearer upstream-token"),
]
)
with (
_patch_delegate_resolver(server, "delegate_test"),
patch(
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
patch( # test-quality-ok: the removed probe call is the security regression under test
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(403, None)),
),
new=AsyncMock(),
) as probe,
):
with pytest.raises(HTTPException) as exc_info:
await _check_passthrough_upstream_auth(
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["delegate_test"],
client_ip=None,
)
probe.assert_not_awaited()
@pytest.mark.parametrize(
"probe_status, expected_status",
[(200, None), (401, 401), (403, 403)],
)
@pytest.mark.asyncio
async def test_oauth_passthrough_preflight_preserves_status_contract(probe_status, expected_status):
from litellm.proxy._experimental.mcp_server.server import (
_check_passthrough_upstream_auth,
)
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="passthrough-id",
name="passthrough_server",
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
oauth_passthrough=True,
extra_headers=["Authorization"],
)
scope = _delegate_scope(
[
(b"x-litellm-api-key", b"sk-litellm-proxy-key"),
(b"authorization", b"Bearer upstream-token"),
]
)
with (
patch( # test-quality-ok: isolate authorized-server resolution so this test exercises the preflight contract
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch( # test-quality-ok: the upstream transport boundary is the behavior being mapped to an HTTP response
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(probe_status, None)),
) as probe,
):
if expected_status is None:
result = await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test"],
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["passthrough_server"],
client_ip=None,
)
assert result is None
else:
with pytest.raises(HTTPException) as exc_info:
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["passthrough_server"],
client_ip=None,
)
assert exc_info.value.status_code == expected_status
if expected_status == 401:
assert "passthrough_server" in exc_info.value.headers["www-authenticate"]
assert exc_info.value.status_code == 403
assert not (exc_info.value.headers or {})
probe.assert_awaited_once_with("https://upstream.example.com/mcp", "Bearer upstream-token")
@pytest.mark.asyncio
@ -6428,123 +6526,6 @@ async def test_delegate_not_probed_when_named_only_via_server_id():
probe.assert_not_awaited()
@pytest.mark.asyncio
async def test_delegate_preflight_with_unpatched_probe():
"""Integration across the preflight and the unpatched ``_probe_upstream_auth``,
mocked only at the httpx-client boundary (tests/test_litellm is mocked-only; the
real-network proof lives in the PR's live-proxy evidence). The mock honors the
``AsyncHTTPHandler.post`` contract by raising ``httpx.HTTPStatusError`` on the
upstream 401, so the production ``except httpx.HTTPStatusError`` branch is the one
exercised. A rejected token surfaces as the connect-time 401 challenge; an
accepted token passes untouched, and the caller's bearer reaches the delegate URL."""
import httpx
from litellm.proxy._experimental.mcp_server.server import (
_check_passthrough_upstream_auth,
)
from litellm.proxy._types import UserAPIKeyAuth
accepted = MagicMock()
accepted.status_code = 200
accepted.headers = {}
rejected = MagicMock()
rejected.status_code = 401
rejected.headers = {"www-authenticate": 'Bearer realm="stub-upstream", error="invalid_token"'}
async def respond_by_token(url=None, headers=None, json=None, timeout=None, **kwargs):
if headers.get("Authorization") == "Bearer good-token":
return accepted
raise httpx.HTTPStatusError(
"401 Unauthorized",
request=httpx.Request("POST", url),
response=rejected,
)
mock_client = MagicMock()
mock_client.post = AsyncMock(side_effect=respond_by_token)
server = _delegate_auth_mcp_server()
with (
_patch_delegate_resolver(server, "delegate_test"),
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
"litellm.proxy._experimental.mcp_server.server.get_async_httpx_client",
return_value=mock_client,
),
):
with pytest.raises(HTTPException) as exc_info:
await _check_passthrough_upstream_auth(
scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]),
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test"],
client_ip=None,
)
await _check_passthrough_upstream_auth(
scope=_delegate_scope([(b"authorization", b"Bearer good-token")]),
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test"],
client_ip=None,
)
assert exc_info.value.status_code == 401
challenge = exc_info.value.headers["www-authenticate"]
assert 'error="invalid_token"' in challenge
assert (
'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge
)
probed_urls = [call.kwargs["url"] for call in mock_client.post.await_args_list]
assert probed_urls == ["http://upstream:9401/mcp", "http://upstream:9401/mcp"]
@pytest.mark.asyncio
async def test_delegate_challenge_echoes_requested_alias():
"""An alias-routed delegate request must be probed, and the challenge must echo
the requested alias (not the canonical server name) so the resource_metadata
URL matches what the tokenless preemptive challenge emits for the same route."""
from litellm.proxy._experimental.mcp_server.server import (
_check_passthrough_upstream_auth,
)
from litellm.proxy._types import UserAPIKeyAuth
server = _delegate_auth_mcp_server().model_copy(update={"alias": "dt-alias"})
scope = {
"type": "http",
"method": "POST",
"path": "/mcp/dt-alias",
"scheme": "http",
"server": ("localhost", 4000),
"headers": [(b"authorization", b"Bearer bogus-token")],
}
with (
_patch_delegate_resolver(server, "dt-alias"),
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')),
),
):
with pytest.raises(HTTPException) as exc_info:
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["dt-alias"],
client_ip=None,
)
challenge = exc_info.value.headers["www-authenticate"]
assert 'error="invalid_token"' in challenge
assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/dt-alias"' in challenge
@pytest.mark.asyncio
async def test_delegate_probe_not_fanned_out_to_access_group_members():
"""A single access-group name passes the one-target route gate but must not fan
@ -6578,41 +6559,6 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members():
probe.assert_not_awaited()
def test_is_delegate_upstream_probe_target_fails_closed_on_m2m_shape():
"""An unstamped M2M-shape row (null ``oauth2_flow`` + client credentials)
resolves to ``client_credentials`` and must not be probed with the caller's
bearer; its stored client credentials drive egress instead."""
from litellm.proxy._experimental.mcp_server.server import (
_is_delegate_upstream_probe_target,
)
assert _is_delegate_upstream_probe_target(_delegate_auth_mcp_server()) is True
m2m_shape = MCPServer(
server_id="delegate-m2m",
name="delegate_m2m",
url="http://upstream:9401/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
oauth2_flow=None,
token_url="http://idp:9000/token",
client_id="client",
client_secret="secret",
)
assert _is_delegate_upstream_probe_target(m2m_shape) is False
non_delegate = MCPServer(
server_id="oauth2-plain",
name="oauth2_plain",
url="http://upstream:9401/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
)
assert _is_delegate_upstream_probe_target(non_delegate) is False
@pytest.mark.asyncio
async def test_create_mcp_client_sampling_disabled_by_default():
"""Sampling callback must be None when allow_sampling is not set (default False)."""

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

@ -3390,6 +3390,72 @@ class TestMCPServerManager:
)
assert not extra_headers or "authorization" not in {k.lower() for k in extra_headers}
@pytest.mark.asyncio
async def test_call_regular_mcp_tool_legacy_delegate_never_forwards_admission_key(self):
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="server-legacy-delegate-leak",
name="legacy-delegate",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
extra_headers = await self._capture_call_extra_headers(
server,
oauth2_headers={"Authorization": "Bearer sk-litellm-key"},
raw_headers={"authorization": "Bearer sk-litellm-key"},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
)
assert not extra_headers or "authorization" not in {k.lower() for k in extra_headers}
@pytest.mark.asyncio
async def test_call_regular_mcp_tool_legacy_delegate_forwards_separate_authorization(self):
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="server-legacy-delegate-dual-credential",
name="legacy-delegate",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
extra_headers = await self._capture_call_extra_headers(
server,
oauth2_headers={"Authorization": "Bearer upstream-token"},
raw_headers={
"x-litellm-api-key": "Bearer sk-litellm-key",
"authorization": "Bearer upstream-token",
},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
)
assert extra_headers == {"Authorization": "Bearer upstream-token"}
@pytest.mark.asyncio
async def test_call_regular_mcp_tool_legacy_delegate_strips_repeated_admission_key(self):
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id="server-legacy-delegate-repeated-key",
name="legacy-delegate",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
extra_headers = await self._capture_call_extra_headers(
server,
oauth2_headers={"Authorization": "Bearer sk-litellm-key"},
raw_headers={
"x-litellm-api-key": "Bearer sk-litellm-key",
"authorization": "Bearer sk-litellm-key",
},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
)
assert not extra_headers or "authorization" not in {k.lower() for k in extra_headers}
def test_should_strip_caller_authorization_new_modes(self):
from litellm.proxy._types import UserAPIKeyAuth
@ -6916,51 +6982,6 @@ class TestMCPServerManager:
== expected_server_ids
)
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_anonymous_delegate_requires_oauth2(self):
"""Anonymous delegated auth listing should only include oauth2 servers."""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
manager = MCPServerManager()
oauth_delegate_server = MCPServer(
server_id="oauth-delegate",
name="oauth_delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
api_key_delegate_server = MCPServer(
server_id="api-key-delegate",
name="api_key_delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.api_key,
delegate_auth_to_upstream=True,
)
oauth_non_delegate_server = MCPServer(
server_id="oauth-non-delegate",
name="oauth_non_delegate",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=False,
)
manager.registry = {
oauth_delegate_server.server_id: oauth_delegate_server,
api_key_delegate_server.server_id: api_key_delegate_server,
oauth_non_delegate_server.server_id: oauth_non_delegate_server,
}
with patch.object(
MCPRequestHandler,
"get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[],
):
result = await manager.get_allowed_mcp_servers(None)
assert set(result) == {"oauth-delegate"}
def test_get_mcp_server_from_tool_name_uses_server_name_not_name(self):
"""
Test that _get_mcp_server_from_tool_name uses server.server_name instead of server.name
@ -8088,9 +8109,9 @@ class TestMCPServerTokenExchangeColumns:
assert rebuilt_table.token_exchange_profile == "entra_obo"
class TestInternalDelegatePkceWarningLog:
class TestLegacyDelegateAuthWarningLog:
@pytest.mark.asyncio
async def test_build_mcp_server_logs_on_internal_delegate_interactive(self, caplog):
async def test_build_mcp_server_logs_deprecation_for_internal_delegate(self, caplog):
caplog.set_level(logging.WARNING, logger="LiteLLM")
manager = MCPServerManager()
table_record = LiteLLM_MCPServerTable(
@ -8106,11 +8127,11 @@ class TestInternalDelegatePkceWarningLog:
)
await manager.build_mcp_server_from_table(table_record)
combined = " ".join(r.getMessage() for r in caplog.records)
assert "internal-only" in combined
assert "deprecated auth_type=oauth2" in combined
assert "delegate_auth_to_upstream=true" in combined
@pytest.mark.asyncio
async def test_build_mcp_server_no_internal_delegate_log_when_public(self, caplog):
async def test_build_mcp_server_logs_deprecation_for_public_delegate(self, caplog):
caplog.set_level(logging.WARNING, logger="LiteLLM")
manager = MCPServerManager()
table_record = LiteLLM_MCPServerTable(
@ -8126,12 +8147,13 @@ class TestInternalDelegatePkceWarningLog:
)
await manager.build_mcp_server_from_table(table_record)
combined = " ".join(r.getMessage() for r in caplog.records)
assert "internal-only" not in combined
assert "deprecated auth_type=oauth2" in combined
assert "auth_type=oauth_delegate" in combined
def test_warn_skipped_for_client_credentials(self, caplog):
caplog.set_level(logging.WARNING, logger="LiteLLM")
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_warn_internal_delegate_pkce_if_applicable,
_warn_legacy_delegate_auth_if_applicable,
)
server = MCPServer(
@ -8144,9 +8166,9 @@ class TestInternalDelegatePkceWarningLog:
available_on_public_internet=False,
delegate_auth_to_upstream=True,
)
_warn_internal_delegate_pkce_if_applicable(server, source="test")
_warn_legacy_delegate_auth_if_applicable(server, source="test")
combined = " ".join(r.getMessage() for r in caplog.records)
assert "internal-only" not in combined
assert "deprecated auth_type=oauth2" not in combined
class TestHasClientCredentialsOAuth2Flow:

93
tests/test_litellm/proxy/db/test_spend_log_tool_index.py
View file

@ -1,15 +1,17 @@
"""
Tests for the tool usage writer: ToolUsageTransaction construction (invoked tools
only) and the flush that writes LiteLLM_SpendLogToolIndex plus the
LiteLLM_DailyToolSpend rollup in one transaction.
only) and the flush that writes LiteLLM_SpendLogToolIndex in bounded statements
plus the LiteLLM_DailyToolSpend rollup in one transaction.
"""
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from litellm.proxy.db import spend_log_tool_index
from litellm.proxy.db.spend_log_tool_index import (
ToolUsageTransaction,
build_tool_usage_transaction,
@ -35,11 +37,24 @@ class _FakeBatcher:
return None
def _prisma(batch_: MagicMock) -> MagicMock:
prisma = MagicMock()
prisma.db.batch_ = batch_
prisma.db.litellm_spendlogtoolindex.create_many = AsyncMock()
return prisma
def _prisma_with_batcher() -> tuple[MagicMock, _FakeBatcher]:
batcher = _FakeBatcher()
prisma = MagicMock()
prisma.db.batch_ = MagicMock(return_value=batcher)
return prisma, batcher
return _prisma(MagicMock(return_value=batcher)), batcher
def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]:
return [
(row["request_id"], row["tool_name"])
for call in prisma.db.litellm_spendlogtoolindex.create_many.call_args_list
for row in call.kwargs["data"]
]
class TestBuildToolUsageTransaction:
@ -228,9 +243,8 @@ class TestFlushToolUsageTransactions:
prisma_client=prisma,
transactions=[_transaction("r1", tool_names=("tool_a", "tool_b"), spend=0.10, total_tokens=100)],
)
index_rows = batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["data"]
assert [(r["request_id"], r["tool_name"]) for r in index_rows] == [("r1", "tool_a"), ("r1", "tool_b")]
assert batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True
assert _index_rows_written(prisma) == [("r1", "tool_a"), ("r1", "tool_b")]
assert prisma.db.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True
upserts = {
c.kwargs["where"]["date_tool_name"]["tool_name"]: c.kwargs["data"]
@ -267,19 +281,46 @@ class TestFlushToolUsageTransactions:
assert data["update"]["request_count"] == {"increment": 2}
@pytest.mark.asyncio
async def test_index_rows_and_rollup_share_one_transaction(self):
# Both writes go through the same batch_() so a failed flush cannot leave
# index rows without their rollup increments (or vice versa); increments
# are not idempotent, so partial states must be unreachable.
async def test_index_rows_are_written_in_bounded_statements_outside_the_rollup_transaction(self, monkeypatch):
monkeypatch.setattr(spend_log_tool_index, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
prisma, batcher = _prisma_with_batcher()
await flush_tool_usage_transactions(
prisma_client=prisma,
transactions=[_transaction("r1")],
)
tool_names = tuple(f"tool_{i}" for i in range(50))
transactions = [_transaction(f"r{i}", tool_names=tool_names) for i in range(5)]
await flush_tool_usage_transactions(prisma_client=prisma, transactions=transactions)
statements = prisma.db.litellm_spendlogtoolindex.create_many.call_args_list
assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50]
assert all(call.kwargs["skip_duplicates"] is True for call in statements)
assert _index_rows_written(prisma) == [
(txn.request_id, tool_name) for txn in transactions for tool_name in tool_names
]
batcher.litellm_spendlogtoolindex.create_many.assert_not_called()
prisma.db.batch_.assert_called_once()
assert batcher.litellm_dailytoolspend.upsert.call_count == len(tool_names)
@pytest.mark.asyncio
async def test_index_connection_error_is_retried_before_the_rollup_is_attempted(self, monkeypatch):
prisma, batcher = _prisma_with_batcher()
prisma.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=[httpx.ConnectError("down"), None])
async def fake_sleep(seconds: float) -> None:
return None
monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep)
await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")])
assert prisma.db.litellm_spendlogtoolindex.create_many.await_count == 2
prisma.db.batch_.assert_called_once()
batcher.litellm_spendlogtoolindex.create_many.assert_called_once()
batcher.litellm_dailytoolspend.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_ambiguous_index_error_drops_the_batch_without_touching_the_rollup(self):
prisma, _ = _prisma_with_batcher()
prisma.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=httpx.ReadTimeout("ambiguous"))
with pytest.raises(httpx.ReadTimeout):
await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")])
prisma.db.litellm_spendlogtoolindex.create_many.assert_awaited_once()
prisma.db.batch_.assert_not_called()
@pytest.mark.asyncio
async def test_empty_batch_touches_nothing(self):
prisma, _ = _prisma_with_batcher()
@ -290,11 +331,8 @@ class TestFlushToolUsageTransactions:
async def test_connection_errors_retry_and_succeed(self, monkeypatch):
# A failed batch commits nothing, so retrying a connection error cannot
# double-count; the flush must retry rather than drop the batch.
import httpx
batcher = _FakeBatcher()
prisma = MagicMock()
prisma.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), batcher])
prisma = _prisma(MagicMock(side_effect=[httpx.ConnectError("down"), batcher]))
sleeps: list[float] = []
async def fake_sleep(seconds: float) -> None:
@ -308,10 +346,7 @@ class TestFlushToolUsageTransactions:
@pytest.mark.asyncio
async def test_connection_errors_exhaust_retries_then_raise(self, monkeypatch):
import httpx
prisma = MagicMock()
prisma.db.batch_ = MagicMock(side_effect=httpx.ConnectError("down"))
prisma = _prisma(MagicMock(side_effect=httpx.ConnectError("down")))
async def fake_sleep(seconds: float) -> None:
return None
@ -325,8 +360,7 @@ class TestFlushToolUsageTransactions:
@pytest.mark.asyncio
async def test_non_connection_errors_do_not_retry(self):
prisma = MagicMock()
prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data"))
prisma = _prisma(MagicMock(side_effect=ValueError("bad data")))
with pytest.raises(ValueError, match="bad data"):
await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")])
prisma.db.batch_.assert_called_once()
@ -338,11 +372,8 @@ class TestFlushToolUsageTransactions:
# unknown; the engine can leave the transaction open on the pooled
# connection, so a retry's statements would stack into it and one
# commit would apply both increment sets. These must never retry.
import httpx
error = getattr(httpx, ambiguous_error)("ambiguous")
prisma = MagicMock()
prisma.db.batch_ = MagicMock(side_effect=error)
prisma = _prisma(MagicMock(side_effect=error))
with pytest.raises((httpx.ReadTimeout, httpx.ReadError)):
await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")])
prisma.db.batch_.assert_called_once()

54
tests/test_litellm/proxy/guardrails/test_usage_tracking.py
View file

@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from litellm.proxy.guardrails import usage_tracking
from litellm.proxy.guardrails.usage_tracking import (
_MAX_PENDING_ROWS,
PendingRollups,
@ -511,3 +512,56 @@ async def test_requeued_cost_is_added_to_the_next_flush():
costs = _cost_upserts(recovered)
assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0)
assert costs["someFutureCounter"] == (0.0, 7)
def _fan_out_payload(request_id: str, guardrail_ids: tuple[str, ...]) -> dict[str, Any]:
return {
"request_id": request_id,
"startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc),
"team_id": "team-a",
"api_key": "hashed-key-1",
"metadata": json.dumps(
{"guardrail_information": [{"guardrail_id": gid, "guardrail_status": "success"} for gid in guardrail_ids]}
),
}
def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]:
return [
(row["request_id"], row["guardrail_id"])
for call in prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list
for row in call.kwargs["data"]
]
@pytest.mark.asyncio
async def test_index_rows_are_written_in_row_bounded_statements(monkeypatch):
"""
LIT-5931: the drain caps logs, not logs x guardrails, so a fan-out must be
split into statements the query engine can afford instead of one create_many.
"""
monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
prisma = _prisma()
guardrail_ids = tuple(f"guard-{i}" for i in range(50))
logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)]
await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups())
statements = prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list
assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50]
assert all(call.kwargs["skip_duplicates"] is True for call in statements)
assert _index_rows_written(prisma) == [(f"r{i}", gid) for i in range(5) for gid in guardrail_ids]
@pytest.mark.asyncio
async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollup(monkeypatch):
monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
prisma = _prisma()
prisma.db.litellm_spendlogguardrailindex.create_many.side_effect = [None, httpx.ReadTimeout("ambiguous"), None]
guardrail_ids = tuple(f"guard-{i}" for i in range(50))
logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)]
await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups())
assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3
assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids)

26
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -19,6 +19,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import (
run_spend_event,
)
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
@ -1912,14 +1913,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit():
("allm_passthrough_route", True),
("aretrieve_batch", True),
("acompletion", False),
("call_mcp_tool", False),
("call_mcp_tool", True),
(None, False),
],
)
def test_should_track_cost_callback_pass_through_without_owner(call_type, expected):
"""Regression for LIT-3782: unauthenticated pass-through requests (auth=false)
carry no key/user/team/end-user, yet must still be tracked so they land in
LiteLLM_SpendLogs. Other call types with no owner stay untracked.
LiteLLM_SpendLogs. Explicit MCP passthrough calls require the same handling.
Other call types with no owner stay untracked.
aretrieve_batch is included for the same reason: CheckBatchCost's synthetic
logging_obj for a completed managed batch only ever carries
@ -1939,10 +1941,26 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect
)
def test_should_track_cost_callback_respects_disabled_spend_updates(monkeypatch):
monkeypatch.setattr(ProxyUpdateSpend, "disable_spend_updates", staticmethod(lambda: True))
assert (
_should_track_cost_callback(
user_api_key="key",
user_id="user",
team_id="team",
end_user_id="end-user",
call_type="call_mcp_tool",
)
is False
)
@pytest.mark.parametrize(
"call_type, expect_spend_log",
[
("pass_through_endpoint", True),
("call_mcp_tool", True),
("aretrieve_batch", True),
("acompletion", False),
(None, False),
@ -1953,8 +1971,8 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(cal
"""Regression for LIT-3782: a pass-through request with auth=false reaches the
cost callback with no key/user/team/end-user. Before the fix the spend-log
write was skipped and the request never appeared in request/usage logs. It
must now be written for pass-through call types while other unauthenticated
calls remain skipped.
must now be written for pass-through and MCP tool call types while other
unauthenticated calls remain skipped.
aretrieve_batch is included because CheckBatchCost's completed-batch cost
event reaches this same callback with no attributable key/user/team when

174
tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py
View file

@ -2527,6 +2527,180 @@ class TestToolTransformation:
"type": "object",
}
@pytest.mark.parametrize(
"model, custom_llm_provider",
[
("bedrock/converse/global.anthropic.claude-sonnet-5", "bedrock_converse"),
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
("claude-sonnet-5", "vertex_ai"),
("gemini-3.1-pro-preview", "vertex_ai"),
("moonshotai.kimi-k2-thinking", "bedrock_mantle"),
],
)
def test_reasoning_summary_still_yields_a_string_reasoning_effort(self, model, custom_llm_provider):
"""
A Responses request carrying ``reasoning.summary`` must still reach a chat provider as a
plain ``reasoning_effort`` string. ``summary`` is Responses-only, and forwarding the whole
object turns reasoning off: Bedrock Converse and Vertex silently discard a non-string
``reasoning_effort``, and Bedrock Mantle rejects the request outright.
"""
responses_api_request = {"reasoning": {"effort": "medium", "summary": "auto"}}
result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=model,
input="hi",
responses_api_request=responses_api_request,
custom_llm_provider=custom_llm_provider,
)
assert result["reasoning_effort"] == "medium"
@pytest.mark.parametrize(
"model, custom_llm_provider",
[
("gpt-5.4-pro", "azure_ai"),
("gpt-5", "openai"),
("gpt-5.1", "openai"),
("gpt-5", "azure"),
],
)
def test_bridged_model_carries_the_summary_as_an_alias(self, model, custom_llm_provider):
"""
``summary`` reaches a bridged model through the ``reasoning_summary`` alias, never smuggled
inside ``reasoning_effort``. ``litellm.completion`` reads that alias back with
``peek_reasoning_summary_aliases`` and reassembles ``{effort, summary}``, so the far end
gets the same object it always did while no chat provider ever sees a non-string effort.
"""
result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=model,
input="hi",
responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}},
custom_llm_provider=custom_llm_provider,
)
assert result["reasoning_effort"] == "medium"
assert result["reasoning_summary"] == "auto"
@pytest.mark.parametrize(
"model, custom_llm_provider",
[
("gpt-5", "openai"),
("gpt-5.1", "openai"),
("gpt-5", "azure"),
],
)
def test_gpt_5_summary_survives_the_bridge_it_claims_to_take(self, model, custom_llm_provider):
"""
Regression for the probe disagreeing with the real decision. The transform asked
``responses_api_bridge_check`` with ``reasoning_summary`` taken straight off the Responses
object, but ``litellm.completion`` reads it from ``optional_params`` via
``peek_reasoning_summary_aliases``, which the bridged request never populated. So these
models answered "bridging" to the probe and "not bridging" for real, and the object landed
on Chat Completions, which only takes a string. Emitting the alias makes the two agree.
"""
from litellm.main import responses_api_bridge_check
from litellm.utils import get_optional_params, peek_reasoning_summary_aliases
result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=model,
input="hi",
responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}},
custom_llm_provider=custom_llm_provider,
)
optional_params = get_optional_params(
model=model,
custom_llm_provider=custom_llm_provider,
reasoning_effort=result["reasoning_effort"],
reasoning_summary=result["reasoning_summary"],
)
model_info, _ = responses_api_bridge_check(
model=model,
custom_llm_provider=custom_llm_provider,
reasoning_effort=result["reasoning_effort"],
reasoning_summary=peek_reasoning_summary_aliases(optional_params),
)
assert model_info.get("mode") == "responses"
def test_a_failing_bridge_probe_falls_back_to_the_string_effort(self, monkeypatch):
"""
The probe is a capability question, so a model-info lookup blowing up must not fail the
request. It degrades to the chat-safe form: a string effort and no alias.
"""
import litellm.main
def _boom(**_kwargs):
raise RuntimeError("model info unavailable")
monkeypatch.setattr(litellm.main, "responses_api_bridge_check", _boom)
result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model="gpt-5.4-pro",
input="hi",
responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}},
custom_llm_provider="azure_ai",
)
assert result["reasoning_effort"] == "medium"
assert "reasoning_summary" not in result
@pytest.mark.parametrize(
"reasoning, expected",
[
({"effort": "high"}, "high"),
("low", "low"),
({"summary": "auto"}, None),
({}, None),
(None, None),
],
)
def test_reasoning_param_shapes_map_to_reasoning_effort(self, reasoning, expected):
"""
An object without ``effort`` carries nothing Chat Completions can use, so no
``reasoning_effort`` is sent at all (the bridge drops None-valued params).
"""
result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model="anthropic.claude-sonnet-4-5-20250929-v1:0",
input="hi",
responses_api_request={"reasoning": reasoning},
custom_llm_provider="bedrock",
)
assert result.get("reasoning_effort") == expected
assert ("reasoning_effort" in result) is (expected is not None)
@pytest.mark.parametrize(
"model, expected_thinking",
[
("global.anthropic.claude-sonnet-5", {"type": "adaptive"}),
(
"anthropic.claude-sonnet-4-5-20250929-v1:0",
{"type": "enabled", "budget_tokens": 2048},
),
],
)
def test_reasoning_summary_still_enables_thinking_on_bedrock(self, model, expected_thinking):
"""
End to end through Bedrock Converse's own param mapping: the effort a Responses request asks
for must survive into ``thinking``, whether the model takes an adaptive effort or a legacy
token budget. Forwarding the object instead leaves ``thinking`` unset and the model never
reasons, which is the failure this guards.
"""
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
bridged = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=model,
input="hi",
responses_api_request={"reasoning": {"effort": "medium", "summary": "auto"}},
custom_llm_provider="bedrock",
)
mapped = AmazonConverseConfig().map_openai_params(
{"reasoning_effort": bridged["reasoning_effort"]}, {}, model, True
)
assert expected_thinking.items() <= mapped["thinking"].items()
def test_bedrock_anthropic_responses_tools_yield_only_function_toolspec(self):
"""
End-to-end (no network) of the LIT-3858 acceptance criterion: the mixed tools array

69
tests/test_litellm/test_utils.py
View file

@ -4449,6 +4449,75 @@ def test_fireworks_models_in_backup_cost_map():
), f"short-form {short_key} does not match long-form {long_key}"
@pytest.fixture
def fireworks_short_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
monkeypatch.setattr(
litellm,
"model_cost",
{
"fireworks_ai/accounts/fireworks/models/glm-5p3": {
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
"litellm_provider": "fireworks_ai",
"mode": "chat",
"max_tokens": 100,
},
"fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": {
"input_cost_per_token": 2.1e-6,
"output_cost_per_token": 6.6e-6,
"litellm_provider": "fireworks_ai",
"mode": "chat",
},
"fireworks_ai/nomic-ai/nomic-embed-text-v1.5": {
"input_cost_per_token": 8e-9,
"output_cost_per_token": 0.0,
"litellm_provider": "fireworks_ai",
"mode": "embedding",
},
},
)
litellm.get_model_info.cache_clear()
yield
litellm.get_model_info.cache_clear()
def test_fireworks_short_model_names_resolve_to_long_cost_map_keys(fireworks_short_model_cost_map: None) -> None:
model_info = litellm.get_model_info("fireworks_ai/glm-5p3")
assert model_info["key"] == "fireworks_ai/accounts/fireworks/models/glm-5p3"
assert model_info["input_cost_per_token"] == 1e-6
assert model_info["max_tokens"] == 100
model_info = litellm.get_model_info("glm-5p3", custom_llm_provider="fireworks_ai")
assert model_info["key"] == "fireworks_ai/accounts/fireworks/models/glm-5p3"
model_info = litellm.get_model_info("fireworks_ai/glm-5p3-fast")
assert model_info["key"] == "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast"
assert model_info["input_cost_per_token"] == 2.1e-6
model_info = litellm.get_model_info("fireworks_ai/nomic-ai/nomic-embed-text-v1.5")
assert model_info["key"] == "fireworks_ai/nomic-ai/nomic-embed-text-v1.5"
with pytest.raises(Exception, match="isn't mapped"):
litellm.get_model_info("fireworks_ai/does-not-exist")
def test_fireworks_short_model_names_price_with_completion_cost(fireworks_short_model_cost_map: None) -> None:
from litellm.types.utils import ModelResponse
response = ModelResponse(
model="fireworks_ai/glm-5p3",
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
cost = litellm.completion_cost(
completion_response=response,
model="fireworks_ai/glm-5p3",
custom_llm_provider="fireworks_ai",
)
assert cost == pytest.approx(10 * 1e-6 + 5 * 2e-6)
class TestBedrockBaseModelLabelKeepsTools:
"""Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly
label must not silently drop ``tools``/``tool_choice`` under ``drop_params``."""