From 0dbf75a66427f1f430095d567595e373fcaf185b Mon Sep 17 00:00:00 2001
From: Paul Carleton <paulc@anthropic.com>
Date: Thu, 29 May 2025 11:50:09 +0100
Subject: [PATCH] store state before redirect

---
 client/src/App.tsx                           | 31 +++++++++++++++-----
 client/src/components/AuthDebugger.tsx       |  6 ++++
 client/src/components/OAuthDebugCallback.tsx | 19 ++++++++++--
 client/src/lib/constants.ts                  |  1 +
 4 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/client/src/App.tsx b/client/src/App.tsx
index 55a4541e..8be23a5b 100644
--- a/client/src/App.tsx
+++ b/client/src/App.tsx
@@ -234,21 +234,36 @@ const App = () => {
     ({
       authorizationCode,
       errorMsg,
+      restoredState,
     }: {
       authorizationCode?: string;
       errorMsg?: string;
+      restoredState?: AuthDebuggerState;
     }) => {
       setIsAuthDebuggerVisible(true);
-      if (authorizationCode) {
+      
+      if (restoredState) {
+        // Restore the previous auth state
         updateAuthState({
-          authorizationCode,
-          oauthStep: "token_request",
-        });
-      }
-      if (errorMsg) {
-        updateAuthState({
-          latestError: new Error(errorMsg),
+          ...restoredState,
+          // Update with the new authorization code if provided
+          authorizationCode: authorizationCode || restoredState.authorizationCode,
+          oauthStep: authorizationCode ? "token_request" : restoredState.oauthStep,
+          latestError: errorMsg ? new Error(errorMsg) : restoredState.latestError,
         });
+      } else {
+        // Fallback to the original behavior if no state was restored
+        if (authorizationCode) {
+          updateAuthState({
+            authorizationCode,
+            oauthStep: "token_request",
+          });
+        }
+        if (errorMsg) {
+          updateAuthState({
+            latestError: new Error(errorMsg),
+          });
+        }
       }
     },
     [],
diff --git a/client/src/components/AuthDebugger.tsx b/client/src/components/AuthDebugger.tsx
index 5565a77a..0ef8e70b 100644
--- a/client/src/components/AuthDebugger.tsx
+++ b/client/src/components/AuthDebugger.tsx
@@ -5,6 +5,7 @@ import { AlertCircle } from "lucide-react";
 import { AuthDebuggerState, EMPTY_DEBUGGER_STATE } from "../lib/auth-types";
 import { OAuthFlowProgress } from "./OAuthFlowProgress";
 import { OAuthStateMachine } from "../lib/oauth-state-machine";
+import { SESSION_KEYS } from "../lib/constants";
 
 export interface AuthDebuggerProps {
   serverUrl: string;
@@ -141,6 +142,11 @@ const AuthDebugger = ({
           currentState.oauthStep === "authorization_code" &&
           currentState.authorizationUrl
         ) {
+          // Store the current auth state before redirecting
+          sessionStorage.setItem(
+            SESSION_KEYS.AUTH_DEBUGGER_STATE,
+            JSON.stringify(currentState)
+          );
           // Open the authorization URL automatically
           window.location.href = currentState.authorizationUrl;
           break;
diff --git a/client/src/components/OAuthDebugCallback.tsx b/client/src/components/OAuthDebugCallback.tsx
index 88d931c0..af9944ac 100644
--- a/client/src/components/OAuthDebugCallback.tsx
+++ b/client/src/components/OAuthDebugCallback.tsx
@@ -9,9 +9,11 @@ interface OAuthCallbackProps {
   onConnect: ({
     authorizationCode,
     errorMsg,
+    restoredState,
   }: {
     authorizationCode?: string;
     errorMsg?: string;
+    restoredState?: any;
   }) => void;
 }
 
@@ -34,6 +36,19 @@ const OAuthDebugCallback = ({ onConnect }: OAuthCallbackProps) => {
       }
 
       const serverUrl = sessionStorage.getItem(SESSION_KEYS.SERVER_URL);
+      
+      // Try to restore the auth state
+      const storedState = sessionStorage.getItem(SESSION_KEYS.AUTH_DEBUGGER_STATE);
+      let restoredState = null;
+      if (storedState) {
+        try {
+          restoredState = JSON.parse(storedState);
+          // Clean up the stored state
+          sessionStorage.removeItem(SESSION_KEYS.AUTH_DEBUGGER_STATE);
+        } catch (e) {
+          console.error("Failed to parse stored auth state:", e);
+        }
+      }
 
       // ServerURL isn't set, this can happen if we've opened the
       // authentication request in a new tab, so we don't have the same
@@ -50,8 +65,8 @@ const OAuthDebugCallback = ({ onConnect }: OAuthCallbackProps) => {
       }
 
       // Instead of storing in sessionStorage, pass the code directly
-      // to the auth state manager through onConnect
-      onConnect({ authorizationCode: params.code });
+      // to the auth state manager through onConnect, along with restored state
+      onConnect({ authorizationCode: params.code, restoredState });
     };
 
     handleCallback().finally(() => {
diff --git a/client/src/lib/constants.ts b/client/src/lib/constants.ts
index 4c3e27aa..3acbda91 100644
--- a/client/src/lib/constants.ts
+++ b/client/src/lib/constants.ts
@@ -7,6 +7,7 @@ export const SESSION_KEYS = {
   TOKENS: "mcp_tokens",
   CLIENT_INFORMATION: "mcp_client_information",
   SERVER_METADATA: "mcp_server_metadata",
+  AUTH_DEBUGGER_STATE: "mcp_auth_debugger_state",
 } as const;
 
 // Generate server-specific session storage keys