fix(mcp): scope gateway OAuth consent to x-mcp-servers and add explicit scope setting (#44957)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-15 15:58:34 -07:00 • committed by GitHub
parent 0a17da59fd
commit 1ddaa54793
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 2045 additions and 432 deletions

View file

@ -307,12 +307,43 @@ def _gateway_dcr_challenge_target(
targets: Final = _parse_mcp_server_names_from_path(route, mcp_servers)
if targets is None:
return None
server: Final = global_mcp_server_manager.get_mcp_server_by_name(targets[0], client_ip=client_ip)
server: Final = global_mcp_server_manager.get_mcp_server_by_identifier(targets[0], client_ip=client_ip)
if server is None or not server.advertises_gateway_authorization_server:
return None
return targets[0]
def gateway_dcr_challenge_scope(
route: str,
mcp_servers: Sequence[str] | None,
client_ip: str | None,
) -> str | None:
if MCPRequestHandler.extract_target_server_names_from_path(route) or not mcp_servers:
return None
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import gateway_server_scope
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
resolved: Final = tuple(
(entry, global_mcp_server_manager.get_mcp_server_by_identifier(entry, client_ip=client_ip))
for entry in dict.fromkeys(mcp_servers)
)
if any(server is None or not server.advertises_gateway_authorization_server for _, server in resolved):
return None
servers: Final = tuple((entry, server) for entry, server in resolved if server is not None)
first_by_id: Final = {server.server_id: (entry, server) for entry, server in reversed(servers)}
unique_servers: Final = tuple(
first_by_id[server_id] for server_id in dict.fromkeys(server.server_id for _, server in servers)
)
tokens: Final = tuple(
gateway_server_scope(entry) or gateway_server_scope(server.server_id) for entry, server in unique_servers
)
if any(token is None for token in tokens):
return None
return " ".join(token for token in tokens if token is not None)
def _is_gateway_dcr_challenge_scope(
route: str,
mcp_servers: list[str] | None,
@ -327,9 +358,10 @@ def _is_gateway_dcr_challenge_scope(
Fires only for a genuine 401 with no client-supplied MCP auth headers (those mean
the caller is not a cold-start DCR client), on the scopes the gateway's keyless
flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request
(the resource the client configured is still ``/mcp``), or a per-server path whose
single target advertises gateway-owned sign-in. Every other named target keeps
its existing behavior, failing closed to the original admission error."""
(the resource the client configured is still ``/mcp``; a single eligible header
target is also carried in the OAuth scope), or a per-server path whose single target
advertises gateway-owned sign-in. Every other named target keeps its existing
behavior, failing closed to the original admission error."""
if not _is_litellm_auth_admission_error(exc):
return False
if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers):
@ -355,20 +387,23 @@ def _gateway_dcr_challenge(
present a bearer that failed admission (expired or revoked), telling
spec-compliant clients to re-authorize rather than retry; a request with
no credentials at all gets the bare challenge per RFC 6750 section 3.1."""
target: Final = _gateway_dcr_challenge_target(route, mcp_servers, IPAddressUtils.get_mcp_client_ip(request))
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
target: Final = _gateway_dcr_challenge_target(route, mcp_servers, client_ip)
server_scope: Final = gateway_dcr_challenge_scope(route, mcp_servers, client_ip) if target is None else None
resource_metadata_url: Final = (
get_passthrough_resource_metadata_url(request.scope, target)
if target is not None
else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp"
)
error_attr: Final = 'error="invalid_token", ' if invalid_token else ""
scope_attr: Final = f', scope="{server_scope}"' if server_scope is not None else ""
return HTTPException(
status_code=401,
detail={
"error": "authentication_required",
"message": "Authenticate with the gateway to use the MCP endpoint.",
},
headers={"WWW-Authenticate": f'Bearer {error_attr}resource_metadata="{resource_metadata_url}"'},
headers={"WWW-Authenticate": f'Bearer {error_attr}resource_metadata="{resource_metadata_url}"{scope_attr}'},
)
@ -1123,7 +1158,7 @@ class MCPRequestHandler:
case SessionBearerAdmitted():
try:
admitted: Final = await MCPRequestHandler.reload_admitted_user(result.principal.user_id)
admitted.mcp_session_resource_server_id = result.principal.resource_server_id
admitted.mcp_session_resource_server_ids = result.principal.resource_server_ids
await MCPRequestHandler._enforce_admitted_live_policy(
admitted=admitted, request=request, route=route
)

View file

@ -3113,10 +3113,28 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict:
return {
"authorization_servers": [f"{request_base_url}/mcp"],
"resource": f"{request_base_url}/mcp",
"scopes_supported": [],
"scopes_supported": _aggregate_header_scopes(request),
}
def _aggregate_header_scopes(request: Request) -> tuple[str, ...]:
"""Gateway server-scope tokens for the ``x-mcp-servers`` header a client sends with its
discovery request, so a client that skips the 401 challenge (Claude Code re-authenticating
a server it already knows needs auth) still carries the selection to ``/authorize`` through
``scopes_supported``. Empty when the header is absent or any entry does not resolve."""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( # noqa: PLC0415 # proxy import cycle
MCPRequestHandler,
gateway_dcr_challenge_scope,
)
header: Final = request.headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME)
if not header:
return ()
names: Final = tuple(dict.fromkeys(entry.strip() for entry in header.split(",") if entry.strip()))
scope: Final = gateway_dcr_challenge_scope("/mcp", names, IPAddressUtils.get_mcp_client_ip(request))
return tuple(scope.split()) if scope else ()
def _build_aggregate_authorization_server_response(
request: Request, token_exchange_available: bool
) -> dict[str, object]:

View file

@ -42,7 +42,8 @@ import hmac
import html
import secrets
from base64 import urlsafe_b64encode
from collections.abc import Awaitable, Callable, Iterable, Mapping
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Final, Literal, Protocol, TypeVar
@ -50,7 +51,8 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
from fastapi import HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, model_serializer, model_validator
from pydantic_core.core_schema import SerializerFunctionWrapHandler
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
from litellm._internal_context import with_service_target
@ -106,6 +108,8 @@ GATEWAY_DCR_CLIENT_ID_PREFIX: Final = "llm_dcrc_"
endpoints can route an aggregate-flow request without decrypting, and existing per-server
flows (whose client_ids are upstream-issued) are never captured by the aggregate arm."""
GATEWAY_SERVER_SCOPE_PREFIX: Final = "litellm:mcp_server:"
GATEWAY_AUTH_CODE_PREFIX: Final = "llm_gcode_"
"""Marker prefix on the gateway-sealed authorization code, distinct from the bridge
``llm_bcode_`` so neither flow can consume the other's codes."""
@ -301,13 +305,55 @@ class GatewayDcrClient(LiteLLMBaseModel):
iat: int
class _ConnectFlow(LiteLLMBaseModel):
_SEALED_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _legacy_server_scope_payload(payload: object) -> object:
if not isinstance(payload, Mapping):
return payload
values: Final = _SEALED_PAYLOAD_ADAPTER.validate_python(payload)
if "resource_server_id" not in values:
return values
if "resource_server_ids" in values:
raise ValueError("sealed payload cannot contain both resource_server_id and resource_server_ids")
resource_server_id: Final = values["resource_server_id"]
return {key: value for key, value in values.items() if key != "resource_server_id"} | (
{"resource_server_ids": (resource_server_id,)} if resource_server_id is not None else {}
)
def _server_scope_wire_payload(
payload: Mapping[str, object], resource_server_ids: tuple[str, ...] | None
) -> Mapping[str, object]:
if resource_server_ids is not None and len(resource_server_ids) != 1:
return payload
unscoped: Final = {key: value for key, value in payload.items() if key != "resource_server_ids"}
if resource_server_ids is None:
return unscoped
return unscoped | {"resource_server_id": resource_server_ids[0]}
class _ServerScopedSealedModel(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
resource_server_ids: tuple[str, ...] | None = Field(default=None, min_length=1)
@model_validator(mode="before")
@classmethod
def _accept_legacy_server_scope(cls, payload: object) -> object:
return _legacy_server_scope_payload(payload)
@model_serializer(mode="wrap")
def _serialize_server_scope(self, handler: SerializerFunctionWrapHandler) -> Mapping[str, object]:
payload: Final = _SEALED_PAYLOAD_ADAPTER.validate_python(handler(self))
return _server_scope_wire_payload(payload, self.resource_server_ids)
class _ConnectFlow(_ServerScopedSealedModel):
"""One in-flight authorize: the SSO user it belongs to and the client parameters
needed to mint the code at the finish step. Sealed into the per-flow cookie. ``jti``
makes the flow single-use at complete; ``extra="forbid"`` rejects cross-type
confusion."""
model_config = ConfigDict(frozen=True, extra="forbid")
user_id: str = Field(min_length=1)
client_id: str = Field(min_length=1)
redirect_uri: str = Field(min_length=1)
@ -315,19 +361,17 @@ class _ConnectFlow(LiteLLMBaseModel):
code_challenge: str = Field(min_length=1)
jti: str = Field(min_length=1)
exp: int
resource_server_id: str | None = None
audience: SessionAudience | None = None
requested_scopes: tuple[str, ...] = ()
upstream_scopes: tuple[UpstreamScope, ...] = ()
class _GatewayAuthCode(LiteLLMBaseModel):
class _GatewayAuthCode(_ServerScopedSealedModel):
"""The gateway-sealed authorization code: the user consent it represents and the
bindings the token endpoint must verify (client, redirect URI, PKCE challenge),
plus a ``jti`` for the single-use guard. ``extra="forbid"`` rejects cross-type
confusion."""
model_config = ConfigDict(frozen=True, extra="forbid")
user_id: str = Field(min_length=1)
client_id: str = Field(min_length=1)
redirect_uri: str = Field(min_length=1)
@ -335,7 +379,6 @@ class _GatewayAuthCode(LiteLLMBaseModel):
jti: str = Field(min_length=1)
iat: int
exp: int
resource_server_id: str | None = None
audience: SessionAudience | None = None
team_id: str | None = None
@ -516,6 +559,34 @@ def relative_request_url(request: Request) -> str:
return f"{path}?{request.url.query}" if request.url.query else path
def gateway_server_scope(server_name: str) -> str | None:
if not server_name or not all(
ord(character) == 0x21 or 0x23 <= ord(character) <= 0x5B or 0x5D <= ord(character) <= 0x7E
for character in server_name
):
return None
return f"{GATEWAY_SERVER_SCOPE_PREFIX}{server_name}"
def _is_aggregate_resource(request: Request, resource: str) -> bool:
canonical: Final = canonical_resource_uri(resource)
if canonical is None:
return False
base: Final = canonicalize_url_identity(get_request_base_url(request))
return canonical in (base, f"{base}/mcp")
def _gateway_flow_server(name: str) -> MCPServer | None:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # proxy import cycle
global_mcp_server_manager,
)
server: Final = global_mcp_server_manager.get_mcp_server_by_identifier(name)
if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server):
return None
return server
def resolve_scoped_resource_server(request: Request, resource: str | None) -> MCPServer | None:
"""Resolve an RFC 8707 ``resource`` value to the single gateway-owned server it
names, or ``None`` for every other shape: absent, the aggregate resource, a foreign
@ -536,22 +607,40 @@ def resolve_scoped_resource_server(request: Request, resource: str | None) -> MC
if canonical is None:
return None
base: Final = canonicalize_url_identity(get_request_base_url(request))
if canonical == f"{base}/mcp" or not canonical.startswith(f"{base}/"):
if _is_aggregate_resource(request, resource) or not canonical.startswith(f"{base}/"):
return None
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( # noqa: PLC0415 # proxy import cycle
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # proxy import cycle
global_mcp_server_manager,
)
names: Final = MCPRequestHandler.extract_target_server_names_from_path(canonical[len(base) :])
if len(names) != 1:
return None
server: Final = global_mcp_server_manager.get_mcp_server_by_name(names[0])
if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server):
return _gateway_flow_server(names[0])
@dataclass(frozen=True, slots=True)
class InvalidGatewayScope:
pass
def resolve_scoped_servers_from_scope(
scope: str | None,
) -> tuple[MCPServer, ...] | InvalidGatewayScope | None:
if scope is None:
return None
return server
prefixed_tokens: Final = tuple(token for token in scope.split() if token.startswith(GATEWAY_SERVER_SCOPE_PREFIX))
if not prefixed_tokens:
return None
server_names: Final = tuple(dict.fromkeys(token[len(GATEWAY_SERVER_SCOPE_PREFIX) :] for token in prefixed_tokens))
if any(not name for name in server_names):
return InvalidGatewayScope()
resolved: Final = tuple(_gateway_flow_server(name) for name in server_names)
if any(server is None for server in resolved):
return InvalidGatewayScope()
servers: Final = tuple(server for server in resolved if server is not None)
first_by_id: Final = {server.server_id: server for server in reversed(servers)}
return tuple(first_by_id[server_id] for server_id in dict.fromkeys(server.server_id for server in servers))
def aggregate_authorize(
@ -570,32 +659,42 @@ def aggregate_authorize(
LiteLLM sign-in, and hand the browser to the connect page with the flow sealed into a
per-flow cookie.
A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the
flow to that one server: the scope is sealed into the flow, carried into the code, and
bound into the session token. The connect URL carries only the flow handle; the page
learns the client origin, the scoped server, and whether its vendor OAuth is done from
:func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry
steers which server the page authorizes or names on the confirmation.
A per-server RFC 8707 ``resource`` or gateway-server scope tokens on an aggregate
request scope the flow to those servers. The scope is sealed into the flow, carried
into the code, and bound into the session token.
Validation failures respond directly with 400 and never redirect: per RFC 6749
section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and
once the client is at fault there is no trusted place to send the browser.
Validation failures respond directly with 400, except for an unknown gateway scope,
which is redirected to the already-validated ``redirect_uri`` as ``invalid_scope`` per
RFC 6749 section 4.1.2.1.
"""
rejected: Final = _rejected_authorize_request(
client_id, redirect_uri, state, code_challenge, code_challenge_method, response_type
)
if rejected is not None:
return rejected
resource_scoped_server: Final = resolve_scoped_resource_server(request, resource)
scope_resolution: Final = None if resource_scoped_server is not None else resolve_scoped_servers_from_scope(scope)
base_url: Final = get_request_base_url(request)
if session_user_id is None:
return _login_redirect(base_url, request)
scoped_server: Final = resolve_scoped_resource_server(request, resource)
if isinstance(scope_resolution, InvalidGatewayScope):
invalid_scope_location: Final = _append_query_params(
redirect_uri,
(
("error", "invalid_scope"),
("error_description", "unknown or unsupported MCP server scope"),
*_state_param(state),
),
)
return RedirectResponse(invalid_scope_location, status_code=303)
scoped_servers: Final = (resource_scoped_server,) if resource_scoped_server is not None else scope_resolution or ()
scope_tokens: Final = tuple(token for token in (scope or "").split() if token.startswith(UPSTREAM_SCOPE_PREFIX))
decoded_scopes: Final = tuple(decode_upstream_scope(token) for token in scope_tokens)
if any(item is None for item in decoded_scopes):
return _oauth_error(400, "invalid_scope", "invalid upstream permission scope")
upstream_scopes: Final = tuple(item for item in decoded_scopes if item is not None)
if scoped_server is not None and any(item.server_id != scoped_server.server_id for item in upstream_scopes):
scoped_server_ids: Final = frozenset(server.server_id for server in scoped_servers)
if scoped_server_ids and any(item.server_id not in scoped_server_ids for item in upstream_scopes):
return _oauth_error(400, "invalid_scope", "upstream permission does not match the requested resource")
handle: Final = secrets.token_urlsafe(24)
flow: Final = _new_connect_flow(
@ -604,9 +703,11 @@ def aggregate_authorize(
redirect_uri=redirect_uri,
state=state,
code_challenge=code_challenge or "",
resource_server_id=scoped_server.server_id if scoped_server is not None else None,
resource_server_ids=tuple(server.server_id for server in scoped_servers) or None,
audience=None,
requested_scopes=tuple((scope or "").split()) if scoped_server is not None and not upstream_scopes else (),
requested_scopes=tuple((scope or "").split())
if resource_scoped_server is not None and not upstream_scopes
else (),
upstream_scopes=upstream_scopes,
)
connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
@ -653,7 +754,7 @@ async def native_client_authorize(
redirect_uri=redirect_uri,
state=state,
code_challenge=code_challenge or "",
resource_server_id=None,
resource_server_ids=None,
audience=PROXY_API_AUDIENCE,
)
page: Final = render_native_client_consent_page(
@ -765,7 +866,7 @@ def _new_connect_flow(
redirect_uri: str,
state: str,
code_challenge: str,
resource_server_id: str | None,
resource_server_ids: tuple[str, ...] | None,
audience: SessionAudience | None,
requested_scopes: tuple[str, ...] = (),
upstream_scopes: tuple[UpstreamScope, ...] = (),
@ -779,7 +880,7 @@ def _new_connect_flow(
code_challenge=code_challenge,
jti=secrets.token_urlsafe(24),
exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS,
resource_server_id=resource_server_id,
resource_server_ids=resource_server_ids,
audience=audience,
requested_scopes=requested_scopes,
upstream_scopes=upstream_scopes,
@ -862,17 +963,82 @@ async def _flow_target(
return state, server
class ConnectFlowServer(TypedDict):
server_id: ReadOnly[str]
server_name: ReadOnly[str]
connected: ReadOnly[bool]
class ConnectFlowDescription(TypedDict):
"""What the connect page is allowed to know about one in-flight flow."""
state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]]
state: ReadOnly[Literal["unscoped", "interactive", "m2m", "multi", "stale"]]
client_origin: ReadOnly[str]
server_id: ReadOnly[str | None]
server_name: ReadOnly[str | None]
connected: ReadOnly[bool | None]
servers: ReadOnly[Sequence[ConnectFlowServer] | None]
requested_scopes: NotRequired[ReadOnly[tuple[str, ...]]]
async def _describe_connect_flow_server(
user_id: str,
server: MCPServer,
interactive: bool,
lookup_vendor_credential: LookupVendorCredential,
) -> ConnectFlowServer | Response:
credential: Final = await lookup_vendor_credential(user_id, server.server_id) if interactive else "present"
if credential == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
return {
"server_id": server.server_id,
"server_name": server.server_name or server.alias or server.name,
"connected": credential == "present",
}
def _unscoped_or_stale_description(flow: _ConnectFlow, state: Literal["unscoped", "stale"]) -> ConnectFlowDescription:
return {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
"server_id": None,
"server_name": None,
"connected": None,
"servers": None,
}
async def _describe_multi_server_flow(
flow: _ConnectFlow,
server_ids: tuple[str, ...],
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> ConnectFlowDescription | Response:
targets: Final = tuple(
[await _flow_target(flow.user_id, server_id, lookup_server_reachability) for server_id in server_ids]
)
if any(server is None for _, server in targets):
return _unscoped_or_stale_description(flow, "stale")
described: Final = tuple(
[
await _describe_connect_flow_server(flow.user_id, server, state == "interactive", lookup_vendor_credential)
for state, server in targets
if server is not None
]
)
unavailable: Final = next((item for item in described if isinstance(item, Response)), None)
if unavailable is not None:
return unavailable
return {
"state": "multi",
"client_origin": _origin_only(flow.redirect_uri),
"server_id": None,
"server_name": None,
"connected": None,
"servers": [item for item in described if not isinstance(item, Response)],
}
async def _describe_opened_flow(
flow: _ConnectFlow,
lookup_vendor_credential: LookupVendorCredential,
@ -880,7 +1046,10 @@ async def _describe_opened_flow(
) -> ConnectFlowDescription | Response:
if flow.upstream_scopes:
return await _describe_upstream_scope_flow(flow, lookup_vendor_credential, lookup_server_reachability)
state, server = await _flow_target(flow.user_id, flow.resource_server_id, lookup_server_reachability)
server_ids: Final = flow.resource_server_ids or ()
if len(server_ids) > 1:
return await _describe_multi_server_flow(flow, server_ids, lookup_vendor_credential, lookup_server_reachability)
state, server = await _flow_target(flow.user_id, next(iter(server_ids), None), lookup_server_reachability)
if state == "interactive" and server is not None:
return await _describe_interactive_target(flow, server, flow.requested_scopes, lookup_vendor_credential)
described: Final[ConnectFlowDescription] = {
@ -889,6 +1058,7 @@ async def _describe_opened_flow(
"server_id": None if server is None else server.server_id,
"server_name": None if server is None else (server.server_name or server.alias or server.name),
"connected": state == "m2m" or None,
"servers": None,
}
return described
@ -912,6 +1082,7 @@ async def _describe_interactive_target(
"server_id": server.server_id,
"server_name": server.server_name or server.alias or server.name,
"connected": credential == "present",
"servers": None,
}
return {**described, "requested_scopes": scopes} if scopes else described
@ -946,15 +1117,12 @@ async def _describe_upstream_scope_flow(
for description in descriptions:
if isinstance(description, Response) or not description["connected"]:
return description
if flow.resource_server_id is not None:
return descriptions[0]
return {
"state": "unscoped",
"client_origin": _origin_only(flow.redirect_uri),
"server_id": None,
"server_name": None,
"connected": None,
}
sealed: Final = flow.resource_server_ids
if sealed is None:
return _unscoped_or_stale_description(flow, "unscoped")
if len(sealed) > 1:
return await _describe_multi_server_flow(flow, sealed, lookup_vendor_credential, lookup_server_reachability)
return descriptions[0]
async def describe_connect_flow(
@ -1006,6 +1174,8 @@ async def complete_connect_flow(
return described
if described["state"] == "stale":
return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available")
if described["state"] == "multi" and not all(server["connected"] for server in described["servers"] or ()):
return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing")
if described["connected"] is False:
return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing")
flow_refusal: Final = _claim_refusal(
@ -1026,12 +1196,12 @@ async def complete_connect_flow(
return response
def _state_param(flow: _ConnectFlow) -> tuple[tuple[str, str], ...]:
return (("state", flow.state),) if flow.state else ()
def _state_param(state: str) -> tuple[tuple[str, str], ...]:
return (("state", state),) if state else ()
def _denied_flow_response(flow: _ConnectFlow) -> Response:
params: Final = (("error", "access_denied"), *_state_param(flow))
params: Final = (("error", "access_denied"), *_state_param(flow.state))
return RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303)
@ -1048,12 +1218,12 @@ def _approved_flow_response(flow: _ConnectFlow, delivery: str | None, team_id: s
jti=secrets.token_urlsafe(24),
iat=int(now.timestamp()),
exp=int(now.timestamp()) + code_ttl,
resource_server_id=flow.resource_server_id,
resource_server_ids=flow.resource_server_ids,
audience=flow.audience,
team_id=(team_id or None) if flow.audience == PROXY_API_AUDIENCE else None,
),
)
callback_url: Final = _append_query_params(flow.redirect_uri, (("code", code), *_state_param(flow)))
callback_url: Final = _append_query_params(flow.redirect_uri, (("code", code), *_state_param(flow.state)))
if manual_delivery:
return _manual_delivery_response(callback_url)
return RedirectResponse(callback_url, status_code=303)
@ -1313,17 +1483,19 @@ def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response:
def _resource_conflicts_with_scope(
request: Request, resource: str | None, sealed_resource_server_id: str | None
request: Request, resource: str | None, sealed_resource_server_ids: tuple[str, ...] | None
) -> bool:
"""True when a scoped grant is being redeemed for a DIFFERENT resource than the one
sealed into it (RFC 8707 section 2.2: reject with ``invalid_target``). An absent
``resource`` never conflicts (the sealed scope still binds the minted session), and an
unscoped grant ignores the parameter entirely, exactly as the endpoint always has, so
no pre-existing client breaks."""
if sealed_resource_server_id is None or resource is None:
``resource`` never conflicts (the sealed scope still binds the minted session), the
aggregate resource does not conflict with a single-server grant, and an unscoped grant
ignores the parameter entirely."""
if sealed_resource_server_ids is None or resource is None:
return False
if _is_aggregate_resource(request, resource):
return False
resolved: Final = resolve_scoped_resource_server(request, resource)
return resolved is None or resolved.server_id != sealed_resource_server_id
return resolved is None or sealed_resource_server_ids != (resolved.server_id,)
async def aggregate_token(
@ -1545,7 +1717,7 @@ async def _authorization_code_grant(
return _oauth_error(400, "invalid_grant", "the authorization code has expired")
if client_id != parsed.client_id or redirect_uri != parsed.redirect_uri:
return _oauth_error(400, "invalid_grant", "the authorization code was issued to a different client")
if _resource_conflicts_with_scope(request, resource, parsed.resource_server_id):
if _resource_conflicts_with_scope(request, resource, parsed.resource_server_ids):
return _oauth_error(400, "invalid_target", "resource does not match the scope this code was issued for")
if not _pkce_verifier_matches(code_verifier, parsed.code_challenge):
return _oauth_error(400, "invalid_grant", "PKCE verification failed")
@ -1555,7 +1727,7 @@ async def _authorization_code_grant(
SessionPrincipal(
user_id=parsed.user_id,
client_id=client_id,
resource_server_id=parsed.resource_server_id,
resource_server_ids=parsed.resource_server_ids,
audience=parsed.audience,
team_id=parsed.team_id,
),
@ -1587,7 +1759,7 @@ async def _refresh_token_grant(
opened: Final = _open_presented_refresh_token(refresh_token, client_id, signing)
if isinstance(opened, Response):
return opened
if _resource_conflicts_with_scope(request, resource, opened.principal.resource_server_id):
if _resource_conflicts_with_scope(request, resource, opened.principal.resource_server_ids):
return _oauth_error(400, "invalid_target", "resource does not match the scope this token was issued for")
# Refresh-token rotation (OAuth 2.0 Security BCP section 4.13): the presented refresh token is
# single-use, so a captured or replayed refresh token cannot mint a second pair after the
@ -1713,7 +1885,18 @@ def _active_introspection_response(opened: OpenedSessionToken) -> Response:
for key, value in (
("token_type", "Bearer" if opened.kind == "session" else None),
("team_id", principal.team_id),
("resource_server_id", principal.resource_server_id),
(
"resource_server_id",
principal.resource_server_ids[0]
if principal.resource_server_ids is not None and len(principal.resource_server_ids) == 1
else None,
),
(
"resource_server_ids",
list(principal.resource_server_ids)
if principal.resource_server_ids is not None and len(principal.resource_server_ids) > 1
else None,
),
("audience", principal.audience),
)
if value is not None

View file

@ -196,6 +196,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
normalize_server_name,
openapi_tool_name,
parse_admin_env_vars,
server_answers_to,
strip_known_server_prefix,
validate_mcp_server_name,
)
@ -3616,8 +3617,8 @@ class MCPServerManager:
return open_ids
@staticmethod
def _admitted_session_resource_scope(user_api_key_auth: UserAPIKeyAuth | None) -> str | None:
"""The single server an admitted session subject's bearer was scoped to at authorize
def _admitted_session_resource_scope(user_api_key_auth: UserAPIKeyAuth | None) -> frozenset[str] | None:
"""The servers an admitted session subject's bearer was scoped to at authorize
time (RFC 8707 resource), or None for every other principal shape and for unscoped
sessions. Read at every return path of :meth:`get_allowed_mcp_servers`, including
the exception fallback, and applied AFTER every union (grants, operator-open,
@ -3625,7 +3626,8 @@ class MCPServerManager:
fault therefore never widens a scoped bearer to the allow-all set."""
if user_api_key_auth is None or not is_mcp_admitted_user_subject(user_api_key_auth):
return None
return user_api_key_auth.mcp_session_resource_server_id
resource_server_ids: Final = user_api_key_auth.mcp_session_resource_server_ids
return None if resource_server_ids is None else frozenset(resource_server_ids)
async def get_allowed_mcp_servers(
self,
@ -3740,7 +3742,7 @@ class MCPServerManager:
if len(combined_servers) == 0:
verbose_logger.debug("No allowed MCP Servers found for user api key auth.")
scope = MCPServerManager._admitted_session_resource_scope(user_api_key_auth)
return [server_id for server_id in combined_servers if scope is None or server_id == scope]
return [server_id for server_id in combined_servers if scope is None or server_id in scope]
except Exception: # noqa: BLE001
verbose_logger.exception(
"Failed to get allowed MCP servers; team-level object_permission "
@ -3750,7 +3752,7 @@ class MCPServerManager:
return [
server_id
for server_id in dict.fromkeys(allow_all_server_ids + submitted_server_ids)
if not explicit_grants_only and (scope is None or server_id == scope)
if not explicit_grants_only and (scope is None or server_id in scope)
]
@with_service_target(MCP_SERVERS_TARGET)
@ -7329,6 +7331,19 @@ class MCPServerManager:
return server
return None
def get_mcp_server_by_identifier(self, identifier: str, client_ip: str | None = None) -> MCPServer | None:
server: Final = self.get_mcp_server_by_name(identifier, client_ip)
if server is not None:
return server
return next(
(
server
for server in self.get_filtered_registry(client_ip).values()
if server_answers_to(server, identifier)
),
None,
)
def get_filtered_registry(self, client_ip: str | None = None) -> Mapping[str, MCPServer]:
"""
Get registry filtered by client IP access control.

View file

@ -146,6 +146,9 @@ from litellm.proxy._experimental.mcp_server.utils import (
split_server_prefix_from_name,
strip_known_server_prefix,
)
from litellm.proxy._experimental.mcp_server.utils import (
server_answers_to as _server_answers_to,
)
from litellm.proxy._types import (
UserAPIKeyAuth,
)
@ -578,9 +581,36 @@ def _http_detail_message(detail: object) -> str:
return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
def _server_answers_to(server: MCPServer, name: str) -> bool:
requested: Final = name.lower()
return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def raise_if_unscoped_aggregate_request(
mcp_servers: Sequence[str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
toolset_id: str | None,
) -> None:
from litellm.proxy import proxy_server # noqa: PLC0415 # proxy import cycle
raw_settings: Final[object] = getattr(proxy_server, "general_settings", None)
if not raw_settings:
return
settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(raw_settings)
if not settings.get("mcp_require_explicit_server_scope", False):
return
if mcp_servers or toolset_id is not None:
return
if user_api_key_auth is not None and user_api_key_auth.mcp_session_resource_server_ids is not None:
return
raise HTTPException(
status_code=400,
detail={
"error": "mcp_server_scope_required",
"message": (
"This gateway requires an explicit MCP server scope: connect to /mcp/<server_name> "
"or send the x-mcp-servers header."
),
},
)
async def raise_denied_scoped_mcp_access(

View file

@ -37,6 +37,7 @@ from __future__ import annotations
import secrets
from collections import Counter
from collections.abc import Sequence
from datetime import datetime, timedelta
from functools import lru_cache
from typing import Final, Literal, TypeAlias
@ -110,20 +111,27 @@ class SessionPrincipal(LiteLLMBaseModel):
gateway-sealed) DCR client identifier the token was issued to; the token endpoint
requires it to match on the refresh grant.
``resource_server_id`` is the single MCP server this session was authorized for when
the client requested a per-server RFC 8707 resource at authorize time, or ``None`` for
the aggregate scope. It is a RESTRICTION carried for admission to intersect against
the live grant resolution, never a grant by itself; the refresh grant re-mints from
this principal so the restriction survives rotation.
``resource_server_ids`` are the MCP servers this session was authorized for when the
client requested an RFC 8707 resource or gateway scope at authorize time, or ``None``
for the aggregate scope. They are a RESTRICTION carried for admission to intersect
against the live grant resolution, never a grant by itself; the refresh grant re-mints
from this principal so the restriction survives rotation.
"""
model_config = ConfigDict(frozen=True)
user_id: str = Field(min_length=1)
client_id: str = Field(min_length=1)
resource_server_id: str | None = None
resource_server_ids: tuple[str, ...] | None = Field(default=None, min_length=1)
audience: SessionAudience | None = None
team_id: str | None = None
@field_validator("resource_server_ids")
@classmethod
def _resource_server_ids_are_nonempty(cls, value: tuple[str, ...] | None) -> tuple[str, ...] | None:
if value is not None and not value:
raise ValueError("resource_server_ids must not be empty")
return value
class SessionKeys(LiteLLMBaseModel):
"""Injected key material: the HS256 signing key.
@ -307,6 +315,7 @@ class _SessionClaims(LiteLLMBaseModel):
user_id: str = Field(min_length=1)
client_id: str = Field(min_length=1)
resource_server_id: str | None = None
resource_server_ids: Sequence[str] | None = Field(default=None, min_length=2)
audience: SessionAudience | None = None
team_id: str | None = None
family: str | None = Field(default=None, min_length=1)
@ -414,7 +423,16 @@ def _mint(
kind=kind,
user_id=principal.user_id,
client_id=principal.client_id,
resource_server_id=principal.resource_server_id,
resource_server_id=(
principal.resource_server_ids[0]
if principal.resource_server_ids is not None and len(principal.resource_server_ids) == 1
else None
),
resource_server_ids=(
list(principal.resource_server_ids)
if principal.resource_server_ids is not None and len(principal.resource_server_ids) >= 2
else None
),
audience=principal.audience,
team_id=principal.team_id,
family=family,
@ -465,13 +483,21 @@ def _open(
return claims
if claims.kind != expected_kind:
return SessionMalformed()
if claims.resource_server_id is not None and claims.resource_server_ids is not None:
return SessionMalformed()
if now.timestamp() >= claims.exp:
return SessionExpired()
return OpenedSessionToken(
principal=SessionPrincipal(
user_id=claims.user_id,
client_id=claims.client_id,
resource_server_id=claims.resource_server_id,
resource_server_ids=(
(claims.resource_server_id,)
if claims.resource_server_id is not None
else tuple(claims.resource_server_ids)
if claims.resource_server_ids is not None
else None
),
audience=claims.audience,
team_id=claims.team_id,
),

View file

@ -1272,6 +1272,14 @@ if MCP_AVAILABLE:
verbose_logger.warning("Rejected MCP request from a disallowed client application: %s", rejection.details)
raise HTTPException(status_code=403, detail=rejection.response_body)
def admit_mcp_request(
scope: Scope, user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: Sequence[str] | None
) -> None:
reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
operations.raise_if_unscoped_aggregate_request(
mcp_servers=mcp_servers, user_api_key_auth=user_api_key_auth, toolset_id=mcp_active_toolset_id.get()
)
async def extract_mcp_auth_context(scope, path):
"""
Extracts mcp_servers from the path and processes the MCP request for auth context.
@ -1776,8 +1784,8 @@ if MCP_AVAILABLE:
detail=f"API key does not have access to toolset '{toolset_id}'.",
)
if is_mcp_admitted_user_subject(acting):
resource_server_id: Final = acting.mcp_session_resource_server_id
if resource_server_id is not None and resource_server_id not in (
resource_server_ids: Final = acting.mcp_session_resource_server_ids
if resource_server_ids is not None and frozenset(resource_server_ids).isdisjoint(
await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy
)
@ -2413,7 +2421,7 @@ if MCP_AVAILABLE:
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
admit_mcp_request(scope, user_api_key_auth, mcp_servers)
body: Final = await peek.body() if peek is not None else b""
is_initialize: Final = _is_initialize_request(body)
scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1
@ -2780,7 +2788,7 @@ if MCP_AVAILABLE:
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
admit_mcp_request(scope, user_api_key_auth, mcp_servers)
scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1
# Extract client IP for MCP access control

View file

@ -379,6 +379,11 @@ def iter_known_server_prefixes(server: _McpServerLike) -> Iterator[str]:
yield from _emit(server_id)
def server_answers_to(server: _McpServerLike, name: str) -> bool:
requested: Final = name.lower()
return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
def iter_known_tool_name_spellings(tool_name: str, server: MCPServer) -> Iterator[str]:
"""Yield every name that denotes the bare ``tool_name`` on ``server``: the bare name,
then its wire spelling under each prefix ``iter_known_server_prefixes`` accepts.

View file

@ -2897,6 +2897,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.",
)
mcp_require_explicit_server_scope: bool = Field(
default=False,
description="Require aggregate MCP requests to include x-mcp-servers unless a toolset or gateway OAuth session already scopes the server. Off by default.",
)
mcp_client_id_header: str | None = Field(
None,
description="Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs.",

View file

@ -168,7 +168,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
# path via post-construction assignment and stripped from validated input like the markers
# above; a forged value could at most narrow, but the stripping keeps the field's provenance
# single-owner so its meaning stays trustworthy.
mcp_session_resource_server_id: str | None = Field(default=None, exclude=True)
mcp_session_resource_server_ids: tuple[str, ...] | None = Field(default=None, exclude=True)
mcp_toolset_id: str | None = Field(default=None, exclude=True)
authenticated_by_custom_auth: bool = Field(default=False, exclude=True)
via_virtual_key: bool = Field(
@ -229,7 +229,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
values.pop("requires_fresh_policy", None)
values.pop("mcp_explicit_grants_only", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("mcp_session_resource_server_id", None)
values.pop("mcp_session_resource_server_ids", None)
values.pop("mcp_toolset_id", None)
values.pop("via_virtual_key", None)
values.pop("authenticated_by_custom_auth", None)

View file

@ -7789,6 +7789,73 @@ class TestAggregateGatewayDcrChallenge:
www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"]
assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}"
@pytest.mark.parametrize(
("header", "expected_scope"),
[
("alpha,beta", 'scope="litellm:mcp_server:alpha litellm:mcp_server:beta"'),
("ALPHA", 'scope="litellm:mcp_server:ALPHA"'),
("alpha-id", 'scope="litellm:mcp_server:alpha-id"'),
("alpha,unknown", None),
("access-group", None),
("token-exchange", None),
],
)
async def test_aggregate_header_challenge_scopes_resolved_gateway_servers(
self, header: str, expected_scope: str | None
) -> None:
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
alpha: Final = MCPServer(
server_id="alpha-id",
name="alpha",
server_name="alpha",
alias="alpha",
url="https://upstream.example/mcp",
transport="http",
auth_type=MCPAuth("oauth2"),
)
beta: Final = MCPServer(
server_id="beta-id",
name="beta",
server_name="beta",
alias="beta",
url="https://upstream.example/mcp",
transport="http",
auth_type=MCPAuth("oauth2"),
)
token_exchange: Final = MCPServer(
server_id="token-exchange-id",
name="token-exchange",
server_name="token-exchange",
alias="token-exchange",
url="https://upstream.example/mcp",
transport="http",
auth_type=MCPAuth("oauth2_token_exchange"),
)
with (
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as manager,
):
manager.get_mcp_server_by_identifier.side_effect = lambda identifier, client_ip=None: {
"alpha": alpha,
"alpha-id": alpha,
"beta": beta,
"token-exchange": token_exchange,
}.get(identifier.lower())
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(
self._scope(extra_headers=((b"x-mcp-servers", header.encode()),))
)
expected: Final = (
f"Bearer {self._EXPECTED_RESOURCE_METADATA}"
if expected_scope is None
else f"Bearer {self._EXPECTED_RESOURCE_METADATA}, {expected_scope}"
)
assert (exc_info.value.headers or {})["WWW-Authenticate"] == expected
assert exc_info.value.status_code == 401
@pytest.mark.parametrize(
"auth_type",
(None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token", "oauth2"),
@ -7898,6 +7965,7 @@ class TestAggregateGatewayDcrChallenge:
) as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = resolved
mock_mgr.get_mcp_server_by_identifier.return_value = resolved
with pytest.raises(ProxyException):
await MCPRequestHandler.process_mcp_request(
self._scope(path=path, extra_headers=((b"authorization", b"Bearer not-a-key"),))
@ -7942,12 +8010,12 @@ class TestAggregateGatewayDcrChallenge:
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr:
mock_mgr.get_mcp_server_by_name.return_value = resolved
mock_mgr.get_mcp_server_by_identifier.return_value = resolved
assert _gateway_dcr_challenge_target("/mcp/srv", None, None) == expected, resolved
assert _gateway_dcr_challenge_target("/mcp/a,b", None, None) is None
assert _gateway_dcr_challenge_target("/mcp", None, None) is None
with patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr:
mock_mgr.get_mcp_server_by_name.return_value = _server(MCPAuth.oauth2)
mock_mgr.get_mcp_server_by_identifier.return_value = _server(MCPAuth.oauth2)
assert _gateway_dcr_challenge_target("/mcp/srv", ["other"], None) is None
async def test_no_challenge_for_path_named_server(self):
@ -10658,7 +10726,7 @@ class TestScopedSessionAdmission:
_MASTER_KEY = "sk-scoped-session-admission-master-key"
def _bearer(self, resource_server_id):
def _bearer(self, resource_server_ids):
from datetime import datetime, timezone
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
@ -10671,11 +10739,11 @@ class TestScopedSessionAdmission:
keys = session_keys_from_master_key(self._MASTER_KEY)
principal = SessionPrincipal(
user_id="scoped-user", client_id="llm_dcrc_abc", resource_server_id=resource_server_id
user_id="scoped-user", client_id="llm_dcrc_abc", resource_server_ids=resource_server_ids
)
return mint_session_token(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
@pytest.mark.parametrize("scope", ["github-server-id", None])
@pytest.mark.parametrize("scope", [("github-server-id",), None])
async def test_admission_carries_sealed_resource_scope(self, scope):
token = self._bearer(scope)
scope_dict = {
@ -10708,11 +10776,11 @@ class TestScopedSessionAdmission:
):
auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope_dict)
assert auth_result.mcp_admitted_user_subject is True
assert auth_result.mcp_session_resource_server_id == scope
assert auth_result.mcp_session_resource_server_ids == scope
def test_scope_field_cannot_be_forged_through_construction(self):
forged = UserAPIKeyAuth(user_id="u1", mcp_session_resource_server_id="any-server")
assert forged.mcp_session_resource_server_id is None
forged = UserAPIKeyAuth(user_id="u1", mcp_session_resource_server_ids=("any-server",))
assert forged.mcp_session_resource_server_ids is None
@pytest.mark.asyncio
@ -10789,13 +10857,13 @@ async def test_catalog_refresh_preserves_non_database_admission_and_resource_sco
caller = UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS if kind == "master" else "custom-subject")
caller.via_virtual_key = kind == "master"
caller.authenticated_by_custom_auth = kind == "custom"
caller.mcp_session_resource_server_id = "only-this-server"
caller.mcp_session_resource_server_ids = ("only-this-server",)
caller.mcp_toolset_id = "only-this-toolset"
refreshed = await MCPRequestHandler.refresh_catalog_authority(caller)
assert refreshed is not caller
assert refreshed.api_key == caller.api_key
assert refreshed.authenticated_by_custom_auth == caller.authenticated_by_custom_auth
assert refreshed.mcp_session_resource_server_id == "only-this-server"
assert refreshed.mcp_session_resource_server_ids == ("only-this-server",)
assert refreshed.mcp_toolset_id == "only-this-toolset"
assert refreshed.requires_fresh_policy is True
assert caller.requires_fresh_policy is False
@ -10818,11 +10886,11 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
caller = UserAPIKeyAuth(user_id="catalog-user", org_id="previous-org", user_role="proxy_admin")
caller.mcp_admitted_user_subject = True
caller.mcp_session_resource_server_id = "scoped-server"
caller.mcp_session_resource_server_ids = ("scoped-server",)
refreshed = await MCPRequestHandler.refresh_catalog_authority(caller)
assert refreshed.org_id == "current-org"
assert refreshed.user_role == "internal_user"
assert refreshed.mcp_session_resource_server_id == "scoped-server"
assert refreshed.mcp_session_resource_server_ids == ("scoped-server",)
assert refreshed.mcp_admitted_user_subject is True
assert caller.org_id == "previous-org"
@ -10851,7 +10919,7 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
access_group_ids=["original-group"],
)
caller.via_virtual_key = True
caller.mcp_session_resource_server_id = "session-server"
caller.mcp_session_resource_server_ids = ("session-server",)
caller.mcp_toolset_id = "session-toolset"
refreshed = await MCPRequestHandler.refresh_catalog_authority(caller)
reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True)
@ -10863,7 +10931,7 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
"new-project",
"new-owner",
)
assert refreshed.mcp_session_resource_server_id == "session-server"
assert refreshed.mcp_session_resource_server_ids == ("session-server",)
assert refreshed.mcp_toolset_id == "session-toolset"
assert refreshed.via_virtual_key and refreshed.requires_fresh_policy
assert caller.team_id == "old-team" and not caller.requires_fresh_policy

View file

@ -219,6 +219,8 @@ def test_empty_principal_fields_rejected_at_construction():
SessionPrincipal(user_id="", client_id="c")
with pytest.raises(ValidationError):
SessionPrincipal(user_id="u", client_id="")
with pytest.raises(ValidationError):
SessionPrincipal(user_id="u", client_id="c", resource_server_ids=())
def test_short_signing_key_rejected_at_construction():
@ -248,6 +250,46 @@ def _decoded_claims(token: str, prefix: str) -> dict:
)
def test_single_server_scope_keeps_the_scalar_wire_claim():
principal = SessionPrincipal(
user_id="user-123",
client_id="llm_client_abc",
resource_server_ids=("alpha-id",),
)
minted = mint_session_token(principal, KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
claims = _decoded_claims(minted.token.get_secret_value(), SESSION_TOKEN_PREFIX)
opened = open_session_token(minted.token.get_secret_value(), KEYS, NOW)
assert claims["resource_server_id"] == "alpha-id"
assert "resource_server_ids" not in claims
assert isinstance(opened, OpenedSessionToken)
assert opened.principal.resource_server_ids == ("alpha-id",)
def test_multi_server_scope_round_trips_through_the_list_claim():
principal = SessionPrincipal(
user_id="user-123",
client_id="llm_client_abc",
resource_server_ids=("alpha-id", "beta-id"),
)
minted = mint_session_token(principal, KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
claims = _decoded_claims(minted.token.get_secret_value(), SESSION_TOKEN_PREFIX)
opened = open_session_token(minted.token.get_secret_value(), KEYS, NOW)
assert claims["resource_server_ids"] == ["alpha-id", "beta-id"]
assert "resource_server_id" not in claims
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == principal
def test_token_with_both_server_scope_claims_is_malformed():
token = _sign_claims(_valid_claims(resource_server_id="alpha-id", resource_server_ids=["alpha-id", "beta-id"]))
assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed)
def test_mcp_principal_wire_claims_carry_no_audience_or_team_keys():
access_claims = _decoded_claims(_mint_access(), SESSION_TOKEN_PREFIX)
refresh_claims = _decoded_claims(_mint_refresh(), SESSION_REFRESH_PREFIX)

View file

@ -22,7 +22,7 @@ def test_guardrail_bearer_stays_out_of_egress_and_context_representations() -> N
def test_operation_context_isolates_nested_headers_and_caller_permissions():
caller = UserAPIKeyAuth(user_id="alpha", models=["allowed"])
caller.mcp_admitted_user_subject = True
caller.mcp_session_resource_server_id = "alpha-server"
caller.mcp_session_resource_server_ids = ("alpha-server",)
caller.mcp_toolset_id = "toolset-alpha"
caller.mcp_source_team_rpm_limits = {"team": {"alpha-server": 2}}
headers = {"x-caller": "alpha"}
@ -37,7 +37,7 @@ def test_operation_context_isolates_nested_headers_and_caller_permissions():
assert captured is not None
assert captured.models == ["allowed"]
assert captured.mcp_admitted_user_subject is True
assert captured.mcp_session_resource_server_id == "alpha-server"
assert captured.mcp_session_resource_server_ids == ("alpha-server",)
assert captured.mcp_toolset_id == "toolset-alpha"
assert captured.mcp_source_team_rpm_limits == {"team": {"alpha-server": 2}}
captured.models.append("also-forbidden")

View file

@ -11,6 +11,7 @@ from urllib.parse import parse_qs, urlparse
import pytest
from starlette.requests import Request
from starlette.responses import Response
from litellm.caching.caching import DualCache
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
@ -27,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
MintedProxyCredential,
SubjectIdentity,
SubjectTokenRefusal,
_ConnectFlow,
_GatewayAuthCode,
_open_sealed,
_seal,
@ -57,6 +59,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
mint_session_refresh_token,
mint_session_token,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
MASTER_KEY = "sk-gateway-dcr-flow-tests"
REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback"
@ -264,10 +267,17 @@ async def test_tampered_client_id_does_not_open():
def _authorize(
client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code"
client_id,
session_user_id,
redirect_uri=REDIRECT_URI,
challenge=CODE_CHALLENGE,
method="S256",
response_type="code",
scope: str | None = None,
):
query: Final = f"client_id={client_id}" if scope is None else f"client_id={client_id}&scope={scope}"
return aggregate_authorize(
request=_request(query=f"client_id={client_id}"),
request=_request(query=query),
client_id=client_id,
redirect_uri=redirect_uri,
state="client-state-123",
@ -275,6 +285,7 @@ def _authorize(
code_challenge_method=method,
response_type=response_type,
session_user_id=session_user_id,
scope=scope,
)
@ -295,11 +306,12 @@ async def test_authorize_validation_failures_never_redirect_to_client():
@pytest.mark.asyncio
async def test_authorize_without_session_redirects_to_login_with_return_to():
client_id = (await _register([REDIRECT_URI]))["client_id"]
response = _authorize(client_id, session_user_id=None)
response = _authorize(client_id, session_user_id=None, scope="litellm:mcp_server:missing")
assert response.status_code == 303
location = response.headers["location"]
assert location.startswith("https://llm.example.com/sso/key/generate?return_to=")
assert "return_to=%2Fauthorize" in location
assert "scope=litellm:mcp_server:missing" in parse_qs(urlparse(location).query)["return_to"][0]
@pytest.mark.asyncio
@ -991,7 +1003,12 @@ SCOPED_RESOURCE = "https://llm.example.com/mcp/github"
_MANAGER_PATCH = "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs):
def _scoped_authorize(
client_id: str,
resource: str | None,
session_user_id: str | None = "u1",
scope: str | None = None,
) -> Response:
return aggregate_authorize(
request=_request(query=f"client_id={client_id}"),
client_id=client_id,
@ -1002,7 +1019,7 @@ def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs):
response_type="code",
session_user_id=session_user_id,
resource=resource,
**kwargs,
scope=scope,
)
@ -1094,6 +1111,105 @@ def _sealed_wire_json(sealed, prefix, debug_key):
return json.loads(raw)
def test_legacy_server_scope_payloads_open_as_single_server_ids():
flow_payload: Final = {
"user_id": "u1",
"client_id": "client-1",
"redirect_uri": REDIRECT_URI,
"state": "state",
"code_challenge": CODE_CHALLENGE,
"jti": "flow-jti",
"exp": 2,
"resource_server_id": "server-1",
}
encrypted_flow_payload: Final = encrypt_value_helper(json.dumps(flow_payload))
assert isinstance(encrypted_flow_payload, str)
opened_flow: Final = _open_sealed(encrypted_flow_payload, "", _ConnectFlow, "gateway_connect_flow")
assert opened_flow is not None
assert opened_flow.resource_server_ids == ("server-1",)
rejected_flow_payload: Final = encrypt_value_helper(
json.dumps(flow_payload | {"resource_server_ids": ["server-1"]})
)
assert isinstance(rejected_flow_payload, str)
assert _open_sealed(rejected_flow_payload, "", _ConnectFlow, "gateway_connect_flow") is None
code_payload: Final = {
"user_id": "u1",
"client_id": "client-1",
"redirect_uri": REDIRECT_URI,
"code_challenge": CODE_CHALLENGE,
"jti": "code-jti",
"iat": 1,
"exp": 2,
"resource_server_id": "server-1",
}
encrypted_code_payload: Final = encrypt_value_helper(json.dumps(code_payload))
assert isinstance(encrypted_code_payload, str)
opened_code: Final = _open_sealed(
f"{GATEWAY_AUTH_CODE_PREFIX}{encrypted_code_payload}",
GATEWAY_AUTH_CODE_PREFIX,
_GatewayAuthCode,
_AUTH_CODE_DEBUG_KEY,
)
assert opened_code is not None
assert opened_code.resource_server_ids == ("server-1",)
rejected_code_payload: Final = encrypt_value_helper(
json.dumps(code_payload | {"resource_server_ids": ["server-1"]})
)
assert isinstance(rejected_code_payload, str)
assert (
_open_sealed(
f"{GATEWAY_AUTH_CODE_PREFIX}{rejected_code_payload}",
GATEWAY_AUTH_CODE_PREFIX,
_GatewayAuthCode,
_AUTH_CODE_DEBUG_KEY,
)
is None
)
def test_single_server_flow_and_code_payloads_keep_the_legacy_wire_key():
flow_payload: Final = _sealed_wire_json(
_seal(
"",
_ConnectFlow(
user_id="u1",
client_id="client-1",
redirect_uri=REDIRECT_URI,
state="state",
code_challenge=CODE_CHALLENGE,
jti="flow-jti",
exp=2,
resource_server_ids=("server-1",),
),
),
"",
"gateway_connect_flow",
)
assert flow_payload["resource_server_id"] == "server-1"
assert "resource_server_ids" not in flow_payload
code_payload: Final = _sealed_wire_json(
_seal(
GATEWAY_AUTH_CODE_PREFIX,
_GatewayAuthCode(
user_id="u1",
client_id="client-1",
redirect_uri=REDIRECT_URI,
code_challenge=CODE_CHALLENGE,
jti="code-jti",
iat=1,
exp=2,
resource_server_ids=("server-1",),
),
),
GATEWAY_AUTH_CODE_PREFIX,
_AUTH_CODE_DEBUG_KEY,
)
assert code_payload["resource_server_id"] == "server-1"
assert "resource_server_ids" not in code_payload
@pytest.mark.asyncio
async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
"""LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2
@ -1106,15 +1222,15 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
client_id = (await _register([REDIRECT_URI]))["client_id"]
github = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = github
manager.get_mcp_server_by_identifier.return_value = github
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
assert response.status_code == 303
location = urlparse(response.headers["location"])
assert location.path == "/ui/connect"
assert set(parse_qs(location.query)) == {"connect_flow"}
_, cookies = _flow_cookie_from(response)
assert (
_sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id"
assert _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == (
"github-id"
)
described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent"))
assert json.loads(described.body) == {
@ -1123,6 +1239,7 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
"server_id": "github-id",
"server_name": "github",
"connected": False,
"servers": None,
}
cache = DualCache()
premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache)
@ -1133,17 +1250,233 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
assert completed.status_code == 303
assert present.calls == [("u1", "github-id")]
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
assert (
_sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"]
== "github-id"
assert _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] == (
"github-id"
)
token_response = await _redeem(code, client_id)
assert token_response.status_code == 200
principal = _opened_principal(json.loads(token_response.body))
assert principal.resource_server_id == "github-id"
assert principal.resource_server_ids == ("github-id",)
assert principal.user_id == "u1"
@pytest.mark.asyncio
@pytest.mark.parametrize("resource", ["https://llm.example.com/mcp", "https://llm.example.com"])
async def test_aggregate_scope_seals_server_through_redemption_and_refresh(resource: str) -> None:
from unittest.mock import patch
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
github: Final = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.return_value = github
response: Final = _scoped_authorize(
client_id,
resource,
scope="litellm:mcp_server:github",
)
assert response.status_code == 303
_, cookies = _flow_cookie_from(response)
sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert sealed_flow["resource_server_id"] == "github-id"
code: Final = await _finish_connect_page(response, scoped_server=github)
cache: Final = DualCache()
token_response: Final = await _redeem(
code,
client_id,
cache=cache,
resource=resource,
)
assert token_response.status_code == 200
payload: Final = json.loads(token_response.body)
assert _opened_principal(payload).resource_server_ids == ("github-id",)
rotated: Final = await _redeem(
None,
client_id,
cache=cache,
grant_type="refresh_token",
refresh_token=payload["refresh_token"],
resource=resource,
)
assert rotated.status_code == 200
assert _opened_principal(json.loads(rotated.body)).resource_server_ids == ("github-id",)
@pytest.mark.asyncio
async def test_multi_server_scope_seals_set_through_consent_redemption_and_refresh() -> None:
from unittest.mock import patch
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
alpha: Final = _scoped_mcp_server(name="alpha")
beta: Final = _scoped_mcp_server(name="beta", oauth2_flow="client_credentials")
servers_by_name: Final = {"alpha": alpha, "beta": beta}
servers_by_id: Final = {alpha.server_id: alpha, beta.server_id: beta}
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.side_effect = servers_by_name.get
response: Final = _scoped_authorize(
client_id,
"https://llm.example.com/mcp",
scope="litellm:mcp_server:alpha litellm:mcp_server:beta",
)
assert response.status_code == 303
handle, cookies = _flow_cookie_from(response)
sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert sealed_flow["resource_server_ids"] == ["alpha-id", "beta-id"]
vendor: Final = _VendorCredential("absent")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.side_effect = servers_by_id.get
described: Final = await describe_connect_flow(
request=_request("/authorize/flow", cookies=cookies),
flow_handle=handle,
session_user_id="u1",
lookup_vendor_credential=vendor,
lookup_server_reachability=_ServerReachability(),
)
assert json.loads(described.body) == {
"state": "multi",
"client_origin": "https://claude.ai",
"server_id": None,
"server_name": None,
"connected": None,
"servers": [
{"server_id": "alpha-id", "server_name": "alpha", "connected": False},
{"server_id": "beta-id", "server_name": "beta", "connected": True},
],
}
assert vendor.calls == [("u1", "alpha-id")]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.side_effect = servers_by_id.get
blocked: Final = await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id="u1",
cache=DualCache(),
lookup_vendor_credential=_VendorCredential("absent"),
lookup_server_reachability=_ServerReachability(),
)
assert blocked.status_code == 400
assert "authorize the requested MCP server" in json.loads(blocked.body)["error_description"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.side_effect = servers_by_id.get
completed: Final = await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id="u1",
cache=DualCache(),
lookup_vendor_credential=_VendorCredential("present"),
lookup_server_reachability=_ServerReachability(),
)
assert completed.status_code == 303
code: Final = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
sealed_code: Final = _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
assert sealed_code["resource_server_ids"] == ["alpha-id", "beta-id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.return_value = alpha
conflict: Final = await _redeem(code, client_id, resource="https://llm.example.com/mcp/alpha")
assert json.loads(conflict.body)["error"] == "invalid_target"
cache: Final = DualCache()
token_response: Final = await _redeem(
code,
client_id,
cache=cache,
resource="https://llm.example.com/mcp",
)
assert token_response.status_code == 200
payload: Final = json.loads(token_response.body)
assert _opened_principal(payload).resource_server_ids == ("alpha-id", "beta-id")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.return_value = alpha
refresh_conflict: Final = await _redeem(
None,
client_id,
cache=cache,
grant_type="refresh_token",
refresh_token=payload["refresh_token"],
resource="https://llm.example.com/mcp/alpha",
)
assert json.loads(refresh_conflict.body)["error"] == "invalid_target"
refreshed: Final = await _redeem(
None,
client_id,
cache=cache,
grant_type="refresh_token",
refresh_token=payload["refresh_token"],
resource="https://llm.example.com/mcp",
)
assert refreshed.status_code == 200
assert _opened_principal(json.loads(refreshed.body)).resource_server_ids == ("alpha-id", "beta-id")
@pytest.mark.asyncio
async def test_per_server_resource_takes_precedence_over_conflicting_scope() -> None:
from unittest.mock import patch
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
github: Final = _scoped_mcp_server()
linear: Final = _scoped_mcp_server(name="linear")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.side_effect = (github, linear)
response: Final = _scoped_authorize(
client_id,
SCOPED_RESOURCE,
scope="litellm:mcp_server:linear",
)
_, cookies = _flow_cookie_from(response)
sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert sealed_flow["resource_server_id"] == "github-id"
@pytest.mark.asyncio
async def test_unknown_server_scope_redirects_with_invalid_scope_and_state() -> None:
from unittest.mock import patch
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.return_value = None
response: Final = _scoped_authorize(
client_id,
"https://llm.example.com/mcp",
session_user_id="u1",
scope="openid litellm:mcp_server:missing",
)
assert response.status_code == 303
location: Final = urlparse(response.headers["location"])
assert location.scheme == "https"
assert location.netloc == "claude.ai"
assert location.path == "/api/mcp/auth_callback"
assert parse_qs(location.query) == {
"error": ["invalid_scope"],
"error_description": ["unknown or unsupported MCP server scope"],
"state": ["client-state-123"],
}
@pytest.mark.parametrize(
("server_name", "expected"),
[
("alpha", "litellm:mcp_server:alpha"),
("server~name", "litellm:mcp_server:server~name"),
("", None),
("bad name", None),
('bad"name', None),
("naïve", None),
],
)
def test_gateway_server_scope_accepts_only_rfc6749_scope_tokens(server_name: str, expected: str | None) -> None:
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import gateway_server_scope
assert gateway_server_scope(server_name) == expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
"resource, resolves",
@ -1168,14 +1501,14 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource,
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server()
manager.get_mcp_server_by_identifier.return_value = None if resolves is None else _scoped_mcp_server()
response = _scoped_authorize(client_id, resource)
assert response.status_code == 303
location = urlparse(response.headers["location"])
assert location.path == "/ui/connect"
assert set(parse_qs(location.query)) == {"connect_flow"}
_, cookies = _flow_cookie_from(response)
assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert "resource_server_ids" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
vendor = _VendorCredential("absent")
described = await _describe_page(response, vendor=vendor)
assert json.loads(described.body)["state"] == "unscoped"
@ -1183,13 +1516,14 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource,
completed = await _complete_page(response, vendor=vendor)
assert vendor.calls == []
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
assert "resource_server_ids" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
token_response = await _redeem(code, client_id)
payload = json.loads(token_response.body)
assert _opened_principal(payload).resource_server_id is None
assert _opened_principal(payload).resource_server_ids is None
jwt_payload_segment = payload["access_token"].removeprefix("llm_session_").split(".")[1]
claims = json.loads(base64.urlsafe_b64decode(jwt_payload_segment + "=" * (-len(jwt_payload_segment) % 4)))
assert "resource_server_id" not in claims
assert "resource_server_ids" not in claims
@pytest.mark.asyncio
@ -1201,7 +1535,7 @@ async def test_scoped_authorize_delegate_server_stays_unscoped():
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True)
manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True)
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
location = urlparse(response.headers["location"])
assert location.path == "/ui/connect"
@ -1209,7 +1543,7 @@ async def test_scoped_authorize_delegate_server_stays_unscoped():
assert json.loads((await _describe_page(response)).body)["server_id"] is None
code = await _finish_connect_page(response)
token_response = await _redeem(code, client_id)
assert _opened_principal(json.loads(token_response.body)).resource_server_id is None
assert _opened_principal(json.loads(token_response.body)).resource_server_ids is None
@pytest.mark.asyncio
@ -1221,7 +1555,7 @@ async def test_m2m_scoped_flow_mints_without_a_user_credential():
client_id = (await _register([REDIRECT_URI]))["client_id"]
m2m = _scoped_mcp_server(oauth2_flow="client_credentials")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = m2m
manager.get_mcp_server_by_identifier.return_value = m2m
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
vendor = _VendorCredential("unavailable")
described = await _describe_page(response, scoped_server=m2m, vendor=vendor)
@ -1241,7 +1575,7 @@ async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow):
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server(oauth2_flow=oauth2_flow)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
manager.get_mcp_server_by_identifier.return_value = server
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
reachable = _ServerReachability(False)
vendor = _VendorCredential("present")
@ -1252,6 +1586,7 @@ async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow):
"server_id": None,
"server_name": None,
"connected": None,
"servers": None,
}
cache = DualCache()
refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache)
@ -1270,7 +1605,7 @@ async def test_stale_scoped_flow_remains_distinct_from_unscoped():
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server()
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
described = await _describe_page(response, scoped_server=None)
assert json.loads(described.body)["state"] == "stale"
@ -1290,7 +1625,7 @@ async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credentia
client_id = (await _register([REDIRECT_URI]))["client_id"]
github = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = github
manager.get_mcp_server_by_identifier.return_value = github
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
cache = DualCache()
skipped_reachability = _ServerReachability(False)
@ -1350,22 +1685,22 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope():
github = _scoped_mcp_server()
linear = _scoped_mcp_server(name="linear")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = github
manager.get_mcp_server_by_identifier.return_value = github
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
code = await _finish_connect_page(response, scoped_server=github)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = linear
manager.get_mcp_server_by_identifier.return_value = linear
mismatched = await _redeem(code, client_id, resource="https://llm.example.com/mcp/linear")
assert json.loads(mismatched.body)["error"] == "invalid_target"
cache = DualCache()
token_response = await _redeem(code, client_id, cache=cache)
payload = json.loads(token_response.body)
assert _opened_principal(payload).resource_server_id == "github-id"
assert _opened_principal(payload).resource_server_ids == ("github-id",)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = linear
manager.get_mcp_server_by_identifier.return_value = linear
refresh_mismatch = await _redeem(
None,
client_id,
@ -1380,7 +1715,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope():
None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"]
)
assert rotated.status_code == 200
assert _opened_principal(json.loads(rotated.body)).resource_server_id == "github-id"
assert _opened_principal(json.loads(rotated.body)).resource_server_ids == ("github-id",)
@pytest.mark.asyncio
@ -1406,7 +1741,7 @@ async def test_resolve_scoped_resource_server_matrix():
(None, github, None),
]:
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = resolved_server
manager.get_mcp_server_by_identifier.return_value = resolved_server
result = resolve_scoped_resource_server(request, resource)
assert (result.server_id if result is not None else None) == expected, resource
@ -1422,10 +1757,10 @@ async def test_resource_resolution_is_identity_not_ip_filtered_access():
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import resolve_scoped_resource_server
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server()
result = resolve_scoped_resource_server(_request(), SCOPED_RESOURCE)
assert result is not None
manager.get_mcp_server_by_name.assert_called_once_with("github")
manager.get_mcp_server_by_identifier.assert_called_once_with("github")
LOOPBACK_REDIRECT_URI = "http://127.0.0.1:51234/callback"
@ -1560,7 +1895,7 @@ async def test_native_authorize_renders_consent_page_and_sets_flow_cookie():
assert flow["client_id"] == client_id
assert flow["redirect_uri"] == LOOPBACK_REDIRECT_URI
assert flow["user_id"] == "u1"
assert "resource_server_id" not in flow
assert "resource_server_ids" not in flow
assert handle not in body.replace(f'value="{handle}"', "")
@ -2214,6 +2549,34 @@ async def test_introspect_active_access_token_reports_rfc7662_claims():
assert body["jti"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("resource_server_ids", "claim", "expected", "other_claim"),
[
(("alpha-id",), "resource_server_id", "alpha-id", "resource_server_ids"),
(("alpha-id", "beta-id"), "resource_server_ids", ["alpha-id", "beta-id"], "resource_server_id"),
],
)
async def test_introspect_reports_server_scope_as_scalar_or_list(
resource_server_ids: tuple[str, ...],
claim: str,
expected: str | list[str],
other_claim: str,
):
keys, now, _ = _introspection_fixtures()
principal = SessionPrincipal(
user_id="u1",
client_id="llm_dcrc_client",
resource_server_ids=resource_server_ids,
)
minted = mint_session_token(principal, keys, now)
status, body = await _introspect(minted.token.get_secret_value())
assert status == 200
assert body[claim] == expected
assert other_claim not in body
@pytest.mark.asyncio
async def test_introspect_invalid_tokens_answer_active_false():
keys, now, principal = _introspection_fixtures()
@ -2380,7 +2743,7 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a
server = _scoped_mcp_server(auth_type=auth_type)
vendor = _VendorCredential("absent")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
manager.get_mcp_server_by_identifier.return_value = server
response = _scoped_authorize(client_id, resource)
described = await _describe_page(response, scoped_server=server, vendor=vendor)
assert json.loads(described.body) == {
@ -2389,6 +2752,7 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a
"server_id": "github-id",
"server_name": "github",
"connected": True,
"servers": None,
}
unreachable = await _complete_page(response, scoped_server=server, reachable=_ServerReachability(False))
assert unreachable.status_code == 400
@ -2398,16 +2762,16 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a
assert vendor.calls == []
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
manager.get_mcp_server_by_identifier.return_value = server
redeemed = await _redeem(code, client_id, cache=cache, resource=resource)
assert redeemed.status_code == 200
payload = json.loads(redeemed.body)
assert _opened_principal(payload).resource_server_id == "github-id"
assert _opened_principal(payload).resource_server_ids == ("github-id",)
renewed = await _redeem(
None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"]
)
assert renewed.status_code == 200
assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id"
assert _opened_principal(json.loads(renewed.body)).resource_server_ids == ("github-id",)
JWT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt"
@ -2594,7 +2958,7 @@ async def test_connect_flow_retains_step_up_scope_until_the_vendor_grants_it():
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
manager.get_mcp_server_by_identifier.return_value = server
response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope="tools.write")
vendor = AsyncMock(return_value="absent")
described = await _describe_page(response, scoped_server=server, vendor=vendor)
@ -2631,7 +2995,7 @@ async def test_aggregate_step_up_requires_the_upstream_grant_without_narrowing_g
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
tokens = await _redeem(code, client_id)
assert tokens.status_code == 200
assert _opened_principal(json.loads(tokens.body)).resource_server_id is None
assert _opened_principal(json.loads(tokens.body)).resource_server_ids is None
@pytest.mark.asyncio
@ -2710,7 +3074,7 @@ async def test_named_resource_rejects_another_upstreams_permission_before_consen
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
manager.get_mcp_server_by_identifier.return_value = _scoped_mcp_server()
response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope=encode_upstream_scope("other-id", "write"))
assert response.status_code == 400
assert json.loads(response.body)["error"] == "invalid_scope"
@ -2825,7 +3189,7 @@ async def test_sdk_completes_aggregate_step_up_discovery_consent_and_pkce(route)
tokens = await _redeem(cache=cache, **params)
assert tokens.status_code == 200
payload = json.loads(tokens.body)
assert _opened_principal(payload).resource_server_id is None
assert _opened_principal(payload).resource_server_ids is None
retry = await flow.asend(httpx2.Response(200, json=payload, request=request))
assert retry.url == initial.url
assert retry.headers["authorization"] == f"Bearer {payload['access_token']}"
@ -2854,7 +3218,7 @@ async def test_named_step_up_finish_describes_and_retains_its_single_server_ceil
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
manager.get_mcp_server_by_identifier.return_value = server
response = _scoped_authorize(
client_id, SCOPED_RESOURCE, scope=encode_upstream_scope(server.server_id, "tools.write")
)
@ -2868,4 +3232,67 @@ async def test_named_step_up_finish_describes_and_retains_its_single_server_ceil
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
tokens = await _redeem(code, client_id)
assert tokens.status_code == 200
assert _opened_principal(json.loads(tokens.body)).resource_server_id == server.server_id
assert _opened_principal(json.loads(tokens.body)).resource_server_ids == (server.server_id,)
@pytest.mark.asyncio
async def test_upstream_step_up_inside_multi_server_scope_describes_every_sealed_server() -> None:
from unittest.mock import AsyncMock, patch
from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
alpha: Final = _scoped_mcp_server(name="alpha")
beta: Final = _scoped_mcp_server(name="beta")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.side_effect = {"alpha": alpha, "beta": beta}.get
response: Final = _scoped_authorize(
client_id,
"https://llm.example.com/mcp",
scope=" ".join(
["litellm:mcp_server:alpha", "litellm:mcp_server:beta", encode_upstream_scope(alpha.server_id, "w")]
),
)
assert response.status_code == 303
handle, cookies = _flow_cookie_from(response)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.side_effect = {alpha.server_id: alpha, beta.server_id: beta}.get
described: Final = await describe_connect_flow(
request=_request("/authorize/flow", cookies=cookies),
flow_handle=handle,
session_user_id="u1",
lookup_vendor_credential=AsyncMock(return_value="present"),
lookup_server_reachability=_ServerReachability(),
)
body: Final = json.loads(described.body)
assert body["state"] == "multi"
assert [item["server_id"] for item in body["servers"]] == ["alpha-id", "beta-id"]
@pytest.mark.asyncio
async def test_unscoped_flow_cookie_keeps_the_legacy_wire_shape() -> None:
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
response: Final = _scoped_authorize(client_id, "https://llm.example.com/mcp")
assert response.status_code == 303
_, cookies = _flow_cookie_from(response)
sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert "resource_server_ids" not in sealed_flow
assert "resource_server_id" not in sealed_flow
@pytest.mark.asyncio
async def test_gateway_scope_still_narrows_when_the_resource_does_not_resolve() -> None:
from unittest.mock import patch
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
alpha: Final = _scoped_mcp_server(name="alpha")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_identifier.side_effect = {"alpha": alpha}.get
manager.get_mcp_server_by_name.return_value = None
response: Final = _scoped_authorize(
client_id, "https://llm.example.com/mcp/deleted-server", scope="litellm:mcp_server:alpha"
)
assert response.status_code == 303
_, cookies = _flow_cookie_from(response)
sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
assert sealed_flow["resource_server_id"] == "alpha-id"

View file

@ -379,6 +379,22 @@ class TestMCPServerManager:
with patch.object(manager, "_get_general_settings", return_value={}):
assert manager.get_mcp_server_by_id(server.server_id, client_ip="8.8.8.8") is None
def test_get_mcp_server_by_identifier_matches_alias_case_and_server_id(self):
manager = MCPServerManager()
server = MCPServer(
server_id="alpha-id",
name="alpha",
server_name="alpha",
alias="alpha-alias",
transport=MCPTransport.http,
)
manager.registry[server.server_id] = server
assert manager.get_mcp_server_by_identifier("ALPHA") is server
assert manager.get_mcp_server_by_identifier("alpha-alias") is server
assert manager.get_mcp_server_by_identifier("alpha-id") is server
assert manager.get_mcp_server_by_identifier("unknown") is None
async def test_create_mcp_client_stdio(self, monkeypatch):
"""Test creating MCP client for stdio transport"""
monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true")
@ -13765,7 +13781,7 @@ class TestSessionResourceScopeIntersect:
auth = UserAPIKeyAuth(user_id="scoped-user")
auth.mcp_admitted_user_subject = True
auth.mcp_session_resource_server_id = scope
auth.mcp_session_resource_server_ids = None if scope is None else (scope,)
return auth
def test_scope_reader_is_none_for_keys_and_unscoped_subjects(self):
@ -13779,7 +13795,14 @@ class TestSessionResourceScopeIntersect:
def test_scope_reader_returns_sealed_scope_for_admitted_subjects(self):
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
assert MCPServerManager._admitted_session_resource_scope(self._admitted_auth("b")) == "b"
assert MCPServerManager._admitted_session_resource_scope(self._admitted_auth("b")) == frozenset({"b"})
def test_scope_reader_returns_all_sealed_server_ids(self):
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
auth = self._admitted_auth("first")
auth.mcp_session_resource_server_ids = ("first", "second")
assert MCPServerManager._admitted_session_resource_scope(auth) == frozenset({"first", "second"})
@pytest.mark.asyncio
async def test_admin_registry_seed_still_bounded_by_session_resource_scope(self):
@ -13811,9 +13834,41 @@ class TestSessionResourceScopeIntersect:
),
):
assert await manager.get_allowed_mcp_servers(auth) == ["granted-id"]
auth.mcp_session_resource_server_id = None
auth.mcp_session_resource_server_ids = None
assert set(await manager.get_allowed_mcp_servers(auth)) == {"granted-id", "other-id"}
@pytest.mark.asyncio
async def test_admin_registry_seed_keeps_both_servers_in_a_multi_server_session(self):
from unittest.mock import AsyncMock, patch
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
from litellm.proxy._types import LitellmUserRoles
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
manager = MCPServerManager()
for server_id in ("alpha-id", "beta-id", "gamma-id"):
manager.registry[server_id] = MCPServer(
server_id=server_id,
name=server_id,
server_name=server_id,
url="https://example.com/mcp",
transport=MCPTransport.http,
)
auth = self._admitted_auth("alpha-id")
auth.mcp_session_resource_server_ids = ("alpha-id", "beta-id")
auth.user_role = LitellmUserRoles.PROXY_ADMIN
with (
patch.object(MCPServerManager, "get_allow_all_keys_server_ids", return_value=[]),
patch.object(
MCPServerManager,
"_get_active_submitted_mcp_server_ids_for_user",
new_callable=AsyncMock,
return_value=[],
),
):
assert set(await manager.get_allowed_mcp_servers(auth)) == {"alpha-id", "beta-id"}
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_scopes_past_operator_open_union(self):
"""The intersect applies AFTER the operator-open (allow_all_keys) union, so a scoped

View file

@ -182,13 +182,12 @@ class TestApplyToolsetScope:
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None)
admitted.mcp_admitted_user_subject = True
admitted.mcp_session_resource_server_id = "server-own"
admitted.mcp_session_resource_server_ids = ("server-own", "server-own-other")
admitted.requires_fresh_policy = True
granted = AsyncMock(return_value=frozenset({"toolset-123"}))
resolve = AsyncMock(return_value={"server-team": ["tool1"]})
with patch(
"litellm.proxy._experimental.mcp_server.server."
"global_mcp_server_manager.resolve_toolset_tool_permissions",
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.resolve_toolset_tool_permissions",
new=resolve,
):
with pytest.raises(HTTPException) as exc_info:
@ -203,18 +202,17 @@ class TestApplyToolsetScope:
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None)
admitted.mcp_admitted_user_subject = True
admitted.mcp_session_resource_server_id = "server-team"
admitted.mcp_session_resource_server_ids = ("server-team",)
granted = AsyncMock(return_value=frozenset({"toolset-123"}))
resolve = AsyncMock(return_value={"server-team": ["tool1"], "server-other": ["tool2"]})
with patch(
"litellm.proxy._experimental.mcp_server.server."
"global_mcp_server_manager.resolve_toolset_tool_permissions",
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.resolve_toolset_tool_permissions",
new=resolve,
):
result = await apply_toolset_scope(admitted, "toolset-123", granted=granted)
assert result.mcp_toolset_id == "toolset-123"
assert result.mcp_session_resource_server_id == "server-team"
assert result.mcp_session_resource_server_ids == ("server-team",)
resolve.assert_awaited_once_with(toolset_ids=["toolset-123"], requires_fresh_policy=False)
@pytest.mark.asyncio

View file

@ -1,6 +1,6 @@
import asyncio
from collections.abc import Callable, Sequence
from typing import Final
from typing import Final, Literal
from unittest.mock import AsyncMock, patch
import pytest
@ -44,6 +44,50 @@ async def _served_catalog_tool() -> str:
return "ok"
@pytest.mark.parametrize(
("settings", "mcp_servers", "toolset_id", "auth_state", "raises"),
[
({"mcp_require_explicit_server_scope": True}, None, None, None, True),
({"mcp_require_explicit_server_scope": True}, None, None, "unscoped", True),
({"mcp_require_explicit_server_scope": True}, ["alpha"], None, None, False),
({"mcp_require_explicit_server_scope": True}, [], None, None, True),
({"mcp_require_explicit_server_scope": True}, None, "toolset-id", None, False),
({"mcp_require_explicit_server_scope": True}, None, None, "sealed", False),
({"mcp_require_explicit_server_scope": False}, None, None, None, False),
({}, None, None, None, False),
],
)
def test_explicit_server_scope_setting_only_rejects_unscoped_aggregate_requests(
settings: dict[str, bool],
mcp_servers: list[str] | None,
toolset_id: str | None,
auth_state: Literal["unscoped", "sealed"] | None,
raises: bool,
) -> None:
from fastapi import HTTPException
auth: Final[UserAPIKeyAuth | None] = (
None if auth_state is None else UserAPIKeyAuth(api_key=None, user_id="scope-test")
)
if auth_state == "sealed" and auth is not None:
auth.mcp_session_resource_server_ids = ("alpha-id",)
with patch("litellm.proxy.proxy_server.general_settings", settings):
if raises:
with pytest.raises(HTTPException) as exc_info:
operations.raise_if_unscoped_aggregate_request(mcp_servers, auth, toolset_id)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == {
"error": "mcp_server_scope_required",
"message": (
"This gateway requires an explicit MCP server scope: connect to /mcp/<server_name> "
"or send the x-mcp-servers header."
),
}
else:
operations.raise_if_unscoped_aggregate_request(mcp_servers, auth, toolset_id)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["direct", "virtual", "proxy"])
@pytest.mark.parametrize("verdict", ["allow", "block", "obo_error", "evaluate_error", "missing"])

View file

@ -1,4 +1,5 @@
import json
from typing import Final
import pytest
from pydantic import ValidationError
@ -32,32 +33,39 @@ SERVER_ONLY_MARKERS = (
"billing_agent_policy",
"mcp_admitted_user_subject",
"mcp_source_team_rpm_limits",
"mcp_session_resource_server_id",
"mcp_session_resource_server_ids",
"via_virtual_key",
)
def _server_only_marker_value(marker: str, value: str) -> str | tuple[str, ...]:
return (value,) if marker == "mcp_session_resource_server_ids" else value
@pytest.mark.parametrize("marker", SERVER_ONLY_MARKERS)
def test_a_caller_cannot_forge_a_server_only_marker_through_the_constructor(marker):
auth = UserAPIKeyAuth(**{marker: "forged-by-caller"})
marker_value: Final = _server_only_marker_value(marker, "forged-by-caller")
auth = UserAPIKeyAuth(**{marker: marker_value})
assert getattr(auth, marker) != "forged-by-caller"
assert getattr(auth, marker) != marker_value
@pytest.mark.parametrize("marker", SERVER_ONLY_MARKERS)
def test_a_caller_cannot_forge_a_server_only_marker_through_model_validate(marker):
auth = UserAPIKeyAuth.model_validate({marker: "forged-by-caller"})
marker_value: Final = _server_only_marker_value(marker, "forged-by-caller")
auth = UserAPIKeyAuth.model_validate({marker: marker_value})
assert getattr(auth, marker) != "forged-by-caller"
assert getattr(auth, marker) != marker_value
@pytest.mark.parametrize("marker", SERVER_ONLY_MARKERS)
def test_the_server_sets_a_marker_by_assignment_after_construction(marker):
auth = UserAPIKeyAuth()
marker_value: Final = _server_only_marker_value(marker, "set-by-the-server")
setattr(auth, marker, "set-by-the-server")
setattr(auth, marker, marker_value)
assert getattr(auth, marker) == "set-by-the-server"
assert getattr(auth, marker) == marker_value
def test_a_virtual_key_is_hashed_out_of_the_auth_object():

View file

@ -266,7 +266,7 @@ def test_an_llm_backed_injection_check_needs_the_call_it_would_make():
[
("mcp_admitted_user_subject", "someone-else", False),
("mcp_source_team_rpm_limits", {"team-1": 10_000}, None),
("mcp_session_resource_server_id", "server-1", None),
("mcp_session_resource_server_ids", ("server-1",), None),
("via_virtual_key", "sk-someone-elses-key", False),
],
)

View file

@ -3,16 +3,19 @@ import { render, screen } from "@testing-library/react";
import type { ConnectFlowStatus } from "@/components/networking";
import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner";
const { startOAuthFlow } = vi.hoisted(() => ({ startOAuthFlow: vi.fn() }));
vi.mock("@/components/networking", () => ({
getProxyBaseUrl: () => "https://gateway.example.com",
}));
vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle" }),
useUserMcpOAuthFlow: () => ({ startOAuthFlow, status: "idle" }),
}));
afterEach(() => {
vi.restoreAllMocks();
startOAuthFlow.mockClear();
});
const unscoped = (client_origin: string): ConnectFlowStatus => ({
@ -21,6 +24,7 @@ const unscoped = (client_origin: string): ConnectFlowStatus => ({
server_id: null,
server_name: null,
connected: null,
servers: null,
});
const renderBanner = (clientOrigin: string) =>
@ -47,6 +51,53 @@ describe("ConnectFlowBanner", () => {
expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument();
});
it("lists multi-server consent and requires every server before finishing", () => {
const flow: ConnectFlowStatus = {
state: "multi",
client_origin: "https://claude.ai",
server_id: null,
server_name: null,
connected: null,
servers: [
{ server_id: "alpha-id", server_name: "alpha", connected: false },
{ server_id: "beta-id", server_name: "beta", connected: false },
],
};
const { rerender } = render(
<ConnectFlowBanner
flowHandle="flow-handle-123"
flow={flow}
accessToken="tok"
onConnected={vi.fn()}
failed={false}
/>,
);
expect(screen.getByText("Allow https://claude.ai to use alpha, beta")).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Connect alpha" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Connect beta" })).toBeInTheDocument();
expect(startOAuthFlow).not.toHaveBeenCalled();
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument();
rerender(
<ConnectFlowBanner
flowHandle="flow-handle-123"
flow={{
...flow,
servers: flow.servers?.map((server) => ({ ...server, connected: true })) ?? null,
}}
accessToken="tok"
onConnected={vi.fn()}
failed={false}
/>,
);
expect(
screen.getByText("Click Finish connecting to give https://claude.ai access to alpha, beta as you."),
).toBeInTheDocument();
});
it("offers manual delivery only for a loopback client, posted only when checked", () => {
const loopback = renderBanner("http://localhost:3118");
const checkbox = loopback.container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement;

View file

@ -27,6 +27,7 @@ export function isLoopbackOrigin(origin: string | null): boolean {
const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly [string, string] => {
const clientLabel = flow?.client_origin ?? "the application";
const serverLabel = flow?.server_name ?? "the requested MCP server";
const serverNames = flow?.servers?.map((server) => server.server_name).join(", ") ?? "";
if (failed || flow === undefined || flow.state === "stale") {
return [
"The connection cannot continue",
@ -39,6 +40,16 @@ const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly
`Authorize the servers you want to use below, then click Finish connecting to return to ${clientLabel}.`,
];
}
if (flow.state === "multi") {
const allServersConnected =
flow.servers !== null && flow.servers.length > 0 && flow.servers.every((server) => server.connected);
return [
`Allow ${clientLabel} to use ${serverNames}`,
allServersConnected
? `Click Finish connecting to give ${clientLabel} access to ${serverNames} as you.`
: `Authorize each requested server below to continue, or cancel to send ${clientLabel} away.`,
];
}
if (flow.state === "interactive" && !flow.connected) {
return [
`Allow ${clientLabel} to use ${serverLabel}`,
@ -51,16 +62,33 @@ const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly
];
};
const canFinishConnectFlow = (state: ConnectFlowStatus["state"], flow: ConnectFlowStatus | undefined): boolean => {
if (state === "unscoped") return true;
if (state === "stale" || flow === undefined) return false;
if (state !== "multi") return flow.connected === true;
return flow.servers !== null && flow.servers.length > 0 && flow.servers.every((server) => server.connected);
};
const vendorServersFor = (
state: ConnectFlowStatus["state"],
flow: ConnectFlowStatus | undefined,
): { server_id: string; server_name: string | null }[] => {
if (state === "multi") {
return (flow?.servers ?? [])
.filter((server) => !server.connected)
.map(({ server_id, server_name }) => ({ server_id, server_name }));
}
if (state !== "interactive" || flow?.connected !== false || flow.server_id === null) return [];
return [{ server_id: flow.server_id, server_name: flow.server_name }];
};
const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, flow, accessToken, onConnected, failed }) => {
const action = `${getProxyBaseUrl()}/authorize/complete`;
const state = failed || flow === undefined ? "stale" : flow.state;
const canFinish = state === "unscoped" || (state !== "stale" && flow?.connected === true);
const canFinish = canFinishConnectFlow(state, flow);
const canCancel = state !== "unscoped";
const loopbackClient = isLoopbackOrigin(flow?.client_origin ?? null);
const vendorServer =
state === "interactive" && flow?.connected === false && flow.server_id !== null
? { server_id: flow.server_id, server_name: flow.server_name }
: null;
const vendorServers = vendorServersFor(state, flow);
const copy = copyFor(flow, failed);
return (
@ -74,17 +102,20 @@ const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, flow, accessToken, onC
</div>
</div>
<div className="flex shrink-0 gap-2">
{vendorServer !== null && (
{vendorServers.map((server) => (
<OAuth2ConnectButton
key={vendorServer.server_id}
server={vendorServer}
key={server.server_id}
server={server}
accessToken={accessToken}
onConnect={onConnected}
scopes={flow?.requested_scopes}
variant="button"
autoStartKey={`litellm-mcp-autostart:${flowHandle}:${vendorServer.server_id}`}
buttonLabel={state === "multi" ? `Connect ${server.server_name ?? server.server_id}` : undefined}
{...(state === "multi"
? {}
: { autoStartKey: `litellm-mcp-autostart:${flowHandle}:${server.server_id}` })}
/>
)}
))}
<form method="POST" action={action}>
<input type="hidden" name="flow" value={flowHandle} />
{canFinish && (

View file

@ -32,12 +32,19 @@ vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
},
}));
const flow = (state: "unscoped" | "interactive" | "m2m" | "stale", connected: boolean | null = null) => ({
const flow = (state: "unscoped" | "interactive" | "m2m" | "multi" | "stale", connected: boolean | null = null) => ({
state,
client_origin: "https://claude.ai",
server_id: state === "interactive" || state === "m2m" ? "s-design" : null,
server_name: state === "interactive" || state === "m2m" ? "design_tool" : null,
connected,
servers:
state === "multi"
? [
{ server_id: "s-alpha", server_name: "alpha", connected: false },
{ server_id: "s-beta", server_name: "beta", connected: false },
]
: null,
});
const renderSurface = () =>

View file

@ -32,6 +32,7 @@ interface OAuth2ConnectButtonProps {
variant?: "badge" | "button";
autoStartKey?: string | null;
scopes?: string[];
buttonLabel?: string;
}
export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
@ -41,6 +42,7 @@ export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
variant = "badge",
autoStartKey = null,
scopes,
buttonLabel = "Connect",
}) => {
const name = server.server_name ?? server.alias ?? server.server_id;
const flowOptions = {
@ -64,7 +66,7 @@ export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
return (
<Button onClick={startOAuthFlow} disabled={loading} className="font-semibold h-[38px] min-w-[110px]">
{loading && <Loader2 className="h-4 w-4 animate-spin mr-1.5" />}
{loading ? "Connecting\u2026" : "Connect"}
{loading ? "Connecting\u2026" : buttonLabel}
</Button>
);
}

View file

@ -4840,12 +4840,13 @@ export const fetchDiscoverableMCPServers = async (accessToken: string) => {
};
export interface ConnectFlowStatus {
state: "unscoped" | "interactive" | "m2m" | "stale";
state: "unscoped" | "interactive" | "m2m" | "multi" | "stale";
client_origin: string;
server_id: string | null;
server_name: string | null;
connected: boolean | null;
requested_scopes?: string[];
servers: { server_id: string; server_name: string; connected: boolean }[] | null;
}
/**

View file

@ -30051,6 +30051,12 @@ export interface components {
* @description When true, a gateway-managed OAuth2 MCP server whose authorization server advertises Client ID Metadata Document support identifies itself with the gateway's public metadata document URL even when that authorization server also offers dynamic client registration. Requires a public HTTPS PROXY_BASE_URL the authorization server can fetch. Default false: dynamic client registration is used whenever the authorization server offers it, and the metadata document only when it does not.
*/
mcp_prefer_client_id_metadata_document?: boolean | null;
/**
* Mcp Require Explicit Server Scope
* @description Require aggregate MCP requests to include x-mcp-servers unless a toolset or gateway OAuth session already scopes the server. Off by default.
* @default false
*/
mcp_require_explicit_server_scope: boolean;
/**
* Mcp Required Fields
* @description List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).