diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5b3d5bd60e2..86a89f16388 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1,25 +1,16 @@ """ LiteLLM MCP Server Routes """ + # pyright: reportInvalidTypeForm=false, reportArgumentType=false, reportOptionalCall=false import asyncio import contextlib - import traceback import uuid from datetime import datetime -from typing import ( - Any, - AsyncIterator, - Callable, - Dict, - List, - Optional, - Tuple, - Union, - cast, -) +from typing import (Any, AsyncIterator, Callable, Dict, List, Optional, Tuple, + Union, cast) from fastapi import FastAPI, HTTPException from pydantic import AnyUrl, ConfigDict @@ -29,19 +20,16 @@ from starlette.types import Receive, Scope, Send from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, -) -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, -) +from litellm.litellm_core_utils.litellm_logging import \ + Logging as LiteLLMLoggingObj +from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import \ + MCPRequestHandler +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import \ + get_request_base_url from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.utils import ( - LITELLM_MCP_SERVER_DESCRIPTION, - LITELLM_MCP_SERVER_NAME, - LITELLM_MCP_SERVER_VERSION, -) + LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, + LITELLM_MCP_SERVER_VERSION) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup @@ -59,12 +47,8 @@ try: from mcp import ReadResourceResult, Resource from mcp.server import Server from mcp.server.lowlevel.helper_types import ReadResourceContents - from mcp.types import ( - BlobResourceContents, - GetPromptResult, - ResourceTemplate, - TextResourceContents, - ) + from mcp.types import (BlobResourceContents, GetPromptResult, + ResourceTemplate, TextResourceContents) except ImportError as e: verbose_logger.debug(f"MCP module not found: {e}") MCP_AVAILABLE = False @@ -88,39 +72,29 @@ _INITIALIZATION_LOCK = asyncio.Lock() if MCP_AVAILABLE: from mcp.server import Server - # Import auth context variables and middleware - from mcp.server.auth.middleware.auth_context import ( - AuthContextMiddleware, - auth_context_var, - ) + from mcp.server.auth.middleware.auth_context import (AuthContextMiddleware, + auth_context_var) try: - from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.server.streamable_http_manager import \ + StreamableHTTPSessionManager except ImportError: StreamableHTTPSessionManager = None # type: ignore - from mcp.types import ( - CallToolResult, - EmbeddedResource, - ImageContent, - Prompt, - TextContent, - ) + from mcp.types import (CallToolResult, EmbeddedResource, ImageContent, + Prompt, TextContent) from mcp.types import Tool as MCPTool - from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( - MCPAuthenticatedUser, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport - from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, - ) - from litellm.proxy._experimental.mcp_server.utils import ( - split_server_prefix_from_name, - ) + from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import \ + MCPAuthenticatedUser + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.sse_transport import \ + SseServerTransport + from litellm.proxy._experimental.mcp_server.tool_registry import \ + global_mcp_tool_registry + from litellm.proxy._experimental.mcp_server.utils import \ + split_server_prefix_from_name ###################################################### ############ MCP Tools List REST API Response Object # @@ -148,7 +122,7 @@ if MCP_AVAILABLE: session_manager = StreamableHTTPSessionManager( app=server, event_store=None, - json_response=False, # enables SSE streaming + json_response=False, # enables SSE streaming stateless=True, ) @@ -286,8 +260,10 @@ if MCP_AVAILABLE: """ 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.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 # Validate arguments @@ -307,9 +283,9 @@ if MCP_AVAILABLE: 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: + 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: float | None): @@ -318,14 +294,20 @@ if MCP_AVAILABLE: await host_session.send_progress_notification( progress_token=host_token, progress=progress, - total=total + total=total, + ) + verbose_logger.debug( + f"Forwarded progress {progress}/{total} to Host" ) - 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}") + 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]}...") + 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: @@ -680,9 +662,8 @@ if MCP_AVAILABLE: Returns: True if the tool name (prefixed or unprefixed) is in the filter list """ - from litellm.proxy._experimental.mcp_server.utils import ( - split_server_prefix_from_name, - ) + from litellm.proxy._experimental.mcp_server.utils import \ + split_server_prefix_from_name # Check if the full name is in the list if tool_name in filter_list: @@ -769,9 +750,7 @@ if MCP_AVAILABLE: ) allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth - ) + await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) ) allowed_mcp_server_ids, _ip_blocked = ( global_mcp_server_manager.filter_server_ids_by_ip_with_info( @@ -780,7 +759,8 @@ if MCP_AVAILABLE: ) verbose_logger.debug( "MCP IP filter: client_ip=%s, allowed_server_ids=%s", - client_ip, allowed_mcp_server_ids, + client_ip, + allowed_mcp_server_ids, ) if _ip_blocked > 0: verbose_logger.debug( @@ -805,7 +785,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) - return allowed_mcp_servers @@ -1087,7 +1066,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - # Get prompts from each allowed server all_prompts = [] for server in allowed_mcp_servers: @@ -1146,7 +1124,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - all_resources: List[Resource] = [] for server in allowed_mcp_servers: if server is None: @@ -1202,7 +1179,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - all_resource_templates: List[ResourceTemplate] = [] for server in allowed_mcp_servers: if server is None: @@ -1491,7 +1467,7 @@ if MCP_AVAILABLE: ): raise HTTPException( status_code=403, - detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}", + detail=f"User not allowed to call this MCP server.\nTrying to call server: {server_name}\nAllowed MCP servers: {allowed_mcp_servers}", ) standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = ( @@ -1505,9 +1481,9 @@ if MCP_AVAILABLE: "litellm_logging_obj", None ) if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) litellm_logging_obj.model = f"MCP: {name}" # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names @@ -1533,9 +1509,9 @@ if MCP_AVAILABLE: ).get("mcp_server_cost_info") # Update model_call_details with the cost info if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) response = await _handle_managed_mcp_tool( server_name=server_name, name=original_tool_name, # Pass the full name (potentially prefixed) @@ -1680,7 +1656,6 @@ if MCP_AVAILABLE: detail="User not allowed to get this prompt.", ) - # Extract server name from prefixed prompt name original_prompt_name, server_name = split_server_prefix_from_name(name) @@ -2224,17 +2199,15 @@ if MCP_AVAILABLE: ) auth_context_var.set(auth_user) - def get_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], - ] - ): + def get_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 the UserAPIKeyAuth from the auth context variable. diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 36efa4ebcaf..064bf2ff650 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -66,10 +66,10 @@ async def batch_upsert_tools( await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" ' "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) " - "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) " + "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8::timestamp, $8::timestamp) " "ON CONFLICT (tool_name) DO UPDATE SET " "call_count = \"LiteLLM_ToolTable\".call_count + 1, " - "updated_at = $8", + "updated_at = $8::timestamp", tool_name, origin, created_by, @@ -143,8 +143,8 @@ async def update_tool_policy( now = datetime.now(timezone.utc) await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) ' - "VALUES ($4, $1, $2, $3, $3, $5, $5) " - "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5", + "VALUES ($4, $1, $2, $3, $3, $5::timestamp, $5::timestamp) " + "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5::timestamp", tool_name, call_policy, _updated_by, diff --git a/ui/litellm-dashboard/src/components/Navbar/ProxySwitcher/ProxySwitcher.tsx b/ui/litellm-dashboard/src/components/Navbar/ProxySwitcher/ProxySwitcher.tsx index 3c5d994ed1d..b9d48bc7478 100644 --- a/ui/litellm-dashboard/src/components/Navbar/ProxySwitcher/ProxySwitcher.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/ProxySwitcher/ProxySwitcher.tsx @@ -13,9 +13,6 @@ interface ProxySwitcherProps { const ProxySwitcher: React.FC = ({ onManageClick }) => { const { connections, activeConnection, switchConnection } = useProxyConnection(); - // Only show when there are multiple connections configured - if (connections.length <= 1) return null; - const items: MenuProps["items"] = [ ...connections.map((conn) => ({ key: conn.id,