store state before redirect

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;

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