From 95a9d1e9b7c907c1da9234151adb0f15864fa092 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 8 Jun 2026 13:35:49 -0700 Subject: [PATCH] fix(mcp): let non-creator users OAuth into OBO-mode MCP servers from the Tools page (#29867) * fix(ui): let non-creator users OAuth into OBO-mode MCP servers from the Tools page * fix(ui): clear OBO Tools-tab one-shot on navigate-back and gate on credential-status errors (cherry picked from commit 1528f43d4c5f6c9981d065761834957cf559f7b5) --- .../components/mcp_tools/mcp_server_view.tsx | 4 +- .../src/components/mcp_tools/mcp_servers.tsx | 41 +++++- .../components/mcp_tools/mcp_tools.test.tsx | 55 +++++++- .../src/components/mcp_tools/mcp_tools.tsx | 119 +++++++++++++++--- .../src/hooks/mcpOAuthUtils.ts | 9 ++ 5 files changed, 209 insertions(+), 19 deletions(-) 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 cafecdafd45..fc1fe7cd405 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 @@ -21,6 +21,7 @@ interface MCPServerViewProps { userRole: string | null; userID: string | null; availableAccessGroups: string[]; + initialTabIndex?: number; } export const MCPServerView: React.FC = ({ @@ -32,11 +33,12 @@ export const MCPServerView: React.FC = ({ userRole, userID, availableAccessGroups, + initialTabIndex = 0, }) => { const [editing, setEditing] = useState(isEditing); const [showFullUrl, setShowFullUrl] = useState(false); const [copiedStates, setCopiedStates] = useState>({}); - const [selectedTabIndex, setSelectedTabIndex] = useState(0); + const [selectedTabIndex, setSelectedTabIndex] = useState(initialTabIndex); const handleSuccess = (updated: MCPServer) => { setEditing(false); 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 4b383f701cb..0a445265e2d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -21,6 +21,7 @@ import MCPNetworkSettings from "./MCPNetworkSettings"; import MCPDiscovery from "./mcp_discovery"; import { ByokCredentialModal } from "./ByokCredentialModal"; import { getSecureItem } from "@/utils/secureStorage"; +import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; import UserEnvVarsModal from "./UserEnvVarsModal"; import { listMCPUserEnvVarStatus } from "../networking"; @@ -71,6 +72,23 @@ const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => { const { Text: AntdText, Title: AntdTitle } = Typography; const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; +// Server id stashed by the Tools tab before an OBO OAuth redirect, read once at +// mount so the redirect returns straight to that server's Tools tab. +const readToolsOAuthServerId = (): string | null => { + if (typeof window === "undefined") { + return null; + } + try { + const stored = getSecureItem(TOOLS_OAUTH_UI_STATE_KEY); + if (!stored) { + return null; + } + return JSON.parse(stored)?.serverId ?? null; + } catch { + return null; + } +}; + const { Option } = Select; const MCPServers: React.FC = ({ accessToken, userRole, userID }) => { @@ -103,7 +121,12 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // state const [serverIdToDelete, setServerToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [selectedServerId, setSelectedServerId] = useState(null); + // Server whose Tools tab should be reopened after an OBO OAuth redirect; read + // once from sessionStorage so the restored server selection is correct on the + // first render. Cleared when the user navigates back to the list (handleBack) + // so a later visit to the same server defaults to Overview, not the Tools tab. + const [toolsTabServerId, setToolsTabServerId] = useState(readToolsOAuthServerId); + const [selectedServerId, setSelectedServerId] = useState(toolsTabServerId); const [editServer, setEditServer] = useState(false); const [selectedTeam, setSelectedTeam] = useState("all"); const [selectedMcpAccessGroup, setSelectedMcpAccessGroup] = useState("all"); @@ -178,6 +201,19 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } }, []); + // The restored server id was consumed by the initializer above; remove the + // one-shot sessionStorage key so a full page reload doesn't reopen the Tools + // tab (removeItem only, no setState). + useEffect(() => { + if (typeof window !== "undefined") { + try { + window.sessionStorage.removeItem(TOOLS_OAUTH_UI_STATE_KEY); + } catch { + // ignore storage errors + } + } + }, []); + // Get unique teams from all servers const uniqueTeams = React.useMemo(() => { if (!serversWithHealth) return []; @@ -338,6 +374,8 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const handleBack = React.useCallback(() => { setEditServer(false); setSelectedServerId(null); + // Drop the post-redirect one-shot so re-selecting that server opens Overview. + setToolsTabServerId(null); refetch(); }, [refetch]); @@ -483,6 +521,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) userID={userID} userRole={userRole} availableAccessGroups={uniqueMcpAccessGroups} + initialTabIndex={selectedServerId === toolsTabServerId ? 1 : 0} /> ) : (
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx index 4917f1fcf1d..72ed3984244 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx @@ -2,12 +2,13 @@ import { render, screen, waitFor } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { describe, expect, it, vi, beforeEach } from "vitest"; import MCPToolsViewer from "./mcp_tools"; -import { listMCPTools } from "../networking"; +import { listMCPTools, getMCPOAuthUserCredentialStatus } from "../networking"; import { isTokenValid, getToken } from "@/utils/mcpTokenStore"; vi.mock("../networking", () => ({ listMCPTools: vi.fn(), callMCPTool: vi.fn(), + getMCPOAuthUserCredentialStatus: vi.fn(), })); vi.mock("@/utils/mcpTokenStore", () => ({ @@ -20,6 +21,10 @@ vi.mock("@/hooks/useToolsOAuthFlow", () => ({ useToolsOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null }), })); +vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ + useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null }), +})); + const GATE_TEXT = "Authentication required"; // Realistic interactive servers carry a token endpoint; the old heuristic // (`oauth2 && !tokenUrl`) mislabeled exactly these as M2M. Setting it here is @@ -42,6 +47,13 @@ const renderViewer = (props: Record) => , ); +const credStatus = (overrides: Record = {}) => ({ + server_id: "srv-1", + has_credential: true, + is_expired: false, + ...overrides, +}); + describe("MCPToolsViewer auth gate routing", () => { beforeEach(() => { vi.mocked(listMCPTools).mockReset().mockResolvedValue({ tools: [], error: null }); @@ -49,6 +61,8 @@ describe("MCPToolsViewer auth gate routing", () => { vi.mocked(getToken) .mockReset() .mockReturnValue(undefined as any); + // Default: the OBO credential exists and is valid, so OBO servers list tools. + vi.mocked(getMCPOAuthUserCredentialStatus).mockReset().mockResolvedValue(credStatus()); }); it("shows the Authorize gate for a passthrough server with a token endpoint and does not list tools", async () => { @@ -57,6 +71,8 @@ describe("MCPToolsViewer auth gate routing", () => { expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument(); expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument(); expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled(); + // Passthrough must not consult the per-user DB credential. + expect(vi.mocked(getMCPOAuthUserCredentialStatus)).not.toHaveBeenCalled(); }); it("forwards the session token via the x-mcp header for a passthrough server that has one", async () => { @@ -75,7 +91,40 @@ describe("MCPToolsViewer auth gate routing", () => { expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument(); }); - it("does not gate an OBO server with a token endpoint; lists with the LiteLLM key and no x-mcp header", async () => { + it("lists tools for an OBO server when the user has a DB credential, with no x-mcp header", async () => { + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); + + await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined)); + expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument(); + }); + + it("shows the Authorize gate for an OBO server when the user has no DB credential and does not list tools", async () => { + vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue(credStatus({ has_credential: false })); + + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); + + expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument(); + expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled(); + }); + + it("shows the Authorize gate for an OBO server when the credential-status check fails", async () => { + vi.mocked(getMCPOAuthUserCredentialStatus).mockRejectedValue(new Error("network down")); + + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); + + expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument(); + expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled(); + }); + + it("does not gate an OBO server whose stored token is expired; the list call refreshes it server-side", async () => { + // has_credential=true with is_expired=true must NOT gate: resolve_valid_user_oauth_token + // refreshes from the stored refresh_token on the list call, so the user never reauthorizes. + vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue( + credStatus({ has_credential: true, is_expired: true }), + ); + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined)); @@ -87,5 +136,7 @@ describe("MCPToolsViewer auth gate routing", () => { await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined)); expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument(); + // M2M uses the backend service token, not a per-user DB credential. + expect(vi.mocked(getMCPOAuthUserCredentialStatus)).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx index 5ec9a3d683a..8957c7b30b8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx @@ -1,11 +1,14 @@ -import React, { useEffect, useState } from "react"; +import React, { useCallback, useEffect, useState } from "react"; import { useQuery, useMutation } from "@tanstack/react-query"; import { ToolTestPanel } from "./ToolTestPanel"; import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse, getMcpOAuthMode } from "./types"; -import { listMCPTools, callMCPTool } from "../networking"; +import { listMCPTools, callMCPTool, getMCPOAuthUserCredentialStatus } from "../networking"; import { isTokenValid, getToken, removeToken } from "@/utils/mcpTokenStore"; import { sanitizeMcpAliasForHeader } from "@/utils/mcpHeaderUtils"; import { useToolsOAuthFlow } from "@/hooks/useToolsOAuthFlow"; +import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; +import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; +import { setSecureItem } from "@/utils/secureStorage"; import { Card, Title, Text } from "@tremor/react"; import { RobotOutlined, ToolOutlined, SearchOutlined, KeyOutlined, LockOutlined } from "@ant-design/icons"; @@ -31,11 +34,14 @@ const MCPToolsViewer = ({ const [passthroughHeaders, setPassthroughHeaders] = useState>({}); const [showHeaderInput, setShowHeaderInput] = useState(false); - // Only PKCE passthrough uses a browser-held session token (sessionStorage, - // cleared on tab/browser close) and a user-facing auth gate. OBO uses the - // backend-stored per-user token and M2M uses the backend's own service token, - // so neither needs a gate — they list tools with just the LiteLLM key. - const isPassthrough = getMcpOAuthMode({ auth_type, oauth2_flow, delegate_auth_to_upstream }) === "passthrough"; + // PKCE passthrough holds a browser-side session token (sessionStorage) and + // gates tool listing behind it. OBO uses a backend-stored per-user token that + // the user must establish once via an interactive login; we gate on whether + // that DB credential exists. M2M uses the backend's own service token and + // needs no gate. + const oauthMode = getMcpOAuthMode({ auth_type, oauth2_flow, delegate_auth_to_upstream }); + const isPassthrough = oauthMode === "passthrough"; + const isObo = oauthMode === "obo"; const [oauthToken, setOauthToken] = useState(() => isPassthrough && isTokenValid(serverId, userID) ? getToken(serverId, userID)?.access_token ?? null : null, ); @@ -61,6 +67,31 @@ const MCPToolsViewer = ({ onSuccess: setOauthToken, }); + // OBO servers list tools using a per-user token the backend stores in the DB; + // check whether the current user has a valid one so we can prompt them to + // authorize when they don't (otherwise the backend silently returns no tools). + const { + data: oboCredStatus, + isLoading: isLoadingOboCred, + isError: isOboCredError, + refetch: refetchOboCred, + } = useQuery({ + queryKey: ["mcpOauthUserCredStatus", serverId, userID], + queryFn: () => getMCPOAuthUserCredentialStatus(accessToken ?? "", serverId), + enabled: !!accessToken && isObo, + staleTime: 30000, + }); + + // A stored credential is sufficient: the backend proactively refreshes an + // expired or near-expiry token from the stored refresh_token on the next list + // call, so the user only needs to authorize when no credential row exists. If + // the status check itself fails we can't confirm a credential, so surface the + // Authorize gate rather than a silent empty tool list; re-authorizing only + // overwrites the user's own row, so it is safe when a credential did exist. + const hasOboCred = !!oboCredStatus?.has_credential; + const oboNeedsAuth = isObo && !isLoadingOboCred && (isOboCredError || (!!oboCredStatus && !hasOboCred)); + const oboStatusLoading = isObo && isLoadingOboCred; + // Check if this server has extra headers configured const hasExtraHeaders = extraHeaders && extraHeaders.length > 0; @@ -135,8 +166,9 @@ const MCPToolsViewer = ({ } return result; }, - // For OAuth servers, block the query until a session token is available - enabled: !!accessToken && (!isPassthrough || oauthToken !== null), + // Passthrough blocks until a browser session token exists; OBO blocks until + // the user has a valid DB credential (else the backend returns no tools). + enabled: !!accessToken && (isPassthrough ? oauthToken !== null : isObo ? hasOboCred : true), staleTime: 30000, // Consider data fresh for 30 seconds retry: (failureCount, error: any) => { // Don't retry on 401 — token is invalid, user must re-authenticate @@ -145,6 +177,33 @@ const MCPToolsViewer = ({ }, }); + // OBO authorize: same redirect+exchange flow as the admin "Authorize & Fetch" + // and the chat "Connect" button, but persists the token to the per-user DB. + const onOboAuthSuccess = useCallback(() => { + refetchOboCred(); + refetchTools(); + }, [refetchOboCred, refetchTools]); + + const { + startOAuthFlow: startDbOAuthFlow, + status: dbOAuthStatus, + error: dbOAuthError, + } = useUserMcpOAuthFlow({ + accessToken: accessToken ?? "", + serverId, + serverAlias, + onSuccess: onOboAuthSuccess, + }); + + // Stash which server started the redirect so the MCP Servers page can reopen + // this Tools tab on return and let the flow resume to persist the credential. + const startOboAuthorize = useCallback(() => { + try { + setSecureItem(TOOLS_OAUTH_UI_STATE_KEY, JSON.stringify({ serverId })); + } catch (_) {} + startDbOAuthFlow(); + }, [serverId, startDbOAuthFlow]); + // If the tools query fails with 401, the cached OAuth token is invalid — // clear it so the auth gate is shown again and the user can re-authenticate. useEffect(() => { @@ -187,6 +246,13 @@ const MCPToolsViewer = ({ const toolsData = mcpToolsResponse?.tools || []; + // An auth gate replaces the tool list when the user must authenticate first: + // passthrough needs a browser token, OBO needs a stored DB credential. + const authGateActive = (isPassthrough && !oauthToken) || oboNeedsAuth; + // Treat OBO credential-status loading as "tools loading" so the empty state + // doesn't flash before we know whether the user needs to authorize. + const toolsAreaLoading = isLoadingTools || oboStatusLoading; + // Filter tools based on search term const filteredTools = toolsData.filter((tool: MCPTool) => { const searchLower = toolSearchTerm.toLowerCase(); @@ -287,7 +353,7 @@ const MCPToolsViewer = ({ )} - {/* OAuth Auth Gate — shown when token is absent for OAuth servers */} + {/* Passthrough auth gate — browser session token absent */} {isPassthrough && !oauthToken && (
@@ -306,8 +372,31 @@ const MCPToolsViewer = ({
)} + {/* OBO auth gate — only when no credential row exists for this user. + An existing-but-expired token is refreshed server-side on the + list call, so the gate never appears for a stored credential. */} + {oboNeedsAuth && ( +
+ +

Authentication required

+

+ Authenticate with the upstream provider to view available tools +

+ + Authorize + + {dbOAuthError &&

{dbOAuthError}

} +
+ )} + {/* Search Bar — only shown when tools are loaded */} - {!isPassthrough || oauthToken ? ( + {!authGateActive ? ( <> {toolsData.length > 0 && (
@@ -324,7 +413,7 @@ const MCPToolsViewer = ({ )} {/* Loading State */} - {isLoadingTools && ( + {toolsAreaLoading && (
@@ -335,7 +424,7 @@ const MCPToolsViewer = ({ )} {/* Error State */} - {(mcpToolsResponse?.error || mcpToolsError) && !isLoadingTools && !toolsData.length && ( + {(mcpToolsResponse?.error || mcpToolsError) && !toolsAreaLoading && !toolsData.length && (

Error: {mcpToolsResponse?.message || (mcpToolsError as Error)?.message} @@ -344,7 +433,7 @@ const MCPToolsViewer = ({ )} {/* No Tools State */} - {!isLoadingTools && + {!toolsAreaLoading && !mcpToolsResponse?.error && !mcpToolsError && (!toolsData || toolsData.length === 0) && ( @@ -370,7 +459,7 @@ const MCPToolsViewer = ({ )} {/* Tools List */} - {!isLoadingTools && !mcpToolsResponse?.error && toolsData.length > 0 && ( + {!toolsAreaLoading && !mcpToolsResponse?.error && toolsData.length > 0 && ( <> {filteredTools.length === 0 ? (

diff --git a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts index 3aff8af6eef..83b2cb21bed 100644 --- a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts +++ b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts @@ -7,6 +7,15 @@ import { getProxyBaseUrl, serverRootPath } from "@/components/networking"; +/** + * sessionStorage key used to restore the MCP server detail view on the Tools + * tab after a full-page OAuth redirect. The OBO authorize flow redirects to the + * IdP and back to the MCP Servers page; without this the user lands on the + * server list and useUserMcpOAuthFlow never re-mounts to persist the credential. + * Mirrors the admin edit flow's EDIT_OAUTH_UI_STATE_KEY. + */ +export const TOOLS_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-tools-state"; + /** * Build the OAuth callback URL for the current UI deployment. *