scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge bb37ed2314 into 6b8788eb27

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-11-04 22:50:47 +00:00 • committed by GitHub
commit c67dbe3f33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 540 additions and 117 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/mcp_server/mcp_server_manager.py
View file

@ -1906,6 +1906,17 @@ def _listing_task_was_cancelled(error: asyncio.CancelledError) -> bool:
return error.args != (SESSION_CLEANUP_TIMEOUT_MESSAGE,)
def is_anonymous_mcp_caller(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
return not (
user_api_key_auth
and (
getattr(user_api_key_auth, "user_id", None)
or getattr(user_api_key_auth, "user_role", None)
or getattr(user_api_key_auth, "api_key", None)
)
)
class MCPServerManager:
_STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
@ -3775,14 +3786,7 @@ class MCPServerManager:
)
)
is_anonymous: Final = not (
user_api_key_auth
and (
getattr(user_api_key_auth, "user_id", None)
or getattr(user_api_key_auth, "user_role", None)
or getattr(user_api_key_auth, "api_key", None)
)
)
is_anonymous: Final = is_anonymous_mcp_caller(user_api_key_auth)
if is_anonymous:
anonymous_server_ids: Final = [
server.server_id

100
litellm/proxy/mcp_server/server.py
View file

@ -593,6 +593,7 @@ if MCP_AVAILABLE:
from litellm.proxy.mcp_server.mcp_server_manager import (
MCPServerManager,
global_mcp_server_manager,
is_anonymous_mcp_caller,
is_legacy_upstream_delegate,
)
@ -1885,6 +1886,24 @@ if MCP_AVAILABLE:
)
return acting.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id})
def _narrow_allowed_server_ids(
allowed_servers: Sequence[MCPServer],
toolset_allowed_server_ids: Set[str] | None,
) -> frozenset[str]:
allowed_ids: Final = frozenset(srv.server_id for srv in allowed_servers)
if toolset_allowed_server_ids is None:
return allowed_ids
return allowed_ids & frozenset(toolset_allowed_server_ids)
def _challenge_server_ids(
user_api_key_auth: UserAPIKeyAuth | None,
allowed_servers: Sequence[MCPServer],
toolset_allowed_server_ids: Set[str] | None,
) -> frozenset[str] | None:
if is_anonymous_mcp_caller(user_api_key_auth):
return None if toolset_allowed_server_ids is None else frozenset(toolset_allowed_server_ids)
return _narrow_allowed_server_ids(allowed_servers, toolset_allowed_server_ids)
_apply_toolset_scope: Final = apply_toolset_scope
async def _toolset_server_ids(toolset_id: str) -> set[str]:
@ -1902,7 +1921,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None,
allowed_server_ids: set[str] | None = None,
allowed_server_ids: Set[str] | None = None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Fail fast with HTTP 401 for MCP servers that need user auth but
@ -1910,11 +1929,11 @@ if MCP_AVAILABLE:
(points clients at the gateway AS metadata) and pass-through OAuth
(points clients at the upstream resource-metadata via our well-known).
``allowed_server_ids`` may be passed by callers that have already
narrowed the authorized server set (e.g. toolset scoping); servers
not in that set are skipped so a client targeting a toolset that
excludes a passthrough server is not pushed into an OAuth flow for
a server it will be 403'd on immediately after authentication.
``allowed_server_ids`` is the set of servers the caller's key may
reach, narrowed further by any active toolset scope; anonymous
cold-start callers are not narrowed by a key allowlist. Servers not in
that set are skipped so a client is not pushed into an OAuth flow
for a server it will be 403'd on immediately after authentication.
"""
for server_name in mcp_servers or []:
server = operations.global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip)
@ -2212,19 +2231,18 @@ if MCP_AVAILABLE:
async def _check_passthrough_upstream_auth(
scope: Scope,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_servers: list[str] | None,
client_ip: str | None,
allowed_servers: Sequence[MCPServer],
*,
oauth2_headers: Mapping[str, str] | None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
raw_headers: Mapping[str, str] | None,
requested_servers: Sequence[str] | None = None,
client_ip: str | None = None,
) -> None:
"""Probe pass-through and legacy delegate 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
trigger an upstream probe against a server their key is not permitted for.
The caller passes the already-resolved allowed-server set so a user cannot
trigger a probe against a server their key is not permitted for.
The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401
can only be returned before that point. This function raises HTTPException(401)
@ -2236,19 +2254,11 @@ if MCP_AVAILABLE:
if not oauth2_headers and not mcp_server_auth_headers:
return
# Use the authorized server set, not the raw user-supplied names, so that
# a caller cannot force a probe to a server their key is not allowed to use.
allowed_servers: Final = await operations._get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_servers=mcp_servers,
client_ip=client_ip,
)
# A legacy delegate is probed only when the single requested name resolves to it the way admission
# resolves it, so neither a multi-server connect nor an access group can be failed by one rejected
# token; its challenge echoes the requested name, matching the tokenless challenge.
requested_scope: Final = requested_servers if requested_servers is not None else mcp_servers
single_requested_name: Final = (
requested_scope[0] if requested_scope is not None and len(requested_scope) == 1 else None
requested_servers[0] if requested_servers is not None and len(requested_servers) == 1 else None
)
requested_server: Final = (
operations.global_mcp_server_manager.get_mcp_server_by_name(single_requested_name, client_ip=client_ip)
@ -2378,14 +2388,14 @@ if MCP_AVAILABLE:
await _check_passthrough_upstream_auth(
scope,
context.user_api_key_auth,
authorized_names,
context.client_ip,
allowed,
oauth2_headers=context.oauth2_headers,
mcp_server_auth_headers={key: dict(value) for key, value in context.mcp_server_auth_headers.items()}
if context.mcp_server_auth_headers is not None
else None,
raw_headers=context.raw_headers,
requested_servers=context.mcp_servers,
client_ip=context.client_ip,
)
return None
@ -2398,10 +2408,15 @@ if MCP_AVAILABLE:
mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None,
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None,
allowed_server_ids: Set[str] | None,
toolset_allowed_server_ids: Set[str] | None,
raw_headers: Mapping[str, str] | None,
) -> None:
if not defer_upstream_probes and not targets_disabled_capability:
allowed_servers: Final = await operations._get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_servers=list(mcp_servers) if mcp_servers is not None else None,
client_ip=client_ip,
)
await _raise_preemptive_401_for_unauthenticated_servers(
scope=scope,
mcp_servers=list(mcp_servers) if mcp_servers is not None else None,
@ -2413,7 +2428,9 @@ if MCP_AVAILABLE:
),
user_api_key_auth=user_api_key_auth,
client_ip=client_ip,
allowed_server_ids=set(allowed_server_ids) if allowed_server_ids is not None else None,
allowed_server_ids=_challenge_server_ids(
user_api_key_auth, allowed_servers, toolset_allowed_server_ids
),
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers. Must run after
@ -2422,8 +2439,7 @@ if MCP_AVAILABLE:
await _check_passthrough_upstream_auth(
scope,
user_api_key_auth,
list(mcp_servers) if mcp_servers is not None else None,
client_ip,
allowed_servers,
oauth2_headers=oauth2_headers,
mcp_server_auth_headers=(
{name: dict(headers) for name, headers in mcp_server_auth_headers.items()}
@ -2431,6 +2447,8 @@ if MCP_AVAILABLE:
else None
),
raw_headers=raw_headers,
requested_servers=mcp_servers,
client_ip=client_ip,
)
async def _targets_disabled_capability(peek: _LazyPeekedBody | None) -> bool:
@ -2543,9 +2561,10 @@ if MCP_AVAILABLE:
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
# Must run after toolset scoping so the challenge set is derived
# from the fully-authorized server set: a passthrough server that
# the active toolset excludes should not trigger an OAuth flow
# for a server the caller will be 403'd on after authentication.
# from the fully-authorized server set (the key's allowed servers,
# narrowed by any active toolset scope): a server the key cannot
# reach should not trigger an OAuth flow for a server the caller
# will be 403'd on after authentication.
await _preflight_streamable_http_upstream_auth(
defer_upstream_probes=defer_upstream_probes,
targets_disabled_capability=targets_disabled_capability,
@ -2555,7 +2574,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers=mcp_server_auth_headers,
user_api_key_auth=user_api_key_auth,
client_ip=_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
toolset_allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
@ -2904,9 +2923,15 @@ if MCP_AVAILABLE:
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
# Must run after toolset scoping so the challenge set is derived
# from the fully-authorized server set: a passthrough server that
# the active toolset excludes should not trigger an OAuth flow
# for a server the caller will be 403'd on after authentication.
# from the fully-authorized server set (the key's allowed servers,
# narrowed by any active toolset scope): a server the key cannot
# reach should not trigger an OAuth flow for a server the caller
# will be 403'd on after authentication.
allowed_servers: Final = await operations._get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_servers=mcp_servers,
client_ip=_sse_client_ip,
)
await _raise_preemptive_401_for_unauthenticated_servers(
scope=scope,
mcp_servers=mcp_servers,
@ -2914,7 +2939,9 @@ if MCP_AVAILABLE:
mcp_server_auth_headers=mcp_server_auth_headers,
user_api_key_auth=user_api_key_auth,
client_ip=_sse_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
allowed_server_ids=_challenge_server_ids(
user_api_key_auth, allowed_servers, toolset_allowed_server_ids
),
raw_headers=raw_headers,
)
@ -2927,11 +2954,12 @@ if MCP_AVAILABLE:
await _check_passthrough_upstream_auth(
scope,
user_api_key_auth,
mcp_servers,
_sse_client_ip,
allowed_servers,
oauth2_headers=oauth2_headers,
mcp_server_auth_headers=mcp_server_auth_headers,
raw_headers=raw_headers,
requested_servers=mcp_servers,
client_ip=_sse_client_ip,
)
set_auth_context(
user_api_key_auth=user_api_key_auth,

493
tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py
View file

@ -15,6 +15,7 @@ from fastapi import HTTPException
from mcp import ReadResourceResult, Resource
from mcp.server.context import ServerRequestContext
from mcp.server.models import InitializationOptions
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.types import (
INVALID_REQUEST,
METHOD_NOT_FOUND,
@ -7312,10 +7313,6 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual
with (
_patch_delegate_resolver(server, "delegate_test"),
patch(
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')),
@ -7324,8 +7321,7 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test"],
client_ip=None,
allowed_servers=[server],
oauth2_headers=None,
mcp_server_auth_headers=None,
raw_headers=None,
@ -7356,18 +7352,14 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token()
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).respond(200)
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["delegate_test"],
client_ip=None,
allowed_servers=[server],
requested_servers=["delegate_test"],
oauth2_headers={"Authorization": "Bearer upstream-token"},
mcp_server_auth_headers=None,
raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"},
@ -7379,7 +7371,9 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token()
assert all("sk-litellm-proxy-key" not in value for value in sent_headers.values())
def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[object, object, None]:
def _expired_delegate_token_preflight(
mcp_servers: Sequence[str], allowed_servers: Sequence[MCPServer]
) -> Coroutine[object, object, None]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_check_passthrough_upstream_auth,
@ -7388,7 +7382,8 @@ def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[o
return _check_passthrough_upstream_auth(
scope=_delegate_scope([(b"authorization", b"Bearer expired-upstream-token")]),
user_api_key_auth=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}),
mcp_servers=list(mcp_servers),
allowed_servers=allowed_servers,
requested_servers=mcp_servers,
client_ip=None,
oauth2_headers={"Authorization": "Bearer expired-upstream-token"},
mcp_server_auth_headers=None,
@ -7397,28 +7392,24 @@ def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[o
@contextlib.contextmanager
def _rejecting_delegate_upstream() -> Iterator[None]:
def _rejecting_delegate_upstream() -> Iterator[MCPServer]:
server: Final = _delegate_auth_mcp_server()
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch( # test-quality-ok: stands in for the upstream rejecting an expired token
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')),
),
):
yield
yield server
@pytest.mark.asyncio
async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_challenge() -> None:
"""An expired upstream token on a single-server route gets the RFC 9728 re-authorization challenge
for the requested name, so the client signs in again instead of seeing an empty tool list."""
with _rejecting_delegate_upstream(), pytest.raises(HTTPException) as exc_info:
await _expired_delegate_token_preflight(["delegate_test"])
with _rejecting_delegate_upstream() as server, pytest.raises(HTTPException) as exc_info:
await _expired_delegate_token_preflight(["delegate_test"], [server])
assert exc_info.value.status_code == 401
challenge: Final = exc_info.value.headers["www-authenticate"]
@ -7431,8 +7422,8 @@ async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_chall
@pytest.mark.asyncio
async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_token() -> None:
"""A multi-server scope is not probed, so one rejected upstream token cannot fail the whole connect."""
with _rejecting_delegate_upstream():
assert await _expired_delegate_token_preflight(["delegate_test", "other_server"]) is None
with _rejecting_delegate_upstream() as server:
assert await _expired_delegate_token_preflight(["delegate_test", "other_server"], [server]) is None
@pytest.mark.asyncio
@ -7447,10 +7438,6 @@ async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_autho
)
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
respx.mock as upstream,
):
upstream.post(server.url).mock(
@ -7459,7 +7446,7 @@ async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_autho
)
)
outcome: Final = await asyncio.gather(
_expired_delegate_token_preflight(["delegate_test"]), return_exceptions=True
_expired_delegate_token_preflight(["delegate_test"], [server]), return_exceptions=True
)
assert getattr(outcome[0], "status_code", None) == 401
@ -7557,10 +7544,6 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu
)
with (
patch( # test-quality-ok: isolate authorized-server resolution so this test exercises the preflight contract
"litellm.proxy.mcp_server.operations._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.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(probe_status, None)),
@ -7570,11 +7553,13 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu
result = await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["passthrough_server"],
client_ip=None,
allowed_servers=[server],
oauth2_headers={"Authorization": "Bearer upstream-token"},
mcp_server_auth_headers=None,
raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"},
raw_headers={
"x-litellm-api-key": "sk-litellm-proxy-key",
"Authorization": "Bearer upstream-token",
},
)
assert result is None
else:
@ -7582,11 +7567,13 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["passthrough_server"],
client_ip=None,
allowed_servers=[server],
oauth2_headers={"Authorization": "Bearer upstream-token"},
mcp_server_auth_headers=None,
raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"},
raw_headers={
"x-litellm-api-key": "sk-litellm-proxy-key",
"Authorization": "Bearer upstream-token",
},
)
assert exc_info.value.status_code == expected_status
if expected_status == 401:
@ -7609,10 +7596,6 @@ async def test_delegate_tokenless_request_not_probed():
with (
_patch_delegate_resolver(server, "delegate_test"),
patch(
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, None)),
@ -7621,8 +7604,7 @@ async def test_delegate_tokenless_request_not_probed():
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test"],
client_ip=None,
allowed_servers=[server],
oauth2_headers=None,
mcp_server_auth_headers=None,
raw_headers=None,
@ -7645,10 +7627,6 @@ async def test_delegate_preflight_skipped_on_multi_server_routes():
with (
_patch_delegate_resolver(servers[0], "delegate_test", "other_server"),
patch(
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=servers),
),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, None)),
@ -7657,8 +7635,7 @@ async def test_delegate_preflight_skipped_on_multi_server_routes():
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["delegate_test", "other_server"],
client_ip=None,
allowed_servers=servers,
oauth2_headers=None,
mcp_server_auth_headers=None,
raw_headers=None,
@ -7691,10 +7668,6 @@ async def test_bare_authorization_never_probes_passthrough_servers():
with (
_patch_delegate_resolver(passthrough_server, "pt_server"),
patch(
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[passthrough_server]),
),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, None)),
@ -7703,8 +7676,7 @@ async def test_bare_authorization_never_probes_passthrough_servers():
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(),
mcp_servers=["pt_server"],
client_ip=None,
allowed_servers=[passthrough_server],
oauth2_headers=None,
mcp_server_auth_headers=None,
raw_headers=None,
@ -7740,10 +7712,6 @@ async def test_delegate_not_probed_when_named_only_via_server_id():
with (
_patch_delegate_resolver(server, "delegate_test"),
patch(
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, None)),
@ -7752,8 +7720,7 @@ async def test_delegate_not_probed_when_named_only_via_server_id():
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="u1", api_key="hashed-sk"),
mcp_servers=["delegate-secret-id"],
client_ip=None,
allowed_servers=[server],
oauth2_headers=None,
mcp_server_auth_headers=None,
raw_headers=None,
@ -7780,10 +7747,6 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members() -> None:
)
with (
_patch_delegate_resolver(rejecting, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[rejecting, accepting]),
),
respx.mock as upstream,
):
upstream.post(rejecting.url).respond(401, headers={"www-authenticate": 'Bearer error="invalid_token"'})
@ -7794,7 +7757,8 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members() -> None:
[(b"x-litellm-api-key", b"sk-litellm-proxy-key"), (b"authorization", b"Bearer upstream-token")]
),
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
mcp_servers=["prod_tools_group"],
allowed_servers=[rejecting, accepting],
requested_servers=["prod_tools_group"],
client_ip=None,
oauth2_headers={"Authorization": "Bearer upstream-token"},
mcp_server_auth_headers=None,
@ -10891,6 +10855,395 @@ class TestPreemptive401ModeAware:
await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
class TestPreemptive401GatedByKeyAllowlist:
"""The preemptive OAuth 401 challenge must be gated on the servers the
caller's key is allowed to reach, not only on toolset scoping. A key whose
object_permission.mcp_servers excludes an oauth2 server must fall through
to the existing 403 denial instead of being pushed into an OAuth flow it
can never complete."""
@staticmethod
def _auth() -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="sk-litellm-virtual-key")
@staticmethod
def _initialize_call(alias: str) -> tuple[dict[str, object], AsyncMock]:
body: Final = (
b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":'
b'{"protocolVersion":"2024-11-05","capabilities":{},'
b'"clientInfo":{"name":"test","version":"0"}}}'
)
receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False})
return {"type": "http", "method": "POST", "path": f"/mcp/{alias}", "headers": []}, receive
@contextlib.contextmanager
def _patches(
self,
server: MCPServer,
*,
allowed_servers: list[MCPServer],
handle_target: StreamableHTTPSessionManager | None,
toolset_auth: UserAPIKeyAuth | None = None,
request_auth: UserAPIKeyAuth | None = None,
mcp_servers: list[str] | None = None,
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> Iterator[tuple[AsyncMock, AsyncMock]]:
def _resolve_allowed(
user_api_key_auth: UserAPIKeyAuth | None = None,
mcp_servers: list[str] | None = None,
client_ip: str | None = None,
) -> list[MCPServer]:
if user_api_key_auth is None or user_api_key_auth.object_permission is None:
return allowed_servers
scoped_ids: Final = user_api_key_auth.object_permission.mcp_servers
if not scoped_ids:
return allowed_servers
return [srv for srv in allowed_servers if srv.server_id in scoped_ids]
registry: Final = mcp_operations.global_mcp_server_manager.registry
registry[server.server_id] = server
try:
with contextlib.ExitStack() as stack:
stack.enter_context(
patch( # test-quality-ok: admission parse is unrelated to the challenge gate; the handler reads it as a module global
"litellm.proxy.mcp_server.server.extract_mcp_auth_context",
new_callable=AsyncMock,
return_value=(
request_auth if request_auth is not None else self._auth(),
None,
mcp_servers if mcp_servers is not None else [server.alias],
None,
oauth2_headers,
raw_headers,
),
)
)
allowed_servers_lookup: Final = stack.enter_context(
patch( # test-quality-ok: allowed-server resolution hits the key/toolset DB; the test must control the allowlist
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
side_effect=_resolve_allowed,
)
)
stack.enter_context(
patch( # test-quality-ok: module-level startup flag gates the transport; there is no injectable seam
"litellm.proxy.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
True,
)
)
stack.enter_context(
patch.object( # test-quality-ok: no stored per-user token is the scenario; the real lookup hits the token cache/DB
mcp_operations.global_mcp_server_manager,
"has_user_oauth_token",
new_callable=AsyncMock,
return_value=False,
)
)
stack.enter_context(
patch.object( # test-quality-ok: the real cache warm-up calls the upstream server over the network
mcp_operations.global_mcp_server_manager,
"_ensure_upstream_initialize_instructions_cached",
new_callable=AsyncMock,
)
)
if toolset_auth is not None:
stack.enter_context(
patch( # test-quality-ok: toolset resolution hits the DB; the test controls the scoped auth directly
"litellm.proxy.mcp_server.server.apply_toolset_scope",
new_callable=AsyncMock,
return_value=toolset_auth,
)
)
stack.enter_context(
patch( # test-quality-ok: toolset resolution hits the DB; the test controls the scoped server ids
"litellm.proxy.mcp_server.server._toolset_server_ids",
new_callable=AsyncMock,
return_value={"id-other"},
)
)
handle_request: Final = AsyncMock()
if handle_target is not None:
stack.enter_context(
patch.object( # test-quality-ok: session dispatch is a module-level global; rejected requests must not reach it
handle_target,
"handle_request",
new=handle_request,
)
)
yield allowed_servers_lookup, handle_request
finally:
registry.pop(server.server_id, None)
@pytest.mark.asyncio
async def test_key_not_allowed_gets_403_not_401_streamable(self):
"""Denied key on an oauth2 server must reach the scoped-access denial
(403) instead of an OAuth challenge (401) it can never satisfy."""
from litellm.proxy.mcp_server import server as server_module
server: Final = _make_oauth2_server("restricted")
scope, receive = self._initialize_call("restricted")
with (
self._patches(server, allowed_servers=[], handle_target=server_module.session_manager_stateful) as (
_,
handle_request,
),
pytest.raises(HTTPException) as exc,
):
await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock())
assert exc.value.status_code == 403
handle_request.assert_not_awaited()
@pytest.mark.asyncio
async def test_key_allowed_still_challenges_streamable(self):
"""Key allowlist includes the oauth2 server with no stored token:
the 401 challenge still fires, so the gate is not over-broad."""
from litellm.proxy.mcp_server import server as server_module
server: Final = _make_oauth2_server("restricted")
scope, receive = self._initialize_call("restricted")
with (
self._patches(server, allowed_servers=[server], handle_target=server_module.session_manager_stateful) as (
_,
handle_request,
),
pytest.raises(HTTPException) as exc,
):
await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock())
assert exc.value.status_code == 401
assert "www-authenticate" in {k.lower() for k in exc.value.headers}
handle_request.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("server_kind", "transport"),
[("pt", "streamable"), ("pt", "sse"), ("dl", "streamable"), ("dl", "sse")],
)
async def test_keyless_cold_start_still_gets_oauth_challenge(self, server_kind: str, transport: str) -> None:
from litellm.proxy.mcp_server import server as server_module
server: Final = MCPServer(
server_id=f"id-{server_kind}",
name=server_kind,
alias=server_kind,
server_name=server_kind,
url=None,
transport=MCPTransport.http,
auth_type=MCPAuth.none if server_kind == "pt" else MCPAuth.oauth2,
extra_headers=["Authorization"] if server_kind == "pt" else [],
oauth_passthrough=server_kind == "pt",
delegate_auth_to_upstream=server_kind == "dl",
oauth2_flow="authorization_code" if server_kind == "dl" else None,
authorization_url="https://idp.test/authorize" if server_kind == "dl" else None,
token_url="https://idp.test/token" if server_kind == "dl" else None,
mcp_info={"server_name": server_kind},
)
body: Final = (
b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":'
b'{"protocolVersion":"2024-11-05","capabilities":{},'
b'"clientInfo":{"name":"test","version":"0"}}}'
)
scope: Final = {
"type": "http",
"method": "POST" if transport == "streamable" else "GET",
"path": f"/mcp/{server_kind}",
"scheme": "http",
"server": ("127.0.0.1", 4000),
"root_path": "",
"query_string": b"",
"headers": [
(b"host", b"127.0.0.1:4000"),
(b"accept", b"application/json, text/event-stream"),
(b"content-type", b"application/json"),
],
}
receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False})
registry: Final = mcp_operations.global_mcp_server_manager.registry
expected_challenge: Final = (
f'Bearer resource_metadata="http://127.0.0.1:4000/.well-known/oauth-protected-resource/mcp/{server_kind}"'
)
with (
patch.dict(registry, {server.server_id: server}),
patch( # test-quality-ok: force the admitted anonymous cold-start auth context without DB or network
"litellm.proxy.mcp_server.server.extract_mcp_auth_context",
new_callable=AsyncMock,
return_value=(UserAPIKeyAuth(), None, [server.alias], None, None, None),
),
patch( # test-quality-ok: module-level startup flag gates the transport, with no injectable seam
"litellm.proxy.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True
),
):
if transport == "streamable":
handle_request: Final = AsyncMock()
with (
patch.object(
server_module.session_manager_stateful,
"handle_request",
new=handle_request,
), # test-quality-ok: rejected requests must not dispatch to the session manager
pytest.raises(HTTPException) as exc,
):
await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock())
handle_request.assert_not_awaited()
else:
with (
patch.object(server_module.sse, "connect_sse") as connect_sse,
pytest.raises(HTTPException) as exc,
):
await server_module.handle_sse_mcp(scope, receive, AsyncMock())
connect_sse.assert_not_called()
assert exc.value.status_code == 401
assert exc.value.headers is not None
assert exc.value.headers["www-authenticate"] == expected_challenge
@pytest.mark.asyncio
async def test_streamable_probe_reuses_allowed_servers( # test-quality-ok: one lookup per request is the contract
self, monkeypatch: pytest.MonkeyPatch
):
from litellm.proxy.mcp_server import server as server_module
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"mcp_advertised_versions": ["2026-07-28"]})
server: Final = MCPServer(
server_id="id-restricted",
name="restricted",
alias="restricted",
server_name="restricted",
url="https://restricted.test/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
extra_headers=["Authorization"],
oauth_passthrough=True,
)
body: Final = json.dumps(
{
"jsonrpc": "2.0",
"id": 7,
"method": "tools/call",
"params": {
"name": "confirm",
"arguments": {},
"_meta": {
"io.modelcontextprotocol/protocolVersion": "2026-07-28",
"io.modelcontextprotocol/clientCapabilities": {},
},
},
}
).encode()
scope: Final = {
"type": "http",
"method": "POST",
"path": "/mcp",
"scheme": "http",
"server": ("localhost", 8000),
"query_string": b"",
"root_path": "",
"headers": [
(b"content-type", b"application/json"),
(b"x-litellm-api-key", b"sk-litellm-virtual-key"),
(b"authorization", b"Bearer upstream-token"),
(b"mcp-protocol-version", b"2026-07-28"),
(b"mcp-method", b"tools/call"),
(b"mcp-name", b"confirm"),
(b"accept", b"application/json, text/event-stream"),
],
}
receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False})
with (
self._patches(
server,
allowed_servers=[server],
handle_target=server_module.session_manager_stateless,
mcp_servers=["restricted", "another"],
oauth2_headers={"Authorization": "Bearer upstream-token"},
raw_headers={
"x-litellm-api-key": "sk-litellm-virtual-key",
"Authorization": "Bearer upstream-token",
},
) as (allowed_servers_lookup, handle_request),
patch(
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new_callable=AsyncMock,
return_value=(200, None),
) as probe,
):
await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock())
allowed_servers_lookup.assert_awaited_once()
probe.assert_awaited_once_with(server.url, "Bearer upstream-token")
handle_request.assert_awaited_once()
@pytest.mark.asyncio
async def test_toolset_excluding_server_gets_403_not_401_streamable(self):
"""Key allowlist includes the server but the active toolset scope
excludes it: the caller-visible denial is the scoped-access 403, not
a 401 challenge."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
from litellm.proxy.mcp_server import server as server_module
server: Final = _make_oauth2_server("restricted")
key_auth: Final = self._auth().model_copy(
update={
"object_permission": LiteLLM_ObjectPermissionTable(
object_permission_id="key-scope",
mcp_servers=[server.server_id],
)
}
)
scoped_auth: Final = key_auth.model_copy(
update={
"object_permission": LiteLLM_ObjectPermissionTable(
object_permission_id="toolset-scope",
mcp_servers=["id-other"],
)
}
)
scope, receive = self._initialize_call("restricted")
token: Final = server_module._mcp_active_toolset_id.set("toolset-1")
try:
with (
self._patches(
server,
allowed_servers=[server],
handle_target=server_module.session_manager_stateful,
toolset_auth=scoped_auth,
request_auth=key_auth,
) as (_, handle_request),
pytest.raises(HTTPException) as exc,
):
await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock())
finally:
server_module._mcp_active_toolset_id.reset(token)
assert exc.value.status_code == 403
assert "www-authenticate" not in {k.lower() for k in exc.value.headers or {}}
handle_request.assert_not_awaited()
@pytest.mark.asyncio
async def test_key_not_allowed_gets_403_not_401_sse(self):
"""Same gate on the SSE transport: a denied key gets 403, not a 401."""
from litellm.proxy.mcp_server import server as server_module
server: Final = _make_oauth2_server("restricted")
scope: Final = {"type": "http", "method": "GET", "path": "/mcp/restricted", "headers": []}
with (
self._patches(server, allowed_servers=[], handle_target=None),
patch.object(server_module.sse, "connect_sse") as connect_sse,
pytest.raises(HTTPException) as exc,
):
await server_module.handle_sse_mcp(scope, AsyncMock(), AsyncMock())
assert exc.value.status_code == 403
connect_sse.assert_not_called()
class TestSingleServerPreflightReachesIdJag:
"""The connect-time preflight is what turns a credential failure into an HTTP status the client
can read. An oauth2_id_jag server has to reach it: its subject comes from the assertion stored at
@ -12265,6 +12618,7 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow
)
for name in ("first", "second")
]
allowed_servers: Final = targets if allowed else []
mcp_operations.global_mcp_server_manager.registry.update({target.server_id: target for target in targets})
with respx.mock(assert_all_called=False) as upstream:
first: Final = upstream.post(targets[0].url).respond(200)
@ -12279,8 +12633,7 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow
mcp_servers=[target.server_id for target in targets] if allowed else [],
),
),
mcp_servers=["first", "second"],
client_ip=None,
allowed_servers=allowed_servers,
oauth2_headers=None,
mcp_server_auth_headers={
"first": {"x-other" if header_shape == "non_authorization" else "aUtHoRiZaTiOn": "Bearer first-token"},

44
tests/unit/proxy/mcp_server/test_mcp_stale_session.py
View file

@ -652,6 +652,11 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[oauth_server],
),
patch.object(
session_manager_stateless,
"handle_request",
@ -740,6 +745,11 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[oauth_server],
),
patch.object(
session_manager_stateless,
"handle_request",
@ -1035,6 +1045,11 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[oauth_server],
),
patch.object(
session_manager_stateless,
"handle_request",
@ -1138,6 +1153,11 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=delegated_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[delegated_server],
),
patch.object(
session_manager_stateful,
"handle_request",
@ -1226,6 +1246,11 @@ async def test_handle_streamable_http_mcp_token_exchange_without_subject_returns
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=obo_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[obo_server],
),
patch.object(
session_manager_stateful,
"handle_request",
@ -1325,6 +1350,11 @@ async def test_handle_streamable_http_mcp_oauth_delegate_without_token_returns_g
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=od_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[od_server],
),
patch.object(
session_manager_stateful,
"handle_request",
@ -1582,6 +1612,11 @@ async def test_handle_streamable_http_mcp_true_passthrough_without_token_surface
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=tp_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[tp_server],
),
patch.object(
session_manager_stateful,
"handle_request",
@ -1650,6 +1685,11 @@ async def test_handle_streamable_http_mcp_true_passthrough_dcr_bridge_challenges
"litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=bridge_server,
),
patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[bridge_server],
),
patch.object(
session_manager_stateful,
"handle_request",
@ -1827,9 +1867,7 @@ async def test_streamable_http_skips_preflight_for_disabled_methods( # test-qua
"tp_server",
extra_headers=request_headers,
)
body: Final[bytes] = json.dumps(
{"jsonrpc": "2.0", "id": 1, "method": method, "params": params}
).encode()
body: Final[bytes] = json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode()
receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False})
send: Final = AsyncMock()
user_auth: Final = UserAPIKeyAuth()