[scenario] Add test for broken oauth client (#12)

* starting on broken auth client

* switch to middleware approach

* add broken example

* fix failure testing

* subset matching

examples/clients/typescript/auth-test-broken1.ts

@@ -0,0 +1,99 @@
+#!/usr/bin/env node
+
+import { Client } from '@modelcontextprotocol/sdk/client/index.js';
+import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { handle401, withOAuthRetry } from './helpers/withOAuthRetry.js';
+import { ConformanceOAuthProvider } from './helpers/ConformanceOAuthProvider.js';
+import { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js';
+import {
+  auth,
+  UnauthorizedError
+} from '@modelcontextprotocol/sdk/client/auth.js';
+
+export const handle401Broken = async (
+  response: Response,
+  provider: ConformanceOAuthProvider,
+  next: FetchLike,
+  serverUrl: string | URL
+): Promise<void> => {
+  // BROKEN: Use root-based PRM discovery exclusively, regardless of input.
+  const resourceMetadataUrl = new URL(
+    '/.well-known/oauth-protected-resource',
+    typeof serverUrl === 'string' ? serverUrl : serverUrl.origin
+  );
+
+  let result = await auth(provider, {
+    serverUrl,
+    resourceMetadataUrl,
+    fetchFn: next
+  });
+
+  if (result === 'REDIRECT') {
+    // Ordinarily, we'd wait for the callback to be handled here,
+    // but in our conformance provider, we get the authorization code
+    // during the redirect handling, so we can go straight to
+    // retrying the auth step.
+    // await provider.waitForCallback();
+
+    const authorizationCode = await provider.getAuthCode();
+
+    // TODO: this retry logic should be incorporated into the typescript SDK
+    result = await auth(provider, {
+      serverUrl,
+      resourceMetadataUrl,
+      authorizationCode,
+      fetchFn: next
+    });
+    if (result !== 'AUTHORIZED') {
+      throw new UnauthorizedError(
+        `Authentication failed with result: ${result}`
+      );
+    }
+  }
+};
+
+async function main(): Promise<void> {
+  const serverUrl = process.argv[2];
+
+  if (!serverUrl) {
+    console.error('Usage: auth-test <server-url>');
+    process.exit(1);
+  }
+
+  console.log(`Connecting to MCP server at: ${serverUrl}`);
+
+  const client = new Client(
+    {
+      name: 'test-auth-client',
+      version: '1.0.0'
+    },
+    {
+      capabilities: {}
+    }
+  );
+
+  // Create a custom fetch that uses the OAuth middleware with retry logic
+  const oauthFetch = withOAuthRetry(
+    'test-auth-client',
+    new URL(serverUrl),
+    handle401Broken
+  )(fetch);
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
+    fetch: oauthFetch
+  });
+
+  // Connect to the server - OAuth is handled automatically by the middleware
+  await client.connect(transport);
+  console.log('✅ Successfully connected to MCP server');
+
+  await client.listTools();
+  console.log('✅ Successfully listed tools');
+
+  await transport.close();
+  console.log('✅ Connection closed successfully');
+
+  process.exit(0);
+}
+
+main();

examples/clients/typescript/auth-test.ts

@@ -2,8 +2,7 @@
 
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { ConformanceOAuthProvider } from './helpers/ConformanceOAuthProvider.js';
-import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js';
+import { withOAuthRetry } from './helpers/withOAuthRetry.js';
 
 async function main(): Promise<void> {
   const serverUrl = process.argv[2];
@@ -25,47 +24,19 @@ async function main(): Promise<void> {
     }
   );
 
-  const authProvider = new ConformanceOAuthProvider(
-    'http://localhost:3000/callback',
-    {
-      client_name: 'test-auth-client',
-      redirect_uris: ['http://localhost:3000/callback']
-    }
-  );
+  // Create a custom fetch that uses the OAuth middleware with retry logic
+  const oauthFetch = withOAuthRetry(
+    'test-auth-client',
+    new URL(serverUrl)
+  )(fetch);
 
-  let transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
-    authProvider
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
+    fetch: oauthFetch
   });
 
-  // Try to connect - handle OAuth if needed
-  try {
-    await client.connect(transport);
-    console.log('✅ Successfully connected to MCP server');
-  } catch (error) {
-    if (error instanceof UnauthorizedError) {
-      console.log('🔐 OAuth required - handling authorization...');
-
-      // The provider will automatically fetch the auth code
-      const authCode = await authProvider.getAuthCode();
-
-      // Complete the auth flow
-      await transport.finishAuth(authCode);
-
-      // Close the old transport
-      await transport.close();
-
-      // Create a new transport with the authenticated provider
-      transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
-        authProvider: authProvider
-      });
-
-      // Connect with the new transport
-      await client.connect(transport);
-      console.log('✅ Successfully connected with authentication');
-    } else {
-      throw error;
-    }
-  }
+  // Connect to the server - OAuth is handled automatically by the middleware
+  await client.connect(transport);
+  console.log('✅ Successfully connected to MCP server');
 
   await client.listTools();
   console.log('✅ Successfully listed tools');

examples/clients/typescript/helpers/withOAuthRetry.ts

@@ -0,0 +1,109 @@
+import {
+  auth,
+  extractResourceMetadataUrl,
+  UnauthorizedError
+} from '@modelcontextprotocol/sdk/client/auth.js';
+import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js';
+import type { Middleware } from '@modelcontextprotocol/sdk/client/middleware.js';
+import { ConformanceOAuthProvider } from './ConformanceOAuthProvider';
+
+export const handle401 = async (
+  response: Response,
+  provider: ConformanceOAuthProvider,
+  next: FetchLike,
+  serverUrl: string | URL
+): Promise<void> => {
+  const resourceMetadataUrl = extractResourceMetadataUrl(response);
+
+  let result = await auth(provider, {
+    serverUrl,
+    resourceMetadataUrl,
+    fetchFn: next
+  });
+
+  if (result === 'REDIRECT') {
+    // Ordinarily, we'd wait for the callback to be handled here,
+    // but in our conformance provider, we get the authorization code
+    // during the redirect handling, so we can go straight to
+    // retrying the auth step.
+    // await provider.waitForCallback();
+
+    const authorizationCode = await provider.getAuthCode();
+
+    // TODO: this retry logic should be incorporated into the typescript SDK
+    result = await auth(provider, {
+      serverUrl,
+      resourceMetadataUrl,
+      authorizationCode,
+      fetchFn: next
+    });
+    if (result !== 'AUTHORIZED') {
+      throw new UnauthorizedError(
+        `Authentication failed with result: ${result}`
+      );
+    }
+  }
+};
+/**
+ * Creates a fetch wrapper that handles OAuth authentication with retry logic.
+ *
+ * Unlike the SDK's withOAuth, this version:
+ * - Automatically handles authorization redirects by retrying with fresh tokens
+ * - Does not throw UnauthorizedError on redirect, but instead retries
+ * - Calls next() instead of throwing for redirect-based auth
+ *
+ * @param provider - OAuth client provider for authentication
+ * @param baseUrl - Base URL for OAuth server discovery (defaults to request URL domain)
+ * @returns A fetch middleware function
+ */
+export const withOAuthRetry = (
+  clientName: string,
+  baseUrl?: string | URL,
+  handle401Fn: typeof handle401 = handle401
+): Middleware => {
+  const provider = new ConformanceOAuthProvider(
+    'http://localhost:3000/callback',
+    {
+      client_name: clientName,
+      redirect_uris: ['http://localhost:3000/callback']
+    }
+  );
+  return (next: FetchLike) => {
+    return async (
+      input: string | URL,
+      init?: RequestInit
+    ): Promise<Response> => {
+      const makeRequest = async (): Promise<Response> => {
+        const headers = new Headers(init?.headers);
+
+        // Add authorization header if tokens are available
+        const tokens = await provider.tokens();
+        if (tokens) {
+          headers.set('Authorization', `Bearer ${tokens.access_token}`);
+        }
+
+        return await next(input, { ...init, headers });
+      };
+
+      let response = await makeRequest();
+
+      // Handle 401 responses by attempting re-authentication
+      if (response.status === 401) {
+        const serverUrl =
+          baseUrl ||
+          (typeof input === 'string' ? new URL(input).origin : input.origin);
+        await handle401Fn(response, provider, next, serverUrl);
+
+        response = await makeRequest();
+      }
+
+      // If we still have a 401 after re-auth attempt, throw an error
+      if (response.status === 401) {
+        const url = typeof input === 'string' ? input : input.toString();
+        throw new UnauthorizedError(`Authentication failed for ${url}`);
+      }
+
+      return response;
+    };
+  };
+};