[auth] Support new auth metadata endpoint from draft spec

client/src/App.tsx

@@ -19,7 +19,8 @@ import {
 } from "@modelcontextprotocol/sdk/types.js";
 import { OAuthTokensSchema } from "@modelcontextprotocol/sdk/shared/auth.js";
 import { SESSION_KEYS, getServerSpecificKey } from "./lib/constants";
-import { AuthDebuggerState } from "./lib/auth-types";
+import { AuthDebuggerState, EMPTY_DEBUGGER_STATE } from "./lib/auth-types";
+import { OAuthStateMachine } from "./lib/oauth-state-machine";
 import { cacheToolOutputSchemas } from "./utils/schemaUtils";
 import React, {
   Suspense,
@@ -121,19 +122,8 @@ const App = () => {
   const [isAuthDebuggerVisible, setIsAuthDebuggerVisible] = useState(false);
 
   // Auth debugger state
-  const [authState, setAuthState] = useState<AuthDebuggerState>({
-    isInitiatingAuth: false,
-    oauthTokens: null,
-    loading: true,
-    oauthStep: "metadata_discovery",
-    oauthMetadata: null,
-    oauthClientInfo: null,
-    authorizationUrl: null,
-    authorizationCode: "",
-    latestError: null,
-    statusMessage: null,
-    validationError: null,
-  });
+  const [authState, setAuthState] =
+    useState<AuthDebuggerState>(EMPTY_DEBUGGER_STATE);
 
   // Helper function to update specific auth state properties
   const updateAuthState = (updates: Partial<AuthDebuggerState>) => {
@@ -243,27 +233,81 @@ const App = () => {
 
   // Update OAuth debug state during debug callback
   const onOAuthDebugConnect = useCallback(
-    ({
+    async ({
       authorizationCode,
       errorMsg,
+      restoredState,
     }: {
       authorizationCode?: string;
       errorMsg?: string;
+      restoredState?: AuthDebuggerState;
     }) => {
       setIsAuthDebuggerVisible(true);
-      if (authorizationCode) {
+
+      if (errorMsg) {
         updateAuthState({
-          authorizationCode,
-          oauthStep: "token_request",
+          latestError: new Error(errorMsg),
         });
+        return;
       }
-      if (errorMsg) {
+
+      if (restoredState && authorizationCode) {
+        // Restore the previous auth state and continue the OAuth flow
+        let currentState: AuthDebuggerState = {
+          ...restoredState,
+          authorizationCode,
+          oauthStep: "token_request",
+          isInitiatingAuth: true,
+          statusMessage: null,
+          latestError: null,
+        };
+
+        try {
+          // Create a new state machine instance to continue the flow
+          const stateMachine = new OAuthStateMachine(sseUrl, (updates) => {
+            currentState = { ...currentState, ...updates };
+          });
+
+          // Continue stepping through the OAuth flow from where we left off
+          while (
+            currentState.oauthStep !== "complete" &&
+            currentState.oauthStep !== "authorization_code"
+          ) {
+            await stateMachine.executeStep(currentState);
+          }
+
+          if (currentState.oauthStep === "complete") {
+            // After the flow completes or reaches a user-input step, update the app state
+            updateAuthState({
+              ...currentState,
+              statusMessage: {
+                type: "success",
+                message: "Authentication completed successfully",
+              },
+              isInitiatingAuth: false,
+            });
+          }
+        } catch (error) {
+          console.error("OAuth continuation error:", error);
+          updateAuthState({
+            latestError:
+              error instanceof Error ? error : new Error(String(error)),
+            statusMessage: {
+              type: "error",
+              message: `Failed to complete OAuth flow: ${error instanceof Error ? error.message : String(error)}`,
+            },
+            isInitiatingAuth: false,
+          });
+        }
+      } else if (authorizationCode) {
+        // Fallback to the original behavior if no state was restored
         updateAuthState({
-          latestError: new Error(errorMsg),
+          authorizationCode,
+          oauthStep: "token_request",
         });
       }
     },
-    [],
+    [sseUrl],
   );
 
   // Load OAuth tokens when sseUrl changes
@@ -285,8 +329,6 @@ const App = () => {
         }
       } catch (error) {
         console.error("Error loading OAuth tokens:", error);
-      } finally {
-        updateAuthState({ loading: false });
       }
     };
 

client/src/components/AuthDebugger.tsx

@@ -1,10 +1,11 @@
-import { useCallback, useMemo } from "react";
+import { useCallback, useMemo, useEffect } from "react";
 import { Button } from "@/components/ui/button";
 import { DebugInspectorOAuthClientProvider } from "../lib/auth";
 import { AlertCircle } from "lucide-react";
-import { AuthDebuggerState } from "../lib/auth-types";
+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;
@@ -59,6 +60,27 @@ const AuthDebugger = ({
   authState,
   updateAuthState,
 }: AuthDebuggerProps) => {
+  // Check for existing tokens on mount
+  useEffect(() => {
+    if (serverUrl && !authState.oauthTokens) {
+      const checkTokens = async () => {
+        try {
+          const provider = new DebugInspectorOAuthClientProvider(serverUrl);
+          const existingTokens = await provider.tokens();
+          if (existingTokens) {
+            updateAuthState({
+              oauthTokens: existingTokens,
+              oauthStep: "complete",
+            });
+          }
+        } catch (error) {
+          console.error("Failed to load existing OAuth tokens:", error);
+        }
+      };
+      checkTokens();
+    }
+  }, [serverUrl, updateAuthState, authState.oauthTokens]);
+
   const startOAuthFlow = useCallback(() => {
     if (!serverUrl) {
       updateAuthState({
@@ -141,6 +163,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;
@@ -178,13 +205,7 @@ const AuthDebugger = ({
       );
       serverAuthProvider.clear();
       updateAuthState({
-        oauthTokens: null,
-        oauthStep: "metadata_discovery",
-        latestError: null,
-        oauthClientInfo: null,
-        authorizationCode: "",
-        validationError: null,
-        oauthMetadata: null,
+        ...EMPTY_DEBUGGER_STATE,
         statusMessage: {
           type: "success",
           message: "OAuth tokens cleared successfully",
@@ -224,52 +245,48 @@ const AuthDebugger = ({
                 <StatusMessage message={authState.statusMessage} />
               )}
 
-              {authState.loading ? (
-                <p>Loading authentication status...</p>
-              ) : (
-                <div className="space-y-4">
-                  {authState.oauthTokens && (
-                    <div className="space-y-2">
-                      <p className="text-sm font-medium">Access Token:</p>
-                      <div className="bg-muted p-2 rounded-md text-xs overflow-x-auto">
-                        {authState.oauthTokens.access_token.substring(0, 25)}...
-                      </div>
+              <div className="space-y-4">
+                {authState.oauthTokens && (
+                  <div className="space-y-2">
+                    <p className="text-sm font-medium">Access Token:</p>
+                    <div className="bg-muted p-2 rounded-md text-xs overflow-x-auto">
+                      {authState.oauthTokens.access_token.substring(0, 25)}...
                     </div>
-                  )}
-
-                  <div className="flex gap-4">
-                    <Button
-                      variant="outline"
-                      onClick={startOAuthFlow}
-                      disabled={authState.isInitiatingAuth}
-                    >
-                      {authState.oauthTokens
-                        ? "Guided Token Refresh"
-                        : "Guided OAuth Flow"}
-                    </Button>
+                  </div>
+                )}
 
-                    <Button
-                      onClick={handleQuickOAuth}
-                      disabled={authState.isInitiatingAuth}
-                    >
-                      {authState.isInitiatingAuth
-                        ? "Initiating..."
-                        : authState.oauthTokens
-                          ? "Quick Refresh"
-                          : "Quick OAuth Flow"}
-                    </Button>
+                <div className="flex gap-4">
+                  <Button
+                    variant="outline"
+                    onClick={startOAuthFlow}
+                    disabled={authState.isInitiatingAuth}
+                  >
+                    {authState.oauthTokens
+                      ? "Guided Token Refresh"
+                      : "Guided OAuth Flow"}
+                  </Button>
 
-                    <Button variant="outline" onClick={handleClearOAuth}>
-                      Clear OAuth State
-                    </Button>
-                  </div>
+                  <Button
+                    onClick={handleQuickOAuth}
+                    disabled={authState.isInitiatingAuth}
+                  >
+                    {authState.isInitiatingAuth
+                      ? "Initiating..."
+                      : authState.oauthTokens
+                        ? "Quick Refresh"
+                        : "Quick OAuth Flow"}
+                  </Button>
 
-                  <p className="text-xs text-muted-foreground">
-                    Choose "Guided" for step-by-step instructions or "Quick" for
-                    the standard automatic flow.
-                  </p>
+                  <Button variant="outline" onClick={handleClearOAuth}>
+                    Clear OAuth State
+                  </Button>
                 </div>
-              )}
+
+                <p className="text-xs text-muted-foreground">
+                  Choose "Guided" for step-by-step instructions or "Quick" for
+                  the standard automatic flow.
+                </p>
+              </div>
             </div>
 
             <OAuthFlowProgress

client/src/components/OAuthDebugCallback.tsx

@@ -4,14 +4,17 @@ import {
   generateOAuthErrorDescription,
   parseOAuthCallbackParams,
 } from "@/utils/oauthUtils.ts";
+import { AuthDebuggerState } from "@/lib/auth-types";
 
 interface OAuthCallbackProps {
   onConnect: ({
     authorizationCode,
     errorMsg,
+    restoredState,
   }: {
     authorizationCode?: string;
     errorMsg?: string;
+    restoredState?: AuthDebuggerState;
   }) => void;
 }
 
@@ -35,6 +38,21 @@ 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
       // session storage
@@ -50,8 +68,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(() => {