feat: add client conformance testing (#1360)

2026-01-09 17:49:49 +00:00 · 2026-01-09 17:49:49 +00:00 · 1014d9f79c
commit 1014d9f79c
parent 3466a9e0e5
7 changed files with 759 additions and 2 deletions
--- a/.github/workflows/conformance.yml
+++ b/.github/workflows/conformance.yml
@ -0,0 +1,29 @@
+name: Conformance Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  client-conformance:
+    runs-on: ubuntu-latest
+    continue-on-error: true  # Non-blocking initially
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install pnpm
+        uses: pnpm/action-setup@v4
+        with:
+          run_install: false
+      - uses: actions/setup-node@v4
+        with:
+          node-version: 24
+          cache: pnpm
+          cache-dependency-path: pnpm-lock.yaml
+      - run: pnpm install
+      - run: pnpm run build:all
+      - run: pnpm run test:conformance:client:all
--- a/package.json
+++ b/package.json
@ -29,13 +29,18 @@
        "lint:all": "pnpm -r lint",
        "lint:fix:all": "pnpm -r lint:fix",
        "check:all": "pnpm -r typecheck && pnpm -r lint",
-        "test:all": "pnpm -r test"
+        "test:all": "pnpm -r test",
+        "test:conformance:client": "conformance client --command 'npx tsx src/conformance/everything-client.ts'",
+        "test:conformance:client:all": "conformance client --command 'npx tsx src/conformance/everything-client.ts' --suite all",
+        "test:conformance:client:run": "npx tsx src/conformance/everything-client.ts"
    },
    "devDependencies": {
        "@cfworker/json-schema": "catalog:runtimeShared",
        "@changesets/changelog-github": "^0.5.2",
        "@changesets/cli": "^2.29.8",
        "@eslint/js": "catalog:devTools",
+        "@modelcontextprotocol/client": "workspace:^",
+        "@modelcontextprotocol/conformance": "0.1.9",
        "@types/content-type": "catalog:devTools",
        "@types/cors": "catalog:devTools",
        "@types/cross-spawn": "catalog:devTools",
@ -56,7 +61,8 @@
        "typescript": "catalog:devTools",
        "typescript-eslint": "catalog:devTools",
        "vitest": "catalog:devTools",
-        "ws": "catalog:devTools"
+        "ws": "catalog:devTools",
+        "zod": "catalog:runtimeShared"
    },
    "resolutions": {
        "strip-ansi": "6.0.1"
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -149,6 +149,12 @@ importers:
      '@eslint/js':
        specifier: catalog:devTools
        version: 9.39.1
+      '@modelcontextprotocol/client':
+        specifier: workspace:^
+        version: link:packages/client
+      '@modelcontextprotocol/conformance':
+        specifier: 0.1.9
+        version: 0.1.9(@cfworker/json-schema@4.1.1)(hono@4.11.1)
      '@types/content-type':
        specifier: catalog:devTools
        version: 1.1.9
@ -212,6 +218,9 @@ importers:
      ws:
        specifier: catalog:devTools
        version: 8.18.3
+      zod:
+        specifier: catalog:runtimeShared
+        version: 3.25.76

  common/eslint-config:
    dependencies:
@ -1066,6 +1075,20 @@ packages:
  '@manypkg/get-packages@1.1.3':
    resolution: {integrity: sha512-fo+QhuU3qE/2TQMQmbVMqaQ6EWbMhi4ABWP+O4AM1NqPBuy0OrApV5LO6BrrgnhtAHS2NH6RrVk9OL181tTi8A==}

+  '@modelcontextprotocol/conformance@0.1.9':
+    resolution: {integrity: sha512-hpR5PoW0feue3LHSi1kJNhQxbySEQNWR6McuB3QCoK0zsxIdoq+id4GxRwWVOnRnjOiTecDKMD1QMfXuurDZPQ==}
+    hasBin: true
+
+  '@modelcontextprotocol/sdk@1.25.2':
+    resolution: {integrity: sha512-LZFeo4F9M5qOhC/Uc1aQSrBHxMrvxett+9KLHt7OhcExtoiRN9DKgbZffMP/nxjutWDQpfMDfP3nkHI4X9ijww==}
+    engines: {node: '>=18'}
+    peerDependencies:
+      '@cfworker/json-schema': ^4.1.1
+      zod: ^3.25 || ^4.0
+    peerDependenciesMeta:
+      '@cfworker/json-schema':
+        optional: true
+
  '@napi-rs/wasm-runtime@0.2.12':
    resolution: {integrity: sha512-ZVWUcfwY4E/yPitQJl481FjFo3K22D6qF0DuFH6Y/nbnE11GY5uguDxZMGXPQ8WQ0128MXQD7TnfHyK4oWoIJQ==}

@ -1773,6 +1796,10 @@ packages:
    resolution: {integrity: sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==}
    engines: {node: '>= 0.8'}

+  commander@14.0.2:
+    resolution: {integrity: sha512-TywoWNNRbhoD0BXs1P3ZEScW8W5iKrnbithIl0YH+uCmBd0QpPOA8yc82DS3BIE5Ma6FnBVUsJ7wVUDz4dvOWQ==}
+    engines: {node: '>=20'}
+
  component-emitter@1.3.1:
    resolution: {integrity: sha512-T0+barUSQRTUQASh8bx02dl+DhF54GtIDY13Y3m9oWTklKbb3Wv974meRpeZ3lp1JpLVECWWNHC4vaG2XHXouQ==}

@ -3730,6 +3757,43 @@ snapshots:
      globby: 11.1.0
      read-yaml-file: 1.1.0

+  '@modelcontextprotocol/conformance@0.1.9(@cfworker/json-schema@4.1.1)(hono@4.11.1)':
+    dependencies:
+      '@modelcontextprotocol/sdk': 1.25.2(@cfworker/json-schema@4.1.1)(hono@4.11.1)(zod@3.25.76)
+      commander: 14.0.2
+      eventsource-parser: 3.0.6
+      express: 5.1.0
+      jose: 6.1.3
+      zod: 3.25.76
+    transitivePeerDependencies:
+      - '@cfworker/json-schema'
+      - hono
+      - supports-color
+
+  '@modelcontextprotocol/sdk@1.25.2(@cfworker/json-schema@4.1.1)(hono@4.11.1)(zod@3.25.76)':
+    dependencies:
+      '@hono/node-server': 1.19.7(hono@4.11.1)
+      ajv: 8.17.1
+      ajv-formats: 3.0.1(ajv@8.17.1)
+      content-type: 1.0.5
+      cors: 2.8.5
+      cross-spawn: 7.0.6
+      eventsource: 3.0.7
+      eventsource-parser: 3.0.6
+      express: 5.1.0
+      express-rate-limit: 7.5.1(express@5.1.0)
+      jose: 6.1.3
+      json-schema-typed: 8.0.2
+      pkce-challenge: 5.0.0
+      raw-body: 3.0.1
+      zod: 3.25.76
+      zod-to-json-schema: 3.25.0(zod@3.25.76)
+    optionalDependencies:
+      '@cfworker/json-schema': 4.1.1
+    transitivePeerDependencies:
+      - hono
+      - supports-color
+
  '@napi-rs/wasm-runtime@0.2.12':
    dependencies:
      '@emnapi/core': 1.7.1
@ -4395,6 +4459,8 @@ snapshots:
    dependencies:
      delayed-stream: 1.0.0

+  commander@14.0.2: {}
+
  component-emitter@1.3.1: {}

  concat-map@0.0.1: {}
--- a/src/conformance/everything-client.ts
+++ b/src/conformance/everything-client.ts
@ -0,0 +1,423 @@
+#!/usr/bin/env node
+
+/**
+ * Everything client - a single conformance test client that handles all scenarios.
+ *
+ * Usage: everything-client <server-url>
+ *
+ * The scenario name is read from the MCP_CONFORMANCE_SCENARIO environment variable,
+ * which is set by the conformance test runner.
+ *
+ * This client routes to the appropriate behavior based on the scenario name,
+ * consolidating all the individual test clients into one.
+ */
+
+import {
+  Client,
+  StreamableHTTPClientTransport,
+  ElicitRequestSchema,
+  ClientCredentialsProvider,
+  PrivateKeyJwtProvider
+} from '@modelcontextprotocol/client';
+import { z } from 'zod';
+import { withOAuthRetry, handle401 } from './helpers/withOAuthRetry.js';
+import { logger } from './helpers/logger.js';
+
+/**
+ * Fixed client metadata URL for CIMD conformance tests.
+ * When server supports client_id_metadata_document_supported, this URL
+ * will be used as the client_id instead of doing dynamic registration.
+ */
+const CIMD_CLIENT_METADATA_URL =
+  'https://conformance-test.local/client-metadata.json';
+
+/**
+ * Schema for client conformance test context passed via MCP_CONFORMANCE_CONTEXT.
+ *
+ * Each variant includes a `name` field matching the scenario name to enable
+ * discriminated union parsing and type-safe access to scenario-specific fields.
+ */
+const ClientConformanceContextSchema = z.discriminatedUnion('name', [
+  z.object({
+    name: z.literal('auth/client-credentials-jwt'),
+    client_id: z.string(),
+    private_key_pem: z.string(),
+    signing_algorithm: z.string().optional()
+  }),
+  z.object({
+    name: z.literal('auth/client-credentials-basic'),
+    client_id: z.string(),
+    client_secret: z.string()
+  })
+]);
+
+/**
+ * Parse the conformance context from MCP_CONFORMANCE_CONTEXT env var.
+ */
+function parseContext() {
+  const raw = process.env.MCP_CONFORMANCE_CONTEXT;
+  if (!raw) {
+    throw new Error('MCP_CONFORMANCE_CONTEXT not set');
+  }
+  return ClientConformanceContextSchema.parse(JSON.parse(raw));
+}
+
+// Scenario handler type
+type ScenarioHandler = (serverUrl: string) => Promise<void>;
+
+// Registry of scenario handlers
+const scenarioHandlers: Record<string, ScenarioHandler> = {};
+
+// Helper to register a scenario handler
+function registerScenario(name: string, handler: ScenarioHandler): void {
+  scenarioHandlers[name] = handler;
+}
+
+// Helper to register multiple scenarios with the same handler
+function registerScenarios(names: string[], handler: ScenarioHandler): void {
+  for (const name of names) {
+    scenarioHandlers[name] = handler;
+  }
+}
+
+// ============================================================================
+// Basic scenarios (initialize, tools_call)
+// ============================================================================
+
+async function runBasicClient(serverUrl: string): Promise<void> {
+  const client = new Client(
+    { name: 'test-client', version: '1.0.0' },
+    { capabilities: {} }
+  );
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl));
+
+  await client.connect(transport);
+  logger.debug('Successfully connected to MCP server');
+
+  await client.listTools();
+  logger.debug('Successfully listed tools');
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+// tools_call scenario needs to actually call a tool
+async function runToolsCallClient(serverUrl: string): Promise<void> {
+  const client = new Client(
+    { name: 'test-client', version: '1.0.0' },
+    { capabilities: {} }
+  );
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl));
+
+  await client.connect(transport);
+  logger.debug('Successfully connected to MCP server');
+
+  const tools = await client.listTools();
+  logger.debug('Successfully listed tools');
+
+  // Call the add_numbers tool
+  const addTool = tools.tools.find((t) => t.name === 'add_numbers');
+  if (addTool) {
+    const result = await client.callTool({
+      name: 'add_numbers',
+      arguments: { a: 5, b: 3 }
+    });
+    logger.debug('Tool call result:', JSON.stringify(result, null, 2));
+  }
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+registerScenario('initialize', runBasicClient);
+registerScenario('tools_call', runToolsCallClient);
+
+// ============================================================================
+// Auth scenarios - well-behaved client
+// ============================================================================
+
+async function runAuthClient(serverUrl: string): Promise<void> {
+  const client = new Client(
+    { name: 'test-auth-client', version: '1.0.0' },
+    { capabilities: {} }
+  );
+
+  const oauthFetch = withOAuthRetry(
+    'test-auth-client',
+    new URL(serverUrl),
+    handle401,
+    CIMD_CLIENT_METADATA_URL
+  )(fetch);
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
+    fetch: oauthFetch
+  });
+
+  await client.connect(transport);
+  logger.debug('Successfully connected to MCP server');
+
+  await client.listTools();
+  logger.debug('Successfully listed tools');
+
+  await client.callTool({ name: 'test-tool', arguments: {} });
+  logger.debug('Successfully called tool');
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+// Register all auth scenarios that should use the well-behaved auth client
+// Note: client-credentials-jwt and client-credentials-basic have their own handlers below
+registerScenarios(
+  [
+    'auth/basic-cimd',
+    'auth/metadata-default',
+    'auth/metadata-var1',
+    'auth/metadata-var2',
+    'auth/metadata-var3',
+    'auth/2025-03-26-oauth-metadata-backcompat',
+    'auth/2025-03-26-oauth-endpoint-fallback',
+    'auth/scope-from-www-authenticate',
+    'auth/scope-from-scopes-supported',
+    'auth/scope-omitted-when-undefined',
+    'auth/scope-step-up',
+    'auth/scope-retry-limit',
+    'auth/token-endpoint-auth-basic',
+    'auth/token-endpoint-auth-post',
+    'auth/token-endpoint-auth-none'
+  ],
+  runAuthClient
+);
+
+// ============================================================================
+// Client Credentials scenarios
+// ============================================================================
+
+/**
+ * Client credentials with private_key_jwt authentication.
+ */
+async function runClientCredentialsJwt(serverUrl: string): Promise<void> {
+  const ctx = parseContext();
+  if (ctx.name !== 'auth/client-credentials-jwt') {
+    throw new Error(`Expected jwt context, got ${ctx.name}`);
+  }
+
+  const provider = new PrivateKeyJwtProvider({
+    clientId: ctx.client_id,
+    privateKey: ctx.private_key_pem,
+    algorithm: ctx.signing_algorithm || 'ES256'
+  });
+
+  const client = new Client(
+    { name: 'conformance-client-credentials-jwt', version: '1.0.0' },
+    { capabilities: {} }
+  );
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
+    authProvider: provider
+  });
+
+  await client.connect(transport);
+  logger.debug('Successfully connected with private_key_jwt auth');
+
+  await client.listTools();
+  logger.debug('Successfully listed tools');
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+registerScenario('auth/client-credentials-jwt', runClientCredentialsJwt);
+
+/**
+ * Client credentials with client_secret_basic authentication.
+ */
+async function runClientCredentialsBasic(serverUrl: string): Promise<void> {
+  const ctx = parseContext();
+  if (ctx.name !== 'auth/client-credentials-basic') {
+    throw new Error(`Expected basic context, got ${ctx.name}`);
+  }
+
+  const provider = new ClientCredentialsProvider({
+    clientId: ctx.client_id,
+    clientSecret: ctx.client_secret
+  });
+
+  const client = new Client(
+    { name: 'conformance-client-credentials-basic', version: '1.0.0' },
+    { capabilities: {} }
+  );
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
+    authProvider: provider
+  });
+
+  await client.connect(transport);
+  logger.debug('Successfully connected with client_secret_basic auth');
+
+  await client.listTools();
+  logger.debug('Successfully listed tools');
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+registerScenario('auth/client-credentials-basic', runClientCredentialsBasic);
+
+// ============================================================================
+// Elicitation defaults scenario
+// ============================================================================
+
+async function runElicitationDefaultsClient(serverUrl: string): Promise<void> {
+  const client = new Client(
+    { name: 'elicitation-defaults-test-client', version: '1.0.0' },
+    {
+      capabilities: {
+        elicitation: {
+          form: {
+            applyDefaults: true
+          }
+        }
+      }
+    }
+  );
+
+  // Register elicitation handler that returns empty content
+  // The SDK should fill in defaults for all omitted fields
+  client.setRequestHandler(ElicitRequestSchema, async (request) => {
+    logger.debug(
+      'Received elicitation request:',
+      JSON.stringify(request.params, null, 2)
+    );
+    logger.debug('Accepting with empty content - SDK should apply defaults');
+
+    // Return empty content - SDK should merge in defaults
+    return {
+      action: 'accept' as const,
+      content: {}
+    };
+  });
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl));
+
+  await client.connect(transport);
+  logger.debug('Successfully connected to MCP server');
+
+  // List available tools
+  const tools = await client.listTools();
+  logger.debug(
+    'Available tools:',
+    tools.tools.map((t) => t.name)
+  );
+
+  // Call the test tool which will trigger elicitation
+  const testTool = tools.tools.find(
+    (t) => t.name === 'test_client_elicitation_defaults'
+  );
+  if (!testTool) {
+    throw new Error('Test tool not found: test_client_elicitation_defaults');
+  }
+
+  logger.debug('Calling test_client_elicitation_defaults tool...');
+  const result = await client.callTool({
+    name: 'test_client_elicitation_defaults',
+    arguments: {}
+  });
+
+  logger.debug('Tool result:', JSON.stringify(result, null, 2));
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+registerScenario('elicitation-sep1034-client-defaults', runElicitationDefaultsClient);
+
+// ============================================================================
+// SSE retry scenario
+// ============================================================================
+
+async function runSSERetryClient(serverUrl: string): Promise<void> {
+  const client = new Client(
+    { name: 'sse-retry-test-client', version: '1.0.0' },
+    { capabilities: {} }
+  );
+
+  const transport = new StreamableHTTPClientTransport(new URL(serverUrl));
+
+  await client.connect(transport);
+  logger.debug('Successfully connected to MCP server');
+
+  // List tools to get the reconnection test tool
+  const tools = await client.listTools();
+  logger.debug(
+    'Available tools:',
+    tools.tools.map((t) => t.name)
+  );
+
+  // Call the test_reconnection tool which triggers stream closure
+  const testTool = tools.tools.find((t) => t.name === 'test_reconnection');
+  if (!testTool) {
+    throw new Error('Test tool not found: test_reconnection');
+  }
+
+  logger.debug('Calling test_reconnection tool...');
+  const result = await client.callTool({
+    name: 'test_reconnection',
+    arguments: {}
+  });
+
+  logger.debug('Tool result:', JSON.stringify(result, null, 2));
+
+  await transport.close();
+  logger.debug('Connection closed successfully');
+}
+
+registerScenario('sse-retry', runSSERetryClient);
+
+// ============================================================================
+// Main entry point
+// ============================================================================
+
+async function main(): Promise<void> {
+  const scenarioName = process.env.MCP_CONFORMANCE_SCENARIO;
+  const serverUrl = process.argv[2];
+
+  if (!scenarioName || !serverUrl) {
+    logger.error(
+      'Usage: MCP_CONFORMANCE_SCENARIO=<scenario> everything-client <server-url>'
+    );
+    logger.error(
+      '\nThe MCP_CONFORMANCE_SCENARIO env var is set automatically by the conformance runner.'
+    );
+    logger.error('\nAvailable scenarios:');
+    for (const name of Object.keys(scenarioHandlers).sort()) {
+      logger.error(`  - ${name}`);
+    }
+    process.exit(1);
+  }
+
+  const handler = scenarioHandlers[scenarioName];
+  if (!handler) {
+    logger.error(`Unknown scenario: ${scenarioName}`);
+    logger.error('\nAvailable scenarios:');
+    for (const name of Object.keys(scenarioHandlers).sort()) {
+      logger.error(`  - ${name}`);
+    }
+    process.exit(1);
+  }
+
+  try {
+    await handler(serverUrl);
+    process.exit(0);
+  } catch (error) {
+    logger.error('Error:', error);
+    process.exit(1);
+  }
+}
+
+main().catch((error) => {
+  logger.error('Unhandled error:', error);
+  process.exit(1);
+});
--- a/src/conformance/helpers/ConformanceOAuthProvider.ts
+++ b/src/conformance/helpers/ConformanceOAuthProvider.ts
@ -0,0 +1,95 @@
+import type {
+  OAuthClientProvider,
+  OAuthClientInformation,
+  OAuthClientInformationFull,
+  OAuthClientMetadata,
+  OAuthTokens
+} from '@modelcontextprotocol/client';
+
+export class ConformanceOAuthProvider implements OAuthClientProvider {
+  private _clientInformation?: OAuthClientInformationFull;
+  private _tokens?: OAuthTokens;
+  private _codeVerifier?: string;
+  private _authCode?: string;
+  private _authCodePromise?: Promise<string>;
+
+  constructor(
+    private readonly _redirectUrl: string | URL,
+    private readonly _clientMetadata: OAuthClientMetadata,
+    private readonly _clientMetadataUrl?: string | URL
+  ) {}
+
+  get redirectUrl(): string | URL {
+    return this._redirectUrl;
+  }
+
+  get clientMetadata(): OAuthClientMetadata {
+    return this._clientMetadata;
+  }
+
+  get clientMetadataUrl(): string | undefined {
+    return this._clientMetadataUrl?.toString();
+  }
+
+  clientInformation(): OAuthClientInformation | undefined {
+    return this._clientInformation;
+  }
+
+  saveClientInformation(clientInformation: OAuthClientInformationFull): void {
+    this._clientInformation = clientInformation;
+  }
+
+  tokens(): OAuthTokens | undefined {
+    return this._tokens;
+  }
+
+  saveTokens(tokens: OAuthTokens): void {
+    this._tokens = tokens;
+  }
+
+  async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
+    try {
+      const response = await fetch(authorizationUrl.toString(), {
+        redirect: 'manual' // Don't follow redirects automatically
+      });
+
+      // Get the Location header which contains the redirect with auth code
+      const location = response.headers.get('location');
+      if (location) {
+        const redirectUrl = new URL(location);
+        const code = redirectUrl.searchParams.get('code');
+        if (code) {
+          this._authCode = code;
+          return;
+        } else {
+          throw new Error('No auth code in redirect URL');
+        }
+      } else {
+        throw new Error(
+          `No redirect location received, from '${authorizationUrl.toString()}'`
+        );
+      }
+    } catch (error) {
+      console.error('Failed to fetch authorization URL:', error);
+      throw error;
+    }
+  }
+
+  async getAuthCode(): Promise<string> {
+    if (this._authCode) {
+      return this._authCode;
+    }
+    throw new Error('No authorization code');
+  }
+
+  saveCodeVerifier(codeVerifier: string): void {
+    this._codeVerifier = codeVerifier;
+  }
+
+  codeVerifier(): string {
+    if (!this._codeVerifier) {
+      throw new Error('No code verifier saved');
+    }
+    return this._codeVerifier;
+  }
+}
--- a/src/conformance/helpers/logger.ts
+++ b/src/conformance/helpers/logger.ts
@ -0,0 +1,27 @@
+/**
+ * Simple logger with configurable log levels.
+ * Set to 'error' in tests to suppress debug output.
+ */
+
+export type LogLevel = 'debug' | 'error';
+
+let currentLogLevel: LogLevel = 'debug';
+
+export function setLogLevel(level: LogLevel): void {
+  currentLogLevel = level;
+}
+
+export function getLogLevel(): LogLevel {
+  return currentLogLevel;
+}
+
+export const logger = {
+  debug: (...args: unknown[]): void => {
+    if (currentLogLevel === 'debug') {
+      console.log(...args);
+    }
+  },
+  error: (...args: unknown[]): void => {
+    console.error(...args);
+  }
+};
--- a/src/conformance/helpers/withOAuthRetry.ts
+++ b/src/conformance/helpers/withOAuthRetry.ts
@ -0,0 +1,111 @@
+import {
+  auth,
+  extractWWWAuthenticateParams,
+  UnauthorizedError
+} from '@modelcontextprotocol/client';
+import type { FetchLike, Middleware } from '@modelcontextprotocol/client';
+import { ConformanceOAuthProvider } from './ConformanceOAuthProvider.js';
+
+export const handle401 = async (
+  response: Response,
+  provider: ConformanceOAuthProvider,
+  next: FetchLike,
+  serverUrl: string | URL
+): Promise<void> => {
+  const { resourceMetadataUrl, scope } = extractWWWAuthenticateParams(response);
+  let result = await auth(provider, {
+    serverUrl,
+    resourceMetadataUrl,
+    scope,
+    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,
+      scope,
+      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,
+  clientMetadataUrl?: string
+): Middleware => {
+  const provider = new ConformanceOAuthProvider(
+    'http://localhost:3000/callback',
+    {
+      client_name: clientName,
+      redirect_uris: ['http://localhost:3000/callback']
+    },
+    clientMetadataUrl
+  );
+  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 || response.status === 403) {
+        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 || response.status === 403) {
+        const url = typeof input === 'string' ? input : input.toString();
+        throw new UnauthorizedError(`Authentication failed for ${url}`);
+      }
+
+      return response;
+    };
+  };
+};