Merge branch 'litellm_internal_staging' into litellm_mcp_server_env_vars

litellm/__init__.py

@@ -1740,6 +1740,9 @@ if TYPE_CHECKING:
     from .llms.openrouter.responses.transformation import (
         OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig,
     )
+    from .llms.bedrock_mantle.responses.transformation import (
+        BedrockMantleResponsesAPIConfig as BedrockMantleResponsesAPIConfig,
+    )
     from .llms.gemini.interactions.transformation import (
         GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
     )

litellm/_lazy_imports_registry.py

@@ -237,6 +237,7 @@ LLM_CONFIG_NAMES = (
     "PerplexityResponsesConfig",
     "DatabricksResponsesAPIConfig",
     "OpenRouterResponsesAPIConfig",
+    "BedrockMantleResponsesAPIConfig",
     "GoogleAIStudioInteractionsConfig",
     "OpenAIOSeriesConfig",
     "AnthropicSkillsConfig",
@@ -958,6 +959,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
         ".llms.openrouter.responses.transformation",
         "OpenRouterResponsesAPIConfig",
     ),
+    "BedrockMantleResponsesAPIConfig": (
+        ".llms.bedrock_mantle.responses.transformation",
+        "BedrockMantleResponsesAPIConfig",
+    ),
     "GoogleAIStudioInteractionsConfig": (
         ".llms.gemini.interactions.transformation",
         "GoogleAIStudioInteractionsConfig",

litellm/experimental_mcp_client/client.py

@@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
 
 import asyncio
 import base64
+import os
 from typing import (
     Any,
     Awaitable,
@@ -16,7 +17,6 @@ from typing import (
     TypeVar,
     Union,
 )
-
 import httpx
 from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
 from mcp.client.sse import sse_client
@@ -42,9 +42,8 @@ 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.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
 from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
 from litellm.types.llms.custom_http import VerifyTypes
 from litellm.types.mcp import (
@@ -67,7 +66,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 +90,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 +139,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 +171,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 +188,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 +203,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 +217,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 +229,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._get_safe_stdio_env(self.stdio_config.get("env")),
             )
             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 +255,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)
@@ -281,6 +274,54 @@ class MCPClient:
         )
         return transport_ctx, http_client
 
+    def _get_safe_stdio_env(
+        self, provided_env: Optional[Dict[str, str]]
+    ) -> Optional[Dict[str, str]]:
+        """
+        Return a safe environment for the stdio subprocess.
+
+        If provided_env is set, we use it as-is.
+        If provided_env is None, we return a minimal allowlist from the parent environment
+        to avoid leaking sensitive LiteLLM keys (OPENAI_API_KEY, etc.) to sub-processes.
+        """
+        if provided_env is not None:
+            return provided_env
+
+        # Minimal allowlist of safe/standard environment variables
+        safe_keys = {
+            "PATH",
+            "HOME",
+            "USER",
+            "LOGNAME",
+            "TMPDIR",
+            "TMP",
+            "TEMP",
+            "SHELL",
+            "LANG",
+            "LC_ALL",
+            # Node/Package manager caches
+            "NPM_CONFIG_CACHE",
+            "PNPM_HOME",
+            "XDG_CACHE_HOME",
+            "XDG_CONFIG_HOME",
+            "XDG_DATA_HOME",
+            # System info
+            "SYSTEMROOT",
+            "COMSPEC",
+            "PATHEXT",
+            "WINDIR",
+        }
+
+        safe_env = {}
+        for key in safe_keys:
+            if key in os.environ:
+                safe_env[key] = os.environ[key]
+
+        if "NPM_CONFIG_CACHE" not in safe_env:
+            safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
+
+        return safe_env
+
     async def _execute_session_operation(
         self,
         transport_ctx: Any,
@@ -288,13 +329,23 @@ 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 +402,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:
@@ -373,17 +423,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
@@ -400,17 +447,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,
@@ -458,7 +502,6 @@ 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(
@@ -491,7 +534,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:
@@ -521,7 +563,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(
@@ -532,14 +573,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=[
@@ -577,14 +616,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 []
 
@@ -617,7 +654,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(
@@ -628,14 +664,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]:
@@ -667,14 +701,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 []
 
@@ -709,14 +741,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 []
 
@@ -742,7 +772,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(
@@ -753,12 +782,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/litellm_core_utils/get_supported_openai_params.py

@@ -22,9 +22,11 @@ def get_supported_openai_params(  # noqa: PLR0915
     ```
 
     Args:
-        base_model: For Azure, the true underlying model (e.g. ``"azure/gpt-5.2"``)
-            when the deployment name differs. Used for model-type detection so that
-            non-standard deployment names route to the correct config.
+        base_model: An optional capability hint for deployments whose ``model``
+            label isn't recognized on its own (e.g. an Azure deployment name, or a
+            friendly Bedrock alias). It is additive: the result is the union of the
+            params supported by ``model`` and by ``base_model``, so a hint can only
+            add capabilities, never strip ones the real model already supports.
 
     Returns:
     - List if custom_llm_provider is mapped
@@ -52,7 +54,15 @@ def get_supported_openai_params(  # noqa: PLR0915
         provider_config = None
 
     if provider_config and request_type == "chat_completion":
-        return provider_config.get_supported_openai_params(model=base_model or model)
+        supported_params = provider_config.get_supported_openai_params(model=model)
+        if base_model and base_model != model:
+            base_model_params = provider_config.get_supported_openai_params(
+                model=base_model
+            )
+            supported_params = list(
+                dict.fromkeys([*supported_params, *base_model_params])
+            )
+        return supported_params
 
     if custom_llm_provider == "bedrock":
         return litellm.AmazonConverseConfig().get_supported_openai_params(model=model)

litellm/litellm_core_utils/streaming_handler.py

@@ -1149,6 +1149,32 @@ class CustomStreamWrapper:
             completion_obj: Dict[str, Any] = {"content": ""}
             from litellm.types.utils import GenericStreamingChunk as GChunk
 
+            if (
+                isinstance(chunk, ModelResponseStream)
+                and self.custom_llm_provider is not None
+                and self.custom_llm_provider in litellm._custom_providers
+            ):
+                _has_content = bool(
+                    chunk.choices
+                    and chunk.choices[0].delta is not None
+                    and (
+                        chunk.choices[0].delta.content
+                        or chunk.choices[0].delta.tool_calls
+                    )
+                )
+                if self.received_finish_reason is not None:
+                    if not _has_content:
+                        raise StopIteration
+                if chunk.choices and chunk.choices[0].finish_reason:
+                    self.received_finish_reason = chunk.choices[0].finish_reason
+                    if not _has_content:
+                        return None
+                    # Strip finish_reason from the content chunk so it appears
+                    # only on the trailing empty-delta chunk (OpenAI spec).
+                    # finish_reason_handler() will emit the proper terminal chunk.
+                    chunk.choices[0].finish_reason = None  # type: ignore[assignment]
+                return chunk
+
             if (
                 isinstance(chunk, dict)
                 and generic_chunk_has_all_required_fields(

litellm/llms/azure/image_edit/transformation.py

@@ -97,8 +97,15 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
             )
         original_url = httpx.URL(api_base)
 
-        # Extract api_version or use default
-        api_version = cast(Optional[str], litellm_params.get("api_version"))
+        # Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default.
+        # Mirrors the fallback chain used by the Azure chat path in common_utils.py,
+        # so callers that set a global / env api_version don't get an unversioned URL.
+        api_version = (
+            cast(Optional[str], litellm_params.get("api_version"))
+            or litellm.api_version
+            or get_secret_str("AZURE_API_VERSION")
+            or litellm.AZURE_DEFAULT_API_VERSION
+        )
 
         # Create a new dictionary with existing params
         query_params = dict(original_url.params)

litellm/llms/bedrock_mantle/responses/transformation.py

@@ -0,0 +1,81 @@
+"""
+Amazon Bedrock Mantle - Responses API backend.
+
+gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses`
+path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI
+Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
+only the endpoint URL and Bearer authentication.
+
+Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
+standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
+"""
+
+from typing import Optional
+
+from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1"
+
+# Checked longest/most-specific first so a full endpoint URL collapses to host
+# in one pass and the appended path never doubles.
+_BASE_SUFFIXES_TO_STRIP = (
+    "/openai/v1/responses",
+    "/v1/responses",
+    "/responses",
+    "/openai/v1",
+    "/v1",
+)
+
+
+class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
+    @property
+    def custom_llm_provider(self) -> LlmProviders:
+        return LlmProviders.BEDROCK_MANTLE
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        region = (
+            get_secret_str("BEDROCK_MANTLE_REGION")
+            or get_secret_str("AWS_REGION")
+            or BEDROCK_MANTLE_DEFAULT_REGION
+        )
+        base = (
+            api_base
+            or get_secret_str("BEDROCK_MANTLE_API_BASE")
+            or f"https://bedrock-mantle.{region}.api.aws"
+        )
+        base = base.rstrip("/")
+        for suffix in _BASE_SUFFIXES_TO_STRIP:
+            if base.endswith(suffix):
+                base = base[: -len(suffix)]
+                break
+        return f"{base}/openai/v1/responses"
+
+    def validate_environment(
+        self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]
+    ) -> dict:
+        litellm_params = litellm_params or GenericLiteLLMParams()
+        api_key = (
+            litellm_params.api_key
+            or get_secret_str("BEDROCK_MANTLE_API_KEY")
+            or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
+        )
+        if not api_key:
+            raise ValueError(
+                "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
+                "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
+            )
+        headers["Authorization"] = f"Bearer {api_key}"
+        return headers
+
+    def supports_native_file_search(self) -> bool:
+        return False
+
+    def supports_native_websocket(self) -> bool:
+        return False

litellm/llms/fireworks_ai/chat/transformation.py

@@ -170,11 +170,6 @@ class FireworksAIConfig(OpenAIGPTConfig):
                         is_response_format_supported=False,
                         enforce_tool_choice=False,  # tools and response_format are both set, don't enforce tool_choice
                     )
-                elif "json_schema" in value:
-                    optional_params["response_format"] = {
-                        "type": "json_object",
-                        "schema": value["json_schema"]["schema"],
-                    }
                 else:
                     optional_params["response_format"] = value
             elif param == "max_completion_tokens":

litellm/llms/gemini/videos/transformation.py

@@ -265,7 +265,11 @@ class GeminiVideoConfig(BaseVideoConfig):
         {
             "instances": [
                 {
-                    "prompt": "A cat playing with a ball of yarn"
+                    "prompt": "A cat playing with a ball of yarn",
+                    "image": {
+                        "bytesBase64Encoded": "...",
+                        "mimeType": "image/jpeg"
+                    }
                 }
             ],
             "parameters": {
@@ -275,13 +279,18 @@ class GeminiVideoConfig(BaseVideoConfig):
             }
         }
         """
-        instance = GeminiVideoGenerationInstance(prompt=prompt)
+        instance: GeminiVideoGenerationInstance = {"prompt": prompt}
 
         params_copy = video_create_optional_request_params.copy()
 
-        if "image" in params_copy and params_copy["image"] is not None:
-            image_data = _convert_image_to_gemini_format(params_copy["image"])
-            params_copy["image"] = image_data
+        if "image" in params_copy:
+            image = params_copy.pop("image")
+            if image is not None:
+                if isinstance(image, dict):
+                    image_data = image
+                else:
+                    image_data = _convert_image_to_gemini_format(image)
+                instance["image"] = image_data
 
         parameters = GeminiVideoGenerationParameters(**params_copy)
 

litellm/llms/huggingface/embedding/handler.py

@@ -239,7 +239,7 @@ class HuggingFaceEmbedding(BaseLLM):
         model_response.model = model
         input_tokens = 0
         for text in input:
-            input_tokens += len(encoding.encode(text))
+            input_tokens += len(encoding.encode(text, disallowed_special=()))
 
         setattr(
             model_response,

litellm/llms/snowflake/utils.py

@@ -25,6 +25,7 @@ class SnowflakeBaseConfig:
             "temperature",
             "max_tokens",
             "top_p",
+            "stream",
             "response_format",
             "tools",
             "tool_choice",

litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py

@@ -337,6 +337,7 @@ class ContextCachingEndpoints(VertexBase):
             return messages, optional_params, None
 
         tools = optional_params.pop("tools", None)
+        tool_choice = optional_params.pop("tool_choice", None)
 
         ## AUTHORIZATION ##
         token, url = self._get_token_and_url_context_caching(
@@ -371,7 +372,7 @@ class ContextCachingEndpoints(VertexBase):
 
         ## CHECK IF CACHED ALREADY
         generated_cache_key = local_cache_obj.get_cache_key(
-            messages=cached_messages, tools=tools, model=model
+            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
         )
         google_cache_name = self.check_cache(
             cache_key=generated_cache_key,
@@ -402,6 +403,8 @@ class ContextCachingEndpoints(VertexBase):
         )
 
         cached_content_request_body["tools"] = tools
+        if tool_choice is not None:
+            cached_content_request_body["toolConfig"] = tool_choice
 
         ## LOGGING
         logging_obj.pre_call(
@@ -487,6 +490,7 @@ class ContextCachingEndpoints(VertexBase):
             return messages, optional_params, None
 
         tools = optional_params.pop("tools", None)
+        tool_choice = optional_params.pop("tool_choice", None)
 
         ## AUTHORIZATION ##
         token, url = self._get_token_and_url_context_caching(
@@ -518,7 +522,7 @@ class ContextCachingEndpoints(VertexBase):
 
         ## CHECK IF CACHED ALREADY
         generated_cache_key = local_cache_obj.get_cache_key(
-            messages=cached_messages, tools=tools, model=model
+            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
         )
         google_cache_name = await self.async_check_cache(
             cache_key=generated_cache_key,
@@ -550,6 +554,8 @@ class ContextCachingEndpoints(VertexBase):
         )
 
         cached_content_request_body["tools"] = tools
+        if tool_choice is not None:
+            cached_content_request_body["toolConfig"] = tool_choice
 
         ## LOGGING
         logging_obj.pre_call(

litellm/main.py

@@ -1322,7 +1322,9 @@ def completion(  # type: ignore # noqa: PLR0915
     preset_cache_key = kwargs.get("preset_cache_key", None)
     hf_model_name = kwargs.get("hf_model_name", None)
     supports_system_message = kwargs.get("supports_system_message", None)
-    base_model = kwargs.get("base_model", None)
+    base_model = kwargs.get("base_model", None) or (
+        model_info.get("base_model") if isinstance(model_info, dict) else None
+    )
     ### DISABLE FLAGS ###
     disable_add_transform_inline_image_block = kwargs.get(
         "disable_add_transform_inline_image_block", None
@@ -1534,11 +1536,7 @@ def completion(  # type: ignore # noqa: PLR0915
             "logit_bias": logit_bias,
             "user": user,
             # params to identify the model
-            "model": (
-                model_info.get("base_model")
-                if isinstance(model_info, dict) and model_info.get("base_model")
-                else model
-            ),
+            "model": model,
             "custom_llm_provider": custom_llm_provider,
             "response_format": response_format,
             "seed": seed,

litellm/model_prices_and_context_window_backup.json

@@ -41223,6 +41223,44 @@
         "supports_response_schema": true,
         "supports_tool_choice": true
     },
+    "bedrock_mantle/openai.gpt-5.5": {
+        "input_cost_per_token": 5.5e-06,
+        "cache_read_input_token_cost": 5.5e-07,
+        "output_cost_per_token": 3.3e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/openai.gpt-5.4": {
+        "input_cost_per_token": 2.75e-06,
+        "cache_read_input_token_cost": 2.75e-07,
+        "output_cost_per_token": 1.65e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "volcengine/doubao-seed-2-0-pro-260215": {
         "litellm_provider": "volcengine",
         "max_input_tokens": 256000,

litellm/proxy/_experimental/mcp_server/elicitation_handler.py

@@ -0,0 +1,163 @@
+"""
+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 — pass an empty schema
+            # since elicit() requires requestedSchema as a positional arg.
+            verbose_logger.info(
+                "MCP elicitation: relaying generic elicitation to downstream"
+            )
+            result = await downstream_session.elicit(
+                message=getattr(params, "message", ""),
+                requestedSchema=getattr(params, "requestedSchema", {}),
+            )
+        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

@@ -49,6 +49,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
     MCPRequestHandler,
 )
 from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
+from litellm.proxy._experimental.mcp_server.elicitation_handler import (
+    MCP_ELICITATION_AVAILABLE,
+)
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    MCP_SAMPLING_AVAILABLE,
+)
 from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth
 from litellm.proxy._experimental.mcp_server.utils import (
     MCP_TOOL_PREFIX_SEPARATOR,
@@ -347,6 +353,82 @@ def _deserialize_json_list(data: Any) -> Optional[List[Dict[str, Any]]]:
     ]
 
 
+def _create_sampling_callback(user_api_key_auth: Optional[Any] = None):
+    """
+    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().
+    """
+    if not MCP_SAMPLING_AVAILABLE:
+        return None
+
+    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,
+        )
+
+        auth_context = get_active_auth_context()
+        resolved_auth = user_api_key_auth or (
+            auth_context.user_api_key_auth if auth_context else None
+        )
+        # Forward original HTTP headers and client IP so that
+        # header-dependent guardrails, tag-based routing, trace
+        # correlation, and forward_llm_provider_auth_headers work
+        # correctly for sampling sub-calls.
+        _raw_headers = getattr(auth_context, "raw_headers", None)
+        _client_ip = getattr(auth_context, "client_ip", None)
+
+        return await handle_sampling_create_message(
+            context=context,
+            params=params,
+            default_model=getattr(litellm, "default_mcp_sampling_model", None),
+            user_api_key_auth=resolved_auth,
+            raw_headers=_raw_headers,
+            client_ip=_client_ip,
+        )
+
+    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.
+    """
+    if not MCP_ELICITATION_AVAILABLE:
+        return None
+
+    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-[^}]+)\}$")
 
@@ -667,6 +749,8 @@ class MCPServerManager:
                     "subject_token_type",
                     "urn:ietf:params:oauth:token-type:access_token",
                 ),
+                allow_sampling=bool(server_config.get("allow_sampling", False)),
+                allow_elicitation=bool(server_config.get("allow_elicitation", False)),
             )
             self._assign_unique_short_prefix(new_server)
             _warn_internal_delegate_pkce_if_applicable(new_server, source="config")
@@ -766,8 +850,7 @@ class MCPServerManager:
             )
 
             verbose_logger.debug(
-                f"Using headers for OpenAPI tools (excluding sensitive values): "
-                f"{list(headers.keys())}"
+                f"Using headers for OpenAPI tools (excluding sensitive values): {list(headers.keys())}"
             )
 
             # Extract and register tools from OpenAPI paths
@@ -1766,6 +1849,7 @@ class MCPServerManager:
         extra_headers: Optional[Dict[str, str]] = None,
         stdio_env: Optional[Dict[str, str]] = None,
         subject_token: Optional[str] = None,
+        user_api_key_auth: Optional[UserAPIKeyAuth] = None,
     ) -> MCPClient:
         """
         Create an MCPClient instance for the given server.
@@ -1782,6 +1866,7 @@ class MCPServerManager:
             extra_headers: Additional headers to forward.
             stdio_env: Environment variables for stdio transport.
             subject_token: Optional user JWT for token exchange (OBO) flow.
+            user_api_key_auth: Optional auth context for sampling callbacks.
 
         Returns:
             Configured MCP client instance.
@@ -1792,23 +1877,44 @@ class MCPServerManager:
 
         transport = server.transport or MCPTransport.sse
 
+        # Create sampling and elicitation callbacks for this client
+        sampling_cb = (
+            _create_sampling_callback(user_api_key_auth=user_api_key_auth)
+            if server.allow_sampling
+            else None
+        )
+        elicitation_cb = (
+            _create_elicitation_callback() if server.allow_elicitation else None
+        )
+
         # 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 is not None 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/.com suffix for Windows compatibility
+                base_command_no_ext = base_command.lower()
+                for ext in [".exe", ".cmd", ".bat", ".com"]:
+                    if base_command.lower().endswith(ext):
+                        base_command_no_ext = base_command[: -len(ext)].lower()
+                        break
+                if (
+                    base_command.lower() 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)}). "
@@ -1831,6 +1937,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
@@ -1857,6 +1965,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(
@@ -1947,6 +2057,7 @@ class MCPServerManager:
                 mcp_auth_header=mcp_auth_header,
                 extra_headers=extra_headers,
                 stdio_env=stdio_env,
+                user_api_key_auth=user_api_key_auth,
             )
 
             ## HANDLE OPENAPI TOOLS
@@ -3316,6 +3427,7 @@ class MCPServerManager:
             extra_headers=extra_headers,
             stdio_env=stdio_env,
             subject_token=subject_token,
+            user_api_key_auth=user_api_key_auth,
         )
 
         call_tool_params = MCPCallToolRequestParams(
@@ -3546,7 +3658,6 @@ class MCPServerManager:
                 )
             )
         else:
-            # For regular MCP servers, use the MCP client
             return await self._call_regular_mcp_tool(
                 mcp_server=mcp_server,
                 original_tool_name=name,

litellm/proxy/_experimental/mcp_server/server.py

@@ -6,6 +6,7 @@ LiteLLM MCP Server Routes
 
 import asyncio
 import contextlib
+import contextvars
 import hashlib
 import json
 import time
@@ -126,6 +127,18 @@ try:
         GetPromptResult,
         ResourceTemplate,
         TextResourceContents,
+        Tool,
+    )
+    from mcp.server.session import ServerSession as _McpServerSession
+    import weakref
+
+    # Robust auth lookup keyed by session_object.
+    _session_obj_auth_storage: (
+        "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]"
+    ) = weakref.WeakKeyDictionary()
+
+    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}")
@@ -161,6 +174,60 @@ def _mcp_session_id_from_headers(
     return None
 
 
+def _jsonrpc_text_has_top_level_method(text: str) -> bool:
+    """Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at
+    the root object's top level.
+
+    Used to tell a request/notification (carries ``method``) apart from a
+    response (carries ``result``/``error`` and no top-level ``method``). A
+    response payload can itself nest a ``method`` field, so only keys at the
+    root object's depth are inspected rather than searching the whole string.
+    Returns ``True`` only when a top-level ``method`` key is positively found;
+    truncation that hides it yields ``False``.
+    """
+    depth = 0
+    in_string = False
+    escaped = False
+    in_object: List[bool] = []
+    reading_key = False
+    expect_key = False
+    key_chars: List[str] = []
+    for ch in text:
+        if in_string:
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == '"':
+                in_string = False
+                if reading_key and depth == 1 and "".join(key_chars) == "method":
+                    return True
+            elif reading_key:
+                key_chars.append(ch)
+            continue
+        if ch == '"':
+            in_string = True
+            reading_key = expect_key and depth >= 1 and in_object[-1]
+            key_chars = []
+            expect_key = False
+        elif ch == "{" or ch == "[":
+            depth += 1
+            in_object.append(ch == "{")
+            expect_key = ch == "{"
+        elif ch == "}" or ch == "]":
+            if in_object:
+                in_object.pop()
+            depth -= 1
+            if depth <= 0:
+                break
+            expect_key = False
+        elif ch == ",":
+            expect_key = bool(in_object) and in_object[-1]
+        elif ch == ":":
+            expect_key = False
+    return False
+
+
 if MCP_AVAILABLE:
     from mcp.server import Server
     from mcp.server.lowlevel.server import NotificationOptions
@@ -484,10 +551,18 @@ if MCP_AVAILABLE:
     ########################################################
 
     @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)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
         try:
             # Get user authentication from context variable
             (
@@ -498,7 +573,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}"
             )
@@ -529,162 +604,188 @@ if MCP_AVAILABLE:
             # Return empty list instead of failing completely
             # This prevents the HTTP stream from failing and allows the client to get a response
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.call_tool()
-    async def mcp_server_tool_call(
-        name: str, arguments: Optional[Dict[str, Any]]
+    async def mcp_server_tool_call(  # noqa: PLR0915
+        name: str, arguments: Dict[str, Any] | None
     ) -> CallToolResult:
         """
         Call a specific tool with the provided arguments
-
         Args:
             name (str): Name of the tool to call
             arguments (Dict[str, Any] | None): Arguments to pass to the tool
-
         Returns:
             List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results
-
         Raises:
             HTTPException: If tool not found or arguments missing
         """
         from 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
+        from mcp.types import CallToolResult
+        from mcp.server.lowlevel.server import request_ctx
 
-        # Validate arguments
-        (
-            user_api_key_auth,
-            mcp_auth_header,
-            mcp_servers,
-            mcp_server_auth_headers,
-            oauth2_headers,
-            raw_headers,
-            _client_ip,
-        ) = get_auth_context()
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
 
-        verbose_logger.debug(
-            f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
-        )
-        host_progress_callback = None
         try:
-            host_ctx = server.request_context
-            if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta:
-                host_token = getattr(host_ctx.meta, "progressToken", None)
-                if host_token and hasattr(host_ctx, "session") and host_ctx.session:
-                    host_session = host_ctx.session
-
-                    async def forward_progress(progress: float, total: Optional[float]):
-                        """Forward progress notifications from external MCP to Host"""
-                        try:
-                            await host_session.send_progress_notification(
-                                progress_token=host_token,
-                                progress=progress,
-                                total=total,
-                            )
-                            verbose_logger.debug(
-                                f"Forwarded progress {progress}/{total} to Host"
-                            )
-                        except Exception as e:
-                            verbose_logger.error(
-                                f"Failed to forward progress to Host: {e}"
-                            )
-
-                    host_progress_callback = forward_progress
-                    verbose_logger.debug(
-                        f"Host progressToken captured: {host_token[:8]}..."
-                    )
-        except Exception as e:
-            verbose_logger.warning(f"Could not capture host progress context: {e}")
-        try:
-            # Create a body date for logging
-            body_data = {"name": name, "arguments": arguments}
-            # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A)
-            chain_id = get_chain_id_from_headers(raw_headers)
-            if chain_id:
-                body_data["litellm_trace_id"] = chain_id
-                body_data["litellm_session_id"] = chain_id
-
-            request = Request(
-                scope={
-                    "type": "http",
-                    "method": "POST",
-                    "path": "/mcp/tools/call",
-                    "headers": [(b"content-type", b"application/json")],
-                }
+            # Validate arguments
+            (
+                user_api_key_auth,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+                _client_ip,
+            ) = await get_or_extract_auth_context()
+            verbose_logger.debug(
+                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')}"
             )
-            if user_api_key_auth is not None:
-                data = await add_litellm_data_to_request(
-                    data=body_data,
-                    request=request,
-                    user_api_key_dict=user_api_key_auth,
-                    proxy_config=proxy_config,
+
+            verbose_logger.debug(
+                f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
+            )
+            host_progress_callback = None
+            try:
+                host_ctx = server.request_context
+                if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta:
+                    host_token = getattr(host_ctx.meta, "progressToken", None)
+                    if host_token and hasattr(host_ctx, "session") and host_ctx.session:
+                        host_session = host_ctx.session
+
+                        async def forward_progress(
+                            progress: float, total: Optional[float]
+                        ):
+                            """Forward progress notifications from external MCP to Host"""
+                            try:
+                                await host_session.send_progress_notification(
+                                    progress_token=host_token,
+                                    progress=progress,
+                                    total=total,
+                                )
+                                verbose_logger.debug(
+                                    f"Forwarded progress {progress}/{total} to Host"
+                                )
+                            except Exception as e:
+                                verbose_logger.error(
+                                    f"Failed to forward progress to Host: {e}"
+                                )
+
+                        host_progress_callback = forward_progress
+                        verbose_logger.debug(
+                            f"Host progressToken captured: {host_token[:8]}..."
+                        )
+            except Exception as e:
+                verbose_logger.warning(f"Could not capture host progress context: {e}")
+            try:
+                # Create a body date for logging
+                body_data = {"name": name, "arguments": arguments}
+                # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A)
+                chain_id = get_chain_id_from_headers(raw_headers)
+                if chain_id:
+                    body_data["litellm_trace_id"] = chain_id
+                    body_data["litellm_session_id"] = chain_id
+
+                request = Request(
+                    scope={
+                        "type": "http",
+                        "method": "POST",
+                        "path": "/mcp/tools/call",
+                        "headers": [(b"content-type", b"application/json")],
+                    }
                 )
-            else:
-                data = body_data
-
-            response = await call_mcp_tool(
-                user_api_key_auth=user_api_key_auth,
-                mcp_auth_header=mcp_auth_header,
-                mcp_servers=mcp_servers,
-                mcp_server_auth_headers=mcp_server_auth_headers,
-                oauth2_headers=oauth2_headers,
-                raw_headers=raw_headers,
-                host_progress_callback=host_progress_callback,
-                **data,  # for logging
-            )
-        except MCPMissingUserEnvVarsError as e:
-            verbose_logger.info(
-                "MCP mcp_server_tool_call missing per-user env vars: server_id=%s missing=%s",
-                e.server_id,
-                e.missing,
-            )
-            return CallToolResult(
-                content=[TextContent(text=str(e), type="text")],
-                isError=True,
-            )
-        except BlockedPiiEntityError as e:
-            verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}")
-            return CallToolResult(
-                content=[
-                    TextContent(
-                        text=f"Error: Blocked PII entity detected - {str(e)}",
-                        type="text",
+                if user_api_key_auth is not None:
+                    data = await add_litellm_data_to_request(
+                        data=body_data,
+                        request=request,
+                        user_api_key_dict=user_api_key_auth,
+                        proxy_config=proxy_config,
                     )
-                ],
-                isError=True,
-            )
-        except GuardrailRaisedException as e:
-            verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}")
-            return CallToolResult(
-                content=[
-                    TextContent(
-                        text=f"Error: Guardrail violation - {str(e)}", type="text"
-                    )
-                ],
-                isError=True,
-            )
-        except HTTPException as e:
-            verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
-            return CallToolResult(
-                content=[TextContent(text=f"Error: {str(e.detail)}", type="text")],
-                isError=True,
-            )
-        except Exception as e:
-            verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}")
-            return CallToolResult(
-                content=[TextContent(text=f"Error: {str(e)}", type="text")],
-                isError=True,
-            )
+                else:
+                    data = body_data
 
-        return response
+                response = await call_mcp_tool(
+                    user_api_key_auth=user_api_key_auth,
+                    mcp_auth_header=mcp_auth_header,
+                    mcp_servers=mcp_servers,
+                    mcp_server_auth_headers=mcp_server_auth_headers,
+                    oauth2_headers=oauth2_headers,
+                    raw_headers=raw_headers,
+                    host_progress_callback=host_progress_callback,
+                    **data,  # for logging
+                )
+            except MCPMissingUserEnvVarsError as e:
+                verbose_logger.info(
+                    "MCP mcp_server_tool_call missing per-user env vars: server_id=%s missing=%s",
+                    e.server_id,
+                    e.missing,
+                )
+                return CallToolResult(
+                    content=[TextContent(text=str(e), type="text")],
+                    isError=True,
+                )
+            except BlockedPiiEntityError as e:
+                verbose_logger.error(
+                    f"BlockedPiiEntityError in MCP tool call: {str(e)}"
+                )
+                return CallToolResult(
+                    content=[
+                        TextContent(
+                            text=f"Error: Blocked PII entity detected - {str(e)}",
+                            type="text",
+                        )
+                    ],
+                    isError=True,
+                )
+            except GuardrailRaisedException as e:
+                verbose_logger.error(
+                    f"GuardrailRaisedException in MCP tool call: {str(e)}"
+                )
+                return CallToolResult(
+                    content=[
+                        TextContent(
+                            text=f"Error: Guardrail violation - {str(e)}", type="text"
+                        )
+                    ],
+                    isError=True,
+                )
+            except HTTPException as e:
+                verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
+                return CallToolResult(
+                    content=[TextContent(text=f"Error: {str(e.detail)}", type="text")],
+                    isError=True,
+                )
+            except Exception as e:
+                verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}")
+                return CallToolResult(
+                    content=[TextContent(text=f"Error: {str(e)}", type="text")],
+                    isError=True,
+                )
+
+            return response
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.list_prompts()
     async def list_prompts() -> List[Prompt]:
         """
         List all available prompts
         """
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
         try:
             # Get user authentication from context variable
             (
@@ -695,7 +796,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}"
             )
@@ -724,6 +825,9 @@ if MCP_AVAILABLE:
             # Return empty list instead of failing completely
             # This prevents the HTTP stream from failing and allows the client to get a response
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.get_prompt()
     async def get_prompt(
@@ -741,33 +845,13 @@ if MCP_AVAILABLE:
         """
 
         # Validate arguments
-        (
-            user_api_key_auth,
-            mcp_auth_header,
-            mcp_servers,
-            mcp_server_auth_headers,
-            oauth2_headers,
-            raw_headers,
-            _client_ip,
-        ) = get_auth_context()
+        from mcp.server.lowlevel.server import request_ctx
 
-        verbose_logger.debug(
-            f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
-        )
-        return await mcp_get_prompt(
-            name=name,
-            arguments=arguments,
-            user_api_key_auth=user_api_key_auth,
-            mcp_auth_header=mcp_auth_header,
-            mcp_servers=mcp_servers,
-            mcp_server_auth_headers=mcp_server_auth_headers,
-            oauth2_headers=oauth2_headers,
-            raw_headers=raw_headers,
-        )
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
 
-    @server.list_resources()
-    async def list_resources() -> List[Resource]:
-        """List all available resources."""
         try:
             (
                 user_api_key_auth,
@@ -777,7 +861,45 @@ 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}"
+            )
+            return await mcp_get_prompt(
+                name=name,
+                arguments=arguments,
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=mcp_auth_header,
+                mcp_servers=mcp_servers,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                oauth2_headers=oauth2_headers,
+                raw_headers=raw_headers,
+            )
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
+
+    @server.list_resources()
+    async def list_resources() -> List[Resource]:
+        """List all available resources."""
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
+        try:
+            (
+                user_api_key_auth,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+                _client_ip,
+            ) = await get_or_extract_auth_context()
             verbose_logger.debug(
                 f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}"
             )
@@ -803,10 +925,20 @@ if MCP_AVAILABLE:
         except Exception as e:
             verbose_logger.exception(f"Error in list_resources endpoint: {str(e)}")
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.list_resource_templates()
     async def list_resource_templates() -> List[ResourceTemplate]:
         """List all available resource templates."""
+        from mcp.server.lowlevel.server import request_ctx
+
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
+
         try:
             (
                 user_api_key_auth,
@@ -816,7 +948,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}"
             )
@@ -836,8 +968,7 @@ if MCP_AVAILABLE:
                 raw_headers=raw_headers,
             )
             verbose_logger.info(
-                "MCP list_resource_templates - Successfully returned "
-                f"{len(resource_templates)} resource templates"
+                f"MCP list_resource_templates - Successfully returned {len(resource_templates)} resource templates"
             )
             return resource_templates
         except Exception as e:
@@ -845,30 +976,44 @@ if MCP_AVAILABLE:
                 f"Error in list_resource_templates endpoint: {str(e)}"
             )
             return []
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     @server.read_resource()
     async def read_resource(url: AnyUrl) -> list[ReadResourceContents]:
-        (
-            user_api_key_auth,
-            mcp_auth_header,
-            mcp_servers,
-            mcp_server_auth_headers,
-            oauth2_headers,
-            raw_headers,
-            _client_ip,
-        ) = get_auth_context()
+        from mcp.server.lowlevel.server import request_ctx
 
-        read_resource_result = await mcp_read_resource(
-            url=url,
-            user_api_key_auth=user_api_key_auth,
-            mcp_auth_header=mcp_auth_header,
-            mcp_servers=mcp_servers,
-            mcp_server_auth_headers=mcp_server_auth_headers,
-            oauth2_headers=oauth2_headers,
-            raw_headers=raw_headers,
-        )
+        req_ctx = request_ctx.get(None)
+        _session_reset_token = None
+        if req_ctx:
+            _session_reset_token = active_mcp_session_var.set(req_ctx.session)
 
-        return _normalize_resource_contents(read_resource_result.contents)
+        try:
+            (
+                user_api_key_auth,
+                mcp_auth_header,
+                mcp_servers,
+                mcp_server_auth_headers,
+                oauth2_headers,
+                raw_headers,
+                _client_ip,
+            ) = await get_or_extract_auth_context()
+
+            read_resource_result = await mcp_read_resource(
+                url=url,
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=mcp_auth_header,
+                mcp_servers=mcp_servers,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                oauth2_headers=oauth2_headers,
+                raw_headers=raw_headers,
+            )
+
+            return _normalize_resource_contents(read_resource_result.contents)
+        finally:
+            if _session_reset_token is not None:
+                active_mcp_session_var.reset(_session_reset_token)
 
     ########################################################
     ############ End of MCP Server Routes ##################
@@ -1191,8 +1336,7 @@ if MCP_AVAILABLE:
                 cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
                 if cached_token is not None:
                     verbose_logger.debug(
-                        "_get_user_oauth_extra_headers_from_db: Redis hit for "
-                        "user=%s server=%s",
+                        "_get_user_oauth_extra_headers_from_db: Redis hit for user=%s server=%s",
                         user_id,
                         server_id,
                     )
@@ -1218,8 +1362,7 @@ if MCP_AVAILABLE:
 
             if is_oauth_credential_expired(cred):
                 verbose_logger.debug(
-                    "_get_user_oauth_extra_headers_from_db: token expired for "
-                    "user=%s server=%s — attempting refresh",
+                    "_get_user_oauth_extra_headers_from_db: token expired for user=%s server=%s — attempting refresh",
                     user_id,
                     server_id,
                 )
@@ -1241,8 +1384,7 @@ if MCP_AVAILABLE:
                         )
                     except Exception as refresh_exc:
                         verbose_logger.warning(
-                            "_get_user_oauth_extra_headers_from_db: refresh failed "
-                            "for user=%s server=%s: %s",
+                            "_get_user_oauth_extra_headers_from_db: refresh failed for user=%s server=%s: %s",
                             user_id,
                             server_id,
                             refresh_exc,
@@ -1286,8 +1428,7 @@ if MCP_AVAILABLE:
             return {"Authorization": f"Bearer {access_token}"}
         except Exception as e:
             verbose_logger.warning(
-                "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
-                "user=%s server=%s: %s",
+                "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for user=%s server=%s: %s",
                 user_id,
                 server_id,
                 e,
@@ -2496,7 +2637,7 @@ if MCP_AVAILABLE:
                 arguments=arguments or {},
                 server_name=server_name or mcp_server.name,
                 user_api_key_auth=user_api_key_auth,
-                proxy_logging_obj=proxy_logging_obj,
+                proxy_logging_obj=proxy_logging_obj,  # type: ignore[arg-type]
                 server=mcp_server,
                 raw_headers=raw_headers,
             )
@@ -2755,8 +2896,7 @@ if MCP_AVAILABLE:
             raise HTTPException(
                 status_code=400,
                 detail=(
-                    "Multiple MCP servers configured; read_resource currently "
-                    "supports exactly one allowed server."
+                    "Multiple MCP servers configured; read_resource currently supports exactly one allowed server."
                 ),
             )
 
@@ -3135,8 +3275,7 @@ if MCP_AVAILABLE:
                 return False
         except Exception:
             verbose_logger.debug(
-                "Unable to inspect active MCP sessions for '%s'. "
-                "Deferring to session manager.",
+                "Unable to inspect active MCP sessions for '%s'. Deferring to session manager.",
                 _session_id,
             )
             return False
@@ -3147,8 +3286,7 @@ if MCP_AVAILABLE:
         if method == "DELETE":
             _remove_stateful_session_tracking(_session_id)
             verbose_logger.info(
-                "DELETE request for non-existent MCP session '%s'. "
-                "Returning success (idempotent DELETE).",
+                "DELETE request for non-existent MCP session '%s'. Returning success (idempotent DELETE).",
                 _session_id,
             )
             success_response = JSONResponse(
@@ -3626,6 +3764,7 @@ if MCP_AVAILABLE:
                     return
                 session_id = _get_session_id_from_scope(scope)
 
+            body = b""
             if scope.get("method") == "POST":
                 consumed_messages, body = await _read_request_body_for_routing(receive)
                 is_initialize = _is_initialize_request(body)
@@ -3650,8 +3789,7 @@ if MCP_AVAILABLE:
                 )
                 if not await _enforce_stateful_session_cap_for_owner(request_owner):
                     verbose_logger.warning(
-                        "Rejecting MCP initialize: caller already holds the maximum "
-                        "number of active stateful sessions."
+                        "Rejecting MCP initialize: caller already holds the maximum number of active stateful sessions."
                     )
                     too_many_response = JSONResponse(
                         status_code=429,
@@ -3683,9 +3821,56 @@ if MCP_AVAILABLE:
             # POST/DELETE are the methods that actually mutate the shared
             # auth context, so serializing those is sufficient for the
             # clobbering race between concurrent JSON-RPC calls.
-            session_lock: Optional[asyncio.Lock] = None
+            #
+            # Also skip the lock for JSON-RPC *responses* (POSTs that carry
+            # a ``result`` or ``error`` but no ``method``). These are replies
+            # to server-initiated requests such as ``elicitation/create`` or
+            # ``sampling/createMessage``. The in-flight tool-call POST that
+            # triggered the server request already holds the session lock, so
+            # trying to acquire it again for the response POST would deadlock.
+            is_jsonrpc_response = False
             request_method = (scope.get("method") or "").upper()
-            if use_stateful and session_id and request_method in ("POST", "DELETE"):
+            if body and request_method == "POST":
+                try:
+                    _peeked = json.loads(body)
+                    if (
+                        isinstance(_peeked, dict)
+                        and _peeked.get("jsonrpc") == "2.0"
+                        and "id" in _peeked
+                        and "method" not in _peeked
+                        and ("result" in _peeked or "error" in _peeked)
+                    ):
+                        is_jsonrpc_response = True
+                        verbose_logger.debug(
+                            "MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
+                            _peeked.get("id"),
+                        )
+                except (json.JSONDecodeError, TypeError):
+                    # Peek cap truncated the body, so it can't be fully parsed.
+                    # Scan the top-level keys (depth-aware) instead of a flat
+                    # substring search: a response's result payload may nest a
+                    # "method" field, and misreading that would acquire the lock
+                    # and deadlock the in-flight tool call awaiting this
+                    # response. A false skip is harmless; a false acquire is not.
+                    _body_str = body.decode("utf-8", errors="replace")
+                    if (
+                        '"jsonrpc"' in _body_str
+                        and ('"result"' in _body_str or '"error"' in _body_str)
+                        and not _jsonrpc_text_has_top_level_method(_body_str)
+                    ):
+                        is_jsonrpc_response = True
+                        verbose_logger.debug(
+                            "MCP: detected truncated JSON-RPC response POST via "
+                            "top-level key scan, skipping session lock to avoid deadlock"
+                        )
+
+            session_lock: Optional[asyncio.Lock] = None
+            if (
+                use_stateful
+                and session_id
+                and request_method in ("POST", "DELETE")
+                and not is_jsonrpc_response
+            ):
                 session_lock = _stateful_session_locks.setdefault(
                     session_id, asyncio.Lock()
                 )
@@ -4110,6 +4295,119 @@ if MCP_AVAILABLE:
             )
         return None, None, None, None, None, None, None
 
+    def _get_current_session():
+        try:
+            from mcp.server.lowlevel.server import request_ctx
+
+            return request_ctx.get().session
+        except (LookupError, ImportError):
+            return None
+
+    def _cache_auth_context_lazily():
+        session = _get_current_session()
+        if session is None:
+            return
+        try:
+            if session in _session_obj_auth_storage:
+                return
+        except TypeError:
+            verbose_logger.debug(
+                "_cache_auth_context_lazily: session object is unhashable (type=%s), cannot cache auth context",
+                type(session).__name__,
+            )
+            return
+
+        auth = auth_context_var.get()
+        if auth and isinstance(auth, MCPAuthenticatedUser):
+            try:
+                _session_obj_auth_storage[session] = auth
+            except TypeError:
+                verbose_logger.debug(
+                    "_cache_auth_context_lazily: could not store auth via "
+                    "session identity — session object is unhashable"
+                )
+
+    def _recover_auth_from_session() -> Optional[MCPAuthenticatedUser]:
+        session = _get_current_session()
+        if session is None:
+            return None
+
+        stored: Optional[MCPAuthenticatedUser] = None
+        try:
+            stored = _session_obj_auth_storage.get(session)
+        except TypeError:
+            verbose_logger.debug(
+                "_recover_auth_from_session: session object is unhashable "
+                "(type=%s), skipping _session_obj_auth_storage lookup",
+                type(session).__name__,
+            )
+
+        return stored
+
+    async def get_or_extract_auth_context() -> Tuple[
+        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 session
+        storage (which survives cross-task boundaries in the MCP SDK).
+        """
+        (
+            user_api_key_auth,
+            mcp_auth_header,
+            mcp_servers,
+            mcp_server_auth_headers,
+            oauth2_headers,
+            raw_headers,
+            _client_ip,
+        ) = get_auth_context()
+
+        if user_api_key_auth is not None:
+            _cache_auth_context_lazily()
+        else:
+            stored = _recover_auth_from_session()
+
+            if stored:
+                user_api_key_auth = stored.user_api_key_auth
+                mcp_auth_header = stored.mcp_auth_header
+                mcp_servers = stored.mcp_servers
+                mcp_server_auth_headers = stored.mcp_server_auth_headers
+                oauth2_headers = stored.oauth2_headers
+                raw_headers = stored.raw_headers
+                _client_ip = stored.client_ip
+        return (
+            user_api_key_auth,
+            mcp_auth_header,
+            mcp_servers,
+            mcp_server_auth_headers,
+            oauth2_headers,
+            raw_headers,
+            _client_ip,
+        )
+
+    def get_active_mcp_session() -> Optional[_McpServerSession]:
+        """Return the active MCP session captured during handler execution."""
+        session = active_mcp_session_var.get()
+        if session is not None:
+            return session
+        return _get_current_session()
+
+    def get_active_auth_context() -> Optional[MCPAuthenticatedUser]:
+        """Return auth context from ContextVar or session storage."""
+        auth = auth_context_var.get()
+        if auth and isinstance(auth, MCPAuthenticatedUser):
+            return auth
+
+        stored = _recover_auth_from_session()
+        if stored is not None:
+            return stored
+        return None
+
     ########################################################
     ############ End of Auth Context Functions #############
     ########################################################

litellm/proxy/caching_routes.py

@@ -60,11 +60,20 @@ async def cache_ping():
     """
     litellm_cache_params: Dict[str, Any] = {}
     cleaned_cache_params: Dict[str, Any] = {}
+    if litellm.cache is None:
+        raise ProxyException(
+            message=safe_dumps(
+                {
+                    "message": "Cache not initialized. litellm.cache is None",
+                    "litellm_cache_params": "{}",
+                    "health_check_cache_params": "{}",
+                }
+            ),
+            type=ProxyErrorTypes.cache_ping_error,
+            param="cache_ping",
+            code=503,
+        )
     try:
-        if litellm.cache is None:
-            raise HTTPException(
-                status_code=503, detail="Cache not initialized. litellm.cache is None"
-            )
         litellm_cache_params = masker.mask_dict(vars(litellm.cache))
         # remove field that might reference itself
         litellm_cache_params.pop("cache", None)
@@ -97,14 +106,14 @@ async def cache_ping():
                 cache_type=str(litellm.cache.type),
                 litellm_cache_params=safe_dumps(litellm_cache_params),
             )
-    except Exception as e:
-        import traceback
-
+    except HTTPException:
+        raise
+    except Exception:
+        verbose_proxy_logger.exception("Cache health check failed")
         error_message = {
-            "message": f"Service Unhealthy ({str(e)})",
+            "message": "Service Unhealthy",
             "litellm_cache_params": safe_dumps(litellm_cache_params),
             "health_check_cache_params": safe_dumps(cleaned_cache_params),
-            "traceback": traceback.format_exc(),
         }
         raise ProxyException(
             message=safe_dumps(error_message),

litellm/proxy/guardrails/guardrail_hooks/tool_permission.py

@@ -16,7 +16,7 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.callback_utils import (
     add_guardrail_to_applied_guardrails_header,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
 from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
     PermissionError,
     ToolPermissionRule,
@@ -60,53 +60,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         super().__init__(**kwargs)
 
-        self.rules: List[ToolPermissionRule] = []
-        self._compiled_rule_patterns: Dict[str, Dict[str, re.Pattern]] = {}
-        self._compiled_rule_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {}
-        if rules:
-            for rule_item in rules:
-                if isinstance(rule_item, ToolPermissionRule):
-                    rule = rule_item
-                else:
-                    rule = ToolPermissionRule(**rule_item)
-                self.rules.append(rule)
-
-                compiled_target_patterns: Dict[str, Optional[re.Pattern]] = {
-                    "tool_name": None,
-                    "tool_type": None,
-                }
-                if rule.tool_name is not None:
-                    try:
-                        compiled_target_patterns["tool_name"] = re.compile(
-                            rule.tool_name
-                        )
-                    except re.error as exc:
-                        raise ValueError(
-                            f"Invalid regex for tool_name in rule '{rule.id}': {exc}"
-                        ) from exc
-                if rule.tool_type is not None:
-                    try:
-                        compiled_target_patterns["tool_type"] = re.compile(
-                            rule.tool_type
-                        )
-                    except re.error as exc:
-                        raise ValueError(
-                            f"Invalid regex for tool_type in rule '{rule.id}': {exc}"
-                        ) from exc
-                self._compiled_rule_targets[rule.id] = compiled_target_patterns
-
-                if rule.allowed_param_patterns:
-                    compiled_patterns: Dict[str, re.Pattern] = {}
-                    for path, pattern in rule.allowed_param_patterns.items():
-                        try:
-                            compiled_patterns[path] = re.compile(pattern)
-                        except re.error as exc:
-                            raise ValueError(
-                                f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}"
-                            ) from exc
-
-                    if compiled_patterns:
-                        self._compiled_rule_patterns[rule.id] = compiled_patterns
+        self._load_rules(rules)
 
         # Normalize to lowercase for case-insensitive handling
         self.default_action = (
@@ -126,6 +80,115 @@ class ToolPermissionGuardrail(CustomGuardrail):
             self.default_action,
         )
 
+    def _load_rules(self, rules: Optional[List[Any]]) -> None:
+        """Parse ``rules`` and (re)build the compiled target/pattern lookups.
+
+        ``self.rules`` plus ``_compiled_rule_targets`` / ``_compiled_rule_patterns``
+        are the state every matching path reads. Centralizing the build here lets
+        both ``__init__`` and ``update_in_memory_litellm_params`` recompile from a
+        single source of truth, so an in-place update (PUT /guardrails, immediate
+        sync) reflects rule changes instead of keeping the construction-time maps.
+        """
+        parsed_rules: List[ToolPermissionRule] = []
+        compiled_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {}
+        compiled_patterns: Dict[str, Dict[str, re.Pattern]] = {}
+
+        for rule_item in rules or []:
+            rule = (
+                rule_item
+                if isinstance(rule_item, ToolPermissionRule)
+                else ToolPermissionRule(**rule_item)
+            )
+
+            target_patterns: Dict[str, Optional[re.Pattern]] = {
+                "tool_name": None,
+                "tool_type": None,
+            }
+            if rule.tool_name is not None:
+                try:
+                    target_patterns["tool_name"] = re.compile(rule.tool_name)
+                except re.error as exc:
+                    raise ValueError(
+                        f"Invalid regex for tool_name in rule '{rule.id}': {exc}"
+                    ) from exc
+            if rule.tool_type is not None:
+                try:
+                    target_patterns["tool_type"] = re.compile(rule.tool_type)
+                except re.error as exc:
+                    raise ValueError(
+                        f"Invalid regex for tool_type in rule '{rule.id}': {exc}"
+                    ) from exc
+
+            rule_patterns: Dict[str, re.Pattern] = {}
+            for path, pattern in (rule.allowed_param_patterns or {}).items():
+                try:
+                    rule_patterns[path] = re.compile(pattern)
+                except re.error as exc:
+                    raise ValueError(
+                        f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}"
+                    ) from exc
+
+            parsed_rules.append(rule)
+            compiled_targets[rule.id] = target_patterns
+            if rule_patterns:
+                compiled_patterns[rule.id] = rule_patterns
+
+        # Swap in the fully-built maps only after every rule compiles, so an
+        # invalid regex raises without leaving a partially-built ruleset (a
+        # missing compiled target is read as a match-all wildcard).
+        self.rules = parsed_rules
+        self._compiled_rule_targets = compiled_targets
+        self._compiled_rule_patterns = compiled_patterns
+
+    def update_in_memory_litellm_params(
+        self, litellm_params: Union[LitellmParams, dict]
+    ) -> None:
+        """Apply updated params in place, rebuilding the compiled rule state.
+
+        The base implementation only ``setattr``s raw fields, which would leave
+        ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` (built in
+        ``__init__``) stale, so a guardrail updated without reinitialization would
+        keep enforcing the old ruleset. Recompile here so PUT /guardrails and the
+        immediate in-memory sync take effect, mirroring the PresidioGuardrail
+        override of this method.
+        """
+        # ``litellm_params`` may arrive as the raw DB dict (the proxy ``cast()``s
+        # it to ``LitellmParams`` without converting), so handle both shapes. The
+        # base ``setattr`` loop is model-only, so apply the dict case here.
+        previous_rules = self.rules
+        if isinstance(litellm_params, dict):
+            params = litellm_params
+            for key, value in params.items():
+                setattr(self, key, value)
+        else:
+            super().update_in_memory_litellm_params(litellm_params)
+            params = vars(litellm_params)
+
+        # The generic update above sets ``self.rules`` from the incoming value
+        # (None on a partial update that omits rules), but never rebuilds the
+        # compiled maps. Rebuild them when rules are provided; otherwise restore
+        # the previous ruleset so a partial update doesn't silently wipe it. An
+        # explicit empty list still clears the rules.
+        rules = params.get("rules")
+        if rules is not None:
+            try:
+                self._load_rules(rules)
+            except Exception:
+                # The generic update above may have overwritten self.rules with
+                # the raw payload; restore the prior consistent ruleset so a
+                # rejected update can't leave the live guardrail enforcing a
+                # broken policy.
+                self.rules = previous_rules
+                raise
+        else:
+            self.rules = previous_rules
+        default_action = params.get("default_action")
+        if isinstance(default_action, str):
+            self.default_action = default_action.lower()
+        on_disallowed_action = params.get("on_disallowed_action")
+        if isinstance(on_disallowed_action, str):
+            self.on_disallowed_action = on_disallowed_action.lower()
+
     @staticmethod
     def get_config_model():
         from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (

litellm/proxy/proxy_server.py

@@ -1910,7 +1910,7 @@ prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = N
 store_model_in_db: bool = False
 open_telemetry_logger: Optional[OpenTelemetry] = None
 ### INITIALIZE GLOBAL LOGGING OBJECT ###
-proxy_logging_obj = ProxyLogging(
+proxy_logging_obj: ProxyLogging = ProxyLogging(
     user_api_key_cache=user_api_key_cache, premium_user=premium_user
 )
 ### REDIS QUEUE ###
@@ -15844,10 +15844,10 @@ async def toolset_mcp_route(toolset_name: str, request: Request):
     except HTTPException as e:
         raise e
     except Exception as e:
-        verbose_proxy_logger.error(
-            f"Error handling toolset MCP route for {toolset_name}: {str(e)}"
+        verbose_proxy_logger.exception(
+            "Error handling toolset MCP route for %s: %s", toolset_name, str(e)
         )
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 
 async def _mcp_forward_as_path(path_segment: str, request: Request):
@@ -16028,7 +16028,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request):
     except HTTPException as e:
         raise e
     except Exception as e:
-        verbose_proxy_logger.error(
-            f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}"
+        verbose_proxy_logger.exception(
+            "Error handling dynamic MCP route for %s: %s", mcp_server_name, str(e)
         )
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")

litellm/types/llms/gemini.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional
 
-from typing_extensions import TypedDict
+from typing_extensions import Required, TypedDict
 
 from .vertex_ai import (
     GenerationConfig,
@@ -233,10 +233,11 @@ class GeminiImageGenerationResponse(TypedDict):
 
 
 # Video Generation Types
-class GeminiVideoGenerationInstance(TypedDict):
+class GeminiVideoGenerationInstance(TypedDict, total=False):
     """Instance data for Gemini video generation request"""
 
-    prompt: str
+    prompt: Required[str]
+    image: Dict[str, Any]
 
 
 class GeminiVideoGenerationParameters(BaseModel):
@@ -264,11 +265,6 @@ class GeminiVideoGenerationParameters(BaseModel):
     negativePrompt: Optional[str] = None
     """Text describing what not to include in the video."""
 
-    image: Optional[Any] = None
-    """
-    An initial image to animate (Image object).
-    """
-
     lastFrame: Optional[Any] = None
     """
     The final image for interpolation video to transition.

litellm/types/mcp_server/mcp_server_manager.py

@@ -119,6 +119,8 @@ class MCPServer(BaseModel):
     # different ``server_id`` values are bumped deterministically.  Left
     # ``None`` in default-prefix mode.
     short_prefix: Optional[str] = None
+    allow_sampling: bool = False
+    allow_elicitation: bool = False
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @property

litellm/utils.py

@@ -4871,7 +4871,7 @@ def add_provider_specific_params_to_optional_params(
             )
             is False
         ):
-            extra_body = passed_params.pop("extra_body", None) or {}
+            extra_body = dict(passed_params.pop("extra_body", None) or {})
             for k in passed_params.keys():
                 if k not in openai_params and passed_params[k] is not None:
                     extra_body[k] = passed_params[k]
@@ -8909,6 +8909,16 @@ class ProviderConfigManager:
             return litellm.OpenRouterResponsesAPIConfig()
         elif litellm.LlmProviders.HOSTED_VLLM == provider:
             return litellm.HostedVLLMResponsesAPIConfig()
+        elif litellm.LlmProviders.BEDROCK_MANTLE == provider:
+            # Only OpenAI gpt frontier models (gpt-5.x, and future gpt-6 etc.) are
+            # served on the /openai/v1/responses path. gpt-oss and every non-OpenAI
+            # model on Mantle (nvidia, mistral, google, zai, ...) are chat-completions
+            # only and 400 on that path, so they fall through to None to keep the
+            # chat-completions emulation (see litellm/responses/main.py "config is None").
+            model_lower = model.lower() if model else ""
+            if "openai.gpt-" in model_lower and "gpt-oss" not in model_lower:
+                return litellm.BedrockMantleResponsesAPIConfig()
+            return None
         return None
 
     @staticmethod

model_prices_and_context_window.json

@@ -30424,21 +30424,32 @@
         "supports_reasoning": true,
         "source": "https://cloud.sambanova.ai/plans/pricing"
     },
-    "snowflake/claude-3-5-sonnet": {
+      "snowflake/claude-3-5-sonnet": {
         "litellm_provider": "snowflake",
-        "max_input_tokens": 18000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
+        "max_input_tokens": 200000,
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
         "mode": "chat",
-        "supports_computer_use": true
+        "input_cost_per_token": 0.000003,
+        "output_cost_per_token": 0.000015,
+        "cache_read_input_token_cost": 0.0000003,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_vision": true,
+        "supports_prompt_caching": true,
+        "supports_system_messages": true,
+        "supports_response_schema": true
     },
-    "snowflake/deepseek-r1": {
+      "snowflake/deepseek-r1": {
         "litellm_provider": "snowflake",
-        "max_input_tokens": 32768,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
+        "max_input_tokens": 128000,
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
         "mode": "chat",
-        "supports_reasoning": true
+        "input_cost_per_token": 0.00000135,
+        "output_cost_per_token": 0.0000054,
+        "supports_reasoning": true,
+        "supports_system_messages": true
     },
     "snowflake/gemma-7b": {
         "litellm_provider": "snowflake",
@@ -30492,23 +30503,34 @@
     "snowflake/llama3.1-405b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.0000012,
+        "output_cost_per_token": 0.0000012,
+        "supports_function_calling": true,
+        "supports_system_messages": true
     },
     "snowflake/llama3.1-70b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.00000072,
+        "output_cost_per_token": 0.00000072,
+        "supports_function_calling": true,
+        "supports_system_messages": true
     },
     "snowflake/llama3.1-8b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.00000024,
+        "output_cost_per_token": 0.00000024,
+        "supports_system_messages": true
     },
     "snowflake/llama3.2-1b": {
         "litellm_provider": "snowflake",
@@ -30524,13 +30546,17 @@
         "max_tokens": 8192,
         "mode": "chat"
     },
-    "snowflake/llama3.3-70b": {
-        "litellm_provider": "snowflake",
+      "snowflake/llama3.3-70b": {
+        "max_tokens": 16384,
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
-    },
+        "max_output_tokens": 16384,
+        "input_cost_per_token": 0.00000072,
+        "output_cost_per_token": 0.00000072,
+        "litellm_provider": "snowflake",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_system_messages": true
+      },
     "snowflake/mistral-7b": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 32000,
@@ -30545,12 +30571,17 @@
         "max_tokens": 8192,
         "mode": "chat"
     },
-    "snowflake/mistral-large2": {
+     "snowflake/mistral-large2": {
         "litellm_provider": "snowflake",
         "max_input_tokens": 128000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
+        "max_output_tokens": 16384,
+        "max_tokens": 16384,
+        "mode": "chat",
+        "input_cost_per_token": 0.000002,
+        "output_cost_per_token": 0.000006,
+        "supports_function_calling": true,
+        "supports_system_messages": true,
+        "supports_response_schema": true
     },
     "snowflake/mixtral-8x7b": {
         "litellm_provider": "snowflake",
@@ -30587,13 +30618,17 @@
         "max_tokens": 8192,
         "mode": "chat"
     },
-    "snowflake/snowflake-llama-3.3-70b": {
+      "snowflake/snowflake-llama-3.3-70b": {
+        "max_tokens": 16384,
+        "max_input_tokens": 128000,
+        "max_output_tokens": 16384,
+        "input_cost_per_token": 0.00000072,
+        "output_cost_per_token": 0.00000072,
         "litellm_provider": "snowflake",
-        "max_input_tokens": 8000,
-        "max_output_tokens": 8192,
-        "max_tokens": 8192,
-        "mode": "chat"
-    },
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_system_messages": true
+      },
     "stability/sd3": {
         "litellm_provider": "stability",
         "mode": "image_generation",
@@ -41223,6 +41258,44 @@
         "supports_response_schema": true,
         "supports_tool_choice": true
     },
+    "bedrock_mantle/openai.gpt-5.5": {
+        "input_cost_per_token": 5.5e-06,
+        "cache_read_input_token_cost": 5.5e-07,
+        "output_cost_per_token": 3.3e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/openai.gpt-5.4": {
+        "input_cost_per_token": 2.75e-06,
+        "cache_read_input_token_cost": 2.75e-07,
+        "output_cost_per_token": 1.65e-05,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 272000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "responses",
+        "supported_endpoints": ["/v1/responses"],
+        "supported_modalities": ["text", "image"],
+        "supported_output_modalities": ["text"],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "volcengine/doubao-seed-2-0-pro-260215": {
         "litellm_provider": "volcengine",
         "max_input_tokens": 256000,
@@ -41501,5 +41574,180 @@
         "supports_vision": true,
         "supports_native_structured_output": true,
         "supports_pdf_input": true
-    }
-}
+    },
+  "snowflake/claude-sonnet-4-5": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-sonnet-4-6": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-4-sonnet": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-4-opus": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000005,
+    "output_cost_per_token": 0.000025,
+    "cache_read_input_token_cost": 0.0000005,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_reasoning": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-haiku-4-5": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000001,
+    "output_cost_per_token": 0.000005,
+    "cache_read_input_token_cost": 0.0000001,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/claude-3-7-sonnet": {
+    "max_tokens": 16384,
+    "max_input_tokens": 200000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000003,
+    "output_cost_per_token": 0.000015,
+    "cache_read_input_token_cost": 0.0000003,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_reasoning": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-4.1": {
+    "max_tokens": 16384,
+    "max_input_tokens": 300000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.000002,
+    "output_cost_per_token": 0.000008,
+    "cache_read_input_token_cost": 0.0000005,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-5": {
+    "max_tokens": 16384,
+    "max_input_tokens": 300000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.00000125,
+    "output_cost_per_token": 0.00001,
+    "cache_read_input_token_cost": 0.000000125,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_vision": true,
+    "supports_prompt_caching": true,
+    "supports_system_messages": true,
+    "supports_reasoning": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-5-mini": {
+    "max_tokens": 16384,
+    "max_input_tokens": 1000000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.0000003,
+    "output_cost_per_token": 0.0000012,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/openai-gpt-5-nano": {
+    "max_tokens": 16384,
+    "max_input_tokens": 5000000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.00000015,
+    "output_cost_per_token": 0.0000006,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_system_messages": true,
+    "supports_response_schema": true
+  },
+  "snowflake/llama4-maverick": {
+    "max_tokens": 16384,
+    "max_input_tokens": 128000,
+    "max_output_tokens": 16384,
+    "input_cost_per_token": 0.00000024,
+    "output_cost_per_token": 0.00000097,
+    "litellm_provider": "snowflake",
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_system_messages": true
+  },
+  "snowflake/snowflake-arctic-embed-l-v2.0": {
+    "max_tokens": 8192,
+    "max_input_tokens": 8192,
+    "input_cost_per_token": 0.00000007,
+    "output_cost_per_token": 0.0,
+    "litellm_provider": "snowflake",
+    "mode": "embedding"
+  },
+  "snowflake/snowflake-arctic-embed-m-v2.0": {
+    "max_tokens": 8192,
+    "max_input_tokens": 8192,
+    "input_cost_per_token": 0.00000007,
+    "output_cost_per_token": 0.0,
+    "litellm_provider": "snowflake",
+    "mode": "embedding"
+  }
+ }
+

tests/llm_translation/test_fireworks_ai_translation.py

@@ -43,12 +43,14 @@ def test_map_openai_params_tool_choice():
 
 def test_map_response_format():
     """
-    Test that the response format is translated correctly.
+    json_schema response_format is passed through to Fireworks unchanged.
 
-    h/t to https://github.com/DaveDeCaprio (@DaveDeCaprio) for the test case
+    Fireworks accepts the OpenAI strict json_schema shape natively. The earlier
+    downgrade to {type: json_object, schema: ...} silently dropped `strict` and
+    `name`, producing a request that Fireworks treats as "any valid JSON" per
+    its docs, disabling grammar-guided decoding.
 
-    Relevant Issue: https://github.com/BerriAI/litellm/issues/6797
-    Fireworks AI Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting#step-1-import-libraries
+    Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting
     """
     response_format = {
         "type": "json_schema",
@@ -65,16 +67,7 @@ def test_map_response_format():
     result = fireworks.map_openai_params(
         {"response_format": response_format}, {}, "some_model", drop_params=False
     )
-    assert result == {
-        "response_format": {
-            "type": "json_object",
-            "schema": {
-                "properties": {"result": {"type": "boolean"}},
-                "required": ["result"],
-                "type": "object",
-            },
-        }
-    }
+    assert result == {"response_format": response_format}
 
 
 class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest):

tests/local_testing/test_custom_llm.py

@@ -44,7 +44,14 @@ from litellm import (
     image_generation,
 )
 from litellm.utils import ModelResponseIterator
-from litellm.types.utils import ImageResponse, ImageObject, EmbeddingResponse
+from litellm.types.utils import (
+    ImageResponse,
+    ImageObject,
+    EmbeddingResponse,
+    ModelResponseStream,
+    StreamingChoices,
+    Delta,
+)
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 
 
@@ -644,3 +651,82 @@ async def test_simple_aembedding():
         "embedding": [0.1, 0.2, 0.3],
         "index": 1,
     }
+
+
+# ── Tests for ModelResponseStream passthrough in custom providers (issue #27389) ──
+
+
+class ModelResponseStreamLLM(MyCustomLLM):
+    """Subclass that overrides streaming/astreaming to yield ModelResponseStream directly."""
+
+    def __init__(self, finish_reason: str = "stop"):
+        self._finish_reason = finish_reason
+
+    def streaming(self, *args, **kwargs) -> Iterator[ModelResponseStream]:  # type: ignore
+        yield ModelResponseStream(
+            id="test-stream-id",
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(content="Hello world"),
+                    finish_reason=self._finish_reason,
+                )
+            ],
+        )
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[ModelResponseStream]:  # type: ignore
+        yield ModelResponseStream(
+            id="test-stream-id",
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(content="Hello world"),
+                    finish_reason=self._finish_reason,
+                )
+            ],
+        )
+
+
+@pytest.mark.parametrize(
+    "finish_reason", ["stop", "tool_calls", "length", "content_filter"]
+)
+def test_custom_llm_streaming_model_response_stream(finish_reason):
+    my_custom_llm = ModelResponseStreamLLM(finish_reason=finish_reason)
+    litellm.custom_provider_map = [
+        {"provider": "custom_llm", "custom_handler": my_custom_llm}
+    ]
+    resp = completion(
+        model="custom_llm/my-fake-model",
+        messages=[{"role": "user", "content": "Hello world!"}],
+        stream=True,
+    )
+
+    for chunk in resp:
+        print(chunk)
+        if chunk.choices[0].finish_reason is None:
+            assert isinstance(chunk.choices[0].delta.content, str)
+        else:
+            assert chunk.choices[0].finish_reason == finish_reason
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "finish_reason", ["stop", "tool_calls", "length", "content_filter"]
+)
+async def test_custom_llm_astreaming_model_response_stream(finish_reason):
+    my_custom_llm = ModelResponseStreamLLM(finish_reason=finish_reason)
+    litellm.custom_provider_map = [
+        {"provider": "custom_llm", "custom_handler": my_custom_llm}
+    ]
+    resp = await litellm.acompletion(
+        model="custom_llm/my-fake-model",
+        messages=[{"role": "user", "content": "Hello world!"}],
+        stream=True,
+    )
+
+    async for chunk in resp:
+        print(chunk)
+        if chunk.choices[0].finish_reason is None:
+            assert isinstance(chunk.choices[0].delta.content, str)
+        else:
+            assert chunk.choices[0].finish_reason == finish_reason

tests/local_testing/test_get_llm_provider.py

@@ -131,6 +131,7 @@ def test_default_api_base():
     from litellm.litellm_core_utils.get_llm_provider_logic import (
         _get_openai_compatible_provider_info,
     )
+    from litellm.types.utils import LlmProviders
 
     # Patch environment variable to remove API base if it's set
     with patch.dict(os.environ, {}, clear=True):
@@ -150,13 +151,13 @@ def test_default_api_base():
             if api_base is None:
                 continue
 
-            for other_provider in litellm.provider_list:
-                if other_provider != provider and provider != "{}_chat".format(
+            for other_provider in LlmProviders:
+                if other_provider.value != provider and provider != "{}_chat".format(
                     other_provider.value
                 ):
-                    if provider == "codestral" and other_provider == "mistral":
+                    if provider == "codestral" and other_provider.value == "mistral":
                         continue
-                    elif provider == "github" and other_provider == "azure":
+                    elif provider == "github" and other_provider.value == "azure":
                         continue
                     assert other_provider.value not in api_base.replace("/openai", "")
 

tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py

@@ -0,0 +1,134 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from litellm.litellm_core_utils.get_supported_openai_params import (
+    get_supported_openai_params,
+)
+
+BEDROCK_REAL_MODEL = "eu.anthropic.claude-haiku-4-5-20251001-v1:0"
+BEDROCK_LABEL = "claude-haiku-4-5"
+
+
+def test_base_model_label_does_not_strip_bedrock_tools():
+    """Regression for #29618.
+
+    A Bedrock deployment whose ``model_info.base_model`` is a friendly label
+    (``claude-haiku-4-5``) must still advertise ``tools``/``tool_choice``. The label
+    on its own resolves to no tool support, so before the fix it stripped the
+    capability the real model id exposes, silently dropping function calling under
+    ``drop_params``."""
+    params = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL,
+        custom_llm_provider="bedrock",
+        base_model=BEDROCK_LABEL,
+    )
+
+    assert params is not None
+    assert "tools" in params
+    assert "tool_choice" in params
+
+
+def test_base_model_label_alone_lacks_bedrock_tools():
+    """The label by itself does not advertise tools; this is what made the union
+    necessary. Guards against the discrepancy disappearing (and the regression test
+    above silently passing for the wrong reason)."""
+    params = get_supported_openai_params(
+        model=BEDROCK_LABEL, custom_llm_provider="bedrock"
+    )
+
+    assert params is not None
+    assert "tools" not in params
+
+
+def test_base_model_is_additive_not_replacement():
+    """``base_model`` may only add capabilities, never remove ones the real model has.
+
+    Bedrock: real id supports ``tools`` but not the label's reasoning hint; the union
+    must contain the real model's ``tools`` regardless of the label being a subset."""
+    real_only = set(
+        get_supported_openai_params(
+            model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
+        )
+    )
+    label_only = set(
+        get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")
+    )
+    combined = set(
+        get_supported_openai_params(
+            model=BEDROCK_REAL_MODEL,
+            custom_llm_provider="bedrock",
+            base_model=BEDROCK_LABEL,
+        )
+    )
+
+    assert combined == real_only | label_only
+    assert real_only - label_only  # the label really is a strict subset here
+    assert real_only <= combined
+
+
+def test_base_model_adds_capabilities_the_real_model_lacks():
+    """Regression for #27717 (the behavior the union must preserve).
+
+    ``gemini-3.1-pro`` isn't in the cost map so it advertises no reasoning support,
+    but the registered ``gemini-3.1-pro-preview`` base_model does. The hint must add
+    ``reasoning_effort``/``thinking`` without the call erroring."""
+    real_only = set(
+        get_supported_openai_params(
+            model="gemini-3.1-pro", custom_llm_provider="gemini"
+        )
+    )
+    assert "reasoning_effort" not in real_only
+
+    combined = set(
+        get_supported_openai_params(
+            model="gemini-3.1-pro",
+            custom_llm_provider="gemini",
+            base_model="gemini-3.1-pro-preview",
+        )
+    )
+    assert "reasoning_effort" in combined
+    assert "thinking" in combined
+
+
+def test_no_base_model_is_unchanged():
+    """Omitting ``base_model`` must resolve purely from ``model``."""
+    with_none = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None
+    )
+    plain = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
+    )
+
+    assert with_none == plain
+
+
+def test_base_model_equal_to_model_is_unchanged():
+    """A ``base_model`` identical to ``model`` must not double-resolve or reorder."""
+    plain = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
+    )
+    same = get_supported_openai_params(
+        model=BEDROCK_REAL_MODEL,
+        custom_llm_provider="bedrock",
+        base_model=BEDROCK_REAL_MODEL,
+    )
+
+    assert same == plain
+
+
+def test_azure_base_model_detection_preserved():
+    """Azure relies on ``base_model`` for model-type detection when the deployment
+    name is opaque; the union must keep advertising the gpt-5 capabilities."""
+    params = get_supported_openai_params(
+        model="my-opaque-deployment",
+        custom_llm_provider="azure",
+        base_model="azure/gpt-5.2",
+    )
+
+    assert params is not None
+    assert "reasoning_effort" in params
+    assert "tools" in params

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -2118,3 +2118,172 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum():
         f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. "
         "STOP enum was not normalised through map_finish_reason()."
     )
+
+
+@pytest.mark.parametrize(
+    "finish_reason", ["stop", "tool_calls", "length", "content_filter"]
+)
+def test_chunk_creator_passes_through_model_response_stream(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+    finish_reason: str,
+):
+    """
+    chunk_creator must pass ModelResponseStream chunks from custom providers
+    straight through and preserve finish_reason exactly — not force-cast to GChunk.
+    Regression test for issue #27389.
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content="Hello", role="assistant"),
+                finish_reason=finish_reason,
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is not None
+    assert initialized_custom_stream_wrapper.received_finish_reason == finish_reason
+
+
+def test_chunk_creator_drops_empty_finish_chunk(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    A ModelResponseStream chunk with finish_reason but no content should return
+    None so finish_reason_handler() synthesises the final chunk — mirrors GChunk
+    behaviour via is_chunk_non_empty.
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content=""),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is None
+    assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+
+
+def test_chunk_creator_stops_iteration_on_trailing_chunk(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    After received_finish_reason is set, any empty trailing chunk (e.g. provider
+    metadata flush) must raise StopIteration to end the stream cleanly.
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    initialized_custom_stream_wrapper.received_finish_reason = "stop"
+    litellm._custom_providers.append("my-custom-provider")
+
+    trailing_chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content=None),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    with pytest.raises(StopIteration):
+        initialized_custom_stream_wrapper.chunk_creator(chunk=trailing_chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+
+def test_chunk_creator_strips_finish_reason_from_content_chunk(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    When content and finish_reason arrive in the same chunk, finish_reason must be
+    stripped so finish_reason_handler() emits it on the synthetic terminal chunk —
+    preventing two terminal chunks (double finish_reason bug).
+    """
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(content="Hello"),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is not None
+    assert (
+        result.choices[0].finish_reason is None
+    ), "finish_reason must be stripped from content chunks to avoid double terminal chunks"
+    assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
+
+
+def test_chunk_creator_tool_calls_not_dropped_on_finish(
+    initialized_custom_stream_wrapper: CustomStreamWrapper,
+):
+    """
+    A terminal chunk with finish_reason="tool_calls" and delta.tool_calls must NOT
+    be silently dropped — tool_calls counts as content so the chunk is passed through
+    (with finish_reason stripped) rather than returning None.
+    """
+    from litellm.types.utils import ChatCompletionDeltaToolCall, Function
+
+    initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-provider"
+    litellm._custom_providers.append("my-custom-provider")
+
+    chunk = ModelResponseStream(
+        id="test-id",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(
+                    content=None,
+                    tool_calls=[
+                        ChatCompletionDeltaToolCall(
+                            id="call_abc",
+                            function=Function(name="get_weather", arguments='{"city":"NYC"}'),
+                            type="function",
+                            index=0,
+                        )
+                    ],
+                ),
+                finish_reason="tool_calls",
+            )
+        ],
+    )
+
+    result = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk)
+
+    litellm._custom_providers.remove("my-custom-provider")
+
+    assert result is not None, "tool_calls chunk must not be dropped"
+    assert result.choices[0].delta.tool_calls is not None
+    assert result.choices[0].finish_reason is None
+    assert initialized_custom_stream_wrapper.received_finish_reason == "tool_calls"

tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py

@@ -1,5 +1,7 @@
+import urllib.parse
 from unittest.mock import patch
 
+import litellm
 from litellm.llms.azure.image_edit.transformation import AzureImageEditConfig
 from litellm.types.router import GenericLiteLLMParams
 
@@ -138,3 +140,96 @@ def test_azure_finalize_image_edit_strips_model_after_openai_transform():
     assert data_out.get("prompt") == prompt
     assert data_out.get("n") == 1
     assert len(files) >= 1
+
+
+# ---------------------------------------------------------------------------
+# api_version fallback chain
+#
+# Pin the resolution order used by ``AzureImageEditConfig.get_complete_url``:
+#   litellm_params["api_version"]
+#     > litellm.api_version (module-global)
+#     > AZURE_API_VERSION env var
+#     > litellm.AZURE_DEFAULT_API_VERSION
+#
+# Before this fallback chain existed, image edit only read ``litellm_params``
+# and produced an unversioned URL when callers set api_version via the global
+# or the env var (Azure then 404s with "Resource not found"). The chat path
+# in ``litellm/llms/azure/common_utils.py`` already had this fallback.
+# ---------------------------------------------------------------------------
+
+
+_FALLBACK_API_BASE = "https://x.openai.azure.com"
+_FALLBACK_MODEL = "gpt-image-1"
+
+
+def _query_params(url: str) -> dict:
+    return dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query))
+
+
+def test_api_version_uses_litellm_params_first(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", "from-global", raising=False)
+    monkeypatch.setenv("AZURE_API_VERSION", "from-env")
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={"api_version": "from-params"},
+    )
+
+    assert _query_params(url) == {"api-version": "from-params"}
+
+
+def test_api_version_falls_back_to_litellm_global(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", "from-global", raising=False)
+    monkeypatch.setenv("AZURE_API_VERSION", "from-env")
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={},
+    )
+
+    assert _query_params(url) == {"api-version": "from-global"}
+
+
+def test_api_version_falls_back_to_env_var(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", None, raising=False)
+    monkeypatch.setenv("AZURE_API_VERSION", "from-env")
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={},
+    )
+
+    assert _query_params(url) == {"api-version": "from-env"}
+
+
+def test_api_version_falls_back_to_azure_default(monkeypatch):
+    monkeypatch.setattr(litellm, "api_version", None, raising=False)
+    monkeypatch.delenv("AZURE_API_VERSION", raising=False)
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=_FALLBACK_API_BASE,
+        litellm_params={},
+    )
+
+    assert _query_params(url) == {"api-version": litellm.AZURE_DEFAULT_API_VERSION}
+
+
+def test_api_version_in_api_base_query_is_preserved(monkeypatch):
+    """``api_base`` already carrying ``?api-version=...`` must not be overridden."""
+    monkeypatch.setattr(litellm, "api_version", None, raising=False)
+    monkeypatch.delenv("AZURE_API_VERSION", raising=False)
+
+    url = AzureImageEditConfig().get_complete_url(
+        model=_FALLBACK_MODEL,
+        api_base=(
+            f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}"
+            "/images/edits?api-version=2024-05-01-preview"
+        ),
+        litellm_params={"api_version": "would-be-overridden"},
+    )
+
+    assert _query_params(url) == {"api-version": "2024-05-01-preview"}

tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py

@@ -0,0 +1,283 @@
+"""
+Unit tests for Amazon Bedrock Mantle Responses API configuration.
+
+Mantle's gpt-5.5 / gpt-5.4 are served ONLY on the non-standard
+`/openai/v1/responses` path. These tests lock the URL construction and
+Bearer auth that make that routing work.
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+
+import pytest
+
+import litellm
+from litellm.llms.bedrock_mantle.responses.transformation import (
+    BedrockMantleResponsesAPIConfig,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+
+class TestBedrockMantleResponsesURL:
+    def test_url_uses_region_from_env(self, monkeypatch):
+        monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2")
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(api_base=None, litellm_params={})
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+
+    def test_url_normalizes_v1_suffix(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/v1",
+            litellm_params={},
+        )
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+        assert "/v1/openai/v1/responses" not in url
+        url_trailing = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/v1/",
+            litellm_params={},
+        )
+        assert (
+            url_trailing
+            == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+        )
+
+    def test_url_does_not_double_openai_v1(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1",
+            litellm_params={},
+        )
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+
+    def test_url_full_endpoint_base_not_doubled(self, monkeypatch):
+        # AWS model card tells users to set OPENAI_BASE_URL to the full endpoint.
+        # If copied into api_base, it must not be doubled.
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(
+            api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
+            litellm_params={},
+        )
+        assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
+        assert url.count("/responses") == 1
+
+    def test_url_region_fallback_to_aws_region(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        monkeypatch.setenv("AWS_REGION", "us-west-2")
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(api_base=None, litellm_params={})
+        assert url == "https://bedrock-mantle.us-west-2.api.aws/openai/v1/responses"
+
+    def test_url_region_default_us_east_1(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
+        monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+        monkeypatch.delenv("AWS_REGION", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        url = cfg.get_complete_url(api_base=None, litellm_params={})
+        assert url == "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses"
+
+
+class TestBedrockMantleResponsesAuth:
+    def test_config_api_key_takes_priority(self, monkeypatch):
+        monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key")
+        cfg = BedrockMantleResponsesAPIConfig()
+        headers = cfg.validate_environment(
+            headers={},
+            model="openai.gpt-5.5",
+            litellm_params=GenericLiteLLMParams(api_key="config-key"),
+        )
+        assert headers["Authorization"] == "Bearer config-key"
+
+    def test_env_key_fallback(self, monkeypatch):
+        monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key")
+        monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        headers = cfg.validate_environment(
+            headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
+        )
+        assert headers["Authorization"] == "Bearer env-key"
+
+    def test_bedrock_bearer_token_fallback(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
+        monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bearer-key")
+        cfg = BedrockMantleResponsesAPIConfig()
+        headers = cfg.validate_environment(
+            headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
+        )
+        assert headers["Authorization"] == "Bearer bearer-key"
+
+    def test_missing_key_raises(self, monkeypatch):
+        monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
+        monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+        cfg = BedrockMantleResponsesAPIConfig()
+        with pytest.raises(ValueError, match="Bedrock Mantle API key"):
+            cfg.validate_environment(
+                headers={},
+                model="openai.gpt-5.5",
+                litellm_params=GenericLiteLLMParams(),
+            )
+
+    def test_custom_llm_provider(self):
+        cfg = BedrockMantleResponsesAPIConfig()
+        assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE
+
+    def test_native_websocket_disabled(self):
+        # Mantle Responses has no realtime/websocket transport, so the config
+        # must opt out; otherwise realtime routing would try a socket Mantle
+        # does not serve.
+        cfg = BedrockMantleResponsesAPIConfig()
+        assert cfg.supports_native_websocket() is False
+
+    def test_file_search_routes_to_emulation(self):
+        # Mantle cannot reach OpenAI's vector stores, so a native file_search
+        # tool forwarded as-is gets a 400. The config must opt out of native
+        # file_search so LiteLLM's emulation handles it instead of forwarding.
+        from litellm.responses.file_search.emulated_handler import (
+            should_use_emulated_file_search,
+        )
+
+        cfg = BedrockMantleResponsesAPIConfig()
+        assert cfg.supports_native_file_search() is False
+        assert (
+            should_use_emulated_file_search(
+                tools=[{"type": "file_search", "vector_store_ids": ["vs_1"]}],
+                provider_config=cfg,
+            )
+            is True
+        )
+
+
+class TestBedrockMantleResponsesRegistry:
+    def test_registry_returns_config_for_gpt_5_5(self):
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-5.5",
+        )
+        assert isinstance(cfg, BedrockMantleResponsesAPIConfig)
+
+    def test_registry_returns_config_for_gpt_5_4_enum(self):
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.BEDROCK_MANTLE,
+            model="openai.gpt-5.4",
+        )
+        assert isinstance(cfg, BedrockMantleResponsesAPIConfig)
+
+    def test_registry_returns_none_for_gpt_oss(self):
+        # Regression guard: gpt-oss must NOT get the native Responses config; it
+        # keeps the chat-completions emulation path (responses/main.py ~line 1109).
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-oss-120b",
+        )
+        assert cfg is None
+
+    def test_registry_returns_none_for_gpt_oss_safeguard(self):
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-oss-safeguard-20b",
+        )
+        assert cfg is None
+
+    def test_registry_returns_config_for_future_frontier_model(self):
+        # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6) must
+        # get the native Responses config without a code change. The gate allow-lists
+        # the openai.gpt- family (minus gpt-oss), so gpt-6 matches automatically.
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model="openai.gpt-6",
+        )
+        assert isinstance(cfg, BedrockMantleResponsesAPIConfig)
+
+    @pytest.mark.parametrize(
+        "model",
+        [
+            "nvidia.nemotron-nano-9b-v2",
+            "mistral.ministral-3-3b-instruct",
+            "google.gemma-3-27b-it",
+            "zai.glm-4.6",
+        ],
+    )
+    def test_registry_returns_none_for_non_openai_models(self, model):
+        # Regression for the chat-only families on Mantle. These models 400 on
+        # /openai/v1/responses and are served on /v1/chat/completions, so the
+        # registry must NOT hand them the Responses config; they fall through to
+        # None and keep the chat-completions emulation.
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model=model,
+        )
+        assert cfg is None
+
+    def test_registry_returns_none_when_model_is_none(self):
+        # By-id operations (delete/get/cancel) call with model=None; keep returning
+        # None so those paths are unchanged.
+        from litellm.utils import ProviderConfigManager
+
+        cfg = ProviderConfigManager.get_provider_responses_api_config(
+            provider="bedrock_mantle",
+            model=None,
+        )
+        assert cfg is None
+
+
+@pytest.fixture
+def local_cost_map(monkeypatch):
+    """Force the bundled backup cost map and re-derive the provider model sets.
+
+    ``litellm.model_cost`` is populated once at import time (here, from the
+    network-fetched ``main`` copy, which lags this branch). ``add_known_models``
+    only re-buckets whatever is already in ``model_cost``, so the cost map must
+    first be reloaded from the local backup before the new keys appear.
+    """
+    original_model_cost = litellm.model_cost
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+    litellm.get_model_info.cache_clear()
+    litellm.add_known_models()
+    try:
+        yield
+    finally:
+        litellm.model_cost = original_model_cost
+        litellm.get_model_info.cache_clear()
+
+
+class TestBedrockMantleResponsesPricing:
+    def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
+        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")
+        assert info["mode"] == "responses"
+        assert info["input_cost_per_token"] == pytest.approx(5.5e-06)
+        assert info["output_cost_per_token"] == pytest.approx(3.3e-05)
+        assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07)
+        assert info["max_input_tokens"] == 272000
+
+    def test_gpt_5_4_pricing_and_mode(self, local_cost_map):
+        info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4")
+        assert info["mode"] == "responses"
+        assert info["input_cost_per_token"] == pytest.approx(2.75e-06)
+        assert info["output_cost_per_token"] == pytest.approx(1.65e-05)
+        assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07)
+        assert info["max_input_tokens"] == 272000
+
+    def test_models_registered(self, local_cost_map):
+        assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
+        assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models

tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py

@@ -496,3 +496,59 @@ def test_transform_tools_skips_non_function_tools():
         "type": "object",
         "properties": {"id": {"type": "string"}},
     }
+
+
+def test_map_response_format_passes_json_schema_through_unchanged():
+    """
+    json_schema response_format must reach Fireworks unchanged.
+
+    Regression guard for the prior downgrade to {type: json_object, schema: ...}
+    which silently dropped `strict` and `name` and disabled grammar-guided
+    decoding on the Fireworks side.
+    """
+    config = FireworksAIConfig()
+    response_format = {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "priority_classification",
+            "strict": True,
+            "schema": {
+                "type": "object",
+                "properties": {
+                    "priority": {
+                        "type": "string",
+                        "enum": ["high", "medium", "low"],
+                    }
+                },
+                "required": ["priority"],
+                "additionalProperties": False,
+            },
+        },
+    }
+
+    result = config.map_openai_params(
+        {"response_format": response_format},
+        {},
+        "fireworks_ai/accounts/fireworks/models/qwen3-32b",
+        drop_params=False,
+    )
+
+    rf = result["response_format"]
+    assert rf["type"] == "json_schema"
+    assert rf["json_schema"]["name"] == "priority_classification"
+    assert rf["json_schema"]["strict"] is True
+    assert rf["json_schema"]["schema"] == response_format["json_schema"]["schema"]
+
+
+def test_map_response_format_json_object_unchanged():
+    """
+    The plain json_object form keeps working as before.
+    """
+    config = FireworksAIConfig()
+    result = config.map_openai_params(
+        {"response_format": {"type": "json_object"}},
+        {},
+        "fireworks_ai/accounts/fireworks/models/qwen3-32b",
+        drop_params=False,
+    )
+    assert result == {"response_format": {"type": "json_object"}}

tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py

@@ -2,6 +2,7 @@
 Tests for Gemini (Veo) video generation transformation.
 """
 
+import io
 import json
 import os
 from unittest.mock import MagicMock, Mock, patch
@@ -132,6 +133,87 @@ class TestGeminiVideoConfig:
         assert data["parameters"]["durationSeconds"] == 8
         assert data["parameters"]["resolution"] == "1080p"
 
+    def test_transform_video_create_request_image_goes_to_instance(self):
+        """Image belongs in instances[0], not in parameters (per Veo API)."""
+        prompt = "Animate this still"
+        api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning"
+        image_dict = {"bytesBase64Encoded": "aGVsbG8=", "mimeType": "image/jpeg"}
+
+        data, _, _ = self.config.transform_video_create_request(
+            model="veo-3.0-generate-preview",
+            prompt=prompt,
+            api_base=api_base,
+            video_create_optional_request_params={
+                "image": image_dict,
+                "aspectRatio": "16:9",
+                "durationSeconds": 4,
+            },
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+
+        assert data["instances"][0]["prompt"] == prompt
+        assert data["instances"][0]["image"] == image_dict
+        assert "image" not in data.get("parameters", {})
+        assert data["parameters"]["aspectRatio"] == "16:9"
+        assert data["parameters"]["durationSeconds"] == 4
+
+    def test_transform_video_create_request_image_filelike_goes_to_instance(self):
+        """File-like image (BytesIO) gets base64-encoded into instances[0]['image']."""
+        prompt = "Animate this still"
+        api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning"
+        # 1x1 PNG (8 bytes after magic + minimal IHDR is not legal — but the
+        # transformer only cares that ImageEditRequestUtils can sniff a MIME and
+        # that .read() returns bytes; an explicit name="image.jpeg" hands the
+        # MIME sniffer a clean answer regardless of payload).
+        image_bytes = b"\xff\xd8\xff\xe0fake-jpeg-bytes"
+        image_file = io.BytesIO(image_bytes)
+        image_file.name = "still.jpeg"
+
+        data, _, _ = self.config.transform_video_create_request(
+            model="veo-3.0-generate-preview",
+            prompt=prompt,
+            api_base=api_base,
+            video_create_optional_request_params={
+                "image": image_file,
+                "aspectRatio": "16:9",
+            },
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+
+        # File-like took the _convert_image_to_gemini_format branch and landed
+        # in instances[0]["image"], not in parameters.
+        instance_image = data["instances"][0]["image"]
+        assert isinstance(instance_image, dict)
+        assert instance_image["mimeType"].startswith("image/")
+        assert instance_image["bytesBase64Encoded"]
+        # Round-trip the base64 — should equal the original bytes.
+        import base64
+
+        assert base64.b64decode(instance_image["bytesBase64Encoded"]) == image_bytes
+        assert "image" not in data.get("parameters", {})
+
+    def test_transform_video_create_request_image_none_is_dropped(self):
+        """Explicit image=None is popped and never reaches parameters."""
+        prompt = "no image at all"
+        api_base = "https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning"
+
+        data, _, _ = self.config.transform_video_create_request(
+            model="veo-3.0-generate-preview",
+            prompt=prompt,
+            api_base=api_base,
+            video_create_optional_request_params={
+                "image": None,
+                "aspectRatio": "16:9",
+            },
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+
+        assert "image" not in data["instances"][0]
+        assert "image" not in data.get("parameters", {})
+
     def test_map_openai_params(self):
         """Test parameter mapping from OpenAI format to Veo format."""
         openai_params = {

tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py

@@ -104,6 +104,23 @@ class TestHuggingFaceEmbedding:
         assert "source_sentence" not in str(request_data)
         assert "sentences" not in str(request_data)
 
+    def test_embedding_allows_special_token_looking_input(self):
+        input_text = ["hello <|fim_prefix|> world"]
+
+        response = litellm.embedding(
+            model=self.model,
+            input=input_text,
+            input_type="embed",
+        )
+
+        self.mock_http.assert_called_once()
+        post_call_args = self.mock_http.call_args
+        request_data = json.loads(post_call_args[1]["data"])
+
+        assert request_data["inputs"] == input_text
+        assert response.usage.prompt_tokens > 0
+        assert response.usage.total_tokens == response.usage.prompt_tokens
+
     def test_embedding_with_sentence_similarity_task(self):
         """Test embedding when task type is sentence-similarity (requires 2+ sentences)"""
 

tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py

@@ -201,9 +201,12 @@ class TestContextCachingEndpoints:
         assert returned_params == optional_params
         assert returned_cache == "existing_cache_name"
 
-        # Verify cache key was generated with tools and model
+        # Verify cache key was generated with tools, tool_choice and model
         mock_cache_obj.get_cache_key.assert_called_once_with(
-            messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=None,
+            model="gemini-1.5-pro",
         )
 
     @pytest.mark.parametrize(
@@ -474,9 +477,12 @@ class TestContextCachingEndpoints:
         assert returned_params == optional_params
         assert returned_cache == "existing_cache_name"
 
-        # Verify cache key was generated with tools and model
+        # Verify cache key was generated with tools, tool_choice and model
         mock_cache_obj.get_cache_key.assert_called_once_with(
-            messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=None,
+            model="gemini-1.5-pro",
         )
 
     @pytest.mark.asyncio
@@ -800,6 +806,546 @@ class TestContextCachingEndpoints:
             # But original tools should still be available for comparison
             assert original_tools == self.sample_tools
 
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    def test_check_and_create_cache_tool_choice_popped_from_optional_params(
+        self, custom_llm_provider
+    ):
+        """tool_choice is popped from optional_params when cached messages exist."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            cached_messages = [self.sample_messages[0]]
+            non_cached_messages = [self.sample_messages[1]]
+            mock_separate.return_value = (cached_messages, non_cached_messages)
+
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}}
+
+            with patch.object(
+                self.context_caching, "check_cache", return_value="existing_cache"
+            ):
+                self.context_caching.check_and_create_cache(
+                    messages=self.sample_messages,
+                    optional_params=optional_params,
+                    api_key="test_key",
+                    api_base=None,
+                    model="gemini-1.5-pro",
+                    client=self.mock_client,
+                    timeout=30.0,
+                    logging_obj=self.mock_logging,
+                    custom_llm_provider=custom_llm_provider,
+                    vertex_project="test_project",
+                    vertex_location="test_location",
+                    vertex_auth_header="vertext_test_token",
+                )
+
+            assert "tool_choice" not in optional_params
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(
+        self, custom_llm_provider
+    ):
+        """tool_choice is NOT popped when there are no cached messages (early return)."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            mock_separate.return_value = ([], self.sample_messages)
+
+            tool_choice = {"functionCallingConfig": {"mode": "AUTO"}}
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = tool_choice
+
+            self.context_caching.check_and_create_cache(
+                messages=self.sample_messages,
+                optional_params=optional_params,
+                api_key="test_key",
+                api_base=None,
+                model="gemini-1.5-pro",
+                client=self.mock_client,
+                timeout=30.0,
+                logging_obj=self.mock_logging,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project="test_project",
+                vertex_location="test_location",
+                vertex_auth_header="vertext_test_token",
+            )
+
+            assert optional_params.get("tool_choice") == tool_choice
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(
+        self, custom_llm_provider
+    ):
+        """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            cached_messages = [self.sample_messages[0]]
+            non_cached_messages = [self.sample_messages[1]]
+            mock_separate.return_value = (cached_messages, non_cached_messages)
+
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}}
+
+            with patch.object(
+                self.context_caching, "async_check_cache", return_value="existing_cache"
+            ):
+                await self.context_caching.async_check_and_create_cache(
+                    messages=self.sample_messages,
+                    optional_params=optional_params,
+                    api_key="test_key",
+                    api_base=None,
+                    model="gemini-1.5-pro",
+                    client=self.mock_async_client,
+                    timeout=30.0,
+                    logging_obj=self.mock_logging,
+                    custom_llm_provider=custom_llm_provider,
+                    vertex_project="test_project",
+                    vertex_location="test_location",
+                    vertex_auth_header="vertext_test_token",
+                )
+
+            assert "tool_choice" not in optional_params
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(
+        self, custom_llm_provider
+    ):
+        """Async equivalent of test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages."""
+        with patch(
+            "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+        ) as mock_separate:
+            mock_separate.return_value = ([], self.sample_messages)
+
+            tool_choice = {"functionCallingConfig": {"mode": "AUTO"}}
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = tool_choice
+
+            await self.context_caching.async_check_and_create_cache(
+                messages=self.sample_messages,
+                optional_params=optional_params,
+                api_key="test_key",
+                api_base=None,
+                model="gemini-1.5-pro",
+                client=self.mock_async_client,
+                timeout=30.0,
+                logging_obj=self.mock_logging,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project="test_project",
+                vertex_location="test_location",
+                vertex_auth_header="vertext_test_token",
+            )
+
+            assert optional_params.get("tool_choice") == tool_choice
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_tool_choice_in_request_body(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """End-to-end: tool_choice ends up as `toolConfig` on the cache-creation HTTP POST body."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None  # cache miss -> create new
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = tool_choice
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        self.mock_client.post.assert_called_once()
+        call_args = self.mock_client.post.call_args
+        assert call_args.kwargs["json"]["tools"] == self.sample_tools
+        assert call_args.kwargs["json"]["toolConfig"] == tool_choice
+        mock_cache_obj.get_cache_key.assert_called_once_with(
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=tool_choice,
+            model="gemini-1.5-pro",
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "async_check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    async def test_async_check_and_create_cache_tool_choice_in_request_body(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """Async equivalent of test_check_and_create_cache_tool_choice_in_request_body."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_async_client.post = AsyncMock(return_value=mock_response)
+
+        tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = tool_choice
+
+        await self.context_caching.async_check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_async_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_async_client.post.call_args
+        assert call_args.kwargs["json"]["tools"] == self.sample_tools
+        assert call_args.kwargs["json"]["toolConfig"] == tool_choice
+        mock_cache_obj.get_cache_key.assert_called_once_with(
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=tool_choice,
+            model="gemini-1.5-pro",
+        )
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_omits_tool_config_when_tool_choice_unset(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """When the caller didn't pass tool_choice, toolConfig must NOT appear in the cache body."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        optional_params = self.sample_optional_params.copy()
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_client.post.call_args
+        assert "tools" in call_args.kwargs["json"]
+        assert "toolConfig" not in call_args.kwargs["json"]
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_tool_choice_function_pin(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """tool_choice as a function-pin dict survives the cache body intact."""
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        function_pin = {
+            "functionCallingConfig": {
+                "mode": "ANY",
+                "allowed_function_names": ["get_current_weather"],
+            }
+        }
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = function_pin
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_client.post.call_args
+        assert call_args.kwargs["json"]["toolConfig"] == function_pin
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj"
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
+    def test_check_and_create_cache_tool_choice_typed_constructor(
+        self,
+        mock_get_token_url,
+        mock_check_cache,
+        mock_transform,
+        mock_cache_obj,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """Exercise the actual ToolConfig(FunctionCallingConfig(...)) constructor that map_tool_choice_values produces.
+
+        ToolConfig / FunctionCallingConfig are TypedDicts (litellm/types/llms/vertex_ai.py:158, 277)
+        so this is functionally identical to the dict-literal tests above at
+        runtime — but exercising the typed constructor pins the test to the
+        same call shape map_tool_choice_values uses and auto-follows if
+        either type ever migrates to a Pydantic model upstream.
+        """
+        from litellm.types.llms.vertex_ai import (
+            FunctionCallingConfig,
+            ToolConfig,
+        )
+
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_cache_obj.get_cache_key.return_value = "test_cache_key"
+        mock_check_cache.return_value = None
+        mock_get_token_url.return_value = ("token", "https://test-url.com")
+        mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []}
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "name": "new_cache_name",
+            "model": "gemini-1.5-pro",
+        }
+        self.mock_client.post.return_value = mock_response
+
+        tool_choice = ToolConfig(
+            functionCallingConfig=FunctionCallingConfig(mode="ANY")
+        )
+        optional_params = self.sample_optional_params.copy()
+        optional_params["tool_choice"] = tool_choice
+
+        self.context_caching.check_and_create_cache(
+            messages=self.sample_messages,
+            optional_params=optional_params,
+            api_key="test_key",
+            api_base=None,
+            model="gemini-1.5-pro",
+            client=self.mock_client,
+            timeout=30.0,
+            logging_obj=self.mock_logging,
+            custom_llm_provider=custom_llm_provider,
+            vertex_project="test_project",
+            vertex_location="test_location",
+            vertex_auth_header="vertext_test_token",
+        )
+
+        call_args = self.mock_client.post.call_args
+        assert call_args.kwargs["json"]["toolConfig"] == tool_choice
+        assert call_args.kwargs["json"]["toolConfig"] == {
+            "functionCallingConfig": {"mode": "ANY"}
+        }
+        mock_cache_obj.get_cache_key.assert_called_once_with(
+            messages=cached_messages,
+            tools=self.sample_tools,
+            tool_choice=tool_choice,
+            model="gemini-1.5-pro",
+        )
+
+    @pytest.mark.parametrize(
+        "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
+    )
+    @patch(
+        "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
+    )
+    @patch.object(ContextCachingEndpoints, "check_cache")
+    def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys(
+        self,
+        mock_check_cache,
+        mock_separate,
+        custom_llm_provider,
+    ):
+        """Two requests with different tool_choice values must produce different cache keys.
+
+        Runs the real local_cache_obj.get_cache_key to verify the hashed
+        output actually differs — mocking it would only prove that distinct
+        arguments are forwarded, not that they produce distinct keys.
+        """
+        cached_messages = [self.sample_messages[0]]
+        non_cached_messages = [self.sample_messages[1]]
+        mock_separate.return_value = (cached_messages, non_cached_messages)
+        mock_check_cache.return_value = "existing_cache"
+
+        auto_tool_choice = {"functionCallingConfig": {"mode": "AUTO"}}
+        any_tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
+        for choice in (auto_tool_choice, any_tool_choice):
+            optional_params = self.sample_optional_params.copy()
+            optional_params["tool_choice"] = choice
+            self.context_caching.check_and_create_cache(
+                messages=self.sample_messages,
+                optional_params=optional_params,
+                api_key="test_key",
+                api_base=None,
+                model="gemini-1.5-pro",
+                client=self.mock_client,
+                timeout=30.0,
+                logging_obj=self.mock_logging,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project="test_project",
+                vertex_location="test_location",
+                vertex_auth_header="vertext_test_token",
+            )
+
+        check_cache_calls = mock_check_cache.call_args_list
+        assert len(check_cache_calls) == 2
+        first_cache_key = check_cache_calls[0].kwargs["cache_key"]
+        second_cache_key = check_cache_calls[1].kwargs["cache_key"]
+        assert first_cache_key != second_cache_key
+
     @pytest.mark.parametrize(
         "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
     )

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py

@@ -0,0 +1,211 @@
+"""
+Tests for the MCP elicitation handler.
+
+Covers the gateway-mode relay logic (`elicitation/create` requests from an
+upstream MCP server being forwarded to the connected downstream client) as
+well as the decline paths used in tool-bridge mode or when the downstream
+client lacks the requested elicitation capability.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from mcp.types import (
+    ElicitRequestFormParams,
+    ElicitRequestURLParams,
+    ElicitResult,
+    ErrorData,
+)
+
+from litellm.proxy._experimental.mcp_server import elicitation_handler
+from litellm.proxy._experimental.mcp_server.elicitation_handler import (
+    _relay_elicitation_to_downstream,
+    handle_elicitation_request,
+)
+
+
+def _form_params(message: str = "fill the form") -> ElicitRequestFormParams:
+    return ElicitRequestFormParams(
+        mode="form",
+        message=message,
+        requestedSchema={"type": "object", "properties": {}},
+    )
+
+
+def _url_params(message: str = "please authorize") -> ElicitRequestURLParams:
+    return ElicitRequestURLParams(
+        mode="url",
+        message=message,
+        url="https://example.com/oauth",
+        elicitationId="elc-1",
+    )
+
+
+def _caps(*, url=True, form=True) -> SimpleNamespace:
+    elicit = SimpleNamespace(
+        url=object() if url else None,
+        form=object() if form else None,
+    )
+    return SimpleNamespace(elicitation=elicit)
+
+
+class TestHandleElicitationRequest:
+    async def test_should_decline_when_no_downstream_session(self):
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_form_params(),
+            downstream_session=None,
+        )
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+
+    async def test_should_relay_to_downstream_when_session_present(self):
+        accepted = ElicitResult(action="accept", content={"name": "ada"})
+        session = SimpleNamespace(elicit_form=AsyncMock(return_value=accepted))
+
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=None,
+        )
+
+        assert result is accepted
+        session.elicit_form.assert_awaited_once()
+
+    async def test_should_return_error_data_when_unavailable(self, monkeypatch):
+        monkeypatch.setattr(elicitation_handler, "MCP_ELICITATION_AVAILABLE", False)
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_form_params(),
+            downstream_session=SimpleNamespace(),
+        )
+        assert isinstance(result, ErrorData)
+        assert "not available" in result.message
+
+    async def test_should_return_error_data_on_unexpected_failure(self):
+        class _ExplodingParams:
+            mode = "form"
+
+            @property
+            def message(self):
+                raise RuntimeError("boom")
+
+        result = await handle_elicitation_request(
+            context=SimpleNamespace(),
+            params=_ExplodingParams(),
+            downstream_session=None,
+        )
+        assert isinstance(result, ErrorData)
+        assert "boom" in result.message
+
+
+class TestRelayElicitationToDownstream:
+    async def test_should_relay_form_mode(self):
+        accepted = ElicitResult(action="accept", content={"name": "ada"})
+        session = SimpleNamespace(elicit_form=AsyncMock(return_value=accepted))
+
+        params = _form_params("collect name")
+        result = await _relay_elicitation_to_downstream(
+            params=params,
+            downstream_session=session,
+            downstream_capabilities=_caps(form=True),
+        )
+
+        assert result is accepted
+        session.elicit_form.assert_awaited_once()
+        _, kwargs = session.elicit_form.call_args
+        assert kwargs["message"] == "collect name"
+        assert kwargs["requestedSchema"] == params.requestedSchema
+
+    async def test_should_relay_url_mode(self):
+        accepted = ElicitResult(action="accept")
+        session = SimpleNamespace(elicit_url=AsyncMock(return_value=accepted))
+
+        result = await _relay_elicitation_to_downstream(
+            params=_url_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(url=True),
+        )
+
+        assert result is accepted
+        session.elicit_url.assert_awaited_once()
+        _, kwargs = session.elicit_url.call_args
+        assert kwargs["url"] == "https://example.com/oauth"
+        assert kwargs["elicitation_id"] == "elc-1"
+
+    async def test_should_use_generic_elicit_for_unknown_param_type(self):
+        accepted = ElicitResult(action="accept")
+        session = SimpleNamespace(elicit=AsyncMock(return_value=accepted))
+
+        # A bare params object that is neither Form nor URL params triggers
+        # the generic fallback path.
+        params = SimpleNamespace(mode="form", message="hi", requestedSchema={})
+        result = await _relay_elicitation_to_downstream(
+            params=params,
+            downstream_session=session,
+            downstream_capabilities=None,
+        )
+
+        assert result is accepted
+        session.elicit.assert_awaited_once()
+
+    async def test_should_decline_when_elicitation_unsupported(self):
+        session = SimpleNamespace(elicit_form=AsyncMock())
+        caps = SimpleNamespace(elicitation=None)
+
+        result = await _relay_elicitation_to_downstream(
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=caps,
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+        session.elicit_form.assert_not_awaited()
+
+    async def test_should_decline_url_mode_when_url_unsupported(self):
+        session = SimpleNamespace(elicit_url=AsyncMock())
+
+        result = await _relay_elicitation_to_downstream(
+            params=_url_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(url=False, form=True),
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+        session.elicit_url.assert_not_awaited()
+
+    async def test_should_decline_form_mode_when_form_unsupported(self):
+        session = SimpleNamespace(elicit_form=AsyncMock())
+
+        result = await _relay_elicitation_to_downstream(
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(url=True, form=False),
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+        session.elicit_form.assert_not_awaited()
+
+    async def test_should_decline_when_downstream_relay_raises(self):
+        session = SimpleNamespace(
+            elicit_form=AsyncMock(side_effect=RuntimeError("transport closed"))
+        )
+
+        result = await _relay_elicitation_to_downstream(
+            params=_form_params(),
+            downstream_session=session,
+            downstream_capabilities=_caps(form=True),
+        )
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "decline"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py

@@ -16,7 +16,6 @@ from typing import Any, Dict, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from fastapi import HTTPException
 
 from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
 from litellm.proxy._types import UserAPIKeyAuth
@@ -549,11 +548,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -593,11 +588,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -643,11 +634,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -703,11 +690,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()
@@ -755,11 +738,7 @@ class TestHookHeaderMergePriority:
         captured_extra_headers: Dict[str, Any] = {}
 
         async def fake_create_mcp_client(
-            server,
-            mcp_auth_header=None,
-            extra_headers=None,
-            stdio_env=None,
-            subject_token=None,
+            server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs
         ):
             captured_extra_headers["value"] = extra_headers
             mock_client = MagicMock()

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py

@@ -0,0 +1,254 @@
+"""
+Tests for the MCP sampling completion pipeline.
+
+Covers building the internal `acompletion` kwargs from MCP request params
+(messages, sampling options, tools, tool choice, metadata), routing the call
+through the proxy router / guardrails, and the end-to-end
+`handle_sampling_create_message` success and error-propagation behaviour.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from mcp.types import CreateMessageResult, ErrorData
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _build_completion_kwargs,
+    _run_guardrails_and_call_llm,
+    handle_sampling_create_message,
+)
+
+
+def _params(**overrides):
+    base = dict(
+        messages=[
+            SimpleNamespace(
+                role="user", content=SimpleNamespace(type="text", text="hi")
+            )
+        ],
+        systemPrompt="be concise",
+        maxTokens=128,
+        temperature=None,
+        stopSequences=None,
+        tools=None,
+        toolChoice=None,
+        metadata=None,
+        modelPreferences=None,
+    )
+    base.update(overrides)
+    return SimpleNamespace(**base)
+
+
+def _passthrough_add_data():
+    async def _add(data, **kwargs):
+        return data
+
+    return _add
+
+
+class TestBuildCompletionKwargs:
+    async def test_should_include_sampling_options_and_tools(self):
+        params = _params(
+            temperature=0.3,
+            stopSequences=["STOP"],
+            tools=[
+                SimpleNamespace(
+                    name="search", description="d", inputSchema={"type": "object"}
+                )
+            ],
+            toolChoice=SimpleNamespace(mode="required"),
+            metadata={"trace": "abc"},
+        )
+        with patch(
+            "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
+            side_effect=_passthrough_add_data(),
+        ):
+            kwargs = await _build_completion_kwargs(
+                params=params,
+                model="gpt-4o",
+                user_api_key_auth=SimpleNamespace(user_id="u1"),
+                raw_headers=None,
+                client_ip=None,
+            )
+
+        assert kwargs["model"] == "gpt-4o"
+        assert kwargs["max_tokens"] == 128
+        assert kwargs["temperature"] == 0.3
+        assert kwargs["stop"] == ["STOP"]
+        assert kwargs["tools"][0]["function"]["name"] == "search"
+        assert kwargs["tool_choice"] == "required"
+        assert kwargs["metadata"]["mcp_metadata"] == {"trace": "abc"}
+        assert kwargs["user"] == "u1"
+        assert kwargs["messages"][0] == {"role": "system", "content": "be concise"}
+
+    async def test_should_omit_optional_fields_when_unset(self):
+        with patch(
+            "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
+            side_effect=_passthrough_add_data(),
+        ):
+            kwargs = await _build_completion_kwargs(
+                params=_params(),
+                model="gpt-4o",
+                user_api_key_auth=SimpleNamespace(user_id=None),
+                raw_headers=None,
+                client_ip=None,
+            )
+
+        assert "temperature" not in kwargs
+        assert "stop" not in kwargs
+        assert "tools" not in kwargs
+        assert "tool_choice" not in kwargs
+        assert kwargs["metadata"] == {}
+
+
+class TestRunGuardrailsAndCallLlm:
+    async def test_should_route_through_llm_router_when_available(self):
+        router = MagicMock()
+        router.acompletion = AsyncMock(return_value="router-response")
+        with (
+            patch("litellm.proxy.proxy_server.proxy_logging_obj", None),
+            patch("litellm.proxy.proxy_server.llm_router", router),
+        ):
+            result = await _run_guardrails_and_call_llm(
+                completion_kwargs={"model": "gpt-4o", "messages": []},
+                user_api_key_auth=SimpleNamespace(),
+            )
+
+        assert result == "router-response"
+        router.acompletion.assert_awaited_once()
+
+    async def test_should_propagate_guardrail_rejection(self):
+        plo = MagicMock()
+        plo.pre_call_hook = AsyncMock(side_effect=ValueError("blocked by guardrail"))
+        with patch("litellm.proxy.proxy_server.proxy_logging_obj", plo):
+            with pytest.raises(ValueError, match="blocked by guardrail"):
+                await _run_guardrails_and_call_llm(
+                    completion_kwargs={"model": "gpt-4o", "messages": []},
+                    user_api_key_auth=SimpleNamespace(),
+                )
+
+
+class TestHandleSamplingCreateMessagePipeline:
+    async def test_should_return_message_result_on_success(self):
+        auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok")
+        response = SimpleNamespace(
+            choices=[
+                SimpleNamespace(
+                    message=SimpleNamespace(
+                        content="the answer is 42", tool_calls=None
+                    ),
+                    finish_reason="stop",
+                )
+            ],
+            model="gpt-4o",
+        )
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs",
+                new_callable=AsyncMock,
+                return_value={"model": "gpt-4o", "messages": []},
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_guardrails_and_call_llm",
+                new_callable=AsyncMock,
+                return_value=response,
+            ),
+        ):
+            result = await handle_sampling_create_message(
+                context=MagicMock(),
+                params=_params(),
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        assert isinstance(result, CreateMessageResult)
+        assert result.content.text == "the answer is 42"
+        assert result.stopReason == "endTurn"
+
+    async def test_should_reraise_known_proxy_exceptions(self):
+        from litellm.exceptions import RateLimitError
+
+        auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok")
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs",
+                new_callable=AsyncMock,
+                side_effect=RateLimitError(
+                    "rate limited", llm_provider="openai", model="gpt-4o"
+                ),
+            ),
+        ):
+            with pytest.raises(RateLimitError):
+                await handle_sampling_create_message(
+                    context=MagicMock(),
+                    params=_params(),
+                    default_model="gpt-4o",
+                    user_api_key_auth=auth,
+                )
+
+    async def test_should_return_error_data_on_unexpected_failure(self):
+        auth = SimpleNamespace(user_id="u1", api_key="sk-test", token="tok")
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._build_completion_kwargs",
+                new_callable=AsyncMock,
+                side_effect=RuntimeError("kaboom"),
+            ),
+        ):
+            result = await handle_sampling_create_message(
+                context=MagicMock(),
+                params=_params(),
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        assert isinstance(result, ErrorData)
+        assert "kaboom" in result.message
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py

@@ -0,0 +1,327 @@
+"""
+Tests for MCP sampling handler model-access enforcement.
+
+Verifies that handle_sampling_create_message and _check_model_access
+enforce the same model-permission checks as regular /chat/completions
+calls, preventing a malicious upstream MCP server from requesting
+inference on models the caller's API key is not authorized to use.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _check_model_access,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_user_api_key_auth(
+    *,
+    models=None,
+    team_id=None,
+    team_model_aliases=None,
+    api_key="sk-test-key",
+    token=None,
+    user_role=None,
+):
+    """Build a minimal UserAPIKeyAuth-like object for tests."""
+    auth = MagicMock()
+    auth.models = models or []
+    auth.team_id = team_id
+    auth.team_model_aliases = team_model_aliases or {}
+    auth.access_group_ids = []
+    auth.api_key = api_key
+    auth.token = token
+    auth.user_role = user_role
+    return auth
+
+
+# ---------------------------------------------------------------------------
+# _check_model_access
+# ---------------------------------------------------------------------------
+
+
+class TestCheckModelAccess:
+    """Tests for the _check_model_access helper that gates sampling requests."""
+
+    @pytest.mark.asyncio
+    async def test_should_return_none_when_no_auth_context(self):
+        """No auth context means no restriction — pass through."""
+        result = await _check_model_access("gpt-4o", user_api_key_auth=None)
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_should_allow_model_when_key_has_access(self):
+        """Key with explicit model access should be allowed."""
+        auth = _make_user_api_key_auth(models=["gpt-4o", "gpt-3.5-turbo"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            return_value=True,
+        ) as mock_check:
+            result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        assert result is None
+        mock_check.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_should_deny_model_when_key_lacks_access(self):
+        """Key without model access should be denied with ErrorData."""
+        from litellm.proxy._types import ProxyException
+
+        auth = _make_user_api_key_auth(models=["gpt-3.5-turbo"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            side_effect=ProxyException(
+                message="key not allowed to access model",
+                type="key_model_access_denied",
+                param="model",
+                code=401,
+            ),
+        ):
+            result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        # Should return ErrorData, not raise
+        assert result is not None
+        assert result.code == -1
+        assert "Model access denied" in result.message
+        assert "gpt-4o" in result.message
+
+    @pytest.mark.asyncio
+    async def test_should_allow_wildcard_model_access(self):
+        """Key with wildcard model access should allow any model."""
+        auth = _make_user_api_key_auth(models=["*"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            return_value=True,
+        ):
+            result = await _check_model_access(
+                "claude-3-opus-20240229", user_api_key_auth=auth
+            )
+
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_should_deny_expensive_model_requested_by_malicious_server(self):
+        """Simulates the attack: malicious MCP server hints at an expensive model
+        the caller's key is restricted from using."""
+        from litellm.proxy._types import ProxyException
+
+        # Key only has access to cheap models
+        auth = _make_user_api_key_auth(models=["gpt-3.5-turbo"])
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            side_effect=ProxyException(
+                message="key not allowed to access model. This key can only access models=['gpt-3.5-turbo']. Tried to access claude-3-opus-20240229",
+                type="key_model_access_denied",
+                param="model",
+                code=401,
+            ),
+        ):
+            result = await _check_model_access(
+                "claude-3-opus-20240229", user_api_key_auth=auth
+            )
+
+        assert result is not None
+        assert result.code == -1
+        assert "claude-3-opus-20240229" in result.message
+
+    @pytest.mark.asyncio
+    async def test_should_deny_empty_oauth_passthrough_placeholder(self):
+        """Regression: process_mcp_request() returns an empty UserAPIKeyAuth()
+        for OAuth2 upstream-token passthrough.  The None check alone is not
+        sufficient — the empty placeholder is truthy but has no api_key, no
+        token, and an empty models list.  can_key_call_model() would treat
+        that as all-model access, letting an OAuth-only user trigger sampling
+        calls on any proxy model without a LiteLLM key or budget."""
+        # Simulate the empty placeholder from process_mcp_request()
+        auth = _make_user_api_key_auth(
+            models=[],
+            api_key=None,
+            token=None,
+            user_role=None,
+        )
+
+        result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        # Must be denied — not passed through to can_key_call_model
+        assert result is not None
+        assert result.code == -1
+        assert "sampling requires a valid LiteLLM" in result.message
+
+    @pytest.mark.asyncio
+    async def test_should_allow_proxy_admin_even_without_api_key(self):
+        """Proxy admins may not have a traditional api_key but should still
+        be allowed to use sampling."""
+        auth = _make_user_api_key_auth(
+            models=[],
+            api_key=None,
+            token=None,
+            user_role="proxy_admin",
+        )
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.can_key_call_model",
+            new_callable=AsyncMock,
+            return_value=True,
+        ):
+            result = await _check_model_access("gpt-4o", user_api_key_auth=auth)
+
+        assert result is None
+
+
+# ---------------------------------------------------------------------------
+# handle_sampling_create_message — auth + budget gating
+# ---------------------------------------------------------------------------
+
+
+class TestSamplingAuthAndBudgetGating:
+
+    @pytest.mark.asyncio
+    async def test_should_deny_when_no_auth_context(self):
+        """Sampling must reject calls with no user_api_key_auth."""
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+
+        params = MagicMock()
+        params.modelPreferences = None
+        params.messages = []
+        params.systemPrompt = None
+        params.maxTokens = 100
+        params.temperature = None
+        params.stopSequences = None
+        params.tools = None
+        params.toolChoice = None
+        params.metadata = None
+
+        result = await handle_sampling_create_message(
+            context=MagicMock(),
+            params=params,
+            default_model="gpt-4o",
+            user_api_key_auth=None,
+        )
+
+        assert result is not None
+        assert result.code == -1
+        assert "authenticated" in result.message.lower()
+
+    @pytest.mark.asyncio
+    async def test_should_run_budget_checks(self):
+        """Sampling must call _run_budget_checks after model access check."""
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+
+        auth = _make_user_api_key_auth(models=["gpt-4o"])
+        params = MagicMock()
+        params.modelPreferences = None
+        params.messages = []
+        params.systemPrompt = None
+        params.maxTokens = 100
+        params.temperature = None
+        params.stopSequences = None
+        params.tools = None
+        params.toolChoice = None
+        params.metadata = None
+
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=None,
+            ) as mock_budget,
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+            patch(
+                "litellm.proxy.proxy_server.llm_router",
+                new=None,
+            ),
+            patch(
+                "litellm.acompletion",
+                new_callable=AsyncMock,
+                return_value=MagicMock(
+                    choices=[
+                        MagicMock(
+                            message=MagicMock(content="hi", tool_calls=None),
+                            finish_reason="stop",
+                        )
+                    ],
+                    model="gpt-4o",
+                ),
+            ),
+        ):
+            await handle_sampling_create_message(
+                context=MagicMock(),
+                params=params,
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        mock_budget.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_should_deny_over_budget_caller(self):
+        """When _run_budget_checks returns ErrorData, sampling must return it."""
+        from mcp.types import ErrorData
+        from litellm.proxy._experimental.mcp_server.sampling_handler import (
+            handle_sampling_create_message,
+        )
+
+        auth = _make_user_api_key_auth(models=["gpt-4o"])
+        params = MagicMock()
+        params.modelPreferences = None
+        params.messages = []
+        params.systemPrompt = None
+        params.maxTokens = 100
+        params.temperature = None
+        params.stopSequences = None
+        params.tools = None
+        params.toolChoice = None
+        params.metadata = None
+
+        budget_error = ErrorData(code=-1, message="ExceededBudget: over limit")
+
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._check_model_access",
+                new_callable=AsyncMock,
+                return_value=None,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._run_budget_checks",
+                new_callable=AsyncMock,
+                return_value=budget_error,
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.sampling_handler._resolve_model_from_preferences",
+                return_value="gpt-4o",
+            ),
+        ):
+            result = await handle_sampling_create_message(
+                context=MagicMock(),
+                params=params,
+                default_model="gpt-4o",
+                user_api_key_auth=auth,
+            )
+
+        assert result is budget_error
+        assert "ExceededBudget" in result.message

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py

@@ -0,0 +1,91 @@
+"""
+Tests for MCP sampling model resolution (hint matching and fallback chain).
+
+`_resolve_model_from_preferences` first tries to match upstream model hints
+against the proxy's available models (direct then substring), then priority
+scoring, then the caller default, the first available model, and finally the
+configured `default_mcp_sampling_model` before raising.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _resolve_model_from_preferences,
+)
+
+
+def _prefs(*, hints=None, cost=None, speed=None, intelligence=None):
+    return SimpleNamespace(
+        hints=hints or [],
+        costPriority=cost,
+        speedPriority=speed,
+        intelligencePriority=intelligence,
+    )
+
+
+class TestHintMatching:
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [{"model_name": "gpt-4o"}, {"model_name": "claude-3"}])
+    def test_should_match_hint_as_substring(self):
+        prefs = _prefs(hints=[SimpleNamespace(name="gpt-4")])
+        assert _resolve_model_from_preferences(prefs) == "gpt-4o"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", ["gpt-4o", "claude-3"])
+    def test_should_match_hint_against_string_model_list_entries(self):
+        prefs = _prefs(hints=[SimpleNamespace(name="claude-3")])
+        assert _resolve_model_from_preferences(prefs) == "claude-3"
+
+    @patch("litellm.model_list", None)
+    def test_should_use_router_model_names(self):
+        router = MagicMock()
+        router.get_model_names.return_value = ["router-gpt", "router-claude"]
+        with patch("litellm.proxy.proxy_server.llm_router", router):
+            prefs = _prefs(hints=[SimpleNamespace(name="router-claude")])
+            assert _resolve_model_from_preferences(prefs) == "router-claude"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [{"model_name": "gpt-4o"}])
+    def test_should_skip_hint_without_name(self):
+        prefs = _prefs(hints=[SimpleNamespace()])  # hint has no `.name`
+        assert (
+            _resolve_model_from_preferences(prefs, default_model="gpt-4o") == "gpt-4o"
+        )
+
+
+class TestFallbackChain:
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list", [{"model_name": "first-model"}, {"model_name": "second"}]
+    )
+    def test_should_fall_back_to_first_available_when_no_default(self):
+        prefs = _prefs(hints=[SimpleNamespace(name="no-such")])
+        assert _resolve_model_from_preferences(prefs) == "first-model"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [])
+    def test_should_use_configured_default_sampling_model(self, monkeypatch):
+        import litellm
+
+        monkeypatch.setattr(
+            litellm, "default_mcp_sampling_model", "fallback-model", raising=False
+        )
+        prefs = _prefs()
+        assert _resolve_model_from_preferences(prefs) == "fallback-model"
+
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch("litellm.model_list", [])
+    def test_should_raise_when_nothing_resolvable(self, monkeypatch):
+        import litellm
+
+        monkeypatch.setattr(litellm, "default_mcp_sampling_model", None, raising=False)
+        prefs = _prefs()
+        with pytest.raises(ValueError, match="No model could be resolved"):
+            _resolve_model_from_preferences(prefs)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py

@@ -0,0 +1,248 @@
+"""
+Tests for MCP sampling handler priority-based model selection.
+
+Verifies that _resolve_model_from_preferences honours costPriority,
+speedPriority, and intelligencePriority when hints don't match,
+per the MCP spec.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _has_priorities,
+    _resolve_model_from_preferences,
+    _select_model_by_priority,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _prefs(*, hints=None, cost=None, speed=None, intelligence=None):
+    """Build a minimal ModelPreferences-like object."""
+    return SimpleNamespace(
+        hints=hints or [],
+        costPriority=cost,
+        speedPriority=speed,
+        intelligencePriority=intelligence,
+    )
+
+
+# Model info stubs keyed by model name
+_MODEL_INFO = {
+    "gpt-3.5-turbo": {
+        "input_cost_per_token": 0.0000005,
+        "output_cost_per_token": 0.0000015,
+        "max_output_tokens": 4096,
+        "max_tokens": 4096,
+        "output_tokens_per_second": 50.0,
+    },
+    "gpt-4o": {
+        "input_cost_per_token": 0.0000025,
+        "output_cost_per_token": 0.0000100,
+        "max_output_tokens": 16384,
+        "max_tokens": 128000,
+        "output_tokens_per_second": 60.0,
+    },
+    "claude-3-opus": {
+        "input_cost_per_token": 0.0000150,
+        "output_cost_per_token": 0.0000750,
+        "max_output_tokens": 4096,
+        "max_tokens": 200000,
+        "output_tokens_per_second": 20.0,
+    },
+    "gpt-4o-mini": {
+        "input_cost_per_token": 0.00000015,
+        "output_cost_per_token": 0.0000006,
+        "max_output_tokens": 16384,
+        "max_tokens": 128000,
+        "output_tokens_per_second": 100.0,
+    },
+}
+
+
+def _mock_get_model_info(model, **kwargs):
+    """Mock litellm.get_model_info using our test data."""
+    if model in _MODEL_INFO:
+        return _MODEL_INFO[model]
+    raise Exception(f"Unknown model: {model}")
+
+
+# ---------------------------------------------------------------------------
+# _has_priorities
+# ---------------------------------------------------------------------------
+
+
+class TestHasPriorities:
+    def test_should_return_false_when_no_priorities_set(self):
+        prefs = _prefs()
+        assert _has_priorities(prefs) is False
+
+    def test_should_return_false_when_all_zero(self):
+        prefs = _prefs(cost=0, speed=0, intelligence=0)
+        assert _has_priorities(prefs) is False
+
+    def test_should_return_true_when_cost_set(self):
+        prefs = _prefs(cost=0.8)
+        assert _has_priorities(prefs) is True
+
+    def test_should_return_true_when_intelligence_set(self):
+        prefs = _prefs(intelligence=0.5)
+        assert _has_priorities(prefs) is True
+
+
+# ---------------------------------------------------------------------------
+# _select_model_by_priority
+# ---------------------------------------------------------------------------
+
+
+class TestSelectModelByPriority:
+    """Tests for the priority-based scoring logic."""
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_prefer_cheapest_when_cost_priority_high(self, _mock):
+        """High costPriority should select the cheapest model."""
+        prefs = _prefs(cost=1.0, speed=0, intelligence=0)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o-mini has the lowest combined cost
+        assert result == "gpt-4o-mini"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_prefer_smartest_when_intelligence_priority_high(self, _mock):
+        """High intelligencePriority should select the model with highest max_output_tokens."""
+        prefs = _prefs(cost=0, speed=0, intelligence=1.0)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o and gpt-4o-mini both have 16384 max_output_tokens (tied)
+        # Either is acceptable
+        assert result in ("gpt-4o", "gpt-4o-mini")
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_balance_cost_and_intelligence(self, _mock):
+        """Balanced priorities should pick a middle-ground model."""
+        prefs = _prefs(cost=0.5, speed=0, intelligence=0.5)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o-mini is cheap AND has high max_output_tokens → best balance
+        assert result == "gpt-4o-mini"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_prefer_fastest_when_speed_priority_high(self, _mock):
+        """High speedPriority should prefer cheaper (faster proxy) models."""
+        prefs = _prefs(cost=0, speed=1.0, intelligence=0)
+        models = ["gpt-3.5-turbo", "gpt-4o", "claude-3-opus", "gpt-4o-mini"]
+        result = _select_model_by_priority(models, prefs)
+        # gpt-4o-mini has lowest cost → fastest proxy
+        assert result == "gpt-4o-mini"
+
+    @patch(
+        "litellm.get_model_info",
+        side_effect=lambda m, **kw: (_ for _ in ()).throw(Exception("no info")),
+    )
+    def test_should_return_none_when_no_model_info(self, _mock):
+        """If get_model_info fails for all models, return None."""
+        prefs = _prefs(cost=1.0)
+        models = ["unknown-model-1", "unknown-model-2"]
+        result = _select_model_by_priority(models, prefs)
+        assert result is None
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    def test_should_handle_single_model(self, _mock):
+        """Single model should always be returned regardless of priorities."""
+        prefs = _prefs(cost=1.0, intelligence=1.0)
+        result = _select_model_by_priority(["gpt-4o"], prefs)
+        assert result == "gpt-4o"
+
+    def test_speed_priority_is_neutral_when_no_tps_data(self):
+        """When no candidate exposes output_tokens_per_second, speedPriority
+        must not fall back to context-window size as a latency proxy: that
+        biased selection toward the smallest-context model regardless of real
+        speed. With a neutral score the tie resolves to the first candidate,
+        so the larger-context model listed first is kept."""
+        no_tps_info = {
+            "big-ctx": {
+                "input_cost_per_token": 0.0,
+                "output_cost_per_token": 0.0,
+                "max_output_tokens": 100000,
+                "max_tokens": 100000,
+            },
+            "small-ctx": {
+                "input_cost_per_token": 0.0,
+                "output_cost_per_token": 0.0,
+                "max_output_tokens": 1000,
+                "max_tokens": 1000,
+            },
+        }
+
+        def info(model, **kwargs):
+            return no_tps_info[model]
+
+        with patch("litellm.get_model_info", side_effect=info):
+            prefs = _prefs(speed=1.0)
+            # The inverse-max_output proxy would pick "small-ctx" here; a
+            # neutral score keeps the first candidate.
+            assert _select_model_by_priority(["big-ctx", "small-ctx"], prefs) == (
+                "big-ctx"
+            )
+
+
+# ---------------------------------------------------------------------------
+# _resolve_model_from_preferences — priority integration
+# ---------------------------------------------------------------------------
+
+
+class TestResolveModelPriorityIntegration:
+    """End-to-end tests for priority selection within _resolve_model_from_preferences."""
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list",
+        [
+            {"model_name": "gpt-3.5-turbo"},
+            {"model_name": "gpt-4o"},
+            {"model_name": "gpt-4o-mini"},
+        ],
+    )
+    def test_should_use_priority_when_hints_empty(self, _mock_info):
+        """With no hints but priorities set, should use priority-based selection."""
+        prefs = _prefs(cost=1.0, speed=0, intelligence=0)
+        result = _resolve_model_from_preferences(prefs, default_model="gpt-4o")
+        # Should pick cheapest, NOT fall through to default_model
+        assert result == "gpt-4o-mini"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list",
+        [
+            {"model_name": "gpt-3.5-turbo"},
+            {"model_name": "gpt-4o"},
+            {"model_name": "gpt-4o-mini"},
+        ],
+    )
+    def test_should_skip_priority_when_no_priorities_set(self, _mock_info):
+        """With no priorities set, should fall through to default_model."""
+        prefs = _prefs()  # no priorities
+        result = _resolve_model_from_preferences(prefs, default_model="gpt-4o")
+        assert result == "gpt-4o"
+
+    @patch("litellm.get_model_info", side_effect=_mock_get_model_info)
+    @patch("litellm.proxy.proxy_server.llm_router", None)
+    @patch(
+        "litellm.model_list",
+        [
+            {"model_name": "gpt-3.5-turbo"},
+            {"model_name": "gpt-4o"},
+            {"model_name": "gpt-4o-mini"},
+        ],
+    )
+    def test_should_prefer_hint_over_priority(self, _mock_info):
+        """Hints should take precedence over priority-based selection."""
+        hints = [SimpleNamespace(name="gpt-4o")]
+        prefs = _prefs(hints=hints, cost=1.0)  # cost says cheap, but hint says gpt-4o
+        result = _resolve_model_from_preferences(prefs, default_model="gpt-3.5-turbo")
+        assert result == "gpt-4o"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py

@@ -0,0 +1,147 @@
+"""
+Tests for _build_sampling_request header forwarding.
+
+Verifies that the synthetic FastAPI Request built for sampling sub-calls
+correctly propagates the original MCP connection's headers and client IP
+so that header-dependent guardrails, routing hooks, and trace correlation
+function correctly.
+"""
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _build_sampling_request,
+)
+
+
+class TestBuildSamplingRequest:
+    """Tests for the _build_sampling_request helper."""
+
+    def test_should_include_content_type_by_default(self):
+        """Even with no raw headers, content-type must be present."""
+        req = _build_sampling_request()
+        headers = dict(req.headers)
+        assert headers.get("content-type") == "application/json"
+
+    def test_should_forward_raw_headers(self):
+        """Headers from the original MCP connection should be forwarded."""
+        raw = {
+            "x-litellm-tags": "tag1,tag2",
+            "x-litellm-trace-id": "trace-abc-123",
+            "user-agent": "MCP-Client/1.0",
+            "authorization": "Bearer sk-test",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+
+        assert headers.get("x-litellm-tags") == "tag1,tag2"
+        assert headers.get("x-litellm-trace-id") == "trace-abc-123"
+        assert headers.get("user-agent") == "MCP-Client/1.0"
+        assert headers.get("authorization") == "Bearer sk-test"
+
+    def test_should_skip_hop_by_hop_headers(self):
+        """content-length and transfer-encoding should not be forwarded."""
+        raw = {
+            "content-length": "42",
+            "transfer-encoding": "chunked",
+            "x-custom": "keep-me",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+
+        assert "content-length" not in headers
+        assert "transfer-encoding" not in headers
+        assert headers.get("x-custom") == "keep-me"
+
+    def test_should_not_duplicate_content_type(self):
+        """If raw_headers includes content-type, don't add it twice."""
+        raw = {"content-type": "text/plain"}
+        req = _build_sampling_request(raw_headers=raw)
+        # Count how many content-type headers are present
+        ct_count = sum(1 for k, _ in req.scope["headers"] if k == b"content-type")
+        assert ct_count == 1
+
+    def test_should_inject_client_ip_as_x_forwarded_for(self):
+        """client_ip should be injected as x-forwarded-for."""
+        req = _build_sampling_request(client_ip="10.0.0.42")
+        headers = dict(req.headers)
+        assert headers.get("x-forwarded-for") == "10.0.0.42"
+
+    def test_should_not_override_existing_x_forwarded_for(self):
+        """Caller-supplied x-forwarded-for is stripped; resolved client_ip wins."""
+        raw = {"x-forwarded-for": "192.168.1.1"}
+        req = _build_sampling_request(raw_headers=raw, client_ip="10.0.0.42")
+        headers = dict(req.headers)
+        assert headers.get("x-forwarded-for") == "10.0.0.42"
+
+    def test_should_set_correct_path(self):
+        """The synthetic request should have the sampling path."""
+        req = _build_sampling_request()
+        assert req.scope["path"] == "/mcp/sampling/createMessage"
+
+    def test_server_should_default_to_litellm_port(self):
+        """Server tuple should use port 4000 (LiteLLM default), not 0."""
+        req = _build_sampling_request()
+        _host, _port = req.scope["server"]
+        assert _port == 4000, f"Expected default LiteLLM port 4000, got {_port}"
+
+    def test_should_populate_client_tuple_from_client_ip(self):
+        """request.client.host must return the real client IP for
+        IP-based routing and guardrails."""
+        req = _build_sampling_request(client_ip="10.0.0.42")
+        assert req.scope.get("client") is not None
+        assert req.scope["client"][0] == "10.0.0.42"
+        # Verify request.client.host works (Starlette Address)
+        assert req.client is not None
+        assert req.client.host == "10.0.0.42"
+
+    def test_should_not_set_client_when_no_ip(self):
+        """If no client_ip is provided, client should not be in scope."""
+        req = _build_sampling_request()
+        assert "client" not in req.scope
+
+    def test_should_skip_all_hop_by_hop_headers(self):
+        """All hop-by-hop headers must be filtered, not just content-length
+        and transfer-encoding."""
+        raw = {
+            "content-length": "42",
+            "transfer-encoding": "chunked",
+            "connection": "keep-alive",
+            "keep-alive": "timeout=5",
+            "upgrade": "websocket",
+            "te": "trailers",
+            "trailer": "Expires",
+            "x-custom": "keep-me",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+
+        for hop_header in [
+            "content-length",
+            "transfer-encoding",
+            "connection",
+            "keep-alive",
+            "upgrade",
+            "te",
+            "trailer",
+        ]:
+            assert (
+                hop_header not in headers
+            ), f"Hop-by-hop header '{hop_header}' should be filtered"
+        assert headers.get("x-custom") == "keep-me"
+
+    def test_should_forward_traceparent_header(self):
+        """traceparent header must be forwarded for trace correlation."""
+        raw = {
+            "traceparent": "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01",
+        }
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+        assert headers.get("traceparent") == (
+            "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01"
+        )
+
+    def test_should_forward_x_litellm_api_key(self):
+        """x-litellm-api-key header must be forwarded for auth."""
+        raw = {"x-litellm-api-key": "sk-proxy-key-123"}
+        req = _build_sampling_request(raw_headers=raw)
+        headers = dict(req.headers)
+        assert headers.get("x-litellm-api-key") == "sk-proxy-key-123"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py

@@ -0,0 +1,180 @@
+"""
+Tests for MCP sampling handler response/tool conversion.
+
+Covers the translation of a LiteLLM completion response back into MCP
+`CreateMessageResult` / `CreateMessageResultWithTools`, plus the helpers that
+convert MCP tool definitions, tool-choice modes, and image/audio content into
+OpenAI request format.
+"""
+
+import json
+from types import SimpleNamespace
+
+from mcp.types import (
+    CreateMessageResult,
+    CreateMessageResultWithTools,
+    ErrorData,
+    TextContent,
+    ToolUseContent,
+)
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _convert_mcp_content_to_openai,
+    _convert_mcp_tool_choice_to_openai,
+    _convert_mcp_tools_to_openai,
+    _convert_openai_response_to_mcp_result,
+    _convert_single_content,
+)
+
+
+def _tool_call(*, call_id: str, name: str, arguments):
+    return SimpleNamespace(
+        id=call_id, function=SimpleNamespace(name=name, arguments=arguments)
+    )
+
+
+def _response(*, content=None, tool_calls=None, finish_reason="stop", model="gpt-4o"):
+    message = SimpleNamespace(content=content, tool_calls=tool_calls)
+    choice = SimpleNamespace(message=message, finish_reason=finish_reason)
+    return SimpleNamespace(choices=[choice], model=model)
+
+
+class TestConvertOpenAIResponseToMcpResult:
+    def test_should_return_error_data_when_no_choices(self):
+        response = SimpleNamespace(choices=[], model="gpt-4o")
+        result = _convert_openai_response_to_mcp_result(response, "gpt-4o")
+        assert isinstance(result, ErrorData)
+        assert "no choices" in result.message.lower()
+
+    def test_should_convert_plain_text_response(self):
+        result = _convert_openai_response_to_mcp_result(
+            _response(content="hello world"), "gpt-4o"
+        )
+        assert isinstance(result, CreateMessageResult)
+        assert isinstance(result.content, TextContent)
+        assert result.content.text == "hello world"
+        assert result.role == "assistant"
+        assert result.stopReason == "endTurn"
+
+    def test_should_map_length_finish_reason_to_max_tokens(self):
+        result = _convert_openai_response_to_mcp_result(
+            _response(content="truncated", finish_reason="length"), "gpt-4o"
+        )
+        assert result.stopReason == "maxTokens"
+
+    def test_should_prefer_actual_model_from_response(self):
+        result = _convert_openai_response_to_mcp_result(
+            _response(content="hi", model="gpt-4o-2024-08-06"), "gpt-4o"
+        )
+        assert result.model == "gpt-4o-2024-08-06"
+
+    def test_should_convert_tool_calls_response(self):
+        tc = _tool_call(
+            call_id="call_1",
+            name="get_weather",
+            arguments=json.dumps({"city": "NYC"}),
+        )
+        result = _convert_openai_response_to_mcp_result(
+            _response(content=None, tool_calls=[tc], finish_reason="tool_calls"),
+            "gpt-4o",
+        )
+        assert isinstance(result, CreateMessageResultWithTools)
+        assert result.stopReason == "toolUse"
+        tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)]
+        assert len(tool_uses) == 1
+        assert tool_uses[0].name == "get_weather"
+        assert tool_uses[0].id == "call_1"
+        assert tool_uses[0].input == {"city": "NYC"}
+
+    def test_should_keep_text_alongside_tool_calls(self):
+        tc = _tool_call(call_id="call_1", name="search", arguments="{}")
+        result = _convert_openai_response_to_mcp_result(
+            _response(
+                content="let me check", tool_calls=[tc], finish_reason="tool_calls"
+            ),
+            "gpt-4o",
+        )
+        texts = [c for c in result.content if isinstance(c, TextContent)]
+        assert texts and texts[0].text == "let me check"
+
+    def test_should_wrap_unparsable_tool_arguments_as_raw(self):
+        tc = _tool_call(call_id="call_1", name="bad", arguments="not-json{")
+        result = _convert_openai_response_to_mcp_result(
+            _response(tool_calls=[tc], finish_reason="tool_calls"), "gpt-4o"
+        )
+        tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)]
+        assert tool_uses[0].input == {"raw": "not-json{"}
+
+
+class TestConvertMcpToolsToOpenAI:
+    def test_should_return_none_when_no_tools(self):
+        assert _convert_mcp_tools_to_openai(None) is None
+
+    def test_should_convert_tool_with_schema(self):
+        schema = {"type": "object", "properties": {"q": {"type": "string"}}}
+        tool = SimpleNamespace(
+            name="search", description="search the web", inputSchema=schema
+        )
+        result = _convert_mcp_tools_to_openai([tool])
+        assert result == [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search",
+                    "description": "search the web",
+                    "parameters": schema,
+                },
+            }
+        ]
+
+    def test_should_default_description_and_parameters(self):
+        tool = SimpleNamespace(name="noop", description=None, inputSchema=None)
+        result = _convert_mcp_tools_to_openai([tool])
+        fn = result[0]["function"]
+        assert fn["description"] == ""
+        assert fn["parameters"] == {"type": "object", "properties": {}}
+
+
+class TestConvertMcpToolChoiceToOpenAI:
+    def test_should_return_none_when_no_choice(self):
+        assert _convert_mcp_tool_choice_to_openai(None) is None
+
+    def test_should_map_known_modes(self):
+        for mode in ("auto", "required", "none"):
+            choice = SimpleNamespace(mode=mode)
+            assert _convert_mcp_tool_choice_to_openai(choice) == mode
+
+    def test_should_default_unknown_mode_to_auto(self):
+        choice = SimpleNamespace(mode="banana")
+        assert _convert_mcp_tool_choice_to_openai(choice) == "auto"
+
+
+class TestConvertImageAndAudioContent:
+    def test_should_convert_image_to_data_uri(self):
+        content = SimpleNamespace(type="image", data="aGVsbG8=", mimeType="image/jpeg")
+        result = _convert_single_content(content)
+        assert result == {
+            "type": "image_url",
+            "image_url": {"url": "data:image/jpeg;base64,aGVsbG8="},
+        }
+
+    def test_should_map_audio_mime_to_format(self):
+        content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/mp3")
+        result = _convert_single_content(content)
+        assert result["type"] == "input_audio"
+        assert result["input_audio"] == {"data": "Zm9v", "format": "mp3"}
+
+    def test_should_default_unknown_audio_mime_to_wav(self):
+        content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/weird")
+        result = _convert_single_content(content)
+        assert result["input_audio"]["format"] == "wav"
+
+    def test_should_flatten_list_content(self):
+        items = [
+            SimpleNamespace(type="text", text="a"),
+            SimpleNamespace(type="image", data="x", mimeType="image/png"),
+        ]
+        result = _convert_mcp_content_to_openai(items)
+        assert isinstance(result, list)
+        assert result[0] == {"type": "text", "text": "a"}
+        assert result[1]["type"] == "image_url"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py

@@ -0,0 +1,312 @@
+"""
+Tests for MCP sampling handler tool_use / tool_result content conversion.
+
+Verifies that multi-turn tool-calling conversations from upstream MCP
+servers are faithfully converted to OpenAI format instead of being
+reduced to lossy plain-text stubs.
+"""
+
+import json
+from types import SimpleNamespace
+from typing import Any, Dict
+
+from litellm.proxy._experimental.mcp_server.sampling_handler import (
+    _convert_mcp_messages_to_openai,
+    _convert_single_content,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers — lightweight MCP type stand-ins
+# ---------------------------------------------------------------------------
+
+
+def _text(text: str) -> SimpleNamespace:
+    return SimpleNamespace(type="text", text=text)
+
+
+def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleNamespace:
+    return SimpleNamespace(type="tool_use", name=name, id=tool_id, input=input_data)
+
+
+def _tool_result(
+    *, tool_use_id: str, content: Any = None, is_error: bool = False
+) -> SimpleNamespace:
+    if content is None:
+        content = []
+    return SimpleNamespace(
+        type="tool_result", toolUseId=tool_use_id, content=content, isError=is_error
+    )
+
+
+def _sampling_msg(role: str, content: Any) -> SimpleNamespace:
+    return SimpleNamespace(role=role, content=content)
+
+
+# ---------------------------------------------------------------------------
+# _convert_single_content — tool_use
+# ---------------------------------------------------------------------------
+
+
+class TestConvertSingleContentToolUse:
+    """Tests for the tool_use branch of _convert_single_content."""
+
+    def test_should_produce_function_call_dict(self):
+        """tool_use must produce a proper function-call dict, not a text stub."""
+        tu = _tool_use(name="get_weather", tool_id="call_123", input_data={"city": "NYC"})
+        result = _convert_single_content(tu)
+
+        assert result["_marker_type"] == "tool_use"
+        assert result["type"] == "function"
+        assert result["id"] == "call_123"
+        assert result["function"]["name"] == "get_weather"
+        assert json.loads(result["function"]["arguments"]) == {"city": "NYC"}
+
+    def test_should_not_produce_text_stub(self):
+        """Regression: the old code produced '[Tool call: get_weather]'."""
+        tu = _tool_use(name="get_weather", tool_id="call_1", input_data={})
+        result = _convert_single_content(tu)
+
+        # Must NOT be a text content part
+        assert result.get("type") != "text"
+        assert "Tool call" not in str(result)
+
+    def test_should_handle_empty_input(self):
+        tu = _tool_use(name="no_args_tool", tool_id="call_2", input_data={})
+        result = _convert_single_content(tu)
+
+        assert json.loads(result["function"]["arguments"]) == {}
+
+
+# ---------------------------------------------------------------------------
+# _convert_single_content — tool_result
+# ---------------------------------------------------------------------------
+
+
+class TestConvertSingleContentToolResult:
+    """Tests for the tool_result branch of _convert_single_content."""
+
+    def test_should_produce_tool_role_message(self):
+        """tool_result must produce a tool-role dict, not a text content part."""
+        tr = _tool_result(
+            tool_use_id="call_123",
+            content=[_text("Temperature: 72°F")],
+        )
+        result = _convert_single_content(tr)
+
+        assert result["_marker_type"] == "tool_result"
+        assert result["role"] == "tool"
+        assert result["tool_call_id"] == "call_123"
+        assert "72°F" in result["content"]
+
+    def test_should_handle_empty_content(self):
+        tr = _tool_result(tool_use_id="call_456", content=[])
+        result = _convert_single_content(tr)
+
+        assert result["role"] == "tool"
+        assert result["tool_call_id"] == "call_456"
+        assert result["content"] == ""
+
+    def test_should_concatenate_multiple_text_parts(self):
+        tr = _tool_result(
+            tool_use_id="call_789",
+            content=[_text("Line 1"), _text("Line 2")],
+        )
+        result = _convert_single_content(tr)
+        assert "Line 1" in result["content"]
+        assert "Line 2" in result["content"]
+
+
+# ---------------------------------------------------------------------------
+# _convert_mcp_messages_to_openai — multi-turn tool calling
+# ---------------------------------------------------------------------------
+
+
+class TestConvertMcpMessagesMultiTurnTools:
+    """End-to-end tests for multi-turn tool-calling message sequences."""
+
+    def test_should_convert_assistant_tool_use_to_tool_calls_array(self):
+        """An assistant message with tool_use content should produce
+        a proper tool_calls array, not a text stub."""
+        messages = [
+            _sampling_msg("assistant", _tool_use(
+                name="search", tool_id="call_1", input_data={"query": "LiteLLM"}
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "assistant"
+        assert "tool_calls" in msg
+        assert len(msg["tool_calls"]) == 1
+        tc = msg["tool_calls"][0]
+        assert tc["function"]["name"] == "search"
+        assert tc["id"] == "call_1"
+
+    def test_should_convert_user_tool_result_to_tool_role_message(self):
+        """A user message with tool_result content should produce
+        a separate role='tool' message."""
+        messages = [
+            _sampling_msg("user", _tool_result(
+                tool_use_id="call_1",
+                content=[_text("Found 42 results")],
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "tool"
+        assert msg["tool_call_id"] == "call_1"
+        assert "42 results" in msg["content"]
+
+    def test_should_handle_full_tool_calling_round_trip(self):
+        """Simulate a complete tool-calling conversation:
+        user → assistant(tool_use) → user(tool_result) → assistant(text)
+        """
+        messages = [
+            _sampling_msg("user", _text("What's the weather in NYC?")),
+            _sampling_msg("assistant", _tool_use(
+                name="get_weather", tool_id="call_w1",
+                input_data={"city": "NYC"},
+            )),
+            _sampling_msg("user", _tool_result(
+                tool_use_id="call_w1",
+                content=[_text("72°F, sunny")],
+            )),
+            _sampling_msg("assistant", _text("It's 72°F and sunny in NYC!")),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 4
+
+        # 1. User message
+        assert result[0]["role"] == "user"
+
+        # 2. Assistant with tool_calls
+        assert result[1]["role"] == "assistant"
+        assert "tool_calls" in result[1]
+        assert result[1]["tool_calls"][0]["function"]["name"] == "get_weather"
+
+        # 3. Tool result
+        assert result[2]["role"] == "tool"
+        assert result[2]["tool_call_id"] == "call_w1"
+
+        # 4. Final assistant text
+        assert result[3]["role"] == "assistant"
+        assert "72°F" in str(result[3]["content"])
+
+    def test_should_handle_mixed_text_and_tool_use_in_assistant(self):
+        """An assistant message with both text and tool_use content."""
+        messages = [
+            _sampling_msg("assistant", [
+                _text("Let me check that for you."),
+                _tool_use(name="lookup", tool_id="call_lu1", input_data={"id": 42}),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "assistant"
+        assert "tool_calls" in msg
+        assert msg["tool_calls"][0]["function"]["name"] == "lookup"
+        # Text content should also be present
+        assert msg.get("content") is not None
+
+    def test_should_handle_multiple_tool_uses_in_single_message(self):
+        """Multiple tool_use items in a single assistant message → multiple tool_calls."""
+        messages = [
+            _sampling_msg("assistant", [
+                _tool_use(name="tool_a", tool_id="call_a", input_data={}),
+                _tool_use(name="tool_b", tool_id="call_b", input_data={"x": 1}),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert len(msg["tool_calls"]) == 2
+        names = {tc["function"]["name"] for tc in msg["tool_calls"]}
+        assert names == {"tool_a", "tool_b"}
+
+    def test_should_handle_multiple_tool_results_in_single_message(self):
+        """Multiple tool_result items in a single user message → multiple tool messages."""
+        messages = [
+            _sampling_msg("user", [
+                _tool_result(tool_use_id="call_a", content=[_text("Result A")]),
+                _tool_result(tool_use_id="call_b", content=[_text("Result B")]),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 2
+        assert all(m["role"] == "tool" for m in result)
+        ids = {m["tool_call_id"] for m in result}
+        assert ids == {"call_a", "call_b"}
+
+    def test_should_preserve_system_prompt(self):
+        """System prompt should still be emitted first."""
+        messages = [_sampling_msg("user", _text("Hi"))]
+        result = _convert_mcp_messages_to_openai(
+            messages, system_prompt="You are helpful."
+        )
+
+        assert result[0]["role"] == "system"
+        assert result[0]["content"] == "You are helpful."
+
+
+# ---------------------------------------------------------------------------
+# _convert_mcp_messages_to_openai — marker hoisting on unexpected roles
+# ---------------------------------------------------------------------------
+
+
+class TestConvertMcpMessagesMarkerHoisting:
+    """The role-matched fast paths only fire for assistant/tool_use and
+    user/tool_result. Content that arrives on an unexpected role must still
+    be hoisted to the correct message position by the generic fallback,
+    not silently dropped or embedded inline as a content part."""
+
+    def test_should_hoist_tool_use_arriving_on_user_role(self):
+        messages = [
+            _sampling_msg("user", _tool_use(
+                name="search", tool_id="call_1", input_data={"q": "x"}
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        assert result[0]["role"] == "assistant"
+        assert result[0]["tool_calls"][0]["function"]["name"] == "search"
+
+    def test_should_hoist_tool_result_arriving_on_assistant_role(self):
+        messages = [
+            _sampling_msg("assistant", _tool_result(
+                tool_use_id="call_1", content=[_text("done")]
+            )),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        assert result[0]["role"] == "tool"
+        assert result[0]["tool_call_id"] == "call_1"
+        assert "done" in result[0]["content"]
+
+    def test_should_keep_text_when_hoisting_tool_use_on_user_role(self):
+        messages = [
+            _sampling_msg("user", [
+                _text("here you go"),
+                _tool_use(name="lookup", tool_id="call_2", input_data={}),
+            ]),
+        ]
+        result = _convert_mcp_messages_to_openai(messages)
+
+        assert len(result) == 1
+        msg = result[0]
+        assert msg["role"] == "assistant"
+        assert msg["tool_calls"][0]["function"]["name"] == "lookup"
+        assert any(
+            isinstance(p, dict) and p.get("text") == "here you go"
+            for p in msg["content"]
+        )

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py

@@ -1,5 +1,4 @@
 import asyncio
-import contextlib
 import contextvars
 from datetime import datetime, timedelta
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -894,7 +893,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         if server.name == "working_server":
             # Working server returns tools
@@ -1000,7 +999,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # All servers fail
         raise Exception(f"Server {server.name} connection failed")
@@ -1122,8 +1121,8 @@ async def test_concurrent_initialize_session_managers():
     # Reset state before test
     original_initialized = mcp_server._SESSION_MANAGERS_INITIALIZED
     original_session_cm = mcp_server._session_manager_cm
-    original_session_stateful_cm = mcp_server._session_manager_stateful_cm
-    original_sse_session_cm = mcp_server._sse_session_manager_cm
+    original_stateful_cm = mcp_server._session_manager_stateful_cm
+    original_sse_cm = mcp_server._sse_session_manager_cm
     original_cleanup_task = mcp_server._stateful_auth_context_cleanup_task
 
     try:
@@ -1131,30 +1130,38 @@ async def test_concurrent_initialize_session_managers():
         mcp_server._session_manager_cm = None
         mcp_server._session_manager_stateful_cm = None
         mcp_server._sse_session_manager_cm = None
-        mcp_server._stateful_auth_context_cleanup_task = None
 
-        # Mock the session managers to avoid actual MCP initialization
+        # Create mock context managers for all three session managers
+        mock_cm_stateless = AsyncMock()
+        mock_cm_stateless.__aenter__ = AsyncMock()
+        mock_cm_stateless.__aexit__ = AsyncMock()
+
+        mock_cm_stateful = AsyncMock()
+        mock_cm_stateful.__aenter__ = AsyncMock()
+        mock_cm_stateful.__aexit__ = AsyncMock()
+
+        mock_cm_sse = AsyncMock()
+        mock_cm_sse.__aenter__ = AsyncMock()
+        mock_cm_sse.__aexit__ = AsyncMock()
+
         with (
-            patch(
-                "litellm.proxy._experimental.mcp_server.server.session_manager_stateless"
-            ) as mock_session_manager_stateless,
-            patch(
-                "litellm.proxy._experimental.mcp_server.server.session_manager_stateful"
-            ) as mock_session_manager_stateful,
-            patch(
-                "litellm.proxy._experimental.mcp_server.server.sse_session_manager"
-            ) as mock_sse_session_manager,
+            patch.object(
+                mcp_server.session_manager_stateless,
+                "run",
+                return_value=mock_cm_stateless,
+            ) as mock_stateless_run,
+            patch.object(
+                mcp_server.session_manager_stateful,
+                "run",
+                return_value=mock_cm_stateful,
+            ) as mock_stateful_run,
+            patch.object(
+                mcp_server.sse_session_manager,
+                "run",
+                return_value=mock_cm_sse,
+            ) as mock_sse_run,
             patch("litellm.proxy._experimental.mcp_server.server.verbose_logger"),
         ):
-            # Mock the run() method to return a mock context manager
-            mock_cm = AsyncMock()
-            mock_cm.__aenter__ = AsyncMock()
-            mock_cm.__aexit__ = AsyncMock()
-
-            mock_session_manager_stateless.run.return_value = mock_cm
-            mock_session_manager_stateful.run.return_value = mock_cm
-            mock_sse_session_manager.run.return_value = mock_cm
-
             # Create multiple concurrent tasks that call initialize_session_managers
             async def init_task():
                 await initialize_session_managers()
@@ -1171,19 +1178,25 @@ async def test_concurrent_initialize_session_managers():
 
             # Each session manager.run() should only be called once due to the lock
             assert (
-                mock_session_manager_stateless.run.call_count == 1
-            ), f"Expected 1 call to session_manager_stateless.run(), got {mock_session_manager_stateless.run.call_count}"
+                mock_stateless_run.call_count == 1
+            ), f"Expected 1 call to session_manager_stateless.run(), got {mock_stateless_run.call_count}"
             assert (
-                mock_session_manager_stateful.run.call_count == 1
-            ), f"Expected 1 call to session_manager_stateful.run(), got {mock_session_manager_stateful.run.call_count}"
+                mock_stateful_run.call_count == 1
+            ), f"Expected 1 call to session_manager_stateful.run(), got {mock_stateful_run.call_count}"
             assert (
-                mock_sse_session_manager.run.call_count == 1
-            ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_session_manager.run.call_count}"
+                mock_sse_run.call_count == 1
+            ), f"Expected 1 call to sse_session_manager.run(), got {mock_sse_run.call_count}"
 
-            # The context managers should only be entered once each (3 managers)
+            # The context managers should only be entered once each
             assert (
-                mock_cm.__aenter__.call_count == 3
-            ), f"Expected 3 calls to __aenter__ (one per session manager), got {mock_cm.__aenter__.call_count}"
+                mock_cm_stateless.__aenter__.call_count == 1
+            ), f"Expected 1 call to stateless __aenter__, got {mock_cm_stateless.__aenter__.call_count}"
+            assert (
+                mock_cm_stateful.__aenter__.call_count == 1
+            ), f"Expected 1 call to stateful __aenter__, got {mock_cm_stateful.__aenter__.call_count}"
+            assert (
+                mock_cm_sse.__aenter__.call_count == 1
+            ), f"Expected 1 call to sse __aenter__, got {mock_cm_sse.__aenter__.call_count}"
 
             # State should be properly set
             assert mcp_server._SESSION_MANAGERS_INITIALIZED is True
@@ -1195,14 +1208,12 @@ async def test_concurrent_initialize_session_managers():
         leaked_task = mcp_server._stateful_auth_context_cleanup_task
         if leaked_task is not None and leaked_task is not original_cleanup_task:
             leaked_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError, Exception):
-                await leaked_task
 
         # Restore original state
         mcp_server._SESSION_MANAGERS_INITIALIZED = original_initialized
         mcp_server._session_manager_cm = original_session_cm
-        mcp_server._session_manager_stateful_cm = original_session_stateful_cm
-        mcp_server._sse_session_manager_cm = original_sse_session_cm
+        mcp_server._session_manager_stateful_cm = original_stateful_cm
+        mcp_server._sse_session_manager_cm = original_sse_cm
         mcp_server._stateful_auth_context_cleanup_task = original_cleanup_task
 
 
@@ -1637,10 +1648,7 @@ async def test_mcp_routing_initialize_rejected_when_owner_at_session_cap():
     active = {f"s{i}": 1 for i in range(cap)}  # all in flight -> cannot evict
     contexts = {f"s{i}": MagicMock() for i in range(cap)}
 
-    init_body = (
-        b'{"jsonrpc":"2.0","id":1,"method":"initialize",'
-        b'"params":{"protocolVersion":"2024-11-05"}}'
-    )
+    init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}'
     scope = {
         "type": "http",
         "method": "POST",
@@ -2587,6 +2595,134 @@ async def test_stateful_mcp_get_stream_does_not_block_post():
         mcp_server._stateful_session_locks.pop(session_id, None)
 
 
+def test_jsonrpc_text_has_top_level_method_ignores_nested_method():
+    """The top-level-key scan must not be fooled by a ``method`` field nested
+    inside a JSON-RPC response's ``result`` payload — a flat substring search
+    would, and that misread is what deadlocks the session lock."""
+    from litellm.proxy._experimental.mcp_server.server import (
+        _jsonrpc_text_has_top_level_method,
+    )
+
+    request = '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}'
+    assert _jsonrpc_text_has_top_level_method(request) is True
+
+    # method key out of order (after params) is still top-level
+    reordered = '{"jsonrpc":"2.0","params":{"x":1},"method":"foo"}'
+    assert _jsonrpc_text_has_top_level_method(reordered) is True
+
+    # response whose result nests a "method" key (and arrays of them)
+    response = (
+        '{"jsonrpc":"2.0","id":1,"result":{"toolResult":{"method":"GET"},'
+        '"steps":[{"method":"x"}]}}'
+    )
+    assert _jsonrpc_text_has_top_level_method(response) is False
+
+    # truncated response: result value never closes, no top-level method seen
+    truncated = '{"jsonrpc":"2.0","id":1,"result":{"text":"' + "q" * 5000
+    assert _jsonrpc_text_has_top_level_method(truncated) is False
+
+
+@pytest.mark.asyncio
+async def test_truncated_jsonrpc_response_with_nested_method_skips_lock():
+    """Regression: a large JSON-RPC *response* POST whose ``result`` payload
+    nests a ``method`` key must skip the per-session lock so it does not
+    deadlock behind the in-flight request POST that is holding the lock while
+    it awaits this very response (e.g. sampling/createMessage)."""
+    try:
+        from litellm.proxy._experimental.mcp_server import server as mcp_server
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager_stateful,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    session_id = "nested-method-response-session"
+    owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner")
+    mcp_server._stateful_session_auth_contexts[session_id] = (
+        mcp_server.MCPAuthenticatedUser(user_api_key_auth=owner_auth)
+    )
+    mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for(
+        owner_auth
+    )
+
+    gate = asyncio.Event()
+    request_in_handle = asyncio.Event()
+    response_handled = asyncio.Event()
+
+    async def handle(s, r, se):
+        msg = await r()
+        body = msg.get("body", b"") or b""
+        if b'"result"' in body:
+            response_handled.set()
+        else:
+            request_in_handle.set()
+            await gate.wait()
+
+    async def call(body: bytes):
+        scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp",
+            "headers": [(b"mcp-session-id", session_id.encode())],
+        }
+        receive = AsyncMock(
+            return_value={
+                "type": "http.request",
+                "body": body,
+                "more_body": False,
+            }
+        )
+        await handle_streamable_http_mcp(scope, receive, AsyncMock())
+
+    # The in-flight request POST holds the session lock while blocked.
+    request_body = b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}'
+    # A JSON-RPC response larger than the routing peek cap so it can't be fully
+    # parsed, with a nested "method" key in the first bytes to trip a flat
+    # substring heuristic.
+    response_body = (
+        '{"jsonrpc":"2.0","id":99,"result":{"toolResult":'
+        '{"method":"GET","payload":"' + ("x" * 5000) + '"}}}'
+    ).encode()
+
+    try:
+        with (
+            patch(
+                "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+                new_callable=AsyncMock,
+                return_value=(owner_auth, None, None, None, None, None),
+            ),
+            patch(
+                "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+                True,
+            ),
+            patch.object(
+                session_manager_stateful, "handle_request", side_effect=handle
+            ),
+            patch.object(
+                session_manager_stateful,
+                "_server_instances",
+                {session_id: MagicMock()},
+            ),
+        ):
+            req_task = asyncio.create_task(call(request_body))
+            await asyncio.wait_for(request_in_handle.wait(), timeout=1.0)
+
+            resp_task = asyncio.create_task(call(response_body))
+            # Under a flat substring heuristic the response would acquire the
+            # lock held by req_task and this wait would time out (deadlock).
+            await asyncio.wait_for(response_handled.wait(), timeout=1.0)
+
+            gate.set()
+            await asyncio.gather(req_task, resp_task)
+    finally:
+        gate.set()
+        mcp_server._stateful_session_auth_contexts.pop(session_id, None)
+        mcp_server._stateful_session_owners.pop(session_id, None)
+        mcp_server._stateful_session_locks.pop(session_id, None)
+        mcp_server._stateful_session_active_request_counts.pop(session_id, None)
+
+
 @pytest.mark.asyncio
 @pytest.mark.no_parallel
 async def test_mcp_routing_with_conflicting_alias_and_group_name():
@@ -2729,7 +2865,7 @@ async def test_oauth2_headers_passed_to_mcp_client():
         mcp_auth_header=None,
         extra_headers=None,
         stdio_env=None,
-        subject_token=None,
+        **kwargs,
     ):
         # Capture the arguments for verification
         captured_client_args.update(
@@ -2738,7 +2874,7 @@ async def test_oauth2_headers_passed_to_mcp_client():
                 "mcp_auth_header": mcp_auth_header,
                 "extra_headers": extra_headers,
                 "stdio_env": stdio_env,
-                "subject_token": subject_token,
+                "kwargs": kwargs,
             }
         )
         # Return a mock client that doesn't actually connect
@@ -2764,6 +2900,16 @@ async def test_oauth2_headers_passed_to_mcp_client():
             "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
             AsyncMock(return_value=[oauth2_server]),
         ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user",
+            new_callable=AsyncMock,
+            return_value={},
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
+            new_callable=AsyncMock,
+            return_value=None,
+        ),
     ):
         # Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client
         await _get_tools_from_mcp_servers(
@@ -2840,7 +2986,7 @@ async def test_list_tools_single_server_unprefixed_names():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         tool = MagicMock()
         tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
@@ -2922,7 +3068,7 @@ async def test_list_tools_multiple_servers_prefixed_names():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         tool = MagicMock()
         # When multiple servers, add_prefix should be True -> prefixed names
@@ -3189,7 +3335,7 @@ async def test_list_tools_filters_by_key_team_permissions():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return 4 tools, but only 2 should be allowed
         tool1 = MagicMock()
@@ -3299,7 +3445,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return 4 tools
         tool1 = MagicMock()
@@ -3395,7 +3541,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all():
         extra_headers=None,
         add_prefix=False,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return 3 tools
         tool1 = MagicMock()
@@ -3494,7 +3640,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions():
         extra_headers=None,
         add_prefix=True,
         raw_headers=None,
-        user_api_key_auth=None,
+        **kwargs,
     ):
         # Return tools WITH prefix (as they come from MCP server)
         tool1 = MagicMock()
@@ -5178,3 +5324,42 @@ def test_get_forwarded_auth_from_scope_skips_when_no_litellm_key_header():
         ]
     }
     assert _get_forwarded_auth_from_scope(scope) is None
+
+
+@pytest.mark.asyncio
+async def test_create_mcp_client_sampling_disabled_by_default():
+    """Sampling callback must be None when allow_sampling is not set (default False)."""
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        MCPServerManager,
+    )
+
+    manager = MCPServerManager()
+    server = MCPServer(
+        server_id="no-sampling",
+        name="no-sampling",
+        url="https://example.com/mcp",
+        transport=MCPTransport.http,
+    )
+
+    client = await manager._create_mcp_client(server=server)
+    assert client._sampling_callback is None
+
+
+@pytest.mark.asyncio
+async def test_create_mcp_client_sampling_enabled():
+    """Sampling callback must be set when allow_sampling=True."""
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        MCPServerManager,
+    )
+
+    manager = MCPServerManager()
+    server = MCPServer(
+        server_id="with-sampling",
+        name="with-sampling",
+        url="https://example.com/mcp",
+        transport=MCPTransport.http,
+        allow_sampling=True,
+    )
+
+    client = await manager._create_mcp_client(server=server)
+    assert client._sampling_callback is not None

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -323,9 +323,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             if server.name == "github":
                 tool1 = MagicMock()
@@ -378,9 +376,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             assert mcp_auth_header == "legacy-token"  # Should use legacy header
             tool = MagicMock()
@@ -417,9 +413,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             assert (
                 mcp_auth_header == "server-specific-token"
@@ -460,7 +454,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -510,7 +504,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -563,7 +557,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -619,7 +613,7 @@ class TestMCPServerManager:
         captured_extra_headers = None
 
         async def capture_create_mcp_client(
-            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
+            server, mcp_auth_header, extra_headers, stdio_env, subject_token=None, **kwargs
         ):  # pragma: no cover - helper
             nonlocal captured_extra_headers
             captured_extra_headers = extra_headers
@@ -1169,9 +1163,7 @@ class TestMCPServerManager:
         async def mock_get_tools_from_server(
             server,
             mcp_auth_header=None,
-            mcp_protocol_version=None,
-            raw_headers=None,
-            user_api_key_auth=None,
+            **kwargs,
         ):
             assert (
                 mcp_auth_header == "server-specific-token"

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py

@@ -9,8 +9,6 @@ they may send a stale `mcp-session-id` header. This test verifies that:
 
 import asyncio
 from unittest.mock import AsyncMock, MagicMock, patch
-
-from fastapi import HTTPException
 from litellm.types.mcp import MCPAuth
 import pytest
 
@@ -600,6 +598,8 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
     Per-user OAuth server with no stored token should fail fast with 401 +
     WWW-Authenticate so PKCE can start.
     """
+    from fastapi import HTTPException
+
     try:
         from litellm.proxy._experimental.mcp_server.server import (
             handle_streamable_http_mcp,
@@ -612,8 +612,13 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
         "type": "http",
         "method": "POST",
         "path": "/mcp",
+        "scheme": "http",
+        "query_string": b"",
+        "root_path": "",
+        "server": ("localhost", 8000),
         "headers": [
             (b"content-type", b"application/json"),
+            (b"host", b"localhost:8000"),
         ],
     }
     receive = AsyncMock()
@@ -660,11 +665,12 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
         with pytest.raises(HTTPException) as exc_info:
             await handle_streamable_http_mcp(scope, receive, send)
 
-    exc = exc_info.value
-    assert exc.status_code == 401
-    assert "www-authenticate" in exc.headers
+    # Verify a 401 was raised
     assert mock_get_stored_token.await_count == 1
     assert mock_handle_request.await_count == 0
+    assert exc_info.value.status_code == 401
+    assert "www-authenticate" in exc_info.value.headers
+    assert "Bearer authorization_uri=" in exc_info.value.headers["www-authenticate"]
 
 
 @pytest.mark.asyncio
@@ -685,11 +691,22 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
         "type": "http",
         "method": "POST",
         "path": "/mcp",
+        "scheme": "http",
+        "query_string": b"",
+        "root_path": "",
+        "server": ("localhost", 8000),
         "headers": [
             (b"content-type", b"application/json"),
+            (b"host", b"localhost:8000"),
         ],
     }
-    receive = AsyncMock()
+    receive = AsyncMock(
+        return_value={
+            "type": "http.request",
+            "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}',
+            "more_body": False,
+        }
+    )
     send = AsyncMock()
     user_auth = MagicMock()
     user_auth.user_id = "test-user-id"
@@ -729,6 +746,11 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
             "handle_request",
             new_callable=AsyncMock,
         ) as mock_handle_request,
+        patch.object(
+            session_manager_stateless,
+            "_server_instances",
+            {},
+        ),
     ):
         await handle_streamable_http_mcp(scope, receive, send)
 

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py

@@ -2,6 +2,7 @@
 Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
 """
 
+import json
 import os
 import re
 import sys
@@ -20,7 +21,7 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.tool_permission import (
     ToolPermissionGuardrail,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
 from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
     PermissionError,
 )
@@ -894,3 +895,153 @@ class TestToolPermissionGuardrailIntegration:
         is_allowed, rule_id, _ = guardrail._check_tool_permission("Read")
         assert is_allowed is False
         assert rule_id == "deny_read"
+
+
+class TestToolPermissionGuardrailInMemoryUpdate:
+    """Regression: an in-memory params update (PUT /guardrails path) must rebuild
+    the compiled rule maps, not just self.rules, so the new rules are enforced
+    without reinitializing the guardrail."""
+
+    def _bash(self, command):
+        return ChatCompletionMessageToolCall(
+            function={"name": "Bash", "arguments": json.dumps({"command": command})},
+            type="function",
+        )
+
+    def test_update_in_memory_recompiles_added_param_pattern(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[{"id": "native-bash", "tool_name": r"^Bash$", "decision": "allow"}],
+            default_action="deny",
+            on_disallowed_action="block",
+        )
+        # No pattern yet: any Bash command is allowed.
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0]
+            is True
+        )
+
+        guardrail.update_in_memory_litellm_params(
+            LitellmParams(
+                guardrail="tool_permission",
+                mode=["pre_call", "post_call"],
+                default_action="deny",
+                on_disallowed_action="block",
+                rules=[
+                    {
+                        "id": "native-bash",
+                        "tool_name": r"^Bash$",
+                        "decision": "allow",
+                        "allowed_param_patterns": {
+                            "command": r"^(?!(echo blockme)$).*$"
+                        },
+                    }
+                ],
+            )
+        )
+
+        # The compiled map must be rebuilt, and enforcement must reflect it.
+        assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {})
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0]
+            is False
+        )
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo hello"))[0] is True
+        )
+
+    def test_update_in_memory_recompiles_tool_name_target(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[],
+            default_action="allow",
+            on_disallowed_action="block",
+        )
+        # No rules: default_action allow lets Bash through.
+        assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is True
+
+        guardrail.update_in_memory_litellm_params(
+            LitellmParams(
+                guardrail="tool_permission",
+                mode=["pre_call", "post_call"],
+                default_action="allow",
+                on_disallowed_action="block",
+                rules=[{"id": "deny-bash", "tool_name": r"^Bash$", "decision": "deny"}],
+            )
+        )
+
+        # A newly added deny rule (new id) must match -> its compiled target was rebuilt.
+        assert "deny-bash" in guardrail._compiled_rule_targets
+        assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is False
+
+    def test_update_in_memory_preserves_rules_when_rules_absent(self):
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[
+                {
+                    "id": "native-bash",
+                    "tool_name": r"^Bash$",
+                    "decision": "allow",
+                    "allowed_param_patterns": {"command": r"^(?!(echo blockme)$).*$"},
+                }
+            ],
+            default_action="deny",
+            on_disallowed_action="block",
+        )
+        assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {})
+
+        # A partial update that does not carry `rules` must NOT wipe the existing
+        # ruleset / compiled maps.
+        guardrail.update_in_memory_litellm_params(
+            LitellmParams(
+                guardrail="tool_permission",
+                mode=["pre_call", "post_call"],
+                default_action="deny",
+                on_disallowed_action="block",
+            )
+        )
+
+        assert len(guardrail.rules) == 1
+        assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {})
+        assert (
+            guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0]
+            is False
+        )
+
+    def test_update_in_memory_rejects_invalid_regex_and_keeps_previous_rules(self):
+        """Regression: a live update whose rules contain an invalid regex must be
+        rejected atomically. The bad rule must not leak in as a compiled-target
+        wildcard (match-all), and the previously enforced ruleset must survive."""
+        guardrail = ToolPermissionGuardrail(
+            guardrail_name="tp",
+            rules=[{"id": "deny-secret", "tool_name": r"^Secret$", "decision": "deny"}],
+            default_action="allow",
+            on_disallowed_action="block",
+        )
+        # Baseline: only "Secret" is denied; any other tool is allowed.
+        assert guardrail._check_tool_permission("Secret")[0] is False
+        assert guardrail._check_tool_permission("Other")[0] is True
+
+        with pytest.raises(ValueError):
+            guardrail.update_in_memory_litellm_params(
+                LitellmParams(
+                    guardrail="tool_permission",
+                    mode=["pre_call", "post_call"],
+                    default_action="allow",
+                    on_disallowed_action="block",
+                    rules=[
+                        {
+                            "id": "deny-secret",
+                            "tool_name": r"^Secret$",
+                            "decision": "deny",
+                        },
+                        {"id": "bad", "tool_name": "[unclosed", "decision": "deny"},
+                    ],
+                )
+            )
+
+        # The bad rule must not have leaked in, and the prior ruleset must hold.
+        assert "bad" not in guardrail._compiled_rule_targets
+        assert all(rule.id != "bad" for rule in guardrail.rules)
+        assert guardrail._check_tool_permission("Other")[0] is True
+        assert guardrail._check_tool_permission("Secret")[0] is False

tests/test_litellm/proxy/test_caching_routes.py

@@ -123,33 +123,100 @@ def test_cache_ping_failure(mock_redis_failure):
     assert "message" in error_details
     assert "litellm_cache_params" in error_details
     assert "health_check_cache_params" in error_details
-    assert "traceback" in error_details
 
-    # Verify specific error message
-    assert "invalid username-password pair" in error_details["message"]
+    # Verify generic static message (exception text must not leak to clients)
+    assert error_details["message"] == "Service Unhealthy"
 
 
-def test_cache_ping_no_cache_initialized():
-    """Test cache ping when no cache is initialized"""
-    # Set cache to None
-    original_cache = litellm.cache
-    litellm.cache = None
-
+def test_cache_ping_failure_does_not_expose_traceback(mock_redis_failure):
+    """CWE-209: Stack trace and exception text must not appear in the HTTP 503 response body."""
     response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
     assert response.status_code == 503
 
     data = response.json()
-    print("response data=", json.dumps(data, indent=4))
-    assert "error" in data
-    error = data["error"]
+    error = data.get("error", {})
+    raw_body = json.dumps(data)
 
-    # Verify error contains all expected fields
-    assert "message" in error
+    # The word "traceback" (case-insensitive) must not appear anywhere in the response
+    assert (
+        "traceback" not in raw_body.lower()
+    ), "CWE-209: Python traceback exposed in HTTP 503 response body"
+    # Internal frame paths should not leak either
+    assert (
+        'File "' not in raw_body
+    ), "CWE-209: Python stack frame paths exposed in HTTP 503 response body"
+    # Exception text (e.g. Redis hostnames/IPs) must not leak either
+    assert (
+        "invalid username-password pair" not in raw_body
+    ), "CWE-209: Exception message text exposed in HTTP 503 response body"
+
+    # The error message should be the safe static string
     error_details = json.loads(error["message"])
-    assert "Cache not initialized. litellm.cache is None" in error_details["message"]
+    assert error_details["message"] == "Service Unhealthy"
 
-    # Restore original cache
-    litellm.cache = original_cache
+
+def test_cache_ping_no_cache_initialized():
+    """Test cache ping when no cache is initialized returns 503 with ProxyException envelope.
+
+    Verifies the exact response structure so that regressions in the error format
+    (e.g. message moving to a different field, or extra internal details leaking)
+    are caught immediately.
+    """
+    original_cache = litellm.cache
+    litellm.cache = None
+
+    try:
+        response = client.get(
+            "/cache/ping", headers={"Authorization": "Bearer sk-1234"}
+        )
+        assert response.status_code == 503
+
+        data = response.json()
+        print("response data=", json.dumps(data, indent=4))
+        # ProxyException is serialised as {"error": {"message": "...", "type": ..., ...}}
+        assert "error" in data
+        error_details = json.loads(data["error"]["message"])
+        assert (
+            error_details["message"] == "Cache not initialized. litellm.cache is None"
+        )
+    finally:
+        litellm.cache = original_cache
+
+
+def test_cache_ping_no_cache_does_not_expose_internals():
+    """CWE-209: No-cache 503 must use the ProxyException envelope with no internal details.
+
+    The null-cache path raises ProxyException directly (not HTTPException), so the
+    response is {"error": {"message": "...", ...}} — same envelope as other 503s from
+    this endpoint — with no tracebacks, source paths, or extra fields leaking.
+    """
+    original_cache = litellm.cache
+    litellm.cache = None
+
+    try:
+        response = client.get(
+            "/cache/ping", headers={"Authorization": "Bearer sk-1234"}
+        )
+        assert response.status_code == 503
+
+        raw_body = response.text
+        # No Python traceback or source-file paths must appear in the response
+        assert "traceback" not in raw_body.lower(), (
+            "CWE-209: Python traceback exposed in /cache/ping no-cache response"
+        )
+        assert 'File "' not in raw_body, (
+            "CWE-209: Python stack frame paths exposed in /cache/ping no-cache response"
+        )
+
+        data = response.json()
+        # Response must use the ProxyException envelope
+        assert "error" in data, f"Expected ProxyException envelope, got: {data}"
+        error_details = json.loads(data["error"]["message"])
+        assert (
+            error_details["message"] == "Cache not initialized. litellm.cache is None"
+        )
+    finally:
+        litellm.cache = original_cache
 
 
 def test_cache_ping_health_check_includes_only_cache_attributes(mock_redis_success):

tests/test_litellm/proxy/test_dynamic_mcp_route.py

@@ -486,3 +486,57 @@ async def test_dynamic_mcp_route_empty_access_group_returns_404():
             await dynamic_mcp_route("empty_group", request)
 
     assert exc_info.value.status_code == 404
+
+
+# ---------------------------------------------------------------------------
+# 6. Unexpected exception → 500 without leaking stack trace (CWE-209)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_dynamic_mcp_route_unexpected_exception_returns_500_without_traceback():
+    """CWE-209: an unexpected exception must return 500 with a generic message,
+    never leaking str(e) or a Python traceback to the caller."""
+    from litellm.proxy.proxy_server import dynamic_mcp_route
+
+    request = _make_request("/boom/mcp")
+
+    fake_mgr = MagicMock()
+    fake_mgr.get_mcp_server_by_name = MagicMock(
+        side_effect=RuntimeError("internal host: redis://10.0.0.1:6379")
+    )
+
+    with patch(_MCP_MANAGER, fake_mgr):
+        with pytest.raises(HTTPException) as exc_info:
+            await dynamic_mcp_route("boom", request)
+
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.detail == "Internal server error"
+    assert "10.0.0.1" not in str(exc_info.value.detail)
+    assert "traceback" not in str(exc_info.value.detail).lower()
+
+
+@pytest.mark.asyncio
+async def test_toolset_mcp_route_unexpected_exception_returns_500_without_traceback():
+    """CWE-209: toolset_mcp_route must return 500 with a generic message on
+    unexpected errors, never leaking exception text to the caller."""
+    from litellm.proxy.proxy_server import toolset_mcp_route
+
+    request = _make_request("/toolset/broken_toolset/mcp")
+
+    fake_mgr = MagicMock()
+    fake_mgr.get_toolset_by_name_cached = AsyncMock(
+        side_effect=RuntimeError("connection to db-host:5432 refused")
+    )
+
+    with (
+        patch(_MCP_MANAGER, fake_mgr),
+        patch(_PRISMA, new=MagicMock()),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await toolset_mcp_route("broken_toolset", request)
+
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.detail == "Internal server error"
+    assert "db-host" not in str(exc_info.value.detail)
+    assert "traceback" not in str(exc_info.value.detail).lower()

tests/test_litellm/test_main.py

@@ -849,12 +849,13 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict(
 
 
 @pytest.mark.parametrize(
-    "model, model_info, expected_model_param",
+    "model, model_info, expected_model_param, expected_base_model_param",
     [
-        ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro"),
+        ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None),
         (
             "gemini/gemini-3.1-pro",
             {"base_model": "gemini-3.1-pro-preview"},
+            "gemini-3.1-pro",
             "gemini-3.1-pro-preview",
         ),
     ],
@@ -863,7 +864,13 @@ def test_completion_optional_params_base_model(
     model: str,
     model_info: dict | None,
     expected_model_param: str,
+    expected_base_model_param: str | None,
 ):
+    """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model``
+    (an additive capability hint), without overwriting ``model`` with the label.
+
+    Regression for #29618: overwriting ``model`` with a friendly ``base_model``
+    label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``."""
     with patch("litellm.main.get_optional_params") as mock_get_optional_params:
         mock_get_optional_params.return_value = MagicMock()
 
@@ -881,10 +888,9 @@ def test_completion_optional_params_base_model(
         litellm.completion(**kwargs)
 
         assert mock_get_optional_params.called is True
-        get_optional_params_model_param = mock_get_optional_params.call_args.kwargs[
-            "model"
-        ]
-        assert get_optional_params_model_param == expected_model_param
+        call_kwargs = mock_get_optional_params.call_args.kwargs
+        assert call_kwargs["model"] == expected_model_param
+        assert call_kwargs["base_model"] == expected_base_model_param
 
 
 @patch("litellm.completion_extras.responses_api_bridge.completion")

tests/test_litellm/test_utils.py

@@ -4144,3 +4144,51 @@ class TestValidateAndFixThinkingParam:
         validate_and_fix_thinking_param(thinking=thinking)
         assert "budgetTokens" in thinking
         assert "budget_tokens" not in thinking
+
+
+class TestBedrockBaseModelLabelKeepsTools:
+    """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly
+    label must not silently drop ``tools``/``tool_choice`` under ``drop_params``."""
+
+    TOOLS = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            },
+        }
+    ]
+
+    def test_base_model_label_keeps_tools_with_drop_params(self):
+        from litellm.utils import get_optional_params
+
+        result = get_optional_params(
+            model="eu.anthropic.claude-haiku-4-5-20251001-v1:0",
+            custom_llm_provider="bedrock",
+            base_model="claude-haiku-4-5",
+            tools=self.TOOLS,
+            tool_choice="auto",
+            drop_params=True,
+        )
+
+        assert "tools" in result
+        assert "tool_choice" in result
+
+    def test_base_model_label_alone_drops_tools(self):
+        """Without the real model id the label resolves to no tool support, so passing
+        the label as ``model`` is exactly what dropped tools before the fix."""
+        from litellm.utils import get_optional_params
+
+        result = get_optional_params(
+            model="claude-haiku-4-5",
+            custom_llm_provider="bedrock",
+            tools=self.TOOLS,
+            tool_choice="auto",
+            drop_params=True,
+        )
+
+        assert "tools" not in result