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:
parent
38edf241a4
commit
d8fe091938
@ -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", () => {
|
||||
|
||||
@ -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);
|
||||
|
||||
|
||||
117
ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx
Normal file
117
ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx
Normal 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));
|
||||
});
|
||||
});
|
||||
@ -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,
|
||||
};
|
||||
};
|
||||
|
||||
@ -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(() => {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user