Add Refresh Token Support for OAuth

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) {

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<OAuthMetadata> {
@@ -60,7 +66,7 @@ export async function startOAuthFlow(serverUrl: string): Promise<string> {
 export async function handleOAuthCallback(
   serverUrl: string,
   code: string,
-): Promise<string> {
+): Promise<OAuthTokens> {
   // 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<OAuthTokens> {
+  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;
 }

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;

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,7 +121,46 @@ export function useConnection({
     }
   };
 
-  const connect = async () => {
+  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);
+      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);
+      await initiateOAuthFlow();
+      throw error;
+    }
+  };
+
+  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<Request, Notification, Result>(
         {
@@ -182,14 +221,15 @@ 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) {
-          // Store the server URL for the callback handler
-          sessionStorage.setItem(SESSION_KEYS.SERVER_URL, sseUrl);
-          const redirectUrl = await startOAuthFlow(sseUrl);
-          window.location.href = redirectUrl;
+          // Don't set error state if we're about to redirect for auth
           return;
         }
-
         throw error;
       }