diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index 042f02251b..e70548d6a9 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -17,15 +17,23 @@ vi.mock("@/utils/mcpTokenStore", () => ({ })); // Mutable holder so individual tests can simulate "Authorize & Fetch" having -// produced a token before submit. -const oauthHook = vi.hoisted(() => ({ tokenResponse: null as Record | null })); +// produced a token before submit, and inspect the reset wiring. +const oauthHook = vi.hoisted(() => ({ + tokenResponse: null as Record | null, + reset: vi.fn(), + onTokenReceived: null as ((token: Record | null) => void) | null, +})); vi.mock("@/hooks/useMcpOAuthFlow", () => ({ - useMcpOAuthFlow: () => ({ - startOAuthFlow: vi.fn(), - status: "idle", - error: null, - tokenResponse: oauthHook.tokenResponse, - }), + useMcpOAuthFlow: (opts: { onTokenReceived: (token: Record | null) => void }) => { + oauthHook.onTokenReceived = opts.onTokenReceived; + return { + startOAuthFlow: vi.fn(), + status: "idle", + error: null, + tokenResponse: oauthHook.tokenResponse, + reset: oauthHook.reset, + }; + }, })); vi.mock("./mcp_server_cost_config", () => ({ @@ -59,7 +67,9 @@ vi.mock("./mcp_tool_configuration", () => ({ })); vi.mock("./mcp_connection_status", () => ({ - default: () =>
, + default: ({ tools }: { tools?: any[] }) => ( +
+ ), })); vi.mock("./StdioConfiguration", () => ({ @@ -121,6 +131,7 @@ describe("CreateMCPServer", () => { beforeEach(() => { vi.clearAllMocks(); oauthHook.tokenResponse = null; + oauthHook.onTokenReceived = null; }); it("should render the modal with title when visible", () => { @@ -614,6 +625,100 @@ describe("CreateMCPServer", () => { expect(defaultProps.setModalVisible).toHaveBeenCalledWith(false); }); + + it("does not leak a previous server's OAuth token into the next add-server session", async () => { + const usedToken = (token: string) => + vi.mocked(networking.testMCPToolsListRequest).mock.calls.some((call) => call[2] === token); + + const { rerender } = render(); + + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + await selectAntOption("Authentication", "OAuth"); + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } }); + }); + + // Simulate "Authorize & Fetch Token" completing for server A. + await act(async () => { + oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 }); + }); + + // Precondition: the freshly fetched token drives the tool preview for server A. + await waitFor(() => { + expect(usedToken("stale-token-A")).toBe(true); + }); + + // Parent hides the modal (Cancel / successful create both flip this prop). + rerender(); + + // The OAuth flow state (source of the "Token fetched" badge) is reset on close. + expect(oauthHook.reset).toHaveBeenCalled(); + + vi.mocked(networking.testMCPToolsListRequest).mockClear(); + + // Reopen for a brand-new server and enter a different URL without re-authorizing. + rerender(); + const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(reopenedUrlInput, { target: { value: "https://server-b.example.com/mcp" } }); + }); + + // The previous server's token must never be replayed for the new session. + expect(usedToken("stale-token-A")).toBe(false); + }); + + it("clears the tool list and form fields when a parent dismisses the modal", async () => { + vi.mocked(networking.testMCPToolsListRequest).mockResolvedValue({ + tools: [{ name: "tool_a" }], + error: null, + }); + const toolCount = () => screen.getByTestId("mcp-connection-status").getAttribute("data-tool-count"); + + const { rerender } = render(); + + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + await selectAntOption("Authentication", "OAuth"); + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } }); + }); + await act(async () => { + oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 }); + }); + + // Precondition: a tool list is shown for server A. + await waitFor(() => { + expect(toolCount()).toBe("1"); + }); + + // Parent dismisses the modal without routing through Cancel or create. + rerender(); + + // Stale tools are cleared even though neither handler ran. + await waitFor(() => { + expect(toolCount()).toBe("0"); + }); + + // Reopening starts clean: the URL the prior server left in the Ant form store is gone. + rerender(); + const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com") as HTMLInputElement; + expect(reopenedUrlInput.value).toBe(""); + }); }); describe("when stdio transport is selected", () => { 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 28c38c459a..ddcc9f65d3 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 @@ -134,6 +134,7 @@ const CreateMCPServer: React.FC = ({ status: oauthStatus, error: oauthError, tokenResponse: oauthTokenResponse, + reset: resetOAuthFlow, } = useMcpOAuthFlow({ accessToken, getCredentials: () => form.getFieldValue("credentials"), @@ -554,12 +555,19 @@ const CreateMCPServer: React.FC = ({ } }, [formValues.server_name]); - // Clear formValues when modal closes to reset child components + // Clear form, tools, and OAuth state when the modal closes so a previous server's + // authorization, credentials, or tool list never bleed into the next "Add New MCP + // Server" session, including when a parent dismisses the modal without routing + // through handleCancel or handleCreate. React.useEffect(() => { if (!isModalVisible) { + form.resetFields(); setFormValues({}); + setOauthAccessToken(null); + clearTools(); + resetOAuthFlow(); } - }, [isModalVisible]); + }, [isModalVisible, form, clearTools, resetOAuthFlow]); const isAdmin = isAdminRole(userRole); diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx new file mode 100644 index 0000000000..d8dce3c29e --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx @@ -0,0 +1,117 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import * as networking from "@/components/networking"; +import { setSecureItem } from "@/utils/secureStorage"; +import { useMcpOAuthFlow } from "./useMcpOAuthFlow"; + +vi.mock("@/components/networking", () => ({ + exchangeMcpOAuthToken: vi.fn(), + cacheTemporaryMcpServer: vi.fn(), + registerMcpOAuthClient: vi.fn(), + buildMcpOAuthAuthorizeUrl: vi.fn(), + getProxyBaseUrl: vi.fn(() => ""), + serverRootPath: "", +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { success: vi.fn(), error: vi.fn() }, +})); + +const FLOW_STATE_KEY = "litellm-mcp-oauth-flow-state"; +const RESULT_KEY = "litellm-mcp-oauth-result"; + +/** Seed the redirect result (the code returned by the IdP callback). */ +function seedResult(code: string) { + setSecureItem(RESULT_KEY, JSON.stringify({ state: "state-1", code })); +} + +/** Seed the flow state stored before the redirect. */ +function seedFlowState() { + setSecureItem( + FLOW_STATE_KEY, + JSON.stringify({ + state: "state-1", + codeVerifier: "verifier-1", + serverId: "server-1", + clientId: "client-1", + redirectUri: "https://app.example.com/ui/mcp/oauth/callback", + flowSource: "create", + }), + ); +} + +/** Seed storage so the hook's on-mount resume flow exchanges a code for a token. */ +function seedCompletedRedirect() { + seedResult("code-1"); + seedFlowState(); +} + +function renderFlow(onTokenReceived = vi.fn()) { + return renderHook( + ({ onTokenReceived: cb }: { onTokenReceived: (t: any) => void }) => + useMcpOAuthFlow({ + accessToken: "admin-token", + getCredentials: () => ({}), + getTemporaryPayload: () => ({ url: "https://server-1.example.com/mcp", transport: "http" }), + onTokenReceived: cb, + flowSource: "create", + }), + { initialProps: { onTokenReceived } }, + ); +} + +describe("useMcpOAuthFlow reset", () => { + beforeEach(() => { + vi.clearAllMocks(); + window.sessionStorage.clear(); + window.localStorage.clear(); + }); + + it("clears a successfully fetched token so it cannot leak into the next session", async () => { + const token = { access_token: "tok-123", expires_in: 3600 }; + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue(token); + seedCompletedRedirect(); + + const onTokenReceived = vi.fn(); + const { result } = renderFlow(onTokenReceived); + + await waitFor(() => expect(result.current.status).toBe("success")); + expect(result.current.tokenResponse).toEqual(token); + expect(onTokenReceived).toHaveBeenCalledWith(token); + + act(() => { + result.current.reset(); + }); + + expect(result.current.status).toBe("idle"); + expect(result.current.tokenResponse).toBeNull(); + expect(result.current.error).toBeNull(); + }); + + it("clears the in-flight guard so a callback after a mid-exchange close is not swallowed", async () => { + // First exchange hangs, mimicking the modal being closed while the token + // endpoint is still in flight. processingRef is left true at that point. + vi.mocked(networking.exchangeMcpOAuthToken).mockReturnValueOnce(new Promise(() => {})); + seedFlowState(); + seedResult("code-1"); + + const onTokenReceived1 = vi.fn(); + const { result, rerender } = renderFlow(onTokenReceived1); + + await waitFor(() => expect(result.current.status).toBe("exchanging")); + + act(() => { + result.current.reset(); + }); + + // The reopened modal receives a fresh callback; it must be processed, not + // dropped by a stale in-flight guard. + const token = { access_token: "tok-2" }; + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValueOnce(token); + seedResult("code-2"); + const onTokenReceived2 = vi.fn(); + rerender({ onTokenReceived: onTokenReceived2 }); + + await waitFor(() => expect(onTokenReceived2).toHaveBeenCalledWith(token)); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index c4561a9f53..ee253577d9 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -40,6 +40,7 @@ interface UseMcpOAuthFlowResult { status: McpOAuthStatus; error: string | null; tokenResponse: Record | null; + reset: () => void; } export const useMcpOAuthFlow = ({ @@ -336,10 +337,18 @@ export const useMcpOAuthFlow = ({ resumeOAuthFlow(); }, [resumeOAuthFlow]); + const reset = useCallback(() => { + setStatus("idle"); + setError(null); + setTokenResponse(null); + processingRef.current = false; + }, []); + return { startOAuthFlow, status, error, tokenResponse, + reset, }; }; diff --git a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx index 3dcd6be888..b160c90224 100644 --- a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx +++ b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx @@ -1,4 +1,4 @@ -import { useState, useEffect } from "react"; +import { useState, useEffect, useCallback } from "react"; import { testMCPToolsListRequest } from "../components/networking"; import { AUTH_TYPE, OAUTH_FLOW, TRANSPORT } from "@/components/mcp_tools/types"; @@ -177,12 +177,12 @@ export const useTestMCPConnection = ({ } }; - const clearTools = () => { + const clearTools = useCallback(() => { setTools([]); setToolsError(null); setToolsErrorStackTrace(null); setHasShownSuccessMessage(false); - }; + }, []); // Auto-fetch tools when form values change and required fields are available useEffect(() => {