From 5010b2b0cd7bdbf692ab1d1323c287c7f0ec4a00 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 11 Mar 2026 20:09:11 -0700 Subject: [PATCH] fix: batch OAuth creds query, hide empty CRUD groups on search, onChange stability --- .../mcp_server/rest_endpoints.py | 45 ++++++++++++++++++- .../src/components/chat/MCPAppsPanel.tsx | 14 +++--- .../mcp_tools/McpCrudPermissionPanel.tsx | 12 +++++ 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index dc846de0bf4..2118804472c 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -108,6 +108,42 @@ if MCP_AVAILABLE: ) return None + async def _get_bulk_user_oauth_headers( + user_api_key_dict: UserAPIKeyAuth, + ) -> Dict[str, Dict[str, str]]: + """ + Fetch ALL OAuth2 credentials for the current user in a single DB query and + return a mapping of server_id → {"Authorization": "Bearer "}. + + This is the batch alternative to calling _get_user_oauth_extra_headers + per-server inside a loop (N+1 DB queries). + """ + user_id = getattr(user_api_key_dict, "user_id", None) + if not user_id: + return {} + try: + from litellm.proxy._experimental.mcp_server.db import ( + list_user_oauth_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + creds = await list_user_oauth_credentials(prisma_client, user_id) + return { + c["server_id"]: {"Authorization": f"Bearer {c['access_token']}"} + for c in creds + if c.get("access_token") and c.get("server_id") + } + except Exception: + verbose_logger.debug( + "Failed to bulk-fetch OAuth credentials for user %r", + str(user_id).replace("\n", "\\n").replace("\r", "\\r"), + exc_info=True, + ) + return {} + def _create_tool_response_objects(tools, server_mcp_info): """Helper function to create tool response objects.""" return [ @@ -333,6 +369,11 @@ if MCP_AVAILABLE: list_tools_result = [] error_message = None + # Bulk-fetch all OAuth credentials for this user in a single DB query + # so per-server calls below can use a dict lookup (O(1)) instead of + # issuing one DB query per server (N+1 pattern). + bulk_oauth_headers = await _get_bulk_user_oauth_headers(user_api_key_dict) + # If server_id is specified, only query that specific server if server_id: # Resolve a server name to its UUID if needed (MCPConnectPicker passes @@ -381,7 +422,7 @@ if MCP_AVAILABLE: server_auth_header = _get_server_auth_header( server, mcp_server_auth_headers, mcp_auth_header ) - user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict) + user_oauth_extra_headers = bulk_oauth_headers.get(server.server_id) try: list_tools_result = await _get_tools_for_single_server( @@ -435,7 +476,7 @@ if MCP_AVAILABLE: server_auth_header = _get_server_auth_header( server, mcp_server_auth_headers, mcp_auth_header ) - user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict) + user_oauth_extra_headers = bulk_oauth_headers.get(server.server_id) try: tools_result = await _get_tools_for_single_server( diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 2afe032e0db..e1ad1d92f84 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -100,12 +100,14 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange const [oauthConnected, setOauthConnected] = useState>(new Set()); // Refs keep the latest values for the auto-enable effect so it always reads - // the current servers/selectedServers without needing them as dependencies - // (which would cause the effect to fire on every render). + // the current servers/selectedServers/onChange without needing them as + // dependencies (which would cause the effect to fire on every render). const serversRef = useRef([]); useEffect(() => { serversRef.current = servers; }, [servers]); const selectedServersRef = useRef(selectedServers); useEffect(() => { selectedServersRef.current = selectedServers; }, [selectedServers]); + const onChangeRef = useRef(onChange); + useEffect(() => { onChangeRef.current = onChange; }, [onChange]); const nameOf = (s: MCPServer) => s.server_name ?? s.alias ?? s.server_id; @@ -165,17 +167,17 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange // Auto-enable oauth2 servers for the current chat session when a valid // credential is detected (either on mount or after a fresh OAuth sign-in). - // Uses refs for servers/selectedServers to avoid stale closures without - // adding them as deps (which would cause the effect to re-fire on every render). + // Uses refs for servers/selectedServers/onChange to avoid stale closures + // without adding them as dependencies (which would re-fire on every render). useEffect(() => { if (oauthConnected.size === 0) return; const namesToAdd = serversRef.current .filter((s) => oauthConnected.has(s.server_id) && !selectedServersRef.current.includes(nameOf(s))) .map(nameOf); if (namesToAdd.length > 0) { - onChange([...selectedServersRef.current, ...namesToAdd]); + onChangeRef.current([...selectedServersRef.current, ...namesToAdd]); } - }, [oauthConnected, onChange]); + }, [oauthConnected]); const handleToggle = async (serverName: string, checked: boolean, serverId?: string) => { if (!checked) { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/McpCrudPermissionPanel.tsx b/ui/litellm-dashboard/src/components/mcp_tools/McpCrudPermissionPanel.tsx index 34985e00cd8..4ad46918797 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/McpCrudPermissionPanel.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/McpCrudPermissionPanel.tsx @@ -147,6 +147,18 @@ const McpCrudPermissionPanel: React.FC = ({ const group = grouped[op]; if (group.length === 0) return null; + // If a search filter is active and no tools in this group match, hide the + // entire group — including its header — to avoid empty visual blocks. + if (searchFilter) { + const lf = searchFilter.toLowerCase(); + const hasMatch = group.some( + (t) => + t.name.toLowerCase().includes(lf) || + (t.description ?? "").toLowerCase().includes(lf) + ); + if (!hasMatch) return null; + } + const meta = CRUD_GROUP_META[op]; const fullyAllowed = isGroupFullyAllowed(op); const partial = isGroupPartiallyAllowed(op);