From 47017e4812492aeef26f9e320c7e281226feafc0 Mon Sep 17 00:00:00 2001 From: Yug Date: Thu, 30 Apr 2026 22:24:34 +0530 Subject: [PATCH] resolve --- .../_experimental/mcp_server/mcp_server_manager.py | 4 ++-- .../_experimental/mcp_server/sampling_handler.py | 9 +++++++-- litellm/proxy/_types.py | 12 ++++++------ tests/mcp_tests/test_mcp_session_binding.py | 10 +++++++--- 4 files changed, 22 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 75705d62219..f3fd1360525 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1229,13 +1229,13 @@ class MCPServerManager: if server.command: base_command = os.path.basename(server.command) # Strip .exe/.cmd/.bat/.com suffix for Windows compatibility - base_command_no_ext = base_command + 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 not in MCP_STDIO_ALLOWED_COMMANDS + base_command.lower() not in MCP_STDIO_ALLOWED_COMMANDS and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS ): raise HTTPException( diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 7a9be95404e..d2071605a87 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -88,8 +88,13 @@ def _resolve_model_from_preferences( # Fall back to first available model if available_model_names: return available_model_names[0] - # Last resort - use LiteLLM default or return None - return getattr(litellm, "default_mcp_sampling_model", None) or "gpt-4o-mini" + # Last resort - use LiteLLM default or raise error + default_sampling_model = getattr(litellm, "default_mcp_sampling_model", None) + if 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 _convert_mcp_content_to_openai( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f47f3c6b5b0..e3f7eb75c78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1217,13 +1217,13 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): # Validate command against allowlist to prevent arbitrary execution base_command = os.path.basename(values["command"]) # Strip .exe/.cmd/.bat/.com suffix for Windows compatibility - base_command_no_ext = base_command + 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() + base_command_no_ext = base_command.lower()[: -len(ext)].lower() break if ( - base_command not in MCP_STDIO_ALLOWED_COMMANDS + base_command.lower() not in MCP_STDIO_ALLOWED_COMMANDS and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS ): raise ValueError( @@ -1294,13 +1294,13 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): # Validate command against allowlist to prevent arbitrary execution base_command = os.path.basename(values["command"]) # Strip .exe/.cmd/.bat/.com suffix for Windows compatibility - base_command_no_ext = base_command + 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() + base_command_no_ext = base_command.lower()[: -len(ext)].lower() break if ( - base_command not in MCP_STDIO_ALLOWED_COMMANDS + base_command.lower() not in MCP_STDIO_ALLOWED_COMMANDS and base_command_no_ext not in MCP_STDIO_ALLOWED_COMMANDS ): raise ValueError( diff --git a/tests/mcp_tests/test_mcp_session_binding.py b/tests/mcp_tests/test_mcp_session_binding.py index c3d2313b046..1fa1c24d064 100644 --- a/tests/mcp_tests/test_mcp_session_binding.py +++ b/tests/mcp_tests/test_mcp_session_binding.py @@ -7,7 +7,7 @@ from contextlib import asynccontextmanager from litellm.proxy._experimental.mcp_server.server import ( handle_sse_mcp_endpoint, handle_sse_post_messages, - _captured_session_id_var, + _captured_session_id_container_var, _session_id_auth_storage, ) from litellm.proxy._types import UserAPIKeyAuth @@ -29,7 +29,9 @@ async def test_session_id_capture_and_binding(): @asynccontextmanager async def mock_connect_sse(scope, receive, send): # Directly set the ContextVar to simulate our capturing dict's behavior - _captured_session_id_var.set(session_id) + container = _captured_session_id_container_var.get() + if container is not None: + container["session_id"] = session_id yield (AsyncMock(), AsyncMock()) mock_sse.connect_sse = MagicMock(side_effect=mock_connect_sse) @@ -174,7 +176,9 @@ async def test_anonymous_session_still_works(): @asynccontextmanager async def mock_connect_sse(scope, receive, send): - _captured_session_id_var.set(session_id) + container = _captured_session_id_container_var.get() + if container is not None: + container["session_id"] = session_id yield (AsyncMock(), AsyncMock()) mock_sse.connect_sse = MagicMock(side_effect=mock_connect_sse)