mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
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:
parent
0a17da59fd
commit
1ddaa54793
26 changed files with 2045 additions and 432 deletions
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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 && (
|
||||
|
|
|
|||
|
|
@ -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 = () =>
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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']).
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue