diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e703a3956b9..52117fbcbac 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -16,7 +16,6 @@ from typing import ( TypeVar, Union, ) - import httpx from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client @@ -42,7 +41,6 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl - from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT from litellm.llms.custom_httpx.http_handler import get_ssl_configuration @@ -67,7 +65,6 @@ TSessionResult = TypeVar("TSessionResult") class MCPSigV4Auth(httpx.Auth): """ httpx Auth class that signs each request with AWS SigV4. - This is used for MCP servers that require AWS SigV4 authentication, such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow() for every outgoing request, enabling per-request signature computation. @@ -92,10 +89,8 @@ class MCPSigV4Auth(httpx.Auth): "Missing botocore to use AWS SigV4 authentication. " "Run 'pip install boto3'." ) - self.service_name = aws_service_name or "bedrock-agentcore" self.region_name = aws_region_name or "us-east-1" - # Note: os.environ/ prefixed values are already resolved by # ProxyConfig._check_for_os_environ_vars() at config load time. # Values arrive here as plain strings. @@ -143,20 +138,17 @@ class MCPSigV4Auth(httpx.Auth): session_name = ( aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" ) - sts_kwargs: dict = {"region_name": aws_region_name} if aws_access_key_id and aws_secret_access_key: sts_kwargs["aws_access_key_id"] = aws_access_key_id sts_kwargs["aws_secret_access_key"] = aws_secret_access_key if aws_session_token: sts_kwargs["aws_session_token"] = aws_session_token - sts_client = boto3.client("sts", **sts_kwargs) sts_response = sts_client.assume_role( RoleArn=aws_role_name, RoleSessionName=session_name, ) - sts_creds = sts_response["Credentials"] return Credentials( access_key=sts_creds["AccessKeyId"], @@ -178,17 +170,14 @@ class MCPSigV4Auth(httpx.Auth): data=request.content, headers=dict(request.headers), ) - # Sign the request — SigV4Auth.add_auth() adds Authorization, # X-Amz-Date, and X-Amz-Security-Token (if session token present). # Host header is derived automatically from the URL. sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name) sigv4.add_auth(aws_request) - # Copy SigV4 headers back to the httpx request for header_name, header_value in aws_request.headers.items(): request.headers[header_name] = header_value - yield request @@ -198,6 +187,8 @@ class MCPClient: SSE and HTTP transports Authentication via Bearer token, Basic Auth, or API Key Tool calling with error handling and result parsing + Sampling callbacks for upstream server LLM requests + Elicitation callbacks for upstream server user-input requests """ def __init__( @@ -211,6 +202,9 @@ class MCPClient: extra_headers: Optional[Dict[str, str]] = None, ssl_verify: Optional[VerifyTypes] = None, aws_auth: Optional[httpx.Auth] = None, + sampling_callback: Optional[Callable] = None, + elicitation_callback: Optional[Callable] = None, + logging_callback: Optional[Callable] = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type @@ -222,6 +216,9 @@ class MCPClient: self.ssl_verify: Optional[VerifyTypes] = ssl_verify self._aws_auth: Optional[httpx.Auth] = aws_auth self._last_initialize_instructions: Optional[str] = None + self._sampling_callback: Optional[Callable] = sampling_callback + self._elicitation_callback: Optional[Callable] = elicitation_callback + self._logging_callback: Optional[Callable] = logging_callback # handle the basic auth value if provided if auth_value: self.update_auth_value(auth_value) @@ -231,23 +228,20 @@ class MCPClient: ) -> Tuple[Any, Optional[httpx.AsyncClient]]: """ Create the appropriate transport context based on transport type. - Returns: Tuple of (transport_context, http_client). http_client is only set for HTTP transport and needs cleanup. """ http_client: Optional[httpx.AsyncClient] = None - if self.transport_type == MCPTransport.stdio: if not self.stdio_config: raise ValueError("stdio_config is required for stdio transport") server_params = StdioServerParameters( command=self.stdio_config.get("command", ""), args=self.stdio_config.get("args", []), - env=self.stdio_config.get("env", {}), + env=self.stdio_config.get("env", None), ) return stdio_client(server_params), None - if self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() @@ -260,14 +254,12 @@ class MCPClient: ), None, ) - # HTTP transport (default) if streamable_http_client is None: raise ImportError( "streamable_http_client is not available. " "Please install mcp with HTTP support." ) - headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) @@ -288,13 +280,24 @@ class MCPClient: ) -> TSessionResult: """ Execute an operation within a transport and session context. - Handles entering/exiting contexts and running the operation. + Passes sampling/elicitation/logging callbacks to the ClientSession + so that upstream MCP servers can request LLM inference (sampling), + user input (elicitation), or send log messages. """ transport = await transport_ctx.__aenter__() try: read_stream, write_stream = transport[0], transport[1] session_ctx = ClientSession(read_stream, write_stream) + # Build session kwargs with optional callbacks + session_kwargs: Dict[str, Any] = {} + if self._sampling_callback is not None: + session_kwargs["sampling_callback"] = self._sampling_callback + if self._elicitation_callback is not None: + session_kwargs["elicitation_callback"] = self._elicitation_callback + if self._logging_callback is not None: + session_kwargs["logging_callback"] = self._logging_callback + session_ctx = ClientSession(read_stream, write_stream, **session_kwargs) session = await session_ctx.__aenter__() try: init_result = await session.initialize() @@ -351,7 +354,6 @@ class MCPClient: def _get_auth_headers(self) -> dict: """Generate authentication headers based on auth type.""" headers = {} - if self._mcp_auth_value: if isinstance(self._mcp_auth_value, str): if self.auth_type == MCPAuth.bearer_token: @@ -371,17 +373,14 @@ class MCPClient: # Note: aws_sigv4 auth is not handled here — SigV4 requires per-request # signing (including the body hash), so it uses httpx.Auth flow instead # of static headers. See MCPSigV4Auth and _create_httpx_client_factory(). - # update the headers with the extra headers if self.extra_headers: headers.update(self.extra_headers) - return headers def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]: """ Create a custom httpx client factory that uses LiteLLM's SSL configuration. - This factory follows the same CA bundle path logic as http_handler.py: 1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle) 2. Check SSL_VERIFY environment variable @@ -398,17 +397,14 @@ class MCPClient: """Create an httpx.AsyncClient with LiteLLM's SSL configuration.""" # Get unified SSL configuration using the same logic as http_handler.py ssl_config = get_ssl_configuration(self.ssl_verify) - verbose_logger.debug( f"MCP client using SSL configuration: {type(ssl_config).__name__}" ) - # Use SigV4 auth if configured and no explicit auth provided. # The MCP SDK's sse_client and streamable_http_client call this # factory without passing auth=, so self._aws_auth is used. # For non-SigV4 clients, self._aws_auth is None — no behavior change. effective_auth = auth if auth is not None else self._aws_auth - return httpx.AsyncClient( headers=headers, timeout=timeout, @@ -448,14 +444,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_tools - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -479,7 +473,6 @@ class MCPClient: f"MCP Tool '{call_tool_request_params.name}' progress: " f"{progress}/{total} ({percentage:.0f}%) - {message or ''}" ) - # Forward to Host if callback provided if host_progress_callback: try: @@ -509,7 +502,6 @@ class MCPClient: error_trace = traceback.format_exc() verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}") - # Log detailed error information error_type = type(e).__name__ verbose_logger.error( @@ -520,14 +512,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream - " "the MCP server may have crashed, disconnected, or timed out." ) - # Return a default error result instead of raising return MCPCallToolResult( content=[ @@ -565,14 +555,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_tools - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -605,7 +593,6 @@ class MCPClient: error_trace = traceback.format_exc() verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}") - # Log detailed error information error_type = type(e).__name__ verbose_logger.error( @@ -616,14 +603,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during get_prompt - " "the MCP server may have crashed, disconnected, or timed out." ) - raise async def list_resources(self) -> list[Resource]: @@ -655,14 +640,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_resources - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -697,14 +680,12 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during list_resource_templates - " "the MCP server may have crashed, disconnected, or timed out" ) - # Return empty list instead of raising to allow graceful degradation return [] @@ -730,7 +711,6 @@ class MCPClient: error_trace = traceback.format_exc() verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}") - # Log detailed error information error_type = type(e).__name__ verbose_logger.error( @@ -741,12 +721,10 @@ class MCPClient: f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) - # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream during read_resource - " "the MCP server may have crashed, disconnected, or timed out." ) - raise diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py new file mode 100644 index 00000000000..323e2fff02c --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py @@ -0,0 +1,161 @@ +""" +MCP Elicitation Handler +Handles `elicitation/create` requests from upstream MCP servers by either: +1. Relaying them to the connected downstream MCP client (if it supports elicitation) +2. Returning a decline/error response (if no downstream client or unsupported) +Supports both Form mode (structured data collection) and URL mode (external URL +navigation for sensitive interactions like OAuth). +MCP Spec Reference: + https://modelcontextprotocol.io/specification/2025-11-25/client/elicitation +""" + +from typing import Any, Optional, Union +from litellm._logging import verbose_logger + +# Guard imports that require the mcp package +try: + from mcp.types import ( + ElicitRequestFormParams, + ElicitRequestParams, + ElicitRequestURLParams, + ElicitResult, + ErrorData, + ) + + MCP_ELICITATION_AVAILABLE = True +except ImportError: + MCP_ELICITATION_AVAILABLE = False + + +async def handle_elicitation_request( + context: Any, + params: "ElicitRequestParams", + downstream_session: Optional[Any] = None, + downstream_capabilities: Optional[Any] = None, +) -> Union["ElicitResult", "ErrorData"]: + """ + Handle an MCP elicitation/create request from an upstream MCP server. + In Gateway mode (Mode A), we relay the elicitation request to the + connected downstream client if they declared elicitation capabilities. + In Tool Bridge mode (Mode B), there's no persistent downstream MCP + client, so we return a decline response. + Args: + context: MCP RequestContext from the upstream server connection. + params: The ElicitRequestParams (either form or URL mode). + downstream_session: The ServerSession to the downstream client, + if available (for relaying). + downstream_capabilities: The downstream client's declared + capabilities, used to check elicitation support. + Returns: + ElicitResult with the user's response, or ErrorData on failure. + """ + if not MCP_ELICITATION_AVAILABLE: + return ErrorData( + code=-1, + message="MCP elicitation is not available (mcp package not installed)", + ) + try: + mode = getattr(params, "mode", "form") + verbose_logger.info( + "MCP elicitation: received request mode=%s, message=%s", + mode, + getattr(params, "message", ""), + ) + # Check if we have a downstream session to relay to + if downstream_session is not None: + return await _relay_elicitation_to_downstream( + params=params, + downstream_session=downstream_session, + downstream_capabilities=downstream_capabilities, + ) + # No downstream session — we're in Tool Bridge mode + # or the client doesn't support elicitation + verbose_logger.info( + "MCP elicitation: no downstream session available, declining" + ) + return ElicitResult( + action="decline", + ) + except Exception as e: + verbose_logger.exception("MCP elicitation handler failed: %s", e) + return ErrorData( + code=-1, + message=f"Elicitation failed: {str(e)}", + ) + + +async def _relay_elicitation_to_downstream( + params: "ElicitRequestParams", + downstream_session: Any, + downstream_capabilities: Optional[Any] = None, +) -> Union["ElicitResult", "ErrorData"]: + """ + Relay an elicitation request to the downstream MCP client. + Uses the ServerSession's elicit_form() or elicit_url() methods to + send the elicitation request back to the connected client. + Args: + params: The elicitation request parameters. + downstream_session: The ServerSession connected to the downstream client. + downstream_capabilities: Client capabilities to check support. + Returns: + ElicitResult from the downstream client. + """ + mode = getattr(params, "mode", "form") + # Check if the downstream client supports the requested mode + if downstream_capabilities is not None: + elicit_caps = getattr(downstream_capabilities, "elicitation", None) + if elicit_caps is None: + verbose_logger.info( + "MCP elicitation: downstream client does not support elicitation" + ) + return ElicitResult(action="decline") + if mode == "url": + url_cap = getattr(elicit_caps, "url", None) + if url_cap is None: + verbose_logger.info( + "MCP elicitation: downstream client does not support URL mode" + ) + return ElicitResult(action="decline") + if mode == "form": + form_cap = getattr(elicit_caps, "form", None) + if form_cap is None: + verbose_logger.info( + "MCP elicitation: downstream client does not support form mode" + ) + return ElicitResult(action="decline") + try: + if mode == "url" and isinstance(params, ElicitRequestURLParams): + # URL mode: relay URL to client for external navigation + verbose_logger.info( + "MCP elicitation: relaying URL mode to downstream, url=%s", + getattr(params, "url", ""), + ) + result = await downstream_session.elicit_url( + message=params.message, + url=params.url, + elicitation_id=getattr(params, "elicitationId", None), + ) + elif isinstance(params, ElicitRequestFormParams): + # Form mode: relay structured form to client + verbose_logger.info("MCP elicitation: relaying form mode to downstream") + result = await downstream_session.elicit_form( + message=params.message, + requestedSchema=getattr(params, "requestedSchema", None), + ) + else: + # Fallback for generic ElicitRequestParams + verbose_logger.info( + "MCP elicitation: relaying generic elicitation to downstream" + ) + result = await downstream_session.elicit( + message=getattr(params, "message", ""), + ) + verbose_logger.info( + "MCP elicitation: downstream responded with action=%s", + getattr(result, "action", "unknown"), + ) + return result + except Exception as e: + verbose_logger.warning("MCP elicitation: failed to relay to downstream: %s", e) + # If relay fails, decline gracefully + return ElicitResult(action="decline") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 251f271903b..618ac4dba49 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -156,6 +156,65 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: return data +def _create_sampling_callback(): + """ + Create a sampling callback for MCP ClientSession. + Returns a callable that handles sampling/createMessage requests from + upstream MCP servers by routing them through litellm.acompletion(). + """ + + async def _sampling_callback(context, params): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + import litellm + from litellm.proxy._experimental.mcp_server.server import ( + get_active_auth_context, + ) + + return await handle_sampling_create_message( + context=context, + params=params, + default_model=getattr(litellm, "default_mcp_sampling_model", None), + user_api_key_auth=get_active_auth_context(), + ) + + return _sampling_callback + + +def _create_elicitation_callback(): + """ + Create an elicitation callback for MCP ClientSession. + Returns a callable that handles elicitation/create requests from + upstream MCP servers. In gateway mode, this relays to the downstream + client; in tool bridge mode, it returns a decline response. + """ + + async def _elicitation_callback(context, params): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session + + # In Gateway mode, we relay the elicitation request to the downstream client + # that triggered the current operation. + downstream_session = get_active_mcp_session() + downstream_capabilities = ( + getattr(downstream_session, "capabilities", None) + if downstream_session + else None + ) + + return await handle_elicitation_request( + context=context, + params=params, + downstream_session=downstream_session, + downstream_capabilities=downstream_capabilities, + ) + + return _elicitation_callback + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -1125,23 +1184,34 @@ class MCPServerManager: transport = server.transport or MCPTransport.sse + # Create sampling and elicitation callbacks for this client + sampling_cb = _create_sampling_callback() + elicitation_cb = _create_elicitation_callback() + # Handle stdio transport if transport == MCPTransport.stdio: resolved_env = ( - stdio_env if stdio_env is not None else dict(server.env or {}) + stdio_env + if stdio_env is not None + else (dict(server.env) if server.env else None) ) # Ensure npm-based STDIO MCP servers have a writable cache dir. # In containers the default (~/.npm or /app/.npm) may not exist # or be read-only, causing npx to fail with ENOENT. - if "NPM_CONFIG_CACHE" not in resolved_env: + if resolved_env is not None and "NPM_CONFIG_CACHE" not in resolved_env: resolved_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR # Defense-in-depth: block commands not in the allowlist. # The Pydantic validator blocks new servers; this catches legacy # config/DB records predating the allowlist. if server.command: base_command = os.path.basename(server.command) - if base_command not in MCP_STDIO_ALLOWED_COMMANDS: + # Strip .exe/.cmd/.bat suffix for Windows compatibility + base_command_no_ext = os.path.splitext(base_command)[0] + if ( + base_command not in MCP_STDIO_ALLOWED_COMMANDS + and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS + ): raise HTTPException( status_code=403, detail=f"MCP stdio command '{server.command}' is not in the allowlist ({sorted(MCP_STDIO_ALLOWED_COMMANDS)}). " @@ -1164,6 +1234,8 @@ class MCPServerManager: timeout=MCP_CLIENT_TIMEOUT, stdio_config=stdio_config, extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) else: # For HTTP/SSE transports @@ -1190,6 +1262,8 @@ class MCPServerManager: timeout=MCP_CLIENT_TIMEOUT, extra_headers=extra_headers, aws_auth=aws_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) async def _get_tools_from_server( diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py new file mode 100644 index 00000000000..3150ebf6e29 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -0,0 +1,525 @@ +""" +MCP Sampling Handler +Handles `sampling/createMessage` requests from upstream MCP servers by +routing them through LiteLLM's internal completion infrastructure. +This allows MCP servers to perform agentic reasoning (e.g., multi-step +tool calling, chain-of-thought) without needing their own LLM API keys — +LiteLLM acts as the LLM provider using its existing 100+ provider support, +cost tracking, rate limiting, and model routing. +MCP Spec Reference: + https://modelcontextprotocol.io/specification/2025-11-25/client/sampling +""" + +from typing import Any, Dict, List, Optional, Union +from litellm._logging import verbose_logger + +# Guard imports that require the mcp package +try: + from mcp.types import ( + CreateMessageRequestParams, + CreateMessageResult, + CreateMessageResultWithTools, + ErrorData, + ModelPreferences, + SamplingMessage, + TextContent, + Tool, + ToolChoice, + ToolUseContent, + ) + + MCP_SAMPLING_AVAILABLE = True +except ImportError: + MCP_SAMPLING_AVAILABLE = False +# Maximum number of sampling iterations to prevent infinite loops +DEFAULT_MAX_SAMPLING_ITERATIONS = 10 + + +def _resolve_model_from_preferences( + model_preferences: Optional["ModelPreferences"], + default_model: Optional[str] = None, +) -> str: + """ + Resolve an LLM model name from MCP ModelPreferences. + Strategy: + 1. Check hints for substring matches against known model names. + 2. Fall back to priority-based selection (cost/speed/intelligence). + 3. Fall back to the configured default model. + Args: + model_preferences: MCP ModelPreferences with hints and priorities. + default_model: Fallback model if no hint matches. + Returns: + A model string suitable for litellm.acompletion(). + """ + import litellm + + # Build list of available model names from proxy Router or litellm.model_list + available_model_names: list = [] + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is not None: + available_model_names = llm_router.get_model_names() + except Exception: + pass + if not available_model_names and litellm.model_list: + for entry in litellm.model_list: + if isinstance(entry, dict): + name = entry.get("model_name") + if name: + available_model_names.append(name) + elif isinstance(entry, str): + available_model_names.append(entry) + if model_preferences and model_preferences.hints: + for hint in model_preferences.hints: + hint_name = getattr(hint, "name", None) + if not hint_name: + continue + # Try direct match first + if hint_name in available_model_names: + return hint_name + # Try substring match against known models + for model_name in available_model_names: + if hint_name.lower() in model_name.lower(): + return model_name + # Use default model from caller + if default_model: + return default_model + # Fall back to first available model + if available_model_names: + return available_model_names[0] + # Last resort + return "gpt-4o-mini" + + +def _convert_mcp_content_to_openai( + content: Any, +) -> Union[str, Dict[str, Any], List[Dict[str, Any]]]: + """ + Convert MCP SamplingMessage content to OpenAI message content format. + Handles: + - TextContent → string or {"type": "text", "text": ...} + - ImageContent → {"type": "image_url", "image_url": {"url": "data:..."}} + - AudioContent → {"type": "input_audio", "input_audio": {...}} + - ToolUseContent → function call representation + - ToolResultContent → tool result representation + - List of mixed content → list of content parts + """ + if isinstance(content, list): + parts = [] + for item in content: + converted = _convert_single_content(item) + if isinstance(converted, list): + parts.extend(converted) + else: + parts.append(converted) + return parts + return _convert_single_content(content) + + +def _convert_single_content(content: Any) -> Union[str, Dict[str, Any]]: + """Convert a single MCP content item to OpenAI format.""" + content_type = getattr(content, "type", None) + if content_type == "text": + return {"type": "text", "text": content.text} + elif content_type == "image": + data = getattr(content, "data", "") + mime_type = getattr(content, "mimeType", "image/png") + return { + "type": "image_url", + "image_url": {"url": f"data:{mime_type};base64,{data}"}, + } + elif content_type == "audio": + data = getattr(content, "data", "") + mime_type = getattr(content, "mimeType", "audio/wav") + # Map MIME type to OpenAI audio format + format_map = { + "audio/wav": "wav", + "audio/mp3": "mp3", + "audio/mpeg": "mp3", + "audio/flac": "flac", + "audio/ogg": "ogg", + } + audio_format = format_map.get(mime_type, "wav") + return { + "type": "input_audio", + "input_audio": {"data": data, "format": audio_format}, + } + elif content_type == "tool_use": + # ToolUseContent → represents the assistant calling a tool + return { + "type": "text", + "text": f"[Tool call: {getattr(content, 'name', 'unknown')}]", + } + elif content_type == "tool_result": + # ToolResultContent → represents tool results + tool_content = getattr(content, "content", []) + if isinstance(tool_content, list) and tool_content: + texts = [ + getattr(c, "text", str(c)) + for c in tool_content + if getattr(c, "type", None) == "text" + ] + return {"type": "text", "text": "\n".join(texts) if texts else ""} + return {"type": "text", "text": str(tool_content)} + # Fallback: treat as text + return {"type": "text", "text": str(content)} + + +def _convert_mcp_messages_to_openai( + messages: List["SamplingMessage"], + system_prompt: Optional[str] = None, +) -> List[Dict[str, Any]]: + """ + Convert MCP SamplingMessage list to OpenAI messages format. + MCP messages use: + - role: "user" | "assistant" + - content: TextContent | ImageContent | AudioContent | ToolUseContent + | ToolResultContent | list[...] + OpenAI messages use: + - role: "system" | "user" | "assistant" | "tool" + - content: str | list[content_part] + """ + openai_messages: List[Dict[str, Any]] = [] + # Add system prompt if provided + if system_prompt: + openai_messages.append({"role": "system", "content": system_prompt}) + for msg in messages: + role = msg.role + content = msg.content + # Handle tool use content from assistant + if role == "assistant" and _has_tool_use(content): + tool_calls = _extract_tool_calls(content) + if tool_calls: + openai_msg: Dict[str, Any] = { + "role": "assistant", + "tool_calls": tool_calls, + } + # Also include any text content alongside tool calls + text_parts = _extract_text_parts(content) + if text_parts: + openai_msg["content"] = text_parts + openai_messages.append(openai_msg) + continue + # Handle tool result content from user + if role == "user" and _has_tool_result(content): + tool_results = _extract_tool_results(content) + for tool_result in tool_results: + openai_messages.append(tool_result) + continue + # Standard text/image/audio message + converted = _convert_mcp_content_to_openai(content) + if isinstance(converted, str): + openai_messages.append({"role": role, "content": converted}) + elif isinstance(converted, dict): + openai_messages.append({"role": role, "content": [converted]}) + elif isinstance(converted, list): + openai_messages.append({"role": role, "content": converted}) + return openai_messages + + +def _has_tool_use(content: Any) -> bool: + """Check if content contains ToolUseContent.""" + if isinstance(content, list): + return any(getattr(c, "type", None) == "tool_use" for c in content) + return getattr(content, "type", None) == "tool_use" + + +def _has_tool_result(content: Any) -> bool: + """Check if content contains ToolResultContent.""" + if isinstance(content, list): + return any(getattr(c, "type", None) == "tool_result" for c in content) + return getattr(content, "type", None) == "tool_result" + + +def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]: + """Extract OpenAI-format tool_calls from MCP ToolUseContent.""" + import json + + items = content if isinstance(content, list) else [content] + tool_calls = [] + for item in items: + if getattr(item, "type", None) == "tool_use": + tool_calls.append( + { + "id": getattr(item, "id", f"call_{id(item)}"), + "type": "function", + "function": { + "name": getattr(item, "name", ""), + "arguments": json.dumps( + getattr(item, "input", {}), default=str + ), + }, + } + ) + return tool_calls + + +def _extract_text_parts(content: Any) -> Optional[str]: + """Extract text parts from mixed content.""" + items = content if isinstance(content, list) else [content] + texts = [] + for item in items: + if getattr(item, "type", None) == "text": + texts.append(getattr(item, "text", "")) + return "\n".join(texts) if texts else None + + +def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: + """Extract OpenAI-format tool messages from MCP ToolResultContent.""" + items = content if isinstance(content, list) else [content] + results = [] + for item in items: + if getattr(item, "type", None) == "tool_result": + tool_use_id = getattr(item, "toolUseId", "") + # Extract text from nested content + nested_content = getattr(item, "content", []) + if isinstance(nested_content, list): + text_parts = [ + getattr(c, "text", str(c)) + for c in nested_content + if getattr(c, "type", None) == "text" + ] + result_text = "\n".join(text_parts) if text_parts else "" + else: + result_text = str(nested_content) + results.append( + { + "role": "tool", + "tool_call_id": tool_use_id, + "content": result_text, + } + ) + return results + + +def _convert_mcp_tools_to_openai( + tools: Optional[List["Tool"]], +) -> Optional[List[Dict[str, Any]]]: + """ + Convert MCP Tool definitions to OpenAI function calling format. + MCP Tool: {name, description, inputSchema} + OpenAI Tool: {type: "function", function: {name, description, parameters}} + """ + if not tools: + return None + openai_tools = [] + for tool in tools: + openai_tool = { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description or "", + "parameters": tool.inputSchema + or { + "type": "object", + "properties": {}, + }, + }, + } + openai_tools.append(openai_tool) + return openai_tools + + +def _convert_mcp_tool_choice_to_openai( + tool_choice: Optional["ToolChoice"], +) -> Optional[Union[str, Dict[str, Any]]]: + """ + Convert MCP ToolChoice to OpenAI tool_choice format. + MCP: {mode: "auto"} | {mode: "required"} | {mode: "none"} + OpenAI: "auto" | "required" | "none" + """ + if not tool_choice: + return None + mode = getattr(tool_choice, "mode", "auto") + if mode == "auto": + return "auto" + elif mode == "required": + return "required" + elif mode == "none": + return "none" + return "auto" + + +def _convert_openai_response_to_mcp_result( + response: Any, + model_name: str, +) -> Union["CreateMessageResult", "CreateMessageResultWithTools"]: + """ + Convert a litellm completion response to MCP CreateMessageResult. + Args: + response: The litellm ModelResponse. + model_name: The model that was used. + Returns: + MCP CreateMessageResult or CreateMessageResultWithTools. + """ + choice = response.choices[0] + message = choice.message + # Determine stop reason + finish_reason = getattr(choice, "finish_reason", "stop") + if finish_reason == "tool_calls": + stop_reason = "toolUse" + elif finish_reason == "length": + stop_reason = "maxTokens" + else: + stop_reason = "endTurn" + actual_model = getattr(response, "model", model_name) or model_name + # Check if response has tool calls + tool_calls = getattr(message, "tool_calls", None) + if tool_calls: + # Build ToolUseContent items + content_parts = [] + # Include text content if present + if message.content: + content_parts.append(TextContent(type="text", text=message.content)) + # Convert tool calls to MCP ToolUseContent + for tc in tool_calls: + import json + + tool_input = tc.function.arguments + if isinstance(tool_input, str): + try: + tool_input = json.loads(tool_input) + except (json.JSONDecodeError, TypeError): + tool_input = {"raw": tool_input} + content_parts.append( + ToolUseContent( + type="tool_use", + id=tc.id, + name=tc.function.name, + input=tool_input, + ) + ) + return CreateMessageResultWithTools( + role="assistant", + content=content_parts, + model=actual_model, + stopReason=stop_reason, + ) + # Simple text response + text = message.content or "" + return CreateMessageResult( + role="assistant", + content=TextContent(type="text", text=text), + model=actual_model, + stopReason=stop_reason, + ) + + +async def handle_sampling_create_message( + context: Any, + params: "CreateMessageRequestParams", + default_model: Optional[str] = None, + user_api_key_auth: Optional[Any] = None, + max_iterations: int = DEFAULT_MAX_SAMPLING_ITERATIONS, +) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: + """ + Handle an MCP sampling/createMessage request by routing through LiteLLM. + This is the main entry point called by the MCP client session when an + upstream MCP server requests LLM inference. + Args: + context: MCP RequestContext (contains session info). + params: The CreateMessageRequestParams from the MCP server. + default_model: Default model to use if no preferences match. + user_api_key_auth: Auth context for the requesting user. + max_iterations: Maximum tool-calling iterations to prevent loops. + Returns: + CreateMessageResult with the LLM's response, or ErrorData on failure. + """ + if not MCP_SAMPLING_AVAILABLE: + return ErrorData( + code=-1, + message="MCP sampling is not available (mcp package not installed)", + ) + try: + import litellm + + # 1. Resolve model from preferences + model = _resolve_model_from_preferences( + model_preferences=params.modelPreferences, + default_model=default_model, + ) + verbose_logger.info( + "MCP sampling: resolved model=%s from preferences=%s", + model, + params.modelPreferences, + ) + # 2. Convert MCP messages to OpenAI format + openai_messages = _convert_mcp_messages_to_openai( + messages=params.messages, + system_prompt=params.systemPrompt, + ) + # 3. Build completion kwargs + completion_kwargs: Dict[str, Any] = { + "model": model, + "messages": openai_messages, + "max_tokens": params.maxTokens, + } + if params.temperature is not None: + completion_kwargs["temperature"] = params.temperature + if params.stopSequences: + completion_kwargs["stop"] = params.stopSequences + # 4. Convert tools and tool_choice if provided + openai_tools = _convert_mcp_tools_to_openai(params.tools) + if openai_tools: + completion_kwargs["tools"] = openai_tools + openai_tool_choice = _convert_mcp_tool_choice_to_openai(params.toolChoice) + if openai_tool_choice is not None: + completion_kwargs["tool_choice"] = openai_tool_choice + # 5. Add metadata for tracking + if params.metadata: + completion_kwargs["metadata"] = params.metadata + + # 6. Inject auth context for cost tracking + if user_api_key_auth: + completion_kwargs["user"] = getattr(user_api_key_auth, "user_id", None) + + # Pass user_api_key_dict directly so proxy hooks can attribute the cost + # litellm_pre_call_utils usually checks for this in kwargs or metadata + if "metadata" not in completion_kwargs: + completion_kwargs["metadata"] = {} + + api_key = getattr(user_api_key_auth, "api_key", None) + if api_key: + completion_kwargs["metadata"]["user_api_key"] = api_key + + team_id = getattr(user_api_key_auth, "team_id", None) + if team_id: + completion_kwargs["metadata"]["user_api_key_team_id"] = team_id + + verbose_logger.debug( + "MCP sampling: calling litellm.acompletion with model=%s, " + "num_messages=%d, has_tools=%s", + model, + len(openai_messages), + bool(openai_tools), + ) + + # 7. Call LiteLLM + # Use proxy's llm_router if available, else fallback to litellm.acompletion + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is not None: + response = await llm_router.acompletion(**completion_kwargs) + else: + response = await litellm.acompletion(**completion_kwargs) + except Exception: + response = await litellm.acompletion(**completion_kwargs) + # 7. Convert response to MCP format + result = _convert_openai_response_to_mcp_result( + response=response, + model_name=model, + ) + verbose_logger.info( + "MCP sampling: completed successfully, model=%s, stopReason=%s", + getattr(result, "model", "unknown"), + getattr(result, "stopReason", "unknown"), + ) + return result + except Exception as e: + verbose_logger.exception("MCP sampling handler failed: %s", e) + return ErrorData( + code=-1, + message=f"Sampling failed: {str(e)}", + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c2e998f01e5..b31103446a5 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3,9 +3,9 @@ LiteLLM MCP Server Routes """ # pyright: reportInvalidTypeForm=false, reportArgumentType=false, reportOptionalCall=false - import asyncio import contextlib +import contextvars import time import types import traceback @@ -22,13 +22,11 @@ from typing import ( Union, cast, ) - 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 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 @@ -72,7 +70,6 @@ _BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth 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. """ @@ -102,6 +99,12 @@ try: GetPromptResult, ResourceTemplate, TextResourceContents, + Tool, + ) + from mcp.server.session import ServerSession as _McpServerSession + + active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = ( + contextvars.ContextVar("active_mcp_session", default=None) ) except ImportError as e: verbose_logger.debug(f"MCP module not found: {e}") @@ -117,13 +120,9 @@ except ImportError as e: ResourceTemplate = None # type: ignore Server = None # type: ignore TextResourceContents = None # type: ignore - - # Global variables to track initialization _SESSION_MANAGERS_INITIALIZED = False _INITIALIZATION_LOCK = asyncio.Lock() - - if MCP_AVAILABLE: from mcp.server import Server from mcp.server.lowlevel.server import NotificationOptions @@ -147,7 +146,6 @@ if MCP_AVAILABLE: TextContent, ) from mcp.types import Tool as MCPTool - from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) @@ -157,7 +155,6 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, ) - from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -224,6 +221,9 @@ if MCP_AVAILABLE: ######################################################## ############ Initialize the MCP Server ################# ######################################################## + from fastapi import APIRouter + + router = APIRouter() server: Server = Server( name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, @@ -231,16 +231,21 @@ if MCP_AVAILABLE: server.create_initialization_options = types.MethodType( # type: ignore[method-assign] _gateway_create_initialization_options, server ) - sse: SseServerTransport = SseServerTransport("/mcp/sse/messages") + # SSE Server Transport — uses the official MCP SDK class. + # The endpoint "/messages" is the path (relative to where the SSE GET + # endpoint is mounted) that the SDK will tell clients to POST messages to. + # Since the mcp_app is mounted at /mcp in the main proxy, the full + # client-visible POST path becomes /mcp/messages. + from mcp.server.sse import SseServerTransport as _McpSseServerTransport - # Create session managers + sse = _McpSseServerTransport("/messages") + # Create session managers (StreamableHTTP — stateless by default) session_manager = StreamableHTTPSessionManager( app=server, event_store=None, json_response=False, # enables SSE streaming stateless=True, ) - # Create SSE session manager sse_session_manager = StreamableHTTPSessionManager( app=server, @@ -248,7 +253,6 @@ if MCP_AVAILABLE: json_response=False, # Use SSE responses for this endpoint stateless=True, ) - # Context managers for proper lifecycle management _session_manager_cm = None _sse_session_manager_cm = None @@ -256,34 +260,27 @@ if MCP_AVAILABLE: async def initialize_session_managers(): """Initialize the session managers. Can be called from main app lifespan.""" global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm - # 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.run() _sse_session_manager_cm = sse_session_manager.run() - # Enter the context managers await _session_manager_cm.__aenter__() await _sse_session_manager_cm.__aenter__() - _SESSION_MANAGERS_INITIALIZED = True verbose_logger.info( - "MCP Server started with StreamableHTTP and SSE session managers!" + "MCP Server started with StreamableHTTP session manager and SSE transport!" ) async def shutdown_session_managers(): """Shutdown the session managers.""" global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _sse_session_manager_cm - if _SESSION_MANAGERS_INITIALIZED: verbose_logger.info("Shutting down MCP session managers...") - try: if _session_manager_cm: await _session_manager_cm.__aexit__(None, None, None) @@ -291,7 +288,6 @@ if MCP_AVAILABLE: await _sse_session_manager_cm.__aexit__(None, None, None) except Exception as e: verbose_logger.exception(f"Error during session manager shutdown: {e}") - _session_manager_cm = None _sse_session_manager_cm = None _SESSION_MANAGERS_INITIALIZED = False @@ -308,12 +304,18 @@ if MCP_AVAILABLE: ######################################################## ############### MCP Server Routes ####################### ######################################################## - @server.list_tools() - async def list_tools() -> List[MCPTool]: + async def handle_list_tools() -> List[Tool]: """ - List all available tools + List all available tools. + Also captures the active session for propagation to callbacks. """ + from mcp.server.lowlevel.server import request_ctx + + req_ctx = request_ctx.get(None) + if req_ctx: + active_mcp_session_var.set(req_ctx.session) + try: # Get user authentication from context variable ( @@ -324,7 +326,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}" ) @@ -362,19 +364,15 @@ if MCP_AVAILABLE: ) -> 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 fastapi import Request - 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 @@ -388,10 +386,9 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() - + ) = await get_or_extract_auth_context() verbose_logger.debug( - f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" + f"MCP mcp_server_tool_call - user_api_key_auth={user_api_key_auth}, user_role={getattr(user_api_key_auth, 'user_role', 'N/A')}" ) host_progress_callback = None try: @@ -431,7 +428,6 @@ if MCP_AVAILABLE: if chain_id: body_data["litellm_trace_id"] = chain_id body_data["litellm_session_id"] = chain_id - request = Request( scope={ "type": "http", @@ -449,7 +445,6 @@ if MCP_AVAILABLE: ) else: data = body_data - response = await call_mcp_tool( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -493,7 +488,6 @@ if MCP_AVAILABLE: content=[TextContent(text=f"Error: {str(e)}", type="text")], isError=True, ) - return response @server.list_prompts() @@ -511,7 +505,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}" ) @@ -547,15 +541,12 @@ if MCP_AVAILABLE: ) -> 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 ( user_api_key_auth, @@ -565,8 +556,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() - + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" ) @@ -593,7 +583,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}" ) @@ -603,7 +593,6 @@ if MCP_AVAILABLE: verbose_logger.debug( f"MCP list_resources - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) - resources = await _list_mcp_resources( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -632,7 +621,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() + ) = await get_or_extract_auth_context() verbose_logger.debug( f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}" ) @@ -642,7 +631,6 @@ if MCP_AVAILABLE: verbose_logger.debug( f"MCP list_resource_templates - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) - resource_templates = await _list_mcp_resource_templates( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -672,8 +660,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, _client_ip, - ) = get_auth_context() - + ) = await get_or_extract_auth_context() read_resource_result = await mcp_read_resource( url=url, user_api_key_auth=user_api_key_auth, @@ -683,17 +670,14 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - return _normalize_resource_contents(read_resource_result.contents) ######################################################## ############ End of MCP Server Routes ################## ######################################################## - ######################################################## ############ Helper Functions ########################## ######################################################## - async def _get_allowed_mcp_servers_from_mcp_server_names( mcp_servers: Optional[List[str]], allowed_mcp_servers: List[MCPServer], @@ -701,13 +685,11 @@ if MCP_AVAILABLE: """ Get the filtered MCP servers from the MCP server names """ - filtered_server: 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 = [ @@ -719,12 +701,10 @@ if MCP_AVAILABLE: ] if s is not None ] - 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 = ( @@ -741,24 +721,19 @@ if MCP_AVAILABLE: verbose_logger.debug( f"Could not resolve '{server_or_group}' as access group: {e}" ) - if filtered_server: return list(filtered_server.values()) - return allowed_mcp_servers def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool: """ Check if a tool name matches any name in the filter list. - Checks both the full tool name and unprefixed version (without server prefix). This allows users to configure simple tool names regardless of prefixing. Comparison is case-insensitive to handle OpenAPI operationIds that may be in camelCase. - Args: tool_name: The tool name to check (may be prefixed like "server-tool_name") filter_list: List of tool names to match against - Returns: True if the tool name (prefixed or unprefixed) is in the filter list """ @@ -768,10 +743,8 @@ if MCP_AVAILABLE: # Normalize filter list to lowercase for case-insensitive comparison filter_list_lower = [f.lower() for f in filter_list] - if tool_name.lower() in filter_list_lower: return True - # Check if the unprefixed name is in the list (case-insensitive) unprefixed_name, _ = split_server_prefix_from_name(tool_name) return unprefixed_name.lower() in filter_list_lower @@ -782,20 +755,16 @@ if MCP_AVAILABLE: ) -> 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 """ tools_to_return = tools - # Filter by allowed_tools (whitelist) if mcp_server.allowed_tools: tools_to_return = [ @@ -803,7 +772,6 @@ if MCP_AVAILABLE: for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools) ] - # Filter by disallowed_tools (blacklist) if mcp_server.disallowed_tools: tools_to_return = [ @@ -811,7 +779,6 @@ if MCP_AVAILABLE: for tool in tools_to_return if not _tool_name_matches(tool.name, mcp_server.disallowed_tools) ] - return tools_to_return def apply_tool_overrides( @@ -819,7 +786,6 @@ if MCP_AVAILABLE: 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. """ @@ -827,7 +793,6 @@ if MCP_AVAILABLE: description_map = mcp_server.tool_name_to_description or {} if not display_name_map and not description_map: return tools - for tool in tools: unprefixed, _ = split_server_prefix_from_name(tool.name) lookup_key = unprefixed or tool.name @@ -856,7 +821,6 @@ if MCP_AVAILABLE: client_ip: Optional[str] = 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. @@ -874,7 +838,6 @@ if MCP_AVAILABLE: "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) ) @@ -906,13 +869,11 @@ if MCP_AVAILABLE: ) if mcp_server is not None: 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 async def _get_user_oauth_extra_headers_from_db( @@ -921,12 +882,10 @@ if MCP_AVAILABLE: prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, ) -> Optional[Dict[str, str]]: """Look up stored OAuth2 token for (user, server) and return as extra_headers dict. - Lookup order: 1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied 2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query 3. Auto-refresh when the stored token is expired and a refresh_token exists - Args: prefetched_creds: Optional dict keyed by server_id with credential payloads. When provided, the Redis and individual DB lookups are @@ -963,7 +922,6 @@ if MCP_AVAILABLE: server_id, ) return {"Authorization": f"Bearer {cached_token}"} - # ── Slow path: DB lookup ────────────────────────────────────────── if prefetched_creds is not None: cred = prefetched_creds.get(server_id) @@ -978,10 +936,8 @@ if MCP_AVAILABLE: cred = await get_user_oauth_credential( prisma_client, user_id, server_id ) - if not cred or not cred.get("access_token"): return None - if is_oauth_credential_expired(cred): verbose_logger.debug( "_get_user_oauth_extra_headers_from_db: token expired for " @@ -1014,16 +970,13 @@ if MCP_AVAILABLE: refresh_exc, ) cred = None - if not cred or not cred.get("access_token"): # Clear stale Redis/cache entry so we don't serve it again. # Do this for both the individual and prefetch paths so the # next request doesn't get a stale cache hit. await mcp_per_user_token_cache.delete(user_id, server_id) return None - access_token: str = cred["access_token"] - # Warm (or re-warm) the Redis cache from the DB result. # Always write regardless of whether expires_at is present — tokens # without an expiry are still valid and should be cached using the @@ -1048,7 +1001,6 @@ if MCP_AVAILABLE: await mcp_per_user_token_cache.set( user_id, server_id, access_token, ttl ) - return {"Authorization": f"Bearer {access_token}"} except Exception as e: verbose_logger.warning( @@ -1064,7 +1016,6 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], ) -> Dict[str, Dict[str, Any]]: """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 = ( @@ -1102,20 +1053,16 @@ if MCP_AVAILABLE: server_auth_header = mcp_server_auth_headers.get(server.alias) elif mcp_server_auth_headers and server.server_name is not None: server_auth_header = mcp_server_auth_headers.get(server.server_name) - extra_headers: Optional[Dict[str, str]] = None if server.auth_type == MCPAuth.oauth2: # Copy to avoid mutating the original dict (important for parallel fetching) extra_headers = oauth2_headers.copy() if oauth2_headers else None - if server.extra_headers and raw_headers: if extra_headers is None: extra_headers = {} - normalized_raw_headers = { str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str) } - for header in server.extra_headers: if not isinstance(header, str): continue @@ -1123,10 +1070,8 @@ if MCP_AVAILABLE: if header_value is None: continue extra_headers[header] = header_value - if server_auth_header is None: server_auth_header = mcp_auth_header - return server_auth_header, extra_headers def _merge_gateway_initialize_instructions( @@ -1135,7 +1080,6 @@ if MCP_AVAILABLE: """YAML/DB override, else in-memory upstream text from list_tools / health_check / call_tool.""" if not allowed_mcp_servers: return None - texts: List[Tuple[str, str]] = [] for server in allowed_mcp_servers: label = ( @@ -1155,7 +1099,6 @@ if MCP_AVAILABLE: ) if cached and cached.strip(): texts.append((label, cached.strip())) - if not texts: return None if len(texts) == 1: @@ -1193,24 +1136,20 @@ if MCP_AVAILABLE: ) -> List[MCPTool]: """ 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: List[MCPTool]: Combined list of tools from filtered servers """ if not MCP_AVAILABLE: return [] - list_tools_start_time = datetime.now() litellm_logging_obj: Optional[LiteLLMLoggingObj] = None list_tools_request_data: Dict[str, Any] = {} - if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook rules_obj = Rules() @@ -1226,7 +1165,6 @@ if MCP_AVAILABLE: 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, @@ -1246,7 +1184,6 @@ if MCP_AVAILABLE: } ], } - # Attach user identifiers using the standard helper if user_api_key_auth is not None: LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( @@ -1254,13 +1191,11 @@ if MCP_AVAILABLE: user_api_key_dict=user_api_key_auth, _metadata_variable_name="metadata", ) - user_identifier = 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", @@ -1276,13 +1211,11 @@ if MCP_AVAILABLE: "Failed to initialize logging for MCP list_tools: %s", logging_error ) litellm_logging_obj = None - try: allowed_mcp_servers = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) - # 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( @@ -1301,7 +1234,6 @@ if MCP_AVAILABLE: """Fetch and filter tools from a single server with error handling.""" if server is None: return [] - server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, @@ -1309,7 +1241,6 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - # If no OAuth2 token came from request headers, fall back to pre-fetched creds if extra_headers is None and server.auth_type == MCPAuth.oauth2: extra_headers = await _get_user_oauth_extra_headers_from_db( @@ -1317,7 +1248,6 @@ if MCP_AVAILABLE: user_api_key_auth, prefetched_creds=_prefetched_oauth_creds, ) - try: tools = await global_mcp_server_manager._get_tools_from_server( server=server, @@ -1327,17 +1257,14 @@ if MCP_AVAILABLE: raw_headers=raw_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( f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering" ) @@ -1353,10 +1280,8 @@ if MCP_AVAILABLE: _fetch_and_filter_server_tools(server) for server in allowed_mcp_servers ] results = await asyncio.gather(*tasks) - # Flatten results into single list all_tools: List[MCPTool] = [tool for tools in results for tool in tools] - # If logging is enabled, enrich spend_logs_metadata with counts if litellm_logging_obj: per_server_tool_counts: Dict[str, int] = {} @@ -1370,7 +1295,6 @@ if MCP_AVAILABLE: or "unknown" ) per_server_tool_counts[str(server_key)] = len(server_tools) - metadata_dict = litellm_logging_obj.model_call_details.get("metadata") if isinstance(metadata_dict, dict): spend_meta = metadata_dict.get("spend_logs_metadata") @@ -1380,18 +1304,15 @@ if MCP_AVAILABLE: 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 - end_time = datetime.now() await litellm_logging_obj.async_success_handler( result=all_tools, start_time=list_tools_start_time, end_time=end_time, ) - verbose_logger.info( f"Successfully fetched {len(all_tools)} tools total from all MCP servers" ) - return all_tools except Exception as e: # Only fire failure hook if logging was requested for this list-tools execution @@ -1426,31 +1347,26 @@ if MCP_AVAILABLE: ) -> 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 = 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 = [] 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, @@ -1458,7 +1374,6 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - try: prompts = await global_mcp_server_manager.get_prompts_from_server( server=server, @@ -1467,9 +1382,7 @@ if MCP_AVAILABLE: add_prefix=True, # Always add server prefix raw_headers=raw_headers, ) - all_prompts.extend(prompts) - verbose_logger.debug( f"Successfully fetched {len(prompts)} prompts from server {server.name}" ) @@ -1478,11 +1391,9 @@ if MCP_AVAILABLE: f"Error getting prompts from server {server.name}: {str(e)}" ) # Continue with other servers instead of failing completely - verbose_logger.info( f"Successfully fetched {len(all_prompts)} prompts total from all MCP servers" ) - return all_prompts async def _get_resources_from_mcp_servers( @@ -1494,20 +1405,16 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, ) -> List[Resource]: """Fetch resources from allowed MCP servers.""" - if not MCP_AVAILABLE: return [] - allowed_mcp_servers = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) - all_resources: 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, @@ -1515,7 +1422,6 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - try: resources = await global_mcp_server_manager.get_resources_from_server( server=server, @@ -1525,7 +1431,6 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) all_resources.extend(resources) - verbose_logger.debug( f"Successfully fetched {len(resources)} resources from server {server.name}" ) @@ -1533,11 +1438,9 @@ if MCP_AVAILABLE: verbose_logger.exception( f"Error getting resources from server {server.name}: {str(e)}" ) - verbose_logger.info( f"Successfully fetched {len(all_resources)} resources total from all MCP servers" ) - return all_resources async def _get_resource_templates_from_mcp_servers( @@ -1549,20 +1452,16 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, ) -> List[ResourceTemplate]: """Fetch resource templates from allowed MCP servers.""" - if not MCP_AVAILABLE: return [] - allowed_mcp_servers = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, ) - all_resource_templates: 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, @@ -1570,7 +1469,6 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - try: resource_templates = ( await global_mcp_server_manager.get_resource_templates_from_server( @@ -1593,12 +1491,10 @@ if MCP_AVAILABLE: 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( @@ -1608,7 +1504,6 @@ if MCP_AVAILABLE: ) -> 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. @@ -1630,7 +1525,6 @@ if MCP_AVAILABLE: else: # No restrictions, return all tools filtered_tools = tools - return filtered_tools async def _merge_toolset_permissions( @@ -1639,7 +1533,6 @@ if MCP_AVAILABLE: """ Resolve mcp_toolsets on the key's object_permission into tool-level permissions and merge them (union) into object_permission.mcp_tool_permissions. - Returns the (possibly mutated copy of) user_api_key_auth. """ if user_api_key_auth is None: @@ -1650,7 +1543,6 @@ if MCP_AVAILABLE: toolset_ids = getattr(op, "mcp_toolsets", None) or [] if not toolset_ids: return user_api_key_auth - toolset_perms = ( await global_mcp_server_manager.resolve_toolset_tool_permissions( toolset_ids=toolset_ids @@ -1658,14 +1550,12 @@ if MCP_AVAILABLE: ) if not toolset_perms: return user_api_key_auth - # Merge toolset_perms into existing mcp_tool_permissions (union) existing = dict(op.mcp_tool_permissions or {}) for server_id, tool_names in toolset_perms.items(): existing_tools = existing.get(server_id, []) merged = list(set(existing_tools) | set(tool_names)) existing[server_id] = merged - # Build updated object_permission with merged tool permissions and server IDs. # Union the toolset's server IDs into mcp_servers so downstream server-level # filtering doesn't silently drop servers that the toolset references but that @@ -1688,23 +1578,19 @@ if MCP_AVAILABLE: ) -> List[MCPTool]: """ 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} - Returns: List[MCPTool]: Combined list of tools from all accessible servers """ if not MCP_AVAILABLE: return [] - # Resolve toolset permissions and merge into the key's object_permission # so that the existing filter_tools_by_key_team_permissions logic picks them up. user_api_key_auth = await _merge_toolset_permissions(user_api_key_auth) - # Get tools from managed MCP servers with error handling managed_tools = [] try: @@ -1726,7 +1612,6 @@ if MCP_AVAILABLE: f"Error getting tools from managed MCP servers: {str(e)}" ) # Continue with empty managed tools list instead of failing completely - return managed_tools async def _list_mcp_prompts( @@ -1739,13 +1624,11 @@ if MCP_AVAILABLE: ) -> 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 """ @@ -1770,7 +1653,6 @@ if MCP_AVAILABLE: f"Error getting tools from managed MCP servers: {str(e)}" ) # Continue with empty managed tools list instead of failing completely - return managed_prompts async def _list_mcp_resources( @@ -1782,10 +1664,8 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = 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( @@ -1803,7 +1683,6 @@ if MCP_AVAILABLE: verbose_logger.exception( f"Error getting resources from managed MCP servers: {str(e)}" ) - return managed_resources async def _list_mcp_resource_templates( @@ -1815,10 +1694,8 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = 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( @@ -1838,7 +1715,6 @@ if MCP_AVAILABLE: "Error getting resource templates from managed MCP servers: %s", str(e), ) - return managed_resource_templates def _resolve_display_name_to_original( @@ -1846,7 +1722,6 @@ if MCP_AVAILABLE: 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. @@ -1866,7 +1741,6 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], ) -> Optional[str]: """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. """ @@ -1875,14 +1749,12 @@ if MCP_AVAILABLE: user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" if not user_id: return None - cache_key = (user_id, mcp_server.server_id) cached = _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 @@ -1908,7 +1780,6 @@ if MCP_AVAILABLE: """ if not mcp_server.is_byok: return - user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" if not user_id: raise HTTPException( @@ -1923,7 +1794,6 @@ if MCP_AVAILABLE: "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' }, ) - # Check shared credential cache before hitting the DB. cache_key = (user_id, mcp_server.server_id) cached = _byok_cred_cache.get(cache_key) @@ -1947,7 +1817,6 @@ if MCP_AVAILABLE: }, ) return - from litellm.proxy._experimental.mcp_server.db import get_user_credential from litellm.proxy.proxy_server import prisma_client @@ -1964,7 +1833,6 @@ if MCP_AVAILABLE: "message": "BYOK credential check requires a database connection.", }, ) - credential = await get_user_credential( prisma_client=prisma_client, user_id=user_id, @@ -2003,9 +1871,7 @@ if MCP_AVAILABLE: ) -> CallToolResult: """ Execute MCP tool. - This function assumes permission checks have already been performed. - Args: name: Tool name (may include server prefix) arguments: Tool arguments @@ -2017,26 +1883,21 @@ if MCP_AVAILABLE: 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: Optional[MCPServer] = None - # 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) - # If tool name is unprefixed, resolve its server so we can enforce permissions if not server_name: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server: server_name = mcp_server.name - # Only enforce server-level permissions when we can resolve a server if server_name: if not MCPRequestHandler.is_tool_allowed( @@ -2047,7 +1908,6 @@ if MCP_AVAILABLE: status_code=403, detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}", ) - standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = ( _get_standard_logging_mcp_tool_call( name=original_tool_name, # Use original name for logging @@ -2067,7 +1927,6 @@ if MCP_AVAILABLE: # 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 {} @@ -2076,7 +1935,6 @@ if MCP_AVAILABLE: 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: @@ -2101,7 +1959,6 @@ if MCP_AVAILABLE: 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 ######################################################### @@ -2129,7 +1986,6 @@ if MCP_AVAILABLE: finally: _request_auth_header.reset(_auth_token) response = CallToolResult(content=cast(Any, local_content), isError=False) - # Try managed MCP server tool (pass the full prefixed name) # Primary and recommended way to use external MCP servers ######################################################### @@ -2146,7 +2002,6 @@ if MCP_AVAILABLE: 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 @@ -2154,7 +2009,6 @@ if MCP_AVAILABLE: else: local_content = await _handle_local_mcp_tool(original_tool_name, arguments) response = CallToolResult(content=cast(Any, local_content), isError=False) - return response @client @@ -2176,20 +2030,17 @@ if MCP_AVAILABLE: litellm_logging_obj: Optional[LiteLLMLoggingObj] = 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 = ( 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( @@ -2197,7 +2048,6 @@ if MCP_AVAILABLE: ) if allowed_server is not None: 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, @@ -2207,7 +2057,6 @@ if MCP_AVAILABLE: 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, @@ -2234,7 +2083,6 @@ if MCP_AVAILABLE: traceback_str=traceback_str, ) raise - if litellm_logging_obj: litellm_logging_obj.post_call(original_response=response) end_time = datetime.now() @@ -2267,23 +2115,19 @@ if MCP_AVAILABLE: 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 = 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, @@ -2291,7 +2135,6 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - return await global_mcp_server_manager.get_prompt_from_server( server=server, prompt_name=original_prompt_name, @@ -2311,18 +2154,15 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, ) -> ReadResourceResult: """Read resource contents from upstream MCP servers.""" - allowed_mcp_servers = 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, @@ -2331,9 +2171,7 @@ if MCP_AVAILABLE: "supports exactly one allowed server." ), ) - server = allowed_mcp_servers[0] - server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, @@ -2341,7 +2179,6 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - return await global_mcp_server_manager.read_resource_from_server( server=server, url=url, @@ -2415,7 +2252,6 @@ if MCP_AVAILABLE: tool = global_mcp_tool_registry.get_tool(name) if not tool: raise HTTPException(status_code=404, detail=f"Tool '{name}' not found") - try: # Check if handler is async or sync if inspect.iscoroutinefunction(tool.handler): @@ -2433,6 +2269,9 @@ if MCP_AVAILABLE: """ import re + # Ignore literal endpoints for the bidirectional MCP server itself + if re.match(r"^/mcp/(sse|messages)(?:\?.*)?(?:#.*)?$", path): + return None mcp_servers_from_path: Optional[List[str]] = None # Match /mcp/ # Where servers can be comma-separated list of server names @@ -2440,7 +2279,6 @@ if MCP_AVAILABLE: mcp_path_match = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path) if mcp_path_match: servers_and_path = mcp_path_match.group(1) - if servers_and_path: # Check if it contains commas (comma-separated servers) if "," in servers_and_path: @@ -2518,12 +2356,10 @@ if MCP_AVAILABLE: 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 = b"mcp-session-id" @@ -2544,16 +2380,13 @@ if MCP_AVAILABLE: else: _session_id = str(header_value) break - if _session_id is None: return False - # Check in-memory session tracking known_sessions = 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: @@ -2565,10 +2398,8 @@ if MCP_AVAILABLE: _session_id, ) return False - # --- Session not in this worker's memory --- method = scope.get("method", "").upper() - if method == "DELETE": verbose_logger.info( "DELETE request for non-existent MCP session '%s'. " @@ -2581,7 +2412,6 @@ if MCP_AVAILABLE: ) 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. " @@ -2601,10 +2431,8 @@ if MCP_AVAILABLE: ) -> 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. @@ -2626,7 +2454,6 @@ if MCP_AVAILABLE: 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] @@ -2666,10 +2493,8 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) - # Extract client IP for MCP access control _client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) - verbose_logger.debug( f"MCP request mcp_servers (header/path): {mcp_servers}" ) @@ -2695,28 +2520,23 @@ if MCP_AVAILABLE: ) if stored_oauth_headers: continue - request = StarletteRequest(scope) base_url = get_request_base_url(request) - authorization_uri = ( f"Bearer authorization_uri=" f"{base_url}/.well-known/oauth-authorization-server/{server_name}" ) - raise HTTPException( status_code=401, detail="Unauthorized", headers={"www-authenticate": authorization_uri}, ) - # Strip any client-supplied x-mcp-toolset-id to prevent forgery. scope["headers"] = [ (k, v) for k, v in scope.get("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 = _mcp_active_toolset_id.get() @@ -2724,7 +2544,6 @@ if MCP_AVAILABLE: user_api_key_auth = await _apply_toolset_scope( user_api_key_auth, active_toolset_id ) - # Inject masked debug headers when client sends x-litellm-mcp-debug: true _debug_headers = MCPDebug.maybe_build_debug_headers( raw_headers=raw_headers, @@ -2737,7 +2556,6 @@ if MCP_AVAILABLE: ) if _debug_headers: send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers) - # Set the auth context variable for easy access in MCP functions set_auth_context( user_api_key_auth=user_api_key_auth, @@ -2748,13 +2566,11 @@ if MCP_AVAILABLE: raw_headers=raw_headers, client_ip=_client_ip, ) - # 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) - # Handle stale session IDs - either strip them for reconnection # or return success for idempotent DELETE operations handled = await _handle_stale_mcp_session( @@ -2763,7 +2579,6 @@ if MCP_AVAILABLE: if handled: # Request was fully handled (e.g., DELETE on non-existent session) return - async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, @@ -2794,7 +2609,18 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" + + async def handle_sse_mcp_endpoint(request: StarletteRequest): + """ + Handle MCP SSE GET requests. + This is a Starlette Route endpoint handler (takes Request, returns Response). + Follows the pattern documented in the official MCP SDK source at + mcp/server/sse.py lines 6-31. + CRITICAL: Must return Response() after the SSE connection ends to prevent + "TypeError: 'NoneType' object is not callable" when the client disconnects. + """ try: + scope = request.scope path = scope.get("path", "") ( user_api_key_auth, @@ -2804,12 +2630,10 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) - # Extract client IP for MCP access control - _sse_client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) - + _sse_client_ip = IPAddressUtils.get_mcp_client_ip(request) verbose_logger.debug( - f"MCP request mcp_servers (header/path): {mcp_servers}" + f"MCP SSE request mcp_servers (header/path): {mcp_servers}" ) verbose_logger.debug( f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" @@ -2823,17 +2647,43 @@ if MCP_AVAILABLE: raw_headers=raw_headers, client_ip=_sse_client_ip, ) - + # Also persist auth context on the server object itself. + # ContextVars are lost when the MCP SDK spawns internal tasks + # (e.g. _receive_loop), so tool handlers can't read auth_context_var. + # Storing it on the server object makes it available everywhere. + server._litellm_auth_context = 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=_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, ): - await sse_session_manager.handle_request(scope, receive, send) + verbose_logger.info("Initializing SSE session...") + options = server.create_initialization_options() + async with sse.connect_sse( + request.scope, request.receive, request._send + ) as streams: + verbose_logger.info( + "SSE connection established, running server loop..." + ) + try: + # Capture the session for propagation to sampling/elicitation callbacks + # Since server.run doesn't return the session, we use a middleware-like + # wrapper if the SDK allows, or we rely on the fact that for SSE, + # there's usually one active session per request. + await server.run(streams[0], streams[1], options) + except Exception as session_e: + verbose_logger.exception(f"Error in SSE session: {session_e}") except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Instead of re-raising, try to send a graceful error response @@ -2846,13 +2696,18 @@ if MCP_AVAILABLE: status_code=HTTP_500_INTERNAL_SERVER_ERROR, content={"error": "MCP request failed", "details": str(e)}, ) - await error_response(scope, receive, send) + await error_response(request.scope, request.receive, request._send) except Exception as response_error: verbose_logger.exception( f"Failed to send error response: {response_error}" ) # If we can't send a proper response, re-raise the original error raise e + # CRITICAL: Return empty Response to prevent NoneType crash. + # See MCP SDK docstring at mcp/server/sse.py lines 25-26. + from starlette.responses import Response as StarletteResponse + + return StarletteResponse() app = FastAPI( title=LITELLM_MCP_SERVER_NAME, @@ -2872,17 +2727,80 @@ if MCP_AVAILABLE: """ return {"enabled": MCP_AVAILABLE} - # Mount the MCP handlers - app.mount("/", handle_streamable_http_mcp) + # Include the MCP router + app.include_router(router) + # Mount SSE handlers using the SDK's documented pattern. + # We use Starlette Route for the SSE GET endpoint (must return Response), + # and a FastAPI POST route for the POST messages endpoint. + from starlette.routing import Route as StarletteRoute + + app.routes.insert( + 0, StarletteRoute("/sse", endpoint=handle_sse_mcp_endpoint, methods=["GET"]) + ) + from starlette.responses import Response as StarletteResponse + + class NoOpResponse(StarletteResponse): + """A response that does nothing. Used when the underlying ASGI app already sent the response.""" + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + pass + + @app.post("/messages", include_in_schema=False) + async def handle_sse_post_messages(request: StarletteRequest): + """Handle SSE POST messages by delegating to the SDK's handle_post_message.""" + try: + scope = request.scope + path = 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) + _sse_client_ip = IPAddressUtils.get_mcp_client_ip(request) + 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, + ) + except Exception as e: + verbose_logger.warning( + f"Failed to extract auth context in POST /messages: {e}" + ) + # The SDK's handler calls `send` directly. + await sse.handle_post_message(request.scope, request.receive, request._send) + # Return NoOpResponse to prevent Starlette from sending a second response. + return NoOpResponse() + + def get_active_mcp_session() -> Optional[_McpServerSession]: + """Get the active downstream MCP session from the current context.""" + return active_mcp_session_var.get() + + def get_active_auth_context() -> Optional[MCPAuthenticatedUser]: + """Get the active auth context from the server object or context var.""" + # Check context var first + auth = auth_context_var.get() + if auth: + return auth + # Fallback to server object + return getattr(server, "_litellm_auth_context", None) + + # StreamableHTTP catch-all mounts (must come after specific routes) app.mount("/mcp", handle_streamable_http_mcp) app.mount("/{mcp_server_name}/mcp", handle_streamable_http_mcp) app.mount("/sse", handle_sse_mcp) + app.mount("/", handle_streamable_http_mcp) app.add_middleware(AuthContextMiddleware) ######################################################## ############ Auth Context Functions #################### ######################################################## - def set_auth_context( user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, @@ -2894,7 +2812,6 @@ if MCP_AVAILABLE: ) -> None: """ 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) @@ -2911,6 +2828,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, client_ip=client_ip, ) + verbose_logger.debug(f"set_auth_context called with mcp_servers={mcp_servers}") auth_context_var.set(auth_user) def get_auth_context() -> Tuple[ @@ -2924,7 +2842,6 @@ if MCP_AVAILABLE: ]: """ 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 @@ -2942,9 +2859,57 @@ if MCP_AVAILABLE: ) return None, None, None, None, None, None, None + async def get_or_extract_auth_context() -> Tuple[ + Optional[UserAPIKeyAuth], + Optional[str], + Optional[List[str]], + Optional[Dict[str, Dict[str, str]]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], + Optional[str], + ]: + """ + Get auth context from ContextVar first, then fall back to the + server object (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() + # Fallback: read from server object if ContextVar was lost + if user_api_key_auth is None: + stored = getattr(server, "_litellm_auth_context", None) + verbose_logger.debug( + f"get_or_extract_auth_context FALLBACK: stored={stored}, type={type(stored)}" + ) + if stored and isinstance(stored, MCPAuthenticatedUser): + verbose_logger.debug( + "get_or_extract_auth_context: Recovered auth from server object" + ) + 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, + ) + ######################################################## ############ End of Auth Context Functions ############# ######################################################## - else: app = FastAPI() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 92c920ca594..e5e2fc82b88 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1215,7 +1215,12 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): raise ValueError("args is required for stdio transport") # Validate command against allowlist to prevent arbitrary execution base_command = os.path.basename(values["command"]) - if base_command not in MCP_STDIO_ALLOWED_COMMANDS: + # Strip .exe/.cmd/.bat suffix for Windows compatibility + base_command_no_ext = os.path.splitext(base_command)[0] + if ( + base_command not in MCP_STDIO_ALLOWED_COMMANDS + and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS + ): raise ValueError( f"Command '{values['command']}' is not in the allowed commands list " f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}" @@ -1283,7 +1288,12 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): raise ValueError("args is required for stdio transport") # Validate command against allowlist to prevent arbitrary execution base_command = os.path.basename(values["command"]) - if base_command not in MCP_STDIO_ALLOWED_COMMANDS: + # Strip .exe/.cmd/.bat suffix for Windows compatibility + base_command_no_ext = os.path.splitext(base_command)[0] + if ( + base_command not in MCP_STDIO_ALLOWED_COMMANDS + and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS + ): raise ValueError( f"Command '{values['command']}' is not in the allowed commands list " f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}" diff --git a/tests/mcp_sampling_elicitation/custom_mcp_server.py b/tests/mcp_sampling_elicitation/custom_mcp_server.py new file mode 100644 index 00000000000..3f90759d2a8 --- /dev/null +++ b/tests/mcp_sampling_elicitation/custom_mcp_server.py @@ -0,0 +1,105 @@ +import asyncio +from mcp.server import Server +from mcp.server.stdio import stdio_server +from mcp.types import ( + Tool, + CallToolResult, + TextContent, + SamplingMessage, +) + +import logging + +logging.basicConfig(filename="custom_server.log", level=logging.DEBUG, force=True) + +server = Server("custom-test-server") + + +@server.list_tools() +async def list_tools() -> list[Tool]: + return [ + Tool( + name="test_complex_pipeline", + description="Tests the full pipeline: Asks user for a topic, then asks AI to write a story about it.", + inputSchema={"type": "object", "properties": {}}, + ) + ] + + +@server.call_tool() +async def call_tool(name: str, arguments: dict) -> CallToolResult: + session = server.request_context.session + + if name == "test_complex_pipeline": + # 1. Elicitation: Ask the user for inputs + elicit_result = await session.elicit_form( + message="Please provide details for the story", + requestedSchema={ + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "What should the story be about?", + }, + "adjective": { + "type": "string", + "description": "What is the tone of the story?", + }, + }, + "required": ["topic", "adjective"], + }, + ) + + # Parse the user's response + topic = "a random thing" + adjective = "weird" + content = getattr(elicit_result, "content", None) + if isinstance(content, dict): + topic = content.get("topic", topic) + adjective = content.get("adjective", adjective) + elif hasattr(content, "topic"): + topic = content.topic + adjective = content.adjective + + logging.info(f"Got elicitation: topic={topic}, adjective={adjective}") + + # 2. Sampling: Ask the AI to write the story based on user input + sample_result = await session.create_message( + messages=[ + SamplingMessage( + role="user", + content=TextContent( + type="text", + text=f"Write a very short, 3-sentence {adjective} story about {topic}.", + ), + ) + ], + max_tokens=150, + ) + + ai_story = sample_result.content.text + + # 3. Return final result + return CallToolResult( + content=[ + TextContent( + type="text", + text=f"Pipeline Complete!\nUser chose: a {adjective} story about '{topic}'.\n\nAI generated story:\n{ai_story}", + ) + ] + ) + + return CallToolResult( + content=[TextContent(type="text", text="Tool not found")], isError=True + ) + + +async def main(): + async with stdio_server() as (read_stream, write_stream): + await server.run( + read_stream, write_stream, server.create_initialization_options() + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/mcp_sampling_elicitation/mcp_test_config.yaml b/tests/mcp_sampling_elicitation/mcp_test_config.yaml new file mode 100644 index 00000000000..e50ad9bf751 --- /dev/null +++ b/tests/mcp_sampling_elicitation/mcp_test_config.yaml @@ -0,0 +1,13 @@ +model_list: + - model_name: groq-model + litellm_params: + model: groq/llama-3.3-70b-versatile + api_key: os.environ/GROQ_API_KEY +litellm_settings: + default_mcp_sampling_model: groq/llama-3.3-70b-versatile +mcp_servers: + test_server: + transport: stdio + command: "c:\\Users\\DELL\\Desktop\\litellm\\.venv\\Scripts\\python.exe" + args: ["c:\\Users\\DELL\\Desktop\\litellm\\tests\\mcp_sampling_elicitation\\custom_mcp_server.py"] + allow_all_keys: true diff --git a/tests/mcp_sampling_elicitation/test_live_mcp.py b/tests/mcp_sampling_elicitation/test_live_mcp.py new file mode 100644 index 00000000000..22d3d6cad8e --- /dev/null +++ b/tests/mcp_sampling_elicitation/test_live_mcp.py @@ -0,0 +1,47 @@ +import asyncio +from mcp.client.sse import sse_client +from mcp.client.session import ClientSession +from mcp.types import ElicitResult + + +async def main(): + print("Connecting to LiteLLM Proxy via SSE...") + + async def my_elicitation_callback(context, params): + print(f"\n[CLIENT] Received elicitation request from upstream!") + + # We will simulate the user filling out the form + user_response = { + "topic": "a time-traveling developer", + "adjective": "suspenseful", + } + + print(f"[CLIENT] User is filling the form with: {user_response}") + + return ElicitResult(action="accept", content=user_response) + + async with sse_client( + "http://localhost:4000/mcp/sse", headers={"Authorization": "Bearer sk-1234"} + ) as (read_stream, write_stream): + print("SSE connection established.") + async with ClientSession( + read_stream, write_stream, elicitation_callback=my_elicitation_callback + ) as session: + await session.initialize() + print("Initialized!") + + print("\n--- Testing Complex Pipeline (Elicitation + Sampling) ---") + print("Calling 'test_server-test_complex_pipeline'...") + try: + result = await session.call_tool( + "test_server-test_complex_pipeline", arguments={} + ) + print("\nFINAL TOOL RESULT:") + print("==================") + print(result.content[0].text) + except Exception as e: + print(f"Error calling test_complex_pipeline: {e}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/scratch/mcp_test_config.yaml b/tests/scratch/mcp_test_config.yaml new file mode 100644 index 00000000000..7795f48e457 --- /dev/null +++ b/tests/scratch/mcp_test_config.yaml @@ -0,0 +1,13 @@ +model_list: + - model_name: groq-model + litellm_params: + model: groq/llama-3.3-70b-versatile + api_key: os.environ/GROQ_API_KEY +litellm_settings: + default_mcp_sampling_model: groq-model +mcp_servers: + test_server: + transport: stdio + command: "c:\\Users\\DELL\\Desktop\\litellm\\.venv\\Scripts\\python.exe" + args: ["c:\\Users\\DELL\\Desktop\\litellm\\tests\\scratch\\custom_mcp_server.py"] + allow_all_keys: true \ No newline at end of file diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py new file mode 100644 index 00000000000..44d9bcf07f7 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py @@ -0,0 +1,188 @@ +""" +Unit tests for the MCP Elicitation Handler. +Tests the elicitation/create handler that relays elicitation requests +from upstream MCP servers to downstream clients or declines them. +""" +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +# ───────────────────────────────────────────────────────────── +# Helper factories +# ───────────────────────────────────────────────────────────── +def _make_form_params(message="Please provide info", schema=None): + """Create a mock ElicitRequestFormParams.""" + from mcp.types import ElicitRequestFormParams + params = MagicMock(spec=ElicitRequestFormParams) + params.mode = "form" + params.message = message + params.requestedSchema = schema + return params +def _make_url_params(message="Click the link", url="https://auth.example.com"): + """Create a mock ElicitRequestURLParams.""" + from mcp.types import ElicitRequestURLParams + params = MagicMock(spec=ElicitRequestURLParams) + params.mode = "url" + params.message = message + params.url = url + params.elicitationId = "elicit-123" + return params +def _make_capabilities(form=True, url=True): + """Create mock client capabilities with elicitation support.""" + caps = MagicMock() + elicit = MagicMock() + elicit.form = MagicMock() if form else None + elicit.url = MagicMock() if url else None + caps.elicitation = elicit + return caps +def _make_capabilities_no_elicitation(): + """Create mock client capabilities without elicitation.""" + caps = MagicMock() + caps.elicitation = None + return caps +# ───────────────────────────────────────────────────────────── +# Tests: No downstream session (Tool Bridge mode) +# ───────────────────────────────────────────────────────────── +class TestElicitationNoDownstream: + """Tests when no downstream client is available.""" + @pytest.mark.asyncio + async def test_should_decline_when_no_downstream_session(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_form_params() + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=None, + ) + assert result.action == "decline" +# ───────────────────────────────────────────────────────────── +# Tests: Downstream session relay +# ───────────────────────────────────────────────────────────── +class TestElicitationRelay: + """Tests for relaying elicitation to downstream clients.""" + @pytest.mark.asyncio + async def test_should_relay_form_mode(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_form_params(message="Enter your name") + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.action = "submit" + mock_result.content = {"name": "Alice"} + mock_session.elicit_form.return_value = mock_result + caps = _make_capabilities(form=True, url=True) + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=caps, + ) + mock_session.elicit_form.assert_called_once() + assert result.action == "submit" + @pytest.mark.asyncio + async def test_should_relay_url_mode(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_url_params( + message="Authenticate", url="https://oauth.example.com" + ) + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.action = "submit" + mock_session.elicit_url.return_value = mock_result + caps = _make_capabilities(form=True, url=True) + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=caps, + ) + mock_session.elicit_url.assert_called_once() + assert result.action == "submit" + @pytest.mark.asyncio + async def test_should_decline_when_client_lacks_elicitation(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_form_params() + mock_session = AsyncMock() + caps = _make_capabilities_no_elicitation() + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=caps, + ) + assert result.action == "decline" + @pytest.mark.asyncio + async def test_should_decline_when_client_lacks_url_mode(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_url_params() + mock_session = AsyncMock() + caps = _make_capabilities(form=True, url=False) + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=caps, + ) + assert result.action == "decline" + @pytest.mark.asyncio + async def test_should_decline_when_client_lacks_form_mode(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_form_params() + mock_session = AsyncMock() + caps = _make_capabilities(form=False, url=True) + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=caps, + ) + assert result.action == "decline" + @pytest.mark.asyncio + async def test_should_decline_gracefully_on_relay_failure(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_form_params() + mock_session = AsyncMock() + mock_session.elicit_form.side_effect = Exception("Connection lost") + caps = _make_capabilities(form=True, url=True) + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=caps, + ) + assert result.action == "decline" +# ───────────────────────────────────────────────────────────── +# Tests: Error handling +# ───────────────────────────────────────────────────────────── +class TestElicitationErrorHandling: + """Tests for error handling in the elicitation handler.""" + @pytest.mark.asyncio + async def test_should_relay_without_capability_check_when_caps_none(self): + from litellm.proxy._experimental.mcp_server.elicitation_handler import ( + handle_elicitation_request, + ) + params = _make_form_params() + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.action = "submit" + mock_session.elicit_form.return_value = mock_result + # No capabilities provided — should still attempt relay + result = await handle_elicitation_request( + context=MagicMock(), + params=params, + downstream_session=mock_session, + downstream_capabilities=None, + ) + mock_session.elicit_form.assert_called_once() + assert result.action == "submit" \ No newline at end of file diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_handler.py new file mode 100644 index 00000000000..681be0056d4 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_handler.py @@ -0,0 +1,385 @@ +""" +Unit tests for the MCP Sampling Handler. +Tests the sampling/createMessage handler that routes MCP sampling +requests through litellm.acompletion(). +""" +import json +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +# ───────────────────────────────────────────────────────────── +# Helper factories +# ───────────────────────────────────────────────────────────── +def _make_text_content(text: str): + """Create a mock TextContent.""" + tc = MagicMock() + tc.type = "text" + tc.text = text + return tc +def _make_image_content(data: str = "base64data", mime_type: str = "image/png"): + """Create a mock ImageContent.""" + ic = MagicMock() + ic.type = "image" + ic.data = data + ic.mimeType = mime_type + return ic +def _make_sampling_message(role: str, content): + """Create a mock SamplingMessage.""" + msg = MagicMock() + msg.role = role + msg.content = content + return msg +def _make_model_preferences(hints=None, cost=None, speed=None, intelligence=None): + """Create a mock ModelPreferences.""" + prefs = MagicMock() + prefs.hints = hints or [] + prefs.costPriority = cost + prefs.speedPriority = speed + prefs.intelligencePriority = intelligence + return prefs +def _make_hint(name: str): + """Create a mock model hint.""" + hint = MagicMock() + hint.name = name + return hint +def _make_params( + messages=None, + model_preferences=None, + system_prompt=None, + max_tokens=100, + temperature=None, + stop_sequences=None, + tools=None, + tool_choice=None, + metadata=None, +): + """Create a mock CreateMessageRequestParams.""" + params = MagicMock() + params.messages = messages or [] + params.modelPreferences = model_preferences + params.systemPrompt = system_prompt + params.maxTokens = max_tokens + params.temperature = temperature + params.stopSequences = stop_sequences + params.tools = tools + params.toolChoice = tool_choice + params.metadata = metadata + return params +def _make_completion_response( + content="Hello!", model="gpt-4o-mini", finish_reason="stop", tool_calls=None +): + """Create a mock litellm completion response.""" + response = MagicMock() + choice = MagicMock() + choice.finish_reason = finish_reason + message = MagicMock() + message.content = content + message.tool_calls = tool_calls + choice.message = message + response.choices = [choice] + response.model = model + return response +# ───────────────────────────────────────────────────────────── +# Tests: Message conversion +# ───────────────────────────────────────────────────────────── +class TestConvertMCPMessagesToOpenAI: + """Tests for _convert_mcp_messages_to_openai.""" + def test_should_convert_simple_text_message(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_messages_to_openai, + ) + tc = _make_text_content("Hello") + msg = _make_sampling_message("user", tc) + result = _convert_mcp_messages_to_openai([msg]) + assert len(result) == 1 + assert result[0]["role"] == "user" + def test_should_add_system_prompt(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_messages_to_openai, + ) + tc = _make_text_content("Hello") + msg = _make_sampling_message("user", tc) + result = _convert_mcp_messages_to_openai([msg], system_prompt="Be helpful") + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[0]["content"] == "Be helpful" + def test_should_convert_image_content(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_messages_to_openai, + ) + ic = _make_image_content("base64imgdata", "image/jpeg") + msg = _make_sampling_message("user", ic) + result = _convert_mcp_messages_to_openai([msg]) + assert len(result) == 1 + content = result[0]["content"] + assert isinstance(content, list) + assert content[0]["type"] == "image_url" + assert "base64imgdata" in content[0]["image_url"]["url"] + def test_should_convert_list_of_mixed_content(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_messages_to_openai, + ) + tc = _make_text_content("Describe this image") + ic = _make_image_content("imgdata") + msg = _make_sampling_message("user", [tc, ic]) + result = _convert_mcp_messages_to_openai([msg]) + assert len(result) == 1 + content = result[0]["content"] + assert isinstance(content, list) + assert len(content) == 2 + def test_should_convert_multiple_messages(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_messages_to_openai, + ) + user_msg = _make_sampling_message("user", _make_text_content("Hi")) + assistant_msg = _make_sampling_message( + "assistant", _make_text_content("Hello!") + ) + result = _convert_mcp_messages_to_openai([user_msg, assistant_msg]) + assert len(result) == 2 + assert result[0]["role"] == "user" + assert result[1]["role"] == "assistant" +# ───────────────────────────────────────────────────────────── +# Tests: Model resolution +# ───────────────────────────────────────────────────────────── +class TestResolveModel: + """Tests for _resolve_model_from_preferences.""" + def test_should_use_default_model_when_no_preferences(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _resolve_model_from_preferences, + ) + result = _resolve_model_from_preferences(None, default_model="claude-3.5-sonnet") + assert result == "claude-3.5-sonnet" + def test_should_fallback_to_gpt4o_mini(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _resolve_model_from_preferences, + ) + with patch("litellm.model_list", []): + result = _resolve_model_from_preferences(None) + assert result == "gpt-4o-mini" + @patch("litellm.model_list", ["gpt-4o", "claude-3.5-sonnet", "gemini-pro"]) + def test_should_match_hint_by_substring(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _resolve_model_from_preferences, + ) + hint = _make_hint("claude") + prefs = _make_model_preferences(hints=[hint]) + result = _resolve_model_from_preferences(prefs) + assert "claude" in result.lower() + @patch("litellm.model_list", ["gpt-4o"]) + def test_should_use_default_when_no_hint_matches(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _resolve_model_from_preferences, + ) + hint = _make_hint("nonexistent-model") + prefs = _make_model_preferences(hints=[hint]) + result = _resolve_model_from_preferences(prefs, default_model="gpt-4o") + assert result == "gpt-4o" +# ───────────────────────────────────────────────────────────── +# Tests: Tool conversion +# ───────────────────────────────────────────────────────────── +class TestConvertMCPToolsToOpenAI: + """Tests for _convert_mcp_tools_to_openai.""" + def test_should_return_none_for_no_tools(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tools_to_openai, + ) + assert _convert_mcp_tools_to_openai(None) is None + assert _convert_mcp_tools_to_openai([]) is None + def test_should_convert_mcp_tool_to_openai_format(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tools_to_openai, + ) + tool = MagicMock() + tool.name = "get_weather" + tool.description = "Get weather for a city" + tool.inputSchema = { + "type": "object", + "properties": {"city": {"type": "string"}}, + } + result = _convert_mcp_tools_to_openai([tool]) + assert result is not None + assert len(result) == 1 + assert result[0]["type"] == "function" + assert result[0]["function"]["name"] == "get_weather" + assert result[0]["function"]["description"] == "Get weather for a city" +# ───────────────────────────────────────────────────────────── +# Tests: Tool choice conversion +# ───────────────────────────────────────────────────────────── +class TestConvertMCPToolChoiceToOpenAI: + """Tests for _convert_mcp_tool_choice_to_openai.""" + def test_should_return_none_for_no_tool_choice(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tool_choice_to_openai, + ) + assert _convert_mcp_tool_choice_to_openai(None) is None + def test_should_convert_auto_mode(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tool_choice_to_openai, + ) + tc = MagicMock() + tc.mode = "auto" + assert _convert_mcp_tool_choice_to_openai(tc) == "auto" + def test_should_convert_required_mode(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tool_choice_to_openai, + ) + tc = MagicMock() + tc.mode = "required" + assert _convert_mcp_tool_choice_to_openai(tc) == "required" + def test_should_convert_none_mode(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_mcp_tool_choice_to_openai, + ) + tc = MagicMock() + tc.mode = "none" + assert _convert_mcp_tool_choice_to_openai(tc) == "none" +# ───────────────────────────────────────────────────────────── +# Tests: Response conversion +# ───────────────────────────────────────────────────────────── +class TestConvertOpenAIResponseToMCPResult: + """Tests for _convert_openai_response_to_mcp_result.""" + def test_should_convert_text_response(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_openai_response_to_mcp_result, + ) + response = _make_completion_response(content="Hello!", model="gpt-4o") + result = _convert_openai_response_to_mcp_result(response, "gpt-4o") + assert result.role == "assistant" + assert result.model == "gpt-4o" + assert result.stopReason == "endTurn" + assert result.content.text == "Hello!" + def test_should_set_max_tokens_stop_reason(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_openai_response_to_mcp_result, + ) + response = _make_completion_response( + content="Partial...", finish_reason="length" + ) + result = _convert_openai_response_to_mcp_result(response, "gpt-4o") + assert result.stopReason == "maxTokens" + def test_should_convert_tool_calls_response(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + _convert_openai_response_to_mcp_result, + ) + tc = MagicMock() + tc.id = "call_123" + tc.function.name = "get_weather" + tc.function.arguments = '{"city": "NYC"}' + response = _make_completion_response( + content=None, finish_reason="tool_calls", tool_calls=[tc] + ) + result = _convert_openai_response_to_mcp_result(response, "gpt-4o") + assert result.stopReason == "toolUse" + assert isinstance(result.content, list) + # Should contain ToolUseContent + tool_use = result.content[0] + assert tool_use.type == "tool_use" + assert tool_use.name == "get_weather" +# ───────────────────────────────────────────────────────────── +# Tests: Full handler +# ───────────────────────────────────────────────────────────── +class TestHandleSamplingCreateMessage: + """Tests for the main handle_sampling_create_message function.""" + @pytest.mark.asyncio + async def test_should_call_litellm_acompletion(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + mock_response = _make_completion_response(content="Test response") + params = _make_params( + messages=[_make_sampling_message("user", _make_text_content("Hello"))], + max_tokens=100, + ) + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_completion: + mock_completion.return_value = mock_response + result = await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o-mini", + ) + mock_completion.assert_called_once() + assert result.role == "assistant" + assert result.content.text == "Test response" + @pytest.mark.asyncio + async def test_should_include_temperature(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + mock_response = _make_completion_response() + params = _make_params( + messages=[_make_sampling_message("user", _make_text_content("Hi"))], + temperature=0.7, + ) + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_completion: + mock_completion.return_value = mock_response + await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o-mini", + ) + call_kwargs = mock_completion.call_args[1] + assert call_kwargs["temperature"] == 0.7 + @pytest.mark.asyncio + async def test_should_include_stop_sequences(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + mock_response = _make_completion_response() + params = _make_params( + messages=[_make_sampling_message("user", _make_text_content("Hi"))], + stop_sequences=["STOP", "END"], + ) + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_completion: + mock_completion.return_value = mock_response + await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o-mini", + ) + call_kwargs = mock_completion.call_args[1] + assert call_kwargs["stop"] == ["STOP", "END"] + @pytest.mark.asyncio + async def test_should_include_tools_and_tool_choice(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + mock_response = _make_completion_response() + tool = MagicMock() + tool.name = "search" + tool.description = "Search the web" + tool.inputSchema = {"type": "object", "properties": {"q": {"type": "string"}}} + tc = MagicMock() + tc.mode = "auto" + params = _make_params( + messages=[_make_sampling_message("user", _make_text_content("Search"))], + tools=[tool], + tool_choice=tc, + ) + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_completion: + mock_completion.return_value = mock_response + await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o-mini", + ) + call_kwargs = mock_completion.call_args[1] + assert "tools" in call_kwargs + assert call_kwargs["tool_choice"] == "auto" + @pytest.mark.asyncio + async def test_should_return_error_on_exception(self): + from litellm.proxy._experimental.mcp_server.sampling_handler import ( + handle_sampling_create_message, + ) + params = _make_params( + messages=[_make_sampling_message("user", _make_text_content("Hi"))], + ) + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_completion: + mock_completion.side_effect = Exception("API error") + result = await handle_sampling_create_message( + context=MagicMock(), + params=params, + default_model="gpt-4o-mini", + ) + assert hasattr(result, "code") + assert result.code == -1 + assert "API error" in result.message \ No newline at end of file