diff --git a/litellm/__init__.py b/litellm/__init__.py index 98c9dcb5dd..e49f4a4699 100644 --- a/litellm/__init__.py +++ b/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, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index bdc3289b87..5df8db7317 100644 --- a/litellm/_lazy_imports_registry.py +++ b/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", diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 7559fe142c..0bc81ece5f 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/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 diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index b8cdc8210f..7c4f994152 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/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) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 55042a733e..f3274151e5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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( diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index a450ee0b21..72f1eef36c 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/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) diff --git a/litellm/llms/bedrock_mantle/responses/__init__.py b/litellm/llms/bedrock_mantle/responses/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py new file mode 100644 index 0000000000..b63fd0ecdb --- /dev/null +++ b/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 diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 9e9d300b58..cca3b3da37 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/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": diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 77a95bfa5a..644e96a7dd 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/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) diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 226f6b2eba..6be885b1f9 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/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, diff --git a/litellm/llms/snowflake/utils.py b/litellm/llms/snowflake/utils.py index d84efdd9fc..4f79006f6f 100644 --- a/litellm/llms/snowflake/utils.py +++ b/litellm/llms/snowflake/utils.py @@ -25,6 +25,7 @@ class SnowflakeBaseConfig: "temperature", "max_tokens", "top_p", + "stream", "response_format", "tools", "tool_choice", diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 3f945adca0..e9f08f403f 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/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( diff --git a/litellm/main.py b/litellm/main.py index 96f81381c8..c8aae0ce85 100644 --- a/litellm/main.py +++ b/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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ed6de4fa6b..fbe6c09720 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py new file mode 100644 index 0000000000..e42270bf10 --- /dev/null +++ b/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") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d9b112f6c2..0d2008cdad 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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, @@ -289,6 +295,82 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: return data +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-[^}]+)\}$") @@ -600,6 +682,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") @@ -699,8 +783,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 @@ -1494,6 +1577,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. @@ -1510,6 +1594,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. @@ -1520,23 +1605,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)}). " @@ -1559,6 +1665,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 @@ -1585,6 +1693,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( @@ -1668,6 +1778,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 @@ -3030,6 +3141,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( @@ -3260,7 +3372,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, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py new file mode 100644 index 0000000000..1637c9eb0b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -0,0 +1,1279 @@ +""" +MCP Sampling Handler +Handles `sampling/createMessage` requests from upstream MCP servers by +routing them through LiteLLM's internal completion infrastructure. +This allows MCP servers to perform agentic reasoning (e.g., multi-step +tool calling, chain-of-thought) without needing their own LLM API keys — +LiteLLM acts as the LLM provider using its existing 100+ provider support, +cost tracking, rate limiting, and model routing. +MCP Spec Reference: + https://modelcontextprotocol.io/specification/2025-11-25/client/sampling +""" + +from typing import Any, Dict, List, Optional, Union +import typing + +if typing.TYPE_CHECKING: + from litellm.proxy.utils import ProxyLogging + +from litellm._logging import verbose_logger + +from fastapi import HTTPException + +# Guard imports that require the mcp package +try: + from mcp.types import ( + CreateMessageRequestParams, + CreateMessageResult, + CreateMessageResultWithTools, + ErrorData, + ModelPreferences, + SamplingMessage, + TextContent, + Tool, + ToolChoice, + ToolUseContent, + ) + + MCP_SAMPLING_AVAILABLE = True +except ImportError as _sampling_import_err: + MCP_SAMPLING_AVAILABLE = False + verbose_logger.warning( + "MCP sampling disabled: failed to import required types from mcp.types — %s. " + "This usually means the 'mcp' package is not installed or is an older version " + "that does not support sampling. Install/upgrade with: pip install 'mcp>=1.1'", + _sampling_import_err, + ) + + +def _resolve_model_from_preferences( + model_preferences: Optional["ModelPreferences"], + default_model: Optional[str] = None, +) -> str: + """ + Resolve an LLM model name from MCP ModelPreferences. + Strategy: + 1. Check hints for substring matches against known model names. + 2. Fall back to priority-based selection (cost/speed/intelligence). + 3. Fall back to the configured default model. + Args: + model_preferences: MCP ModelPreferences with hints and priorities. + default_model: Fallback model if no hint matches. + Returns: + A model string suitable for litellm.acompletion(). + """ + import litellm + + # Build list of available model names from proxy Router or litellm.model_list + available_model_names: list = [] + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is not None: + available_model_names = llm_router.get_model_names() + except Exception: + pass + if not available_model_names and litellm.model_list: + for entry in litellm.model_list: + if isinstance(entry, dict): + name = entry.get("model_name") + if name: + available_model_names.append(name) + elif isinstance(entry, str): + available_model_names.append(entry) + if model_preferences and model_preferences.hints: + for hint in model_preferences.hints: + hint_name = getattr(hint, "name", None) + if not hint_name: + continue + # Try direct match first + if hint_name in available_model_names: + verbose_logger.debug( + "MCP sampling model resolution: direct hint match '%s'", + hint_name, + ) + return hint_name + # Try substring match against known models + for model_name in available_model_names: + if hint_name.lower() in model_name.lower(): + verbose_logger.debug( + "MCP sampling model resolution: substring hint match " + "'%s' -> '%s'", + hint_name, + model_name, + ) + return model_name + verbose_logger.debug( + "MCP sampling model resolution: no hint matched from %s " + "against %d available models", + [getattr(h, "name", None) for h in model_preferences.hints], + len(available_model_names), + ) + + # 2. Priority-based selection (cost/speed/intelligence) + if ( + model_preferences + and available_model_names + and _has_priorities(model_preferences) + ): + best = _select_model_by_priority(available_model_names, model_preferences) + if best is not None: + verbose_logger.debug( + "MCP sampling model resolution: priority-based selection chose '%s'", + best, + ) + return best + + # 3. Use default model from caller + if default_model: + verbose_logger.debug( + "MCP sampling model resolution: using caller-provided default '%s'", + default_model, + ) + return default_model + # Fall back to first available model + if available_model_names: + verbose_logger.debug( + "MCP sampling model resolution: no default configured, " + "falling back to first available model '%s'", + available_model_names[0], + ) + return available_model_names[0] + # Last resort - use LiteLLM default or raise error + default_sampling_model = getattr(litellm, "default_mcp_sampling_model", None) + if default_sampling_model: + verbose_logger.debug( + "MCP sampling model resolution: using litellm.default_mcp_sampling_model='%s'", + default_sampling_model, + ) + return default_sampling_model + raise ValueError( + "No model could be resolved for MCP sampling. Please configure 'default_mcp_sampling_model' in your LiteLLM configuration." + ) + + +def _has_priorities(model_preferences: "ModelPreferences") -> bool: + """Return True if any priority weight is set (non-None and > 0).""" + return any( + (getattr(model_preferences, attr, None) or 0) > 0 + for attr in ("costPriority", "speedPriority", "intelligencePriority") + ) + + +def _select_model_by_priority( + model_names: List[str], + model_preferences: "ModelPreferences", +) -> Optional[str]: + """Score available models by MCP priority weights and return the best. + + Scoring strategy (per the MCP spec, priorities are 0-1 floats): + + * **costPriority** — higher means "prefer cheaper models". + Metric: combined (input + output) cost per token from + ``model_prices_and_context_window.json``. Lower cost → higher score. + + * **speedPriority** — higher means "prefer faster models". + Metric: ``output_tokens_per_second`` from model info when available; + otherwise a neutral score for every candidate, since no reliable + latency proxy exists (context-window size does not track speed). + + * **intelligencePriority** — higher means "prefer smarter models". + Metric: ``max_output_tokens`` is used as a rough capability proxy + (frontier models expose larger context windows). + + Each metric is min-max normalised across the candidate set so that + every model gets a 0-1 score per dimension. The final score is the + weighted sum of the three normalised dimensions. + + Returns the highest-scoring model name, or None if scoring fails for + all candidates (e.g. no model_info available). + """ + import litellm as _litellm + + cost_weight = getattr(model_preferences, "costPriority", None) or 0.0 + speed_weight = getattr(model_preferences, "speedPriority", None) or 0.0 + intel_weight = getattr(model_preferences, "intelligencePriority", None) or 0.0 + + # Gather raw metrics for each model + scored: List[Dict[str, Any]] = [] + for name in model_names: + try: + info = _litellm.get_model_info(name) + except Exception: + continue + input_cost = info.get("input_cost_per_token") or 0.0 + output_cost = info.get("output_cost_per_token") or 0.0 + total_cost = input_cost + output_cost + max_output = info.get("max_output_tokens") or info.get("max_tokens") or 0 + output_tps = info.get("output_tokens_per_second") or 0.0 + scored.append( + { + "name": name, + "cost": total_cost, + "max_output": max_output, + "output_tps": output_tps, + } + ) + + if not scored: + return None + + # Min-max normalisation helpers + def _normalise(values: List[float], invert: bool = False) -> List[float]: + """Normalise to [0, 1]. If *invert*, lower raw → higher score.""" + lo, hi = min(values), max(values) + if hi == lo: + return [0.5] * len(values) # all equal → neutral score + normed = [(v - lo) / (hi - lo) for v in values] + if invert: + normed = [1.0 - n for n in normed] + return normed + + costs = [s["cost"] for s in scored] + max_outputs = [float(s["max_output"]) for s in scored] + output_tps_values = [s["output_tps"] for s in scored] + + # costPriority: lower cost → higher score (invert) + cost_scores = _normalise(costs, invert=True) + # speedPriority: use output_tokens_per_second if any model has it, + # otherwise a neutral score (no reliable latency proxy is available). + if any(v > 0 for v in output_tps_values): + speed_scores = _normalise(output_tps_values, invert=False) + else: + speed_scores = [0.5] * len(scored) + # intelligencePriority: higher max_output → smarter + intel_scores = _normalise(max_outputs, invert=False) + + best_name = None + best_score = -1.0 + for i, entry in enumerate(scored): + score = ( + cost_weight * cost_scores[i] + + speed_weight * speed_scores[i] + + intel_weight * intel_scores[i] + ) + verbose_logger.debug( + "MCP priority scoring: model=%s cost_score=%.3f speed_score=%.3f " + "intel_score=%.3f → weighted=%.3f", + entry["name"], + cost_scores[i], + speed_scores[i], + intel_scores[i], + score, + ) + if score > best_score: + best_score = score + best_name = entry["name"] + + return best_name + + +def _convert_mcp_content_to_openai( + content: Any, +) -> Union[str, Dict[str, Any], List[Dict[str, Any]]]: + """ + Convert MCP SamplingMessage content to OpenAI message content format. + Handles: + - TextContent → string or {"type": "text", "text": ...} + - ImageContent → {"type": "image_url", "image_url": {"url": "data:..."}} + - AudioContent → {"type": "input_audio", "input_audio": {...}} + - ToolUseContent → function call representation + - ToolResultContent → tool result representation + - List of mixed content → list of content parts + """ + if isinstance(content, list): + parts = [] + for item in content: + converted = _convert_single_content(item) + if isinstance(converted, list): + parts.extend(converted) + else: + parts.append(converted) + return parts + return _convert_single_content(content) + + +def _convert_single_content( + content: Any, +) -> Union[Dict[str, Any], List[Dict[str, Any]]]: + """Convert a single MCP content item to OpenAI format. + + For text/image/audio content, returns a single content-part dict. + For tool_use/tool_result, returns a dict with a ``_marker_type`` key + so the caller (``_convert_mcp_messages_to_openai``) can hoist it to + the correct message-level position (``tool_calls`` array or a + separate ``role: "tool"`` message). + """ + import json + + content_type = getattr(content, "type", None) + if content_type == "text": + return {"type": "text", "text": content.text} + elif content_type == "image": + data = getattr(content, "data", "") + mime_type = getattr(content, "mimeType", "image/png") + return { + "type": "image_url", + "image_url": {"url": f"data:{mime_type};base64,{data}"}, + } + elif content_type == "audio": + data = getattr(content, "data", "") + mime_type = getattr(content, "mimeType", "audio/wav") + # Map MIME type to OpenAI audio format + format_map = { + "audio/wav": "wav", + "audio/mp3": "mp3", + "audio/mpeg": "mp3", + "audio/flac": "flac", + "audio/ogg": "ogg", + } + audio_format = format_map.get(mime_type, "wav") + return { + "type": "input_audio", + "input_audio": {"data": data, "format": audio_format}, + } + elif content_type == "tool_use": + # ToolUseContent → proper OpenAI function-call representation. + # The ``_marker_type`` key lets the message-level converter + # hoist this into the ``tool_calls`` array on the assistant + # message instead of embedding it inline as a content part. + return { + "_marker_type": "tool_use", + "id": getattr(content, "id", f"call_{id(content)}"), + "type": "function", + "function": { + "name": getattr(content, "name", ""), + "arguments": json.dumps(getattr(content, "input", {}), default=str), + }, + } + elif content_type == "tool_result": + # ToolResultContent → proper OpenAI tool-role message. + # Marked so the message-level converter can emit it as a + # separate ``{"role": "tool", ...}`` message. + tool_use_id = getattr(content, "toolUseId", "") + nested_content = getattr(content, "content", []) + if isinstance(nested_content, list): + text_parts = [ + getattr(c, "text", str(c)) + for c in nested_content + if getattr(c, "type", None) == "text" + ] + result_text = "\n".join(text_parts) if text_parts else "" + else: + result_text = str(nested_content) + return { + "_marker_type": "tool_result", + "role": "tool", + "tool_call_id": tool_use_id, + "content": result_text, + } + # Fallback: treat as text + return {"type": "text", "text": str(content)} + + +def _convert_mcp_messages_to_openai( + messages: List["SamplingMessage"], + system_prompt: Optional[str] = None, +) -> List[Dict[str, Any]]: + """ + Convert MCP SamplingMessage list to OpenAI messages format. + MCP messages use: + - role: "user" | "assistant" + - content: TextContent | ImageContent | AudioContent | ToolUseContent + | ToolResultContent | list[...] + OpenAI messages use: + - role: "system" | "user" | "assistant" | "tool" + - content: str | list[content_part] + """ + openai_messages: List[Dict[str, Any]] = [] + # Add system prompt if provided + if system_prompt: + openai_messages.append({"role": "system", "content": system_prompt}) + for msg in messages: + role = msg.role + content = msg.content + # Handle tool use content from assistant + if role == "assistant" and _has_tool_use(content): + tool_calls = _extract_tool_calls(content) + if tool_calls: + openai_msg: Dict[str, Any] = { + "role": "assistant", + "tool_calls": tool_calls, + } + # Also include any text content alongside tool calls + text_parts = _extract_text_parts(content) + if text_parts: + openai_msg["content"] = text_parts + openai_messages.append(openai_msg) + continue + # Handle tool result content from user + if role == "user" and _has_tool_result(content): + tool_results = _extract_tool_results(content) + for tool_result in tool_results: + openai_messages.append(tool_result) + continue + # Standard text/image/audio message — also handles any stray + # tool_use / tool_result that slipped past the fast-path checks + # above (e.g. unexpected role, single non-list content). + converted = _convert_mcp_content_to_openai(content) + converted_parts = ( + converted + if isinstance(converted, list) + else ([converted] if isinstance(converted, dict) else []) + ) + + # Separate marker items from regular content parts + tool_call_markers = [] + tool_result_markers = [] + regular_parts = [] + for part in converted_parts: + marker = part.get("_marker_type") if isinstance(part, dict) else None + if marker == "tool_use": + # Strip the internal marker before emitting + tc = {k: v for k, v in part.items() if k != "_marker_type"} + tool_call_markers.append(tc) + elif marker == "tool_result": + tr = {k: v for k, v in part.items() if k != "_marker_type"} + tool_result_markers.append(tr) + else: + regular_parts.append(part) + + # Emit assistant message with tool_calls if any were found + if tool_call_markers: + openai_msg_tc: Dict[str, Any] = { + "role": "assistant", + "tool_calls": tool_call_markers, + } + if regular_parts: + openai_msg_tc["content"] = regular_parts + openai_messages.append(openai_msg_tc) + elif regular_parts: + if isinstance(converted, str): + openai_messages.append({"role": role, "content": converted}) + else: + openai_messages.append({"role": role, "content": regular_parts}) + + # Emit separate tool-result messages + for tr in tool_result_markers: + openai_messages.append(tr) + + return openai_messages + + +def _has_tool_use(content: Any) -> bool: + """Check if content contains ToolUseContent.""" + if isinstance(content, list): + return any(getattr(c, "type", None) == "tool_use" for c in content) + return getattr(content, "type", None) == "tool_use" + + +def _has_tool_result(content: Any) -> bool: + """Check if content contains ToolResultContent.""" + if isinstance(content, list): + return any(getattr(c, "type", None) == "tool_result" for c in content) + return getattr(content, "type", None) == "tool_result" + + +def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]: + """Extract OpenAI-format tool_calls from MCP ToolUseContent.""" + import json + + items = content if isinstance(content, list) else [content] + tool_calls = [] + for item in items: + if getattr(item, "type", None) == "tool_use": + tool_calls.append( + { + "id": getattr(item, "id", f"call_{id(item)}"), + "type": "function", + "function": { + "name": getattr(item, "name", ""), + "arguments": json.dumps( + getattr(item, "input", {}), default=str + ), + }, + } + ) + return tool_calls + + +def _extract_text_parts(content: Any) -> Optional[str]: + """Extract text parts from mixed content.""" + items = content if isinstance(content, list) else [content] + texts = [] + for item in items: + if getattr(item, "type", None) == "text": + texts.append(getattr(item, "text", "")) + return "\n".join(texts) if texts else None + + +def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: + """Extract OpenAI-format tool messages from MCP ToolResultContent.""" + items = content if isinstance(content, list) else [content] + results = [] + for item in items: + if getattr(item, "type", None) == "tool_result": + tool_use_id = getattr(item, "toolUseId", "") + # Extract text from nested content + nested_content = getattr(item, "content", []) + if isinstance(nested_content, list): + text_parts = [ + getattr(c, "text", str(c)) + for c in nested_content + if getattr(c, "type", None) == "text" + ] + result_text = "\n".join(text_parts) if text_parts else "" + else: + result_text = str(nested_content) + results.append( + { + "role": "tool", + "tool_call_id": tool_use_id, + "content": result_text, + } + ) + return results + + +def _convert_mcp_tools_to_openai( + tools: Optional[List["Tool"]], +) -> Optional[List[Dict[str, Any]]]: + """ + Convert MCP Tool definitions to OpenAI function calling format. + MCP Tool: {name, description, inputSchema} + OpenAI Tool: {type: "function", function: {name, description, parameters}} + """ + if not tools: + return None + openai_tools = [] + for tool in tools: + openai_tool = { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description or "", + "parameters": tool.inputSchema + or { + "type": "object", + "properties": {}, + }, + }, + } + openai_tools.append(openai_tool) + return openai_tools + + +def _convert_mcp_tool_choice_to_openai( + tool_choice: Optional["ToolChoice"], +) -> Optional[Union[str, Dict[str, Any]]]: + """ + Convert MCP ToolChoice to OpenAI tool_choice format. + MCP: {mode: "auto"} | {mode: "required"} | {mode: "none"} + OpenAI: "auto" | "required" | "none" + """ + if not tool_choice: + return None + mode = getattr(tool_choice, "mode", "auto") + if mode == "auto": + return "auto" + elif mode == "required": + return "required" + elif mode == "none": + return "none" + return "auto" + + +def _convert_openai_response_to_mcp_result( + response: Any, + model_name: str, +) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: + """ + Convert a litellm completion response to MCP CreateMessageResult. + Args: + response: The litellm ModelResponse. + model_name: The model that was used. + Returns: + MCP CreateMessageResult or CreateMessageResultWithTools. + """ + if not response.choices: + verbose_logger.warning( + "MCP sampling: LLM returned empty choices list for model=%s " + "(possible content filter or provider error)", + model_name, + ) + return ErrorData( + code=-1, + message=( + f"LLM returned no choices for model '{model_name}'. " + "This may indicate content filtering or a provider-side error." + ), + ) + choice = response.choices[0] + message = choice.message + # Determine stop reason + finish_reason = getattr(choice, "finish_reason", "stop") + if finish_reason == "tool_calls": + stop_reason = "toolUse" + elif finish_reason == "length": + stop_reason = "maxTokens" + else: + stop_reason = "endTurn" + actual_model = getattr(response, "model", model_name) or model_name + # Check if response has tool calls + tool_calls = getattr(message, "tool_calls", None) + if tool_calls: + # Build ToolUseContent items + content_parts: "List[Any]" = [] + # Include text content if present + if message.content: + content_parts.append(TextContent(type="text", text=message.content)) + # Convert tool calls to MCP ToolUseContent + for tc in tool_calls: + import json + + tool_input = tc.function.arguments + if isinstance(tool_input, str): + try: + tool_input = json.loads(tool_input) + except (json.JSONDecodeError, TypeError): + tool_input = {"raw": tool_input} + content_parts.append( + ToolUseContent( + type="tool_use", + id=tc.id, + name=tc.function.name, + input=tool_input, + ) + ) + return CreateMessageResultWithTools( + role="assistant", + content=content_parts, + model=actual_model, + stopReason=stop_reason, + ) + # Simple text response + text = message.content or "" + return CreateMessageResult( + role="assistant", + content=TextContent(type="text", text=text), + model=actual_model, + stopReason=stop_reason, + ) + + +async def _check_model_access( # noqa: PLR0915 + model: str, user_api_key_auth: Any +) -> Optional["ErrorData"]: + """Enforce model-permission checks for MCP sampling requests. + + Runs the same authorization checks as ``/chat/completions``: + key-level, team-level, per-member, user-level, and project-level + model restrictions. The model name comes from the upstream MCP + server (untrusted input). + + Returns None if authorized, or an ErrorData describing the denial. + """ + if user_api_key_auth is None: + return None + + _api_key = getattr(user_api_key_auth, "api_key", None) + _token = getattr(user_api_key_auth, "token", None) + _user_role = getattr(user_api_key_auth, "user_role", None) + + _has_real_credential = bool(_api_key) or bool(_token) + _is_admin = ( + _user_role in ("proxy_admin", "proxy_admin_viewer") if _user_role else False + ) + + if not _has_real_credential and not _is_admin: + verbose_logger.warning( + "MCP sampling: denying model access for model=%s — " + "auth context has no real LiteLLM credential (possible " + "OAuth passthrough placeholder). api_key=%s, token=%s, role=%s", + model, + bool(_api_key), + bool(_token), + _user_role, + ) + return ErrorData( + code=-1, + message=( + "Model access denied: sampling requires a valid LiteLLM " + "API key or admin credential. OAuth-only sessions cannot " + "trigger proxy model calls without explicit authorization." + ), + ) + + try: + import litellm + from litellm.proxy.auth.auth_checks import ( + can_key_call_model, + can_team_access_model, + can_user_call_model, + can_project_access_model, + _check_team_member_model_access, + get_team_object, + get_user_object, + get_project_object, + ) + + try: + from litellm.proxy.proxy_server import llm_router as _llm_router + except ImportError: + _llm_router = None + + await can_key_call_model( + model=model, + llm_model_list=getattr(litellm, "model_list", None), + valid_token=user_api_key_auth, + llm_router=_llm_router, + ) + + _team_id = getattr(user_api_key_auth, "team_id", None) + _user_id = getattr(user_api_key_auth, "user_id", None) + _project_id = getattr(user_api_key_auth, "project_id", None) + + try: + from litellm.proxy.proxy_server import ( + prisma_client as _prisma_client, + user_api_key_cache as _user_api_key_cache, + proxy_logging_obj as _proxy_logging_obj, + ) + except ImportError: + _prisma_client = None + _user_api_key_cache = None # type: ignore[assignment] + _proxy_logging_obj = None # type: ignore[assignment] + + if _team_id and _prisma_client and _user_api_key_cache: + try: + team_obj = await get_team_object( + team_id=_team_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + team_obj = None + + if team_obj: + await can_team_access_model( + model=model, + team_object=team_obj, + llm_router=_llm_router, + team_model_aliases=getattr( + user_api_key_auth, "team_model_aliases", None + ), + ) + if _user_id and _proxy_logging_obj: + await _check_team_member_model_access( + model=model, + team_object=team_obj, + valid_token=user_api_key_auth, + llm_router=_llm_router, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + elif not _team_id and _user_id and _prisma_client and _user_api_key_cache: + try: + user_obj = await get_user_object( + user_id=_user_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + user_obj = None + + if user_obj: + await can_user_call_model( + model=model, + llm_router=_llm_router, + user_object=user_obj, + ) + + if _project_id and _prisma_client and _user_api_key_cache: + try: + project_obj = await get_project_object( + project_id=_project_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + project_obj = None + + if project_obj: + can_project_access_model( + model=model, + project_object=project_obj, + llm_router=_llm_router, + ) + + verbose_logger.debug( + "MCP sampling: model access check passed for model=%s", + model, + ) + return None + except Exception as access_err: + verbose_logger.warning( + "MCP sampling: model access denied for model=%s: %s", + model, + access_err, + ) + return ErrorData( + code=-1, + message=( + f"Model access denied: the API key is not authorized " + f"to use model '{model}'. {access_err}" + ), + ) + + +async def _run_budget_checks( + model: str, + user_api_key_auth: Any, + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, +) -> Optional["ErrorData"]: + """Enforce key/team/user/org/global budget checks for sampling requests. + + Runs the same ``common_checks`` path that ``/chat/completions`` uses, + so sampling cannot bypass budget limits. + + Returns None if all checks pass, or an ErrorData describing the denial. + """ + try: + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.proxy_server import ( + general_settings, + llm_router as _llm_router, + prisma_client as _prisma_client, + proxy_logging_obj as _proxy_logging_obj, + user_api_key_cache as _user_api_key_cache, + ) + from litellm.proxy.auth.auth_checks import ( + get_team_object, + get_user_object, + ) + import litellm + except ImportError as import_err: + verbose_logger.warning( + "MCP sampling: budget check imports unavailable: %s", import_err + ) + return None # Can't enforce budgets without the modules + + _team_id = getattr(user_api_key_auth, "team_id", None) + _user_id = getattr(user_api_key_auth, "user_id", None) + + team_obj = None + if _team_id and _prisma_client and _user_api_key_cache: + try: + team_obj = await get_team_object( + team_id=_team_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + pass + + user_obj = None + if _user_id and _prisma_client and _user_api_key_cache: + try: + user_obj = await get_user_object( + user_id=_user_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + pass + + dummy_request = _build_sampling_request( + raw_headers=raw_headers, + client_ip=client_ip, + ) + + # Enforce virtual-key route restrictions: a key limited to MCP routes + # must not be able to trigger a /chat/completions call via sampling. + # This mirrors the RouteChecks.should_call_route gate that runs in + # user_api_key_auth before common_checks for regular requests. + try: + from litellm.proxy.auth.route_checks import RouteChecks + + RouteChecks.should_call_route( + route="/chat/completions", + valid_token=user_api_key_auth, + request=dummy_request, + ) + except HTTPException as route_err: + verbose_logger.warning( + "MCP sampling: route check denied /chat/completions for key: %s", + route_err.detail, + ) + return ErrorData( + code=-1, + message=f"Sampling denied: virtual key is not allowed to call /chat/completions. {route_err.detail}", + ) + + global_proxy_spend = getattr(litellm, "_global_proxy_spend", None) + + # Build request body and merge x-litellm-tags from MCP headers BEFORE + # common_checks runs. _tag_max_budget_check inside common_checks only + # inspects request_body; without this pre-merge, header-supplied tags + # bypass per-tag budget enforcement (mirroring the regular auth path). + request_body: Dict[str, Any] = {"model": model} + try: + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=dummy_request, + request_data=request_body, + user_api_key_dict=user_api_key_auth, + ) + except Exception: + # Non-fatal: tag merge is defense-in-depth; don't block sampling + # if the merge utility is unavailable or fails. + pass + + try: + await common_checks( + request_body=request_body, + team_object=team_obj, + user_object=user_obj, + end_user_object=None, + global_proxy_spend=global_proxy_spend, + general_settings=general_settings or {}, + route="/chat/completions", + llm_router=_llm_router, + proxy_logging_obj=typing.cast("ProxyLogging", _proxy_logging_obj), + valid_token=user_api_key_auth, + request=dummy_request, + ) + except Exception as budget_err: + verbose_logger.warning( + "MCP sampling: budget check failed for model=%s: %s", + model, + budget_err, + ) + return ErrorData( + code=-1, + message=f"Sampling denied: {budget_err}", + ) + + verbose_logger.debug("MCP sampling: budget checks passed for model=%s", model) + return None + + +def _build_sampling_request( + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, +) -> Any: + """Build a synthetic FastAPI Request for sampling sub-calls. + + Converts the original MCP connection's HTTP headers into ASGI + scope format so that ``add_litellm_data_to_request`` can apply + header-dependent guardrails, tag-based routing, trace correlation, + and ``forward_llm_provider_auth_headers``. + + Key fields populated: + - **headers**: All original HTTP headers are forwarded (except + hop-by-hop: content-length, transfer-encoding). This ensures + ``traceparent``, ``authorization``, ``user-agent``, and + ``x-litellm-api-key`` are visible to pre-call utils. + - **client**: The ASGI ``(host, port)`` tuple so that + ``request.client.host`` returns the real client IP for + IP-based routing and guardrails. + - **server**: Derived from the running proxy's ``server_host`` + / ``server_port`` when available, avoiding the misleading + ``127.0.0.1:0`` placeholder. + - **x-forwarded-for**: Injected from ``client_ip`` if the + original headers don't already carry it, as a fallback for + IP attribution. + """ + from fastapi import Request + + # --- Build ASGI headers --- + _scope_headers: list = [(b"content-type", b"application/json")] + # Hop-by-hop headers that must NOT be forwarded into the + # synthetic request (they describe the original HTTP framing, + # not the logical request). + _HOP_BY_HOP = frozenset( + { + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "te", + "trailer", + } + ) + if raw_headers: + for hdr_name, hdr_value in raw_headers.items(): + _key = hdr_name.lower() + # Skip content-type (already set), x-forwarded-for (use resolved + # client_ip instead to prevent spoofing), and hop-by-hop headers + if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP: + continue + _scope_headers.append( + ( + _key.encode("latin-1", errors="replace"), + hdr_value.encode("utf-8"), + ) + ) + + # Inject x-forwarded-for from captured client_ip if the + # original headers don't already carry it + if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers): + _scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8"))) + + # --- Derive server (host, port) from the running proxy --- + _server_host = "127.0.0.1" + _server_port = 4000 # LiteLLM default + try: + import litellm.proxy.proxy_server as proxy_server + + _proxy_host = getattr(proxy_server, "server_host", None) + _proxy_port = getattr(proxy_server, "server_port", None) + + if _proxy_host: + _server_host = str(_proxy_host) + if _proxy_port: + _server_port = int(_proxy_port) + except (ImportError, AttributeError, TypeError, ValueError): + pass + + # --- Build ASGI client tuple for request.client.host --- + _client_tuple = None + if client_ip: + _client_tuple = (client_ip, 0) + + scope: Dict[str, Any] = { + "type": "http", + "method": "POST", + "path": "/mcp/sampling/createMessage", + "scheme": "http", + "server": (_server_host, _server_port), + "query_string": b"", + "root_path": "", + "headers": _scope_headers, + } + if _client_tuple is not None: + scope["client"] = _client_tuple + + return Request(scope=scope) + + +async def _build_completion_kwargs( + params: "CreateMessageRequestParams", + model: str, + user_api_key_auth: Any, + raw_headers: Optional[Dict[str, str]], + client_ip: Optional[str], +) -> Dict[str, Any]: + openai_messages = _convert_mcp_messages_to_openai( + messages=params.messages, + system_prompt=params.systemPrompt, + ) + completion_kwargs: Dict[str, Any] = { + "model": model, + "messages": openai_messages, + "max_tokens": params.maxTokens, + } + if params.temperature is not None: + completion_kwargs["temperature"] = params.temperature + if params.stopSequences: + completion_kwargs["stop"] = params.stopSequences + openai_tools = _convert_mcp_tools_to_openai(params.tools) + if openai_tools: + completion_kwargs["tools"] = openai_tools + openai_tool_choice = _convert_mcp_tool_choice_to_openai(params.toolChoice) + if openai_tool_choice is not None: + completion_kwargs["tool_choice"] = openai_tool_choice + completion_kwargs["metadata"] = {} + if params.metadata: + completion_kwargs["metadata"]["mcp_metadata"] = params.metadata + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import proxy_config + + completion_kwargs["user"] = getattr(user_api_key_auth, "user_id", None) + _dummy_request = _build_sampling_request( + raw_headers=raw_headers, client_ip=client_ip + ) + completion_kwargs = await add_litellm_data_to_request( + data=completion_kwargs, + request=_dummy_request, + user_api_key_dict=user_api_key_auth, + proxy_config=proxy_config, + ) + return completion_kwargs + + +async def _run_guardrails_and_call_llm( + completion_kwargs: Dict[str, Any], + user_api_key_auth: Any, +) -> Any: + try: + from litellm.proxy.proxy_server import proxy_logging_obj as _plo + + if _plo is not None: + completion_kwargs = await typing.cast("ProxyLogging", _plo).pre_call_hook( + user_api_key_dict=user_api_key_auth, + data=completion_kwargs, + call_type="acompletion", + ) + except ImportError: + pass + except Exception as guardrail_err: + verbose_logger.warning( + "MCP sampling: pre-call guardrail rejected request: %s", + guardrail_err, + ) + raise + + import litellm + + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is not None: + return await llm_router.acompletion(**completion_kwargs) + return await litellm.acompletion(**completion_kwargs) + except ImportError: + return await litellm.acompletion(**completion_kwargs) + + +async def handle_sampling_create_message( + context: Any, + params: "CreateMessageRequestParams", + default_model: Optional[str] = None, + user_api_key_auth: Optional[Any] = None, + raw_headers: Optional[Dict[str, str]] = None, + client_ip: Optional[str] = None, +) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: + """ + Handle an MCP sampling/createMessage request by routing through LiteLLM. + This is the main entry point called by the MCP client session when an + upstream MCP server requests LLM inference. + Args: + context: MCP RequestContext (contains session info). + params: The CreateMessageRequestParams from the MCP server. + default_model: Default model to use if no preferences match. + user_api_key_auth: Auth context for the requesting user. + raw_headers: Original HTTP headers from the MCP connection. + Forwarded into the internal acompletion call so that + header-dependent guardrails, IP-routing, trace-id + correlation, and forward_llm_provider_auth_headers + work correctly for sampling sub-calls. + client_ip: Original client IP address for IP-based guardrails. + Returns: + CreateMessageResult with the LLM's response, or ErrorData on failure. + """ + if not MCP_SAMPLING_AVAILABLE: + return ErrorData( + code=-1, + message="MCP sampling is not available (mcp package not installed)", + ) + + if user_api_key_auth is None: + return ErrorData( + code=-1, + message=( + "Sampling requires an authenticated user context. " + "Internal or unauthenticated sessions cannot trigger " + "upstream-initiated model calls." + ), + ) + + try: + model = _resolve_model_from_preferences( + model_preferences=params.modelPreferences, + default_model=default_model, + ) + verbose_logger.info( + "MCP sampling: resolved model=%s from preferences=%s", + model, + params.modelPreferences, + ) + + access_denial = await _check_model_access(model, user_api_key_auth) + if access_denial is not None: + return access_denial + + budget_denial = await _run_budget_checks( + model=model, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if budget_denial is not None: + return budget_denial + + completion_kwargs = await _build_completion_kwargs( + params=params, + model=model, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + + openai_messages = completion_kwargs["messages"] + openai_tools = completion_kwargs.get("tools") + verbose_logger.debug( + "MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s", + model, + len(openai_messages), + bool(openai_tools), + ) + + response = await _run_guardrails_and_call_llm( + completion_kwargs=completion_kwargs, + user_api_key_auth=user_api_key_auth, + ) + + result = _convert_openai_response_to_mcp_result( + response=response, model_name=model + ) + verbose_logger.info( + "MCP sampling: completed successfully, model=%s, stopReason=%s", + getattr(result, "model", "unknown"), + getattr(result, "stopReason", "unknown"), + ) + return result + except Exception as e: + from litellm.exceptions import ( + AuthenticationError, + BudgetExceededError, + ContextWindowExceededError, + PermissionDeniedError, + RateLimitError, + ServiceUnavailableError, + ) + + from litellm.proxy._types import ProxyException + + if isinstance( + e, + ( + HTTPException, + BudgetExceededError, + RateLimitError, + AuthenticationError, + PermissionDeniedError, + ContextWindowExceededError, + ServiceUnavailableError, + ProxyException, + ), + ): + raise + + verbose_logger.exception("MCP sampling handler failed: %s", e) + return ErrorData( + code=-1, + message=f"Sampling failed: {str(e)}", + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6e33a105ec..df6cb22fda 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/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 @@ -125,6 +126,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}") @@ -160,6 +173,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 @@ -483,10 +550,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 ( @@ -497,7 +572,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}" ) @@ -528,152 +603,178 @@ 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 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 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 ( @@ -684,7 +785,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}" ) @@ -713,6 +814,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( @@ -730,33 +834,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, @@ -766,7 +850,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}" ) @@ -792,10 +914,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, @@ -805,7 +937,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}" ) @@ -825,8 +957,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: @@ -834,30 +965,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 ################## @@ -1180,8 +1325,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, ) @@ -1207,8 +1351,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, ) @@ -1230,8 +1373,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, @@ -1275,8 +1417,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, @@ -2485,7 +2626,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, ) @@ -2744,8 +2885,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." ), ) @@ -3124,8 +3264,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 @@ -3136,8 +3275,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( @@ -3615,6 +3753,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) @@ -3639,8 +3778,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, @@ -3672,9 +3810,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() ) @@ -4099,6 +4284,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 ############# ######################################################## diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 20c951d350..f0d8ddf97d 100644 --- a/litellm/proxy/caching_routes.py +++ b/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), diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 27fa685eaa..b0932015ab 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/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 ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a092df6cdf..7aed9ad894 100644 --- a/litellm/proxy/proxy_server.py +++ b/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") diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index 38e6d53344..e24eb4aebb 100644 --- a/litellm/types/llms/gemini.py +++ b/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. diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 6aa62c3510..2108fe8990 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -115,6 +115,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 diff --git a/litellm/utils.py b/litellm/utils.py index 7cac830b2c..d010391229 100644 --- a/litellm/utils.py +++ b/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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ed6de4fa6b..4c227656e5 100644 --- a/model_prices_and_context_window.json +++ b/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" + } + } + diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py index 47b95c27ab..c4f15ac4c3 100644 --- a/tests/llm_translation/test_fireworks_ai_translation.py +++ b/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): diff --git a/tests/local_testing/test_custom_llm.py b/tests/local_testing/test_custom_llm.py index 34ab6c043b..ea15c3db9d 100644 --- a/tests/local_testing/test_custom_llm.py +++ b/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 diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 46f2f89b3c..1c041be094 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/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", "") diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py new file mode 100644 index 0000000000..3c280c6ba9 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 63e2cb7f35..b2002f9a0f 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/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" diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index 9d2787f78a..59472d1a49 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/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"} diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py new file mode 100644 index 0000000000..e2133d56f8 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 2061522fef..ca340b5f27 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/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"}} diff --git a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py b/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py index 4cf2429d73..6f215deed4 100644 --- a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py +++ b/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 = { diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 560796ea58..8a072fa509 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/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)""" diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 6f32c4ca34..74888e6cd9 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/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"] ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py new file mode 100644 index 0000000000..b93f0d56f8 --- /dev/null +++ b/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"]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 04ff1e4be2..363948ff4e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/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() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py new file mode 100644 index 0000000000..78aee7b534 --- /dev/null +++ b/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"]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py new file mode 100644 index 0000000000..f141cb2e31 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_resolution.py new file mode 100644 index 0000000000..0c8f7bd481 --- /dev/null +++ b/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"]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_priority_selection.py new file mode 100644 index 0000000000..24309ed046 --- /dev/null +++ b/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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_request_builder.py new file mode 100644 index 0000000000..d5c636baea --- /dev/null +++ b/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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py new file mode 100644 index 0000000000..bb17a8f710 --- /dev/null +++ b/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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py new file mode 100644 index 0000000000..b4b219e958 --- /dev/null +++ b/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"] + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 6b6c7bc37d..227cf3f4bc 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 32e4ec1931..e7d0ee6247 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -320,9 +320,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() @@ -375,9 +373,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() @@ -414,9 +410,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" @@ -457,7 +451,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 @@ -507,7 +501,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 @@ -560,7 +554,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 @@ -616,7 +610,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 @@ -1166,9 +1160,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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 549afd774b..d52af94c47 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/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) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index d31dbdd434..6804ea9f8f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/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 diff --git a/tests/test_litellm/proxy/test_caching_routes.py b/tests/test_litellm/proxy/test_caching_routes.py index 3e842d118d..840ba054cc 100644 --- a/tests/test_litellm/proxy/test_caching_routes.py +++ b/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): diff --git a/tests/test_litellm/proxy/test_dynamic_mcp_route.py b/tests/test_litellm/proxy/test_dynamic_mcp_route.py index 2462aff211..592cebd957 100644 --- a/tests/test_litellm/proxy/test_dynamic_mcp_route.py +++ b/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() diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index b03579c2db..113e1bc0df 100644 --- a/tests/test_litellm/test_main.py +++ b/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") diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2d75671f1c..f2b9bef923 100644 --- a/tests/test_litellm/test_utils.py +++ b/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