mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
3391 lines
153 KiB
Python
3391 lines
153 KiB
Python
"""
|
|
LiteLLM MCP Server Routes
|
|
"""
|
|
|
|
# pyright: reportInvalidTypeForm=false, reportArgumentType=false, reportOptionalCall=false
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import contextvars
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
import types
|
|
from collections import Counter
|
|
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING, Final, Literal, NoReturn, Protocol
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
|
from starlette.requests import Request as StarletteRequest
|
|
from starlette.responses import JSONResponse, Response
|
|
from starlette.routing import Route
|
|
from starlette.types import Message, Receive, Scope, Send
|
|
|
|
from litellm._logging import verbose_logger
|
|
from litellm.constants import (
|
|
MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS,
|
|
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
|
|
MCP_PEEKED_BODY_SCOPE_KEY,
|
|
)
|
|
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
|
from litellm.llms.custom_httpx.http_handler import (
|
|
get_async_httpx_client,
|
|
httpxSpecialProvider,
|
|
)
|
|
from litellm.proxy._types import (
|
|
ProxyException,
|
|
SpecialMCPServerNames,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
|
from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( # noqa: F401 # legacy module exports
|
|
MCPRequestHandler,
|
|
_is_mcp_admitted_user_subject, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
|
is_mcp_admitted_user_subject,
|
|
)
|
|
from litellm.proxy.mcp_server.client_allowlist import (
|
|
MCPClientAllowlist,
|
|
check_mcp_client_allowed,
|
|
load_mcp_client_allowlist,
|
|
)
|
|
from litellm.proxy.mcp_server.discoverable_endpoints import (
|
|
get_request_base_url,
|
|
)
|
|
from litellm.proxy.mcp_server.exceptions import (
|
|
MCPUpstreamAuthError,
|
|
)
|
|
from litellm.proxy.mcp_server.mcp_context import ( # noqa: F401 # legacy module exports
|
|
_mcp_active_toolset_id, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
|
_mcp_gateway_initialize_instructions, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
|
_mcp_gateway_server_name, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
|
_mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode
|
|
active_mcp_request_ctx_var,
|
|
get_active_mcp_request_ctx,
|
|
mcp_active_toolset_id,
|
|
mcp_gateway_initialize_instructions,
|
|
mcp_gateway_server_name,
|
|
)
|
|
from litellm.proxy.mcp_server.mcp_debug import (
|
|
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY,
|
|
MCPAuthDiagnostics,
|
|
MCPDebug,
|
|
)
|
|
from litellm.proxy.mcp_server.oauth_utils import (
|
|
get_passthrough_www_authenticate,
|
|
get_route_relative_request_path,
|
|
redact_mcp_resource_url,
|
|
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 (
|
|
ActingUser,
|
|
GrantedToolsetIds,
|
|
acting_user_auth,
|
|
granted_toolset_ids,
|
|
is_ui_session_credential,
|
|
)
|
|
from litellm.proxy.mcp_server.utils import (
|
|
LITELLM_MCP_SERVER_DESCRIPTION,
|
|
LITELLM_MCP_SERVER_NAME,
|
|
LITELLM_MCP_SERVER_VERSION,
|
|
asgi_scope_fields,
|
|
merge_mcp_headers,
|
|
)
|
|
from litellm.types.mcp import (
|
|
MCPAuth,
|
|
MCPGatewaySession,
|
|
MCPGatewaySessionGroupCount,
|
|
MCPGatewaySessionsResponse,
|
|
MCPGatewaySessionsTerminateResponse,
|
|
MCPSpecVersion,
|
|
)
|
|
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp.server.session import ServerSession as _McpServerSession
|
|
|
|
|
|
_redact_mcp_resource_url: Final = redact_mcp_resource_url
|
|
|
|
_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 30 * 60
|
|
# Upper bound on concurrent stateful sessions a single caller may hold. Each
|
|
# `initialize` creates a session that survives until the idle timeout, so
|
|
# without a cap an authenticated client could spam `initialize` and exhaust
|
|
# memory. The caller's own oldest idle sessions are evicted to make room; if
|
|
# the cap is still hit (every session in flight), the new `initialize` is
|
|
# rejected with 429.
|
|
_MAX_STATEFUL_SESSIONS_PER_OWNER: Final = 100
|
|
# Maximum bytes to peek when sniffing the JSON-RPC method on a POST.
|
|
# An `initialize` envelope is a few hundred bytes; capping the peek
|
|
# prevents an authenticated client from forcing the proxy to buffer an
|
|
# arbitrarily large body just to make a routing decision.
|
|
_MCP_ROUTING_PEEK_MAX_BYTES: Final = 4096
|
|
# ASGI scope keys carrying OTel request state into a stateful MCP message handler.
|
|
_MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
|
|
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
|
|
_MCP_PROTOCOL_VERSION_HEADER: Final = b"mcp-protocol-version"
|
|
_DECODED_JSON: Final = TypeAdapter(object)
|
|
|
|
|
|
def reject_disallowed_mcp_origin(request: StarletteRequest) -> None:
|
|
from litellm.proxy.proxy_server import origins # noqa: PLC0415 # proxy imports this module during startup
|
|
|
|
if "*" not in origins and any(origin not in origins for origin in request.headers.getlist("origin")):
|
|
raise HTTPException(status_code=403, detail="Invalid Origin header")
|
|
|
|
|
|
def _scope_path(scope: Scope) -> str:
|
|
return asgi_scope_fields(scope).get("path", "")
|
|
|
|
|
|
def _scope_headers(scope: Scope) -> Sequence[tuple[bytes, bytes]]:
|
|
return asgi_scope_fields(scope).get("headers", [])
|
|
|
|
|
|
class _RawHeaderPairsSource(Protocol):
|
|
def get(
|
|
self, key: Literal["headers"], default: Sequence[tuple[bytes | str, bytes | str]], /
|
|
) -> Sequence[tuple[bytes | str, bytes | str]]: ...
|
|
|
|
|
|
def _raw_header_pairs(source: _RawHeaderPairsSource) -> Sequence[tuple[bytes | str, bytes | str]]:
|
|
return source.get("headers", [])
|
|
|
|
|
|
def unsupported_protocol_version(scope: Scope) -> str | None:
|
|
"""Admit configured HTTP revisions while keeping SSE on the legacy protocol."""
|
|
from litellm.proxy.mcp_server.capabilities import configured_versions
|
|
|
|
headers: Final[Iterable[tuple[bytes, bytes]]] = scope.get("headers") or ()
|
|
values: Final = tuple(
|
|
raw.decode("latin-1").strip() for key, raw in headers if key.lower() == _MCP_PROTOCOL_VERSION_HEADER
|
|
)
|
|
for value in values:
|
|
if value == "2026-07-28" and _scope_path(scope).rstrip("/").endswith("/sse"):
|
|
return value
|
|
if value and value not in configured_versions():
|
|
return value
|
|
return None
|
|
|
|
|
|
# Check if MCP is available
|
|
# "mcp" requires python 3.10 or higher, but several litellm users use python 3.8
|
|
# We're making this conditional import to avoid breaking users who use python 3.8.
|
|
# TODO: Make this a util function for litellm client usage
|
|
MCP_AVAILABLE: bool = True
|
|
try:
|
|
import weakref
|
|
|
|
from mcp import ReadResourceResult, Resource
|
|
from mcp.server import Server
|
|
from mcp.server.runner import serve_loop
|
|
from mcp.server.session import ServerSession as _McpServerSession
|
|
from mcp.types import (
|
|
BlobResourceContents,
|
|
DiscoverRequest,
|
|
DiscoverResult,
|
|
GetPromptResult,
|
|
RequestParams,
|
|
ResourceTemplate,
|
|
TextResourceContents,
|
|
)
|
|
|
|
# Robust auth lookup keyed by session_object.
|
|
_session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary()
|
|
except ImportError as e:
|
|
verbose_logger.debug("MCP module not found: %s", e)
|
|
MCP_AVAILABLE = False
|
|
# When MCP is not available, we set these to None at module level
|
|
# All code using these types is inside `if MCP_AVAILABLE:` blocks
|
|
# so they will never be accessed at runtime
|
|
BlobResourceContents = None
|
|
GetPromptResult = None
|
|
ReadResourceResult = None
|
|
Resource = None
|
|
ResourceTemplate = None
|
|
Server = None
|
|
TextResourceContents = None
|
|
|
|
active_mcp_session_var: Final[contextvars.ContextVar["_McpServerSession | None"]] = contextvars.ContextVar(
|
|
"active_mcp_session", default=None
|
|
)
|
|
|
|
|
|
# Global variables to track initialization
|
|
_SESSION_MANAGERS_INITIALIZED = False
|
|
_INITIALIZATION_LOCK: Final = asyncio.Lock()
|
|
|
|
|
|
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.
|
|
|
|
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 an index only when the key is positively found; truncation that
|
|
hides it yields ``None``.
|
|
"""
|
|
depth = 0
|
|
in_string = False
|
|
escaped = False
|
|
in_object: Final[list[bool]] = []
|
|
reading_key = False
|
|
expect_key = False
|
|
key_chars: list[str] = []
|
|
for index, ch in enumerate(text):
|
|
if in_string:
|
|
if escaped:
|
|
escaped = False
|
|
elif ch == "\\":
|
|
escaped = True
|
|
elif ch == '"':
|
|
in_string = False
|
|
if reading_key and depth == 1 and "".join(key_chars) == key:
|
|
return index + 1
|
|
elif reading_key:
|
|
key_chars.append(ch)
|
|
continue
|
|
if ch == '"':
|
|
in_string = True
|
|
reading_key = expect_key and depth >= 1 and in_object[-1]
|
|
key_chars = []
|
|
expect_key = False
|
|
elif ch == "{" or ch == "[":
|
|
depth += 1
|
|
in_object.append(ch == "{")
|
|
expect_key = ch == "{"
|
|
elif ch == "}" or ch == "]":
|
|
if in_object:
|
|
in_object.pop()
|
|
depth -= 1
|
|
if depth <= 0:
|
|
break
|
|
expect_key = False
|
|
elif ch == ",":
|
|
expect_key = bool(in_object) and in_object[-1]
|
|
elif ch == ":":
|
|
expect_key = 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:
|
|
"""The W3C trace context (``traceparent``/``tracestate``) the MCP client
|
|
propagated in the request's ``params._meta`` (SEP-414), or ``None``.
|
|
|
|
When present, the MCP span records this propagated context as a span *link*,
|
|
never the parent — a remote parent would root the span in a trace whose root
|
|
never reaches the gateway's tracing backend. The span itself nests under the
|
|
transport span of the request carrying this specific message, so a
|
|
streamable-HTTP session that multiplexes many messages still does not glue
|
|
every message under the session's first request;
|
|
see ``resolve_mcp_span_context``. The client's W3C Baggage is
|
|
deliberately excluded: it is caller-controlled, and the otel baggage processor
|
|
stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``,
|
|
...) onto the span, so honoring remote baggage would let a client spoof a
|
|
span's identity attribution.
|
|
"""
|
|
meta: Final = getattr(req_ctx, "meta", None)
|
|
extra: Final = meta if isinstance(meta, Mapping) else getattr(meta, "model_extra", None)
|
|
if not isinstance(extra, Mapping):
|
|
return None
|
|
carrier: Final = {key: extra[key] for key in ("traceparent", "tracestate") if isinstance(extra.get(key), str)}
|
|
return carrier or None
|
|
|
|
|
|
def _otel_set_mcp_trace_carrier(carrier: dict[str, str] | None) -> object:
|
|
"""Stash ``carrier`` for the otel_v2 MCP span and return a reset token, or
|
|
``None`` when otel_v2 is unavailable. Lazily imported so opentelemetry stays an
|
|
optional dependency."""
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import (
|
|
set_mcp_message_trace_carrier,
|
|
)
|
|
|
|
return set_mcp_message_trace_carrier(carrier)
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def _otel_reset_mcp_trace_carrier(token: object) -> None:
|
|
"""Clear the per-message trace carrier so it never leaks to the next message on
|
|
the same session task. Paired with ``_otel_set_mcp_trace_carrier``."""
|
|
if token is None:
|
|
return
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import (
|
|
reset_mcp_message_trace_carrier,
|
|
)
|
|
|
|
reset_mcp_message_trace_carrier(token)
|
|
except ImportError:
|
|
return
|
|
|
|
|
|
def _otel_publish_transport_span_on_scope(scope: Scope) -> None:
|
|
"""Record this request's tracing span on its own ASGI scope.
|
|
|
|
Resolved on the ASGI request task, where the proxy's server span is anchored,
|
|
and read back by the MCP message handler through ``req_ctx.request`` — the
|
|
``Request`` the transport attaches to each message. A stateful streamable-HTTP
|
|
session handles every message on the task spawned by its ``initialize`` POST, so
|
|
the handler's own task cannot see later requests' spans.
|
|
|
|
The scope, not the shared session auth context: a JSON-RPC *response* POST
|
|
deliberately skips the per-session lock (it can arrive while the tool call that
|
|
awaits it is still in flight), so a field on that shared object would be
|
|
overwritten mid-call and the tool call would attribute itself to the response's
|
|
request. A scope belongs to exactly one request and dies with it, which also
|
|
keeps a finished span from being retained by an idle session.
|
|
|
|
The live span, not just its context: a failed tool call stamps ``error.*`` on it,
|
|
which needs a span still open for writes. Lazily imported so opentelemetry stays
|
|
an optional dependency; a no-op when otel_v2 is unavailable or no request span is
|
|
anchored."""
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import (
|
|
request_root_span,
|
|
)
|
|
|
|
span: Final = request_root_span()
|
|
except ImportError:
|
|
return
|
|
if span is not None:
|
|
scope[_MCP_TRANSPORT_SPAN_SCOPE_KEY] = span
|
|
|
|
|
|
def _otel_value_from_message_scope(req_ctx: object, key: str) -> object:
|
|
request: Final = getattr(req_ctx, "request", None)
|
|
scope: Final = getattr(request, "scope", None)
|
|
if not isinstance(scope, Mapping):
|
|
return None
|
|
return scope.get(key)
|
|
|
|
|
|
def _otel_transport_span_from_message(req_ctx: object) -> object:
|
|
"""The tracing span of the HTTP request that carried this MCP message."""
|
|
return _otel_value_from_message_scope(req_ctx, _MCP_TRANSPORT_SPAN_SCOPE_KEY)
|
|
|
|
|
|
def _otel_set_mcp_transport_span(span: object) -> object:
|
|
"""Publish the current message's transport span, which the otel_v2 MCP span
|
|
attaches to and a failed tool call stamps its error on. Returns a reset token,
|
|
or ``None`` when otel_v2 is unavailable."""
|
|
if span is None:
|
|
return None
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import (
|
|
set_mcp_message_transport_span,
|
|
)
|
|
|
|
return set_mcp_message_transport_span(span)
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def _otel_reset_mcp_transport_span(token: object) -> None:
|
|
"""Paired with ``_otel_set_mcp_transport_span``."""
|
|
if token is None:
|
|
return
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import (
|
|
reset_mcp_message_transport_span,
|
|
)
|
|
|
|
reset_mcp_message_transport_span(token)
|
|
except ImportError:
|
|
return
|
|
|
|
|
|
def _otel_publish_request_destinations_on_scope(scope: Scope) -> None:
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import request_destinations
|
|
|
|
scope[_MCP_DESTINATIONS_SCOPE_KEY] = request_destinations()
|
|
except ImportError:
|
|
return
|
|
|
|
|
|
def _otel_set_mcp_request_destinations(req_ctx: object) -> object:
|
|
destinations: Final = _otel_value_from_message_scope(req_ctx, _MCP_DESTINATIONS_SCOPE_KEY)
|
|
if not isinstance(destinations, tuple):
|
|
return None
|
|
try:
|
|
from litellm.integrations.otel.model.destination import OtelDestination
|
|
from litellm.integrations.otel.plumbing.context import set_request_destinations
|
|
|
|
destination_adapter: Final[TypeAdapter[tuple[OtelDestination, ...]]] = TypeAdapter(
|
|
tuple[OtelDestination, ...],
|
|
config=ConfigDict(revalidate_instances="always"),
|
|
)
|
|
validated_destinations: Final = destination_adapter.validate_python(destinations, strict=True)
|
|
return set_request_destinations(validated_destinations)
|
|
except (ImportError, ValidationError):
|
|
return None
|
|
|
|
|
|
def _otel_reset_mcp_request_destinations(token: object) -> None:
|
|
if token is None:
|
|
return
|
|
try:
|
|
from litellm.integrations.otel.plumbing.context import reset_request_destinations
|
|
|
|
reset_request_destinations(token)
|
|
except ImportError:
|
|
return
|
|
|
|
|
|
def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException:
|
|
"""Map a ``ProxyException`` to an ``HTTPException`` that preserves its real
|
|
status code and headers.
|
|
|
|
``user_api_key_auth`` raises ``ProxyException`` (not ``HTTPException``) on
|
|
auth failures. The MCP ASGI handlers re-raise ``HTTPException`` to keep the
|
|
status and any ``WWW-Authenticate`` challenge, but a ``ProxyException`` would
|
|
otherwise fall through to their generic handler and be flattened to a 500 —
|
|
dropping the 401 + challenge an OAuth client needs to re-authenticate, so the
|
|
tool call surfaces as a cancelled/terminated session instead.
|
|
"""
|
|
try:
|
|
status_code = int(exc.code)
|
|
except (TypeError, ValueError):
|
|
status_code = 500
|
|
return HTTPException(
|
|
status_code=status_code,
|
|
detail=exc.message,
|
|
headers=exc.headers or None,
|
|
)
|
|
|
|
|
|
if MCP_AVAILABLE:
|
|
__all__ = (
|
|
"_MCP_CREDENTIAL_REQUEST_FIELDS",
|
|
"BlobResourceContents",
|
|
"ListMCPToolsRestAPIResponseObject",
|
|
"ResourceTemplate",
|
|
"TextResourceContents",
|
|
"_McpDeniedDetail",
|
|
"_aggregate_server_key",
|
|
"_build_virtual_call_logging_obj",
|
|
"_check_byok_credential",
|
|
"_client_has_passthrough_authorization",
|
|
"_client_has_per_server_auth_header",
|
|
"_dispatch_virtual_mcp_tool",
|
|
"_fire_mcp_tool_call_logging",
|
|
"_get_allowed_mcp_servers",
|
|
"_get_allowed_mcp_servers_from_mcp_server_names",
|
|
"_get_byok_credential",
|
|
"_get_prompts_from_mcp_servers",
|
|
"_get_resource_templates_from_mcp_servers",
|
|
"_get_resources_from_mcp_servers",
|
|
"_get_standard_logging_mcp_tool_call",
|
|
"_get_tools_from_mcp_servers",
|
|
"_get_user_oauth_extra_headers_from_db",
|
|
"_handle_local_mcp_tool",
|
|
"_handle_managed_mcp_tool",
|
|
"_http_detail_message",
|
|
"_invalidate_byok_cred_cache",
|
|
"_list_mcp_prompts",
|
|
"_list_mcp_resource_templates",
|
|
"_list_mcp_resources",
|
|
"_list_mcp_tools",
|
|
"_list_tools_before_first_call",
|
|
"_mcp_session_id_from_headers",
|
|
"_merge_gateway_initialize_instructions",
|
|
"_prefetch_oauth_creds_for_user",
|
|
"_prepare_mcp_server_headers",
|
|
"_raise_if_initialize_grants_no_mcp_servers",
|
|
"_redact_mcp_resource_url",
|
|
"_resolve_display_name_to_original",
|
|
"_run_post_mcp_call_guardrails",
|
|
"_server_answers_to",
|
|
"_tool_name_matches",
|
|
"apply_display_name_overrides",
|
|
"call_mcp_tool",
|
|
"execute_mcp_tool",
|
|
"filter_tools_by_allowed_tools",
|
|
"filter_tools_by_key_team_permissions",
|
|
"fire_mcp_tool_call_failure_logging",
|
|
"global_mcp_server_manager",
|
|
"mcp_get_prompt",
|
|
"mcp_read_resource",
|
|
"raise_denied_scoped_mcp_access",
|
|
"raise_generic_mcp_access_denial",
|
|
"redact_mcp_resource_url",
|
|
)
|
|
from mcp.server import Server
|
|
|
|
# Import auth context variables and middleware
|
|
from mcp.server.auth.middleware.auth_context import (
|
|
AuthContextMiddleware,
|
|
auth_context_var,
|
|
)
|
|
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
|
from mcp.server.auth.provider import AccessToken
|
|
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
|
|
from mcp.server.lowlevel.server import NotificationOptions
|
|
from mcp.server.lowlevel.server import Server as MCPProtocolServer
|
|
from mcp.server.models import InitializationOptions
|
|
from mcp.shared.exceptions import MCPError
|
|
from mcp.shared.inbound import InboundLadderRejection, classify_inbound_request, find_duplicated_routing_header
|
|
from mcp.types import (
|
|
INVALID_PARAMS,
|
|
METHOD_NOT_FOUND,
|
|
CallToolRequest,
|
|
GetPromptRequest,
|
|
JSONRPCRequest,
|
|
ListPromptsRequest,
|
|
ListResourcesRequest,
|
|
ListResourceTemplatesRequest,
|
|
ListToolsRequest,
|
|
ReadResourceRequest,
|
|
)
|
|
from mcp_types.methods import CLIENT_REQUESTS, SPEC_CLIENT_METHODS
|
|
|
|
from litellm.proxy.mcp_server import operations
|
|
from litellm.proxy.mcp_server.contracts import OperationContext
|
|
from litellm.proxy.mcp_server.interactions import InteractionOperation
|
|
from litellm.proxy.mcp_server.operations import (
|
|
_invalidate_byok_cred_cache,
|
|
_mcp_session_id_from_headers,
|
|
)
|
|
from litellm.proxy.mcp_server.result_conversion import wire_compat_for
|
|
|
|
try:
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
|
except ImportError:
|
|
StreamableHTTPSessionManager = None
|
|
from mcp.types import (
|
|
INVALID_REQUEST,
|
|
CallToolRequestParams,
|
|
CallToolResult,
|
|
GetPromptRequestParams,
|
|
Implementation,
|
|
InitializeRequest,
|
|
InputRequiredResult,
|
|
ListPromptsResult,
|
|
ListResourcesResult,
|
|
ListResourceTemplatesResult,
|
|
ListToolsResult,
|
|
PaginatedRequestParams,
|
|
ReadResourceRequestParams,
|
|
)
|
|
|
|
from litellm.proxy.mcp_server.auth.litellm_auth_handler import (
|
|
MCPAuthenticatedUser,
|
|
)
|
|
from litellm.proxy.mcp_server.capabilities import (
|
|
GatewayVersionPolicy,
|
|
configured_versions,
|
|
disabled_capabilities,
|
|
disabled_operations,
|
|
)
|
|
from litellm.proxy.mcp_server.mcp_server_manager import (
|
|
MCPServerManager,
|
|
global_mcp_server_manager,
|
|
is_legacy_upstream_delegate,
|
|
)
|
|
|
|
######################################################
|
|
############ MCP Tools List REST API Response Object #
|
|
# Defined here because we don't want to add `mcp` as a
|
|
# required dependency for `litellm` pip package
|
|
######################################################
|
|
from litellm.proxy.mcp_server.operations import (
|
|
ListMCPToolsRestAPIResponseObject,
|
|
)
|
|
from litellm.proxy.mcp_server.sse_transport import SseServerTransport
|
|
|
|
def _gateway_create_initialization_options(
|
|
self,
|
|
notification_options: NotificationOptions | None = None,
|
|
experimental_capabilities: dict[str, dict[str, object]] | None = None,
|
|
extensions: dict[str, dict[str, object]] | None = None,
|
|
) -> InitializationOptions:
|
|
base_options: Final = Server.create_initialization_options(
|
|
self,
|
|
notification_options=notification_options
|
|
or NotificationOptions(
|
|
tools_changed=not _mcp_proxy_mode.get(),
|
|
prompts_changed=not _mcp_proxy_mode.get(),
|
|
resources_changed=not _mcp_proxy_mode.get(),
|
|
),
|
|
experimental_capabilities=experimental_capabilities or {},
|
|
extensions=extensions,
|
|
)
|
|
opts: Final = (
|
|
base_options.model_copy(
|
|
update={
|
|
"capabilities": base_options.capabilities.model_copy(update={"prompts": None, "resources": None})
|
|
}
|
|
)
|
|
if _mcp_proxy_mode.get()
|
|
else base_options
|
|
)
|
|
updates: Final[dict[str, str]] = {}
|
|
merged: Final = mcp_gateway_initialize_instructions.get()
|
|
if merged is not None:
|
|
updates["instructions"] = merged
|
|
scoped_server_name: Final = mcp_gateway_server_name.get()
|
|
if scoped_server_name is not None:
|
|
updates["server_name"] = scoped_server_name
|
|
return opts.model_copy(update=updates) if updates else opts
|
|
|
|
########################################################
|
|
############ Initialize the MCP Server #################
|
|
########################################################
|
|
server: Final = MCPProtocolServer[object](
|
|
name=LITELLM_MCP_SERVER_NAME,
|
|
version=LITELLM_MCP_SERVER_VERSION,
|
|
)
|
|
|
|
def _sdk_rejection_code(ctx: ServerRequestContext, error: MCPError | ValidationError) -> int | None:
|
|
if ctx.request_id is None:
|
|
return None
|
|
entry: Final = server.get_request_handler(ctx.method)
|
|
if entry is None:
|
|
if ctx.method == "initialize":
|
|
return None
|
|
if isinstance(error, ValidationError):
|
|
return INVALID_PARAMS
|
|
if error.code == METHOD_NOT_FOUND:
|
|
return METHOD_NOT_FOUND
|
|
return None
|
|
if not isinstance(error, ValidationError):
|
|
unavailable_for_version: Final = (
|
|
ctx.method in SPEC_CLIENT_METHODS and (ctx.method, ctx.protocol_version) not in CLIENT_REQUESTS
|
|
)
|
|
return METHOD_NOT_FOUND if error.code == METHOD_NOT_FOUND and unavailable_for_version else None
|
|
try:
|
|
entry.params_type.model_validate({} if ctx.params is None else ctx.params, by_name=False)
|
|
except ValidationError:
|
|
return INVALID_PARAMS
|
|
return None
|
|
|
|
async def _log_sdk_rejection(ctx: ServerRequestContext, call_next: CallNext) -> HandlerResult:
|
|
start_time: Final = datetime.now() # noqa: DTZ005 # failure logging uses naive timestamps
|
|
try:
|
|
return await call_next(ctx)
|
|
except (MCPError, ValidationError) as error:
|
|
try:
|
|
code: Final = _sdk_rejection_code(ctx, error)
|
|
if code is not None:
|
|
token: Final = active_mcp_request_ctx_var.set(ctx)
|
|
try:
|
|
auth, *_ = await get_or_extract_auth_context()
|
|
finally:
|
|
active_mcp_request_ctx_var.reset(token)
|
|
if auth is not None:
|
|
await operations.log_mcp_sdk_rejection(
|
|
method=ctx.method,
|
|
request_id=ctx.request_id,
|
|
error_code=code,
|
|
start_time=start_time,
|
|
user_api_key_auth=auth,
|
|
)
|
|
except Exception: # noqa: BLE001 # logging must not replace the SDK response or expose rejected input
|
|
verbose_logger.warning("Failed to log MCP SDK rejection")
|
|
raise
|
|
|
|
server.middleware.extend((finish_scope_response, GatewayVersionPolicy(), _log_sdk_rejection))
|
|
server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server)
|
|
sse: Final[SseServerTransport] = SseServerTransport("/sse/messages")
|
|
|
|
# Create session managers
|
|
session_manager_stateless: Final = StreamableHTTPSessionManager(
|
|
app=server,
|
|
event_store=None,
|
|
json_response=False, # enables SSE streaming
|
|
stateless=True,
|
|
)
|
|
|
|
session_manager_stateful: Final = StreamableHTTPSessionManager(
|
|
app=server,
|
|
event_store=None, # TODO: Add EventStore for reconnection/event replay if needed
|
|
json_response=False, # enables SSE streaming
|
|
stateless=False,
|
|
)
|
|
_stateful_session_auth_contexts: Final[dict[str, MCPAuthenticatedUser]] = {}
|
|
_stateful_session_auth_context_last_seen: Final[dict[str, float]] = {}
|
|
# Maps session_id -> owner identifier (hashed API key/token) so we can
|
|
# reject requests that supply a session_id created by a different caller.
|
|
# Without this, a leaked mcp-session-id could be driven (or terminated)
|
|
# by any other authenticated proxy user.
|
|
_stateful_session_owners: Final[dict[str, str]] = {}
|
|
# Per-session lock that serializes ``handle_request`` for the same
|
|
# mcp-session-id. The stored ``MCPAuthenticatedUser`` is mutated in place
|
|
# by ``_update_auth_context`` each request; without this lock, two
|
|
# concurrent requests on the same session would clobber each other's
|
|
# auth headers / mcp_servers / oauth state while in-flight callbacks are
|
|
# still reading the shared object.
|
|
_stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
|
|
_stateful_session_active_request_counts: Final[dict[str, int]] = {}
|
|
_stateful_session_client_info: Final[dict[str, Implementation]] = {} # mutable-ok: cleared on session teardown
|
|
_admin_terminated_session_ids: Final[dict[str, float]] = {} # mutable-ok: admin-closed id -> last replay
|
|
|
|
class _TerminableTransport(Protocol):
|
|
async def terminate(self) -> None: ...
|
|
|
|
class _TransportRegistry(Protocol):
|
|
def __contains__(self, session_id: object, /) -> bool: ...
|
|
|
|
def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ...
|
|
|
|
def _stateful_server_instances() -> _TransportRegistry:
|
|
return getattr(session_manager_stateful, "_server_instances", {})
|
|
|
|
def _remove_stateful_session_tracking(session_id: str) -> None:
|
|
_stateful_session_auth_contexts.pop(session_id, None)
|
|
_stateful_session_auth_context_last_seen.pop(session_id, None)
|
|
_stateful_session_owners.pop(session_id, None)
|
|
_stateful_session_locks.pop(session_id, None)
|
|
_stateful_session_active_request_counts.pop(session_id, None)
|
|
_stateful_session_client_info.pop(session_id, None)
|
|
operations.global_mcp_server_manager.release_upstream_sessions(session_id)
|
|
|
|
# Keep this alias so existing references to session_manager still work
|
|
session_manager: Final = session_manager_stateless
|
|
|
|
# Context managers for proper lifecycle management
|
|
_session_manager_cm = None
|
|
_session_manager_stateful_cm = None
|
|
_stateful_auth_context_cleanup_task: asyncio.Task | None = None
|
|
|
|
async def _purge_expired_stateful_session_auth_contexts(
|
|
now: float | None = None,
|
|
) -> None:
|
|
"""Terminate expired stateful sessions and drop their auth contexts."""
|
|
now = time.monotonic() if now is None else now
|
|
server_instances: Final = _stateful_server_instances()
|
|
expired_session_ids: Final[list[str]] = []
|
|
for session_id, last_seen in _stateful_session_auth_context_last_seen.items():
|
|
if _stateful_session_active_request_counts.get(session_id, 0) > 0:
|
|
continue
|
|
if now - last_seen >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS or session_id not in server_instances:
|
|
expired_session_ids.append(session_id)
|
|
|
|
for session_id in expired_session_ids:
|
|
# Re-check the active-request count immediately before tearing
|
|
# the session down. ``await transport.terminate()`` yields to
|
|
# the event loop, so a request that started after the first
|
|
# collection pass could otherwise observe its transport being
|
|
# ripped out from under it mid-flight.
|
|
if _stateful_session_active_request_counts.get(session_id, 0) > 0:
|
|
continue
|
|
# Pop transport + terminate BEFORE removing owner/auth tracking.
|
|
# Reversing the order avoids a window where ``_stateful_session_owners``
|
|
# is empty but ``server_instances`` still serves the session — a
|
|
# concurrent request in that window would observe ``expected_owner
|
|
# is None`` and bypass the owner-binding check.
|
|
transport = server_instances.pop(session_id, None)
|
|
if transport is not None:
|
|
await transport.terminate()
|
|
_remove_stateful_session_tracking(session_id)
|
|
|
|
for session_id in list(_stateful_session_auth_context_last_seen):
|
|
if session_id not in _stateful_session_auth_contexts:
|
|
_remove_stateful_session_tracking(session_id)
|
|
_forget_expired_admin_terminated_session_ids(now)
|
|
|
|
async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool:
|
|
"""
|
|
Bound the number of concurrent stateful sessions a single caller holds
|
|
before routing a new ``initialize`` to the stateful manager.
|
|
|
|
Evicts the caller's *own* oldest idle sessions (no in-flight requests)
|
|
to make room, so a busy-but-legitimate client keeps its newest sessions
|
|
and other callers are never affected. Returns ``True`` if the new
|
|
session may proceed, or ``False`` when the caller is already at the cap
|
|
with every session in flight (the new ``initialize`` should be rejected).
|
|
"""
|
|
server_instances: Final = _stateful_server_instances()
|
|
|
|
def _owned_live_session_ids() -> list[str]:
|
|
return [
|
|
session_id
|
|
for session_id, session_owner in _stateful_session_owners.items()
|
|
if session_owner == owner and session_id in server_instances
|
|
]
|
|
|
|
owned: Final = _owned_live_session_ids()
|
|
if len(owned) < _MAX_STATEFUL_SESSIONS_PER_OWNER:
|
|
return True
|
|
|
|
for session_id in sorted(
|
|
owned,
|
|
key=lambda sid: _stateful_session_auth_context_last_seen.get(sid, 0.0),
|
|
):
|
|
if len(_owned_live_session_ids()) < _MAX_STATEFUL_SESSIONS_PER_OWNER:
|
|
break
|
|
if _stateful_session_active_request_counts.get(session_id, 0) > 0:
|
|
continue
|
|
transport = server_instances.pop(session_id, None)
|
|
if transport is not None:
|
|
await transport.terminate()
|
|
_remove_stateful_session_tracking(session_id)
|
|
|
|
return len(_owned_live_session_ids()) < _MAX_STATEFUL_SESSIONS_PER_OWNER
|
|
|
|
async def _cleanup_expired_stateful_session_auth_contexts() -> None:
|
|
while True:
|
|
await asyncio.sleep(_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS)
|
|
try:
|
|
await _purge_expired_stateful_session_auth_contexts()
|
|
except Exception as e:
|
|
verbose_logger.exception("Error cleaning up expired MCP stateful sessions: %s", e)
|
|
|
|
async def initialize_session_managers():
|
|
"""Initialize the session managers. Can be called from main app lifespan."""
|
|
global \
|
|
_SESSION_MANAGERS_INITIALIZED, \
|
|
_session_manager_cm, \
|
|
_session_manager_stateful_cm, \
|
|
_stateful_auth_context_cleanup_task
|
|
|
|
# Use async lock to prevent concurrent initialization
|
|
async with _INITIALIZATION_LOCK:
|
|
if _SESSION_MANAGERS_INITIALIZED:
|
|
return
|
|
|
|
verbose_logger.info("Initializing MCP session managers...")
|
|
|
|
# Start the session managers with context managers
|
|
_session_manager_cm = session_manager_stateless.run()
|
|
_session_manager_stateful_cm = session_manager_stateful.run()
|
|
|
|
# Enter the context managers
|
|
await _session_manager_cm.__aenter__()
|
|
await _session_manager_stateful_cm.__aenter__()
|
|
_stateful_auth_context_cleanup_task = asyncio.create_task(_cleanup_expired_stateful_session_auth_contexts())
|
|
|
|
_SESSION_MANAGERS_INITIALIZED = True
|
|
verbose_logger.info("MCP Server started with StreamableHTTP and SSE session managers!")
|
|
|
|
async def shutdown_session_managers():
|
|
"""Shutdown the session managers."""
|
|
global \
|
|
_SESSION_MANAGERS_INITIALIZED, \
|
|
_session_manager_cm, \
|
|
_session_manager_stateful_cm, \
|
|
_stateful_auth_context_cleanup_task
|
|
|
|
if _SESSION_MANAGERS_INITIALIZED:
|
|
verbose_logger.info("Shutting down MCP session managers...")
|
|
|
|
try:
|
|
if _stateful_auth_context_cleanup_task:
|
|
_stateful_auth_context_cleanup_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await _stateful_auth_context_cleanup_task
|
|
await operations.global_mcp_server_manager.close_all_upstream_sessions()
|
|
if _session_manager_stateful_cm:
|
|
await _session_manager_stateful_cm.__aexit__(None, None, None)
|
|
if _session_manager_cm:
|
|
await _session_manager_cm.__aexit__(None, None, None)
|
|
except Exception as e:
|
|
verbose_logger.exception("Error during session manager shutdown: %s", e)
|
|
|
|
_session_manager_cm = None
|
|
_session_manager_stateful_cm = None
|
|
_stateful_auth_context_cleanup_task = None
|
|
_SESSION_MANAGERS_INITIALIZED = False
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(app) -> AsyncIterator[None]:
|
|
"""Application lifespan context manager."""
|
|
await initialize_session_managers()
|
|
try:
|
|
yield
|
|
finally:
|
|
await shutdown_session_managers()
|
|
|
|
########################################################
|
|
############### MCP Server Routes #######################
|
|
########################################################
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _legacy_operation_context(ctx: ServerRequestContext, *, trace: bool) -> AsyncGenerator[OperationContext]:
|
|
if ctx.protocol_version != "2026-07-28" and ctx.method in ("tools/list", "prompts/list", "resources/list"):
|
|
await subscriptions.wait_initialized(ctx)
|
|
with contextlib.ExitStack() as cleanup:
|
|
cleanup.callback(active_mcp_request_ctx_var.reset, active_mcp_request_ctx_var.set(ctx))
|
|
cleanup.callback(active_mcp_session_var.reset, active_mcp_session_var.set(ctx.session))
|
|
if trace:
|
|
cleanup.callback(
|
|
_otel_reset_mcp_trace_carrier, _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(ctx))
|
|
)
|
|
cleanup.callback(
|
|
_otel_reset_mcp_transport_span, _otel_set_mcp_transport_span(_otel_transport_span_from_message(ctx))
|
|
)
|
|
cleanup.callback(_otel_reset_mcp_request_destinations, _otel_set_mcp_request_destinations(ctx))
|
|
(
|
|
auth,
|
|
token,
|
|
servers,
|
|
server_headers,
|
|
oauth_headers,
|
|
headers,
|
|
client_ip,
|
|
) = await get_or_extract_auth_context()
|
|
request: Final[object] = getattr(ctx, "request", None)
|
|
yield operations.prepare_context(
|
|
auth,
|
|
token,
|
|
servers,
|
|
server_headers,
|
|
oauth_headers,
|
|
headers,
|
|
client_ip,
|
|
_mcp_proxy_mode.get(),
|
|
wire_compat_for(ctx.protocol_version),
|
|
ctx.protocol_version,
|
|
incoming_bearer_token=(
|
|
MCPRequestHandler.get_guardrail_bearer_token(request.headers, auth)
|
|
if isinstance(request, StarletteRequest)
|
|
else None
|
|
),
|
|
)
|
|
|
|
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult:
|
|
try:
|
|
async with _legacy_operation_context(ctx, trace=True) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
ListToolsRequest(params=params), context
|
|
)
|
|
except MCPError:
|
|
raise
|
|
except HTTPException as exc:
|
|
raise MCPError(code=INVALID_REQUEST, message=operations._http_detail_message(exc.detail)) from exc
|
|
except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures
|
|
verbose_logger.exception("Error in list_tools endpoint: %s", exc)
|
|
return ListToolsResult(tools=[])
|
|
|
|
def _capture_host_progress_callback(ctx: ServerRequestContext) -> Callable | None:
|
|
"""Return a progress-forwarding callback bound to the host MCP session.
|
|
|
|
Returns ``None`` when the host did not supply a progress token.
|
|
"""
|
|
host_ctx: Final = ctx
|
|
|
|
if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta):
|
|
return None
|
|
host_token: Final = host_ctx.meta.get("progress_token")
|
|
if host_token is None or not (hasattr(host_ctx, "session") and host_ctx.session):
|
|
return None
|
|
host_session: Final = host_ctx.session
|
|
|
|
async def forward_progress(progress: float, total: float | None, message: str | None = None):
|
|
"""Forward progress notifications from external MCP to Host"""
|
|
try:
|
|
await host_session.report_progress(progress=progress, total=total, message=message)
|
|
verbose_logger.debug("Forwarded progress %s/%s to Host", progress, total)
|
|
except Exception as e:
|
|
verbose_logger.error("Failed to forward progress to Host: %s", e)
|
|
|
|
verbose_logger.debug("Host progressToken captured: %s...", str(host_token)[:8])
|
|
return forward_progress
|
|
|
|
def _reject_mcp_proxy_operation() -> NoReturn:
|
|
from mcp.shared.exceptions import MCPError
|
|
from mcp.types import METHOD_NOT_FOUND
|
|
|
|
raise MCPError(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy")
|
|
|
|
from litellm.proxy.mcp_server.operations import (
|
|
_build_virtual_call_logging_obj,
|
|
_dispatch_virtual_mcp_tool,
|
|
)
|
|
|
|
async def mcp_server_tool_call(
|
|
ctx: ServerRequestContext, params: CallToolRequestParams
|
|
) -> CallToolResult | InputRequiredResult:
|
|
async with _legacy_operation_context(ctx, trace=True) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
CallToolRequest(params=params), context
|
|
)
|
|
|
|
async def list_prompts(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListPromptsResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
try:
|
|
async with _legacy_operation_context(ctx, trace=False) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
ListPromptsRequest(params=params), context
|
|
)
|
|
except MCPError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures
|
|
verbose_logger.exception("Error in list_prompts endpoint: %s", exc)
|
|
return ListPromptsResult(prompts=[])
|
|
|
|
async def get_prompt(
|
|
ctx: ServerRequestContext, params: GetPromptRequestParams
|
|
) -> GetPromptResult | InputRequiredResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
async with _legacy_operation_context(ctx, trace=False) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
GetPromptRequest(params=params), context
|
|
)
|
|
|
|
async def list_resources(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListResourcesResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
try:
|
|
async with _legacy_operation_context(ctx, trace=False) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
ListResourcesRequest(params=params), context
|
|
)
|
|
except MCPError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures
|
|
verbose_logger.exception("Error in list_resources endpoint: %s", exc)
|
|
return ListResourcesResult(resources=[])
|
|
|
|
async def list_resource_templates(
|
|
ctx: ServerRequestContext, params: PaginatedRequestParams
|
|
) -> ListResourceTemplatesResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
try:
|
|
async with _legacy_operation_context(ctx, trace=False) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
ListResourceTemplatesRequest(params=params), context
|
|
)
|
|
except MCPError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures
|
|
verbose_logger.exception("Error in list_resource_templates endpoint: %s", exc)
|
|
return ListResourceTemplatesResult(resource_templates=[])
|
|
|
|
async def read_resource(
|
|
ctx: ServerRequestContext, params: ReadResourceRequestParams
|
|
) -> ReadResourceResult | InputRequiredResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
async with _legacy_operation_context(ctx, trace=False) as context:
|
|
return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute(
|
|
ReadResourceRequest(params=params), context
|
|
)
|
|
|
|
async def discover(ctx: ServerRequestContext, params: RequestParams) -> DiscoverResult:
|
|
async with _legacy_operation_context(ctx, trace=False) as context:
|
|
return await operations.GatewayOperations().execute(DiscoverRequest(params=params), context)
|
|
|
|
from mcp.types import (
|
|
EmptyResult,
|
|
NotificationParams,
|
|
SubscribeRequestParams,
|
|
SubscriptionsListenRequestParams,
|
|
SubscriptionsListenResult,
|
|
UnsubscribeRequestParams,
|
|
)
|
|
|
|
from litellm.proxy.mcp_server.subscriptions import GatewaySubscriptions, SubscriptionAuthorization
|
|
|
|
subscriptions: Final = GatewaySubscriptions()
|
|
legacy_subscription_request: Final[contextvars.ContextVar[StarletteRequest | None]] = contextvars.ContextVar(
|
|
"legacy_subscription_request", default=None
|
|
)
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _subscription_authorization(
|
|
ctx: ServerRequestContext[object, StarletteRequest],
|
|
) -> AsyncGenerator[SubscriptionAuthorization]:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
request: Final[StarletteRequest | None] = (
|
|
ctx.request if isinstance(ctx.request, StarletteRequest) else legacy_subscription_request.get()
|
|
)
|
|
if not isinstance(request, StarletteRequest):
|
|
raise HTTPException(status_code=400, detail="Subscriptions require an HTTP request")
|
|
async with _legacy_operation_context(ctx, trace=False) as initial:
|
|
|
|
async def authorize() -> OperationContext:
|
|
auth, token, servers, server_headers, oauth_headers, headers = await extract_mcp_auth_context(
|
|
request.scope, request.url.path
|
|
)
|
|
return operations.prepare_context(
|
|
auth,
|
|
token,
|
|
servers,
|
|
server_headers,
|
|
oauth_headers,
|
|
headers,
|
|
initial.client_ip,
|
|
False,
|
|
initial.wire_compat,
|
|
initial.protocol_version,
|
|
incoming_bearer_token=MCPRequestHandler.get_guardrail_bearer_token(request.headers, auth),
|
|
)
|
|
|
|
authorization: Final = SubscriptionAuthorization(initial, authorize, _owner_fingerprint_for)
|
|
yield authorization
|
|
|
|
async def listen_subscriptions(
|
|
ctx: ServerRequestContext, params: SubscriptionsListenRequestParams
|
|
) -> SubscriptionsListenResult:
|
|
async with _subscription_authorization(ctx) as authorization:
|
|
return await subscriptions.listen(ctx, params, authorization, authorization.owner)
|
|
|
|
async def legacy_notifications(ctx: ServerRequestContext, params: NotificationParams) -> None:
|
|
if _mcp_proxy_mode.get():
|
|
return
|
|
async with subscriptions.initializing(ctx) as start:
|
|
if start:
|
|
async with _subscription_authorization(ctx) as authorization:
|
|
await subscriptions.legacy_session(ctx, authorization)
|
|
|
|
async def subscribe_resource(ctx: ServerRequestContext, params: SubscribeRequestParams) -> EmptyResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
await subscriptions.legacy_resource(ctx, str(params.uri), subscribe=True)
|
|
return EmptyResult()
|
|
|
|
async def unsubscribe_resource(ctx: ServerRequestContext, params: UnsubscribeRequestParams) -> EmptyResult:
|
|
if _mcp_proxy_mode.get():
|
|
_reject_mcp_proxy_operation()
|
|
await subscriptions.legacy_resource(ctx, str(params.uri), subscribe=False)
|
|
return EmptyResult()
|
|
|
|
server.add_notification_handler("notifications/initialized", NotificationParams, legacy_notifications)
|
|
server.add_request_handler("resources/subscribe", SubscribeRequestParams, subscribe_resource)
|
|
server.add_request_handler("resources/unsubscribe", UnsubscribeRequestParams, unsubscribe_resource)
|
|
|
|
server.add_request_handler("subscriptions/listen", SubscriptionsListenRequestParams, listen_subscriptions)
|
|
|
|
server.add_request_handler("server/discover", RequestParams, discover)
|
|
server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools)
|
|
server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call)
|
|
server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts)
|
|
server.add_request_handler("prompts/get", GetPromptRequestParams, get_prompt)
|
|
server.add_request_handler("resources/list", PaginatedRequestParams, list_resources)
|
|
server.add_request_handler("resources/templates/list", PaginatedRequestParams, list_resource_templates)
|
|
server.add_request_handler("resources/read", ReadResourceRequestParams, read_resource)
|
|
|
|
########################################################
|
|
############ End of MCP Server Routes ##################
|
|
########################################################
|
|
|
|
########################################################
|
|
############ Helper Functions ##########################
|
|
########################################################
|
|
|
|
from litellm.proxy.mcp_server.operations import (
|
|
_client_has_passthrough_authorization,
|
|
_client_has_per_server_auth_header,
|
|
_get_allowed_mcp_servers,
|
|
_get_allowed_mcp_servers_from_mcp_server_names,
|
|
_get_user_oauth_extra_headers_from_db,
|
|
_http_detail_message,
|
|
_McpDeniedDetail,
|
|
_merge_gateway_initialize_instructions,
|
|
_prefetch_oauth_creds_for_user,
|
|
_prepare_mcp_server_headers,
|
|
_raise_if_initialize_grants_no_mcp_servers,
|
|
_server_answers_to,
|
|
_tool_name_matches,
|
|
apply_display_name_overrides,
|
|
filter_tools_by_allowed_tools,
|
|
raise_denied_scoped_mcp_access,
|
|
raise_generic_mcp_access_denial,
|
|
)
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _gateway_initialize_instructions_request_scope(
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
mcp_servers: list[str] | None,
|
|
client_ip: str | None,
|
|
scoped_server_endpoint: bool = False,
|
|
is_initialize: bool = False,
|
|
) -> AsyncIterator[None]:
|
|
allowed: Final = await operations._get_allowed_mcp_servers(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_servers=mcp_servers,
|
|
client_ip=client_ip,
|
|
)
|
|
if is_initialize:
|
|
await operations._raise_if_initialize_grants_no_mcp_servers(
|
|
allowed, user_api_key_auth, mcp_servers, client_ip
|
|
)
|
|
if allowed:
|
|
# return_exceptions=True: a per-server probe failure (incl. CancelledError
|
|
# bubbled from anyio task group teardown on connection refused) must not
|
|
# cancel sibling probes or 500 the gateway initialize request.
|
|
await asyncio.gather(
|
|
*[
|
|
operations.global_mcp_server_manager.ensure_upstream_initialize_instructions_cached(s)
|
|
for s in allowed
|
|
if s is not None
|
|
],
|
|
return_exceptions=True,
|
|
)
|
|
merged: Final = operations._merge_gateway_initialize_instructions(allowed_mcp_servers=allowed)
|
|
scoped_server_name = None
|
|
if scoped_server_endpoint and len(allowed) == 1:
|
|
scoped_server: Final = allowed[0]
|
|
scoped_server_name = (
|
|
scoped_server.alias or scoped_server.server_name or scoped_server.name or scoped_server.server_id
|
|
)
|
|
instructions_token: Final = mcp_gateway_initialize_instructions.set(merged)
|
|
server_name_token: Final = mcp_gateway_server_name.set(scoped_server_name)
|
|
try:
|
|
yield
|
|
finally:
|
|
mcp_gateway_initialize_instructions.reset(instructions_token)
|
|
mcp_gateway_server_name.reset(server_name_token)
|
|
|
|
from litellm.proxy.mcp_server.operations import (
|
|
_MCP_CREDENTIAL_REQUEST_FIELDS,
|
|
_aggregate_server_key,
|
|
_check_byok_credential,
|
|
_fire_mcp_tool_call_logging,
|
|
_get_byok_credential,
|
|
_get_prompts_from_mcp_servers,
|
|
_get_resource_templates_from_mcp_servers,
|
|
_get_resources_from_mcp_servers,
|
|
_get_standard_logging_mcp_tool_call,
|
|
_get_tools_from_mcp_servers,
|
|
_handle_local_mcp_tool,
|
|
_handle_managed_mcp_tool,
|
|
_list_mcp_prompts,
|
|
_list_mcp_resource_templates,
|
|
_list_mcp_resources,
|
|
_list_mcp_tools,
|
|
_list_tools_before_first_call,
|
|
_resolve_display_name_to_original,
|
|
_run_post_mcp_call_guardrails,
|
|
call_mcp_tool,
|
|
execute_mcp_tool,
|
|
filter_tools_by_key_team_permissions,
|
|
fire_mcp_tool_call_failure_logging,
|
|
mcp_get_prompt,
|
|
mcp_read_resource,
|
|
)
|
|
|
|
def _get_mcp_servers_in_path(path: str) -> list[str] | None:
|
|
"""
|
|
Get the MCP servers from the path
|
|
"""
|
|
import re
|
|
|
|
if path.rstrip("/") in ("/mcp/sse", "/mcp/sse/messages"):
|
|
return None
|
|
mcp_servers_from_path: list[str] | None = None
|
|
segments: Final = [s for s in path.split("/") if s]
|
|
if len(segments) >= 2 and segments[1] == "mcp" and segments[0] != "mcp":
|
|
return [segments[0]]
|
|
|
|
# Match /mcp/<servers_and_maybe_path>
|
|
# Where servers can be comma-separated list of server names
|
|
# Server names can contain slashes (e.g., "custom_solutions/user_123")
|
|
mcp_path_match: Final = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path)
|
|
if mcp_path_match:
|
|
servers_and_path: Final = mcp_path_match.group(1)
|
|
|
|
if servers_and_path:
|
|
# Check if it contains commas (comma-separated servers)
|
|
if "," in servers_and_path:
|
|
# For comma-separated, look for a path at the end
|
|
# Common patterns: /tools, /chat/completions, etc.
|
|
path_match: Final = re.search(r"/([^/,]+(?:/[^/,]+)*)$", servers_and_path)
|
|
if path_match:
|
|
# Path found at the end, remove it from servers
|
|
path_part: Final = "/" + path_match.group(1)
|
|
servers_part: Final = servers_and_path[: -len(path_part)]
|
|
mcp_servers_from_path = [s.strip() for s in servers_part.split(",") if s.strip()]
|
|
else:
|
|
# No path, just comma-separated servers
|
|
mcp_servers_from_path = [s.strip() for s in servers_and_path.split(",") if s.strip()]
|
|
else:
|
|
# Single server case - use regex approach for server/path separation
|
|
# This handles cases like "custom_solutions/user_123/chat/completions"
|
|
# where we want to extract "custom_solutions/user_123" as the server name
|
|
single_server_match: Final = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path)
|
|
if single_server_match:
|
|
server_name: Final = single_server_match.group(1)
|
|
mcp_servers_from_path = [server_name]
|
|
else:
|
|
mcp_servers_from_path = [servers_and_path]
|
|
return mcp_servers_from_path
|
|
|
|
def _load_mcp_client_allowlist() -> MCPClientAllowlist | None:
|
|
from litellm.proxy.proxy_server import general_settings
|
|
|
|
return load_mcp_client_allowlist(general_settings)
|
|
|
|
def reject_disallowed_mcp_client(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth | None) -> None:
|
|
"""Gate every MCP tool surface on ``mcp_allowed_clients``; the dashboard's own session is not a client app."""
|
|
if user_api_key_auth is not None and is_ui_session_credential(user_api_key_auth):
|
|
return
|
|
rejection: Final = check_mcp_client_allowed(
|
|
allowlist=_load_mcp_client_allowlist(),
|
|
jwt_claims=user_api_key_auth.jwt_claims if user_api_key_auth is not None else None,
|
|
headers=headers,
|
|
)
|
|
if rejection is None:
|
|
return
|
|
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.
|
|
Returns: (user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers)
|
|
"""
|
|
mcp_servers_from_path: Final = _get_mcp_servers_in_path(path)
|
|
if mcp_servers_from_path is not None:
|
|
(
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
_,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
) = await MCPRequestHandler.process_mcp_request(scope)
|
|
mcp_servers = mcp_servers_from_path
|
|
else:
|
|
(
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
mcp_servers,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
) = await MCPRequestHandler.process_mcp_request(scope)
|
|
return (
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
mcp_servers,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
)
|
|
|
|
def _get_session_id_from_scope(scope: Scope) -> str | None:
|
|
"""
|
|
Extract mcp-session-id from ASGI scope headers.
|
|
Returns None if not present.
|
|
"""
|
|
scope_headers: Final = _raw_header_pairs(scope)
|
|
for header_name, header_value in scope_headers:
|
|
name = header_name if isinstance(header_name, bytes) else header_name.encode()
|
|
if name.lower() == b"mcp-session-id":
|
|
return header_value.decode() if isinstance(header_value, bytes) else str(header_value)
|
|
return None
|
|
|
|
def _owner_fingerprint_for(
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
oauth2_headers: Mapping[str, str] | None = None,
|
|
client_ip: str | None = None,
|
|
) -> str:
|
|
"""
|
|
Stable, non-reversible identifier for the caller used to bind an
|
|
mcp-session-id to its creator. Hash the resolved credential before
|
|
using it so custom key formats are never stored in cleartext.
|
|
|
|
For OAuth2 passthrough (``UserAPIKeyAuth()`` with no key/user_id),
|
|
the caller's identity is the upstream OAuth bearer; hash it so two
|
|
OAuth callers with different tokens don't both fingerprint to
|
|
``anonymous`` and end up sharing a session.
|
|
|
|
When no caller-identifying credentials are available at all
|
|
(e.g. proxy running without master key, or an unauthenticated
|
|
passthrough path), fall back to the client IP so two unrelated
|
|
anonymous callers from different sources do not collapse to a
|
|
single ``anonymous`` owner and end up able to drive each other's
|
|
stateful sessions. Note: when even client IP is unavailable
|
|
(exotic deployments without trusted X-Forwarded-For and direct
|
|
socket info), the fingerprint degrades to the ``anonymous``
|
|
sentinel and cannot meaningfully protect against another
|
|
unauthenticated caller who learns the session id — owner-binding
|
|
is best-effort in that mode.
|
|
"""
|
|
|
|
def _bytes_for_hash(value: object) -> bytes | None:
|
|
"""Only hash str/bytes secrets; skip mocks and other unexpected types."""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, (bytes, bytearray)):
|
|
return bytes(value)
|
|
if isinstance(value, str):
|
|
return value.encode("utf-8")
|
|
return None
|
|
|
|
if user_api_key_auth is not None:
|
|
key_material: Final = _bytes_for_hash(getattr(user_api_key_auth, "api_key", None))
|
|
if key_material:
|
|
api_key_hash: Final = hashlib.sha256(key_material).hexdigest()
|
|
return f"key:{api_key_hash}"
|
|
uid_material: Final = _bytes_for_hash(getattr(user_api_key_auth, "user_id", None))
|
|
if uid_material:
|
|
user_id_hash: Final = hashlib.sha256(uid_material).hexdigest()
|
|
return f"user:{user_id_hash}"
|
|
if oauth2_headers:
|
|
authz: Final = oauth2_headers.get("Authorization") or oauth2_headers.get("authorization")
|
|
authz_bytes: Final = _bytes_for_hash(authz)
|
|
if authz_bytes:
|
|
return f"oauth:{hashlib.sha256(authz_bytes).hexdigest()}"
|
|
if client_ip and isinstance(client_ip, str):
|
|
return f"ip:{hashlib.sha256(client_ip.encode('utf-8')).hexdigest()}"
|
|
return "anonymous"
|
|
|
|
def _admission_auth_body(body: bytes) -> bytes:
|
|
"""Return only the JSON-RPC method for admission auth."""
|
|
try:
|
|
data: Final = TypeAdapter(dict[str, object]).validate_json(body)
|
|
except (ValueError, TypeError):
|
|
return b"{}"
|
|
method: Final = data.get("method")
|
|
if not isinstance(method, str):
|
|
return b"{}"
|
|
return json.dumps({"method": method}).encode()
|
|
|
|
def _is_initialize_request(body: bytes) -> bool:
|
|
"""
|
|
Check if the request body is a JSON-RPC initialize method.
|
|
Returns True if method is "initialize", False otherwise or on parse error.
|
|
"""
|
|
if not body:
|
|
return False
|
|
try:
|
|
data: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(body)))
|
|
return data is not None and data.get("method") == "initialize"
|
|
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
|
|
return False
|
|
|
|
def _jsonrpc_request_method(body: bytes) -> str | None:
|
|
try:
|
|
return JSONRPCRequest.model_validate_json(body).method
|
|
except ValidationError:
|
|
return None
|
|
|
|
def _extract_initialize_client_info(body: bytes) -> Implementation | None:
|
|
try:
|
|
return InitializeRequest.model_validate_json(body, by_name=False).params.client_info
|
|
except ValidationError:
|
|
return None
|
|
|
|
def _group_session_counts(
|
|
sessions: Sequence[MCPGatewaySession],
|
|
label_for: Callable[[MCPGatewaySession], str | None],
|
|
) -> tuple[MCPGatewaySessionGroupCount, ...]:
|
|
counts: Final = types.MappingProxyType(Counter(label_for(session) for session in sessions))
|
|
return tuple(
|
|
sorted(
|
|
(MCPGatewaySessionGroupCount(label=label, count=count) for label, count in counts.items()),
|
|
key=lambda group: (-group.count, group.label is None, group.label or ""),
|
|
)
|
|
)
|
|
|
|
def _gateway_session_for(session_id: str, auth_user: MCPAuthenticatedUser, now: float) -> MCPGatewaySession:
|
|
client_info: Final = _stateful_session_client_info.get(session_id)
|
|
key_auth: Final = auth_user.user_api_key_auth
|
|
return MCPGatewaySession(
|
|
session_id_prefix=session_id[:MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH],
|
|
client_name=client_info.name if client_info is not None else None,
|
|
client_version=client_info.version if client_info is not None else None,
|
|
user_id=key_auth.user_id if key_auth is not None else None,
|
|
user_email=key_auth.user_email if key_auth is not None else None,
|
|
key_alias=key_auth.key_alias if key_auth is not None else None,
|
|
team_id=key_auth.team_id if key_auth is not None else None,
|
|
team_alias=key_auth.team_alias if key_auth is not None else None,
|
|
client_ip=auth_user.client_ip,
|
|
idle_seconds=max(0.0, now - _stateful_session_auth_context_last_seen.get(session_id, now)),
|
|
in_flight_requests=_stateful_session_active_request_counts.get(session_id, 0),
|
|
)
|
|
|
|
def get_mcp_gateway_sessions_report(now: float | None = None) -> MCPGatewaySessionsResponse:
|
|
"""Live stateful Streamable HTTP sessions held by this worker process.
|
|
|
|
Only sessions whose transport is still registered with the stateful
|
|
session manager are reported; SSE and stateless requests hold no
|
|
session and are never counted.
|
|
"""
|
|
report_time: Final = time.monotonic() if now is None else now
|
|
live_session_ids: Final = frozenset(_stateful_server_instances())
|
|
sessions: Final = tuple(
|
|
_gateway_session_for(session_id, auth_user, report_time)
|
|
for session_id, auth_user in tuple(_stateful_session_auth_contexts.items())
|
|
if session_id in live_session_ids
|
|
)
|
|
return MCPGatewaySessionsResponse(
|
|
worker_pid=os.getpid(),
|
|
total_sessions=len(sessions),
|
|
by_client=_group_session_counts(sessions, lambda session: session.client_name),
|
|
by_user=_group_session_counts(sessions, lambda session: session.user_id),
|
|
sessions=sessions,
|
|
)
|
|
|
|
def _session_matches_admin_selector(
|
|
session_id: str,
|
|
auth_user: MCPAuthenticatedUser,
|
|
session_id_prefix: str | None,
|
|
user_id: str | None,
|
|
) -> bool:
|
|
if session_id_prefix is not None and not session_id.startswith(session_id_prefix):
|
|
return False
|
|
if user_id is None:
|
|
return True
|
|
key_auth: Final = auth_user.user_api_key_auth
|
|
return key_auth is not None and key_auth.user_id == user_id
|
|
|
|
def _forget_expired_admin_terminated_session_ids(now: float) -> None:
|
|
for session_id in [
|
|
session_id
|
|
for session_id, last_replayed in _admin_terminated_session_ids.items()
|
|
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
|
|
]:
|
|
del _admin_terminated_session_ids[session_id]
|
|
|
|
def _is_admin_terminated_session_id(session_id: str, now: float) -> bool:
|
|
last_replayed: Final = _admin_terminated_session_ids.get(session_id)
|
|
if last_replayed is None:
|
|
return False
|
|
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS:
|
|
del _admin_terminated_session_ids[session_id]
|
|
return False
|
|
_admin_terminated_session_ids[session_id] = now
|
|
return True
|
|
|
|
async def terminate_mcp_gateway_sessions(
|
|
*,
|
|
session_id_prefix: str | None = None,
|
|
user_id: str | None = None,
|
|
) -> MCPGatewaySessionsTerminateResponse:
|
|
"""Force-close every live stateful session on this worker matching the selector.
|
|
|
|
The transport is terminated (open streams close), all per-session
|
|
tracking is dropped, and the id is remembered so a client that keeps
|
|
sending it receives 404 and has to ``initialize`` again, which re-runs
|
|
admission. Only sessions held by this worker process are affected.
|
|
"""
|
|
now: Final = time.monotonic()
|
|
_forget_expired_admin_terminated_session_ids(now)
|
|
server_instances: Final = _stateful_server_instances()
|
|
targets: Final = tuple(
|
|
(session_id, auth_user)
|
|
for session_id, auth_user in tuple(_stateful_session_auth_contexts.items())
|
|
if session_id in server_instances
|
|
and _session_matches_admin_selector(session_id, auth_user, session_id_prefix, user_id)
|
|
)
|
|
terminated: Final = tuple(_gateway_session_for(session_id, auth_user, now) for session_id, auth_user in targets)
|
|
for session_id, _ in targets:
|
|
_admin_terminated_session_ids[session_id] = now
|
|
transport = server_instances.pop(session_id, None)
|
|
_remove_stateful_session_tracking(session_id)
|
|
if transport is not None:
|
|
await transport.terminate()
|
|
verbose_logger.warning("MCP session '%s' terminated by an administrator.", session_id)
|
|
return MCPGatewaySessionsTerminateResponse(
|
|
worker_pid=os.getpid(),
|
|
terminated_sessions=len(terminated),
|
|
sessions=terminated,
|
|
)
|
|
|
|
async def _read_request_body_for_routing(
|
|
receive: Receive,
|
|
*,
|
|
full_body: bool = False,
|
|
already_read: int = 0,
|
|
) -> tuple[list[Message], bytes]:
|
|
"""
|
|
Read just enough of the request body to decide whether this is a
|
|
JSON-RPC ``initialize`` call. Returns the consumed ASGI messages so
|
|
the caller can replay them faithfully to the downstream handler, and
|
|
the peeked body bytes (capped at ``_MCP_ROUTING_PEEK_MAX_BYTES``).
|
|
|
|
Stops reading from the wire as soon as either (a) we have peeked
|
|
``_MCP_ROUTING_PEEK_MAX_BYTES`` of body, or (b) the body is complete.
|
|
The remainder of an oversized body is streamed lazily through
|
|
``_LazyPeekedBody.receive`` in the caller — so an authenticated client cannot
|
|
force the proxy to buffer an arbitrarily large payload just to make a
|
|
routing decision. Modern interaction preflight requests the full body,
|
|
which the SDK's single-exchange transport also requires.
|
|
"""
|
|
consumed_messages: Final[list[Message]] = []
|
|
body_chunks: Final[list[bytes]] = []
|
|
peeked_bytes = already_read
|
|
|
|
while True:
|
|
message = await receive()
|
|
if (
|
|
full_body
|
|
and peeked_bytes + len(message.get("body", b"") or b"")
|
|
> session_manager_stateless.max_request_body_size
|
|
):
|
|
raise HTTPException(status_code=413, detail="Request body too large")
|
|
consumed_messages.append(message)
|
|
|
|
if message.get("type") != "http.request":
|
|
break
|
|
|
|
body: bytes = message.get("body", b"") or b""
|
|
if body:
|
|
# Only retain up to the remaining peek budget for sniffing.
|
|
# The full ``message`` is already in memory (delivered by
|
|
# the ASGI server) and must round-trip to the downstream
|
|
# handler via ``consumed_messages``, but ``body_chunks`` is
|
|
# purely for the JSON-RPC method check — there is no reason
|
|
# to copy a large body frame into a second buffer.
|
|
remaining = len(body) if full_body else _MCP_ROUTING_PEEK_MAX_BYTES - peeked_bytes
|
|
if remaining > 0:
|
|
body_chunks.append(body[:remaining])
|
|
peeked_bytes += min(len(body), remaining)
|
|
|
|
if not message.get("more_body", False):
|
|
break
|
|
|
|
if not full_body and peeked_bytes >= _MCP_ROUTING_PEEK_MAX_BYTES:
|
|
# Stop draining; downstream replay will pull remaining chunks
|
|
# directly from the original `receive` via wrapped_receive.
|
|
break
|
|
|
|
return consumed_messages, b"".join(body_chunks)
|
|
|
|
_PeekedRead = tuple[list[Message], bytes]
|
|
|
|
class _LazyPeekedBody:
|
|
"""Reads the POST body for routing and auth only when first asked, so an
|
|
allowlist rejection never touches the receive channel. Whatever it
|
|
consumed is replayed to the downstream handler through ``receive``,
|
|
which falls through to the wire once the replay buffer is drained."""
|
|
|
|
__slots__ = ("_consumed_messages", "_peeked_body", "_read", "_receive")
|
|
|
|
def __init__(self, receive: Receive) -> None:
|
|
self._receive = receive
|
|
self._read: asyncio.Future[_PeekedRead] | None = None
|
|
self._consumed_messages: list[Message] | None = None
|
|
self._peeked_body: bytes = b""
|
|
|
|
def _start_read(self) -> asyncio.Future[_PeekedRead]:
|
|
if self._read is None:
|
|
self._read = asyncio.ensure_future(_read_request_body_for_routing(self._receive))
|
|
return self._read
|
|
|
|
async def body(self) -> bytes:
|
|
if self._consumed_messages is None:
|
|
self._consumed_messages, self._peeked_body = await self._start_read()
|
|
return self._peeked_body
|
|
|
|
async def admission_body(self) -> bytes:
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.shield(self._start_read()), timeout=MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS
|
|
)
|
|
except asyncio.TimeoutError:
|
|
return b""
|
|
return await self.body()
|
|
|
|
async def full_body(self) -> bytes:
|
|
await self.body()
|
|
consumed: Final = self._consumed_messages or []
|
|
consumed_body: Final = b"".join(message.get("body", b"") or b"" for message in consumed)
|
|
if len(consumed_body) > session_manager_stateless.max_request_body_size:
|
|
raise HTTPException(status_code=413, detail="Request body too large")
|
|
if not consumed or consumed[-1].get("type") != "http.request" or not consumed[-1].get("more_body", False):
|
|
return consumed_body
|
|
rest_messages, rest_body = await _read_request_body_for_routing(
|
|
self._receive, full_body=True, already_read=len(consumed_body)
|
|
)
|
|
self._consumed_messages = [*consumed, *rest_messages]
|
|
return consumed_body + rest_body
|
|
|
|
async def receive(self) -> Message:
|
|
if self._consumed_messages:
|
|
return self._consumed_messages.pop(0)
|
|
return await self._receive()
|
|
|
|
async def _handle_stale_mcp_session(
|
|
scope: Scope,
|
|
receive: Receive,
|
|
send: Send,
|
|
mgr: "StreamableHTTPSessionManager",
|
|
) -> bool:
|
|
"""
|
|
Inspect the incoming ``mcp-session-id`` header **before** the
|
|
request reaches the MCP SDK. If the session is stale (not known
|
|
to this worker), strip the header so the SDK creates a fresh
|
|
stateless session instead of returning a 400.
|
|
|
|
Returns:
|
|
True if the request was fully handled (e.g. DELETE on
|
|
non-existent session). False if the request should continue
|
|
to the session manager.
|
|
|
|
Fixes https://github.com/BerriAI/litellm/issues/20992
|
|
"""
|
|
_mcp_session_header: Final = b"mcp-session-id"
|
|
_headers: Final = _raw_header_pairs(scope)
|
|
|
|
def _normalize_header_name(header_name: object) -> bytes | None:
|
|
if isinstance(header_name, bytes):
|
|
return header_name.lower()
|
|
if isinstance(header_name, str):
|
|
return header_name.lower().encode("utf-8", errors="replace")
|
|
return None
|
|
|
|
_session_id: str | None = None
|
|
for header_name, header_value in _headers:
|
|
if _normalize_header_name(header_name) == _mcp_session_header:
|
|
if isinstance(header_value, bytes):
|
|
_session_id = header_value.decode("utf-8", errors="replace")
|
|
else:
|
|
_session_id = str(header_value)
|
|
break
|
|
|
|
if _session_id is None:
|
|
return False
|
|
|
|
# Check in-memory session tracking
|
|
known_sessions: Final = getattr(mgr, "_server_instances", None)
|
|
# If we cannot inspect known_sessions, let the manager handle it
|
|
if known_sessions is None:
|
|
return False
|
|
|
|
# If session exists in this worker's memory, let the manager handle it
|
|
try:
|
|
if _session_id in known_sessions:
|
|
return False
|
|
except Exception:
|
|
verbose_logger.debug(
|
|
"Unable to inspect active MCP sessions for '%s'. Deferring to session manager.",
|
|
_session_id,
|
|
)
|
|
return False
|
|
|
|
# --- Session not in this worker's memory ---
|
|
method: Final = asgi_scope_fields(scope).get("method", "").upper()
|
|
|
|
if method == "DELETE":
|
|
_remove_stateful_session_tracking(_session_id)
|
|
verbose_logger.info(
|
|
"DELETE request for non-existent MCP session '%s'. Returning success (idempotent DELETE).",
|
|
_session_id,
|
|
)
|
|
success_response: Final = JSONResponse(
|
|
status_code=200,
|
|
content={"message": "Session terminated successfully"},
|
|
)
|
|
await success_response(scope, receive, send)
|
|
return True
|
|
|
|
if _is_admin_terminated_session_id(_session_id, time.monotonic()):
|
|
terminated_response: Final = JSONResponse(
|
|
status_code=404,
|
|
content={
|
|
"error": "Not Found",
|
|
"details": "mcp-session-id was terminated by an administrator. Send initialize to start a new session.",
|
|
},
|
|
)
|
|
await terminated_response(scope, receive, send)
|
|
return True
|
|
|
|
# Non-DELETE: strip stale session ID to allow new session creation
|
|
verbose_logger.warning(
|
|
"MCP session ID '%s' not found in this worker's memory. "
|
|
"Stripping stale header to force new session creation.",
|
|
_session_id,
|
|
)
|
|
scope["headers"] = [(k, v) for k, v in _headers if _normalize_header_name(k) != _mcp_session_header]
|
|
return False
|
|
|
|
async def apply_toolset_scope(
|
|
user_api_key_auth: UserAPIKeyAuth,
|
|
toolset_id: str,
|
|
acting_user: ActingUser = acting_user_auth,
|
|
granted: GrantedToolsetIds = granted_toolset_ids,
|
|
) -> UserAPIKeyAuth:
|
|
"""
|
|
Pin a principal's MCP permissions to a single toolset for /toolset/{name}/mcp.
|
|
|
|
A virtual key (and an admin session) has its object_permission rewritten to
|
|
the toolset's servers and tools. A keyless subject resolves per grant source,
|
|
so a non-admin dashboard session first becomes its admitted user and the
|
|
toolset rides along as ``mcp_toolset_id``, which every source's grant is
|
|
intersected with; a team-granted toolset is served without the user's own
|
|
row capping it.
|
|
|
|
Raises HTTPException(403) unless the principal holds toolset_id through one
|
|
of its grant sources. Admins pass unless the key is bound to a real team.
|
|
"""
|
|
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
|
|
from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import has_team_unbound_mcp_admin_view
|
|
|
|
# A key scoped to no MCP servers opts out of every MCP path. Enforce it
|
|
# here too, since toolset scoping replaces mcp_servers and would otherwise
|
|
# drop the sentinel. Checked before the admin branch, mirroring
|
|
# get_allowed_mcp_servers.
|
|
original_op: Final = user_api_key_auth.object_permission
|
|
if original_op is not None and SpecialMCPServerNames.no_mcp_servers.value in (original_op.mcp_servers or []):
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="API key is scoped to no MCP servers; toolset access is denied.",
|
|
)
|
|
|
|
acting: Final = await acting_user(user_api_key_auth)
|
|
is_admin: Final = has_team_unbound_mcp_admin_view(acting)
|
|
if not is_admin and toolset_id not in await granted(acting):
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail=f"API key does not have access to toolset '{toolset_id}'.",
|
|
)
|
|
if is_mcp_admitted_user_subject(acting):
|
|
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
|
|
)
|
|
):
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail=f"API key does not have access to toolset '{toolset_id}'.",
|
|
)
|
|
return acting.model_copy(update={"mcp_toolset_id": toolset_id})
|
|
|
|
tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(
|
|
toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy
|
|
)
|
|
server_ids: Final = list(tool_permissions.keys())
|
|
existing_op: Final = acting.object_permission
|
|
if existing_op is not None:
|
|
updated_op = existing_op.model_copy(
|
|
update={
|
|
"mcp_servers": server_ids,
|
|
"mcp_tool_permissions": tool_permissions,
|
|
"mcp_toolsets": [],
|
|
# mcp_access_groups is preserved: a key's access-group grants
|
|
# remain valid even when the request is scoped to a single toolset.
|
|
}
|
|
)
|
|
else:
|
|
updated_op = LiteLLM_ObjectPermissionTable(
|
|
object_permission_id="toolset-scope",
|
|
mcp_servers=server_ids,
|
|
mcp_tool_permissions=tool_permissions,
|
|
)
|
|
return acting.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id})
|
|
|
|
_apply_toolset_scope: Final = apply_toolset_scope
|
|
|
|
async def _toolset_server_ids(toolset_id: str) -> set[str]:
|
|
return set(
|
|
await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id])
|
|
)
|
|
|
|
from litellm.proxy.mcp_server.catalog import catalog_operation
|
|
|
|
@catalog_operation(lambda: operations.global_mcp_server_manager)
|
|
async def _raise_preemptive_401_for_unauthenticated_servers(
|
|
scope: Scope,
|
|
mcp_servers: list[str] | None,
|
|
oauth2_headers: dict[str, str] | None,
|
|
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
client_ip: str | None,
|
|
allowed_server_ids: set[str] | None = None,
|
|
raw_headers: Mapping[str, str] | None = None,
|
|
) -> None:
|
|
"""Fail fast with HTTP 401 for MCP servers that need user auth but
|
|
didn't receive it on this request. Covers both gateway-managed OAuth2
|
|
(points clients at the gateway AS metadata) and pass-through OAuth
|
|
(points clients at the upstream resource-metadata via our well-known).
|
|
|
|
``allowed_server_ids`` may be passed by callers that have already
|
|
narrowed the authorized server set (e.g. toolset scoping); servers
|
|
not in that set are skipped so a client targeting a toolset that
|
|
excludes a passthrough server is not pushed into an OAuth flow for
|
|
a server it will be 403'd on immediately after authentication.
|
|
"""
|
|
for server_name in mcp_servers or []:
|
|
server = operations.global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip)
|
|
if server is not None and allowed_server_ids is not None and server.server_id not in allowed_server_ids:
|
|
# Caller's narrowed scope excludes this server — skip the
|
|
# preemptive challenge and let downstream authorization
|
|
# return 403.
|
|
continue
|
|
if server is not None and server.auth_type == MCPAuth.oauth2 and server.oauth2_flow == "client_credentials":
|
|
# Stamped M2M: the challenge decision below never reads discovered
|
|
# metadata, so deferred-discovery failures must not 503 this loop.
|
|
# Unstamped rows stay on the discover-first path because filling
|
|
# authorization_url/token_url can change their inferred flow.
|
|
continue
|
|
if server is not None:
|
|
server = await operations.global_mcp_server_manager.ensure_oauth_metadata_discovered(server)
|
|
if server and server.auth_type == MCPAuth.oauth2:
|
|
# The challenge decision is per oauth2 sub-mode, not per header:
|
|
# gateway-managed modes (M2M and interactive authorization_code)
|
|
# never receive a client-supplied upstream token, so a bearer in
|
|
# Authorization is a LiteLLM key (surfaced here as oauth2_headers)
|
|
# and must not suppress the challenge. Only the delegate mode
|
|
# treats a present bearer as the upstream token. The sub-mode is
|
|
# resolved the same way egress resolves it, via
|
|
# effective_oauth2_flow: an unstamped (null oauth2_flow) row with
|
|
# the M2M shape resolves to client_credentials, so the bare
|
|
# has_client_credentials column is never trusted here.
|
|
if MCPServerManager.effective_oauth2_flow(server) == "client_credentials":
|
|
# M2M: the gateway mints its own token at egress from the
|
|
# stored client credentials, so there is nothing to challenge.
|
|
continue
|
|
|
|
if getattr(server, "delegate_auth_to_upstream", False) is not True:
|
|
# Gateway-managed interactive (authorization_code): the only
|
|
# thing that authorizes egress is a stored per-user token, so
|
|
# challenge whenever one is absent, regardless of any bearer.
|
|
# The v2 resolver owns the existence check, so every
|
|
# authorization_code resolution (egress and this discovery
|
|
# challenge) runs through it. A keyless admitted subject is
|
|
# challenged with the per-server resource_metadata (whose
|
|
# authorization server is the gateway itself, vaulting via the
|
|
# authorize interlude); the per-server relay advertised below
|
|
# cannot vault without a litellm key on its token request.
|
|
if await operations.global_mcp_server_manager.has_user_oauth_token(server, user_api_key_auth):
|
|
continue
|
|
|
|
if is_mcp_admitted_user_subject(user_api_key_auth):
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={
|
|
"www-authenticate": get_passthrough_www_authenticate(
|
|
scope=scope,
|
|
server_name=server_name,
|
|
)
|
|
},
|
|
)
|
|
|
|
request = StarletteRequest(scope)
|
|
base_url = get_request_base_url(request)
|
|
_path = get_route_relative_request_path(scope)
|
|
|
|
# Pick the well-known AS-metadata form that matches the inbound route
|
|
# so strict RFC 9728 §3.2 clients can resolve it correctly.
|
|
as_metadata_root = f"{base_url}/.well-known/oauth-authorization-server{well_known_root_suffix()}"
|
|
if _path.startswith(f"/mcp/{server_name}"):
|
|
_as_url = f"{as_metadata_root}/mcp/{server_name}"
|
|
else:
|
|
_as_url = f"{as_metadata_root}/{server_name}"
|
|
authorization_uri = f'Bearer authorization_uri="{_as_url}"'
|
|
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={"www-authenticate": authorization_uri},
|
|
)
|
|
|
|
if not oauth2_headers:
|
|
# Delegate-auth servers run upstream PKCE: a present bearer is
|
|
# the upstream token, so only challenge when it is absent, with
|
|
# the proxied resource_metadata (RFC 9728), not the gateway
|
|
# authorization_uri above which would authorize against the
|
|
# gateway instead of the upstream IdP.
|
|
www_authenticate = get_passthrough_www_authenticate(
|
|
scope=scope,
|
|
server_name=server_name,
|
|
)
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={"www-authenticate": www_authenticate},
|
|
)
|
|
# Delegate server with a bearer present: it is the upstream token,
|
|
# so admit the session and move to the next target. Every oauth2
|
|
# sub-mode is terminal here (continue or raise) so no oauth2 server
|
|
# reaches the token_exchange / pass-through blocks below.
|
|
continue
|
|
|
|
# token_exchange (OBO): the caller supplied no subject token. Challenge at connect
|
|
# (transport level, where WWW-Authenticate survives) with the RFC 9728 resource_metadata
|
|
# so the client discovers the IdP, SSOs, and retries with a subject token, which LiteLLM
|
|
# then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the
|
|
# header lost, so the discovery flow needs this pre-emptive challenge.
|
|
if (
|
|
server
|
|
and server.auth_type == MCPAuth.oauth2_token_exchange
|
|
and not oauth2_headers
|
|
and operations.global_mcp_server_manager.extract_subject_token(
|
|
oauth2_headers, raw_headers, user_api_key_auth
|
|
)
|
|
is None
|
|
):
|
|
from litellm.proxy.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph
|
|
raise_token_exchange_challenge,
|
|
)
|
|
from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports proxy utils
|
|
get_request_root_path,
|
|
)
|
|
|
|
raise_token_exchange_challenge(server, root_path=get_request_root_path())
|
|
|
|
# Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run
|
|
# the exchange here at the transport edge, so a rejected subject raises the RFC 9728
|
|
# challenge and any other failure its public status, instead of the session opening and
|
|
# list_tools masking it as an empty tool list. The manager owns which modes pre-flight
|
|
# and what each mints from. Gated to single-server routes the key may reach; the
|
|
# multi-server aggregate keeps absorbing per-server auth failures so one bad server
|
|
# cannot 401 the whole connect.
|
|
if (
|
|
server
|
|
and len(mcp_servers or []) == 1
|
|
and server.server_id
|
|
in frozenset(
|
|
allowed.server_id
|
|
for allowed in await operations._get_allowed_mcp_servers(
|
|
user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
|
|
)
|
|
)
|
|
):
|
|
await operations.global_mcp_server_manager.preflight_token_exchange(
|
|
server=server,
|
|
oauth2_headers=oauth2_headers,
|
|
user_api_key_auth=user_api_key_auth,
|
|
raw_headers=raw_headers,
|
|
)
|
|
|
|
# Pass-through OAuth: when the admin has opted a server into
|
|
# forwarding the client's bearer token (is_oauth_passthrough) and
|
|
# the client hasn't supplied one, fail fast with 401 and point
|
|
# them at the gateway's oauth-protected-resource well-known URL.
|
|
# That endpoint proxies the upstream's metadata so the client
|
|
# kicks off OAuth against the real upstream IdP, not the gateway.
|
|
if (
|
|
server
|
|
and server.is_oauth_passthrough
|
|
and not operations._client_has_passthrough_authorization(
|
|
server, oauth2_headers, mcp_server_auth_headers
|
|
)
|
|
):
|
|
www_authenticate = get_passthrough_www_authenticate(
|
|
scope=scope,
|
|
server_name=server_name,
|
|
)
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={"www-authenticate": www_authenticate},
|
|
)
|
|
|
|
if (
|
|
server
|
|
and server.is_oauth_delegate
|
|
and len(mcp_servers or []) == 1
|
|
and _get_forwarded_auth_from_scope(scope) is None
|
|
and not operations._client_has_per_server_auth_header(server, mcp_server_auth_headers)
|
|
):
|
|
www_authenticate = get_passthrough_www_authenticate(
|
|
scope=scope,
|
|
server_name=server_name,
|
|
)
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={"www-authenticate": www_authenticate},
|
|
)
|
|
|
|
if (
|
|
server
|
|
and server.is_true_passthrough
|
|
and len(mcp_servers or []) == 1
|
|
and not _scope_has_authorization_header(scope)
|
|
and not operations._client_has_per_server_auth_header(server, mcp_server_auth_headers)
|
|
):
|
|
if server.is_dcr_bridge:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={
|
|
"www-authenticate": get_passthrough_www_authenticate(
|
|
scope=scope,
|
|
server_name=server_name,
|
|
)
|
|
},
|
|
)
|
|
upstream_status, upstream_www_authenticate = await _probe_upstream_auth(server.url or "", "")
|
|
if upstream_status == 401 and upstream_www_authenticate:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={"www-authenticate": upstream_www_authenticate},
|
|
)
|
|
|
|
def _get_authorization_header_from_scope(scope: Scope) -> str | None:
|
|
"""First ``Authorization`` header value in the ASGI scope, or None."""
|
|
scope_headers: Final = _scope_headers(scope)
|
|
for key, value in scope_headers:
|
|
if key.lower() == b"authorization":
|
|
return value.decode("latin-1")
|
|
return None
|
|
|
|
def _scope_has_authorization_header(scope: Scope) -> bool:
|
|
return _get_authorization_header_from_scope(scope) is not None
|
|
|
|
def _get_forwarded_auth_from_scope(scope: Scope) -> str | None:
|
|
"""Return the upstream-bound ``Authorization`` header value, or None.
|
|
|
|
Only returns the ``Authorization`` header when ``x-litellm-api-key`` is
|
|
also present. In that case ``Authorization`` is unambiguously the
|
|
upstream token the caller wants forwarded to the MCP server. When
|
|
``x-litellm-api-key`` is absent the ``Authorization`` header may itself
|
|
be the LiteLLM proxy API key (backward-compat path in
|
|
``MCPRequestHandler.process_mcp_request``), and forwarding it upstream
|
|
would leak the proxy key to a third-party MCP server.
|
|
"""
|
|
scope_headers: Final = _scope_headers(scope)
|
|
has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers)
|
|
if not has_litellm_key_header:
|
|
return None
|
|
return _get_authorization_header_from_scope(scope)
|
|
|
|
async def _probe_upstream_auth(
|
|
url: str,
|
|
auth_header: str,
|
|
timeout: float = 5.0,
|
|
) -> tuple[int, str | None]:
|
|
"""JSON-RPC initialize-probe the upstream URL to check whether the token is accepted.
|
|
|
|
Uses POST so StreamableHTTP MCP servers run the same auth path as a
|
|
real client request. Returns (status_code, www_authenticate).
|
|
Fails-open with (200, None) on network errors so a transient hiccup
|
|
does not block valid requests.
|
|
|
|
Uses the public ``AsyncHTTPHandler.post()`` interface and catches
|
|
``httpx.HTTPStatusError`` separately so the 401/403 we want to surface
|
|
is not swallowed by the broad fail-open ``except Exception`` below.
|
|
"""
|
|
client: Final = get_async_httpx_client(
|
|
llm_provider=httpxSpecialProvider.MCP,
|
|
params={"timeout": timeout},
|
|
)
|
|
probe_payload: Final = {
|
|
"jsonrpc": "2.0",
|
|
"id": "litellm-mcp-auth-probe",
|
|
"method": "initialize",
|
|
"params": {
|
|
"protocolVersion": MCPSpecVersion.jun_2025.value,
|
|
"capabilities": {},
|
|
"clientInfo": {
|
|
"name": "litellm-mcp-auth-probe",
|
|
"version": "1.0.0",
|
|
},
|
|
},
|
|
}
|
|
probe_headers: Final = {
|
|
"Accept": "application/json, text/event-stream",
|
|
**({"Authorization": auth_header} if auth_header else {}),
|
|
}
|
|
try:
|
|
resp: Final = await client.post(
|
|
url=url,
|
|
headers=probe_headers,
|
|
json=probe_payload,
|
|
timeout=timeout,
|
|
)
|
|
return resp.status_code, resp.headers.get("www-authenticate")
|
|
except httpx.HTTPStatusError as exc:
|
|
# AsyncHTTPHandler.post() calls raise_for_status(); a 401/403 from
|
|
# upstream lands here. Return its status so the caller can map it
|
|
# to the appropriate response.
|
|
return exc.response.status_code, exc.response.headers.get("www-authenticate")
|
|
except Exception as exc:
|
|
verbose_logger.debug("_probe_upstream_auth: probe to %s failed (%s), allowing request through", url, exc)
|
|
return 200, None
|
|
|
|
async def _check_passthrough_upstream_auth(
|
|
scope: Scope,
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
mcp_servers: list[str] | None,
|
|
client_ip: str | None,
|
|
*,
|
|
oauth2_headers: Mapping[str, str] | None,
|
|
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
|
|
raw_headers: Mapping[str, str] | None,
|
|
requested_servers: Sequence[str] | None = None,
|
|
) -> None:
|
|
"""Probe pass-through and legacy delegate upstream servers in parallel before the MCP session starts.
|
|
|
|
Only servers the caller's key is already authorized to reach are probed —
|
|
the list is derived from _get_allowed_mcp_servers so that a user cannot
|
|
trigger an upstream probe against a server their key is not permitted for.
|
|
|
|
The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401
|
|
can only be returned before that point. This function raises HTTPException(401)
|
|
with a WWW-Authenticate header if any upstream rejects the client token, or 403
|
|
if the upstream accepts it but forbids the caller.
|
|
Fails-open: network errors are logged and the request is allowed through.
|
|
|
|
"""
|
|
if not oauth2_headers and not mcp_server_auth_headers:
|
|
return
|
|
|
|
# Use the authorized server set, not the raw user-supplied names, so that
|
|
# a caller cannot force a probe to a server their key is not allowed to use.
|
|
allowed_servers: Final = await operations._get_allowed_mcp_servers(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_servers=mcp_servers,
|
|
client_ip=client_ip,
|
|
)
|
|
# A legacy delegate is probed only when the single requested name resolves to it the way admission
|
|
# resolves it, so neither a multi-server connect nor an access group can be failed by one rejected
|
|
# token; its challenge echoes the requested name, matching the tokenless challenge.
|
|
requested_scope: Final = requested_servers if requested_servers is not None else mcp_servers
|
|
single_requested_name: Final = (
|
|
requested_scope[0] if requested_scope is not None and len(requested_scope) == 1 else None
|
|
)
|
|
requested_server: Final = (
|
|
operations.global_mcp_server_manager.get_mcp_server_by_name(single_requested_name, client_ip=client_ip)
|
|
if single_requested_name is not None
|
|
else None
|
|
)
|
|
prepared_headers: Final = tuple(
|
|
(
|
|
srv,
|
|
operations._prepare_mcp_server_headers(
|
|
server=srv,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
mcp_auth_header=None,
|
|
oauth2_headers=dict(oauth2_headers) if oauth2_headers is not None else None,
|
|
raw_headers=dict(raw_headers) if raw_headers is not None else None,
|
|
user_api_key_auth=user_api_key_auth,
|
|
scope_servers=allowed_servers,
|
|
),
|
|
)
|
|
for srv in allowed_servers
|
|
if srv.is_oauth_passthrough
|
|
or (
|
|
requested_server is not None
|
|
and srv.server_id == requested_server.server_id
|
|
and is_legacy_upstream_delegate(srv)
|
|
)
|
|
)
|
|
probe_headers: Final = tuple(
|
|
(
|
|
srv,
|
|
merge_mcp_headers(
|
|
extra_headers=auth if isinstance(auth, dict) else None,
|
|
static_headers=extra,
|
|
),
|
|
)
|
|
for srv, (auth, extra) in prepared_headers
|
|
)
|
|
probe_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple(
|
|
(srv, auth_header, srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name)
|
|
for srv, headers in probe_headers
|
|
if (
|
|
auth_header := next(
|
|
(value for name, value in (headers or {}).items() if name.lower() == "authorization"), None
|
|
)
|
|
)
|
|
)
|
|
if not probe_targets:
|
|
return
|
|
|
|
probe_results: Final = await asyncio.gather(
|
|
*[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets]
|
|
)
|
|
for (_srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results):
|
|
if probe_status == 401:
|
|
# Token is missing or expired: keep pass-through clients on the
|
|
# protected-resource discovery flow so they re-authorize against
|
|
# the upstream IdP metadata proxied by LiteLLM.
|
|
www_authenticate = get_passthrough_www_authenticate(
|
|
scope=scope,
|
|
server_name=challenge_server_name,
|
|
invalid_token=True,
|
|
)
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Unauthorized",
|
|
headers={"www-authenticate": www_authenticate},
|
|
)
|
|
if probe_status == 403:
|
|
# Token is valid but the caller lacks permission — do not hint
|
|
# at re-authorization (RFC 9110: a fresh token with the same
|
|
# scopes would just hit 403 again and loop indefinitely).
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Forbidden",
|
|
)
|
|
|
|
_INTERACTION_REQUEST: Final[TypeAdapter[InteractionOperation]] = TypeAdapter(InteractionOperation)
|
|
|
|
@catalog_operation(lambda: operations.global_mcp_server_manager)
|
|
async def _preflight_modern_interaction(
|
|
scope: Scope, body: bytes, context: OperationContext
|
|
) -> JSONResponse | None:
|
|
try:
|
|
envelope: Final = JSONRPCRequest.model_validate_json(body)
|
|
operation: Final = _INTERACTION_REQUEST.validate_json(body)
|
|
except ValidationError:
|
|
return None
|
|
headers: Final = StarletteRequest(scope).headers
|
|
if find_duplicated_routing_header(headers.items()) is not None or isinstance(
|
|
classify_inbound_request(envelope.model_dump(by_alias=True), headers=dict(headers)), InboundLadderRejection
|
|
):
|
|
return None
|
|
try:
|
|
state: Final = operations.validate_continuation(operation, context)
|
|
except MCPError as error:
|
|
return JSONResponse(
|
|
status_code=400,
|
|
content={"jsonrpc": "2.0", "id": envelope.id, "error": error.error.model_dump(exclude_none=True)},
|
|
)
|
|
if state is None and context.mcp_servers is None:
|
|
return None
|
|
targets: Final = (
|
|
[state.target_id]
|
|
if state is not None
|
|
else list(context.mcp_servers)
|
|
if context.mcp_servers is not None
|
|
else None
|
|
)
|
|
allowed: Final = await operations._get_allowed_mcp_servers(
|
|
user_api_key_auth=context.user_api_key_auth, mcp_servers=targets, client_ip=context.client_ip
|
|
)
|
|
if state is not None and not any(target.server_id == state.target_id for target in allowed):
|
|
raise HTTPException(status_code=403, detail="MCP continuation target is no longer authorized")
|
|
authorized_names: Final = [target.alias or target.name for target in allowed]
|
|
await _raise_preemptive_401_for_unauthenticated_servers(
|
|
scope=scope,
|
|
mcp_servers=list(context.mcp_servers) if context.mcp_servers is not None else authorized_names,
|
|
oauth2_headers=dict(context.oauth2_headers) if context.oauth2_headers is not None else None,
|
|
mcp_server_auth_headers={key: dict(value) for key, value in context.mcp_server_auth_headers.items()}
|
|
if context.mcp_server_auth_headers is not None
|
|
else None,
|
|
user_api_key_auth=context.user_api_key_auth,
|
|
client_ip=context.client_ip,
|
|
allowed_server_ids={target.server_id for target in allowed},
|
|
raw_headers=context.raw_headers,
|
|
)
|
|
await _check_passthrough_upstream_auth(
|
|
scope,
|
|
context.user_api_key_auth,
|
|
authorized_names,
|
|
context.client_ip,
|
|
oauth2_headers=context.oauth2_headers,
|
|
mcp_server_auth_headers={key: dict(value) for key, value in context.mcp_server_auth_headers.items()}
|
|
if context.mcp_server_auth_headers is not None
|
|
else None,
|
|
raw_headers=context.raw_headers,
|
|
requested_servers=context.mcp_servers,
|
|
)
|
|
return None
|
|
|
|
async def _preflight_streamable_http_upstream_auth(
|
|
defer_upstream_probes: bool,
|
|
targets_disabled_capability: bool,
|
|
scope: Scope,
|
|
mcp_servers: Sequence[str] | None,
|
|
oauth2_headers: Mapping[str, str] | None,
|
|
mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None,
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
client_ip: str | None,
|
|
allowed_server_ids: Set[str] | None,
|
|
raw_headers: Mapping[str, str] | None,
|
|
) -> None:
|
|
if not defer_upstream_probes and not targets_disabled_capability:
|
|
await _raise_preemptive_401_for_unauthenticated_servers(
|
|
scope=scope,
|
|
mcp_servers=list(mcp_servers) if mcp_servers is not None else None,
|
|
oauth2_headers=dict(oauth2_headers) if oauth2_headers is not None else None,
|
|
mcp_server_auth_headers=(
|
|
{name: dict(headers) for name, headers in mcp_server_auth_headers.items()}
|
|
if mcp_server_auth_headers is not None
|
|
else None
|
|
),
|
|
user_api_key_auth=user_api_key_auth,
|
|
client_ip=client_ip,
|
|
allowed_server_ids=set(allowed_server_ids) if allowed_server_ids is not None else None,
|
|
raw_headers=raw_headers,
|
|
)
|
|
# Pre-flight auth check for pass-through servers. Must run after
|
|
# toolset scoping so the probe list is derived from the fully-authorized
|
|
# server set, not the raw user-supplied names.
|
|
await _check_passthrough_upstream_auth(
|
|
scope,
|
|
user_api_key_auth,
|
|
list(mcp_servers) if mcp_servers is not None else None,
|
|
client_ip,
|
|
oauth2_headers=oauth2_headers,
|
|
mcp_server_auth_headers=(
|
|
{name: dict(headers) for name, headers in mcp_server_auth_headers.items()}
|
|
if mcp_server_auth_headers is not None
|
|
else None
|
|
),
|
|
raw_headers=raw_headers,
|
|
)
|
|
|
|
async def _targets_disabled_capability(peek: _LazyPeekedBody | None) -> bool:
|
|
if peek is None:
|
|
return False
|
|
disabled_ops: Final = disabled_operations(disabled_capabilities())
|
|
if not disabled_ops:
|
|
return False
|
|
return _jsonrpc_request_method(await peek.full_body()) in disabled_ops
|
|
|
|
async def _preflight_streamable_http_modern_interaction(
|
|
peek: _LazyPeekedBody | None,
|
|
defer_upstream_probes: bool,
|
|
targets_disabled_capability: bool,
|
|
scope: Scope,
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
mcp_auth_header: str | None,
|
|
mcp_servers: Sequence[str] | None,
|
|
mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None,
|
|
oauth2_headers: Mapping[str, str] | None,
|
|
raw_headers: Mapping[str, str] | None,
|
|
client_ip: str | None,
|
|
) -> JSONResponse | None:
|
|
if peek is None or not defer_upstream_probes or targets_disabled_capability:
|
|
return None
|
|
return await _preflight_modern_interaction(
|
|
scope,
|
|
await peek.full_body(),
|
|
OperationContext(
|
|
_caller=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=client_ip,
|
|
),
|
|
)
|
|
|
|
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:
|
|
reject_disallowed_mcp_origin(StarletteRequest(scope))
|
|
bad_version: Final = unsupported_protocol_version(scope)
|
|
if bad_version is not None:
|
|
supported: Final = ", ".join(configured_versions())
|
|
await JSONResponse(
|
|
status_code=400,
|
|
content={
|
|
"jsonrpc": "2.0",
|
|
"id": None,
|
|
"error": {
|
|
"code": INVALID_REQUEST,
|
|
"message": f"Unsupported MCP-Protocol-Version {bad_version}; supported: {supported}",
|
|
},
|
|
},
|
|
)(scope, receive, send)
|
|
return
|
|
if scope.get("method") == "OPTIONS":
|
|
await Response(status_code=204, headers={"Allow": "GET, POST, DELETE, OPTIONS"})(scope, receive, send)
|
|
return
|
|
path: Final = _scope_path(scope)
|
|
peek: Final = _LazyPeekedBody(receive) if scope.get("method") == "POST" else None
|
|
if peek is not None:
|
|
receive = peek.receive # rebind-ok: replay peeked ASGI messages to the downstream handler
|
|
|
|
async def admission_auth_body() -> bytes:
|
|
return _admission_auth_body(await peek.admission_body())
|
|
|
|
scope[MCP_PEEKED_BODY_SCOPE_KEY] = admission_auth_body
|
|
(
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
mcp_servers,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
) = await extract_mcp_auth_context(scope, path)
|
|
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
|
|
request_headers: Final = StarletteRequest(scope).headers
|
|
defer_upstream_probes: Final = request_headers.get(
|
|
"mcp-protocol-version"
|
|
) == "2026-07-28" and request_headers.get("mcp-method") in {"tools/call", "prompts/get", "resources/read"}
|
|
targets_disabled_capability: Final = await _targets_disabled_capability(peek)
|
|
|
|
# Extract client IP for MCP access control
|
|
_client_ip: Final = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope))
|
|
|
|
verbose_logger.debug("MCP request mcp_servers (header/path): %s", mcp_servers)
|
|
verbose_logger.debug(
|
|
"MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None
|
|
)
|
|
|
|
# Strip any client-supplied x-mcp-toolset-id to prevent forgery.
|
|
scope_headers: Final = _scope_headers(scope)
|
|
scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"]
|
|
|
|
# Apply toolset scope if set server-side via ContextVar (set by
|
|
# /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py).
|
|
active_toolset_id: Final = mcp_active_toolset_id.get()
|
|
toolset_allowed_server_ids: set[str] | None = None
|
|
if active_toolset_id and user_api_key_auth is not None:
|
|
user_api_key_auth = ( # rebind-ok: pre-existing rebinding on a rename-only line
|
|
await apply_toolset_scope(user_api_key_auth, active_toolset_id)
|
|
)
|
|
toolset_allowed_server_ids = await _toolset_server_ids(active_toolset_id)
|
|
|
|
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
|
|
# Must run after toolset scoping so the challenge set is derived
|
|
# from the fully-authorized server set: a passthrough server that
|
|
# the active toolset excludes should not trigger an OAuth flow
|
|
# for a server the caller will be 403'd on after authentication.
|
|
await _preflight_streamable_http_upstream_auth(
|
|
defer_upstream_probes=defer_upstream_probes,
|
|
targets_disabled_capability=targets_disabled_capability,
|
|
scope=scope,
|
|
mcp_servers=mcp_servers,
|
|
oauth2_headers=oauth2_headers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
user_api_key_auth=user_api_key_auth,
|
|
client_ip=_client_ip,
|
|
allowed_server_ids=toolset_allowed_server_ids,
|
|
raw_headers=raw_headers,
|
|
)
|
|
|
|
# Inject masked debug headers when client sends x-litellm-mcp-debug: true
|
|
_debug_headers: Final = MCPDebug.maybe_build_debug_headers(
|
|
raw_headers=raw_headers,
|
|
scope=dict(scope),
|
|
mcp_servers=mcp_servers,
|
|
oauth2_headers=oauth2_headers,
|
|
client_ip=_client_ip,
|
|
)
|
|
diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None
|
|
if diagnostics is not None:
|
|
scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics
|
|
send = MCPDebug.wrap_send_with_debug_headers(
|
|
send, _debug_headers, diagnostics.headers, request_method=scope.get("method")
|
|
)
|
|
|
|
# Ensure session managers are initialized
|
|
if not _SESSION_MANAGERS_INITIALIZED:
|
|
await initialize_session_managers()
|
|
# Give it a moment to start up
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Route based on mcp-session-id and request method:
|
|
# - Has session ID → stateful (Claude Code, Cursor, VSCode)
|
|
# - No session ID + initialize → stateful (so client gets mcp-session-id)
|
|
# - No session ID + other → stateless (curl, Inspector, Notion)
|
|
session_id = _get_session_id_from_scope(scope)
|
|
|
|
# Owner-binding: a live stateful session may only be driven by the
|
|
# caller that created it. Reject mismatches with 403 so a leaked
|
|
# mcp-session-id cannot be hijacked by another authenticated user.
|
|
#
|
|
# Run before ``_handle_stale_mcp_session`` so a non-owner cannot
|
|
# force-clean another caller's residual tracking entries via a
|
|
# stale DELETE.
|
|
if session_id:
|
|
expected_owner: Final = _stateful_session_owners.get(session_id)
|
|
request_owner = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip)
|
|
if expected_owner is not None and expected_owner != request_owner:
|
|
verbose_logger.warning(
|
|
"Rejecting MCP request: session '%s' owner mismatch.",
|
|
session_id,
|
|
)
|
|
forbidden_response: Final = JSONResponse(
|
|
status_code=403,
|
|
content={
|
|
"error": "Forbidden",
|
|
"details": "mcp-session-id is bound to a different caller.",
|
|
},
|
|
)
|
|
await forbidden_response(scope, receive, send)
|
|
return
|
|
|
|
# Handle stale session IDs before choosing a target manager. Stale
|
|
# non-DELETE requests have their session header stripped and should
|
|
# be routed as no-session requests.
|
|
if session_id:
|
|
handled: Final = await _handle_stale_mcp_session(scope, receive, send, session_manager_stateful)
|
|
if handled:
|
|
# Request was fully handled (e.g., DELETE on non-existent session)
|
|
return
|
|
session_id = _get_session_id_from_scope(scope)
|
|
|
|
rejection: Final = await _preflight_streamable_http_modern_interaction(
|
|
peek=peek,
|
|
defer_upstream_probes=defer_upstream_probes,
|
|
targets_disabled_capability=targets_disabled_capability,
|
|
scope=scope,
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=_client_ip,
|
|
)
|
|
if rejection is not None:
|
|
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 or reauthorize_as is not None
|
|
else None
|
|
)
|
|
if scope_response is not None:
|
|
scope[SCOPE_RESPONSE_KEY] = scope_response
|
|
|
|
use_stateful: Final = bool(session_id or is_initialize)
|
|
target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless
|
|
|
|
verbose_logger.debug(
|
|
f"MCP routing to {'stateful' if use_stateful else 'stateless'} manager"
|
|
+ (f" (session={session_id[:8]}...)" if session_id else "")
|
|
+ (" (initialize)" if is_initialize else "")
|
|
)
|
|
|
|
# A new `initialize` (no session id) is about to create a stateful
|
|
# session. Cap how many a single caller can hold so an authenticated
|
|
# client cannot spam `initialize` and exhaust memory.
|
|
if is_initialize and not session_id:
|
|
request_owner = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip)
|
|
if not await _enforce_stateful_session_cap_for_owner(request_owner):
|
|
verbose_logger.warning(
|
|
"Rejecting MCP initialize: caller already holds the maximum number of active stateful sessions."
|
|
)
|
|
too_many_response: Final = JSONResponse(
|
|
status_code=429,
|
|
content={
|
|
"error": "Too Many Requests",
|
|
"details": "Too many active MCP sessions for this caller.",
|
|
},
|
|
)
|
|
await too_many_response(scope, receive, send)
|
|
return
|
|
|
|
# Serialize requests on the same stateful session so concurrent
|
|
# callers don't clobber each other's auth context mid-flight.
|
|
#
|
|
# Skip the lock for streaming GETs (SSE channels held open for the
|
|
# life of the session): holding a per-session lock for a long-lived
|
|
# stream would block every subsequent POST on the same session.
|
|
# POST/DELETE are the methods that actually mutate the shared
|
|
# auth context, so serializing those is sufficient for the
|
|
# clobbering race between concurrent JSON-RPC calls.
|
|
#
|
|
# Also skip the lock for JSON-RPC *responses* (POSTs that carry
|
|
# a ``result`` or ``error`` but no ``method``). These are replies
|
|
# to server-initiated requests such as ``elicitation/create`` or
|
|
# ``sampling/createMessage``. The in-flight tool-call POST that
|
|
# triggered the server request already holds the session lock, so
|
|
# trying to acquire it again for the response POST would deadlock.
|
|
is_jsonrpc_response = False
|
|
request_method: Final = (scope.get("method") or "").upper()
|
|
if body and request_method == "POST":
|
|
try:
|
|
_peeked: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(body)))
|
|
if (
|
|
_peeked is not None
|
|
and _peeked.get("jsonrpc") == "2.0"
|
|
and "id" in _peeked
|
|
and "method" not in _peeked
|
|
and ("result" in _peeked or "error" in _peeked)
|
|
):
|
|
is_jsonrpc_response = True
|
|
verbose_logger.debug(
|
|
"MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
|
|
_peeked.get("id"),
|
|
)
|
|
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
|
|
# Peek cap truncated the body, so it can't be fully parsed.
|
|
# Scan the top-level keys (depth-aware) instead of a flat
|
|
# substring search: a response's result payload may nest a
|
|
# "method" field, and misreading that would acquire the lock
|
|
# and deadlock the in-flight tool call awaiting this
|
|
# response. A false skip is harmless; a false acquire is not.
|
|
_body_str: Final = body.decode("utf-8", errors="replace")
|
|
if (
|
|
'"jsonrpc"' in _body_str
|
|
and ('"result"' in _body_str or '"error"' in _body_str)
|
|
and not _jsonrpc_text_has_top_level_method(_body_str)
|
|
):
|
|
is_jsonrpc_response = True
|
|
verbose_logger.debug(
|
|
"MCP: detected truncated JSON-RPC response POST via "
|
|
"top-level key scan, skipping session lock to avoid deadlock"
|
|
)
|
|
|
|
is_cancellation: Final = request_method == "POST" and _is_cancellation_notification(body)
|
|
session_lock: Final = _stateful_request_lock(
|
|
use_stateful, session_id, request_method, is_jsonrpc_response, is_cancellation
|
|
)
|
|
|
|
active_request_session_ids: Final[list[str]] = []
|
|
|
|
def _increment_active_request_session(session_id_to_track: str) -> None:
|
|
if session_id_to_track in active_request_session_ids:
|
|
return
|
|
active_request_session_ids.append(session_id_to_track)
|
|
_stateful_session_active_request_counts[session_id_to_track] = (
|
|
_stateful_session_active_request_counts.get(session_id_to_track, 0) + 1
|
|
)
|
|
|
|
if use_stateful and session_id:
|
|
_increment_active_request_session(session_id)
|
|
|
|
def _track_initialized_stateful_session(
|
|
initialized_session_id: str,
|
|
) -> None:
|
|
_increment_active_request_session(initialized_session_id)
|
|
|
|
async def _dispatch() -> None:
|
|
_otel_publish_transport_span_on_scope(scope)
|
|
_otel_publish_request_destinations_on_scope(scope)
|
|
auth_user: Final = _set_or_update_auth_context(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=_client_ip,
|
|
session_id=session_id if use_stateful else None,
|
|
touch_last_seen=(scope.get("method") or "").upper() != "DELETE",
|
|
copy_existing_session_auth_context=is_initialize or is_cancellation,
|
|
)
|
|
local_send = scope_response or send
|
|
if use_stateful and is_initialize:
|
|
local_send = _wrap_send_with_stateful_session_auth_context(
|
|
local_send,
|
|
auth_user,
|
|
_owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip),
|
|
_track_initialized_stateful_session,
|
|
client_info=_extract_initialize_client_info(body),
|
|
)
|
|
|
|
async with (
|
|
contextlib.nullcontext()
|
|
if defer_upstream_probes
|
|
else _gateway_initialize_instructions_request_scope(
|
|
user_api_key_auth,
|
|
mcp_servers,
|
|
_client_ip,
|
|
scoped_server_endpoint=scoped_server_endpoint,
|
|
is_initialize=is_initialize,
|
|
)
|
|
):
|
|
await target_manager.handle_request(scope, receive, local_send)
|
|
if use_stateful and session_id and scope.get("method") == "DELETE":
|
|
_remove_stateful_session_tracking(session_id)
|
|
|
|
try:
|
|
if session_lock is not None:
|
|
async with session_lock:
|
|
await _dispatch()
|
|
else:
|
|
await _dispatch()
|
|
finally:
|
|
for active_request_session_id in active_request_session_ids:
|
|
active_request_count = _stateful_session_active_request_counts.get(active_request_session_id, 0) - 1
|
|
if active_request_count > 0:
|
|
_stateful_session_active_request_counts[active_request_session_id] = active_request_count
|
|
else:
|
|
_stateful_session_active_request_counts.pop(active_request_session_id, None)
|
|
|
|
if scope.get("method") != "DELETE" and active_request_session_id in _stateful_session_auth_contexts:
|
|
_stateful_session_auth_context_last_seen[active_request_session_id] = time.monotonic()
|
|
|
|
# Periodic cleanup iterates _stateful_session_auth_context_last_seen,
|
|
# so locks for untracked sessions must be dropped here.
|
|
if active_request_count <= 0 and active_request_session_id not in _stateful_session_auth_contexts:
|
|
_stateful_session_locks.pop(active_request_session_id, None)
|
|
except MCPUpstreamAuthError as e:
|
|
# Upstream delegated auth returned 401; surface it to the client so
|
|
# standards-compliant MCP clients trigger the upstream OAuth flow.
|
|
raise e.to_http_exception(
|
|
base_url=get_request_base_url(StarletteRequest(scope)),
|
|
request_path=scope.get("_original_path") or scope.get("path"),
|
|
)
|
|
except HTTPException:
|
|
# Re-raise HTTP exceptions to preserve status codes and details
|
|
raise
|
|
except ProxyException as e:
|
|
# Auth failures from user_api_key_auth arrive as ProxyException, not
|
|
# HTTPException. Preserve the real status (e.g. 401 + WWW-Authenticate)
|
|
# so OAuth clients can re-authenticate instead of receiving a generic
|
|
# 500 that surfaces as a cancelled tool call.
|
|
raise _proxy_exception_to_http_exception(e)
|
|
except Exception as e:
|
|
verbose_logger.exception("Error handling MCP request: %s", e)
|
|
# Try to send a graceful error response for non-HTTP exceptions
|
|
try:
|
|
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
|
|
|
|
error_response: Final = JSONResponse(
|
|
status_code=HTTP_500_INTERNAL_SERVER_ERROR,
|
|
content={"error": "MCP request failed", "details": str(e)},
|
|
)
|
|
await error_response(scope, receive, send)
|
|
except Exception as response_error:
|
|
verbose_logger.exception("Failed to send error response: %s", response_error)
|
|
# If we can't send a proper response, re-raise the original error
|
|
raise e
|
|
|
|
async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None:
|
|
"""Handle MCP requests through SSE."""
|
|
try:
|
|
reject_disallowed_mcp_origin(StarletteRequest(scope))
|
|
bad_version: Final = unsupported_protocol_version(scope)
|
|
if bad_version is not None:
|
|
supported: Final = ", ".join(configured_versions())
|
|
await JSONResponse(
|
|
status_code=400,
|
|
content={
|
|
"jsonrpc": "2.0",
|
|
"id": None,
|
|
"error": {
|
|
"code": INVALID_REQUEST,
|
|
"message": f"Unsupported MCP-Protocol-Version {bad_version}; supported: {supported}",
|
|
},
|
|
},
|
|
)(scope, receive, send)
|
|
return
|
|
from litellm.proxy.auth.auth_utils import get_request_route
|
|
|
|
path: Final = get_request_route(StarletteRequest(scope))
|
|
(
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
mcp_servers,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
) = await extract_mcp_auth_context(scope, path)
|
|
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
|
|
_sse_client_ip: Final = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope))
|
|
|
|
verbose_logger.debug("MCP request mcp_servers (header/path): %s", mcp_servers)
|
|
verbose_logger.debug(
|
|
"MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None
|
|
)
|
|
|
|
# Strip any client-supplied x-mcp-toolset-id to prevent forgery.
|
|
scope_headers: Final = _scope_headers(scope)
|
|
scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"]
|
|
|
|
# Apply toolset scope if set server-side via ContextVar so the
|
|
# downstream probe list matches the fully-authorized server set
|
|
# (mirrors the streamable HTTP handler).
|
|
active_toolset_id: Final = mcp_active_toolset_id.get()
|
|
toolset_allowed_server_ids: set[str] | None = None
|
|
if active_toolset_id and user_api_key_auth is not None:
|
|
user_api_key_auth = ( # rebind-ok: pre-existing rebinding on a rename-only line
|
|
await apply_toolset_scope(user_api_key_auth, active_toolset_id)
|
|
)
|
|
toolset_allowed_server_ids = await _toolset_server_ids(active_toolset_id)
|
|
|
|
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
|
|
# Must run after toolset scoping so the challenge set is derived
|
|
# from the fully-authorized server set: a passthrough server that
|
|
# the active toolset excludes should not trigger an OAuth flow
|
|
# for a server the caller will be 403'd on after authentication.
|
|
await _raise_preemptive_401_for_unauthenticated_servers(
|
|
scope=scope,
|
|
mcp_servers=mcp_servers,
|
|
oauth2_headers=oauth2_headers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
user_api_key_auth=user_api_key_auth,
|
|
client_ip=_sse_client_ip,
|
|
allowed_server_ids=toolset_allowed_server_ids,
|
|
raw_headers=raw_headers,
|
|
)
|
|
|
|
# Pre-flight auth check for pass-through servers: surface upstream
|
|
# 401/403 as a proper challenge before the SSE session commits 200
|
|
# headers, so clients can refresh their OAuth token instead of
|
|
# being stuck with a silently empty tool list. Must run after
|
|
# toolset scoping so the probe list is derived from the fully-
|
|
# authorized server set, not the raw user-supplied names.
|
|
await _check_passthrough_upstream_auth(
|
|
scope,
|
|
user_api_key_auth,
|
|
mcp_servers,
|
|
_sse_client_ip,
|
|
oauth2_headers=oauth2_headers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
raw_headers=raw_headers,
|
|
)
|
|
set_auth_context(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=_sse_client_ip,
|
|
)
|
|
|
|
owner: Final = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _sse_client_ip)
|
|
transport_scope: Final[Scope] = {
|
|
**scope,
|
|
"user": AuthenticatedUser(AccessToken(token=owner, client_id=owner, scopes=[])),
|
|
}
|
|
if scope["method"] == "POST":
|
|
await sse.handle_post_message(transport_scope, receive, send)
|
|
return
|
|
|
|
async with _gateway_initialize_instructions_request_scope(
|
|
user_api_key_auth,
|
|
mcp_servers,
|
|
_sse_client_ip,
|
|
scoped_server_endpoint=scoped_server_endpoint,
|
|
is_initialize=asgi_scope_fields(scope).get("method", "") == "GET",
|
|
):
|
|
async with (
|
|
sse.connect_sse(transport_scope, receive, send) as (read_stream, write_stream),
|
|
server.lifespan(server) as lifespan_state,
|
|
):
|
|
request_token: Final = legacy_subscription_request.set(StarletteRequest(scope))
|
|
try:
|
|
await serve_loop(
|
|
server,
|
|
read_stream,
|
|
write_stream,
|
|
lifespan_state=lifespan_state,
|
|
init_options=server.create_initialization_options(),
|
|
)
|
|
finally:
|
|
legacy_subscription_request.reset(request_token)
|
|
|
|
except MCPUpstreamAuthError as e:
|
|
# Upstream delegated auth returned 401; surface it to the client so
|
|
# standards-compliant MCP clients trigger the upstream OAuth flow.
|
|
raise e.to_http_exception(
|
|
base_url=get_request_base_url(StarletteRequest(scope)),
|
|
request_path=scope.get("_original_path") or scope.get("path"),
|
|
)
|
|
except HTTPException:
|
|
# Re-raise HTTP exceptions to preserve status codes and details
|
|
# (e.g. 401 + WWW-Authenticate challenges from OAuth pass-through).
|
|
raise
|
|
except ProxyException as e:
|
|
# Auth failures from user_api_key_auth arrive as ProxyException, not
|
|
# HTTPException. Preserve the real status (e.g. 401 + WWW-Authenticate)
|
|
# so OAuth clients can re-authenticate instead of receiving a generic
|
|
# 500 that surfaces as a cancelled tool call.
|
|
raise _proxy_exception_to_http_exception(e)
|
|
except Exception as e:
|
|
verbose_logger.exception("Error handling MCP request: %s", e)
|
|
# Try to send a graceful error response for non-HTTP exceptions
|
|
try:
|
|
# Send a proper HTTP error response instead of letting the exception bubble up
|
|
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
|
|
|
|
error_response: Final = JSONResponse(
|
|
status_code=HTTP_500_INTERNAL_SERVER_ERROR,
|
|
content={"error": "MCP request failed", "details": str(e)},
|
|
)
|
|
await error_response(scope, receive, send)
|
|
except Exception as response_error:
|
|
verbose_logger.exception("Failed to send error response: %s", response_error)
|
|
# If we can't send a proper response, re-raise the original error
|
|
raise e
|
|
|
|
app = FastAPI(
|
|
title=LITELLM_MCP_SERVER_NAME,
|
|
description=LITELLM_MCP_SERVER_DESCRIPTION,
|
|
version=LITELLM_MCP_SERVER_VERSION,
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
# Routes
|
|
@app.get(
|
|
"/enabled",
|
|
description="Returns if the MCP server is enabled",
|
|
)
|
|
def get_mcp_server_enabled() -> dict[str, bool]:
|
|
"""
|
|
Returns if the MCP server is enabled
|
|
"""
|
|
return {"enabled": MCP_AVAILABLE}
|
|
|
|
class _LegacySseEndpoint:
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
await handle_sse_mcp(scope, receive, send)
|
|
|
|
for sse_path, sse_method in (
|
|
("/sse", "GET"),
|
|
("/sse/", "GET"),
|
|
("/sse/messages", "POST"),
|
|
("/sse/messages/", "POST"),
|
|
):
|
|
app.router.routes.append(Route(sse_path, endpoint=_LegacySseEndpoint(), methods=[sse_method]))
|
|
|
|
# Mount the MCP handlers
|
|
app.mount("/", handle_streamable_http_mcp)
|
|
app.mount("/mcp", handle_streamable_http_mcp)
|
|
app.mount("/{mcp_server_name}/mcp", handle_streamable_http_mcp)
|
|
app.add_middleware(AuthContextMiddleware)
|
|
|
|
########################################################
|
|
############ Auth Context Functions ####################
|
|
########################################################
|
|
|
|
def _update_auth_context(
|
|
auth_user: MCPAuthenticatedUser,
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
mcp_auth_header: str | None = None,
|
|
mcp_servers: list[str] | None = None,
|
|
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
|
|
oauth2_headers: dict[str, str] | None = None,
|
|
raw_headers: dict[str, str] | None = None,
|
|
client_ip: str | None = None,
|
|
) -> None:
|
|
auth_user.user_api_key_auth = user_api_key_auth
|
|
auth_user.mcp_auth_header = mcp_auth_header
|
|
auth_user.mcp_servers = mcp_servers
|
|
auth_user.mcp_server_auth_headers = mcp_server_auth_headers or {}
|
|
auth_user.oauth2_headers = oauth2_headers
|
|
auth_user.raw_headers = raw_headers
|
|
auth_user.client_ip = client_ip
|
|
|
|
def set_auth_context(
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
mcp_auth_header: str | None = None,
|
|
mcp_servers: list[str] | None = None,
|
|
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
|
|
oauth2_headers: dict[str, str] | None = None,
|
|
raw_headers: dict[str, str] | None = None,
|
|
client_ip: str | None = None,
|
|
) -> MCPAuthenticatedUser:
|
|
"""
|
|
Set the UserAPIKeyAuth in the auth context variable.
|
|
|
|
Args:
|
|
user_api_key_auth: UserAPIKeyAuth object
|
|
mcp_auth_header: MCP auth header to be passed to the MCP server (deprecated)
|
|
mcp_servers: Optional list of server names and access groups to filter by
|
|
mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value}
|
|
client_ip: Client IP address for MCP access control
|
|
"""
|
|
auth_user: Final = MCPAuthenticatedUser(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=client_ip,
|
|
)
|
|
auth_context_var.set(auth_user)
|
|
return auth_user
|
|
|
|
def _stateful_request_lock(
|
|
use_stateful: bool,
|
|
session_id: str | None,
|
|
method: str,
|
|
is_response: bool,
|
|
is_cancellation: bool,
|
|
) -> asyncio.Lock | None:
|
|
if not use_stateful or not session_id or method not in ("POST", "DELETE") or is_response or is_cancellation:
|
|
return None
|
|
return _stateful_session_locks.setdefault(session_id, asyncio.Lock())
|
|
|
|
def _is_cancellation_notification(body: bytes) -> bool:
|
|
try:
|
|
payload: Final = TypeAdapter(dict[str, object]).validate_json(body)
|
|
if payload.get("jsonrpc") != "2.0" or payload.get("method") != "notifications/cancelled" or "id" in payload:
|
|
return False
|
|
from mcp.types import CancelledNotification
|
|
|
|
CancelledNotification.model_validate(payload, strict=True)
|
|
except (ValueError, TypeError):
|
|
return False
|
|
return True
|
|
|
|
def _set_or_update_auth_context(
|
|
user_api_key_auth: UserAPIKeyAuth | None,
|
|
mcp_auth_header: str | None = None,
|
|
mcp_servers: list[str] | None = None,
|
|
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
|
|
oauth2_headers: dict[str, str] | None = None,
|
|
raw_headers: dict[str, str] | None = None,
|
|
client_ip: str | None = None,
|
|
session_id: str | None = None,
|
|
touch_last_seen: bool = True,
|
|
copy_existing_session_auth_context: bool = False,
|
|
) -> MCPAuthenticatedUser:
|
|
auth_user: Final = _stateful_session_auth_contexts.get(session_id) if session_id else None
|
|
if auth_user is not None and session_id is not None:
|
|
if touch_last_seen:
|
|
_stateful_session_auth_context_last_seen[session_id] = time.monotonic()
|
|
if copy_existing_session_auth_context:
|
|
return set_auth_context(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=client_ip,
|
|
)
|
|
_update_auth_context(
|
|
auth_user=auth_user,
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=client_ip,
|
|
)
|
|
auth_context_var.set(auth_user)
|
|
return auth_user
|
|
return set_auth_context(
|
|
user_api_key_auth=user_api_key_auth,
|
|
mcp_auth_header=mcp_auth_header,
|
|
mcp_servers=mcp_servers,
|
|
mcp_server_auth_headers=mcp_server_auth_headers,
|
|
oauth2_headers=oauth2_headers,
|
|
raw_headers=raw_headers,
|
|
client_ip=client_ip,
|
|
)
|
|
|
|
def _wrap_send_with_stateful_session_auth_context(
|
|
send: Send,
|
|
auth_user: MCPAuthenticatedUser,
|
|
owner_fingerprint: str,
|
|
on_session_registered: Callable[[str], None] | None = None,
|
|
client_info: Implementation | None = None,
|
|
) -> Send:
|
|
async def wrapped_send(message: Message) -> None:
|
|
if message.get("type") == "http.response.start":
|
|
response_headers: Final = _raw_header_pairs(message)
|
|
for key, value in response_headers:
|
|
header_name = key if isinstance(key, bytes) else str(key).encode()
|
|
if header_name.lower() == b"mcp-session-id":
|
|
session_id = value.decode() if isinstance(value, bytes) else str(value)
|
|
if on_session_registered is not None:
|
|
on_session_registered(session_id)
|
|
auth_context_var.set(auth_user)
|
|
_stateful_session_auth_contexts[session_id] = auth_user
|
|
_stateful_session_auth_context_last_seen[session_id] = time.monotonic()
|
|
_stateful_session_owners[session_id] = owner_fingerprint
|
|
operations.global_mcp_server_manager.track_gateway_session(session_id)
|
|
if client_info is not None:
|
|
_stateful_session_client_info[session_id] = client_info
|
|
break
|
|
await send(message)
|
|
|
|
return wrapped_send
|
|
|
|
def get_auth_context() -> tuple[
|
|
UserAPIKeyAuth | None,
|
|
str | None,
|
|
list[str] | None,
|
|
dict[str, dict[str, str]] | None,
|
|
dict[str, str] | None,
|
|
dict[str, str] | None,
|
|
str | None,
|
|
]:
|
|
"""
|
|
Get the UserAPIKeyAuth from the auth context variable.
|
|
|
|
Returns:
|
|
Tuple containing: UserAPIKeyAuth, MCP auth header (deprecated),
|
|
MCP servers, server-specific auth headers, OAuth2 headers, raw headers, client IP
|
|
"""
|
|
auth_user: Final = auth_context_var.get()
|
|
if auth_user and isinstance(auth_user, MCPAuthenticatedUser):
|
|
return (
|
|
auth_user.user_api_key_auth,
|
|
auth_user.mcp_auth_header,
|
|
auth_user.mcp_servers,
|
|
auth_user.mcp_server_auth_headers,
|
|
auth_user.oauth2_headers,
|
|
auth_user.raw_headers,
|
|
auth_user.client_ip,
|
|
)
|
|
return None, None, None, None, None, None, None
|
|
|
|
def _get_current_session():
|
|
ctx: Final = get_active_mcp_request_ctx()
|
|
return ctx.session if ctx is not None else None
|
|
|
|
def _cache_auth_context_lazily():
|
|
session: Final = _get_current_session()
|
|
if session is None:
|
|
return
|
|
try:
|
|
if session in _session_obj_auth_storage:
|
|
return
|
|
except TypeError:
|
|
verbose_logger.debug(
|
|
"_cache_auth_context_lazily: session object is unhashable (type=%s), cannot cache auth context",
|
|
type(session).__name__,
|
|
)
|
|
return
|
|
|
|
auth: Final = auth_context_var.get()
|
|
if auth and isinstance(auth, MCPAuthenticatedUser):
|
|
try:
|
|
_session_obj_auth_storage[session] = auth
|
|
except TypeError:
|
|
verbose_logger.debug(
|
|
"_cache_auth_context_lazily: could not store auth via "
|
|
"session identity — session object is unhashable"
|
|
)
|
|
|
|
def _recover_auth_from_session() -> MCPAuthenticatedUser | None:
|
|
session: Final = _get_current_session()
|
|
if session is None:
|
|
return None
|
|
|
|
stored: MCPAuthenticatedUser | None = None
|
|
try:
|
|
stored = _session_obj_auth_storage.get(session)
|
|
except TypeError:
|
|
verbose_logger.debug(
|
|
"_recover_auth_from_session: session object is unhashable "
|
|
"(type=%s), skipping _session_obj_auth_storage lookup",
|
|
type(session).__name__,
|
|
)
|
|
|
|
return stored
|
|
|
|
async def get_or_extract_auth_context() -> tuple[
|
|
UserAPIKeyAuth | None,
|
|
str | None,
|
|
list[str] | None,
|
|
dict[str, dict[str, str]] | None,
|
|
dict[str, str] | None,
|
|
dict[str, str] | None,
|
|
str | None,
|
|
]:
|
|
"""
|
|
Get auth context from ContextVar first, then fall back to session
|
|
storage (which survives cross-task boundaries in the MCP SDK).
|
|
"""
|
|
(
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
mcp_servers,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
_client_ip,
|
|
) = get_auth_context()
|
|
|
|
if user_api_key_auth is not None:
|
|
_cache_auth_context_lazily()
|
|
else:
|
|
stored: Final = _recover_auth_from_session()
|
|
|
|
if stored:
|
|
user_api_key_auth = stored.user_api_key_auth
|
|
mcp_auth_header = stored.mcp_auth_header
|
|
mcp_servers = stored.mcp_servers
|
|
mcp_server_auth_headers = stored.mcp_server_auth_headers
|
|
oauth2_headers = stored.oauth2_headers
|
|
raw_headers = stored.raw_headers
|
|
_client_ip = stored.client_ip
|
|
return (
|
|
user_api_key_auth,
|
|
mcp_auth_header,
|
|
mcp_servers,
|
|
mcp_server_auth_headers,
|
|
oauth2_headers,
|
|
raw_headers,
|
|
_client_ip,
|
|
)
|
|
|
|
def get_active_mcp_session() -> _McpServerSession | None:
|
|
"""Return the active MCP session captured during handler execution."""
|
|
session: Final = active_mcp_session_var.get()
|
|
if session is not None:
|
|
return session
|
|
return _get_current_session()
|
|
|
|
def get_active_auth_context() -> MCPAuthenticatedUser | None:
|
|
"""Return auth context from ContextVar or session storage."""
|
|
auth: Final = auth_context_var.get()
|
|
if auth and isinstance(auth, MCPAuthenticatedUser):
|
|
return auth
|
|
|
|
stored: Final = _recover_auth_from_session()
|
|
if stored is not None:
|
|
return stored
|
|
return None
|
|
|
|
########################################################
|
|
############ End of Auth Context Functions #############
|
|
########################################################
|
|
|
|
else:
|
|
app = FastAPI()
|