mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix: fix connection switching and return more clearer errors
This commit is contained in:
Krrish Dholakia
parent
36ca05739c
commit
bde027f0d4
3 changed files with 72 additions and 102 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -13,9 +13,6 @@ interface ProxySwitcherProps {
|
|||
const ProxySwitcher: React.FC<ProxySwitcherProps> = ({ 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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue