scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(mcp): require admission for delegated OAuth

Browse source
This commit is contained in:
Joshua Valluru 2026-09-12 16:25:57 -07:00
parent 566f026c1c
commit 94e4dd725f
12 changed files with 433 additions and 702 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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,

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

@ -1453,22 +1453,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 +2637,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 +3182,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 +3476,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 +3485,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

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

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

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

226
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
@ -4453,7 +4341,9 @@ class TestAgentMCPPermissions:
stack.enter_context(patcher)
stack.enter_context(
patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"])
MCPRequestHandler,
"_get_allowed_mcp_servers_for_key",
AsyncMock(return_value=["server-a", "server-b"]),
)
)
stack.enter_context(
@ -4479,7 +4369,9 @@ class TestAgentMCPPermissions:
await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth)
stack.enter_context(
patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"])
MCPRequestHandler,
"_get_allowed_mcp_servers_for_key",
AsyncMock(return_value=["server-a", "server-b"]),
)
)
stack.enter_context(
@ -4503,9 +4395,15 @@ class TestAgentMCPPermissions:
with contextlib.ExitStack() as stack:
for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager):
stack.enter_context(patcher)
server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth)
server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth)
server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth)
server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server(
"server-a", user_api_key_auth
)
server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server(
"server-b", user_api_key_auth
)
server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server(
"server-c", user_api_key_auth
)
assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"]
assert server_b_tools == ["tool_b"]

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

@ -6157,15 +6157,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 +6177,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 +6185,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 +6439,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 +6472,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)."""

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

@ -4469,7 +4469,9 @@ class TestMCPServerManager:
@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2])
@pytest.mark.parametrize("is_byok", [False, True])
@pytest.mark.parametrize("scheme", ["http", "https"])
async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme):
async def test_openapi_health_loads_spec_without_mcp_handshake(
self, respx_mock, monkeypatch, auth_type, is_byok, scheme
):
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
manager = MCPServerManager()
server = MCPServer(
@ -4519,14 +4521,28 @@ class TestMCPServerManager:
@pytest.mark.parametrize(
("failure", "expected_status", "expected_error"),
[
(httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"),
(
httpx.Response(401, text="secret response content"),
"unhealthy",
"OpenAPI specification request failed (HTTP 401)",
),
(httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"),
(httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"),
(httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"),
(httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"),
(
httpx.ConnectError("secret network details"),
"unhealthy",
"OpenAPI specification could not be loaded (ConnectError)",
),
(
httpx.Response(200, text="secret invalid JSON body"),
"unhealthy",
"OpenAPI specification could not be loaded (JSONDecodeError)",
),
],
)
async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error):
async def test_openapi_health_reports_safe_failures(
self, respx_mock, monkeypatch, failure, expected_status, expected_error
):
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
manager = MCPServerManager()
server = MCPServer(
@ -6916,51 +6932,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 +8059,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 +8077,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 +8097,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 +8116,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:
@ -12697,7 +12669,12 @@ async def test_debug_resolution_matches_final_header_conflict_winner(
from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser
from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics
from litellm.proxy._experimental.mcp_server.outbound_credentials import (
ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider,
ApiKeyConfig,
AuthorizationCodeConfig,
NoneConfig,
ServerSpec,
SharedKey,
UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -12713,10 +12690,15 @@ async def test_debug_resolution_matches_final_header_conflict_winner(
store = Store()
context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice"))
diagnostics = MCPAuthDiagnostics()
token = request_ctx.set(RequestContext(
request_id=1, meta=None, session=MagicMock(), lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
))
token = request_ctx.set(
RequestContext(
request_id=1,
meta=None,
session=MagicMock(),
lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
)
)
selected = {
"stored": AuthorizationCodeConfig(),
"static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))),
@ -12725,7 +12707,10 @@ async def test_debug_resolution_matches_final_header_conflict_winner(
try:
auth, remaining = await MCPServerManager()._resolve_v2_auth(
server=MCPServer(
server_id="s", name="s", transport="http", url="https://up.example/mcp",
server_id="s",
name="s",
transport="http",
url="https://up.example/mcp",
static_headers={"Authorization": "Bearer configured"},
),
spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected),
@ -12755,17 +12740,28 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li
from litellm.types.mcp_server.mcp_server_manager import MCPServer
diagnostics = MCPAuthDiagnostics()
token = request_ctx.set(RequestContext(
request_id=1, meta=None, session=MagicMock(), lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
))
token = request_ctx.set(
RequestContext(
request_id=1,
meta=None,
session=MagicMock(),
lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
)
)
try:
server = MCPServer(
server_id="signed", name="signed", transport=transport,
url="https://up.example/mcp", auth_type="aws_sigv4",
aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret",
aws_region_name="us-east-1", aws_service_name="execute-api",
command="python", args=["-c", "pass"],
server_id="signed",
name="signed",
transport=transport,
url="https://up.example/mcp",
auth_type="aws_sigv4",
aws_access_key_id="AKIDEXAMPLE",
aws_secret_access_key="test-signing-secret",
aws_region_name="us-east-1",
aws_service_name="execute-api",
command="python",
args=["-c", "pass"],
)
client = await MCPServerManager()._create_mcp_client(server)
if transport == "stdio":
@ -12784,12 +12780,16 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li
async def test_temporary_server_discovery_reuses_resolved_metadata_without_publishing() -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(
server_id="temporary-oauth-discovery", name="temporary", url="https://idp.example.com/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough,
server_id="temporary-oauth-discovery",
name="temporary",
url="https://idp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
)
manager._set_oauth_discovery_deferred(server.server_id, True)
metadata: Final = MCPOAuthMetadata(
authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token",
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
registration_url="https://idp.example.com/register",
)
with patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery:
@ -12809,13 +12809,18 @@ async def test_temporary_server_discovery_reuses_resolved_metadata_without_publi
async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(
server_id="repeated-stale", name="stale", url="https://idp.example.com/mcp",
transport=MCPTransport.http, auth_type=auth_type, oauth2_flow="authorization_code",
server_id="repeated-stale",
name="stale",
url="https://idp.example.com/mcp",
transport=MCPTransport.http,
auth_type=auth_type,
oauth2_flow="authorization_code",
)
manager.registry[server.server_id] = server
manager._set_oauth_discovery_deferred(server.server_id, True)
metadata: Final = MCPOAuthMetadata(
authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token",
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
)
with (
patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery,
@ -12835,13 +12840,20 @@ async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) ->
async def test_stale_discovery_falls_back_to_resolved_registered_server() -> None:
manager: Final = MCPServerManager()
original: Final = MCPServer(
server_id="resolved-replacement", name="replacement", url="https://old.example.com/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code",
server_id="resolved-replacement",
name="replacement",
url="https://old.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
)
replacement: Final = original.model_copy(
update={
"url": "https://new.example.com/mcp",
"authorization_url": "https://new.example.com/authorize",
"token_url": "https://new.example.com/token",
}
)
replacement: Final = original.model_copy(update={
"url": "https://new.example.com/mcp", "authorization_url": "https://new.example.com/authorize",
"token_url": "https://new.example.com/token",
})
manager.registry[original.server_id] = replacement
assert await manager._rejoin_oauth_metadata_discovery(original, retry_stale=False) is replacement
@ -12849,8 +12861,11 @@ async def test_stale_discovery_falls_back_to_resolved_registered_server() -> Non
def test_stale_discovery_cannot_overwrite_new_registered_server() -> None:
manager: Final = MCPServerManager()
original: Final = MCPServer(
server_id="stale-publication", name="publication", url="https://old.example.com/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.oauth2,
server_id="stale-publication",
name="publication",
url="https://old.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
)
manager._set_oauth_discovery_deferred(original.server_id, True)
original_slot: Final = manager._oauth_discovery_slot(original.server_id)
@ -12866,9 +12881,13 @@ def test_stale_discovery_cannot_overwrite_new_registered_server() -> None:
async def test_temporary_oauth_discovery_expires_without_more_requests() -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(
server_id="expiring-session", name="temporary", url="https://idp.example.com/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough,
authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token",
server_id="expiring-session",
name="temporary",
url="https://idp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
)
manager._set_oauth_discovery_deferred(server.server_id, True)
resolved: Final = await manager.ensure_oauth_metadata_discovered(server)
@ -12969,7 +12988,9 @@ async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(r
result = await manager.health_check_server(server.server_id)
cached = await manager.health_check_server(server.server_id)
assert result.status == "unknown"
assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit"
assert (
result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit"
)
assert cached.health_check_error == result.health_check_error
assert cached.last_health_check == result.last_health_check
assert route.call_count == 1
@ -12981,8 +13002,11 @@ async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, mon
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
manager = MCPServerManager()
server = MCPServer(
server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http,
spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none,
server_id="cancelled-cache",
name="cancelled-cache",
transport=MCPTransport.http,
spec_path="https://93.184.216.34/cancelled-cache.json",
auth_type=MCPAuth.none,
)
manager.registry = {server.server_id: server}
started = asyncio.Event()
@ -13040,11 +13064,18 @@ class _DiscoveryUpstream:
return httpx.Response(202)
self.requests = (*self.requests, (payload.method, request.headers.get("authorization", "")))
if payload.method == "initialize":
return httpx.Response(200, json={
"jsonrpc": "2.0", "id": payload.id,
"result": {"protocolVersion": "2025-03-26", "serverInfo": {"name": "discovery", "version": "1"},
"capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}}},
})
return httpx.Response(
200,
json={
"jsonrpc": "2.0",
"id": payload.id,
"result": {
"protocolVersion": "2025-03-26",
"serverInfo": {"name": "discovery", "version": "1"},
"capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}},
},
},
)
self.entered.set()
await self.release.wait()
if self.outcome == "failure":
@ -13052,12 +13083,15 @@ class _DiscoveryUpstream:
if self.outcome == "cancelled":
raise asyncio.CancelledError()
if self.outcome == "rejected":
return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id,
"error": {"code": -32601, "message": "Unsupported"}})
return httpx.Response(
200, json={"jsonrpc": "2.0", "id": payload.id, "error": {"code": -32601, "message": "Unsupported"}}
)
result: Final = {
"prompts/list": {"prompts": [{"name": "example", "description": "original"}]},
"resources/list": {"resources": [{"name": "example", "uri": "test://example", "description": "original"}]},
"resources/templates/list": {"resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}]},
"resources/templates/list": {
"resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}]
},
"tools/list": {"tools": []},
}[payload.method]
return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result})
@ -13068,7 +13102,9 @@ class _DiscoveryUpstream:
def _discovery_server() -> MCPServer:
return MCPServer(server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http)
return MCPServer(
server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http
)
@pytest.mark.asyncio
@ -13079,8 +13115,11 @@ async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None
clock: Final = _DiscoveryClock()
manager: Final = MCPServerManager(discovery_clock=clock)
upstream: Final = _DiscoveryUpstream()
operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server,
"templates": manager.get_resource_templates_from_server}[kind]
operation: Final = {
"prompts": manager.get_prompts_from_server,
"resources": manager.get_resources_from_server,
"templates": manager.get_resource_templates_from_server,
}[kind]
server: Final = _discovery_server()
with respx.mock(base_url="https://discovery.example") as router:
router.route().mock(side_effect=upstream.respond)
@ -13109,8 +13148,11 @@ async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: st
manager: Final = MCPServerManager()
upstream: Final = _DiscoveryUpstream()
upstream.outcome = outcome
operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server,
"templates": manager.get_resource_templates_from_server}[kind]
operation: Final = {
"prompts": manager.get_prompts_from_server,
"resources": manager.get_resources_from_server,
"templates": manager.get_resource_templates_from_server,
}[kind]
with respx.mock(base_url="https://discovery.example") as router:
router.route().mock(side_effect=upstream.respond)
assert await operation(_discovery_server(), None) == []
@ -13137,9 +13179,20 @@ async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_
assert len(await manager.get_prompts_from_server(server, user)) == 1
assert upstream.initializes == 1
for credential in ("first-secret", "second-secret", "first-secret"):
assert len(await manager.get_prompts_from_server(server, first_user, extra_headers={"Authorization": credential})) == 1
assert (
len(
await manager.get_prompts_from_server(
server, first_user, extra_headers={"Authorization": credential}
)
)
== 1
)
assert upstream.initializes == 3
assert {auth for method, auth in upstream.requests if method == "prompts/list"} == {"", "first-secret", "second-secret"}
assert {auth for method, auth in upstream.requests if method == "prompts/list"} == {
"",
"first-secret",
"second-secret",
}
@pytest.mark.asyncio
@ -13151,7 +13204,9 @@ async def test_discovery_cache_coalesces_and_survives_waiter_cancellation() -> N
upstream.release.clear()
with respx.mock(base_url="https://discovery.example") as router:
router.route().mock(side_effect=upstream.respond)
tasks: Final = tuple(asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10))
tasks: Final = tuple(
asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10)
)
await asyncio.wait_for(upstream.entered.wait(), timeout=5)
tasks[0].cancel()
with pytest.raises(asyncio.CancelledError):
@ -13200,7 +13255,9 @@ async def test_discovery_cache_can_be_disabled(monkeypatch: pytest.MonkeyPatch)
assert upstream.initializes == 2
@pytest.mark.parametrize("value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5)))
@pytest.mark.parametrize(
"value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5))
)
def test_discovery_cache_ttl_validation(value: str, expected: float, monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _mcp_discovery_cache_ttl
@ -13318,9 +13375,15 @@ async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> N
source: Final = CredentialSource()
managers: Final = (MCPServerManager(cred_provider=source), MCPServerManager(cred_provider=source))
server: Final = MCPServer(
server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client",
authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token",
server_id="discovery",
name="discovery",
url="https://discovery.example/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
client_id="discovery-client",
authorization_url="https://discovery.example/authorize",
token_url="https://discovery.example/token",
)
user: Final = UserAPIKeyAuth(user_id="same-user", api_key="same-key")
upstream: Final = _DiscoveryUpstream()
@ -13339,11 +13402,15 @@ async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> N
with respx.mock(base_url="https://discovery.example") as router:
router.route().mock(side_effect=respond)
for manager in managers:
assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-a"]
assert [item.name for item in await manager.get_prompts_from_server(server, user)] == [
"discovery-account-a"
]
assert upstream.initializes == 2
source.token = "token-b"
for manager in managers:
assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-b"]
assert [item.name for item in await manager.get_prompts_from_server(server, user)] == [
"discovery-account-b"
]
assert upstream.initializes == 4
source.token = None
for manager in managers:
@ -13370,9 +13437,15 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None
store: Final = TokenStore()
manager: Final = MCPServerManager(per_user_oauth_token_store=store)
server: Final = MCPServer(
server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client",
authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token",
server_id="discovery",
name="discovery",
url="https://discovery.example/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
client_id="discovery-client",
authorization_url="https://discovery.example/authorize",
token_url="https://discovery.example/token",
)
user: Final = UserAPIKeyAuth(user_id="requesting-user")
upstream: Final = _DiscoveryUpstream()

38
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
@ -1885,9 +1886,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit():
}
with (
patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam
patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: same function-body import, no injection seam
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam
patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam
"litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
) as mock_increment,
patch( # test-quality-ok: same function-body import, no injection seam
"litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock
),
patch( # test-quality-ok: same function-body import, no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj"
) as mock_proxy_logging,
):
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
@ -1912,14 +1919,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 +1947,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 +1977,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