mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(mcp): challenge a single-server listing whose credential the upstream rejected (#50389)
This commit is contained in:
shrey
GitHub
parent
05c04d6261
commit
37144b5e27
10 changed files with 665 additions and 43 deletions
|
|
@ -826,6 +826,20 @@ async def prepare_tool_listing_headers(
|
|||
return server_auth_header, extra_headers, catalog_auth_header
|
||||
|
||||
|
||||
def challenge_rejected_credential(server: MCPServer, outcome: ServerOutcome, listed: Sequence[MCPServer]) -> None:
|
||||
"""Answer a listing of this one server with the gateway's 401 when the caller can fix the rejection."""
|
||||
from litellm.proxy.mcp_server.legacy_callbacks import get_listing_scope_response
|
||||
|
||||
scope_response: Final = get_listing_scope_response()
|
||||
if (
|
||||
scope_response is not None
|
||||
and len(listed) == 1
|
||||
and outcome.tag == "auth_required"
|
||||
and server.reauthorizes_in_client
|
||||
):
|
||||
scope_response.deny_reauth()
|
||||
|
||||
|
||||
async def get_filtered_server_tools(
|
||||
server: MCPServer,
|
||||
*,
|
||||
|
|
@ -943,14 +957,11 @@ async def get_filtered_server_tools(
|
|||
)
|
||||
return page.model_copy(update={"tools": filtered_tools}), ServerListOk(tool_count=len(filtered_tools))
|
||||
except MCPUpstreamAuthError as e:
|
||||
# Absorb so one unauthenticated server does not empty every other server's
|
||||
# tools. Surfacing the upstream 401 to the client as a re-auth challenge is
|
||||
# intentionally not done here: raising from this list handler cannot produce a
|
||||
# 401 + WWW-Authenticate (the MCP session manager serializes it as a JSON-RPC
|
||||
# error). Single-server routes surface it via the request-scope preemptive
|
||||
# check in _raise_preemptive_401_for_unauthenticated_servers instead.
|
||||
# Absorbed so peers keep their tools; a raise here would reach the client as a JSON-RPC error, not a 401
|
||||
verbose_logger.debug("MCP list_tools: omitting %s; it needs upstream auth", server.name)
|
||||
return ListToolsResult(tools=[]), classify_list_exception(e)
|
||||
outcome: Final = classify_list_exception(e)
|
||||
challenge_rejected_credential(server, outcome, allowed_mcp_servers)
|
||||
return ListToolsResult(tools=[]), outcome
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error getting tools from server %s: %s", server.name, e)
|
||||
return ListToolsResult(tools=[]), classify_list_exception(e)
|
||||
|
|
@ -1199,6 +1210,8 @@ async def list_gateway_catalog(
|
|||
update={"meta": {key: value for key, value in page_meta.items() if key != SERVER_OUTCOMES_META_KEY}}
|
||||
)
|
||||
except Exception as error:
|
||||
outcome: Final = classify_list_exception(error)
|
||||
challenge_rejected_credential(server, outcome, allowed)
|
||||
if cursor is not None:
|
||||
raise MCPError(
|
||||
code=INVALID_PARAMS, message="Upstream continuation failed; start a fresh listing"
|
||||
|
|
@ -1207,11 +1220,7 @@ async def list_gateway_catalog(
|
|||
request,
|
||||
(),
|
||||
None,
|
||||
{
|
||||
"litellm.ai/server_outcomes": {
|
||||
_aggregate_server_key(server): classify_list_exception(error).model_dump(mode="json")
|
||||
}
|
||||
},
|
||||
{"litellm.ai/server_outcomes": {_aggregate_server_key(server): outcome.model_dump(mode="json")}},
|
||||
)
|
||||
|
||||
pages, next_cursor, outcomes = await paginate_catalog(
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
"""Exceptions raised by the LiteLLM MCP proxy."""
|
||||
|
||||
from typing import Final
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.mcp_server.oauth_utils import encode_upstream_scope
|
||||
from litellm.proxy.mcp_server.oauth_utils import encode_upstream_scope, well_known_root_suffix
|
||||
|
||||
TOKEN_EXCHANGE_SUBJECT_ERROR: Final = (
|
||||
"Missing or invalid subject token. Authenticate with your identity provider and send its token "
|
||||
|
|
@ -101,12 +102,14 @@ class MCPUpstreamAuthError(Exception):
|
|||
challenge: str | None = self.www_authenticate
|
||||
if challenge is None and (self.status_code == 401 or self.required_scope is not None) and base_url:
|
||||
prefix: Final = base_url.rstrip("/")
|
||||
if aggregate_step_up:
|
||||
resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource"
|
||||
elif request_path and request_path.startswith(f"/{self.server_name}/mcp"):
|
||||
resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp"
|
||||
else:
|
||||
resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}"
|
||||
name: Final = quote(self.server_name, safe="/:@!$&'()*+,;=")
|
||||
resource_metadata_url: Final = (
|
||||
f"{prefix}/.well-known/oauth-protected-resource"
|
||||
if aggregate_step_up
|
||||
else f"{prefix}/.well-known/oauth-protected-resource{well_known_root_suffix()}/{name}/mcp"
|
||||
if request_path and request_path.startswith(f"/{self.server_name}/mcp")
|
||||
else f"{prefix}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp/{name}"
|
||||
)
|
||||
challenge = f'Bearer resource_metadata="{resource_metadata_url}"'
|
||||
scoped_challenge: Final = (
|
||||
(f'{challenge}, error="insufficient_scope"' if challenge else 'Bearer error="insufficient_scope"')
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.mcp_server.contracts import OperationContext
|
||||
from litellm.proxy.mcp_server.mcp_context import get_active_mcp_request_ctx
|
||||
from litellm.proxy.mcp_server.scope_challenge import (
|
||||
LISTING_METHODS,
|
||||
OAuthScopeResponse,
|
||||
is_tool_call_request,
|
||||
scope_response_from_request,
|
||||
|
|
@ -113,6 +114,12 @@ def get_scope_response() -> OAuthScopeResponse | None:
|
|||
return scope_response_from_request(context.request) if context is not None else None
|
||||
|
||||
|
||||
def get_listing_scope_response() -> OAuthScopeResponse | None:
|
||||
"""The response gate of the active request, when the SDK parsed that request as a listing."""
|
||||
context: Final = get_active_mcp_request_ctx()
|
||||
return get_scope_response() if context is not None and context.method in LISTING_METHODS else None
|
||||
|
||||
|
||||
async def record_upstream_tool_authorization(response: httpx2.Response, *, pending: OAuthScopeResponse | None) -> None:
|
||||
if pending is None or not pending.pending or not response.is_success or response.request.method != "POST":
|
||||
return
|
||||
|
|
|
|||
|
|
@ -19,18 +19,31 @@ SCOPE_RESPONSE_KEY: Final = "litellm.mcp.scope_response"
|
|||
_REQUEST: Final = TypeAdapter[dict[str, object]](dict[str, object])
|
||||
|
||||
|
||||
def is_tool_call_request(body: bytes) -> bool:
|
||||
LISTING_METHODS: Final = frozenset({"tools/list", "prompts/list", "resources/list", "resources/templates/list"})
|
||||
|
||||
|
||||
def _request_method(body: bytes) -> object:
|
||||
try:
|
||||
return _REQUEST.validate_json(body).get("method") == "tools/call"
|
||||
return _REQUEST.validate_json(body).get("method")
|
||||
except ValidationError:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
def is_tool_call_request(body: bytes) -> bool:
|
||||
return _request_method(body) == "tools/call"
|
||||
|
||||
|
||||
def is_listing_request(body: bytes) -> bool:
|
||||
method: Final = _request_method(body)
|
||||
return isinstance(method, str) and method in LISTING_METHODS
|
||||
|
||||
|
||||
class OAuthScopeResponse:
|
||||
def __init__(self, send: Send, base_url: str, request_path: str | None) -> None:
|
||||
def __init__(self, send: Send, base_url: str, request_path: str | None, reauthorize_as: str | None = None) -> None:
|
||||
self._send = send
|
||||
self._base_url = base_url
|
||||
self._request_path = request_path
|
||||
self._reauthorize_as = reauthorize_as
|
||||
self._ready = asyncio.Event()
|
||||
self._lock = asyncio.Lock()
|
||||
self._challenge: MCPUpstreamAuthError | None = None
|
||||
|
|
@ -44,7 +57,15 @@ class OAuthScopeResponse:
|
|||
self._ready.set()
|
||||
|
||||
def deny(self, error: MCPUpstreamAuthError) -> None:
|
||||
if self.pending and error.status_code == 403 and error.required_scope is not None:
|
||||
if error.status_code == 403 and error.required_scope is not None:
|
||||
self._refuse(error)
|
||||
|
||||
def deny_reauth(self) -> None:
|
||||
if self._reauthorize_as is not None:
|
||||
self._refuse(MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name=self._reauthorize_as))
|
||||
|
||||
def _refuse(self, error: MCPUpstreamAuthError) -> None:
|
||||
if self.pending:
|
||||
self._challenge = error
|
||||
self._ready.set()
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import contextvars
|
|||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import types
|
||||
from collections import Counter
|
||||
|
|
@ -81,9 +82,11 @@ from litellm.proxy.mcp_server.oauth_utils import (
|
|||
well_known_root_suffix,
|
||||
)
|
||||
from litellm.proxy.mcp_server.scope_challenge import (
|
||||
LISTING_METHODS,
|
||||
SCOPE_RESPONSE_KEY,
|
||||
OAuthScopeResponse,
|
||||
finish_scope_response,
|
||||
is_listing_request,
|
||||
is_tool_call_request,
|
||||
)
|
||||
from litellm.proxy.mcp_server.ui_session_utils import (
|
||||
|
|
@ -225,16 +228,13 @@ _SESSION_MANAGERS_INITIALIZED = False
|
|||
_INITIALIZATION_LOCK: Final = asyncio.Lock()
|
||||
|
||||
|
||||
def _jsonrpc_text_has_top_level_method(text: str) -> bool:
|
||||
"""Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at
|
||||
the root object's top level.
|
||||
def _jsonrpc_text_top_level_key_end(text: str, key: str) -> int | None:
|
||||
"""Where the root object's top-level ``key`` ends in a (possibly truncated) JSON-RPC envelope.
|
||||
|
||||
Used to tell a request/notification (carries ``method``) apart from a
|
||||
response (carries ``result``/``error`` and no top-level ``method``). A
|
||||
response payload can itself nest a ``method`` field, so only keys at the
|
||||
A response payload can itself nest a ``method`` field, so only keys at the
|
||||
root object's depth are inspected rather than searching the whole string.
|
||||
Returns ``True`` only when a top-level ``method`` key is positively found;
|
||||
truncation that hides it yields ``False``.
|
||||
Returns an index only when the key is positively found; truncation that
|
||||
hides it yields ``None``.
|
||||
"""
|
||||
depth = 0
|
||||
in_string = False
|
||||
|
|
@ -243,7 +243,7 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool:
|
|||
reading_key = False
|
||||
expect_key = False
|
||||
key_chars: list[str] = []
|
||||
for ch in text:
|
||||
for index, ch in enumerate(text):
|
||||
if in_string:
|
||||
if escaped:
|
||||
escaped = False
|
||||
|
|
@ -251,8 +251,8 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool:
|
|||
escaped = True
|
||||
elif ch == '"':
|
||||
in_string = False
|
||||
if reading_key and depth == 1 and "".join(key_chars) == "method":
|
||||
return True
|
||||
if reading_key and depth == 1 and "".join(key_chars) == key:
|
||||
return index + 1
|
||||
elif reading_key:
|
||||
key_chars.append(ch)
|
||||
continue
|
||||
|
|
@ -276,7 +276,32 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool:
|
|||
expect_key = bool(in_object) and in_object[-1]
|
||||
elif ch == ":":
|
||||
expect_key = False
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
def _jsonrpc_text_has_top_level_method(text: str) -> bool:
|
||||
"""Whether a (possibly truncated) JSON-RPC envelope has a top-level ``method`` key.
|
||||
|
||||
Used to tell a request/notification (carries ``method``) apart from a
|
||||
response (carries ``result``/``error`` and no top-level ``method``).
|
||||
"""
|
||||
return _jsonrpc_text_top_level_key_end(text, "method") is not None
|
||||
|
||||
|
||||
_JSONRPC_STRING_VALUE: Final = re.compile(r'\s*:\s*("(?:[^"\\]|\\.)*")')
|
||||
|
||||
|
||||
def _jsonrpc_text_method(text: str) -> str | None:
|
||||
"""The top-level ``method`` of a (possibly truncated) JSON-RPC envelope, when the cut leaves it whole."""
|
||||
key_end: Final = _jsonrpc_text_top_level_key_end(text, "method")
|
||||
value: Final = _JSONRPC_STRING_VALUE.match(text, key_end) if key_end is not None else None
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
method: Final = json.loads(value.group(1))
|
||||
except ValueError:
|
||||
return None
|
||||
return method if isinstance(method, str) else None
|
||||
|
||||
|
||||
def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None:
|
||||
|
|
@ -2470,6 +2495,27 @@ if MCP_AVAILABLE:
|
|||
),
|
||||
)
|
||||
|
||||
def _is_listing_body(body: bytes) -> bool:
|
||||
"""A listing request, including one the routing peek cut off after its method."""
|
||||
if is_listing_request(body):
|
||||
return True
|
||||
peek_cut: Final = len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES
|
||||
return peek_cut and _jsonrpc_text_method(body.decode("utf-8", errors="replace")) in LISTING_METHODS
|
||||
|
||||
def _server_to_reauthorize(
|
||||
mcp_servers: Sequence[str] | None, scoped_server_endpoint: bool, body: bytes, client_ip: str | None
|
||||
) -> str | None:
|
||||
"""The server a client addressed alone for a listing it can fix by authenticating again."""
|
||||
if not scoped_server_endpoint or mcp_servers is None or len(mcp_servers) != 1 or not _is_listing_body(body):
|
||||
return None
|
||||
server: Final = operations.global_mcp_server_manager.get_mcp_server_by_identifier(
|
||||
mcp_servers[0], client_ip=client_ip
|
||||
)
|
||||
reauthorizes: Final = (
|
||||
server is not None and MCPServerManager.resolve_oauth2_flow_for_request(server).reauthorizes_in_client
|
||||
)
|
||||
return mcp_servers[0] if reauthorizes else None
|
||||
|
||||
async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
"""Handle MCP requests through StreamableHTTP."""
|
||||
try:
|
||||
|
|
@ -2638,13 +2684,15 @@ if MCP_AVAILABLE:
|
|||
await rejection(scope, receive, send)
|
||||
return
|
||||
|
||||
reauthorize_as: Final = _server_to_reauthorize(mcp_servers, scoped_server_endpoint, body, _client_ip)
|
||||
scope_response: Final = (
|
||||
OAuthScopeResponse(
|
||||
send,
|
||||
get_request_base_url(StarletteRequest(scope)),
|
||||
get_route_relative_request_path(scope),
|
||||
reauthorize_as=reauthorize_as,
|
||||
)
|
||||
if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES
|
||||
if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES or reauthorize_as is not None
|
||||
else None
|
||||
)
|
||||
if scope_response is not None:
|
||||
|
|
|
|||
|
|
@ -383,6 +383,12 @@ class MCPServer(LiteLLMBaseModel):
|
|||
door and is never a precondition for opening a session."""
|
||||
return self.is_true_passthrough or self.is_oauth_delegate
|
||||
|
||||
@property
|
||||
def reauthorizes_in_client(self) -> bool:
|
||||
"""True when a caller can fix an upstream 401 by authenticating again: the client-forwarded
|
||||
token modes and gateway-vaulted per-user OAuth."""
|
||||
return self.is_client_forwarded_token or (self.is_gateway_managed_oauth2 and self.needs_user_oauth_token)
|
||||
|
||||
@property
|
||||
def is_dcr_bridge(self) -> bool:
|
||||
"""True when this client-forwarded-token server serves the gateway-hosted DCR front door
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@ import asyncio
|
|||
from collections.abc import Mapping, Sequence
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx2
|
||||
import pytest
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.types import (
|
||||
ListPromptsRequest,
|
||||
|
|
@ -18,17 +20,24 @@ from mcp.types import (
|
|||
PaginatedRequestParams,
|
||||
Tool,
|
||||
)
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
from litellm.proxy.mcp_server import catalog
|
||||
from litellm.proxy.mcp_server import operations as mcp_operations
|
||||
from litellm.proxy.mcp_server.contracts import OperationContext
|
||||
from litellm.proxy.mcp_server.exceptions import MCPUpstreamAuthError
|
||||
from litellm.proxy.mcp_server.faults.list_outcomes import (
|
||||
SERVER_OUTCOMES_META_KEY,
|
||||
AggregateToolListing,
|
||||
ServerListFault,
|
||||
ServerListOk,
|
||||
)
|
||||
from litellm.types.mcp import MCPTransport
|
||||
from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var, mcp_gateway_server_name
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
from litellm.proxy.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
CatalogKind = Literal["tools", "prompts", "resources", "templates"]
|
||||
|
|
@ -772,3 +781,293 @@ def test_partial_optional_catalog_never_advertises_freshness() -> None:
|
|||
assert result.ttl_ms == 0
|
||||
assert result.cache_scope == "private"
|
||||
|
||||
|
||||
_GATEWAY: Final = "https://gateway.example"
|
||||
_UNAUTHORIZED: Final = ServerListFault(tag="auth_required", status_code=401)
|
||||
|
||||
|
||||
def _server(
|
||||
auth_type: MCPAuth,
|
||||
name: str = "docs",
|
||||
oauth2_flow: str | None = None,
|
||||
delegate_auth_to_upstream: bool = False,
|
||||
authentication_token: str | None = None,
|
||||
alias: str | None = None,
|
||||
) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id=f"{name}-id",
|
||||
name=name,
|
||||
alias=alias or name,
|
||||
url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=auth_type,
|
||||
oauth2_flow=oauth2_flow,
|
||||
delegate_auth_to_upstream=delegate_auth_to_upstream,
|
||||
authentication_token=authentication_token,
|
||||
)
|
||||
|
||||
|
||||
_VAULTED: Final = _server(MCPAuth.oauth2, oauth2_flow="authorization_code")
|
||||
_PEER: Final = _server(MCPAuth.none, name="wiki")
|
||||
_STATIC: Final = _server(MCPAuth.api_key, authentication_token="k")
|
||||
_REAUTHORIZABLE: Final = [
|
||||
pytest.param(_VAULTED, id="vaulted-per-user-oauth"),
|
||||
pytest.param(_server(MCPAuth.true_passthrough), id="true-passthrough"),
|
||||
pytest.param(_server(MCPAuth.oauth_delegate), id="oauth-delegate"),
|
||||
]
|
||||
_NOT_REAUTHORIZABLE: Final = [
|
||||
pytest.param(_server(MCPAuth.none), id="no-auth"),
|
||||
pytest.param(_STATIC, id="static-api-key"),
|
||||
pytest.param(_server(MCPAuth.bearer_token, authentication_token="k"), id="static-bearer"),
|
||||
pytest.param(_server(MCPAuth.oauth2, oauth2_flow="client_credentials"), id="client-credentials"),
|
||||
pytest.param(_server(MCPAuth.oauth2, delegate_auth_to_upstream=True), id="legacy-upstream-delegate"),
|
||||
]
|
||||
|
||||
|
||||
def _request_context(gate: OAuthScopeResponse, method: str) -> ServerRequestContext:
|
||||
return ServerRequestContext(
|
||||
session=MagicMock(),
|
||||
lifespan_context={},
|
||||
protocol_version="2025-11-25",
|
||||
method=method,
|
||||
request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}),
|
||||
)
|
||||
|
||||
|
||||
def _challenged(
|
||||
server: MCPServer,
|
||||
outcome: ServerListFault | ServerListOk,
|
||||
listed: Sequence[MCPServer] | None = None,
|
||||
reauthorize_as: str | None = "docs",
|
||||
method: str = "tools/list",
|
||||
) -> bool:
|
||||
gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as=reauthorize_as)
|
||||
token: Final = active_mcp_request_ctx_var.set(_request_context(gate, method))
|
||||
try:
|
||||
catalog.challenge_rejected_credential(server, outcome, [server] if listed is None else listed)
|
||||
finally:
|
||||
active_mcp_request_ctx_var.reset(token)
|
||||
return not gate.pending
|
||||
|
||||
|
||||
@pytest.mark.parametrize("server", _REAUTHORIZABLE)
|
||||
def test_a_rejected_credential_the_caller_can_renew_is_challenged(server: MCPServer) -> None:
|
||||
assert _challenged(server, _UNAUTHORIZED)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("server", _NOT_REAUTHORIZABLE)
|
||||
def test_a_rejected_credential_the_caller_cannot_renew_stays_absorbed(server: MCPServer) -> None:
|
||||
assert not _challenged(server, _UNAUTHORIZED)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"outcome",
|
||||
[
|
||||
ServerListOk(tool_count=3),
|
||||
ServerListFault(tag="forbidden", status_code=403),
|
||||
ServerListFault(tag="upstream_error", status_code=500),
|
||||
ServerListFault(tag="unreachable"),
|
||||
],
|
||||
)
|
||||
def test_only_an_unauthorized_outcome_is_challenged(outcome: ServerListFault | ServerListOk) -> None:
|
||||
assert not _challenged(_VAULTED, outcome)
|
||||
|
||||
|
||||
def test_a_listing_that_spans_several_servers_is_not_challenged() -> None:
|
||||
assert not _challenged(_VAULTED, _UNAUTHORIZED, [_VAULTED, _PEER])
|
||||
|
||||
|
||||
def test_a_request_that_is_not_a_single_server_listing_is_not_challenged() -> None:
|
||||
assert not _challenged(_VAULTED, _UNAUTHORIZED, reauthorize_as=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["tools/call", "server/discover", "initialize"])
|
||||
def test_a_request_the_sdk_parsed_as_something_other_than_a_listing_is_not_challenged(method: str) -> None:
|
||||
assert not _challenged(_VAULTED, _UNAUTHORIZED, method=method)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rejected_stored_token_is_answered_with_the_gateways_own_challenge() -> None:
|
||||
send: Final = AsyncMock()
|
||||
gate: Final = OAuthScopeResponse(send, _GATEWAY, "/docs/mcp", reauthorize_as="docs")
|
||||
token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "tools/list"))
|
||||
try:
|
||||
started: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
|
||||
await asyncio.sleep(0)
|
||||
catalog.challenge_rejected_credential(_VAULTED, _UNAUTHORIZED, [_VAULTED])
|
||||
await asyncio.wait_for(started, 1)
|
||||
finally:
|
||||
active_mcp_request_ctx_var.reset(token)
|
||||
|
||||
start: Final = send.await_args_list[0].args[0]
|
||||
assert start["status"] == 401
|
||||
assert dict(start["headers"])[b"www-authenticate"].decode() == (
|
||||
'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp"'
|
||||
)
|
||||
|
||||
|
||||
def _upstream_rejecting_the_token(_request: httpx2.Request) -> httpx2.Response:
|
||||
return httpx2.Response(401, headers={"www-authenticate": 'Bearer error="invalid_token"'})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("server", "peers", "challenged"),
|
||||
[
|
||||
pytest.param(_server(MCPAuth.true_passthrough), (), True, id="true-passthrough"),
|
||||
pytest.param(_server(MCPAuth.oauth_delegate), (), True, id="oauth-delegate"),
|
||||
pytest.param(_STATIC, (), False, id="static-api-key"),
|
||||
pytest.param(_server(MCPAuth.true_passthrough), (_PEER,), False, id="selected-group-of-two"),
|
||||
],
|
||||
)
|
||||
async def test_tools_listing_challenges_a_single_server_whose_token_the_upstream_rejected(
|
||||
monkeypatch: pytest.MonkeyPatch, server: MCPServer, peers: Sequence[MCPServer], challenged: bool
|
||||
) -> None:
|
||||
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(_upstream_rejecting_the_token))
|
||||
monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager)
|
||||
gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as="docs")
|
||||
context: Final = mcp_operations.prepare_context(
|
||||
UserAPIKeyAuth(token="h", user_id="u1"), oauth2_headers={"Authorization": "Bearer rejected"}
|
||||
)
|
||||
request_token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "tools/list"))
|
||||
name_token: Final = mcp_gateway_server_name.set("docs")
|
||||
try:
|
||||
listing, outcome = await catalog.get_filtered_server_tools(
|
||||
server, context=context, allowed_mcp_servers=[server, *peers], prefetched_oauth_creds={}
|
||||
)
|
||||
finally:
|
||||
mcp_gateway_server_name.reset(name_token)
|
||||
active_mcp_request_ctx_var.reset(request_token)
|
||||
|
||||
assert listing.tools == []
|
||||
assert outcome.tag == "auth_required"
|
||||
assert (not gate.pending) is challenged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("request_name", ["ListPromptsRequest", "ListResourcesRequest", "ListResourceTemplatesRequest"])
|
||||
@pytest.mark.parametrize(
|
||||
("server", "challenged"),
|
||||
[
|
||||
pytest.param(_VAULTED, True, id="vaulted-per-user-oauth"),
|
||||
pytest.param(_STATIC, False, id="static-api-key"),
|
||||
],
|
||||
)
|
||||
async def test_optional_listing_challenges_a_server_whose_stored_token_was_rejected(
|
||||
monkeypatch: pytest.MonkeyPatch, request_name: str, server: MCPServer, challenged: bool
|
||||
) -> None:
|
||||
from mcp import types
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "listing-reauth-test")
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(mcp_operations.global_mcp_server_manager, "registry", {server.server_id: server})
|
||||
monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server]))
|
||||
monkeypatch.setattr(
|
||||
catalog, "fetch_optional_catalog_page", AsyncMock(side_effect=MCPUpstreamAuthError(401, None, "docs"))
|
||||
)
|
||||
gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as="docs")
|
||||
token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "prompts/list"))
|
||||
try:
|
||||
result: Final = await catalog.list_gateway_catalog(
|
||||
mcp_operations.prepare_context(), getattr(types, request_name)()
|
||||
)
|
||||
finally:
|
||||
active_mcp_request_ctx_var.reset(token)
|
||||
|
||||
assert next(iter(result.meta[SERVER_OUTCOMES_META_KEY].values())) == {"status": "auth_required", "http_status": 401}
|
||||
assert (not gate.pending) is challenged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_name", "result_name", "items"),
|
||||
[
|
||||
("ListPromptsRequest", "ListPromptsResult", "prompts"),
|
||||
("ListResourcesRequest", "ListResourcesResult", "resources"),
|
||||
("ListResourceTemplatesRequest", "ListResourceTemplatesResult", "resource_templates"),
|
||||
],
|
||||
)
|
||||
async def test_optional_listing_challenges_when_the_token_is_rejected_on_a_later_page(
|
||||
monkeypatch: pytest.MonkeyPatch, request_name: str, result_name: str, items: str
|
||||
) -> None:
|
||||
from mcp import types
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "listing-reauth-test")
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(mcp_operations.global_mcp_server_manager, "registry", {_VAULTED.server_id: _VAULTED})
|
||||
monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[_VAULTED]))
|
||||
first_page: Final = getattr(types, result_name)(**{items: [], "next_cursor": "next"})
|
||||
monkeypatch.setattr(
|
||||
catalog,
|
||||
"fetch_optional_catalog_page",
|
||||
AsyncMock(side_effect=[first_page, MCPUpstreamAuthError(401, None, "docs")]),
|
||||
)
|
||||
gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as="docs")
|
||||
token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "prompts/list"))
|
||||
request: Final = getattr(types, request_name)
|
||||
try:
|
||||
first: Final = await catalog.list_gateway_catalog(mcp_operations.prepare_context(), request())
|
||||
assert first.next_cursor
|
||||
assert gate.pending
|
||||
with pytest.raises(MCPError, match="Upstream continuation failed"):
|
||||
await catalog.list_gateway_catalog(
|
||||
mcp_operations.prepare_context(),
|
||||
request(params=types.PaginatedRequestParams(cursor=first.next_cursor)),
|
||||
)
|
||||
finally:
|
||||
active_mcp_request_ctx_var.reset(token)
|
||||
|
||||
assert not gate.pending
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("request_path", "resource"),
|
||||
[("/mcp/docs", "litellm/mcp/docs"), ("/docs/mcp", "litellm/docs/mcp")],
|
||||
)
|
||||
def test_the_challenge_names_the_metadata_route_under_a_server_root_path(
|
||||
monkeypatch: pytest.MonkeyPatch, request_path: str, resource: str
|
||||
) -> None:
|
||||
monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm")
|
||||
|
||||
challenge: Final = MCPUpstreamAuthError(401, None, "docs").to_http_exception(_GATEWAY, request_path)
|
||||
|
||||
assert (challenge.headers or {})["www-authenticate"] == (
|
||||
f'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/{resource}"'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_path", "addressed", "resource"),
|
||||
[
|
||||
pytest.param("/docs/mcp", "docs", "docs/mcp", id="per-server-path"),
|
||||
pytest.param("/Docs/mcp", "Docs", "Docs/mcp", id="other-case"),
|
||||
pytest.param("/docs_internal/mcp", "docs_internal", "docs_internal/mcp", id="name-not-alias"),
|
||||
pytest.param("/mcp/Docs", "Docs", "mcp/Docs", id="standard-path-other-case"),
|
||||
pytest.param("/mcp/a:b@c+d,e", "a:b@c+d,e", "mcp/a:b@c+d,e", id="path-legal-characters-stay-as-typed"),
|
||||
pytest.param("/mcp/a b%c", "a b%c", "mcp/a%20b%25c", id="path-illegal-characters-are-encoded"),
|
||||
pytest.param("/mcp/日本", "日本", "mcp/%E6%97%A5%E6%9C%AC", id="non-ascii-name"),
|
||||
pytest.param("/日本/mcp", "日本", "%E6%97%A5%E6%9C%AC/mcp", id="non-ascii-name-legacy-path"),
|
||||
],
|
||||
)
|
||||
async def test_the_challenge_names_the_url_the_client_addressed(request_path: str, addressed: str, resource: str) -> None:
|
||||
send: Final = AsyncMock()
|
||||
gate: Final = OAuthScopeResponse(send, _GATEWAY, request_path, reauthorize_as=addressed)
|
||||
aliased: Final = _server(MCPAuth.oauth2, name="docs_internal", oauth2_flow="authorization_code", alias="docs")
|
||||
token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "tools/list"))
|
||||
try:
|
||||
started: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
|
||||
await asyncio.sleep(0)
|
||||
catalog.challenge_rejected_credential(aliased, _UNAUTHORIZED, [aliased])
|
||||
await asyncio.wait_for(started, 1)
|
||||
finally:
|
||||
active_mcp_request_ctx_var.reset(token)
|
||||
|
||||
start: Final = send.await_args_list[0].args[0]
|
||||
assert dict(start["headers"])[b"www-authenticate"].decode() == (
|
||||
f'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/{resource}"'
|
||||
)
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.proxy.mcp_server.mcp_server_manager import (
|
|||
MCPServerManager,
|
||||
_extract_upstream_auth_failure,
|
||||
)
|
||||
from litellm.proxy.mcp_server.scope_challenge import OAuthScopeResponse, is_listing_request
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
|
@ -375,12 +376,10 @@ async def test_aggregate_list_tools_absorbs_one_unauthenticated_server():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_server_route_also_absorbs_upstream_auth_error():
|
||||
async def test_single_server_route_also_absorbs_upstream_auth_error() -> None:
|
||||
"""A single-server route (/<server>/mcp) absorbs an upstream-auth error just like the aggregate:
|
||||
the failing server contributes no tools and an "auth_required" outcome rather than re-raising.
|
||||
Surfacing it to the client as a 401 + WWW-Authenticate challenge cannot be done from this list
|
||||
handler — the MCP session manager serializes a raise into a JSON-RPC error, not an HTTP 401 — so
|
||||
re-auth surfacing is handled by a request-scope preemptive check, tracked separately."""
|
||||
The 401 a single-server request receives is sent through its scope response, not raised from here."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
@ -836,3 +835,73 @@ def test_managed_step_up_metadata_matches_the_sdk_configured_resource(path):
|
|||
requested_resource=f"https://gateway.test{path}",
|
||||
configured_resource=f"https://gateway.test{resource_path}",
|
||||
), "the SDK must accept the advertised resource before it can reauthorize"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "expected"),
|
||||
[
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"tools/list"}', True),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"prompts/list"}', True),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"resources/list"}', True),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"resources/templates/list"}', True),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"tools/call"}', False),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"initialize"}', False),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":["tools/list"]}', False),
|
||||
(b'{"jsonrpc":"2.0","id":1}', False),
|
||||
(b"not json", False),
|
||||
],
|
||||
)
|
||||
def test_listing_requests_are_recognised_by_method(body: bytes, expected: bool) -> None:
|
||||
assert is_listing_request(body) is expected
|
||||
|
||||
|
||||
async def _pending_listing_response(gate: OAuthScopeResponse, send: AsyncMock) -> "asyncio.Task[None]":
|
||||
task: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
|
||||
await asyncio.sleep(0)
|
||||
send.assert_not_awaited()
|
||||
return task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reauth_denial_replaces_the_pending_listing_response_with_the_gateway_challenge() -> None:
|
||||
send: Final = AsyncMock()
|
||||
gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp", reauthorize_as="docs")
|
||||
started: Final = await _pending_listing_response(gate, send)
|
||||
|
||||
gate.deny_reauth()
|
||||
await asyncio.wait_for(started, 1)
|
||||
|
||||
start, body = [call.args[0] for call in send.await_args_list]
|
||||
assert start["status"] == 401
|
||||
assert dict(start["headers"])[b"www-authenticate"].decode() == (
|
||||
'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp"'
|
||||
)
|
||||
assert body["body"] == b'{"detail":"Unauthorized"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_gate_that_does_not_reauthorize_ignores_a_reauth_denial() -> None:
|
||||
send: Final = AsyncMock()
|
||||
gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp")
|
||||
started: Final = await _pending_listing_response(gate, send)
|
||||
|
||||
gate.deny_reauth()
|
||||
assert gate.pending
|
||||
gate.allow()
|
||||
await asyncio.wait_for(started, 1)
|
||||
|
||||
assert [call.args[0]["status"] for call in send.await_args_list] == [200]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scope_step_up_denial_still_ignores_an_unauthorized_response() -> None:
|
||||
send: Final = AsyncMock()
|
||||
gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp", reauthorize_as="docs")
|
||||
started: Final = await _pending_listing_response(gate, send)
|
||||
|
||||
gate.deny(MCPUpstreamAuthError(401, None, "docs"))
|
||||
assert gate.pending
|
||||
gate.allow()
|
||||
await asyncio.wait_for(started, 1)
|
||||
|
||||
assert [call.args[0]["status"] for call in send.await_args_list] == [200]
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
from collections.abc import Sequence
|
||||
from contextlib import asynccontextmanager, nullcontext
|
||||
from typing import Final
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
|
@ -17,6 +19,7 @@ from litellm.proxy.mcp_server.mcp_server_manager import (
|
|||
MCPServerManager,
|
||||
MCPTransport,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
mcp_server_manager = MCPServerManager()
|
||||
|
||||
|
|
@ -3209,6 +3212,140 @@ async def test_mcp_routing_stashes_only_method_for_auth_and_replays_full_body(me
|
|||
)
|
||||
|
||||
|
||||
def _docs_server(auth_type: str, legacy_client_credentials: bool = False) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id="docs-id",
|
||||
name="docs",
|
||||
alias="docs",
|
||||
url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth(auth_type),
|
||||
token_url="https://idp.example/token" if legacy_client_credentials else None,
|
||||
client_id="client" if legacy_client_credentials else None,
|
||||
client_secret="secret" if legacy_client_credentials else None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "method", "shape", "selected_servers", "docs", "expected"),
|
||||
[
|
||||
("/mcp/docs", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"),
|
||||
("/mcp/docs", "prompts/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"),
|
||||
("/mcp/docs", "resources/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"),
|
||||
("/mcp/docs", "resources/templates/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"),
|
||||
("/mcp/DOCS", "tools/list", "plain", ["DOCS"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/DOCS"),
|
||||
("/mcp/docs", "tools/list", "plain", ["docs"], _docs_server("api_key"), "none"),
|
||||
("/mcp/docs", "tools/list", "plain", ["docs"], _docs_server("oauth2", legacy_client_credentials=True), "none"),
|
||||
("/mcp/docs", "tools/call", "plain", ["docs"], _docs_server("oauth_delegate"), "holds"),
|
||||
("/mcp/docs", "tools/list", "plain", ["docs", "wiki"], _docs_server("oauth_delegate"), "none"),
|
||||
("/mcp/docs", "tools/list", "plain", ["unregistered_group"], _docs_server("oauth_delegate"), "none"),
|
||||
("/mcp/docs", "tools/list", "plain", None, _docs_server("oauth_delegate"), "none"),
|
||||
("/mcp/docs", "ping", "plain", ["docs"], _docs_server("oauth_delegate"), "none"),
|
||||
("/docs/mcp", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/docs/mcp"),
|
||||
("/Docs/mcp", "tools/list", "plain", ["Docs"], _docs_server("oauth_delegate"), "reauthorizes:/Docs/mcp"),
|
||||
("/mcp/docs-id", "tools/list", "plain", ["docs-id"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs-id"),
|
||||
("/mcp", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "none"),
|
||||
("/toolset/team/mcp", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "none"),
|
||||
("/mcp/docs", "tools/list", "cursor-after-method", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"),
|
||||
("/mcp/docs", "prompts/list", "cursor-after-method", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"),
|
||||
("/mcp/docs", "tools/list", "cursor-before-method", ["docs"], _docs_server("oauth_delegate"), "holds"),
|
||||
("/mcp/docs", "tools/call", "cursor-after-method", ["docs"], _docs_server("oauth_delegate"), "holds"),
|
||||
("/mcp/docs", "ping", "duplicate-method", ["docs"], _docs_server("oauth_delegate"), "none"),
|
||||
],
|
||||
)
|
||||
async def test_mcp_routing_holds_a_single_servers_listing_response_for_a_reauth_challenge(
|
||||
path: str, method: str, shape: str, selected_servers: Sequence[str] | None, docs: MCPServer, expected: str
|
||||
) -> None:
|
||||
try:
|
||||
from litellm.proxy.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse
|
||||
from litellm.proxy.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
session_manager_stateful,
|
||||
session_manager_stateless,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
async def gate_kind(gate: object) -> str:
|
||||
if not isinstance(gate, OAuthScopeResponse):
|
||||
return "none"
|
||||
gate.deny_reauth()
|
||||
if gate.pending:
|
||||
return "holds"
|
||||
await gate({"type": "http.response.start", "status": 200, "headers": []})
|
||||
challenge: Final = dict(sent.await_args_list[0].args[0]["headers"])[b"www-authenticate"].decode()
|
||||
return "reauthorizes:" + challenge.split("oauth-protected-resource")[1].rstrip('"')
|
||||
|
||||
from litellm.proxy.mcp_server import operations
|
||||
|
||||
envelope: Final = {"jsonrpc": "2.0", "id": 2, "method": method, "params": {"cursor": "c" * 6000}}
|
||||
bodies: Final = {
|
||||
"plain": json.dumps({**envelope, "params": {}}).encode(),
|
||||
"cursor-after-method": json.dumps(envelope).encode(),
|
||||
"cursor-before-method": json.dumps({"params": envelope["params"], **envelope}).encode(),
|
||||
"duplicate-method": b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{},"method":"ping"}',
|
||||
}
|
||||
jsonrpc_body: Final = bodies[shape]
|
||||
scope: Final = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": path,
|
||||
"scheme": "https",
|
||||
"server": ("gateway.example", 443),
|
||||
"headers": [
|
||||
(b"host", b"gateway.example"),
|
||||
(b"content-type", b"application/json"),
|
||||
(b"authorization", b"Bearer test-key"),
|
||||
],
|
||||
}
|
||||
receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": jsonrpc_body, "more_body": False}])
|
||||
held: Final = queue.SimpleQueue[object]()
|
||||
sent: Final = AsyncMock()
|
||||
|
||||
async def stateless_handle(handled_scope: Scope, _receive: Receive, _send: Send) -> None:
|
||||
held.put(handled_scope.get(SCOPE_RESPONSE_KEY))
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam
|
||||
"litellm.proxy.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(MagicMock(), None, selected_servers, None, None, None),
|
||||
),
|
||||
patch( # test-quality-ok: registry is empty in unit tests; key owns one server
|
||||
"litellm.proxy.mcp_server.server._get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[MagicMock()],
|
||||
),
|
||||
patch( # test-quality-ok: init flag is a module global; no injection seam
|
||||
"litellm.proxy.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
|
||||
True,
|
||||
),
|
||||
patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam
|
||||
session_manager_stateless, "handle_request", side_effect=stateless_handle
|
||||
),
|
||||
patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam
|
||||
session_manager_stateful, "handle_request", side_effect=stateless_handle
|
||||
),
|
||||
patch.object( # test-quality-ok: the registry is a module-level singleton's attribute; the suite's only seam
|
||||
operations.global_mcp_server_manager, "registry", {docs.server_id: docs}
|
||||
),
|
||||
patch( # test-quality-ok: the missing-token challenge runs before the gate and is not under test here
|
||||
"litellm.proxy.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam
|
||||
session_manager_stateless, "_server_instances", {}
|
||||
),
|
||||
patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam
|
||||
session_manager_stateful, "_server_instances", {}
|
||||
),
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, sent)
|
||||
|
||||
assert await gate_kind(held.get_nowait()) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_routing_batch_body_is_not_stashed_for_auth():
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -4190,6 +4190,29 @@ def test_jsonrpc_text_has_top_level_method_ignores_nested_method():
|
|||
assert _jsonrpc_text_has_top_level_method(truncated) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "method"),
|
||||
[
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', "tools/list", id="complete"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"cursor":"' + "c" * 5000, "tools/list", id="cut-after-method"),
|
||||
pytest.param('{"jsonrpc":"2.0","params":{"cursor":"' + "c" * 5000, None, id="cut-before-method"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools/li', None, id="cut-inside-method"),
|
||||
pytest.param('{"jsonrpc":"2.0","params":{"method":"tools/list"},"method":"ping"}', "ping", id="nested-method-first"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools\\/list","params":{"cursor":"' + "c" * 5000, "tools/list", id="escaped-slash"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":"\\u0074ools/list","params":{"cursor":"' + "c" * 5000, "tools/list", id="unicode-escape"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":"bad\\x","params":{}}', None, id="invalid-escape"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"method":7,"params":{}}', None, id="not-a-string"),
|
||||
pytest.param('{"jsonrpc":"2.0","id":1,"result":{"method":"tools/list"}}', None, id="response"),
|
||||
],
|
||||
)
|
||||
def test_jsonrpc_text_method_reads_the_top_level_method_when_the_cut_leaves_it_whole(
|
||||
text: str, method: str | None
|
||||
) -> None:
|
||||
from litellm.proxy.mcp_server.server import _jsonrpc_text_method
|
||||
|
||||
assert _jsonrpc_text_method(text) == method
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("response_field", ("result", "error"))
|
||||
@pytest.mark.parametrize(("character", "bytes_before_cap"), (("", 0), ("x", 0), ("é", 1), ("中", 2), ("😀", 3)))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue