From c6fad3f92343ccce99ac2ad9e635a08279e09ec6 Mon Sep 17 00:00:00 2001
From: Kamal Fariz Mahyuddin <kamal@envoy.com>
Date: Fri, 16 May 2025 08:22:39 -0700
Subject: [PATCH] fix: quick oauth flow could not transition to token_request
 due to missing server metadata (#407)

* fix: quick oauth flow could not transition to token_request due to missing server metadata

* chore: prettier fix
---
 client/src/components/AuthDebugger.tsx        | 14 ++++-
 .../__tests__/AuthDebugger.test.tsx           | 59 +++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/client/src/components/AuthDebugger.tsx b/client/src/components/AuthDebugger.tsx
index fa863373..a2f20e26 100644
--- a/client/src/components/AuthDebugger.tsx
+++ b/client/src/components/AuthDebugger.tsx
@@ -1,7 +1,11 @@
 import { useCallback, useMemo } from "react";
 import { Button } from "@/components/ui/button";
 import { DebugInspectorOAuthClientProvider } from "../lib/auth";
-import { auth } from "@modelcontextprotocol/sdk/client/auth.js";
+import {
+  auth,
+  discoverOAuthMetadata,
+} from "@modelcontextprotocol/sdk/client/auth.js";
+import { OAuthMetadataSchema } from "@modelcontextprotocol/sdk/shared/auth.js";
 import { AlertCircle } from "lucide-react";
 import { AuthDebuggerState } from "../lib/auth-types";
 import { OAuthFlowProgress } from "./OAuthFlowProgress";
@@ -123,6 +127,14 @@ const AuthDebugger = ({
       const serverAuthProvider = new DebugInspectorOAuthClientProvider(
         serverUrl,
       );
+      // First discover OAuth metadata separately so we can save it
+      const metadata = await discoverOAuthMetadata(serverUrl);
+      if (!metadata) {
+        throw new Error("Failed to discover OAuth metadata");
+      }
+      const parsedMetadata = await OAuthMetadataSchema.parseAsync(metadata);
+      serverAuthProvider.saveServerMetadata(parsedMetadata);
+
       await auth(serverAuthProvider, { serverUrl: serverUrl });
       updateAuthState({
         statusMessage: {
diff --git a/client/src/components/__tests__/AuthDebugger.test.tsx b/client/src/components/__tests__/AuthDebugger.test.tsx
index 469c1ba1..aa206ecc 100644
--- a/client/src/components/__tests__/AuthDebugger.test.tsx
+++ b/client/src/components/__tests__/AuthDebugger.test.tsx
@@ -48,6 +48,7 @@ import {
   registerClient,
   startAuthorization,
   exchangeAuthorization,
+  auth,
 } from "@modelcontextprotocol/sdk/client/auth.js";
 import { OAuthMetadata } from "@modelcontextprotocol/sdk/shared/auth.js";
 
@@ -64,6 +65,7 @@ const mockStartAuthorization = startAuthorization as jest.MockedFunction<
 const mockExchangeAuthorization = exchangeAuthorization as jest.MockedFunction<
   typeof exchangeAuthorization
 >;
+const mockAuth = auth as jest.MockedFunction<typeof auth>;
 
 const sessionStorageMock = {
   getItem: jest.fn(),
@@ -186,6 +188,63 @@ describe("AuthDebugger", () => {
         },
       });
     });
+
+    it("should start quick OAuth flow and properly fetch and save metadata", async () => {
+      // Setup the auth mock
+      mockAuth.mockResolvedValue("AUTHORIZED");
+
+      const updateAuthState = jest.fn();
+      await act(async () => {
+        renderAuthDebugger({ updateAuthState });
+      });
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Quick OAuth Flow"));
+      });
+
+      // Should first discover and save OAuth metadata
+      expect(mockDiscoverOAuthMetadata).toHaveBeenCalledWith(
+        "https://example.com",
+      );
+
+      // Then should call auth with the server provider
+      expect(mockAuth).toHaveBeenCalled();
+
+      // Check that updateAuthState was called with the right info message
+      expect(updateAuthState).toHaveBeenCalledWith(
+        expect.objectContaining({
+          statusMessage: {
+            type: "info",
+            message: "Starting OAuth authentication process...",
+          },
+        }),
+      );
+    });
+
+    it("should show error when quick OAuth flow fails to discover metadata", async () => {
+      mockDiscoverOAuthMetadata.mockRejectedValue(
+        new Error("Metadata discovery failed"),
+      );
+
+      const updateAuthState = jest.fn();
+      await act(async () => {
+        renderAuthDebugger({ updateAuthState });
+      });
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Quick OAuth Flow"));
+      });
+
+      // Check that updateAuthState was called with an error message
+      expect(updateAuthState).toHaveBeenCalledWith(
+        expect.objectContaining({
+          statusMessage: {
+            type: "error",
+            message: expect.stringContaining("Failed to start OAuth flow"),
+          },
+        }),
+      );
+    });
   });
 
   describe("Session Storage Integration", () => {