From d36c8d6bbebb792623bc57347716602e71497809 Mon Sep 17 00:00:00 2001 From: rishiganesh2002 <98856261+rishiganesh2002@users.noreply.github.com> Date: Fri, 3 Oct 2025 10:16:29 -0700 Subject: [PATCH] [Feat] MCP Gateway Fine-grained Tools Addition (#15153) * feat: UI to add specific tools under creating MCP connection * chore: pydantic + prisma changes * feat: adding specific MCP tools now works * fix: allowed tools filtering * chore: filtered list to mcp server cost config * chore: update Readme * chore: refactor the filtering * test: Added tests When the allowed_tests is null, empty list or populated * chore: resolve the proxy issue * feat: updating MCP tool filtering --- .../mcp_server/mcp_server_manager.py | 5 + .../mcp_server/rest_endpoints.py | 7 + .../proxy/_experimental/mcp_server/server.py | 57 ++++- litellm/proxy/_types.py | 2 + .../mcp_management_endpoints.py | 1 - litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + tests/README.MD | 6 +- .../mcp_server/test_mcp_server_manager.py | 152 +++++++++++- .../mcp_tools/create_mcp_server.tsx | 52 +++- .../mcp_tools/mcp_connection_status.tsx | 167 +++---------- .../components/mcp_tools/mcp_server_view.tsx | 21 +- .../src/components/mcp_tools/mcp_servers.tsx | 21 +- .../mcp_tools/mcp_tool_configuration.tsx | 177 ++++++++++++++ .../src/components/mcp_tools/types.tsx | 229 +++++++++--------- .../src/hooks/useTestMCPConnection.tsx | 114 +++++++++ 16 files changed, 735 insertions(+), 278 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx create mode 100644 ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a88f94be06f..9172568f304 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1197,6 +1197,11 @@ class MCPServerManager: if server.mcp_access_groups is not None else [] ), + allowed_tools=( + server.allowed_tools + if server.allowed_tools is not None + else [] + ), mcp_info=server.mcp_info, teams=cast( List[Dict[str, str | None]], diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index ecb960ebcf3..6a9c425a81b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -28,6 +28,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.server import ( ListMCPToolsRestAPIResponseObject, call_mcp_tool, + filter_tools_by_allowed_tools, ) ######################################################## @@ -75,6 +76,12 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, add_prefix=False, ) + + # Filter tools based on allowed_tools configuration + # Only filter if allowed_tools is explicitly configured (not None and not empty) + if server.allowed_tools is not None and len(server.allowed_tools) > 0: + tools = filter_tools_by_allowed_tools(tools, server) + return _create_tool_response_objects(tools, server.mcp_info) ######################################################## diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 96e4d47a914..3e7c291810a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -361,22 +361,66 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool: + """ + Check if a tool name matches any name in the filter list. + + Checks both the full tool name and unprefixed version (without server prefix). + This allows users to configure simple tool names regardless of prefixing. + + Args: + tool_name: The tool name to check (may be prefixed like "server-tool_name") + filter_list: List of tool names to match against + + Returns: + True if the tool name (prefixed or unprefixed) is in the filter list + """ + from litellm.proxy._experimental.mcp_server.utils import ( + get_server_name_prefix_tool_mcp, + ) + + # Check if the full name is in the list + if tool_name in filter_list: + return True + + # Check if the unprefixed name is in the list + unprefixed_name, _ = get_server_name_prefix_tool_mcp(tool_name) + return unprefixed_name in filter_list + def filter_tools_by_allowed_tools( tools: List[MCPTool], mcp_server: MCPServer, ) -> List[MCPTool]: """ - Filter tools by allowed tools + Filter tools by allowed/disallowed tools configuration. + + If allowed_tools is set, only tools in that list are returned. + If disallowed_tools is set, tools in that list are excluded. + Tool names are matched with and without server prefixes for flexibility. + + Args: + tools: List of tools to filter + mcp_server: Server configuration with allowed_tools/disallowed_tools + + Returns: + Filtered list of tools """ tools_to_return = tools + + # Filter by allowed_tools (whitelist) if mcp_server.allowed_tools: tools_to_return = [ - tool for tool in tools if tool.name in mcp_server.allowed_tools + tool for tool in tools + if _tool_name_matches(tool.name, mcp_server.allowed_tools) ] + + # Filter by disallowed_tools (blacklist) if mcp_server.disallowed_tools: tools_to_return = [ - tool for tool in tools if tool.name not in mcp_server.disallowed_tools + tool for tool in tools_to_return + if not _tool_name_matches(tool.name, mcp_server.disallowed_tools) ] + return tools_to_return async def _get_tools_from_mcp_servers( @@ -453,9 +497,12 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=add_prefix, ) - all_tools.extend(filter_tools_by_allowed_tools(tools, server)) + + filtered_tools = filter_tools_by_allowed_tools(tools, server) + all_tools.extend(filtered_tools) + verbose_logger.debug( - f"Successfully fetched {len(tools)} tools from server {server.name}" + f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering" ) except Exception as e: verbose_logger.exception( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c5370eb7d70..efe7ff90973 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -917,6 +917,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): url: Optional[str] = None mcp_info: Optional[MCPInfo] = None mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: Optional[List[str]] = None # Stdio-specific fields command: Optional[str] = None args: List[str] = Field(default_factory=list) @@ -985,6 +986,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): updated_by: Optional[str] = None teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: List[str] = Field(default_factory=list) mcp_info: Optional[MCPInfo] = None # Health check status status: Optional[str] = Field( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 1912c54920c..b0c09b1e2e8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -5,7 +5,6 @@ Endpoints here: - GET `/v1/mcp/server` - Returns all of the configured mcp servers in the db filtered by requestor's access - GET `/v1/mcp/server/{server_id}` - Returns the the specific mcp server in the db given `server_id` filtered by requestor's access -- GET `/v1/mcp/server/{server_id}/tools` - Get all the tools from the mcp server specified by the `server_id` - POST `/v1/mcp/server` - Add a new external mcp server. - PUT `/v1/mcp/server` - Edits an existing mcp server. - DELETE `/v1/mcp/server/{server_id}` - Deletes the mcp server given `server_id`. diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 766625145f6..5a79e171438 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -178,6 +178,7 @@ model LiteLLM_MCPServerTable { updated_by String? mcp_info Json? @default("{}") mcp_access_groups String[] + allowed_tools String[] @default([]) // Health check status status String? @default("unknown") last_health_check DateTime? diff --git a/schema.prisma b/schema.prisma index 766625145f6..5a79e171438 100644 --- a/schema.prisma +++ b/schema.prisma @@ -178,6 +178,7 @@ model LiteLLM_MCPServerTable { updated_by String? mcp_info Json? @default("{}") mcp_access_groups String[] + allowed_tools String[] @default([]) // Health check status status String? @default("unknown") last_health_check DateTime? diff --git a/tests/README.MD b/tests/README.MD index ed9ac10e9dc..57275a031f7 100644 --- a/tests/README.MD +++ b/tests/README.MD @@ -1,9 +1,9 @@ -**In total litellm runs 1000+ tests** +**In total litellm runs 1000+ tests** [02/20/2025] Update: To make it easier to contribute and map what behavior is tested, -we've started mapping the litellm directory in `tests/litellm` +we've started mapping the litellm directory in `tests/test_litellm` -This folder can only run mock tests. \ No newline at end of file +This folder can only run mock tests. 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 80ea95a2210..146f9434b8c 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 @@ -1,6 +1,6 @@ import sys from datetime import datetime -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException @@ -759,6 +759,156 @@ class TestMCPServerManager: assert resolved_server_pref is not None assert resolved_server_pref.server_id == server.server_id + @pytest.mark.asyncio + async def test_rest_endpoint_filters_by_allowed_tools(self): + """Test that REST endpoint _get_tools_for_single_server respects allowed_tools configuration""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + + # Create server with allowed_tools configured + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + allowed_tools=["allowed_tool_1", "allowed_tool_2"], + ) + server.mcp_info = {"server_name": "test-server"} + + # Mock tools returned from manager (3 tools, but only 2 are allowed) + tool1 = MagicMock() + tool1.name = "allowed_tool_1" + tool1.description = "This tool is allowed" + tool1.inputSchema = {} + + tool2 = MagicMock() + tool2.name = "blocked_tool" + tool2.description = "This tool is not allowed" + tool2.inputSchema = {} + + tool3 = MagicMock() + tool3.name = "allowed_tool_2" + tool3.description = "This tool is also allowed" + tool3.inputSchema = {} + + # Mock the global_mcp_server_manager._get_tools_from_server + from litellm.proxy._experimental.mcp_server import rest_endpoints + + with patch.object( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + new=AsyncMock(return_value=[tool1, tool2, tool3]), + ): + # Call the REST endpoint helper + filtered_response = await _get_tools_for_single_server( + server, server_auth_header=None + ) + + # Verify only allowed tools are in the response + assert len(filtered_response) == 2 + tool_names = [t.name for t in filtered_response] + assert "allowed_tool_1" in tool_names + assert "allowed_tool_2" in tool_names + assert "blocked_tool" not in tool_names + + @pytest.mark.asyncio + async def test_rest_endpoint_shows_all_when_allowed_tools_is_none(self): + """Test that REST endpoint shows all tools when allowed_tools is None (backwards compatibility)""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + + # Create server with allowed_tools as None + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + allowed_tools=None, # No filtering + ) + server.mcp_info = {"server_name": "test-server"} + + # Mock tools returned from manager + tool1 = MagicMock() + tool1.name = "tool_1" + tool1.description = "Tool 1" + tool1.inputSchema = {} + + tool2 = MagicMock() + tool2.name = "tool_2" + tool2.description = "Tool 2" + tool2.inputSchema = {} + + tool3 = MagicMock() + tool3.name = "tool_3" + tool3.description = "Tool 3" + tool3.inputSchema = {} + + # Mock the global_mcp_server_manager._get_tools_from_server + from litellm.proxy._experimental.mcp_server import rest_endpoints + + with patch.object( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + new=AsyncMock(return_value=[tool1, tool2, tool3]), + ): + # Call the REST endpoint helper + all_tools_response = await _get_tools_for_single_server( + server, server_auth_header=None + ) + + # Verify all tools are returned (no filtering) + assert len(all_tools_response) == 3 + tool_names = [t.name for t in all_tools_response] + assert "tool_1" in tool_names + assert "tool_2" in tool_names + assert "tool_3" in tool_names + + @pytest.mark.asyncio + async def test_rest_endpoint_shows_all_when_allowed_tools_is_empty_list(self): + """Test that REST endpoint shows all tools when allowed_tools is empty list (backwards compatibility)""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + + # Create server with allowed_tools as empty list + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + allowed_tools=[], # Empty list means no filtering + ) + server.mcp_info = {"server_name": "test-server"} + + # Mock tools returned from manager + tool1 = MagicMock() + tool1.name = "tool_1" + tool1.description = "Tool 1" + tool1.inputSchema = {} + + tool2 = MagicMock() + tool2.name = "tool_2" + tool2.description = "Tool 2" + tool2.inputSchema = {} + + # Mock the global_mcp_server_manager._get_tools_from_server + from litellm.proxy._experimental.mcp_server import rest_endpoints + + with patch.object( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + new=AsyncMock(return_value=[tool1, tool2]), + ): + # Call the REST endpoint helper + all_tools_response = await _get_tools_for_single_server( + server, server_auth_header=None + ) + + # Verify all tools are returned (no filtering) + assert len(all_tools_response) == 2 + tool_names = [t.name for t in all_tools_response] + assert "tool_1" in tool_names + assert "tool_2" in tool_names + if __name__ == "__main__": pytest.main([__file__]) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 97132daa466..d2ba0fbfcaf 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -6,6 +6,7 @@ import { createMCPServer } from "../networking" import { MCPServer, MCPServerCostInfo } from "./types" import MCPServerCostConfig from "./mcp_server_cost_config" import MCPConnectionStatus from "./mcp_connection_status" +import MCPToolConfiguration from "./mcp_tool_configuration" import StdioConfiguration from "./StdioConfiguration" import { isAdminRole } from "@/utils/roles" import { validateMCPServerUrl, validateMCPServerName } from "./utils" @@ -37,7 +38,8 @@ const CreateMCPServer: React.FC = ({ const [formValues, setFormValues] = useState>({}) const [aliasManuallyEdited, setAliasManuallyEdited] = useState(false) const [tools, setTools] = useState([]) - const [transportType, setTransportType] = useState("sse") + const [allowedTools, setAllowedTools] = useState([]) + const [transportType, setTransportType] = useState("") const [searchValue, setSearchValue] = useState("") const [urlWarning, setUrlWarning] = useState("") @@ -103,7 +105,7 @@ const CreateMCPServer: React.FC = ({ } } - // Prepare the payload with cost configuration + // Prepare the payload with cost configuration and allowed tools const payload = { ...formValues, ...stdioFields, @@ -116,6 +118,7 @@ const CreateMCPServer: React.FC = ({ }, mcp_access_groups: accessGroups, alias: formValues.alias, + allowed_tools: allowedTools.length > 0 ? allowedTools : null, } console.log(`Payload: ${JSON.stringify(payload)}`) @@ -127,7 +130,9 @@ const CreateMCPServer: React.FC = ({ form.resetFields() setCostConfig({}) setTools([]) + setAllowedTools([]) setUrlWarning("") + setAliasManuallyEdited(false) setModalVisible(false) onCreateSuccess(response) } @@ -143,7 +148,9 @@ const CreateMCPServer: React.FC = ({ form.resetFields() setCostConfig({}) setTools([]) + setAllowedTools([]) setUrlWarning("") + setAliasManuallyEdited(false) setModalVisible(false) } @@ -176,7 +183,10 @@ const CreateMCPServer: React.FC = ({ })) // If search value doesn't match any existing group and is not empty, add "create new group" option - if (searchValue && !availableAccessGroups.some(group => group.toLowerCase().includes(searchValue.toLowerCase()))) { + if ( + searchValue && + !availableAccessGroups.some((group) => group.toLowerCase().includes(searchValue.toLowerCase())) + ) { existingOptions.push({ value: searchValue, label: ( @@ -201,6 +211,13 @@ const CreateMCPServer: React.FC = ({ } }, [formValues.server_name]) + // Clear formValues when modal closes to reset child components + React.useEffect(() => { + if (!isModalVisible) { + setFormValues({}) + } + }, [isModalVisible]) + // rendering if (!isAdminRole(userRole)) { return null @@ -297,7 +314,7 @@ const CreateMCPServer: React.FC = ({ rules={[ { required: false, - message: "Please enter a server description", + message: "Please enter a server description!!!!!!!!!", }, ]} > @@ -341,11 +358,7 @@ const CreateMCPServer: React.FC = ({ className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" onChange={(e) => checkUrlFormat(e.target.value, transportType)} /> - {urlWarning && ( -
- {urlWarning} -
- )} + {urlWarning &&
{urlWarning}
} )} @@ -386,9 +399,7 @@ const CreateMCPServer: React.FC = ({ showSearch placeholder="Select existing groups or type to create new ones" optionFilterProp="value" - filterOption={(input, option) => - (option?.value ?? '').toLowerCase().includes(input.toLowerCase()) - } + filterOption={(input, option) => (option?.value ?? "").toLowerCase().includes(input.toLowerCase())} onSearch={(value) => setSearchValue(value)} tokenSeparators={[","]} options={getAccessGroupOptions()} @@ -403,9 +414,24 @@ const CreateMCPServer: React.FC = ({ + {/* Tool Configuration Section */} +
+ +
+ {/* Cost Configuration Section */}
- + allowedTools.includes(tool.name))} + disabled={false} + />
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx index adeae707af9..acba932c70c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx @@ -1,92 +1,30 @@ -import React, { useState, useEffect } from "react"; -import { Button, message, Spin, Alert, Collapse, Badge } from "antd"; -import { CheckCircleOutlined, ExclamationCircleOutlined, ReloadOutlined, ToolOutlined, InfoCircleOutlined } from "@ant-design/icons"; -import { Card, Title, Text } from "@tremor/react"; -import { testMCPToolsListRequest } from "../networking"; - -const { Panel } = Collapse; +import React, { useEffect } from "react" +import { Button, Spin, Alert } from "antd" +import { CheckCircleOutlined, ExclamationCircleOutlined, ReloadOutlined, ToolOutlined } from "@ant-design/icons" +import { Card, Title, Text } from "@tremor/react" +import { useTestMCPConnection } from "../../hooks/useTestMCPConnection" interface MCPConnectionStatusProps { - accessToken: string | null; - formValues: Record; - onToolsLoaded?: (tools: any[]) => void; + accessToken: string | null + formValues: Record + onToolsLoaded?: (tools: any[]) => void } -const MCPConnectionStatus: React.FC = ({ - accessToken, - formValues, - onToolsLoaded -}) => { - const [tools, setTools] = useState([]); - const [isLoadingTools, setIsLoadingTools] = useState(false); - const [toolsError, setToolsError] = useState(null); - const [hasShownSuccessMessage, setHasShownSuccessMessage] = useState(false); +const MCPConnectionStatus: React.FC = ({ accessToken, formValues, onToolsLoaded }) => { + const { tools, isLoadingTools, toolsError, canFetchTools, fetchTools } = useTestMCPConnection({ + accessToken, + formValues, + enabled: true, // Auto-fetch when required fields are available + }) - // Check if we have the minimum required fields to fetch tools - const canFetchTools = formValues.url && formValues.transport && formValues.auth_type && accessToken; - - const fetchTools = async () => { - if (!accessToken || !formValues.url) { - return; - } - - setIsLoadingTools(true); - setToolsError(null); - - try { - // Prepare the MCP server config from form values - const mcpServerConfig = { - server_id: formValues.server_id || "", - server_name: formValues.server_name || "", - url: formValues.url, - transport: formValues.transport, - auth_type: formValues.auth_type, - mcp_info: formValues.mcp_info, - }; - - const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig); - - if (toolsResponse.tools && !toolsResponse.error) { - setTools(toolsResponse.tools); - setToolsError(null); - onToolsLoaded?.(toolsResponse.tools); - if (toolsResponse.tools.length > 0 && !hasShownSuccessMessage) { - setHasShownSuccessMessage(true); - } - } else { - const errorMessage = toolsResponse.message || "Failed to retrieve tools list"; - setToolsError(errorMessage); - setTools([]); - onToolsLoaded?.([]); - setHasShownSuccessMessage(false); - } - } catch (error) { - console.error("Tools fetch error:", error); - setToolsError(error instanceof Error ? error.message : String(error)); - setTools([]); - onToolsLoaded?.([]); - setHasShownSuccessMessage(false); - } finally { - setIsLoadingTools(false); - } - }; - - // Auto-fetch tools when form values change and required fields are available + // Notify parent component when tools change useEffect(() => { - if (canFetchTools) { - fetchTools(); - } else { - // Clear tools if required fields are missing - setTools([]); - setToolsError(null); - setHasShownSuccessMessage(false); - onToolsLoaded?.([]); - } - }, [formValues.url, formValues.transport, formValues.auth_type, accessToken]); + onToolsLoaded?.(tools) + }, [tools, onToolsLoaded]) // Don't show anything if required fields aren't filled if (!canFetchTools && !formValues.url) { - return null; + return null } return ( @@ -102,9 +40,7 @@ const MCPConnectionStatus: React.FC = ({ Complete required fields to test connection
- - Fill in URL, Transport, and Authentication to test MCP server connection - + Fill in URL, Transport, and Authentication to test MCP server connection
)} @@ -113,34 +49,32 @@ const MCPConnectionStatus: React.FC = ({
- {isLoadingTools - ? "Testing connection to MCP server..." - : tools.length > 0 + {isLoadingTools + ? "Testing connection to MCP server..." + : tools.length > 0 ? "Connection successful" : toolsError ? "Connection failed" : "Ready to test connection"}
- - Server: {formValues.url} - + Server: {formValues.url}
- + {isLoadingTools && (
Connecting...
)} - + {!isLoadingTools && !toolsError && tools.length > 0 && (
Connected
)} - + {toolsError && (
@@ -163,54 +97,13 @@ const MCPConnectionStatus: React.FC = ({ type="error" showIcon action={ - } /> )} - {!isLoadingTools && tools.length > 0 && ( - - - Available Tools - -
- ), - children: ( -
- {tools.map((tool, index) => ( -
- {tool.name} - {tool.description && ( - - {tool.description} - - )} -
- ))} -
- ), - }, - ]} - /> - )} - {!isLoadingTools && tools.length === 0 && !toolsError && (
@@ -223,7 +116,7 @@ const MCPConnectionStatus: React.FC = ({ )}
- ); -}; + ) +} -export default MCPConnectionStatus; \ No newline at end of file +export default MCPConnectionStatus diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx index 0fc9d834b50..02329936f77 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx @@ -108,7 +108,7 @@ export const MCPServerView: React.FC = ({ ? "text-green-600 bg-green-50 border-green-200" : "text-gray-500 hover:text-gray-700 hover:bg-gray-100" }`} - /> + />
@@ -241,6 +241,25 @@ export const MCPServerView: React.FC = ({ )} +
+ Allowed Tools +
+ {mcpServer.allowed_tools && mcpServer.allowed_tools.length > 0 ? ( +
+ {mcpServer.allowed_tools.map((tool: string, index: number) => ( + + {tool} + + ))} +
+ ) : ( + All tools enabled + )} +
+
Cost Configuration diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 95073ad1dbd..fd55dba604d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -50,6 +50,17 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) enabled: !!accessToken, }) as { data: MCPServer[]; isLoading: boolean; refetch: () => void; dataUpdatedAt: number } + // Log allowed_tools from fetched servers + React.useEffect(() => { + if (mcpServers) { + console.log("MCP Servers fetched:", mcpServers) + mcpServers.forEach((server) => { + console.log(`Server: ${server.server_name || server.server_id}`) + console.log(` allowed_tools:`, server.allowed_tools) + }) + } + }, [mcpServers]) + // state const [serverIdToDelete, setServerToDelete] = useState(null) const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false) @@ -60,7 +71,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const [filteredServers, setFilteredServers] = useState([]) const [isModalVisible, setModalVisible] = useState(false) - const isInternalUser = userRole === "Internal User"; + const isInternalUser = userRole === "Internal User" // Get unique teams from all servers const uniqueTeams = React.useMemo(() => { @@ -84,7 +95,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // Get unique MCP access groups from all servers const uniqueMcpAccessGroups = React.useMemo(() => { if (!mcpServers) return [] - return Array.from(new Set(mcpServers.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null))) + return Array.from( + new Set( + mcpServers.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null), + ), + ) }, [mcpServers]) // Handle team filter change @@ -171,7 +186,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } if (!accessToken || !userRole || !userID) { - console.log("Missing required authentication parameters", { accessToken, userRole, userID }); + console.log("Missing required authentication parameters", { accessToken, userRole, userID }) return
Missing required authentication parameters.
} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx new file mode 100644 index 00000000000..fbf8b04048b --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx @@ -0,0 +1,177 @@ +import React, { useEffect, useRef } from "react" +import { Card, Title, Text } from "@tremor/react" +import { ToolOutlined, CheckCircleOutlined } from "@ant-design/icons" +import { Badge, Spin, Checkbox } from "antd" +import { useTestMCPConnection } from "../../hooks/useTestMCPConnection" + +interface MCPToolConfigurationProps { + accessToken: string | null + formValues: Record + allowedTools: string[] + onAllowedToolsChange: (tools: string[]) => void +} + +const MCPToolConfiguration: React.FC = ({ + accessToken, + formValues, + allowedTools, + onAllowedToolsChange, +}) => { + const previousToolsLengthRef = useRef(0) + + const { tools, isLoadingTools, toolsError, canFetchTools } = useTestMCPConnection({ + accessToken, + formValues, + enabled: true, + }) + + // Auto-select all tools when tools are first loaded + useEffect(() => { + // Only auto-select if: + // 1. We have tools + // 2. Tools length changed (new tools loaded) + // 3. No tools are currently selected (initial state) + if (tools.length > 0 && tools.length !== previousToolsLengthRef.current && allowedTools.length === 0) { + const allToolNames = tools.map((tool) => tool.name) + onAllowedToolsChange(allToolNames) + } + // Update ref to track tools length (will be 0 when tools clear) + previousToolsLengthRef.current = tools.length + }, [tools, allowedTools.length, onAllowedToolsChange]) + + const handleToolToggle = (toolName: string) => { + if (allowedTools.includes(toolName)) { + onAllowedToolsChange(allowedTools.filter((name) => name !== toolName)) + } else { + onAllowedToolsChange([...allowedTools, toolName]) + } + } + + const handleSelectAll = () => { + const allToolNames = tools.map((tool) => tool.name) + onAllowedToolsChange(allToolNames) + } + + const handleDeselectAll = () => { + onAllowedToolsChange([]) + } + + // Don't show anything if required fields aren't filled + if (!canFetchTools && !formValues.url) { + return null + } + + return ( + +
+
+
+ + Tool Configuration + {tools.length > 0 && ( + + )} +
+
+ + {/* Loading state */} + {isLoadingTools && ( +
+ + Loading tools... +
+ )} + + {/* Error state */} + {toolsError && !isLoadingTools && ( +
+ + Unable to load tools +
+ {toolsError} +
+ )} + + {/* No tools state */} + {!isLoadingTools && !toolsError && tools.length === 0 && canFetchTools && ( +
+ + No tools available for configuration +
+ Connect to an MCP server with tools to configure them +
+ )} + + {/* Incomplete form state */} + {!canFetchTools && formValues.url && ( +
+ + Complete required fields to configure tools +
+ Fill in URL, Transport, and Authentication to load available tools +
+ )} + + {/* Tools loaded successfully */} + {!isLoadingTools && !toolsError && tools.length > 0 && ( +
+
+
+ + + {allowedTools.length} of {tools.length} {tools.length === 1 ? "tool" : "tools"} selected + +
+
+ + +
+
+ + {/* Tool list with checkboxes */} +
+ {tools.map((tool, index) => ( +
handleToolToggle(tool.name)} + > +
+ handleToolToggle(tool.name)} /> +
+ {tool.name} + {tool.description && {tool.description}} +
+
+
+ ))} +
+
+ )} +
+
+ ) +} + +export default MCPToolConfiguration diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 5213c2610eb..9bf5725ac0b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -1,7 +1,7 @@ export interface Team { - team_id: string; - team_alias?: string; - organization_id?: string | null; + team_id: string + team_alias?: string + organization_id?: string | null } // Default no auth value @@ -10,140 +10,141 @@ export const AUTH_TYPE = { API_KEY: "api_key", BEARER_TOKEN: "bearer_token", BASIC: "basic", -}; +} export const TRANSPORT = { SSE: "sse", HTTP: "http", -}; +} export const handleTransport = (transport?: string | null): string => { console.log(transport) if (transport === null || transport === undefined) { - return TRANSPORT.SSE; + return TRANSPORT.SSE } - return transport; -}; + return transport +} export const handleAuth = (authType?: string | null): string => { if (authType === null || authType === undefined) { - return AUTH_TYPE.NONE; + return AUTH_TYPE.NONE } - return authType; -}; + return authType +} export const mcpServerHasAuth = (authType?: string | null): boolean => { - return handleAuth(authType) !== AUTH_TYPE.NONE; -} + return handleAuth(authType) !== AUTH_TYPE.NONE +} // Define the structure for tool input schema properties export interface InputSchemaProperty { - type: string; - description?: string; - properties?: Record; // For nested object properties - required?: string[]; // For required fields in nested objects - enum?: string[]; // For enum values - default?: any; // For default values - } - - // Define the structure for the input schema of a tool - export interface InputSchema { - type: "object"; - properties: Record; - required?: string[]; - } - - // Define MCPServerCostInfo for cost tracking - export interface MCPServerCostInfo { - default_cost_per_query?: number | null; - tool_name_to_cost_per_query?: Record; - } + type: string + description?: string + properties?: Record // For nested object properties + required?: string[] // For required fields in nested objects + enum?: string[] // For enum values + default?: any // For default values +} - // Define MCP provider info - export interface MCPInfo { - server_name: string; - description?: string; - logo_url?: string; - mcp_server_cost_info?: MCPServerCostInfo | null; - } - - // Define the structure for a single MCP tool - export interface MCPTool { - name: string; - description?: string; - inputSchema: InputSchema | string; // API returns string "tool_input_schema" or the actual schema - mcp_info: MCPInfo; - // Function to select a tool (added in the component) - onToolSelect?: (tool: MCPTool) => void; - } - - // Define the response structure for the listMCPTools endpoint - now a flat array - export type ListMCPToolsResponse = MCPTool[]; - - // Define the argument structure for calling an MCP tool - export interface CallMCPToolArgs { - name: string; - arguments: Record | null; - server_name?: string; // Now using server_name from mcp_info - } - - // Define the possible content types in the response - export interface MCPTextContent { - type: "text"; - text: string; - annotations?: any; - } - - export interface MCPImageContent { - type: "image"; - url?: string; - data?: string; - } - - export interface MCPEmbeddedResource { - type: "embedded_resource"; - resource_type?: string; - url?: string; - data?: any; - } - - // Define the union type for the content array in the response - export type MCPContent = MCPTextContent | MCPImageContent | MCPEmbeddedResource; - - // Define the response structure for the callMCPTool endpoint - export type CallMCPToolResponse = MCPContent[]; - - // Props for the main component - export interface MCPToolsViewerProps { - serverId: string; - accessToken: string | null; - auth_type?: string | null; - userRole: string | null; - userID: string | null; - serverAlias?: string | null; - } +// Define the structure for the input schema of a tool +export interface InputSchema { + type: "object" + properties: Record + required?: string[] +} + +// Define MCPServerCostInfo for cost tracking +export interface MCPServerCostInfo { + default_cost_per_query?: number | null + tool_name_to_cost_per_query?: Record +} + +// Define MCP provider info +export interface MCPInfo { + server_name: string + description?: string + logo_url?: string + mcp_server_cost_info?: MCPServerCostInfo | null +} + +// Define the structure for a single MCP tool +export interface MCPTool { + name: string + description?: string + inputSchema: InputSchema | string // API returns string "tool_input_schema" or the actual schema + mcp_info: MCPInfo + // Function to select a tool (added in the component) + onToolSelect?: (tool: MCPTool) => void +} + +// Define the response structure for the listMCPTools endpoint - now a flat array +export type ListMCPToolsResponse = MCPTool[] + +// Define the argument structure for calling an MCP tool +export interface CallMCPToolArgs { + name: string + arguments: Record | null + server_name?: string // Now using server_name from mcp_info +} + +// Define the possible content types in the response +export interface MCPTextContent { + type: "text" + text: string + annotations?: any +} + +export interface MCPImageContent { + type: "image" + url?: string + data?: string +} + +export interface MCPEmbeddedResource { + type: "embedded_resource" + resource_type?: string + url?: string + data?: any +} + +// Define the union type for the content array in the response +export type MCPContent = MCPTextContent | MCPImageContent | MCPEmbeddedResource + +// Define the response structure for the callMCPTool endpoint +export type CallMCPToolResponse = MCPContent[] + +// Props for the main component +export interface MCPToolsViewerProps { + serverId: string + accessToken: string | null + auth_type?: string | null + userRole: string | null + userID: string | null + serverAlias?: string | null +} export interface MCPServer { - server_id: string; - server_name?: string | null; - alias?: string | null; - description?: string | null; - url: string; - transport?: string | null; - auth_type?: string | null; - mcp_info?: MCPInfo | null; - created_at: string; - created_by: string; - updated_at: string; - updated_by: string; - teams?: Team[]; - mcp_access_groups?: string[]; + server_id: string + server_name?: string | null + alias?: string | null + description?: string | null + url: string + transport?: string | null + auth_type?: string | null + mcp_info?: MCPInfo | null + created_at: string + created_by: string + updated_at: string + updated_by: string + teams?: Team[] + mcp_access_groups?: string[] + allowed_tools?: string[] } export interface MCPServerProps { - accessToken: string | null; - userRole: string | null; - userID: string | null; -} \ No newline at end of file + accessToken: string | null + userRole: string | null + userID: string | null +} diff --git a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx new file mode 100644 index 00000000000..de768501ffd --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx @@ -0,0 +1,114 @@ +import { useState, useEffect } from "react" +import { testMCPToolsListRequest } from "../components/networking" + +interface MCPServerConfig { + server_id?: string + server_name?: string + url?: string + transport?: string + auth_type?: string + mcp_info?: any +} + +interface UseTestMCPConnectionProps { + accessToken: string | null + formValues: Record + enabled?: boolean // Optional flag to enable/disable auto-fetching +} + +interface UseTestMCPConnectionReturn { + tools: any[] + isLoadingTools: boolean + toolsError: string | null + hasShownSuccessMessage: boolean + canFetchTools: boolean + fetchTools: () => Promise + clearTools: () => void +} + +export const useTestMCPConnection = ({ + accessToken, + formValues, + enabled = true, +}: UseTestMCPConnectionProps): UseTestMCPConnectionReturn => { + const [tools, setTools] = useState([]) + const [isLoadingTools, setIsLoadingTools] = useState(false) + const [toolsError, setToolsError] = useState(null) + const [hasShownSuccessMessage, setHasShownSuccessMessage] = useState(false) + + // Check if we have the minimum required fields to fetch tools + const canFetchTools = !!(formValues.url && formValues.transport && formValues.auth_type && accessToken) + + const fetchTools = async () => { + if (!accessToken || !formValues.url) { + return + } + + setIsLoadingTools(true) + setToolsError(null) + + try { + // Prepare the MCP server config from form values + const mcpServerConfig: MCPServerConfig = { + server_id: formValues.server_id || "", + server_name: formValues.server_name || "", + url: formValues.url, + transport: formValues.transport, + auth_type: formValues.auth_type, + mcp_info: formValues.mcp_info, + } + + const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig) + + if (toolsResponse.tools && !toolsResponse.error) { + setTools(toolsResponse.tools) + setToolsError(null) + if (toolsResponse.tools.length > 0 && !hasShownSuccessMessage) { + setHasShownSuccessMessage(true) + } + } else { + const errorMessage = toolsResponse.message || "Failed to retrieve tools list" + setToolsError(errorMessage) + setTools([]) + setHasShownSuccessMessage(false) + } + } catch (error) { + console.error("Tools fetch error:", error) + setToolsError(error instanceof Error ? error.message : String(error)) + setTools([]) + setHasShownSuccessMessage(false) + } finally { + setIsLoadingTools(false) + } + } + + const clearTools = () => { + setTools([]) + setToolsError(null) + setHasShownSuccessMessage(false) + } + + // Auto-fetch tools when form values change and required fields are available + useEffect(() => { + if (!enabled) { + return + } + + if (canFetchTools) { + fetchTools() + } else { + clearTools() + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [formValues.url, formValues.transport, formValues.auth_type, accessToken, enabled, canFetchTools]) + + return { + tools, + isLoadingTools, + toolsError, + hasShownSuccessMessage, + canFetchTools, + fetchTools, + clearTools, + } +}