diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 2eea28f6cc1..9cdcd3894e0 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -205,7 +205,15 @@ class LiteLLM_Proxy_MCP_Handler: else: tool_name = getattr(mcp_tool, "name", None) - if tool_name and tool_name in allowed_tool_names: + if not tool_name: + continue + + if tool_name in allowed_tool_names: + filtered_tools.append(mcp_tool) + continue + + unprefixed_name, _ = split_server_prefix_from_name(tool_name) + if unprefixed_name in allowed_tool_names: filtered_tools.append(mcp_tool) return filtered_tools diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 1cf9c7eb8df..7acc723dd56 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -45,8 +45,8 @@ import { makeOpenAIAudioSpeechRequest } from "../llm_calls/audio_speech"; import { makeOpenAIAudioTranscriptionRequest } from "../llm_calls/audio_transcriptions"; import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion"; import { makeOpenAIEmbeddingsRequest } from "../llm_calls/embeddings_api"; -import type { MCPTool } from "../llm_calls/fetch_mcp_tools"; -import { fetchAvailableMCPTools } from "../llm_calls/fetch_mcp_tools"; +import { listMCPTools, fetchMCPServers } from "../../networking"; +import { MCPServer } from "../../mcp_tools/types"; import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models"; import { makeOpenAIImageEditsRequest } from "../llm_calls/image_edits"; import { makeOpenAIImageGenerationRequest } from "../llm_calls/image_generation"; @@ -97,20 +97,27 @@ const ChatUI: React.FC = ({ disabledPersonalKeyCreation, proxySettings, }) => { - const [isMCPToolsModalVisible, setIsMCPToolsModalVisible] = useState(false); - const [mcpTools, setMCPTools] = useState([]); - const [selectedMCPTools, setSelectedMCPTools] = useState(() => { - const saved = sessionStorage.getItem("selectedMCPTools"); + const [mcpServers, setMCPServers] = useState([]); + const [selectedMCPServers, setSelectedMCPServers] = useState(() => { + const saved = sessionStorage.getItem("selectedMCPServers"); try { - const parsed = saved ? JSON.parse(saved) : []; - // Convert from single string to array if needed for backward compatibility - return Array.isArray(parsed) ? parsed : parsed ? [parsed] : []; + return saved ? JSON.parse(saved) : []; } catch (error) { - console.error("Error parsing selectedMCPTools from sessionStorage", error); + console.error("Error parsing selectedMCPServers from sessionStorage", error); return []; } }); - const [isLoadingMCPTools, setIsLoadingMCPTools] = useState(false); + const [isLoadingMCPServers, setIsLoadingMCPServers] = useState(false); + const [serverToolsMap, setServerToolsMap] = useState>({}); + const [mcpServerToolRestrictions, setMCPServerToolRestrictions] = useState>(() => { + const saved = sessionStorage.getItem("mcpServerToolRestrictions"); + try { + return saved ? JSON.parse(saved) : {}; + } catch (error) { + console.error("Error parsing mcpServerToolRestrictions from sessionStorage", error); + return {}; + } + }); const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => { const saved = sessionStorage.getItem("apiKeySource"); if (saved) { @@ -211,27 +218,38 @@ const ChatUI: React.FC = ({ const chatEndRef = useRef(null); - // Fetch MCP tools - const loadMCPTools = async () => { + // Fetch MCP servers + const loadMCPServers = async () => { const userApiKey = apiKeySource === "session" ? accessToken : apiKey; if (!userApiKey) return; - setIsLoadingMCPTools(true); + setIsLoadingMCPServers(true); try { - const tools = await fetchAvailableMCPTools(userApiKey); - setMCPTools(tools); + const servers = await fetchMCPServers(userApiKey); + setMCPServers(Array.isArray(servers) ? servers : servers.data || []); } catch (error) { - console.error("Error fetching MCP tools:", error); + console.error("Error fetching MCP servers:", error); } finally { - setIsLoadingMCPTools(false); + setIsLoadingMCPServers(false); } }; - useEffect(() => { - if (isMCPToolsModalVisible) { - loadMCPTools(); + // Fetch tools for a specific server + const loadServerTools = async (serverId: string) => { + const userApiKey = apiKeySource === "session" ? accessToken : apiKey; + if (!userApiKey || serverToolsMap[serverId]) return; + + try { + const response = await listMCPTools(userApiKey, serverId); + setServerToolsMap(prev => ({ + ...prev, + [serverId]: response.tools || [] + })); + } catch (error) { + console.error(`Error fetching tools for server ${serverId}:`, error); } - }, [isMCPToolsModalVisible, accessToken, apiKey, apiKeySource]); + }; + useEffect(() => { if (isGetCodeModalVisible) { @@ -244,7 +262,9 @@ const ChatUI: React.FC = ({ selectedTags, selectedVectorStores, selectedGuardrails, - selectedMCPTools, + selectedMCPServers, + mcpServers, + mcpServerToolRestrictions, endpointType, selectedModel, selectedSdk, @@ -264,7 +284,9 @@ const ChatUI: React.FC = ({ selectedTags, selectedVectorStores, selectedGuardrails, - selectedMCPTools, + selectedMCPServers, + mcpServers, + mcpServerToolRestrictions, endpointType, selectedModel, proxySettings, @@ -287,8 +309,10 @@ const ChatUI: React.FC = ({ sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags)); sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores)); sessionStorage.setItem("selectedGuardrails", JSON.stringify(selectedGuardrails)); - sessionStorage.setItem("selectedMCPTools", JSON.stringify(selectedMCPTools)); + sessionStorage.setItem("selectedMCPServers", JSON.stringify(selectedMCPServers)); + sessionStorage.setItem("mcpServerToolRestrictions", JSON.stringify(mcpServerToolRestrictions)); sessionStorage.setItem("selectedVoice", selectedVoice); + sessionStorage.removeItem("selectedMCPTools"); // Clean up old key if (selectedModel) { sessionStorage.setItem("selectedModel", selectedModel); @@ -318,7 +342,8 @@ const ChatUI: React.FC = ({ messageTraceId, responsesSessionId, useApiSessionManagement, - selectedMCPTools, + selectedMCPServers, + mcpServerToolRestrictions, selectedVoice, ]); @@ -355,7 +380,7 @@ const ChatUI: React.FC = ({ }; loadModels(); - loadMCPTools(); + loadMCPServers(); }, [accessToken, userID, userRole, apiKeySource, apiKey, token]); // Fetch agents when A2A endpoint is selected @@ -869,12 +894,14 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools, + selectedMCPServers, updateChatImageUI, updateSearchResults, useAdvancedParams ? temperature : undefined, useAdvancedParams ? maxTokens : undefined, updateTotalLatency, + mcpServers, + mcpServerToolRestrictions, ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -940,12 +967,14 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools, // Pass the selected tools array + selectedMCPServers, // Pass the selected servers array useApiSessionManagement ? responsesSessionId : null, // Only pass session ID if API mode is enabled handleResponseId, // Pass callback to capture new response ID handleMCPEvent, // Pass MCP event handler codeInterpreter.enabled, // Enable Code Interpreter tool codeInterpreter.setResult, // Handle code interpreter output + mcpServers, + mcpServerToolRestrictions, ); } else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) { const apiChatHistory = [ @@ -968,7 +997,9 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools, // Pass the selected tools array + selectedMCPServers, // Pass the selected servers array + mcpServers, + mcpServerToolRestrictions, ); } else if (endpointType === EndpointType.EMBEDDINGS) { await makeOpenAIEmbeddingsRequest( @@ -1329,13 +1360,13 @@ const ChatUI: React.FC = ({ /> - {/* MCP Tool Selection */} + {/* MCP Server Selection */}
- MCP Tool + MCP Servers @@ -1343,30 +1374,103 @@ const ChatUI: React.FC = ({ + + {/* Tool restrictions UI (optional) */} + {selectedMCPServers.length > 0 && + !selectedMCPServers.includes("__all__") && + MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType) && ( +
+ {selectedMCPServers.map(serverId => { + const server = mcpServers.find(s => s.server_id === serverId); + const tools = serverToolsMap[serverId] || []; + if (tools.length === 0) return null; + + return ( +
+ + Limit tools for {server?.alias || server?.server_name || serverId}: + + setSelectedMCPTools(value)} - optionLabelProp="label" - allowClear - maxTagCount="responsive" - > - {mcpTools.map((tool) => ( - {tool.name}
} - > -
- {tool.name} - {tool.description} -
- - ))} - -
- )} - - )}
); }; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx index 60ce6879bea..84f7f6b9e1c 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx @@ -1,5 +1,6 @@ import { MessageType } from "./types"; import { EndpointType } from "./mode_endpoint_mapping"; +import { MCPServer } from "../../mcp_tools/types"; interface CodeGenMetadata { tags?: string[]; @@ -16,7 +17,9 @@ interface GenerateCodeParams { selectedTags: string[]; selectedVectorStores: string[]; selectedGuardrails: string[]; - selectedMCPTools: string[]; + selectedMCPServers: string[]; + mcpServers?: MCPServer[]; + mcpServerToolRestrictions?: Record; selectedVoice?: string; endpointType: string; selectedModel: string | undefined; @@ -37,7 +40,9 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => { selectedTags, selectedVectorStores, selectedGuardrails, - selectedMCPTools, + selectedMCPServers, + mcpServers, + mcpServerToolRestrictions, selectedVoice, endpointType, selectedModel, diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx index 3f8c90424c6..3aef2f2d5ef 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx @@ -17,7 +17,6 @@ export async function makeAnthropicMessagesRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTools?: string[], ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -47,23 +46,6 @@ export async function makeAnthropicMessagesRequest( const startTime = Date.now(); let firstTokenReceived = false; - // Format MCP tools if selected - const tools = - selectedMCPTools && selectedMCPTools.length > 0 - ? [ - { - type: "mcp", - server_label: "litellm", - server_url: `${proxyBaseUrl}/mcp`, - require_approval: "never", - allowed_tools: selectedMCPTools, - headers: { - "x-litellm-api-key": `Bearer ${accessToken}`, - }, - }, - ] - : undefined; - const requestBody: any = { model: selectedModel, messages: messages.map((m) => ({ role: m.role, content: m.content })), @@ -75,11 +57,6 @@ export async function makeAnthropicMessagesRequest( if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids; if (guardrails) requestBody.guardrails = guardrails; - if (tools) { - requestBody.tools = tools; - requestBody.tool_choice = "auto"; - } - // Use the streaming helper method for cleaner async iteration // @ts-ignore - The SDK types might not include all litellm-specific parameters const stream = client.messages.stream(requestBody, { signal }); diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index 8bdaa94ec65..906c72245fd 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -3,6 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; +import { MCPServer } from "../../mcp_tools/types"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], @@ -17,12 +18,14 @@ export async function makeOpenAIChatCompletionRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTools?: string[], + selectedMCPServers?: string[], onImageGenerated?: (imageUrl: string, model?: string) => void, onSearchResults?: (searchResults: VectorStoreSearchResponse[]) => void, temperature?: number, max_tokens?: number, onTotalLatency?: (latency: number) => void, + mcpServers?: MCPServer[], + mcpServerToolRestrictions?: Record, ) { // base url should be the current base_url const isLocal = process.env.NODE_ENV === "development"; @@ -53,22 +56,36 @@ export async function makeOpenAIChatCompletionRequest( let fullResponseContent = ""; let fullReasoningContent = ""; - // Format MCP tools if selected - const tools = - selectedMCPTools && selectedMCPTools.length > 0 - ? [ - { - type: "mcp", - server_label: "litellm", - server_url: 'litellm_proxy/mcp', - require_approval: "never", - allowed_tools: selectedMCPTools, - headers: { - "x-litellm-api-key": `Bearer ${accessToken}`, - }, - }, - ] - : undefined; + // Build tools array + const tools: any[] = []; + + // Add MCP servers if selected + if (selectedMCPServers && selectedMCPServers.length > 0) { + if (selectedMCPServers.includes("__all__")) { + // All MCP Servers selected + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp", + require_approval: "never", + }); + } else { + // Individual servers selected - create one entry per server + selectedMCPServers.forEach((serverId) => { + const server = mcpServers?.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: `litellm_proxy/mcp/${serverName}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + }); + } + } // @ts-ignore const response = await client.chat.completions.create( @@ -82,7 +99,7 @@ export async function makeOpenAIChatCompletionRequest( messages: chatHistory as ChatCompletionMessageParam[], ...(vector_store_ids ? { vector_store_ids } : {}), ...(guardrails ? { guardrails } : {}), - ...(tools ? { tools, tool_choice: "auto" } : {}), + ...(tools.length > 0 ? { tools, tool_choice: "auto" } : {}), ...(temperature !== undefined ? { temperature } : {}), ...(max_tokens !== undefined ? { max_tokens } : {}), }, diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx deleted file mode 100644 index d89ecbfb481..00000000000 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx +++ /dev/null @@ -1,29 +0,0 @@ -import { mcpToolsCall } from "../../networking"; - -export interface MCPTool { - name: string; - description: string; - title: string | null; - inputSchema: { - type: string; - properties: Record; - required: string[]; - }; - outputSchema: any; - annotations: any; - _meta: any; -} - -interface MCPToolsResponse { - tools: MCPTool[]; -} - -export async function fetchAvailableMCPTools(accessToken: string): Promise { - try { - const data = (await mcpToolsCall(accessToken)) as MCPToolsResponse; - return data.tools || []; - } catch (error) { - console.error("Error fetching MCP tools:", error); - return []; - } -} diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx index 8488f012e72..e30c14111fd 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx @@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; import { MCPEvent } from "../chat_ui/MCPEventsDisplay"; +import { MCPServer } from "../../mcp_tools/types"; import { CodeInterpreterResult, CodeInterpreterState, @@ -26,12 +27,14 @@ export async function makeOpenAIResponsesRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTools?: string[], + selectedMCPServers?: string[], previousResponseId?: string | null, onResponseId?: (responseId: string) => void, onMCPEvent?: (event: MCPEvent) => void, codeInterpreterEnabled?: boolean, onCodeInterpreterResult?: (result: CodeInterpreterResult) => void, + mcpServers?: MCPServer[], + mcpServerToolRestrictions?: Record, ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -86,15 +89,32 @@ export async function makeOpenAIResponsesRequest( // Build tools array const tools: any[] = []; - // Add MCP tools if selected - if (selectedMCPTools && selectedMCPTools.length > 0) { - tools.push({ - type: "mcp", - server_label: "litellm", - server_url: `litellm_proxy/mcp`, - require_approval: "never", - allowed_tools: selectedMCPTools, - }); + // Add MCP servers if selected + if (selectedMCPServers && selectedMCPServers.length > 0) { + if (selectedMCPServers.includes("__all__")) { + // All MCP Servers selected + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp", + require_approval: "never", + }); + } else { + // Individual servers selected - create one entry per server + selectedMCPServers.forEach((serverId) => { + const server = mcpServers?.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: `litellm_proxy/mcp/${serverName}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + }); + } } // Add code_interpreter tool if enabled (OpenAI auto-creates container)