fix(ui/mcp): reset OAuth state on create-server modal close so a prior server's token no longer leaks into the next add-server session (#30000)

* fix(ui/mcp): reset OAuth hook state on modal close so a prior server's token no longer leaks into the next add-server session

* fix(ui/mcp): clear in-flight OAuth guard on reset and reset form/tools on modal close so nothing leaks on a parent-driven dismiss
This commit is contained in:
tin-berri 2026-06-09 14:18:28 -07:00 committed by GitHub
parent 38edf241a4
commit d8fe091938
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 253 additions and 14 deletions

View File

@ -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<string, unknown> | null }));
// produced a token before submit, and inspect the reset wiring.
const oauthHook = vi.hoisted(() => ({
tokenResponse: null as Record<string, unknown> | null,
reset: vi.fn(),
onTokenReceived: null as ((token: Record<string, unknown> | null) => void) | null,
}));
vi.mock("@/hooks/useMcpOAuthFlow", () => ({
useMcpOAuthFlow: () => ({
startOAuthFlow: vi.fn(),
status: "idle",
error: null,
tokenResponse: oauthHook.tokenResponse,
}),
useMcpOAuthFlow: (opts: { onTokenReceived: (token: Record<string, unknown> | 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: () => <div data-testid="mcp-connection-status" />,
default: ({ tools }: { tools?: any[] }) => (
<div data-testid="mcp-connection-status" data-tool-count={tools?.length ?? 0} />
),
}));
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(<CreateMCPServer {...defaultProps} />);
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(<CreateMCPServer {...defaultProps} isModalVisible={false} />);
// 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(<CreateMCPServer {...defaultProps} isModalVisible={true} />);
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(<CreateMCPServer {...defaultProps} />);
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(<CreateMCPServer {...defaultProps} isModalVisible={false} />);
// 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(<CreateMCPServer {...defaultProps} isModalVisible={true} />);
const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com") as HTMLInputElement;
expect(reopenedUrlInput.value).toBe("");
});
});
describe("when stdio transport is selected", () => {

View File

@ -134,6 +134,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
status: oauthStatus,
error: oauthError,
tokenResponse: oauthTokenResponse,
reset: resetOAuthFlow,
} = useMcpOAuthFlow({
accessToken,
getCredentials: () => form.getFieldValue("credentials"),
@ -554,12 +555,19 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
}
}, [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);

View File

@ -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<any>(() => {}));
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));
});
});

View File

@ -40,6 +40,7 @@ interface UseMcpOAuthFlowResult {
status: McpOAuthStatus;
error: string | null;
tokenResponse: Record<string, any> | 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,
};
};

View File

@ -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(() => {