scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

refactor(ui): extract shared MCP OAuth PKCE flow hook

Browse source 
useToolsOAuthFlow and useUserMcpOAuthFlow ran the identical authorize
-> redirect -> exchange PKCE sequence and diverged only in their
sessionStorage keys, how the return URL is built, and where the
exchanged token is persisted (sessionStorage vs backend). Pull the
shared logic into useMcpOAuthPkceFlow and reduce both hooks to thin
configs, removing ~150 lines of duplication.

The admin create-server flow (useMcpOAuthFlow) is intentionally left
separate; it targets a not-yet-persisted server and needs an extra
temp-session cache step.

Adds useMcpOAuthPkceFlow.test.tsx pinning the three divergence points
per wrapper (token persistence target, return-URL tagging, and the
per-flow storage-key isolation guard); these hooks previously had no
direct coverage.
This commit is contained in:
ryan-crabbe-berri 2026-07-08 14:25:29 -07:00
parent 4b0ac8b352
commit ae32cc9803
6 changed files with 455 additions and 391 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
ui/litellm-dashboard/eslint-metrics.json
View file

@ -1,6 +1,6 @@
{
"@typescript-eslint/no-explicit-any": 1990,
"complexity": 128,
"complexity": 127,
"max-depth": 59,
"no-console": 15
}

9
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -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
}

160
ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.test.tsx Normal file
View file

@ -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<string, unknown> = {}) {
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");
});
});
});

236
ui/litellm-dashboard/src/hooks/useMcpOAuthPkceFlow.ts Normal file
View file

@ -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> | void;
/** Invoked after the token is persisted so the caller can refresh UI state. */
onSuccess: (accessToken: string) => void;
}
interface UseMcpOAuthPkceFlowResult {
startOAuthFlow: () => Promise<void>;
status: McpOAuthPkceStatus;
error: string | null;
}
export const useMcpOAuthPkceFlow = ({
accessToken,
serverId,
serverAlias,
scopes,
clientId: preClientId,
storageKeys,
buildReturnUrl,
persistToken,
onSuccess,
}: UseMcpOAuthPkceFlowConfig): UseMcpOAuthPkceFlowResult => {
const [status, setStatus] = useState<McpOAuthPkceStatus>("idle");
const [error, setError] = useState<string | null>(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<string, unknown> | 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 };
};

203
ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx
View file

@ -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<ToolsOAuthStatus>("idle");
const [error, setError] = useState<string | null>(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<string, unknown> | 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,
});

236
ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
View file

@ -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<UserMcpOAuthStatus>("idle");
const [error, setError] = useState<string | null>(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<string, unknown> | 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,
});