""" LiteLLM MCP Server Routes """ # pyright: reportInvalidTypeForm=false, reportArgumentType=false, reportOptionalCall=false import asyncio import contextlib import contextvars import hashlib import json import time import traceback import types import uuid from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from fastapi import FastAPI, HTTPException from pydantic import AnyUrl, ConfigDict from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse from starlette.types import Message, Receive, Scope, Send from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, _is_mcp_admitted_user_subject, ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPToolResultError, MCPUpstreamAuthError, ) from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_active_toolset_id, _mcp_gateway_initialize_instructions, _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, get_passthrough_www_authenticate, get_route_relative_request_path, well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, LITELLM_MCP_SERVER_VERSION, MCPMissingUserEnvVarsError, add_server_prefix_to_name, build_synthetic_mcp_request, extract_mcp_tool_result_error_message, get_server_prefix, iter_known_server_prefixes, logging_safe_mcp_headers, match_known_tool_name, ) from litellm.proxy._types import ( ProxyException, SpecialMCPServerNames, UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, get_chain_id_from_headers, ) from litellm.types.mcp import MCPAuth, MCPSpecVersion from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall from litellm.utils import Rules, client, function_setup if TYPE_CHECKING: from mcp.server.session import ServerSession as _McpServerSession from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload # Short-lived in-memory cache for BYOK credentials. # Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp). # Storing the credential value (not just a bool) means _get_byok_credential and # _check_byok_credential share a single DB round-trip per TTL window. _byok_cred_cache: Final[dict[tuple[str, str], tuple[str | None, float]]] = {} _BYOK_CRED_CACHE_TTL: Final = 60 # seconds _BYOK_CRED_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth _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 key holding the tracing span of the request carrying an MCP # message, written on the request task and read back by the message handler. _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span" def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. Call this after storing or deleting a credential so subsequent calls see the fresh value rather than a stale cached result. """ _byok_cred_cache.pop((user_id, server_id), None) def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) -> None: """Write a credential value to the cache, evicting all entries if at capacity.""" if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE: _byok_cred_cache.clear() _byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic()) # 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.lowlevel.helper_types import ReadResourceContents from mcp.server.session import ServerSession as _McpServerSession from mcp.types import ( BlobResourceContents, GetPromptResult, ResourceTemplate, TextResourceContents, Tool, ) # 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 ReadResourceContents = 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 _mcp_session_id_from_headers( raw_headers: dict[str, str] | None, ) -> str | None: """The ``mcp-session-id`` of a stateful MCP session, read case-insensitively from the request headers. ``None`` for stateless calls (no such header).""" if not raw_headers: return None for key, value in raw_headers.items(): if isinstance(key, str) and key.lower() == "mcp-session-id": return value or None return None def _jsonrpc_text_has_top_level_method(text: str) -> bool: """Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at the root object's top level. Used to tell a request/notification (carries ``method``) apart from a response (carries ``result``/``error`` and no top-level ``method``). A response payload can itself nest a ``method`` field, so only keys at the root object's depth are inspected rather than searching the whole string. Returns ``True`` only when a top-level ``method`` key is positively found; truncation that hides it yields ``False``. """ depth = 0 in_string = False escaped = False in_object: Final[list[bool]] = [] reading_key = False expect_key = False key_chars: list[str] = [] for ch in 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) == "method": return True 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 False 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 = getattr(meta, "model_extra", None) if not isinstance(extra, dict): 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_transport_span_from_message(req_ctx: object) -> object: """The tracing span of the HTTP request that carried this MCP message. Read off that request's ASGI scope, reached through the ``Request`` the streamable-HTTP transport attaches to each message, so it is this message's transport and not whichever request happens to have touched the session last. Returns whatever the scope holds; the otel plumbing validates it.""" request: Final = getattr(req_ctx, "request", None) scope: Final = getattr(request, "scope", None) if not isinstance(scope, Mapping): return None return scope.get(_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 _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: 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.lowlevel.server import NotificationOptions from mcp.server.models import InitializationOptions try: from mcp.server.streamable_http_manager import StreamableHTTPSessionManager except ImportError: StreamableHTTPSessionManager = None from mcp.types import ( CallToolResult, ListToolsResult, Prompt, TextContent, ) from mcp.types import Tool as MCPTool from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( SERVER_OUTCOMES_META_KEY, AggregateToolListing, ServerListOk, ServerOutcome, classify_list_exception, outcome_wire_value, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _caller_authorization_fans_out, _client_forwarded_authorization_headers, _resolve_openapi_tool_auth, _should_strip_caller_authorization, global_mcp_server_manager, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, _request_resolved_auth_headers, ) from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, is_tool_name_prefixed, normalize_server_name, split_server_prefix_from_name, strip_known_server_prefix, ) from litellm.types.mcp import DEFAULT_CREDENTIAL_HEADER, without_header ###################################################### ############ 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 ###################################################### class ListMCPToolsRestAPIResponseObject(MCPTool): """ Object returned by the /tools/list REST API route. """ mcp_info: MCPInfo | None = None model_config = ConfigDict(arbitrary_types_allowed=True) def _normalize_resource_contents(contents: list) -> list[ReadResourceContents]: """Normalize ResourceContents to ReadResourceContents, preserving meta (MCP 1.26.0+).""" normalized: Final[list[ReadResourceContents]] = [] for content in contents: meta = getattr(content, "meta", None) if meta is None and hasattr(content, "model_dump"): d = content.model_dump() meta = d.get("meta") if meta is None: meta = d.get("_meta") if isinstance(content, TextResourceContents): normalized.append( ReadResourceContents( content=content.text, mime_type=content.mimeType, meta=meta, ) ) elif isinstance(content, BlobResourceContents): normalized.append( ReadResourceContents( content=content.blob, mime_type=content.mimeType, meta=meta, ) ) return normalized def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: opts: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) 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[Server] = Server( name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server) sse: Final[SseServerTransport] = SseServerTransport("/mcp/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]] = {} 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) # Keep this alias so existing references to session_manager still work session_manager: Final = session_manager_stateless # Create SSE session manager sse_session_manager: Final = StreamableHTTPSessionManager( app=server, event_store=None, json_response=False, # Use SSE responses for this endpoint stateless=True, ) # Context managers for proper lifecycle management _session_manager_cm = None _session_manager_stateful_cm = None _sse_session_manager_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) 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, \ _sse_session_manager_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() _sse_session_manager_cm = sse_session_manager.run() # Enter the context managers await _session_manager_cm.__aenter__() await _session_manager_stateful_cm.__aenter__() await _sse_session_manager_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, \ _sse_session_manager_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 if _session_manager_cm: await _session_manager_cm.__aexit__(None, None, None) if _session_manager_stateful_cm: await _session_manager_stateful_cm.__aexit__(None, None, None) if _sse_session_manager_cm: await _sse_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 _sse_session_manager_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 ####################### ######################################################## @server.list_tools() async def handle_list_tools() -> "ListToolsResult | list[Tool]": """ List all available tools, with each server's listing outcome attached to the result's ``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy server with no tools. Returning a ListToolsResult (rather than a bare list) makes the MCP SDK pass the result through unwrapped, which is what lets the ``_meta`` survive to the client. Also captures the active session for propagation to callbacks. """ from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) _transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx)) # Get user authentication from context variable ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() verbose_logger.debug("MCP list_tools - User API Key Auth from context: %s", user_api_key_auth) verbose_logger.debug("MCP list_tools - MCP servers from context: %s", mcp_servers) verbose_logger.debug( "MCP list_tools - MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, ) if getattr( getattr(user_api_key_auth, "object_permission", None), "mcp_tool_search_enabled", False, ): from mcp.types import Tool from litellm.proxy._experimental.mcp_server.tool_search import ( get_virtual_tool_definitions, ) return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") listing: Final = await _list_mcp_tools( 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, log_list_tools_to_spendlogs=True, list_tools_log_source="mcp_protocol", ) verbose_logger.info("MCP list_tools - Successfully returned %s tools", len(listing.tools)) if not listing.outcomes: return listing.tools outcome_meta: Final = { SERVER_OUTCOMES_META_KEY: { key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items() } } return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) except Exception as e: verbose_logger.exception("Error in list_tools endpoint: %s", e) # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] finally: _otel_reset_mcp_transport_span(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) def _capture_host_progress_callback(host_server) -> Callable | None: """Return a progress-forwarding callback bound to the host MCP session. Returns ``None`` when the host did not supply a progress token. """ try: host_ctx: Final = host_server.request_context except Exception as e: verbose_logger.warning("Could not capture host progress context: %s", e) return None if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta): return None host_token: Final = getattr(host_ctx.meta, "progressToken", None) 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): """Forward progress notifications from external MCP to Host""" try: await host_session.send_progress_notification( progress_token=host_token, progress=progress, total=total, ) 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 async def _build_virtual_call_logging_obj( name: str, arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual mcp_tool_call so the SSE path spend-logs like the REST path.""" from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) from litellm.proxy.proxy_server import ( general_settings, proxy_config, proxy_logging_obj, ) request: Final = build_synthetic_mcp_request( path="/mcp/tools/call", raw_headers=raw_headers, client_ip=client_ip, ) _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing( data={"name": name, "arguments": arguments} ).common_processing_pre_call_logic( request=request, user_api_key_dict=user_api_key_auth, proxy_config=proxy_config, route_type=CallTypes.call_mcp_tool.value, proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, ) return virtual_logging_obj async def _dispatch_virtual_mcp_tool( name: str, arguments: dict[str, object] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, mcp_servers: list[str] | None = None, mcp_auth_header: 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, ) -> CallToolResult | None: """Handle the mcp_tool_search / mcp_tool_call virtual tools. Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so the caller falls through to normal tool routing. """ from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, MCP_TOOL_SEARCH_TOOL_NAME, VIRTUAL_TOOL_NAMES, coerce_top_k, handle_agent_search, handle_mcp_tool_call, handle_mcp_tool_search, ) if name not in VIRTUAL_TOOL_NAMES: return None if not getattr( getattr(user_api_key_auth, "object_permission", None), "mcp_tool_search_enabled", False, ): return CallToolResult( content=[ TextContent( type="text", text=f"Tool {name} requires mcp_tool_search_enabled on the key", ) ], isError=True, ) args: Final = arguments or {} if name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( query=args.get("query", ""), top_k=coerce_top_k(args.get("top_k", 5)), user_api_key_dict=user_api_key_auth, client_ip=client_ip, mcp_servers=mcp_servers, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) assert user_api_key_auth is not None # guaranteed by the flag check above if name == AGENT_SEARCH_TOOL_NAME: return await handle_agent_search( query=str(args.get("query", "")), top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K), user_api_key_dict=user_api_key_auth, ) virtual_logging_obj: Final = await _build_virtual_call_logging_obj( name=name, arguments=args, user_api_key_auth=user_api_key_auth, raw_headers=raw_headers, client_ip=client_ip, ) return await handle_mcp_tool_call( tool_name=args.get("tool_name", ""), arguments=args.get("arguments") or {}, user_api_key_dict=user_api_key_auth, client_ip=client_ip, mcp_servers=mcp_servers, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, litellm_logging_obj=virtual_logging_obj, ) @server.call_tool() async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: name (str): Name of the tool to call arguments (Dict[str, Any] | None): Arguments to pass to the tool Returns: List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results Raises: HTTPException: If tool not found or arguments missing """ from mcp.server.lowlevel.server import request_ctx from mcp.types import CallToolResult from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import proxy_config req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) _transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx)) # Validate arguments ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() verbose_logger.debug( "MCP mcp_server_tool_call - user_api_key_auth=%s, user_role=%s", user_api_key_auth, getattr(user_api_key_auth, "user_role", "N/A"), ) verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) try: # Inside this try so virtual-tool errors convert to isError # CallToolResult instead of raising out of the protocol handler. virtual_tool_result: Final = await _dispatch_virtual_mcp_tool( name=name, arguments=arguments, user_api_key_auth=user_api_key_auth, client_ip=_client_ip, mcp_servers=mcp_servers, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) if virtual_tool_result is not None: return virtual_tool_result host_progress_callback: Final = _capture_host_progress_callback(server) # Create a body date for logging body_data: Final = {"name": name, "arguments": arguments} # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) chain_id: Final = get_chain_id_from_headers(raw_headers) if chain_id: body_data["litellm_trace_id"] = chain_id body_data["litellm_session_id"] = chain_id request: Final = build_synthetic_mcp_request( path="/mcp/tools/call", raw_headers=raw_headers, client_ip=_client_ip, ) if user_api_key_auth is not None: data = await add_litellm_data_to_request( data=body_data, request=request, # Bill a team-derived call to the team that granted it. A keyless admitted # subject carries no team_id, so spend skipped team updates entirely and # charged the user's PRIMARY org — the granting team's budget never # accumulated (so it could never begin to block) and, cross-org, the wrong # organization was charged. This is the ACCOUNTING half; the enforcement # half (an already-over-budget team stops granting) lives in the source gate. # Authorization is unaffected: it ran before this, and the union is resolved # from the untouched auth object passed to call_mcp_tool below. user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call( user_api_key_auth, tool_name=name ), proxy_config=proxy_config, ) else: data = body_data response: Final = await call_mcp_tool( 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, host_progress_callback=host_progress_callback, **data, # for logging ) except MCPMissingUserEnvVarsError as e: verbose_logger.info( "MCP mcp_server_tool_call missing per-user env vars: server_id=%s missing=%s", e.server_id, e.missing, ) return CallToolResult( content=[TextContent(text=str(e), type="text")], isError=True, ) except BlockedPiiEntityError as e: verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) return CallToolResult( content=[ TextContent( text=f"Error: Blocked PII entity detected - {e}", type="text", ) ], isError=True, ) except GuardrailRaisedException as e: verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) return CallToolResult( content=[TextContent(text=f"Error: Guardrail violation - {e}", type="text")], isError=True, ) except HTTPException as e: verbose_logger.error("HTTPException in MCP tool call: %s", e) return CallToolResult( content=[TextContent(text=f"Error: {e.detail}", type="text")], isError=True, ) except MCPUpstreamAuthError as e: # The MCP session manager serializes handler exceptions as JSON-RPC errors, so a # mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST # call path and the connect-time preemptive check do. Return an explicit isError # naming the upstream status (at info level, not a traceback) so the client still # learns it must re-authenticate upstream and expected pass-through 401s don't spam. verbose_logger.info("Upstream auth failure calling MCP tool: HTTP %s", e.status_code) return CallToolResult( content=[ TextContent( text=f"Error: upstream authentication required (HTTP {e.status_code})", type="text", ) ], isError=True, ) except Exception as e: verbose_logger.exception("MCP mcp_server_tool_call - error: %s", e) return CallToolResult( content=[TextContent(text=f"Error: {e}", type="text")], isError=True, ) return response finally: _otel_reset_mcp_transport_span(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @server.list_prompts() async def list_prompts() -> list[Prompt]: """ List all available prompts """ from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) try: # Get user authentication from context variable ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() verbose_logger.debug("MCP list_prompts - User API Key Auth from context: %s", user_api_key_auth) verbose_logger.debug("MCP list_prompts - MCP servers from context: %s", mcp_servers) verbose_logger.debug( "MCP list_prompts - MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, ) # Get mcp_servers from context variable verbose_logger.debug("MCP list_prompts - Calling _list_prompts") prompts: Final = await _list_mcp_prompts( 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, ) verbose_logger.info("MCP list_prompts - Successfully returned %s prompts", len(prompts)) return prompts except Exception as e: verbose_logger.exception("Error in list_prompts endpoint: %s", e) # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] finally: if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @server.get_prompt() async def get_prompt(name: str, arguments: dict[str, str] | None) -> GetPromptResult: """ Get a specific prompt with the provided arguments Args: name (str): Name of the prompt to get arguments (Dict[str, Any] | None): Arguments to pass to the prompt Returns: GetPromptResult: Getting prompt execution results """ # Validate arguments from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) try: ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) return await mcp_get_prompt( name=name, arguments=arguments, 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, ) finally: if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @server.list_resources() async def list_resources() -> list[Resource]: """List all available resources.""" from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) try: ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() verbose_logger.debug("MCP list_resources - User API Key Auth from context: %s", user_api_key_auth) verbose_logger.debug("MCP list_resources - MCP servers from context: %s", mcp_servers) verbose_logger.debug( "MCP list_resources - MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, ) resources: Final = await _list_mcp_resources( 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, ) verbose_logger.info("MCP list_resources - Successfully returned %s resources", len(resources)) return resources except Exception as e: verbose_logger.exception("Error in list_resources endpoint: %s", e) return [] finally: if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @server.list_resource_templates() async def list_resource_templates() -> list[ResourceTemplate]: """List all available resource templates.""" from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) try: ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() verbose_logger.debug("MCP list_resource_templates - User API Key Auth from context: %s", user_api_key_auth) verbose_logger.debug("MCP list_resource_templates - MCP servers from context: %s", mcp_servers) verbose_logger.debug( "MCP list_resource_templates - MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, ) resource_templates: Final = await _list_mcp_resource_templates( 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, ) verbose_logger.info( "MCP list_resource_templates - Successfully returned %s resource templates", len(resource_templates) ) return resource_templates except Exception as e: verbose_logger.exception("Error in list_resource_templates endpoint: %s", e) return [] finally: if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @server.read_resource() async def read_resource(url: AnyUrl) -> list[ReadResourceContents]: from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) _session_reset_token = None if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) try: ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, _client_ip, ) = await get_or_extract_auth_context() read_resource_result: Final = await mcp_read_resource( url=url, 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, ) return _normalize_resource_contents(read_resource_result.contents) finally: if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) ######################################################## ############ End of MCP Server Routes ################## ######################################################## ######################################################## ############ Helper Functions ########################## ######################################################## async def _get_allowed_mcp_servers_from_mcp_server_names( mcp_servers: list[str] | None, allowed_mcp_servers: list[MCPServer], ) -> list[MCPServer]: """ Get the filtered MCP servers from the MCP server names. Fails closed when ``mcp_servers`` is explicitly provided (path- or header-derived) but none of the names resolve to a server alias or access group the caller can access. The previous behavior returned the full ``allowed_mcp_servers`` set, which silently widened scope when a client targeted ``/mcp//`` and made URL/header namespacing appear to work when it did not. """ filtered_server: Final[dict[str, MCPServer]] = {} # Filter servers based on mcp_servers parameter if provided if mcp_servers is not None: for server_or_group in mcp_servers: server_name_matched = False for server in allowed_mcp_servers: if server: match_list = [s.lower() for s in iter_known_server_prefixes(server) if s] if server_or_group.lower() in match_list: filtered_server[server.server_id] = server server_name_matched = True break if not server_name_matched: try: access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups( [server_or_group] ) # Only include servers that the user has access to for server_id in access_group_server_ids: for server in allowed_mcp_servers: if server_id == server.server_id: filtered_server[server.server_id] = server except Exception as e: verbose_logger.debug("Could not resolve '%s' as access group: %s", server_or_group, e) if filtered_server: return list(filtered_server.values()) if mcp_servers is not None: # Caller asked for a specific scope but nothing resolved. Fail # closed so URL/header namespacing cannot silently fall back to # the caller's full allowed-server set. verbose_logger.debug( "MCP scope filter resolved to no servers for requested names %s; returning empty list (fail-closed).", mcp_servers, ) return [] return allowed_mcp_servers def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: """ Check if a tool name matches any name in the filter list. Reads the same owner the server-level permission checks use, so discovery hides exactly what dispatch refuses. ``mcp_server`` is required: guessing the boundary at the first separator mismatches every tool on a server whose prefix contains the separator. """ bare_name: Final = strip_known_server_prefix(tool_name, mcp_server) return match_known_tool_name(bare_name, mcp_server, filter_list) is not None def filter_tools_by_allowed_tools( tools: list[MCPTool], mcp_server: MCPServer, ) -> list[MCPTool]: """ Filter tools by allowed/disallowed tools configuration. If allowed_tools is set, only tools in that list are returned. If disallowed_tools is set, tools in that list are excluded. Tool names are matched with and without server prefixes for flexibility. Args: tools: List of tools to filter mcp_server: Server configuration with allowed_tools/disallowed_tools Returns: Filtered list of tools """ from litellm.proxy._experimental.mcp_server.utils import ( server_applies_tool_allowlist, ) tools_to_return = tools # Filter by allowed_tools (whitelist) if server_applies_tool_allowlist(mcp_server): if not mcp_server.allowed_tools: return [] tools_to_return = [ tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server) ] # Filter by disallowed_tools (blacklist) if mcp_server.disallowed_tools: tools_to_return = [ tool for tool in tools_to_return if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server) ] return tools_to_return def apply_tool_overrides( tools: list[MCPTool], mcp_server: MCPServer, ) -> list[MCPTool]: """Apply admin-configured display name/description overrides to tools. Overrides are keyed by the unprefixed tool name, same convention as allowed_tools configuration. """ display_name_map: Final = mcp_server.tool_name_to_display_name or {} description_map: Final = mcp_server.tool_name_to_description or {} if not display_name_map and not description_map: return tools for tool in tools: unprefixed = strip_known_server_prefix(tool.name, mcp_server) lookup_key = unprefixed or tool.name if lookup_key in display_name_map: tool.name = display_name_map[lookup_key] if lookup_key in description_map: tool.description = description_map[lookup_key] return tools def _get_client_ip_from_context() -> str | None: """ Extract client_ip from auth context. Returns None if context not set (caller should handle this as "no IP filtering"). """ try: auth_user: Final = auth_context_var.get() if auth_user and isinstance(auth_user, MCPAuthenticatedUser): return auth_user.client_ip except Exception: pass return None async def _get_allowed_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: list[str] | None, client_ip: str | None = None, ) -> list[MCPServer]: """Return allowed MCP servers for a request after applying filters. Args: user_api_key_auth: The authenticated user's API key info. mcp_servers: Optional list of server names to filter to. client_ip: Client IP for IP-based access control. If None, falls back to auth context. Pass explicitly from request handlers for safety. Note: If client_ip is None and auth context is not set, IP filtering is skipped. This is intentional for internal callers but may indicate a bug if called from a request handler without proper context setup. """ # Use explicit client_ip if provided, otherwise try auth context if client_ip is None: client_ip = _get_client_ip_from_context() if client_ip is None: verbose_logger.debug( "MCP _get_allowed_mcp_servers called without client_ip and no auth context. " "IP filtering will be skipped. This is expected for internal calls." ) allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) ( allowed_mcp_server_ids, _ip_blocked, ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info(allowed_mcp_server_ids, client_ip) verbose_logger.debug( "MCP IP filter: client_ip=%s, allowed_server_ids=%s", client_ip, allowed_mcp_server_ids, ) if _ip_blocked > 0: verbose_logger.debug( "MCP IP filtering: %d server(s) are not accessible from client IP %s " "because they are restricted to internal networks. " "No tools from those servers will be returned. " "To expose a server externally, set 'available_on_public_internet: true' " "in its configuration.", _ip_blocked, client_ip, ) allowed_mcp_servers: list[MCPServer] = [] for allowed_mcp_server_id in allowed_mcp_server_ids: mcp_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) if mcp_server is not None: # Apply the request-time oauth2_flow backstop for legacy null rows. mcp_server = MCPServerManager.resolve_oauth2_flow_for_request(mcp_server) allowed_mcp_servers.append(mcp_server) if mcp_servers is not None: allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) return allowed_mcp_servers def _client_has_per_server_auth_header( server: MCPServer, mcp_server_auth_headers: dict[str, dict[str, str]] | None, ) -> bool: """True if the request carries a per-server ``x-mcp-{alias}-authorization`` header for this server. This is the multi-server binding: it names one upstream, so it is unambiguously the caller's upstream token regardless of auth mode (never the LiteLLM admission credential). Resolves through the same ``lookup_mcp_server_auth_in_headers`` egress uses, so the connect gate and egress agree on which per-server header names match: a dashboard client sends ``x-mcp-{sanitize_mcp_alias_for_header(alias)}-authorization``, and matching only the raw alias here would 401 a token egress would forward. """ if not mcp_server_auth_headers: return False from litellm.proxy._experimental.mcp_server.utils import ( lookup_mcp_server_auth_in_headers, ) server_headers: Final = lookup_mcp_server_auth_in_headers( mcp_server_auth_headers, alias=server.alias, server_name=server.server_name ) if isinstance(server_headers, str): return bool(server_headers.strip()) if isinstance(server_headers, dict): return any(isinstance(hk, str) and hk.lower() == "authorization" for hk in server_headers) return False def _client_has_passthrough_authorization( server: MCPServer, oauth2_headers: dict[str, str] | None, mcp_server_auth_headers: dict[str, dict[str, str]] | None, ) -> bool: """True if the incoming request already carries an ``Authorization`` header the gateway will forward to this pass-through server. The client may supply the bearer as either the top-level ``Authorization`` header (surfaced via ``oauth2_headers``) or a per-server ``x-mcp-auth-`` style header (surfaced via ``mcp_server_auth_headers``). Either form skips the pre-emptive 401. """ if oauth2_headers: for k in oauth2_headers: if k.lower() == "authorization": return True return _client_has_per_server_auth_header(server, mcp_server_auth_headers) async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. Thin wrapper over ``resolve_user_oauth_access_token`` (Redis cache, else DB + refresh); ``prefetched_creds`` skips the per-server Redis/DB lookups for the batch path. """ if server.auth_type != MCPAuth.oauth2 or user_api_key_auth is None: return None from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 resolve_user_oauth_access_token, ) token: Final = await resolve_user_oauth_access_token( getattr(user_api_key_auth, "user_id", None), server, prefetched_creds ) return {"Authorization": f"Bearer {token}"} if token else None async def _prefetch_oauth_creds_for_user( user_api_key_auth: UserAPIKeyAuth | None, ) -> dict[str, "OAuthCredentialPayload"]: """Fetch all OAuth2 credentials for the user in one DB query. Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not user_id: return {} try: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 list_user_oauth_credentials, ) from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 prisma_client: Final = get_prisma_client_or_throw( "Database not connected. Connect a database to use OAuth2 MCP tools." ) creds: Final = await list_user_oauth_credentials(prisma_client, user_id) return {c["server_id"]: c for c in creds if "server_id" in c} except Exception as e: verbose_logger.warning("_prefetch_oauth_creds_for_user: failed to prefetch for user=%s: %s", user_id, e) return {} def _prepare_mcp_server_headers( server: MCPServer, mcp_server_auth_headers: dict[str, dict[str, str]] | None, mcp_auth_header: str | None, oauth2_headers: dict[str, str] | None, raw_headers: dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None = None, scope_servers: list[MCPServer] | None = None, ) -> tuple[dict[str, str] | str | None, dict[str, str] | None]: """Build auth and extra headers for a server. ``scope_servers`` is the full server list a fan-out handler iterates. Passing it lets the client-forwarded token modes withhold the caller's request-wide ``Authorization`` when another server in the scope would also receive it (``_caller_authorization_fans_out``); explicitly-addressed operations leave it None. Per-server ``x-mcp-{alias}-authorization`` headers are unaffected — they bind one token to one server and are the multi-server shape. """ server_auth_header: dict[str, str] | str | None = None if mcp_server_auth_headers: from litellm.proxy._experimental.mcp_server.utils import ( lookup_mcp_server_auth_in_headers, ) server_auth_header = lookup_mcp_server_auth_in_headers( mcp_server_auth_headers, alias=server.alias, server_name=server.server_name, ) extra_headers: dict[str, str] | None = None is_client_forwarded_mode: Final = server.is_client_forwarded_token # In a multi-server listing scope the request-wide Authorization can only carry one token, # so it is withheld from a client-forwarded server when another server in scope also consumes # it (RFC 9700 cross-resource replay); such scopes must bind per-server via # x-mcp-{alias}-authorization. The decision is computed once so BOTH the forwarding branch and # the extra_headers copy loop below honor it — otherwise a server that lists Authorization in # extra_headers would re-copy the withheld bearer from raw_headers and replay it anyway. withhold_forwarded_authorization: Final = is_client_forwarded_mode and _caller_authorization_fans_out( server, scope_servers ) if server.auth_type == MCPAuth.oauth2: # For OAuth2 M2M servers, upstream Authorization must come from # client_credentials token fetch, never from caller headers. if server.has_client_credentials: extra_headers = None else: # Copy to avoid mutating the original dict (important for parallel fetching) extra_headers = oauth2_headers.copy() if oauth2_headers else None # Migrated authorization_code: the v2 resolver injects the stored per-user # token, so drop the caller-forwarded Authorization (apply-if-absent would # otherwise let it shadow the resolved token). Delegate keeps it. Centralized # via _should_strip_caller_authorization to match _call_regular_mcp_tool. if extra_headers and _should_strip_caller_authorization( mcp_server=server, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, ): extra_headers = without_header(extra_headers, DEFAULT_CREDENTIAL_HEADER) elif is_client_forwarded_mode: if not withhold_forwarded_authorization: extra_headers = _client_forwarded_authorization_headers( mcp_server=server, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, ) if server.extra_headers and raw_headers: if extra_headers is None: extra_headers = {} normalized_raw_headers: Final = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} # Centralized strip decision shared with # ``MCPServerManager._call_regular_mcp_tool`` so the two # code paths cannot drift on this security-sensitive choice. # See ``_should_strip_caller_authorization`` for the rules. strip_caller_authorization: Final = _should_strip_caller_authorization( mcp_server=server, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, ) for header in server.extra_headers: if not isinstance(header, str): continue if header.lower() == "authorization" and ( strip_caller_authorization or withhold_forwarded_authorization ): continue header_value = normalized_raw_headers.get(header.lower()) if header_value is None: continue extra_headers[header] = header_value # Reset to None if no headers were actually added if extra_headers is not None and len(extra_headers) == 0: extra_headers = None if server_auth_header is None: server_auth_header = mcp_auth_header return server_auth_header, extra_headers def _merge_gateway_initialize_instructions( allowed_mcp_servers: list[MCPServer], ) -> str | None: """YAML/DB override, else upstream text (prefetch on init, or list_tools / health_check / call_tool cache).""" if not allowed_mcp_servers: return None texts: Final[list[tuple[str, str]]] = [] for server in allowed_mcp_servers: label = server.alias or server.server_name or server.name or server.server_id or "mcp" if server.instructions and server.instructions.strip(): texts.append((label, server.instructions.strip())) continue if server.spec_path: continue cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(server.server_id) if cached and cached.strip(): texts.append((label, cached.strip())) if not texts: return None if len(texts) == 1: return texts[0][1] return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) @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, ) -> AsyncIterator[None]: allowed: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=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( *[ global_mcp_server_manager._ensure_upstream_initialize_instructions_cached(s) for s in allowed if s is not None ], return_exceptions=True, ) merged: Final = _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) def _aggregate_server_key(server: MCPServer) -> str: """The client-visible key for a server in listing outcomes and spend metadata: the same display prefix (alias, or the short prefix when that mode is enabled) the caller already sees on the tool names. Canonical internal server names never key a caller-readable surface; when the display naming deliberately hides them, the outcome keys must too.""" return get_server_prefix(server) or "unknown" async def _get_tools_from_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | None, mcp_servers: list[str] | 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, log_list_tools_to_spendlogs: bool = False, list_tools_log_source: str | None = None, litellm_trace_id: str | None = None, request_tags: list[str] | None = None, client_ip: str | None = None, ) -> AggregateToolListing: """ Helper method to fetch tools from MCP servers based on server filtering criteria. Args: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers oauth2_headers: Optional dict of oauth2 headers Returns: AggregateToolListing: Combined tools from filtered servers plus each server's classified listing outcome """ if not MCP_AVAILABLE: return AggregateToolListing(tools=[], outcomes={}) list_tools_start_time: Final = datetime.now() litellm_logging_obj: LiteLLMLoggingObj | None = None list_tools_request_data: dict[str, object] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook rules_obj: Final = Rules() list_tools_call_id: Final = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) spend_logs_metadata: Final[dict[str, object]] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): spend_logs_metadata["source"] = list_tools_log_source if isinstance(mcp_servers, list): spend_logs_metadata["requested_mcp_servers"] = mcp_servers list_tools_request_data = { "model": "MCP: list_tools", "call_type": CallTypes.list_mcp_tools.value, "litellm_call_id": list_tools_call_id, "litellm_trace_id": effective_litellm_trace_id, "metadata": { "spend_logs_metadata": spend_logs_metadata, "headers": logging_safe_mcp_headers(raw_headers), **({"tags": request_tags} if request_tags else {}), }, # Provide a small input payload for standard logging "input": [ { "role": "system", "content": { "mcp_operation": "list_tools", "requested_mcp_servers": mcp_servers, }, } ], } # Attach user identifiers using the standard helper if user_api_key_auth is not None: LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=list_tools_request_data, user_api_key_dict=user_api_key_auth, _metadata_variable_name="metadata", ) user_identifier: Final = getattr(user_api_key_auth, "end_user_id", None) or getattr( user_api_key_auth, "user_id", None ) if user_identifier: list_tools_request_data["user"] = user_identifier try: litellm_logging_obj, _ = function_setup( original_function="list_mcp_tools", rules_obj=rules_obj, start_time=list_tools_start_time, **list_tools_request_data, ) if litellm_logging_obj: litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value litellm_logging_obj.model = "MCP: list_tools" except Exception as logging_error: verbose_logger.debug("Failed to initialize logging for MCP list_tools: %s", logging_error) litellm_logging_obj = None try: allowed_mcp_servers: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) # Pre-fetch OAuth credentials only when at least one server uses OAuth2, # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. _has_oauth2_server = any(getattr(s, "auth_type", None) == MCPAuth.oauth2 for s in allowed_mcp_servers) _prefetched_oauth_creds: Final = ( await _prefetch_oauth_creds_for_user(user_api_key_auth) if _has_oauth2_server else {} ) async def _fetch_and_filter_server_tools( server: MCPServer, ) -> "tuple[list[MCPTool], ServerOutcome]": """Fetch and filter tools from a single server, classifying any failure into that server's outcome so the aggregate can keep serving the healthy subset without a broken server masquerading as an empty one.""" if server is None: return [], ServerListOk(tool_count=0) server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, scope_servers=allowed_mcp_servers, ) # Prefer server-stored per-user OAuth when configured, so a stale # Authorization header from the MCP client cannot override Redis/DB # (same issue as call_tool in mcp_server_manager: VS Code caches tokens). from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 to_server_spec, ) # A server migrated to the v2 resolver gets its token from the resolver at connect # time; building it here would double-resolve and be shadowed by the v2 graft. The # preemptive 401 already challenged a missing token, so one exists for the connect. migrated_to_v2: Final = to_server_spec(server) is not None if ( not migrated_to_v2 and server.auth_type == MCPAuth.oauth2 and getattr(server, "needs_user_oauth_token", False) and user_api_key_auth is not None ): db_headers: Final = await _get_user_oauth_extra_headers_from_db( server, user_api_key_auth, prefetched_creds=_prefetched_oauth_creds, ) if db_headers: extra_headers = db_headers # If still no OAuth2 token, fall back to pre-fetched creds (non-stale-client path) elif not migrated_to_v2 and extra_headers is None and server.auth_type == MCPAuth.oauth2: extra_headers = await _get_user_oauth_extra_headers_from_db( server, user_api_key_auth, prefetched_creds=_prefetched_oauth_creds, ) if server.is_byok and server.auth_type != MCPAuth.oauth2 and server_auth_header is None: server_auth_header = await _get_byok_credential(server, user_api_key_auth) try: tools: Final = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, oauth2_headers=oauth2_headers, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) filtered_tools = await filter_tools_by_key_team_permissions( tools=filtered_tools, server_id=server.server_id, user_api_key_auth=user_api_key_auth, ) # Apply display-name/description overrides last so that # permission filtering always works against original names. filtered_tools = apply_tool_overrides(filtered_tools, server) verbose_logger.debug( "Successfully fetched %s tools from server %s, %s after filtering", len(tools), server.name, len(filtered_tools), ) return filtered_tools, ServerListOk(tool_count=len(filtered_tools)) except MCPUpstreamAuthError as e: # Absorb so one unauthenticated server does not empty every other server's # tools. Surfacing the upstream 401 to the client as a re-auth challenge is # intentionally not done here: raising from this list handler cannot produce a # 401 + WWW-Authenticate (the MCP session manager serializes it as a JSON-RPC # error). Single-server routes surface it via the request-scope preemptive # check in _raise_preemptive_401_for_unauthenticated_servers instead. verbose_logger.debug("MCP list_tools: omitting %s; it needs upstream auth", server.name) return [], classify_list_exception(e) except Exception as e: verbose_logger.exception("Error getting tools from server %s: %s", server.name, e) return [], classify_list_exception(e) # Fetch tools from all servers in parallel tasks: Final = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] results: Final = await asyncio.gather(*tasks) # Flatten results into single list all_tools: Final[list[MCPTool]] = [tool for tools, _ in results for tool in tools] server_outcomes: Final[dict[str, ServerOutcome]] = { _aggregate_server_key(server): outcome for server, (_, outcome) in zip(allowed_mcp_servers, results) if server is not None } # If logging is enabled, enrich spend_logs_metadata with counts if litellm_logging_obj: per_server_tool_counts: Final[dict[str, int]] = { _aggregate_server_key(server): len(server_tools) for server, (server_tools, _) in zip(allowed_mcp_servers, results) if server is not None } metadata_dict: Final = litellm_logging_obj.model_call_details.get("metadata") if isinstance(metadata_dict, dict): spend_meta = metadata_dict.get("spend_logs_metadata") if not isinstance(spend_meta, dict): spend_meta = {} metadata_dict["spend_logs_metadata"] = spend_meta spend_meta["allowed_server_count"] = len(allowed_mcp_servers) spend_meta["tool_count_total"] = len(all_tools) spend_meta["per_server_tool_counts"] = per_server_tool_counts spend_meta["per_server_list_outcomes"] = { key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items() } end_time: Final = datetime.now() try: await litellm_logging_obj.async_success_handler( result=[ tool.model_dump(mode="json") if isinstance(tool, MCPTool) else tool for tool in all_tools ], start_time=list_tools_start_time, end_time=end_time, ) except Exception as log_exc: # list_tools responses must not be dropped due to non-blocking # observability/serialization failures. verbose_logger.warning( "MCP list_tools success logging failed (continuing): %s", log_exc, ) verbose_logger.info("Successfully fetched %s tools total from all MCP servers", len(all_tools)) return AggregateToolListing(tools=all_tools, outcomes=server_outcomes) except Exception as e: # Only fire failure hook if logging was requested for this list-tools execution if log_list_tools_to_spendlogs and user_api_key_auth is not None: try: from litellm.proxy.proxy_server import proxy_logging_obj if proxy_logging_obj: traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) await proxy_logging_obj.post_call_failure_hook( request_data=list_tools_request_data or {}, original_exception=e, user_api_key_dict=user_api_key_auth, route="/mcp/list_tools", traceback_str=traceback_str, ) except Exception: verbose_logger.debug("Failed to log MCP list_tools failure via post_call_failure_hook") raise async def _get_prompts_from_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | None, mcp_servers: list[str] | 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, ) -> list[Prompt]: """ Helper method to fetch prompt from MCP servers based on server filtering criteria. Args: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers oauth2_headers: Optional dict of oauth2 headers Returns: List[Prompt]: Combined list of prompts from filtered servers """ if not MCP_AVAILABLE: return [] allowed_mcp_servers: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) # Get prompts from each allowed server all_prompts: Final = [] for server in allowed_mcp_servers: if server is None: continue server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, scope_servers=allowed_mcp_servers, ) try: prompts = await global_mcp_server_manager.get_prompts_from_server( server=server, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, ) all_prompts.extend(prompts) verbose_logger.debug("Successfully fetched %s prompts from server %s", len(prompts), server.name) except Exception as e: verbose_logger.exception("Error getting prompts from server %s: %s", server.name, e) # Continue with other servers instead of failing completely verbose_logger.info("Successfully fetched %s prompts total from all MCP servers", len(all_prompts)) return all_prompts async def _get_resources_from_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | None, mcp_servers: list[str] | 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, ) -> list[Resource]: """Fetch resources from allowed MCP servers.""" if not MCP_AVAILABLE: return [] allowed_mcp_servers: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) all_resources: Final[list[Resource]] = [] for server in allowed_mcp_servers: if server is None: continue server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, scope_servers=allowed_mcp_servers, ) try: resources = await global_mcp_server_manager.get_resources_from_server( server=server, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, ) all_resources.extend(resources) verbose_logger.debug("Successfully fetched %s resources from server %s", len(resources), server.name) except Exception as e: verbose_logger.exception("Error getting resources from server %s: %s", server.name, e) verbose_logger.info("Successfully fetched %s resources total from all MCP servers", len(all_resources)) return all_resources async def _get_resource_templates_from_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | None, mcp_servers: list[str] | 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, ) -> list[ResourceTemplate]: """Fetch resource templates from allowed MCP servers.""" if not MCP_AVAILABLE: return [] allowed_mcp_servers: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) all_resource_templates: Final[list[ResourceTemplate]] = [] for server in allowed_mcp_servers: if server is None: continue server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, scope_servers=allowed_mcp_servers, ) try: resource_templates = await global_mcp_server_manager.get_resource_templates_from_server( server=server, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, ) all_resource_templates.extend(resource_templates) verbose_logger.debug( "Successfully fetched %s resource templates from server %s", len(resource_templates), server.name, ) except Exception as e: verbose_logger.exception( "Error getting resource templates from server %s: %s", server.name, str(e), ) verbose_logger.info( "Successfully fetched %s resource templates total from all MCP servers", len(all_resource_templates), ) return all_resource_templates async def filter_tools_by_key_team_permissions( tools: list[MCPTool], server_id: str, user_api_key_auth: UserAPIKeyAuth | None, ) -> list[MCPTool]: """ Filter tools based on key/team mcp_tool_permissions. Note: Tool names in the DB are stored without server prefixes, but tool names from MCP servers are prefixed. We need to strip the prefix before comparing. """ # Filter by key/team tool-level permissions allowed_tool_names: Final = await MCPRequestHandler.get_allowed_tools_for_server( server_id=server_id, user_api_key_auth=user_api_key_auth, ) # Tools arrive prefixed with the server's own prefix; strip exactly that # prefix (resolved from the server) rather than the first separator, so a # prefix containing the separator still reduces to the stored bare name. server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id) return [ t for t in tools if MCPRequestHandler.tool_is_granted(strip_known_server_prefix(t.name, server), allowed_tool_names) ] async def _list_mcp_tools( user_api_key_auth: UserAPIKeyAuth | None = 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, log_list_tools_to_spendlogs: bool = False, list_tools_log_source: str | None = None, client_ip: str | None = None, ) -> AggregateToolListing: """ List all available MCP tools. Args: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} client_ip: Client IP for IP-based server access control Returns: AggregateToolListing: Combined tools from all accessible servers plus each server's classified listing outcome """ if not MCP_AVAILABLE: return AggregateToolListing(tools=[], outcomes={}) try: listing: Final = await _get_tools_from_mcp_servers( 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, log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, list_tools_log_source=list_tools_log_source, client_ip=client_ip, ) verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools)) return listing except Exception as e: verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) # Continue with an empty listing instead of failing completely return AggregateToolListing(tools=[], outcomes={}) async def _list_mcp_prompts( user_api_key_auth: UserAPIKeyAuth | None = 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, ) -> list[Prompt]: """ List all available MCP prompts. Args: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} Returns: List[Prompt]: Combined list of tools from all accessible servers """ if not MCP_AVAILABLE: return [] # Get tools from managed MCP servers with error handling managed_prompts = [] try: managed_prompts = await _get_prompts_from_mcp_servers( 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, ) verbose_logger.debug("Successfully fetched %s prompts from managed MCP servers", len(managed_prompts)) except Exception as e: verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) # Continue with empty managed tools list instead of failing completely return managed_prompts async def _list_mcp_resources( user_api_key_auth: UserAPIKeyAuth | None = 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, ) -> list[Resource]: """List all available MCP resources.""" if not MCP_AVAILABLE: return [] managed_resources: list[Resource] = [] try: managed_resources = await _get_resources_from_mcp_servers( 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, ) verbose_logger.debug("Successfully fetched %s resources from managed MCP servers", len(managed_resources)) except Exception as e: verbose_logger.exception("Error getting resources from managed MCP servers: %s", e) return managed_resources async def _list_mcp_resource_templates( user_api_key_auth: UserAPIKeyAuth | None = 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, ) -> list[ResourceTemplate]: """List all available MCP resource templates.""" if not MCP_AVAILABLE: return [] managed_resource_templates: list[ResourceTemplate] = [] try: managed_resource_templates = await _get_resource_templates_from_mcp_servers( 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, ) verbose_logger.debug( "Successfully fetched %s resource templates from managed MCP servers", len(managed_resource_templates), ) except Exception as e: verbose_logger.exception( "Error getting resource templates from managed MCP servers: %s", str(e), ) return managed_resource_templates def _resolve_display_name_to_original( name: str, allowed_mcp_servers: list[MCPServer], ) -> str: """Translate a display-name override back to the original prefixed tool name. When a client received a customised display name from tools/list (e.g. "Get Pet") it will call tools/call with that same string. We need to reverse-map it to the original prefixed name (e.g. "petstore_mcp-getPetById") before any routing or permission logic runs. """ for server in allowed_mcp_servers: display_map = server.tool_name_to_display_name or {} for unprefixed_name, display_name in display_map.items(): if display_name == name: return add_server_prefix_to_name(unprefixed_name, get_server_prefix(server)) return name async def _get_byok_credential( mcp_server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, ) -> str | None: """Retrieve the stored BYOK credential for a user+server pair. Uses the shared _byok_cred_cache to avoid a DB round-trip on every tool call within the TTL window. """ if not mcp_server.is_byok: return None user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or "" if not user_id: return None cache_key: Final = (user_id, mcp_server.server_id) cached: Final = _byok_cred_cache.get(cache_key) if cached is not None: credential, ts = cached if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL: return credential from litellm.proxy._experimental.mcp_server.db import get_user_credential from litellm.proxy.proxy_server import prisma_client if prisma_client is None: return None credential = await get_user_credential( prisma_client=prisma_client, user_id=user_id, server_id=mcp_server.server_id, ) _write_byok_cred_cache(user_id, mcp_server.server_id, credential) return credential async def _check_byok_credential( mcp_server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, ) -> None: """ If the MCP server is BYOK-enabled, verify that the requesting user has a stored credential. When no credential is found, raise an HTTP 401 with a WWW-Authenticate header that points the MCP client to our OAuth metadata endpoint so it can drive the authorization flow. """ if not mcp_server.is_byok: return user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or "" if not user_id: raise HTTPException( status_code=401, detail={ "error": "byok_auth_required", "server_id": mcp_server.server_id, "server_name": mcp_server.server_name or mcp_server.name, "message": "User identity is required for BYOK servers", }, headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, ) # Check shared credential cache before hitting the DB. cache_key: Final = (user_id, mcp_server.server_id) cached: Final = _byok_cred_cache.get(cache_key) if cached is not None: cached_cred, ts = cached if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL: if cached_cred is None: raise HTTPException( status_code=401, detail={ "error": "byok_auth_required", "server_id": mcp_server.server_id, "server_name": mcp_server.server_name or mcp_server.name, "message": ( "No stored credential found for this BYOK server. " "Complete the OAuth authorization flow to provide your API key." ), }, headers={ "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' }, ) return from litellm.proxy._experimental.mcp_server.db import get_user_credential from litellm.proxy.proxy_server import prisma_client if prisma_client is None: # Fail closed on DB unavailability: returning here previously # bypassed the ownership check and let any proxy-authenticated # caller invoke BYOK tools during outage windows. raise HTTPException( status_code=503, detail={ "error": "byok_auth_unavailable", "server_id": mcp_server.server_id, "server_name": mcp_server.server_name or mcp_server.name, "message": "BYOK credential check requires a database connection.", }, ) credential: Final = await get_user_credential( prisma_client=prisma_client, user_id=user_id, server_id=mcp_server.server_id, ) _write_byok_cred_cache(user_id, mcp_server.server_id, credential) if credential is None: raise HTTPException( status_code=401, detail={ "error": "byok_auth_required", "server_id": mcp_server.server_id, "server_name": mcp_server.server_name or mcp_server.name, "message": ( "No stored credential found for this BYOK server. " "Complete the OAuth authorization flow to provide your API key." ), }, headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, ) async def execute_mcp_tool( name: str, arguments: dict[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: 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, host_progress_callback: Callable | None = None, **kwargs: Any, ) -> CallToolResult: """ Execute MCP tool. This function assumes permission checks have already been performed. Args: name: Tool name (may include server prefix) arguments: Tool arguments allowed_mcp_servers: Pre-validated list of servers the user can access start_time: Start time for logging user_api_key_auth: Optional user API key auth for logging mcp_auth_header: Optional MCP auth header mcp_server_auth_headers: Optional server-specific auth headers oauth2_headers: Optional OAuth2 headers raw_headers: Optional raw HTTP headers **kwargs: Additional arguments (e.g., litellm_logging_obj) Returns: CallToolResult: Tool execution result """ # Track resolved MCP server for both permission checks and dispatch mcp_server: MCPServer | None = None requested_server_id: Final[str | None] = kwargs.get("requested_server_id") # If the client called with a display-name override (e.g. "Get Pet"), # translate it back to the original prefixed name before any routing. name = _resolve_display_name_to_original(name, allowed_mcp_servers) # Remove prefix from tool name for logging and processing original_tool_name, server_name = split_server_prefix_from_name(name) requested_server: MCPServer | None = None if requested_server_id: requested_server = next( (s for s in allowed_mcp_servers if s.server_id == requested_server_id), None, ) name_is_prefixed = False if requested_server is not None and MCP_TOOL_PREFIX_SEPARATOR in name: all_registry_prefixes: Final[set[str]] = set() for registry_server in global_mcp_server_manager.get_registry().values(): for known_prefix in iter_known_server_prefixes(registry_server): all_registry_prefixes.add(normalize_server_name(known_prefix)) name_is_prefixed = is_tool_name_prefixed(name, known_server_prefixes=all_registry_prefixes) if requested_server is not None and not name_is_prefixed: # REST callers may pass server_id with the upstream tool name (no # LiteLLM prefix). The first segment is not a registered server # prefix, so the whole string is the upstream tool name and may # legitimately contain the separator (e.g. "text-to-speech"). # server_id is authoritative for routing and auth. mcp_server = requested_server server_name = requested_server.name original_tool_name = name else: # Resolve from tool name (MCP JSON-RPC or prefixed REST tool names). mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server is None and requested_server is not None: for known_prefix in iter_known_server_prefixes(requested_server): candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( add_server_prefix_to_name(name, known_prefix) ) if candidate is not None: mcp_server = candidate break if mcp_server is not None: server_name = mcp_server.name original_tool_name = strip_known_server_prefix(name, mcp_server) if requested_server is not None: if mcp_server is not None and mcp_server.server_id != requested_server.server_id: raise HTTPException( status_code=403, detail={ "error": "tool_server_mismatch", "message": ( f"Tool '{name}' belongs to MCP server " f"'{mcp_server.name}' but request specified " f"server_id for '{requested_server.name}'." ), }, ) if mcp_server is None: mcp_server = requested_server server_name = requested_server.name original_tool_name = strip_known_server_prefix(name, requested_server) # Only enforce server-level permissions when we can resolve a server if server_name: if not MCPRequestHandler.is_tool_allowed( allowed_mcp_servers=[server.name for server in allowed_mcp_servers], server_name=server_name, ): raise HTTPException( status_code=403, detail="User not allowed to call this tool.", ) standard_logging_mcp_tool_call: Final[StandardLoggingMCPToolCall] = _get_standard_logging_mcp_tool_call( name=original_tool_name, # Use original name for logging arguments=arguments, server_name=server_name, session_id=_mcp_session_id_from_headers(raw_headers), ) litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) if litellm_logging_obj: litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call litellm_logging_obj.model = f"MCP: {name}" litellm_logging_obj.model_call_details["model"] = f"MCP: {name}" # Resolve the MCP server early so BYOK checks and credential injection # apply to ALL dispatch paths (local tool registry AND managed MCP server). if mcp_server is None: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server: standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get( "mcp_server_cost_info" ) if litellm_logging_obj: litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call # BYOK: retrieve the stored per-user credential. A single DB call # both checks existence and fetches the value, avoiding a double query. if mcp_server.is_byok and not mcp_auth_header: byok_cred: Final = await _get_byok_credential(mcp_server, user_api_key_auth) if byok_cred is None: raise HTTPException( status_code=401, detail={ "error": "byok_auth_required", "server_id": mcp_server.server_id, "server_name": mcp_server.server_name or mcp_server.name, "message": ( "No stored credential found for this BYOK server. " "Complete the OAuth authorization flow to provide your API key." ), }, headers={ "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' }, ) mcp_auth_header = byok_cred elif mcp_server.is_byok: # External auth header supplied; still enforce user-identity check. await _check_byok_credential(mcp_server, user_api_key_auth) # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names ######################################################### local_tool: Final = global_mcp_tool_registry.get_tool(name) if local_tool: # OpenAPI-backed tools used to bypass `pre_call_tool_check` — # only the managed path ran allowed/banned-tool checks, key/team # tool permissions, and parameter validation. Run the same checks # before dispatching to the local registry. Refuse the call if # we cannot resolve a server: tools registered via # openapi_to_mcp_generator are always tied to a server, so a # missing mcp_server here means the tool->server mapping has # not finished initializing or the registry entry is orphaned. # Skipping the check would re-open the same authorization gap. if mcp_server is None: raise HTTPException( status_code=503, detail=( f"MCP server for tool '{name}' is not available; " "refusing to dispatch without authorization checks. " "Retry once the server is registered." ), ) # `pre_call_tool_check` calls into `proxy_logging_obj` for the # pre-call guardrail hooks, so source it from the canonical # `proxy_server` module the same way `_handle_managed_mcp_tool` # does. `kwargs.get("proxy_logging_obj")` is None on the MCP # entry path and would crash with AttributeError after the # security checks pass. from litellm.proxy.proxy_server import proxy_logging_obj hook_result = await global_mcp_server_manager.pre_call_tool_check( name=original_tool_name, arguments=arguments or {}, server_name=server_name or mcp_server.name, user_api_key_auth=user_api_key_auth, proxy_logging_obj=proxy_logging_obj, server=mcp_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. if isinstance(hook_result, dict) and "arguments" in hook_result: arguments = hook_result["arguments"] verbose_logger.debug("Executing local registry tool: %s", name) # The credential rides ContextVars because the tool function has its # headers baked into the closure at registration time. auth_header_value, openapi_forwarded_headers, upstream_credential = _resolve_openapi_tool_auth( mcp_server=mcp_server, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, ) ( resolved_auth_headers, forwarded_headers, ) = await global_mcp_server_manager.resolve_openapi_upstream_auth( mcp_server=mcp_server, oauth2_headers=oauth2_headers, raw_headers=raw_headers, mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, ) _auth_token: Final = _request_auth_header.set(auth_header_value) _extra_token: Final = _request_extra_headers.set(forwarded_headers) _resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) try: response = await _handle_local_mcp_tool(name, arguments) finally: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) # Primary and recommended way to use external MCP servers ######################################################### elif mcp_server: response = await _handle_managed_mcp_tool( server_name=server_name, name=original_tool_name, arguments=arguments, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, host_progress_callback=host_progress_callback, ) # Fall back to local tool registry with original name (legacy support) ######################################################### # Deprecated: Local MCP Server Tool ######################################################### else: # Gate only what can actually dispatch. When the unprefixed name is # not in the registry either, `_handle_local_mcp_tool` below reports # 404 and nothing runs, so demanding a server here would turn every # unknown tool name into a misleading 503. if global_mcp_tool_registry.get_tool(original_tool_name) is not None: # `mcp_server` is None here because the tool name is not in the # tool -> server mapping, but the name still carries a prefix # that the server-level check above compared against the # caller's `allowed_mcp_servers` by exact `name`. So the named # server is in that list and can carry the tool-level checks, # even with the mapping cold. Resolve it from # `allowed_mcp_servers` rather than the registry: the registry # would happily return a server the caller holds no grant for, # and matching anything other than `name` would accept a server # the check never validated. prefix_server: Final = next( (candidate for candidate in allowed_mcp_servers if candidate.name == server_name), None, ) if prefix_server is None: # A non-empty prefix that passed the server-level check # always matches here, so this arm only fires when the # prefix was empty, which is exactly the case that check # skips. Fail closed rather than dispatch with no server to # evaluate a tool ceiling against. raise HTTPException( status_code=503, detail=( f"MCP server for tool '{original_tool_name}' is not available; " "refusing to dispatch without authorization checks. " "Retry once the server is registered." ), ) from litellm.proxy.proxy_server import proxy_logging_obj hook_result = await global_mcp_server_manager.pre_call_tool_check( name=original_tool_name, arguments=arguments, server_name=server_name, user_api_key_auth=user_api_key_auth, proxy_logging_obj=proxy_logging_obj, server=prefix_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args response = await _handle_local_mcp_tool(original_tool_name, arguments) return await _run_post_mcp_call_guardrails( result=response, litellm_logging_obj=litellm_logging_obj, user_api_key_auth=user_api_key_auth, request_data=kwargs, ) async def _run_post_mcp_call_guardrails( result: CallToolResult, litellm_logging_obj: LiteLLMLoggingObj | None, user_api_key_auth: UserAPIKeyAuth | None, request_data: Mapping[str, object], ) -> CallToolResult: """Run ``post_mcp_call`` guardrails over an executed tool result. Lives on ``execute_mcp_tool``'s return path rather than inside ``_fire_mcp_tool_call_logging`` so enforcement never depends on logging being configured, and so every dispatch route gets it: the MCP protocol handler, the REST endpoint, and tool search all funnel through here. A guardrail that rejects the result raises, matching ``pre_mcp_call``. """ from litellm.proxy.proxy_server import proxy_logging_obj if proxy_logging_obj is None: return result return await proxy_logging_obj.post_mcp_call_hook( response=result, request_data=( litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data) ), user_api_key_dict=user_api_key_auth, ) _MCP_CREDENTIAL_REQUEST_FIELDS: Final = frozenset( { "raw_headers", "mcp_auth_header", "mcp_server_auth_headers", "oauth2_headers", "user_api_key_auth", } ) async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, result: CallToolResult, start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, request_data: Mapping[str, object] | None = None, ) -> CallToolResult: """Fire post-call logging for an executed MCP tool call, returning the result to send. The returned result is what the caller must forward to the client: a ``post_mcp_call`` guardrail may rewrite the tool output (e.g. mask sensitive values) or reject it, in which case its exception propagates. Guardrails run before the success/failure logging so the masked text, not the raw one, is what gets logged. A result with ``isError=True`` is logged as a failure (``status="failure"`` payload, so OTel marks the span ERROR) while the HTTP wire behavior stays 200 + ``isError: true`` per the MCP spec. The error check runs after ``async_post_mcp_tool_call_hook`` because guardrails may flip the result to ``isError=True`` in that hook. Raised exceptions never reach here (the ``@client`` wrapper and ``call_mcp_tool``'s except path log those), so this cannot double-log a failure. ``request_data`` may carry credential-bearing fields (the REST path puts ``raw_headers``, ``mcp_auth_header``, ``mcp_server_auth_headers``, and ``oauth2_headers`` at the top level of its data dict), so those are stripped before the dict is handed to ``post_call_failure_hook`` callbacks. """ from litellm.proxy.proxy_server import proxy_logging_obj logging_obj.post_call(original_response=result) await logging_obj.async_post_mcp_tool_call_hook( kwargs=logging_obj.model_call_details, response_obj=result, start_time=start_time, end_time=end_time, ) logging_obj.call_type = CallTypes.call_mcp_tool.value error_message: Final = extract_mcp_tool_result_error_message(result) if error_message is None: await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time) return result logging_obj.has_run_logging(event_type="sync_success") logging_obj.has_run_logging(event_type="async_success") tool_error: Final = MCPToolResultError(error_message) logging_obj.failure_handler(tool_error, "", start_time, end_time) await logging_obj.async_failure_handler(tool_error, "", start_time, end_time) if user_api_key_auth is None: return result if proxy_logging_obj: sanitized_request_data: Final = { key: value for key, value in (request_data or {}).items() if key not in _MCP_CREDENTIAL_REQUEST_FIELDS } await proxy_logging_obj.post_call_failure_hook( request_data=sanitized_request_data, original_exception=tool_error, user_api_key_dict=user_api_key_auth, route="/mcp/call_tool", ) return result @client async def call_mcp_tool( name: str, arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = 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, **kwargs: Any, ) -> CallToolResult: """ Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time: Final = datetime.now() litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) try: if arguments is None: raise HTTPException(status_code=400, detail="Request arguments are required") ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL allowed_mcp_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, ) allowed_mcp_servers: list[MCPServer] = [] for allowed_mcp_server_id in allowed_mcp_server_ids: allowed_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) if allowed_server is not None: # Same request-time oauth2_flow backstop the listing path applies, # so a null-flow M2M-shape row is treated as M2M on tool calls too. allowed_server = MCPServerManager.resolve_oauth2_flow_for_request(allowed_server) allowed_mcp_servers.append(allowed_server) allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) if not allowed_mcp_servers: raise HTTPException( status_code=403, detail="User not allowed to call this tool.", ) # Delegate to execute_mcp_tool for execution response = await execute_mcp_tool( name=name, arguments=arguments, allowed_mcp_servers=allowed_mcp_servers, start_time=start_time, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, **kwargs, ) except MCPUpstreamAuthError: # A client-forwarded pass-through upstream 401 is an expected caller-must-reauth signal, so # re-raise it without post_call_failure_hook, which fires the proxy's llm_exceptions alert. # mcp_server_tool_call then downgrades it to an informational isError result for the # streamable client. Note: this function is @client-decorated, so the decorator's standard # failure logging still records the event (spend log / OTel); only the extra alert sink is # skipped here. raise except Exception as e: traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) from litellm.proxy.proxy_server import proxy_logging_obj # Ordering is load-bearing. ``_ProxyDBLogger.async_post_call_failure_hook``, # reached below, writes the failure spend-log row from this logger's # ``standard_logging_object``, which only exists once the failure handlers # have run. Flush them first or the row lands with # ``guardrail_information=None`` and a guardrail block is never counted. # # Not double-logged: both handlers gate on ``should_run_logging`` and then # mark it, so the ``@client`` wrapper's own post-raise logging no-ops on this # logger, same as ``_fire_mcp_tool_call_logging`` does for ``isError=True``. if litellm_logging_obj is not None: end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from litellm_logging_obj.failure_handler(e, traceback_str, start_time, end_time) await litellm_logging_obj.async_failure_handler(e, traceback_str, start_time, end_time) if proxy_logging_obj and user_api_key_auth: await proxy_logging_obj.post_call_failure_hook( request_data=kwargs, original_exception=e, user_api_key_dict=user_api_key_auth, route="/mcp/call_tool", traceback_str=traceback_str, ) raise if litellm_logging_obj: response = await _fire_mcp_tool_call_logging( logging_obj=litellm_logging_obj, result=response, start_time=start_time, end_time=datetime.now(), user_api_key_auth=user_api_key_auth, request_data=kwargs, ) return response async def mcp_get_prompt( name: str, arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = 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, ) -> GetPromptResult: """ Fetch a specific MCP prompt, handling both prefixed and unprefixed names. """ allowed_mcp_servers: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) if not allowed_mcp_servers: raise HTTPException( status_code=403, detail="User not allowed to get this prompt.", ) # Extract server name from prefixed prompt name original_prompt_name, server_name = split_server_prefix_from_name(name) server: Final = next((s for s in allowed_mcp_servers if s.name == server_name), None) if server is None: raise HTTPException( status_code=403, detail="User not allowed to get this prompt.", ) server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, ) return await global_mcp_server_manager.get_prompt_from_server( server=server, prompt_name=original_prompt_name, arguments=arguments, mcp_auth_header=server_auth_header, extra_headers=extra_headers, raw_headers=raw_headers, ) async def mcp_read_resource( url: AnyUrl, user_api_key_auth: UserAPIKeyAuth | None = 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, ) -> ReadResourceResult: """Read resource contents from upstream MCP servers.""" allowed_mcp_servers: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) if not allowed_mcp_servers: raise HTTPException( status_code=403, detail="User not allowed to read this resource.", ) if len(allowed_mcp_servers) != 1: raise HTTPException( status_code=400, detail=( "Multiple MCP servers configured; read_resource currently supports exactly one allowed server." ), ) server: Final = allowed_mcp_servers[0] server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, ) return await global_mcp_server_manager.read_resource_from_server( server=server, url=url, mcp_auth_header=server_auth_header, extra_headers=extra_headers, raw_headers=raw_headers, ) def _get_standard_logging_mcp_tool_call( name: str, arguments: dict[str, object], server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name(name) namespaced_tool_name: Final = f"{server_name}/{name}" if server_name else name if mcp_server: mcp_info: Final = mcp_server.mcp_info or {} return StandardLoggingMCPToolCall( name=name, arguments=arguments, mcp_server_name=mcp_info.get("server_name"), mcp_server_logo_url=mcp_info.get("logo_url"), namespaced_tool_name=namespaced_tool_name, mcp_session_id=session_id, mcp_auth_mode=mcp_server.auth_type, mcp_server_resource=_redact_mcp_resource_url(mcp_server.url), ) else: return StandardLoggingMCPToolCall( name=name, arguments=arguments, namespaced_tool_name=namespaced_tool_name, mcp_session_id=session_id, ) async def _handle_managed_mcp_tool( server_name: str, name: str, arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: 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, litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import from litellm.proxy.proxy_server import proxy_logging_obj call_tool_result: Final = await global_mcp_server_manager.call_tool( server_name=server_name, name=name, arguments=arguments, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, litellm_logging_obj=litellm_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> CallToolResult: """Execute a local-registry tool and report whether it succeeded. Returns the result rather than bare content because the verdict is part of it: the content alone cannot say whether the handler failed, so callers used to stamp isError=False on every outcome and an upstream rejection was served as tool output. A failure is reported as ``isError=True`` here rather than raised, because the REST surface turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash. ``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to re-authenticate, which both renderers already know how to say. Note: Local tools don't use prefixes, so we use the original name """ import inspect tool: Final = global_mcp_tool_registry.get_tool(name) if not tool: raise HTTPException(status_code=404, detail=f"Tool '{name}' not found") try: if inspect.iscoroutinefunction(tool.handler): result = await tool.handler(**arguments) else: result = tool.handler(**arguments) except MCPUpstreamAuthError: raise except Exception as e: verbose_logger.exception("Error executing local tool %s: %s", name, e) return CallToolResult(content=[TextContent(text=f"Error: {e}", type="text")], isError=True) return CallToolResult(content=[TextContent(text=str(result), type="text")], isError=False) def _get_mcp_servers_in_path(path: str) -> list[str] | None: """ Get the MCP servers from the path """ import re 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/ # 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 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[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) 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: dict[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: Any) -> 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 _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 = json.loads(body) return isinstance(data, dict) and data.get("method") == "initialize" except (json.JSONDecodeError, TypeError): return False async def _read_request_body_for_routing( receive: Receive, ) -> 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 ``wrapped_receive`` in the caller — so an authenticated client cannot force the proxy to buffer an arbitrarily large payload just to make a routing decision. """ consumed_messages: Final[list[Message]] = [] body_chunks: Final[list[bytes]] = [] peeked_bytes = 0 while True: message = await receive() 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 = _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 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) 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[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) 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 = 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 # 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, ) -> UserAPIKeyAuth: """ Restrict a key's MCP permissions to a single toolset. When a request arrives via /toolset/{name}/mcp we override the key's object_permission so that only the toolset's tools are visible. Raises HTTPException(403) if the key has an explicit toolset grant list that does not include toolset_id (i.e. mcp_toolsets is set but empty, or set to a list that omits this toolset). Admin keys always pass. """ from litellm.proxy._types import LiteLLM_ObjectPermissionTable from litellm.proxy.management_endpoints.common_utils import _user_has_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.", ) # Access control: non-admin keys must have this toolset in their grant list. # Use _user_has_admin_view so that PROXY_ADMIN_VIEW_ONLY is also treated as admin. is_admin: Final = _user_has_admin_view(user_api_key_auth) if not is_admin: op: Final = user_api_key_auth.object_permission granted: Final = getattr(op, "mcp_toolsets", None) if op else None # granted=None → key has no explicit toolset grants → deny (same semantics as # fetch_mcp_toolsets which returns [] for non-admin keys with no grants configured). # granted=[] or list without toolset_id → also deny. if granted is None or toolset_id not in granted: raise HTTPException( status_code=403, detail=f"API key does not have access to toolset '{toolset_id}'.", ) tool_permissions = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id]) server_ids: Final = list(tool_permissions.keys()) existing_op: Final = user_api_key_auth.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 user_api_key_auth.model_copy(update={"object_permission": updated_op}) 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, ) -> 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 = 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 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 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: from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 raise_token_exchange_challenge, ) from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 raise_token_exchange_challenge(server, root_path=get_server_root_path()) # token_exchange (OBO) with a subject present: run the exchange here at the transport # edge, so a rejected subject raises the RFC 9728 challenge (and a gateway fault its # public status) instead of the session opening and list_tools masking the failure as # an empty tool list. Gated to single-server routes; the multi-server aggregate keeps # absorbing per-server auth failures so one bad server cannot 401 the whole connect. if ( server and server.auth_type == MCPAuth.oauth2_token_exchange and oauth2_headers and len(mcp_servers or []) == 1 ): await global_mcp_server_manager.preflight_token_exchange( server=server, oauth2_headers=oauth2_headers, user_api_key_auth=user_api_key_auth, ) # 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 _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 _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 _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[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) 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[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) 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) def _is_delegate_upstream_probe_target(server: MCPServer) -> bool: """Whether ``server`` is an interactive delegate-auth server whose client-supplied token should be preflighted upstream. Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed (its stored client credentials drive egress; the caller's bearer is irrelevant). """ return ( server.auth_type == MCPAuth.oauth2 and server.delegate_auth_to_upstream is True and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" ) 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, ) -> None: """Probe pass-through and delegate-auth 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. Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``) are probed with the caller's bare ``Authorization`` bearer. That bearer is only an upstream token (never a LiteLLM key) when admission took the delegate bypass, so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same resolver admission used -- rather than the wider allowed-server prefix/access-group matching. A name that only reaches a delegate server via server_id or an access group would have been admitted as a real LiteLLM key, so probing it would leak that key upstream; requiring the admission-resolver match closes that gap. Without the probe a rejected token is absorbed by the tools/list handler and masked as an empty tool list. Gated to single-server routes so one rejected token cannot 401 a multi-server aggregate connect, matching the OBO preflight gating; the challenge echoes the requested name so aliased routes get the same resource_metadata URL as the tokenless preemptive challenge. """ forwarded_auth: Final = _get_forwarded_auth_from_scope(scope) requested_single_target: Final = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None # The bare Authorization header (no x-litellm-api-key) is a valid upstream token # only when admission classified it as one, i.e. the single requested name resolves # to a delegate server under admission's own resolver. Resolve it the same way here # so a server_id- or access-group-named delegate (which admission would have treated # as a LiteLLM key) is never probed with that key. delegate_server: Final = ( global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip) if requested_single_target else None ) delegate_auth: Final = ( _get_authorization_header_from_scope(scope) if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server) else None ) if not forwarded_auth and not delegate_auth: 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 _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = ( tuple( (srv, forwarded_auth, srv.name) for srv in allowed_servers # Restrict to genuine OAuth pass-through servers (auth_type none + # Authorization in extra_headers). Gateway-managed OAuth2 servers # must not receive the ``resource_metadata=`` challenge emitted # below — they require ``authorization_uri=`` pointing at the # gateway AS metadata. ``is_oauth_passthrough`` already requires # ``auth_type in (None, MCPAuth.none)``, which is mutually # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), # so M2M servers are implicitly excluded here. if srv.is_oauth_passthrough ) if forwarded_auth else () ) # Probe the admission-resolved delegate server only when the caller is actually # authorized for it (present in the IP-filtered allowed set), keyed by server_id. delegate_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = ( tuple( (srv, delegate_auth, requested_single_target) for srv in allowed_servers if delegate_server is not None and srv.server_id == delegate_server.server_id ) if delegate_auth and requested_single_target else () ) probe_targets: Final = passthrough_targets + delegate_targets 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", ) async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 # 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[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) 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 = await _apply_toolset_scope(user_api_key_auth, active_toolset_id) op: Final = user_api_key_auth.object_permission toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set() # 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=_client_ip, allowed_server_ids=toolset_allowed_server_ids, ) # 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, mcp_servers, _client_ip) # 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, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, client_ip=_client_ip, ) if _debug_headers: send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers) # 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) is_initialize = False consumed_messages: list[Message] = [] # 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, and before peeking the request body so the 403 # response sees a pristine ``receive`` channel. 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) body = b"" if scope.get("method") == "POST": consumed_messages, body = await _read_request_body_for_routing(receive) is_initialize = _is_initialize_request(body) 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 # Replay body messages if we consumed them for peeking original_receive: Final = receive if consumed_messages: async def wrapped_receive(): if consumed_messages: return consumed_messages.pop(0) return await original_receive() receive = wrapped_receive # 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 = json.loads(body) if ( isinstance(_peeked, dict) 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, 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" ) session_lock: asyncio.Lock | None = None if use_stateful and session_id and request_method in ("POST", "DELETE") and not is_jsonrpc_response: session_lock = _stateful_session_locks.setdefault(session_id, asyncio.Lock()) 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) 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, ) local_send = 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, ) async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, _client_ip, scoped_server_endpoint=scoped_server_endpoint, ): 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: path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) 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[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) 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 = await _apply_toolset_scope(user_api_key_auth, active_toolset_id) op: Final = user_api_key_auth.object_permission toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set() # 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, ) # 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) 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, ) if not _SESSION_MANAGERS_INITIALIZED: await initialize_session_managers() await asyncio.sleep(0.1) async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, _sse_client_ip, scoped_server_endpoint=scoped_server_endpoint, ): await sse_session_manager.handle_request(scope, receive, send) 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.responses import JSONResponse 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} # 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.mount("/sse", handle_sse_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 _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, ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) 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 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(): try: from mcp.server.lowlevel.server import request_ctx return request_ctx.get().session except (LookupError, ImportError): return 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()