From 0de013a06e84916e45966c81e81c3f9959a3cf1e Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 21 Mar 2026 19:01:20 -0700 Subject: [PATCH] feat(mcp): pass mcpToolsets to playground chat and responses API calls --- .../components/playground/chat_ui/ChatUI.tsx | 270 ++++++++++++++---- 1 file changed, 219 insertions(+), 51 deletions(-) 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 c0eb89a5980..9ace6ed8680 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -34,7 +34,8 @@ import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/M import { MCPServer } from "../../mcp_tools/types"; import { ByokCredentialModal } from "../../mcp_tools/ByokCredentialModal"; import NotificationsManager from "../../molecules/notifications_manager"; -import { callMCPTool, fetchMCPServers, listMCPTools } from "../../networking"; +import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "../../networking"; +import { MCPToolset } from "../../mcp_tools/types"; import TagSelector from "../../tag_management/TagSelector"; import VectorStoreSelector from "../../vector_store_management/VectorStoreSelector"; import { makeA2ASendMessageRequest } from "../llm_calls/a2a_send_message"; @@ -111,6 +112,8 @@ const ChatUI: React.FC = ({ fixedModel, }) => { const [mcpServers, setMCPServers] = useState([]); + const [mcpToolsets, setMCPToolsets] = useState([]); + const [isToolsetsInfoModalVisible, setIsToolsetsInfoModalVisible] = useState(false); const [byokModalServer, setByokModalServer] = useState(null); const [selectedMCPServers, setSelectedMCPServers] = useState(() => { const saved = sessionStorage.getItem("selectedMCPServers"); @@ -255,15 +258,19 @@ const ChatUI: React.FC = ({ const chatEndRef = useRef(null); - // Fetch MCP servers + // Fetch MCP servers and toolsets const loadMCPServers = async () => { const userApiKey = apiKeySource === "session" ? accessToken : apiKey; if (!userApiKey) return; setIsLoadingMCPServers(true); try { - const servers = await fetchMCPServers(userApiKey); + const [servers, toolsets] = await Promise.all([ + fetchMCPServers(userApiKey), + fetchMCPToolsets(userApiKey).catch(() => []), + ]); setMCPServers(Array.isArray(servers) ? servers : servers.data || []); + setMCPToolsets(Array.isArray(toolsets) ? toolsets : []); } catch (error) { console.error("Error fetching MCP servers:", error); } finally { @@ -412,17 +419,29 @@ const ChatUI: React.FC = ({ loadMCPServers(); }, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]); - // Load tools when MCP direct mode has a server selected + // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { if ( endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && - selectedMCPServers[0] !== "__all__" && - !serverToolsMap[selectedMCPServers[0]] + selectedMCPServers[0] !== "__all__" ) { - loadServerTools(selectedMCPServers[0]); + const selected = selectedMCPServers[0]; + if (selected.startsWith("toolset:")) { + // For a toolset, load tools for each server in it + const toolsetId = selected.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + if (toolset) { + const uniqueServerIds = [...new Set(toolset.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + if (!serverToolsMap[sid]) loadServerTools(sid); + }); + } + } else if (!serverToolsMap[selected]) { + loadServerTools(selected); + } } - }, [endpointType, selectedMCPServers, serverToolsMap]); + }, [endpointType, selectedMCPServers, serverToolsMap, mcpToolsets]); // Fetch agents when A2A endpoint is selected useEffect(() => { @@ -568,19 +587,34 @@ const ChatUI: React.FC = ({ // For MCP direct mode, require server and tool selection, and get form values early let mcpToolArguments: Record = {}; if (endpointType === EndpointType.MCP) { - const mcpServerId = + const rawSelected = selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" ? selectedMCPServers[0] : null; - if (!mcpServerId) { + if (!rawSelected) { NotificationsManager.fromBackend("Please select an MCP server to test"); return; } + // Resolve the real server ID (toolsets use toolset: prefix) + const mcpServerId = rawSelected.startsWith("toolset:") ? rawSelected : rawSelected; if (!selectedMCPDirectTool) { NotificationsManager.fromBackend("Please select an MCP tool to call"); return; } - const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find( + // For toolsets, find the tool in the servers that back this toolset + const toolsetForSelected = rawSelected.startsWith("toolset:") + ? mcpToolsets.find((t) => t.toolset_id === rawSelected.slice("toolset:".length)) + : null; + let searchPool: any[] = []; + if (toolsetForSelected) { + const uniqueServerIds = [...new Set(toolsetForSelected.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + searchPool = searchPool.concat(serverToolsMap[sid] || []); + }); + } else { + searchPool = serverToolsMap[rawSelected] || []; + } + const mcpTool = searchPool.find( (t: any) => t.name === selectedMCPDirectTool, ); if (!mcpTool) { @@ -738,6 +772,7 @@ const ChatUI: React.FC = ({ mcpServerToolRestrictions, handleMCPEvent, mockTestFallbacks, + mcpToolsets, ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -818,6 +853,7 @@ const ChatUI: React.FC = ({ customProxyBaseUrl || undefined, mcpServers, mcpServerToolRestrictions, + mcpToolsets, ); } else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) { const apiChatHistory = [ @@ -875,14 +911,22 @@ const ChatUI: React.FC = ({ // Handle MCP direct tool calls (no chat completions) if (endpointType === EndpointType.MCP) { - const mcpServerId = + const rawSelected = selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" ? selectedMCPServers[0] : null; - if (mcpServerId && selectedMCPDirectTool) { + // For toolsets, resolve the real server_id from the toolset's tool list + let resolvedServerId = rawSelected; + if (rawSelected?.startsWith("toolset:")) { + const toolsetId = rawSelected.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + const toolEntry = toolset?.tools.find((t) => t.tool_name === selectedMCPDirectTool); + resolvedServerId = toolEntry?.server_id ?? rawSelected; + } + if (resolvedServerId && !resolvedServerId.startsWith("toolset:") && selectedMCPDirectTool) { const result = await callMCPTool( effectiveApiKey, - mcpServerId, + resolvedServerId, selectedMCPDirectTool, mcpToolArguments, selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined, @@ -1297,11 +1341,14 @@ const ChatUI: React.FC = ({ className="ml-1" title={ endpointType === EndpointType.MCP - ? "Select an MCP server to test tools directly." - : "Select MCP servers to use in your conversation." + ? "Select an MCP server or toolset to test tools directly." + : "Select MCP servers or toolsets to use in your conversation." } > - + setIsToolsetsInfoModalVisible(true)} + /> {/* MCP Tool selector - only for MCP direct mode */} {endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && - selectedMCPServers[0] !== "__all__" && ( -
- Select Tool - setSelectedMCPDirectTool(value)} + options={toolOptions} + allowClear + className="rounded-md" + /> +
+ ); + })()} {/* Tool restrictions UI (optional) - hidden for MCP direct mode */} {selectedMCPServers.length > 0 && @@ -1920,7 +2033,21 @@ const ChatUI: React.FC = ({ selectedMCPDirectTool ? (
{(() => { - const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find( + const rawSel = selectedMCPServers[0]; + let toolPool: any[] = []; + if (rawSel.startsWith("toolset:")) { + const toolsetId = rawSel.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + if (toolset) { + const uniqueServerIds = [...new Set(toolset.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + toolPool = toolPool.concat(serverToolsMap[sid] || []); + }); + } + } else { + toolPool = serverToolsMap[rawSel] || []; + } + const mcpTool = toolPool.find( (t: any) => t.name === selectedMCPDirectTool, ); return mcpTool ? ( @@ -2066,6 +2193,47 @@ const ChatUI: React.FC = ({ accessToken={accessToken || ""} /> )} + + {/* Toolsets info modal */} + setIsToolsetsInfoModalVisible(false)} + footer={[ + , + ]} + width={600} + > +
+

+ Toolsets are named collections of specific tools from one or more MCP servers. + Instead of exposing all tools from a server, a toolset gives an agent exactly the tools it needs. +

+
+

How to use a toolset:

+
    +
  1. Select a Toolset (purple badge) from the MCP Servers dropdown.
  2. +
  3. The tool picker will show only the tools included in that toolset.
  4. +
  5. Select a tool and fill in its parameters, then send.
  6. +
  7. The tool call is routed to the correct underlying MCP server automatically.
  8. +
+
+
+

+ Example: A "GitHub Read-only" toolset might include only list_repos and get_file from a GitHub MCP server — preventing agents from making writes. +

+
+
+

Creating toolsets:

+

+ Admins can create and manage toolsets from the MCP page → Toolsets tab. + Toolsets can then be assigned to keys and teams to scope their tool access. +

+
+
+
); };