fix(mcp): harden sampling and elicitation flows with proper context propagation

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

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")

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(

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)}",
+        )

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)}"

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())

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

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())

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

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"

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