Merge pull request #35835 from BerriAI/litellm_/elated-margulis-7f300f

refactor(ui): route MCP session tokens through the shared storage helper

(cherry picked from commit e4fd790f1c)
This commit is contained in:
yuneng-jiang 2026-08-04 15:03:29 -07:00 committed by Yuneng Jiang
parent 065bf83f0c
commit 2f2e477207
No known key found for this signature in database
5 changed files with 41 additions and 9 deletions

View file

@ -610,7 +610,6 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
const browserHeldToken = {
access_token: oauthTokenResponse.access_token,
expires_in: oauthTokenResponse.expires_in,
refresh_token: oauthTokenResponse.refresh_token,
token_type: oauthTokenResponse.token_type,
};
setToken(response.server_id, browserHeldToken, userID);

View file

@ -213,7 +213,6 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
const browserHeldToken = {
access_token: token.access_token,
expires_in: token.expires_in,
refresh_token: token.refresh_token,
token_type: token.token_type,
};
setToken(mcpServer.server_id, browserHeldToken, userID);
@ -969,7 +968,6 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
const browserHeldToken = {
access_token: oauthTokenResponse.access_token,
expires_in: oauthTokenResponse.expires_in,
refresh_token: oauthTokenResponse.refresh_token,
token_type: oauthTokenResponse.token_type,
};
setToken(mcpServer.server_id, browserHeldToken, userID);

View file

@ -199,7 +199,6 @@ export const useToolsOAuthFlow = ({
{
access_token: token.access_token,
expires_in: token.expires_in,
refresh_token: token.refresh_token,
token_type: token.token_type,
},
userId,

View file

@ -1,6 +1,20 @@
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { clearAllMcpTokens, getToken, isTokenValid, removeToken, setToken } from "./mcpTokenStore";
const decodeMaybeBase64 = (raw: string): string => {
try {
return atob(raw);
} catch {
return raw;
}
};
const allStoredValues = (): string =>
Array.from({ length: sessionStorage.length }, (_, i) => sessionStorage.key(i) ?? "")
.map((key) => sessionStorage.getItem(key) ?? "")
.flatMap((raw) => [raw, decodeMaybeBase64(raw)])
.join("\n");
describe("mcpTokenStore", () => {
beforeEach(() => {
sessionStorage.clear();
@ -10,6 +24,29 @@ describe("mcpTokenStore", () => {
sessionStorage.clear();
});
it("never persists a refresh token, even when a caller supplies one", () => {
const callerPayload = {
access_token: "access-value",
expires_in: 3600,
refresh_token: "refresh-value-must-not-persist",
token_type: "bearer",
};
setToken("server-a", callerPayload, "user-1");
expect(getToken("server-a", "user-1")?.access_token).toBe("access-value");
expect(allStoredValues()).not.toContain("refresh-value-must-not-persist");
});
it("does not write the token payload as readable text", () => {
setToken("server-a", { access_token: "plain-access-value" }, "user-1");
const raw = sessionStorage.getItem("mcp-session-token:user-1:server-a");
expect(raw).not.toBeNull();
expect(raw).not.toContain("plain-access-value");
expect(getToken("server-a", "user-1")?.access_token).toBe("plain-access-value");
});
it("scopes tokens by user id", () => {
setToken("server-a", { access_token: "user1-token" }, "user-1");
setToken("server-a", { access_token: "user2-token" }, "user-2");

View file

@ -4,19 +4,19 @@
* session ends (tab/window close). Never written to localStorage.
*/
import { getSecureItem, setSecureItem } from "./secureStorage";
const KEY_PREFIX = "mcp-session-token:";
interface StoredToken {
access_token: string;
expires_at: number;
refresh_token?: string;
token_type: string;
}
interface TokenInput {
access_token: string;
expires_in?: number;
refresh_token?: string;
token_type?: string;
}
@ -33,10 +33,9 @@ export function setToken(serverId: string, data: TokenInput, userId?: string | n
access_token: data.access_token,
expires_at: Date.now() + (data.expires_in != null ? data.expires_in * 1000 : DEFAULT_TTL_MS),
token_type: data.token_type ?? "bearer",
...(data.refresh_token ? { refresh_token: data.refresh_token } : {}),
};
try {
window.sessionStorage.setItem(storageKey(serverId, userId), JSON.stringify(stored));
setSecureItem(storageKey(serverId, userId), JSON.stringify(stored));
} catch {
// Silently ignore storage errors (private browsing, quota exceeded, etc.)
}
@ -45,7 +44,7 @@ export function setToken(serverId: string, data: TokenInput, userId?: string | n
export function getToken(serverId: string, userId?: string | null): StoredToken | null {
if (typeof window === "undefined") return null;
try {
const raw = window.sessionStorage.getItem(storageKey(serverId, userId));
const raw = getSecureItem(storageKey(serverId, userId));
if (!raw) return null;
return JSON.parse(raw) as StoredToken;
} catch {