From 79547143a832175bfb8856fb7919d6eb3f0d1b69 Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Mon, 3 Feb 2025 19:53:53 -0800 Subject: [PATCH 01/10] Add refresh token handling if returned from server --- client/src/components/OAuthCallback.tsx | 9 ++++-- client/src/lib/auth.ts | 38 +++++++++++++++++++++++-- client/src/lib/constants.ts | 1 + client/src/lib/hooks/useConnection.ts | 34 ++++++++++++++++++++-- 4 files changed, 74 insertions(+), 8 deletions(-) diff --git a/client/src/components/OAuthCallback.tsx b/client/src/components/OAuthCallback.tsx index a7439df..2a9e27a 100644 --- a/client/src/components/OAuthCallback.tsx +++ b/client/src/components/OAuthCallback.tsx @@ -24,9 +24,12 @@ const OAuthCallback = () => { } try { - const accessToken = await handleOAuthCallback(serverUrl, code); - // Store the access token for future use - sessionStorage.setItem(SESSION_KEYS.ACCESS_TOKEN, accessToken); + const tokens = await handleOAuthCallback(serverUrl, code); + // Store both access and refresh tokens + sessionStorage.setItem(SESSION_KEYS.ACCESS_TOKEN, tokens.access_token); + if (tokens.refresh_token) { + sessionStorage.setItem(SESSION_KEYS.REFRESH_TOKEN, tokens.refresh_token); + } // Redirect back to the main app with server URL to trigger auto-connect window.location.href = `/?serverUrl=${encodeURIComponent(serverUrl)}`; } catch (error) { diff --git a/client/src/lib/auth.ts b/client/src/lib/auth.ts index 0417731..7d70a31 100644 --- a/client/src/lib/auth.ts +++ b/client/src/lib/auth.ts @@ -6,6 +6,12 @@ export interface OAuthMetadata { token_endpoint: string; } +export interface OAuthTokens { + access_token: string; + refresh_token?: string; + expires_in?: number; +} + export async function discoverOAuthMetadata( serverUrl: string, ): Promise { @@ -60,7 +66,7 @@ export async function startOAuthFlow(serverUrl: string): Promise { export async function handleOAuthCallback( serverUrl: string, code: string, -): Promise { +): Promise { // Get stored code verifier const codeVerifier = sessionStorage.getItem(SESSION_KEYS.CODE_VERIFIER); if (!codeVerifier) { @@ -69,7 +75,6 @@ export async function handleOAuthCallback( // Discover OAuth endpoints const metadata = await discoverOAuthMetadata(serverUrl); - // Exchange code for tokens const response = await fetch(metadata.token_endpoint, { method: "POST", @@ -89,5 +94,32 @@ export async function handleOAuthCallback( } const data = await response.json(); - return data.access_token; + return data; +} + +export async function refreshAccessToken(serverUrl: string): Promise { + const refreshToken = sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN); + if (!refreshToken) { + throw new Error("No refresh token available"); + } + + const metadata = await discoverOAuthMetadata(serverUrl); + + const response = await fetch(metadata.token_endpoint, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + grant_type: "refresh_token", + refresh_token: refreshToken + }), + }); + + if (!response.ok) { + throw new Error("Token refresh failed"); + } + + const data = await response.json(); + return data; } diff --git a/client/src/lib/constants.ts b/client/src/lib/constants.ts index e302b52..13a2370 100644 --- a/client/src/lib/constants.ts +++ b/client/src/lib/constants.ts @@ -3,4 +3,5 @@ export const SESSION_KEYS = { CODE_VERIFIER: "mcp_code_verifier", SERVER_URL: "mcp_server_url", ACCESS_TOKEN: "mcp_access_token", + REFRESH_TOKEN: "mcp_refresh_token", } as const; diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index de2d29e..58ea0a8 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -16,7 +16,7 @@ import { import { useState } from "react"; import { toast } from "react-toastify"; import { z } from "zod"; -import { startOAuthFlow } from "../auth"; +import { startOAuthFlow, refreshAccessToken } from "../auth"; import { SESSION_KEYS } from "../constants"; import { Notification, StdErrNotificationSchema } from "../notificationTypes"; @@ -121,6 +121,24 @@ export function useConnection({ } }; + const handleTokenRefresh = async () => { + try { + const tokens = await refreshAccessToken(sseUrl); + sessionStorage.setItem(SESSION_KEYS.ACCESS_TOKEN, tokens.access_token); + if (tokens.refresh_token) { + sessionStorage.setItem(SESSION_KEYS.REFRESH_TOKEN, tokens.refresh_token); + } + return tokens.access_token; + } catch (error) { + console.error("Token refresh failed:", error); + // Clear tokens and redirect to home + sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN); + sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN); + window.location.href = "/"; + throw error; + } + }; + const connect = async () => { try { const client = new Client( @@ -157,7 +175,19 @@ export function useConnection({ const clientTransport = new SSEClientTransport(backendUrl, { eventSourceInit: { - fetch: (url, init) => fetch(url, { ...init, headers }), + fetch: async (url, init) => { + const response = await fetch(url, { ...init, headers }); + + if (response.status === 401 && sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { + // Try to refresh the token + const newAccessToken = await handleTokenRefresh(); + headers["Authorization"] = `Bearer ${newAccessToken}`; + // Retry the request with new token + return fetch(url, { ...init, headers }); + } + + return response; + }, }, requestInit: { headers, From 4c89aed4d9f64f4a62d30b49a6ea441188909d77 Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Mon, 3 Feb 2025 20:04:17 -0800 Subject: [PATCH 02/10] Add check for expired refresh or session token that exists --- client/src/lib/hooks/useConnection.ts | 34 ++++++++++++++++++++------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index 58ea0a8..e08e0ae 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -131,10 +131,13 @@ export function useConnection({ return tokens.access_token; } catch (error) { console.error("Token refresh failed:", error); - // Clear tokens and redirect to home + // If refresh token is expired/invalid (401) or any other error, + // clear tokens and redirect to home to trigger re-authentication sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN); sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN); - window.location.href = "/"; + sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); + const redirectUrl = await startOAuthFlow(sseUrl); + window.location.href = redirectUrl; throw error; } }; @@ -178,12 +181,27 @@ export function useConnection({ fetch: async (url, init) => { const response = await fetch(url, { ...init, headers }); - if (response.status === 401 && sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { - // Try to refresh the token - const newAccessToken = await handleTokenRefresh(); - headers["Authorization"] = `Bearer ${newAccessToken}`; - // Retry the request with new token - return fetch(url, { ...init, headers }); + if (response.status === 401) { + // First try to refresh if we have a refresh token + if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { + try { + const newAccessToken = await handleTokenRefresh(); + headers["Authorization"] = `Bearer ${newAccessToken}`; + // Retry the request with new token + return fetch(url, { ...init, headers }); + } catch (error) { + console.error("Token refresh failed:", error); + } + } + + // If we have an access token but refresh failed or wasn't available, + // we need to re-authenticate since the token is invalid + if (sessionStorage.getItem(SESSION_KEYS.ACCESS_TOKEN)) { + sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); + const redirectUrl = await startOAuthFlow(sseUrl); + window.location.href = redirectUrl; + return new Response(); // This won't actually be used due to redirect + } } return response; From 7957d9f577aeed88844c843977dbb1dc6c64912d Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Mon, 3 Feb 2025 20:06:21 -0800 Subject: [PATCH 03/10] Make OAuth start call modular --- client/src/lib/hooks/useConnection.ts | 32 ++++++++++----------------- 1 file changed, 12 insertions(+), 20 deletions(-) diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index e08e0ae..9ceafa2 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -121,6 +121,14 @@ export function useConnection({ } }; + const initiateOAuthFlow = async () => { + sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN); + sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN); + sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); + const redirectUrl = await startOAuthFlow(sseUrl); + window.location.href = redirectUrl; + }; + const handleTokenRefresh = async () => { try { const tokens = await refreshAccessToken(sseUrl); @@ -131,13 +139,7 @@ export function useConnection({ return tokens.access_token; } catch (error) { console.error("Token refresh failed:", error); - // If refresh token is expired/invalid (401) or any other error, - // clear tokens and redirect to home to trigger re-authentication - sessionStorage.removeItem(SESSION_KEYS.ACCESS_TOKEN); - sessionStorage.removeItem(SESSION_KEYS.REFRESH_TOKEN); - sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); - const redirectUrl = await startOAuthFlow(sseUrl); - window.location.href = redirectUrl; + await initiateOAuthFlow(); throw error; } }; @@ -182,25 +184,19 @@ export function useConnection({ const response = await fetch(url, { ...init, headers }); if (response.status === 401) { - // First try to refresh if we have a refresh token if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { try { const newAccessToken = await handleTokenRefresh(); headers["Authorization"] = `Bearer ${newAccessToken}`; - // Retry the request with new token return fetch(url, { ...init, headers }); } catch (error) { console.error("Token refresh failed:", error); } } - // If we have an access token but refresh failed or wasn't available, - // we need to re-authenticate since the token is invalid if (sessionStorage.getItem(SESSION_KEYS.ACCESS_TOKEN)) { - sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); - const redirectUrl = await startOAuthFlow(sseUrl); - window.location.href = redirectUrl; - return new Response(); // This won't actually be used due to redirect + await initiateOAuthFlow(); + return new Response(); } } @@ -231,13 +227,9 @@ export function useConnection({ } catch (error) { console.error("Failed to connect to MCP server:", error); if (error instanceof SseError && error.code === 401) { - // Store the server URL for the callback handler - sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl); - const redirectUrl = await startOAuthFlow(sseUrl); - window.location.href = redirectUrl; + await initiateOAuthFlow(); return; } - throw error; } From b5762d53fd6c543524cf5d9d867a253b60c0d12b Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Tue, 4 Feb 2025 14:53:41 -0800 Subject: [PATCH 04/10] Handle infinite loop if server keeps returning 401 --- client/src/lib/hooks/useConnection.ts | 51 ++++++++++++++------------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index 9ceafa2..cd90d7f 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -144,7 +144,23 @@ export function useConnection({ } }; - const connect = async () => { + const handleAuthError = async (error: unknown) => { + if (error instanceof SseError && error.code === 401) { + if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { + try { + await handleTokenRefresh(); + return true; + } catch (error) { + console.error("Token refresh failed:", error); + } + } else { + await initiateOAuthFlow(); + } + } + return false; + }; + + const connect = async (_e?: unknown, retryCount: number = 0) => { try { const client = new Client( { @@ -180,28 +196,7 @@ export function useConnection({ const clientTransport = new SSEClientTransport(backendUrl, { eventSourceInit: { - fetch: async (url, init) => { - const response = await fetch(url, { ...init, headers }); - - if (response.status === 401) { - if (sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN)) { - try { - const newAccessToken = await handleTokenRefresh(); - headers["Authorization"] = `Bearer ${newAccessToken}`; - return fetch(url, { ...init, headers }); - } catch (error) { - console.error("Token refresh failed:", error); - } - } - - if (sessionStorage.getItem(SESSION_KEYS.ACCESS_TOKEN)) { - await initiateOAuthFlow(); - return new Response(); - } - } - - return response; - }, + fetch: (url, init) => fetch(url, { ...init, headers }), }, requestInit: { headers, @@ -226,11 +221,17 @@ export function useConnection({ await client.connect(clientTransport); } catch (error) { console.error("Failed to connect to MCP server:", error); + const shouldRetry = await handleAuthError(error); + if (shouldRetry) { + return connect(undefined, retryCount + 1); + } + if (error instanceof SseError && error.code === 401) { - await initiateOAuthFlow(); + // Don't set error state if we're about to redirect for auth return; } - throw error; + setConnectionStatus("error"); + return; } const capabilities = client.getServerCapabilities(); From b4ae1327b51887cf031f7900d7258bc81f4e8e99 Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Tue, 4 Feb 2025 15:00:14 -0800 Subject: [PATCH 05/10] Update useConnection.ts --- client/src/lib/hooks/useConnection.ts | 6 ------ 1 file changed, 6 deletions(-) diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index cd90d7f..0e843b8 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -225,12 +225,6 @@ export function useConnection({ if (shouldRetry) { return connect(undefined, retryCount + 1); } - - if (error instanceof SseError && error.code === 401) { - // Don't set error state if we're about to redirect for auth - return; - } - setConnectionStatus("error"); return; } From dd47b574b3ee897e463b700baffb7c2c3a408512 Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Tue, 4 Feb 2025 15:02:12 -0800 Subject: [PATCH 06/10] Update useConnection.ts --- client/src/lib/hooks/useConnection.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index 0e843b8..7225184 100644 --- a/client/src/lib/hooks/useConnection.ts +++ b/client/src/lib/hooks/useConnection.ts @@ -225,7 +225,12 @@ export function useConnection({ if (shouldRetry) { return connect(undefined, retryCount + 1); } - return; + + if (error instanceof SseError && error.code === 401) { + // Don't set error state if we're about to redirect for auth + return; + } + throw error; } const capabilities = client.getServerCapabilities(); From 8592cf2d0716063e8e422215f8a610d342bf31ec Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Wed, 5 Feb 2025 11:22:11 -0800 Subject: [PATCH 07/10] Run prettier-fix --- client/src/components/OAuthCallback.tsx | 5 ++++- client/src/lib/auth.ts | 8 +++++--- client/src/lib/hooks/useConnection.ts | 5 ++++- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/client/src/components/OAuthCallback.tsx b/client/src/components/OAuthCallback.tsx index 2a9e27a..869eef1 100644 --- a/client/src/components/OAuthCallback.tsx +++ b/client/src/components/OAuthCallback.tsx @@ -28,7 +28,10 @@ const OAuthCallback = () => { // Store both access and refresh tokens sessionStorage.setItem(SESSION_KEYS.ACCESS_TOKEN, tokens.access_token); if (tokens.refresh_token) { - sessionStorage.setItem(SESSION_KEYS.REFRESH_TOKEN, tokens.refresh_token); + sessionStorage.setItem( + SESSION_KEYS.REFRESH_TOKEN, + tokens.refresh_token, + ); } // Redirect back to the main app with server URL to trigger auto-connect window.location.href = `/?serverUrl=${encodeURIComponent(serverUrl)}`; diff --git a/client/src/lib/auth.ts b/client/src/lib/auth.ts index 7d70a31..918ac4d 100644 --- a/client/src/lib/auth.ts +++ b/client/src/lib/auth.ts @@ -97,14 +97,16 @@ export async function handleOAuthCallback( return data; } -export async function refreshAccessToken(serverUrl: string): Promise { +export async function refreshAccessToken( + serverUrl: string, +): Promise { const refreshToken = sessionStorage.getItem(SESSION_KEYS.REFRESH_TOKEN); if (!refreshToken) { throw new Error("No refresh token available"); } const metadata = await discoverOAuthMetadata(serverUrl); - + const response = await fetch(metadata.token_endpoint, { method: "POST", headers: { @@ -112,7 +114,7 @@ export async function refreshAccessToken(serverUrl: string): Promise Date: Wed, 5 Feb 2025 12:38:26 -0800 Subject: [PATCH 08/10] Convert OAuthMetadata and OAuthTokens to zod --- client/src/lib/auth.ts | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/client/src/lib/auth.ts b/client/src/lib/auth.ts index 918ac4d..040d35a 100644 --- a/client/src/lib/auth.ts +++ b/client/src/lib/auth.ts @@ -1,16 +1,21 @@ import pkceChallenge from "pkce-challenge"; import { SESSION_KEYS } from "./constants"; +import { z } from "zod"; -export interface OAuthMetadata { - authorization_endpoint: string; - token_endpoint: string; -} +export const OAuthMetadataSchema = z.object({ + authorization_endpoint: z.string(), + token_endpoint: z.string() +}); -export interface OAuthTokens { - access_token: string; - refresh_token?: string; - expires_in?: number; -} +export type OAuthMetadata = z.infer; + +export const OAuthTokensSchema = z.object({ + access_token: z.string(), + refresh_token: z.string().optional(), + expires_in: z.number().optional() +}); + +export type OAuthTokens = z.infer; export async function discoverOAuthMetadata( serverUrl: string, @@ -93,8 +98,7 @@ export async function handleOAuthCallback( throw new Error("Token exchange failed"); } - const data = await response.json(); - return data; + return await response.json(); } export async function refreshAccessToken( @@ -122,6 +126,5 @@ export async function refreshAccessToken( throw new Error("Token refresh failed"); } - const data = await response.json(); - return data; + return await response.json(); } From 95bbd60a38f4f34b7f5ae40146bee7488c5383c4 Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Wed, 5 Feb 2025 12:42:09 -0800 Subject: [PATCH 09/10] Add zod parsing for OAuthMetadataSchema and OAuthTokensSchema --- client/src/lib/auth.ts | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/client/src/lib/auth.ts b/client/src/lib/auth.ts index 040d35a..d7ddc8c 100644 --- a/client/src/lib/auth.ts +++ b/client/src/lib/auth.ts @@ -26,10 +26,11 @@ export async function discoverOAuthMetadata( if (response.ok) { const metadata = await response.json(); - return { + const validatedMetadata = OAuthMetadataSchema.parse({ authorization_endpoint: metadata.authorization_endpoint, token_endpoint: metadata.token_endpoint, - }; + }); + return validatedMetadata; } } catch (error) { console.warn("OAuth metadata discovery failed:", error); @@ -37,10 +38,11 @@ export async function discoverOAuthMetadata( // Fall back to default endpoints const baseUrl = new URL(serverUrl); - return { + const defaultMetadata = { authorization_endpoint: new URL("/authorize", baseUrl).toString(), token_endpoint: new URL("/token", baseUrl).toString(), }; + return OAuthMetadataSchema.parse(defaultMetadata); } export async function startOAuthFlow(serverUrl: string): Promise { @@ -98,7 +100,8 @@ export async function handleOAuthCallback( throw new Error("Token exchange failed"); } - return await response.json(); + const tokens = await response.json(); + return OAuthTokensSchema.parse(tokens); } export async function refreshAccessToken( @@ -126,5 +129,6 @@ export async function refreshAccessToken( throw new Error("Token refresh failed"); } - return await response.json(); + const tokens = await response.json(); + return OAuthTokensSchema.parse(tokens); } From 1b13b574f8aafd96887823d6dfa7aaddde5d9b72 Mon Sep 17 00:00:00 2001 From: Allen Zhou <46854522+allenzhou101@users.noreply.github.com> Date: Wed, 5 Feb 2025 12:45:11 -0800 Subject: [PATCH 10/10] Update auth.ts --- client/src/lib/auth.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/client/src/lib/auth.ts b/client/src/lib/auth.ts index d7ddc8c..592dc17 100644 --- a/client/src/lib/auth.ts +++ b/client/src/lib/auth.ts @@ -4,7 +4,7 @@ import { z } from "zod"; export const OAuthMetadataSchema = z.object({ authorization_endpoint: z.string(), - token_endpoint: z.string() + token_endpoint: z.string(), }); export type OAuthMetadata = z.infer; @@ -12,7 +12,7 @@ export type OAuthMetadata = z.infer; export const OAuthTokensSchema = z.object({ access_token: z.string(), refresh_token: z.string().optional(), - expires_in: z.number().optional() + expires_in: z.number().optional(), }); export type OAuthTokens = z.infer;