diff --git a/ui/litellm-dashboard/eslint-metrics.json b/ui/litellm-dashboard/eslint-metrics.json index 219cb0580e7..0a4f4cec3b9 100644 --- a/ui/litellm-dashboard/eslint-metrics.json +++ b/ui/litellm-dashboard/eslint-metrics.json @@ -1,6 +1,6 @@ { "@typescript-eslint/no-explicit-any": 1990, - "complexity": 128, + "complexity": 127, "max-depth": 59, "no-console": 15 } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 1c8f92b720f..ee5a90ad3d7 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2117,12 +2117,7 @@ "count": 1 } }, - "src/hooks/useTestMCPConnection.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/hooks/useToolsOAuthFlow.tsx": { + "src/hooks/useMcpOAuthPkceFlow.ts": { "react-hooks/refs": { "count": 1 }, @@ -2130,7 +2125,7 @@ "count": 1 } }, - "src/hooks/useUserMcpOAuthFlow.tsx": { + "src/hooks/useTestMCPConnection.tsx": { "react-hooks/set-state-in-effect": { "count": 1 } diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.test.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.test.tsx new file mode 100644 index 00000000000..5fe550f669f --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.test.tsx @@ -0,0 +1,160 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import * as networking from "@/components/networking"; +import * as mcpTokenStore from "@/utils/mcpTokenStore"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { useToolsOAuthFlow } from "./useToolsOAuthFlow"; +import { useUserMcpOAuthFlow } from "./useUserMcpOAuthFlow"; + +vi.mock("@/components/networking", () => ({ + exchangeMcpOAuthToken: vi.fn(), + registerMcpOAuthClient: vi.fn(), + buildMcpOAuthAuthorizeUrl: vi.fn(() => "https://idp.example.com/authorize"), + storeMCPOAuthUserCredential: vi.fn(), + getProxyBaseUrl: vi.fn(() => ""), + serverRootPath: "", +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { success: vi.fn(), error: vi.fn() }, +})); + +const TOOLS_FLOW_STATE_KEY = "litellm-tools-mcp-oauth-flow-state"; +const TOOLS_RESULT_KEY = "litellm-tools-mcp-oauth-result"; +const USER_FLOW_STATE_KEY = "litellm-user-mcp-oauth-flow-state"; +const USER_RESULT_KEY = "litellm-user-mcp-oauth-result"; + +/** Seed the storage for a completed IdP redirect, keyed for a specific flow. */ +function seedCompletedRedirect(flowStateKey: string, resultKey: string, extra: Record = {}) { + setSecureItem(resultKey, JSON.stringify({ state: "state-1", code: "code-1" })); + setSecureItem( + flowStateKey, + JSON.stringify({ + state: "state-1", + codeVerifier: "verifier-1", + serverId: "server-1", + clientId: "client-1", + redirectUri: "https://app.example.com/ui/mcp/oauth/callback", + ...extra, + }), + ); +} + +describe("MCP OAuth PKCE flow wrappers", () => { + beforeEach(() => { + vi.clearAllMocks(); + window.sessionStorage.clear(); + window.localStorage.clear(); + }); + + describe("useToolsOAuthFlow", () => { + it("persists the exchanged token to sessionStorage (with userId) and never to the backend", async () => { + const token = { access_token: "tok-tools", expires_in: 3600, refresh_token: "ref-tools", token_type: "bearer" }; + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue(token); + const setTokenSpy = vi.spyOn(mcpTokenStore, "setToken"); + seedCompletedRedirect(TOOLS_FLOW_STATE_KEY, TOOLS_RESULT_KEY); + + const onSuccess = vi.fn(); + const { result } = renderHook(() => + useToolsOAuthFlow({ accessToken: "user-token", serverId: "server-1", userId: "user-42", onSuccess }), + ); + + await waitFor(() => expect(result.current.status).toBe("success")); + + expect(setTokenSpy).toHaveBeenCalledWith( + "server-1", + { + access_token: "tok-tools", + expires_in: 3600, + refresh_token: "ref-tools", + token_type: "bearer", + }, + "user-42", + ); + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(onSuccess).toHaveBeenCalledWith("tok-tools"); + }); + + it("does not consume a result written under the user flow's key", async () => { + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue({ access_token: "x" }); + seedCompletedRedirect(USER_FLOW_STATE_KEY, USER_RESULT_KEY); + + const onSuccess = vi.fn(); + const { result } = renderHook(() => + useToolsOAuthFlow({ accessToken: "user-token", serverId: "server-1", onSuccess }), + ); + + // Give the on-mount resume a chance to (not) run. + await Promise.resolve(); + expect(networking.exchangeMcpOAuthToken).not.toHaveBeenCalled(); + expect(result.current.status).toBe("idle"); + }); + + it("stores the raw current URL as the post-redirect return target", async () => { + vi.mocked(networking.registerMcpOAuthClient).mockResolvedValue({ client_id: "c" }); + const { result } = renderHook(() => + useToolsOAuthFlow({ accessToken: "user-token", serverId: "server-1", onSuccess: vi.fn() }), + ); + + await act(async () => { + await result.current.startOAuthFlow(); + }); + + expect(getSecureItem("litellm-mcp-oauth-return-url")).toBe(window.location.href); + }); + }); + + describe("useUserMcpOAuthFlow", () => { + it("persists the exchanged token to the backend (with flow scopes) and never to sessionStorage", async () => { + const token = { access_token: "tok-user", refresh_token: "ref-user", expires_in: 1800 }; + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue(token); + const setTokenSpy = vi.spyOn(mcpTokenStore, "setToken"); + seedCompletedRedirect(USER_FLOW_STATE_KEY, USER_RESULT_KEY, { scopes: ["repo", "read:user"] }); + + const onSuccess = vi.fn(); + const { result } = renderHook(() => + useUserMcpOAuthFlow({ accessToken: "user-token", serverId: "server-1", onSuccess }), + ); + + await waitFor(() => expect(result.current.status).toBe("success")); + + expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledWith("user-token", "server-1", { + access_token: "tok-user", + refresh_token: "ref-user", + expires_in: 1800, + scopes: ["repo", "read:user"], + }); + expect(setTokenSpy).not.toHaveBeenCalled(); + expect(onSuccess).toHaveBeenCalledTimes(1); + }); + + it("does not consume a result written under the tools flow's key", async () => { + vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue({ access_token: "x" }); + seedCompletedRedirect(TOOLS_FLOW_STATE_KEY, TOOLS_RESULT_KEY); + + const { result } = renderHook(() => + useUserMcpOAuthFlow({ accessToken: "user-token", serverId: "server-1", onSuccess: vi.fn() }), + ); + + await Promise.resolve(); + expect(networking.exchangeMcpOAuthToken).not.toHaveBeenCalled(); + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(result.current.status).toBe("idle"); + }); + + it("tags the return URL with mcpOauthReturn=apps so the redirect lands back on the apps view", async () => { + vi.mocked(networking.registerMcpOAuthClient).mockResolvedValue({ client_id: "c" }); + const { result } = renderHook(() => + useUserMcpOAuthFlow({ accessToken: "user-token", serverId: "server-1", onSuccess: vi.fn() }), + ); + + await act(async () => { + await result.current.startOAuthFlow(); + }); + + const returnUrl = getSecureItem("litellm-mcp-oauth-return-url"); + expect(returnUrl).not.toBeNull(); + expect(new URL(returnUrl as string).searchParams.get("mcpOauthReturn")).toBe("apps"); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.ts b/ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.ts new file mode 100644 index 00000000000..fa476ed5e90 --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.ts @@ -0,0 +1,236 @@ +"use client"; + +/** + * Shared OAuth2 PKCE flow for MCP servers that already exist in the database. + * + * Both the Tools re-auth path (useToolsOAuthFlow) and the user connect path + * (useUserMcpOAuthFlow) run the identical authorize -> redirect -> exchange + * sequence; they diverge only in three places, injected via config: + * - the sessionStorage keys they read/write (so the flows don't collide), + * - how the post-redirect return URL is built, + * - what they do with the exchanged token (sessionStorage vs backend DB). + * + * The admin create-server flow (useMcpOAuthFlow) is deliberately NOT built on + * this: it targets a not-yet-persisted server and needs an extra temp-session + * cache step, so its logic is materially different. + */ + +import { useCallback, useEffect, useRef, useState } from "react"; +import { buildMcpOAuthAuthorizeUrl, exchangeMcpOAuthToken, registerMcpOAuthClient } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { extractErrorMessage } from "@/utils/errorUtils"; +import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils"; + +export type McpOAuthPkceStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; + +export interface McpOAuthFlowState { + state: string; + codeVerifier: string; + serverId: string; + redirectUri: string; + clientId?: string; + clientSecret?: string; + scopes?: string[]; +} + +export interface McpOAuthTokenResult { + access_token: string; + expires_in?: number; + refresh_token?: string; + token_type?: string; +} + +export interface McpOAuthStorageKeys { + flowState: string; + result: string; + returnUrl: string; +} + +interface UseMcpOAuthPkceFlowConfig { + accessToken: string; + serverId: string; + serverAlias?: string | null; + scopes?: string[]; + clientId?: string | null; + storageKeys: McpOAuthStorageKeys; + /** URL to return to after the OAuth redirect completes. */ + buildReturnUrl: () => string; + /** Persist the exchanged token (sessionStorage, backend DB, etc.). */ + persistToken: (token: McpOAuthTokenResult, flowState: McpOAuthFlowState) => Promise | void; + /** Invoked after the token is persisted so the caller can refresh UI state. */ + onSuccess: (accessToken: string) => void; +} + +interface UseMcpOAuthPkceFlowResult { + startOAuthFlow: () => Promise; + status: McpOAuthPkceStatus; + error: string | null; +} + +export const useMcpOAuthPkceFlow = ({ + accessToken, + serverId, + serverAlias, + scopes, + clientId: preClientId, + storageKeys, + buildReturnUrl, + persistToken, + onSuccess, +}: UseMcpOAuthPkceFlowConfig): UseMcpOAuthPkceFlowResult => { + const [status, setStatus] = useState("idle"); + const [error, setError] = useState(null); + const processingRef = useRef(false); + + const callbacksRef = useRef({ buildReturnUrl, persistToken, onSuccess }); + callbacksRef.current = { buildReturnUrl, persistToken, onSuccess }; + + const { flowState: FLOW_STATE_KEY, result: RESULT_KEY, returnUrl: RETURN_URL_KEY } = storageKeys; + + const startOAuthFlow = useCallback(async () => { + if (typeof window === "undefined") return; + try { + setStatus("authorizing"); + setError(null); + + let clientId: string | undefined = preClientId ?? undefined; + let clientSecret: string | undefined; + + if (!clientId) { + try { + const reg = await registerMcpOAuthClient(accessToken, serverId, { + client_name: serverAlias || serverId, + grant_types: ["authorization_code", "refresh_token"], + response_types: ["code"], + token_endpoint_auth_method: "none", + }); + clientId = reg?.client_id; + clientSecret = reg?.client_secret; + } catch (_) { + // Registration is optional; proceed without client_id + } + } + + const verifier = generateCodeVerifier(); + const challenge = await generateCodeChallenge(verifier); + const state = crypto.randomUUID(); + const redirectUri = buildCallbackUrl(); + const scopeString = scopes?.filter((s) => s.trim()).join(" "); + + const authorizeUrl = buildMcpOAuthAuthorizeUrl({ + serverId, + clientId, + redirectUri, + state, + codeChallenge: challenge, + scope: scopeString, + }); + + const flowState: McpOAuthFlowState = { + state, + codeVerifier: verifier, + serverId, + redirectUri, + clientId, + clientSecret, + scopes, + }; + + setSecureItem(FLOW_STATE_KEY, JSON.stringify(flowState)); + setSecureItem(RETURN_URL_KEY, callbacksRef.current.buildReturnUrl()); + + window.location.href = authorizeUrl; + } catch (err) { + const msg = extractErrorMessage(err); + setError(msg); + setStatus("error"); + NotificationsManager.error(msg); + } + }, [accessToken, serverId, serverAlias, scopes, preClientId, FLOW_STATE_KEY, RETURN_URL_KEY]); + + const resumeOAuthFlow = useCallback(async () => { + if (typeof window === "undefined" || processingRef.current) return; + + const storedResult = getSecureItem(RESULT_KEY); + if (!storedResult) return; + + // The callback page writes a result for every OAuth flow (Tools, user, admin). + // Peek at this hook's flow state first and bail unless it exists and its + // serverId matches: only the hook instance that initiated the flow should + // consume the result. Without this, a stale result from another flow would + // surface "OAuth session state was lost" on unrelated hook instances. + const rawFlowState = getSecureItem(FLOW_STATE_KEY); + if (!rawFlowState) return; + + let flowState: McpOAuthFlowState | null = null; + try { + flowState = JSON.parse(rawFlowState) as McpOAuthFlowState; + if (flowState.serverId && flowState.serverId !== serverId) return; + } catch (_) {} + + processingRef.current = true; + clearStorage(RESULT_KEY); + + let payload: Record | null = null; + try { + payload = JSON.parse(storedResult); + } catch (_) { + setError("Failed to resume OAuth flow. Please retry."); + setStatus("error"); + processingRef.current = false; + clearStorage(FLOW_STATE_KEY); + return; + } + + try { + if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) { + throw new Error("OAuth session state was lost. Please retry."); + } + if (!payload?.state || payload.state !== flowState.state) { + throw new Error("OAuth state mismatch. Please retry."); + } + if (payload.error) { + throw new Error((payload.error_description as string) || (payload.error as string)); + } + if (!payload.code) { + throw new Error("Authorization code missing in callback."); + } + + setStatus("exchanging"); + const token: McpOAuthTokenResult = await exchangeMcpOAuthToken({ + serverId: flowState.serverId, + code: payload.code as string, + clientId: flowState.clientId, + clientSecret: flowState.clientSecret, + codeVerifier: flowState.codeVerifier, + redirectUri: flowState.redirectUri, + accessToken, + }); + + await callbacksRef.current.persistToken(token, flowState); + + setStatus("success"); + setError(null); + NotificationsManager.success("Connected successfully"); + callbacksRef.current.onSuccess(token.access_token); + } catch (err) { + const msg = extractErrorMessage(err); + setError(msg); + setStatus("error"); + NotificationsManager.error(msg); + } finally { + clearStorage(FLOW_STATE_KEY); + setTimeout(() => { + processingRef.current = false; + }, 1000); + } + }, [accessToken, serverId, FLOW_STATE_KEY, RESULT_KEY]); + + useEffect(() => { + resumeOAuthFlow(); + }, [resumeOAuthFlow]); + + return { startOAuthFlow, status, error }; +}; diff --git a/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx index a538bb640c6..b07efe42e87 100644 --- a/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx @@ -3,7 +3,8 @@ /** * OAuth2 PKCE flow for the Tools screen re-authentication path. * - * Unlike useUserMcpOAuthFlow (used in the chat panel), this hook: + * Thin wrapper over useMcpOAuthPkceFlow. Unlike useUserMcpOAuthFlow (chat panel), + * this flow: * - stores the resulting token in sessionStorage via mcpTokenStore only * - does NOT call storeMCPOAuthUserCredential (no backend DB write) * - uses "litellm-tools-mcp-oauth-result" as its result key to avoid @@ -13,16 +14,10 @@ * to this key so this hook can pick up the result after the redirect. */ -import { useCallback, useEffect, useRef, useState } from "react"; -import { buildMcpOAuthAuthorizeUrl, exchangeMcpOAuthToken, registerMcpOAuthClient } from "@/components/networking"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { extractErrorMessage } from "@/utils/errorUtils"; -import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce"; -import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; import { setToken } from "@/utils/mcpTokenStore"; -import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils"; +import { McpOAuthStorageKeys, McpOAuthPkceStatus, useMcpOAuthPkceFlow } from "./useMcpOAuthPkceFlow"; -export type ToolsOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; +export type ToolsOAuthStatus = McpOAuthPkceStatus; interface UseToolsOAuthFlowOptions { accessToken: string; @@ -40,18 +35,10 @@ interface UseToolsOAuthFlowResult { error: string | null; } -const FLOW_STATE_KEY = "litellm-tools-mcp-oauth-flow-state"; -const RESULT_KEY = "litellm-tools-mcp-oauth-result"; -const RETURN_URL_KEY = "litellm-mcp-oauth-return-url"; - -type StoredFlowState = { - state: string; - codeVerifier: string; - serverId: string; - redirectUri: string; - clientId?: string; - clientSecret?: string; - scopes?: string[]; +const STORAGE_KEYS: McpOAuthStorageKeys = { + flowState: "litellm-tools-mcp-oauth-flow-state", + result: "litellm-tools-mcp-oauth-result", + returnUrl: "litellm-mcp-oauth-return-url", }; export const useToolsOAuthFlow = ({ @@ -60,140 +47,20 @@ export const useToolsOAuthFlow = ({ serverAlias, userId, scopes, - clientId: preClientId, + clientId, onSuccess, -}: UseToolsOAuthFlowOptions): UseToolsOAuthFlowResult => { - const [status, setStatus] = useState("idle"); - const [error, setError] = useState(null); - const processingRef = useRef(false); - const onSuccessRef = useRef(onSuccess); - onSuccessRef.current = onSuccess; - - const startOAuthFlow = useCallback(async () => { - if (typeof window === "undefined") return; - try { - setStatus("authorizing"); - setError(null); - - let clientId: string | undefined = preClientId ?? undefined; - let clientSecret: string | undefined; - - if (!clientId) { - try { - const reg = await registerMcpOAuthClient(accessToken, serverId, { - client_name: serverAlias || serverId, - grant_types: ["authorization_code", "refresh_token"], - response_types: ["code"], - token_endpoint_auth_method: "none", - }); - clientId = reg?.client_id; - clientSecret = reg?.client_secret; - } catch (_) { - // Registration is optional; proceed without client_id - } - } - - const verifier = generateCodeVerifier(); - const challenge = await generateCodeChallenge(verifier); - const state = crypto.randomUUID(); - const redirectUri = buildCallbackUrl(); - const scopeString = scopes?.filter((s) => s.trim()).join(" "); - - const authorizeUrl = buildMcpOAuthAuthorizeUrl({ - serverId, - clientId, - redirectUri, - state, - codeChallenge: challenge, - scope: scopeString, - }); - - const flowState: StoredFlowState = { - state, - codeVerifier: verifier, - serverId, - redirectUri, - clientId, - clientSecret, - scopes, - }; - - setSecureItem(FLOW_STATE_KEY, JSON.stringify(flowState)); - // Return to the current page (Tools tab) after the OAuth redirect - setSecureItem(RETURN_URL_KEY, window.location.href); - - window.location.href = authorizeUrl; - } catch (err) { - const msg = extractErrorMessage(err); - setError(msg); - setStatus("error"); - NotificationsManager.error(msg); - } - }, [accessToken, serverId, serverAlias, scopes, preClientId]); - - const resumeOAuthFlow = useCallback(async () => { - if (typeof window === "undefined" || processingRef.current) return; - - const storedResult = getSecureItem(RESULT_KEY); - if (!storedResult) return; - - // The callback page writes to this result key for every OAuth flow (including - // the admin server-creation flow). Guard: only proceed if *this* hook's flow - // state exists, meaning startOAuthFlow() was actually called from the Tools screen. - // Without this guard, a stale result written during server creation would trigger - // "OAuth session state was lost" when the user navigates to the Tools tab. - const rawFlowState = getSecureItem(FLOW_STATE_KEY); - if (!rawFlowState) return; - - let peeked: StoredFlowState | null = null; - try { - peeked = JSON.parse(rawFlowState) as StoredFlowState; - if (peeked.serverId && peeked.serverId !== serverId) return; - } catch (_) {} - - processingRef.current = true; - clearStorage(RESULT_KEY); - - let payload: Record | null = null; - let flowState: StoredFlowState | null = null; - - try { - payload = JSON.parse(storedResult); - flowState = peeked; - } catch (_) { - setError("Failed to resume OAuth flow. Please retry."); - setStatus("error"); - processingRef.current = false; - clearStorage(FLOW_STATE_KEY); - return; - } - - try { - if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) { - throw new Error("OAuth session state was lost. Please retry."); - } - if (!payload?.state || payload.state !== flowState.state) { - throw new Error("OAuth state mismatch. Please retry."); - } - if (payload.error) { - throw new Error((payload.error_description as string) || (payload.error as string)); - } - if (!payload.code) { - throw new Error("Authorization code missing in callback."); - } - - setStatus("exchanging"); - const token = await exchangeMcpOAuthToken({ - serverId: flowState.serverId, - code: payload.code as string, - clientId: flowState.clientId, - clientSecret: flowState.clientSecret, - codeVerifier: flowState.codeVerifier, - redirectUri: flowState.redirectUri, - accessToken, - }); - - // Store in sessionStorage only — no backend DB write +}: UseToolsOAuthFlowOptions): UseToolsOAuthFlowResult => + useMcpOAuthPkceFlow({ + accessToken, + serverId, + serverAlias, + scopes, + clientId, + storageKeys: STORAGE_KEYS, + // Return to the current page (Tools tab) after the OAuth redirect. + buildReturnUrl: () => window.location.href, + // Store in sessionStorage only — no backend DB write. + persistToken: (token, flowState) => setToken( flowState.serverId, { @@ -203,28 +70,6 @@ export const useToolsOAuthFlow = ({ token_type: token.token_type, }, userId, - ); - - setStatus("success"); - setError(null); - NotificationsManager.success("Connected successfully"); - onSuccessRef.current(token.access_token); - } catch (err) { - const msg = extractErrorMessage(err); - setError(msg); - setStatus("error"); - NotificationsManager.error(msg); - } finally { - clearStorage(FLOW_STATE_KEY); - setTimeout(() => { - processingRef.current = false; - }, 1000); - } - }, [accessToken, serverId, userId]); - - useEffect(() => { - resumeOAuthFlow(); - }, [resumeOAuthFlow]); - - return { startOAuthFlow, status, error }; -}; + ), + onSuccess, + }); diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx index 13bad227297..4dc4a3d3d43 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -3,29 +3,16 @@ /** * OAuth2 PKCE flow for the *user* connect path. * - * Unlike useMcpOAuthFlow (used in the admin create-server form), this hook - * targets a server that already exists in the database. It therefore skips - * the temp-session cache step and calls /server/oauth/{serverId}/authorize - * directly with the real server_id. - * - * On success it calls storeMCPOAuthUserCredential to persist the token for - * the user and then invokes onSuccess so the caller can refresh UI state. + * Thin wrapper over useMcpOAuthPkceFlow. Unlike useMcpOAuthFlow (admin + * create-server form), this targets a server that already exists in the + * database, so it skips the temp-session cache step. On success it calls + * storeMCPOAuthUserCredential to persist the token for the user. */ -import { useCallback, useEffect, useRef, useState } from "react"; -import { - buildMcpOAuthAuthorizeUrl, - exchangeMcpOAuthToken, - registerMcpOAuthClient, - storeMCPOAuthUserCredential, -} from "@/components/networking"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { extractErrorMessage } from "@/utils/errorUtils"; -import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce"; -import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; -import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils"; +import { storeMCPOAuthUserCredential } from "@/components/networking"; +import { McpOAuthStorageKeys, McpOAuthPkceStatus, useMcpOAuthPkceFlow } from "./useMcpOAuthPkceFlow"; -export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; +export type UserMcpOAuthStatus = McpOAuthPkceStatus; interface UseUserMcpOAuthFlowOptions { accessToken: string; @@ -44,28 +31,12 @@ interface UseUserMcpOAuthFlowResult { error: string | null; } -const FLOW_STATE_KEY = "litellm-user-mcp-oauth-flow-state"; -// Use a user-flow-specific key to avoid collisions with the admin OAuth flow -// (useMcpOAuthFlow) which uses "litellm-mcp-oauth-result". -const RESULT_KEY = "litellm-user-mcp-oauth-result"; -const RETURN_URL_KEY = "litellm-mcp-oauth-return-url"; - -type StoredFlowState = { - state: string; - codeVerifier: string; - serverId: string; - redirectUri: string; - clientId?: string; - clientSecret?: string; - scopes?: string[]; -}; - -const setStorage = (key: string, value: string) => { - setSecureItem(key, value); -}; - -const getStorage = (key: string): string | null => { - return getSecureItem(key); +const STORAGE_KEYS: McpOAuthStorageKeys = { + flowState: "litellm-user-mcp-oauth-flow-state", + // User-flow-specific keys to avoid collisions with the admin OAuth flow + // (useMcpOAuthFlow), which uses "litellm-mcp-oauth-result". + result: "litellm-user-mcp-oauth-result", + returnUrl: "litellm-mcp-oauth-return-url", }; export const useUserMcpOAuthFlow = ({ @@ -73,172 +44,29 @@ export const useUserMcpOAuthFlow = ({ serverId, serverAlias, scopes, - clientId: preClientId, + clientId, onSuccess, -}: UseUserMcpOAuthFlowOptions): UseUserMcpOAuthFlowResult => { - const [status, setStatus] = useState("idle"); - const [error, setError] = useState(null); - const processingRef = useRef(false); - - const startOAuthFlow = useCallback(async () => { - if (typeof window === "undefined") return; - try { - setStatus("authorizing"); - setError(null); - - let clientId: string | undefined = preClientId ?? undefined; - let clientSecret: string | undefined; - - if (!clientId) { - // Attempt dynamic client registration against the server's registration endpoint. - try { - const reg = await registerMcpOAuthClient(accessToken, serverId, { - client_name: serverAlias || serverId, - grant_types: ["authorization_code", "refresh_token"], - response_types: ["code"], - token_endpoint_auth_method: "none", - }); - clientId = reg?.client_id; - clientSecret = reg?.client_secret; - } catch (_) { - // Registration is optional; proceed without client_id - } - } - - const verifier = generateCodeVerifier(); - const challenge = await generateCodeChallenge(verifier); - const state = crypto.randomUUID(); - const redirectUri = buildCallbackUrl(); - const scopeString = scopes?.filter((s) => s.trim()).join(" "); - - const authorizeUrl = buildMcpOAuthAuthorizeUrl({ - serverId, - clientId, - redirectUri, - state, - codeChallenge: challenge, - scope: scopeString, - }); - - const flowState: StoredFlowState = { - state, - codeVerifier: verifier, - serverId, - redirectUri, - clientId, - clientSecret, - scopes, - }; - - setStorage(FLOW_STATE_KEY, JSON.stringify(flowState)); +}: UseUserMcpOAuthFlowOptions): UseUserMcpOAuthFlowResult => + useMcpOAuthPkceFlow({ + accessToken, + serverId, + serverAlias, + scopes, + clientId, + storageKeys: STORAGE_KEYS, + buildReturnUrl: () => { const returnUrl = new URL(window.location.href); returnUrl.searchParams.set("mcpOauthReturn", "apps"); - setStorage(RETURN_URL_KEY, returnUrl.toString()); - - window.location.href = authorizeUrl; - } catch (err) { - const msg = extractErrorMessage(err); - setError(msg); - setStatus("error"); - NotificationsManager.error(msg); - } - }, [accessToken, serverId, serverAlias, scopes, preClientId]); - - const resumeOAuthFlow = useCallback(async () => { - if (typeof window === "undefined" || processingRef.current) return; - - const storedResult = getStorage(RESULT_KEY); - if (!storedResult) return; - - // When multiple OAuth2ConnectButton components are mounted (one per server - // card), each holds its own hook instance. All run resumeOAuthFlow() on - // mount and would compete for the same RESULT_KEY. Peek at the stored - // flow state first: only the hook instance whose serverId matches the one - // that initiated the OAuth flow should consume the result. - // Guard: only proceed if this hook's flow state exists (startOAuthFlow was - // called from this hook). Without the guard, a tools re-auth redirect writes - // to the user result key too, and every OAuth2ConnectButton instance would try - // to resume a flow that was never started here. - const rawFlowState = getStorage(FLOW_STATE_KEY); - if (!rawFlowState) return; - - try { - const peeked = JSON.parse(rawFlowState) as StoredFlowState; - if (peeked.serverId && peeked.serverId !== serverId) return; - } catch (_) {} - - processingRef.current = true; - clearStorage(RESULT_KEY); - - let payload: Record | null = null; - let flowState: StoredFlowState | null = null; - - try { - payload = JSON.parse(storedResult); - const raw = getStorage(FLOW_STATE_KEY); - flowState = raw ? JSON.parse(raw) : null; - } catch (_) { - setError("Failed to resume OAuth flow. Please retry."); - setStatus("error"); - processingRef.current = false; - clearStorage(FLOW_STATE_KEY); - return; - } - - try { - if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) { - throw new Error("OAuth session state was lost. Please retry."); - } - if (!payload?.state || payload.state !== flowState.state) { - throw new Error("OAuth state mismatch. Please retry."); - } - if (payload.error) { - throw new Error((payload.error_description as string) || (payload.error as string)); - } - if (!payload.code) { - throw new Error("Authorization code missing in callback."); - } - - setStatus("exchanging"); - const token = await exchangeMcpOAuthToken({ - serverId: flowState.serverId, - code: payload.code as string, - clientId: flowState.clientId, - clientSecret: flowState.clientSecret, - codeVerifier: flowState.codeVerifier, - redirectUri: flowState.redirectUri, - accessToken, - }); - - // Persist the token for this user via the backend. - // accessToken comes from props — it is never stored in sessionStorage. - await storeMCPOAuthUserCredential(accessToken, flowState.serverId, { + return returnUrl.toString(); + }, + // Persist the token for this user via the backend. + // accessToken comes from props — it is never stored in sessionStorage. + persistToken: (token, flowState) => + storeMCPOAuthUserCredential(accessToken, flowState.serverId, { access_token: token.access_token, refresh_token: token.refresh_token, expires_in: token.expires_in, scopes: flowState.scopes, - }); - - setStatus("success"); - setError(null); - NotificationsManager.success("Connected successfully"); - onSuccess(); - } catch (err) { - const msg = extractErrorMessage(err); - setError(msg); - setStatus("error"); - NotificationsManager.error(msg); - } finally { - clearStorage(FLOW_STATE_KEY); - setTimeout(() => { - processingRef.current = false; - }, 1000); - } - }, [accessToken, serverId, onSuccess]); - - useEffect(() => { - resumeOAuthFlow(); - }, [resumeOAuthFlow]); - - return { startOAuthFlow, status, error }; -}; + }), + onSuccess, + });