diff --git a/client/src/App.tsx b/client/src/App.tsx index f3f1180e..74cdca19 100644 --- a/client/src/App.tsx +++ b/client/src/App.tsx @@ -122,7 +122,8 @@ const App = () => { const [isAuthDebuggerVisible, setIsAuthDebuggerVisible] = useState(false); // Auth debugger state - const [authState, setAuthState] = useState(EMPTY_DEBUGGER_STATE); + const [authState, setAuthState] = + useState(EMPTY_DEBUGGER_STATE); // Helper function to update specific auth state properties const updateAuthState = (updates: Partial) => { @@ -268,7 +269,10 @@ const App = () => { }); // Continue stepping through the OAuth flow from where we left off - while (currentState.oauthStep !== "complete" && currentState.oauthStep !== "authorization_code") { + while ( + currentState.oauthStep !== "complete" && + currentState.oauthStep !== "authorization_code" + ) { await stateMachine.executeStep(currentState); } @@ -286,7 +290,8 @@ const App = () => { } catch (error) { console.error("OAuth continuation error:", error); updateAuthState({ - latestError: error instanceof Error ? error : new Error(String(error)), + 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)}`, diff --git a/client/src/components/AuthDebugger.tsx b/client/src/components/AuthDebugger.tsx index a90b1061..b88758d4 100644 --- a/client/src/components/AuthDebugger.tsx +++ b/client/src/components/AuthDebugger.tsx @@ -67,13 +67,13 @@ const AuthDebugger = ({ const checkTokens = async () => { const provider = new DebugInspectorOAuthClientProvider(serverUrl); const existingTokens = await provider.tokens(); - + updateAuthState({ loading: false, oauthTokens: existingTokens || null, }); }; - + checkTokens(); } }, [serverUrl, authState.loading, updateAuthState]); @@ -162,7 +162,7 @@ const AuthDebugger = ({ // Store the current auth state before redirecting sessionStorage.setItem( SESSION_KEYS.AUTH_DEBUGGER_STATE, - JSON.stringify(currentState) + JSON.stringify(currentState), ); // Open the authorization URL automatically window.location.href = currentState.authorizationUrl; diff --git a/client/src/components/OAuthDebugCallback.tsx b/client/src/components/OAuthDebugCallback.tsx index af9944ac..17705abb 100644 --- a/client/src/components/OAuthDebugCallback.tsx +++ b/client/src/components/OAuthDebugCallback.tsx @@ -36,9 +36,11 @@ 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); + const storedState = sessionStorage.getItem( + SESSION_KEYS.AUTH_DEBUGGER_STATE, + ); let restoredState = null; if (storedState) { try { diff --git a/client/src/components/OAuthFlowProgress.tsx b/client/src/components/OAuthFlowProgress.tsx index f86ceec8..b9b67c6c 100644 --- a/client/src/components/OAuthFlowProgress.tsx +++ b/client/src/components/OAuthFlowProgress.tsx @@ -139,7 +139,13 @@ export const OAuthFlowProgress = ({

Resource Metadata:

- From {new URL('/.well-known/oauth-protected-resource', serverUrl).href} + From{" "} + { + new URL( + "/.well-known/oauth-protected-resource", + serverUrl, + ).href + } 

                     {JSON.stringify(authState.resourceMetadata, null, 2)}
@@ -150,17 +156,35 @@ export const OAuthFlowProgress = ({
               {authState.resourceMetadataError && (
                 

                   

-                    ℹ️ No resource metadata available from {' '}
-                    
-                      {new URL('/.well-known/oauth-protected-resource', serverUrl).href}
+                    ℹ️ No resource metadata available from{" "}
+                    
+                      {
+                        new URL(
+                          "/.well-known/oauth-protected-resource",
+                          serverUrl,
+                        ).href
+                      }
                     
                   

                   

-                    Resource metadata was added in the 2025-DRAFT-v2 specification update
+                    Resource metadata was added in the{" "}
+                    
+                      2025-DRAFT-v2 specification update
+                    
                     

                     {authState.resourceMetadataError.message}
-                    {authState.resourceMetadataError instanceof TypeError
-                      && " (This could indicate the endpoint doesn't exist or does not have CORS configured)"}
+                    {authState.resourceMetadataError instanceof TypeError &&
+                      " (This could indicate the endpoint doesn't exist or does not have CORS configured)"}
                   

                 

               )}
@@ -168,9 +192,17 @@ export const OAuthFlowProgress = ({
               {authState.oauthMetadata && (
                 

                   
Authorization Server Metadata:

-                  {authState.authServerUrl && 

-                    From {new URL('/.well-known/oauth-authorization-server', authState.authServerUrl).href}
-                  
}
+                  {authState.authServerUrl && (
+                    

+                      From{" "}
+                      {
+                        new URL(
+                          "/.well-known/oauth-authorization-server",
+                          authState.authServerUrl,
+                        ).href
+                      }
+                    

+                  )}
                   
                     {JSON.stringify(authState.oauthMetadata, null, 2)}
                   

diff --git a/client/src/components/__tests__/AuthDebugger.test.tsx b/client/src/components/__tests__/AuthDebugger.test.tsx
index 42a289ff..fe5a94f9 100644
--- a/client/src/components/__tests__/AuthDebugger.test.tsx
+++ b/client/src/components/__tests__/AuthDebugger.test.tsx
@@ -40,6 +40,7 @@ jest.mock("@modelcontextprotocol/sdk/client/auth.js", () => ({
   registerClient: jest.fn(),
   startAuthorization: jest.fn(),
   exchangeAuthorization: jest.fn(),
+  discoverOAuthProtectedResourceMetadata: jest.fn(),
 }));
 
 // Import the functions to get their types
@@ -49,6 +50,7 @@ import {
   startAuthorization,
   exchangeAuthorization,
   auth,
+  discoverOAuthProtectedResourceMetadata,
 } from "@modelcontextprotocol/sdk/client/auth.js";
 import { OAuthMetadata } from "@modelcontextprotocol/sdk/shared/auth.js";
 import { EMPTY_DEBUGGER_STATE } from "@/lib/auth-types";
@@ -67,6 +69,10 @@ const mockExchangeAuthorization = exchangeAuthorization as jest.MockedFunction<
   typeof exchangeAuthorization
 >;
 const mockAuth = auth as jest.MockedFunction;
+const mockDiscoverOAuthProtectedResourceMetadata =
+  discoverOAuthProtectedResourceMetadata as jest.MockedFunction<
+    typeof discoverOAuthProtectedResourceMetadata
+  >;
 
 const sessionStorageMock = {
   getItem: jest.fn(),
@@ -100,6 +106,7 @@ describe("AuthDebugger", () => {
 
     mockDiscoverOAuthMetadata.mockResolvedValue(mockOAuthMetadata);
     mockRegisterClient.mockResolvedValue(mockOAuthClientInfo);
+    mockDiscoverOAuthProtectedResourceMetadata.mockResolvedValue(null);
     mockStartAuthorization.mockImplementation(async (_sseUrl, options) => {
       const authUrl = new URL("https://oauth.example.com/authorize");
 
@@ -421,4 +428,63 @@ describe("AuthDebugger", () => {
       });
     });
   });
+
+  describe("OAuth State Persistence", () => {
+    it("should store auth state to sessionStorage before redirect in Quick OAuth Flow", async () => {
+      const updateAuthState = jest.fn();
+
+      // Mock window.location.href setter
+      delete (window as any).location;
+      window.location = { href: "" } as any;
+
+      // Setup mocks for OAuth flow
+      mockStartAuthorization.mockResolvedValue({
+        authorizationUrl: new URL(
+          "https://oauth.example.com/authorize?client_id=test_client_id&redirect_uri=http%3A%2F%2Flocalhost%3A3000%2Foauth%2Fcallback%2Fdebug",
+        ),
+        codeVerifier: "test_verifier",
+      });
+
+      await act(async () => {
+        renderAuthDebugger({
+          updateAuthState,
+          authState: { ...defaultAuthState, loading: false },
+        });
+      });
+
+      // Click Quick OAuth Flow
+      await act(async () => {
+        fireEvent.click(screen.getByText("Quick OAuth Flow"));
+      });
+
+      // Wait for the flow to reach the authorization step
+      await waitFor(() => {
+        expect(sessionStorage.setItem).toHaveBeenCalledWith(
+          SESSION_KEYS.AUTH_DEBUGGER_STATE,
+          expect.stringContaining('"oauthStep":"authorization_code"'),
+        );
+      });
+
+      // Verify the stored state includes all the accumulated data
+      const storedStateCall = (
+        sessionStorage.setItem as jest.Mock
+      ).mock.calls.find((call) => call[0] === SESSION_KEYS.AUTH_DEBUGGER_STATE);
+
+      expect(storedStateCall).toBeDefined();
+      const storedState = JSON.parse(storedStateCall![1]);
+
+      expect(storedState).toMatchObject({
+        oauthStep: "authorization_code",
+        authorizationUrl: expect.stringMatching(
+          /^https:\/\/oauth\.example\.com\/authorize/,
+        ),
+        oauthMetadata: expect.objectContaining({
+          token_endpoint: "https://oauth.example.com/token",
+        }),
+        oauthClientInfo: expect.objectContaining({
+          client_id: "test_client_id",
+        }),
+      });
+    });
+  });
 });
diff --git a/client/src/lib/auth-types.ts b/client/src/lib/auth-types.ts
index 2541e822..e69dcd12 100644
--- a/client/src/lib/auth-types.ts
+++ b/client/src/lib/auth-types.ts
@@ -30,7 +30,10 @@ export interface AuthDebuggerState {
   loading: boolean;
   oauthStep: OAuthStep;
   resourceMetadata: OAuthProtectedResourceMetadata | null;
-  resourceMetadataError: Error | { status: number; statusText: string; message: string } | null;
+  resourceMetadataError:
+    | Error
+    | { status: number; statusText: string; message: string }
+    | null;
   authServerUrl: URL | null;
   oauthMetadata: OAuthMetadata | null;
   oauthClientInfo: OAuthClientInformationFull | OAuthClientInformation | null;
@@ -56,4 +59,4 @@ export const EMPTY_DEBUGGER_STATE: AuthDebuggerState = {
   latestError: null,
   statusMessage: null,
   validationError: null,
-}
+};
diff --git a/client/src/lib/oauth-state-machine.ts b/client/src/lib/oauth-state-machine.ts
index 939e5d6c..d05369c4 100644
--- a/client/src/lib/oauth-state-machine.ts
+++ b/client/src/lib/oauth-state-machine.ts
@@ -7,7 +7,10 @@ import {
   exchangeAuthorization,
   discoverOAuthProtectedResourceMetadata,
 } from "@modelcontextprotocol/sdk/client/auth.js";
-import { OAuthMetadataSchema, OAuthProtectedResourceMetadata } from "@modelcontextprotocol/sdk/shared/auth.js";
+import {
+  OAuthMetadataSchema,
+  OAuthProtectedResourceMetadata,
+} from "@modelcontextprotocol/sdk/shared/auth.js";
 
 export interface StateMachineContext {
   state: AuthDebuggerState;
@@ -30,8 +33,13 @@ export const oauthTransitions: Record = {
       let resourceMetadata: OAuthProtectedResourceMetadata | null = null;
       let resourceMetadataError: Error | null = null;
       try {
-        resourceMetadata = await discoverOAuthProtectedResourceMetadata(context.serverUrl);
-        if (resourceMetadata && resourceMetadata.authorization_servers?.length) {
+        resourceMetadata = await discoverOAuthProtectedResourceMetadata(
+          context.serverUrl,
+        );
+        if (
+          resourceMetadata &&
+          resourceMetadata.authorization_servers?.length
+        ) {
           authServerUrl = new URL(resourceMetadata.authorization_servers[0]);
         }
       } catch (e) {