From 21ae1b0fbd31cd1c09d53236dca7f4ab85dbecfa Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Feb 2026 17:48:09 -0800 Subject: [PATCH] feat(mcp): add configurable timeouts for MCP tool listing Make the MCP tool listing timeout configurable via query parameter on /mcp-rest/tools/list and /v1/mcp/tools endpoints. Default remains 30s but can be overridden per-request. Fixes LIT-1789 where slow MCP servers would silently timeout and tools wouldn't appear in the UI. Changes: - Add optional timeout parameter to _fetch_tools_with_timeout() - Thread timeout through _get_tools_from_server() - Add timeout query parameter to REST and management endpoints - Update tests to support new timeout parameter Co-Authored-By: Claude Haiku 4.5 --- ...odel_prices_and_context_window_backup.json | 2 +- .../mcp_server/mcp_server_manager.py | 19 ++-- .../mcp_server/rest_endpoints.py | 8 ++ .../proxy/_experimental/mcp_server/server.py | 7 +- .../out/{404.html => 404/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{budgets.html => budgets/index.html} | 0 .../{caching.html => caching/index.html} | 0 .../index.html} | 0 .../{old-usage.html => old-usage/index.html} | 0 .../{prompts.html => prompts/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{login.html => login/index.html} | 0 .../out/{logs.html => logs/index.html} | 0 .../{callback.html => callback/index.html} | 0 .../{model-hub.html => model-hub/index.html} | 0 .../{model_hub.html => model_hub/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{policies.html => policies/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{ui-theme.html => ui-theme/index.html} | 0 .../out/{teams.html => teams/index.html} | 0 .../{test-key.html => test-key/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{usage.html => usage/index.html} | 0 .../out/{users.html => users/index.html} | 0 .../index.html} | 0 .../mcp_management_endpoints.py | 9 +- .../mcp_server/test_mcp_server.py | 15 ++- .../mcp_server/test_mcp_server_manager.py | 96 +++++++++++++++++++ .../mcp_server/test_rest_endpoints.py | 6 +- 41 files changed, 145 insertions(+), 17 deletions(-) rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{_not-found.html => _not-found/index.html} (100%) rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%) rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%) rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%) rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%) rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%) rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%) rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%) rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%) rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%) rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%) rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%) rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%) rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%) rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%) rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f52288ea72a..cbd64a178b8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16289,7 +16289,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "supports_reasoning": false, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 08213f40b43..11b829c1547 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -71,9 +71,7 @@ try: from mcp.shared.tool_name_validation import ( validate_tool_name, # pyright: ignore[reportAssignmentType] ) - from mcp.shared.tool_name_validation import ( - SEP_986_URL, - ) + from mcp.shared.tool_name_validation import SEP_986_URL except ImportError: from pydantic import BaseModel @@ -966,6 +964,7 @@ class MCPServerManager: extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, raw_headers: Optional[Dict[str, str]] = None, + timeout: Optional[float] = None, ) -> List[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -973,6 +972,7 @@ class MCPServerManager: Args: server (MCPServer): The server to query tools from mcp_auth_header: Optional auth header for MCP server + timeout: Optional timeout in seconds for tool listing (default 30.0) Returns: List[MCPTool]: List of tools available on the server with prefixed names @@ -1008,7 +1008,9 @@ class MCPServerManager: _tools ) else: - tools = await self._fetch_tools_with_timeout(client, server.name) + tools = await self._fetch_tools_with_timeout( + client, server.name, timeout=timeout + ) prefixed_or_original_tools = self._create_prefixed_tools( tools, server, add_prefix=add_prefix @@ -1473,7 +1475,10 @@ class MCPServerManager: return None async def _fetch_tools_with_timeout( - self, client: MCPClient, server_name: str + self, + client: MCPClient, + server_name: str, + timeout: Optional[float] = None, ) -> List[MCPTool]: """ Fetch tools from MCP client with timeout and error handling. @@ -1484,12 +1489,14 @@ class MCPServerManager: Args: client: MCP client instance server_name: Name of the server for logging + timeout: Optional timeout in seconds (default 30.0) Returns: List of tools from the server """ + _timeout = timeout or 30.0 try: - with anyio.fail_after(30.0): + with anyio.fail_after(_timeout): tools = await client.list_tools() verbose_logger.debug(f"Tools from {server_name}: {tools}") return tools diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 16f8f835430..548c250f0d6 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -162,6 +162,7 @@ if MCP_AVAILABLE: server_auth_header, raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, + timeout: Optional[float] = None, ): """Helper function to get tools for a single server.""" tools = await global_mcp_server_manager._get_tools_from_server( @@ -169,6 +170,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, add_prefix=False, raw_headers=raw_headers, + timeout=timeout, ) # Filter tools based on allowed_tools configuration @@ -235,6 +237,10 @@ if MCP_AVAILABLE: server_id: Optional[str] = Query( None, description="The server id to list tools for" ), + timeout: Optional[float] = Query( + None, + description="Timeout in seconds for listing tools from each MCP server. Default is 30 seconds.", + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> dict: """ @@ -340,6 +346,7 @@ if MCP_AVAILABLE: server_auth_header, raw_headers_from_request, user_api_key_dict, + timeout=timeout, ) except Exception as e: verbose_logger.exception( @@ -392,6 +399,7 @@ if MCP_AVAILABLE: server_auth_header, raw_headers_from_request, user_api_key_dict, + timeout=timeout, ) list_tools_result.extend(tools_result) except Exception as e: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5b3d5bd60e2..115f376e383 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -5,7 +5,6 @@ LiteLLM MCP Server Routes import asyncio import contextlib - import traceback import uuid from datetime import datetime @@ -859,6 +858,7 @@ if MCP_AVAILABLE: log_list_tools_to_spendlogs: bool = False, list_tools_log_source: Optional[str] = None, litellm_trace_id: Optional[str] = None, + timeout: Optional[float] = None, ) -> List[MCPTool]: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -869,6 +869,7 @@ if MCP_AVAILABLE: mcp_servers: Optional list of server names/aliases to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers oauth2_headers: Optional dict of oauth2 headers + timeout: Optional timeout in seconds for each server's tool listing (default 30.0) Returns: List[MCPTool]: Combined list of tools from filtered servers @@ -971,6 +972,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + timeout=timeout, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) @@ -1287,6 +1289,7 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, log_list_tools_to_spendlogs: bool = False, list_tools_log_source: Optional[str] = None, + timeout: Optional[float] = None, ) -> List[MCPTool]: """ List all available MCP tools. @@ -1296,6 +1299,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} + timeout: Optional timeout in seconds for each server's tool listing (default 30.0) Returns: List[MCPTool]: Combined list of tools from all accessible servers @@ -1314,6 +1318,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, list_tools_log_source=list_tools_log_source, + timeout=timeout, ) verbose_logger.debug( f"Successfully fetched {len(managed_tools)} tools from managed MCP servers" diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 8c4d4e7937e..ccc0d97886f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -38,9 +38,7 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME -from litellm.proxy._experimental.mcp_server.utils import ( - get_server_prefix, -) +from litellm.proxy._experimental.mcp_server.utils import get_server_prefix from litellm.proxy._experimental.mcp_server.utils import ( validate_and_normalize_mcp_server_payload as _base_validate_and_normalize_mcp_server_payload, ) @@ -430,6 +428,10 @@ if MCP_AVAILABLE: ) async def get_mcp_tools( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + timeout: Optional[float] = Query( + None, + description="Timeout in seconds for listing tools from each MCP server. Default is 30 seconds.", + ), ): """ Get all MCP tools available for the current key, including those from access groups @@ -441,6 +443,7 @@ if MCP_AVAILABLE: mcp_auth_header=None, mcp_servers=None, mcp_server_auth_headers=None, + timeout=timeout, ) dumped_tools = [dict(tool) for tool in tools] 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 1fcdeb627d0..0179ab7237f 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 @@ -8,11 +8,7 @@ from fastapi import HTTPException from mcp import ReadResourceResult, Resource from mcp.types import Prompt, ResourceTemplate, TextResourceContents -from litellm.proxy._types import ( - LiteLLM_MCPServerTable, - MCPTransport, - UserAPIKeyAuth, -) +from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport, UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -28,6 +24,7 @@ def cleanup_mcp_global_state(): from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + # Clear before test global_mcp_server_manager.registry.clear() global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() @@ -494,6 +491,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): extra_headers=None, add_prefix=True, raw_headers=None, + timeout=None, ): if server.name == "working_server": # Working server returns tools @@ -596,6 +594,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): extra_headers=None, add_prefix=True, raw_headers=None, + timeout=None, ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -1043,6 +1042,7 @@ async def test_list_tools_single_server_unprefixed_names(): extra_headers=None, add_prefix=False, raw_headers=None, + timeout=None, ): tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" @@ -1121,6 +1121,7 @@ async def test_list_tools_multiple_servers_prefixed_names(): extra_headers=None, add_prefix=True, raw_headers=None, + timeout=None, ): tool = MagicMock() # When multiple servers, add_prefix should be True -> prefixed names @@ -1372,6 +1373,7 @@ async def test_list_tools_filters_by_key_team_permissions(): extra_headers=None, add_prefix=False, raw_headers=None, + timeout=None, ): # Return 4 tools, but only 2 should be allowed tool1 = MagicMock() @@ -1479,6 +1481,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): extra_headers=None, add_prefix=False, raw_headers=None, + timeout=None, ): # Return 4 tools tool1 = MagicMock() @@ -1571,6 +1574,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): extra_headers=None, add_prefix=False, raw_headers=None, + timeout=None, ): # Return 3 tools tool1 = MagicMock() @@ -1666,6 +1670,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): extra_headers=None, add_prefix=True, raw_headers=None, + timeout=None, ): # Return tools WITH prefix (as they come from MCP server) tool1 = MagicMock() 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 c105052479d..ec0013eae22 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 @@ -2307,5 +2307,101 @@ class TestMCPServerManager: assert resolved_server.server_name == "test_server" # server_name matches +class TestFetchToolsWithTimeout: + """Tests for configurable timeout in _fetch_tools_with_timeout.""" + + @pytest.mark.asyncio + async def test_fetch_tools_default_timeout(self): + """_fetch_tools_with_timeout should use 30s default when no timeout is provided.""" + manager = MCPServerManager() + mock_client = MagicMock() + mock_client.list_tools = AsyncMock( + return_value=[ + MCPTool(name="tool1", description="A tool", inputSchema={}) + ] + ) + + # Should succeed with default timeout + tools = await manager._fetch_tools_with_timeout(mock_client, "test_server") + assert len(tools) == 1 + assert tools[0].name == "tool1" + + @pytest.mark.asyncio + async def test_fetch_tools_custom_timeout(self): + """_fetch_tools_with_timeout should use the provided timeout value.""" + manager = MCPServerManager() + mock_client = MagicMock() + mock_client.list_tools = AsyncMock( + return_value=[ + MCPTool(name="tool1", description="A tool", inputSchema={}) + ] + ) + + # Should succeed with custom timeout + tools = await manager._fetch_tools_with_timeout( + mock_client, "test_server", timeout=120.0 + ) + assert len(tools) == 1 + assert tools[0].name == "tool1" + + @pytest.mark.asyncio + async def test_get_tools_from_server_passes_timeout(self): + """_get_tools_from_server should pass timeout to _fetch_tools_with_timeout.""" + manager = MCPServerManager() + + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + ) + + manager._create_mcp_client = AsyncMock(return_value=object()) + + upstream_tool = MCPTool( + name="my_tool", + description="A tool", + inputSchema={}, + ) + + manager._fetch_tools_with_timeout = AsyncMock(return_value=[upstream_tool]) + + # Call with custom timeout + await manager._get_tools_from_server(server, timeout=120.0) + + # Verify _fetch_tools_with_timeout was called with the timeout + manager._fetch_tools_with_timeout.assert_called_once() + call_kwargs = manager._fetch_tools_with_timeout.call_args + assert call_kwargs.kwargs.get("timeout") == 120.0 + + @pytest.mark.asyncio + async def test_get_tools_from_server_default_timeout(self): + """_get_tools_from_server should pass None timeout by default.""" + manager = MCPServerManager() + + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + ) + + manager._create_mcp_client = AsyncMock(return_value=object()) + + upstream_tool = MCPTool( + name="my_tool", + description="A tool", + inputSchema={}, + ) + + manager._fetch_tools_with_timeout = AsyncMock(return_value=[upstream_tool]) + + # Call without timeout + await manager._get_tools_from_server(server) + + # Verify _fetch_tools_with_timeout was called with timeout=None + manager._fetch_tools_with_timeout.assert_called_once() + call_kwargs = manager._fetch_tools_with_timeout.call_args + assert call_kwargs.kwargs.get("timeout") is None + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 4f93270c162..89c7e0951f3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -484,7 +484,11 @@ class TestListToolsRestAPI: captured = {"called": False} async def fake_get_tools( - server, server_auth_header, raw_headers=None, user_api_key_auth=None + server, + server_auth_header, + raw_headers=None, + user_api_key_auth=None, + timeout=None, ): captured["called"] = True captured["server"] = server