fix callback validation in mcp client (#696)

* fix: callback urls were being matched by multiple servers

* feat: server url in state param

* fix: validate nonces

* remove legacy setting a client id manually

* rename function to validateState

* feat: add multiple servers to playground

* fix some ui

* fix: add check state and consume state rather than validate

* more tests

* changeset

* add cleanup of expired keys

* fix: changeset

* fix: readme to remove custom authProvider  wtffffff hahaha

Author: mattzcarey

.changeset/nine-bottles-heal.md

@@ -0,0 +1,7 @@
+---
+"agents": patch
+---
+
+Enables connecting to multiple MCP servers simultaneously and hardens OAuth state handling against replay/DoS attacks.
+
+**Note:** Inflight OAuth flows that were initiated on a previous version will not complete after upgrading, as the state parameter format has changed. Users will need to restart the authentication flow.

examples/mcp-client/README.md

@@ -22,34 +22,17 @@ Then, follow the steps below to setup the client:
 
 Tap "O + enter" to open the front end. It should list out all the tools, prompts, and resources available for each server added.
 
-## Transport Configuration
+## Usage
 
-The MCP client defaults to HTTP Streamable transport for better performance. You can specify transport type explicitly:
+The recommended way to add MCP servers is via `Agent.addMcpServer()`:
 
 ```typescript
-// Using MCPClientManager directly
-const mcpClient = new MCPClientManager("my-app", "1.0.0");
-
-// HTTP Streamable transport (default, recommended)
-await mcpClient.connect(serverUrl, {
-  transport: {
-    type: "streamable-http",
-    authProvider: myAuthProvider
-  }
-});
-
-// SSE transport (legacy compatibility)
-await mcpClient.connect(serverUrl, {
-  transport: {
-    type: "sse",
-    authProvider: myAuthProvider
-  }
-});
-
-// Or using Agent.addMcpServer() (as shown in the example)
 export class MyAgent extends Agent<Env, never> {
   async addServer(name: string, url: string, callbackHost: string) {
+    // Uses HTTP Streamable transport by default
     await this.addMcpServer(name, url, callbackHost);
   }
 }
 ```
+
+The MCP client handles OAuth authentication automatically using the built-in `DurableObjectOAuthClientProvider`.

packages/agents/src/mcp/client.ts

@@ -635,46 +635,81 @@ export class MCPClientManager {
     }
   }
 
+  private extractServerIdFromState(state: string | null): string | null {
+    if (!state) return null;
+    const parts = state.split(".");
+    return parts.length === 2 ? parts[1] : null;
+  }
+
   isCallbackRequest(req: Request): boolean {
     if (req.method !== "GET") {
       return false;
     }
 
-    // Quick heuristic check: most callback URLs contain "/callback"
-    // This avoids DB queries for obviously non-callback requests
     if (!req.url.includes("/callback")) {
       return false;
     }
 
-    // Check database for matching callback URL
+    const url = new URL(req.url);
+    const state = url.searchParams.get("state");
+    const serverId = this.extractServerIdFromState(state);
+    if (!serverId) {
+      return false;
+    }
+
     const servers = this.getServersFromStorage();
-    return servers.some(
-      (server) => server.callback_url && req.url.startsWith(server.callback_url)
-    );
+    return servers.some((server) => server.id === serverId);
   }
 
   async handleCallbackRequest(req: Request) {
     const url = new URL(req.url);
+    const code = url.searchParams.get("code");
+    const state = url.searchParams.get("state");
+    const error = url.searchParams.get("error");
+    const errorDescription = url.searchParams.get("error_description");
+
+    if (!state) {
+      throw new Error("Unauthorized: no state provided");
+    }
+
+    const serverId = this.extractServerIdFromState(state);
+
+    if (!serverId) {
+      throw new Error(
+        "No serverId found in state parameter. Expected format: {nonce}.{serverId}"
+      );
+    }
 
-    // Find the matching server from database
     const servers = this.getServersFromStorage();
-    const matchingServer = servers.find((server: MCPServerRow) => {
-      return server.callback_url && req.url.startsWith(server.callback_url);
-    });
+    const serverExists = servers.some((server) => server.id === serverId);
 
-    if (!matchingServer) {
+    if (!serverExists) {
       throw new Error(
-        `No callback URI match found for the request url: ${req.url}. Was the request matched with \`isCallbackRequest()\`?`
+        `No server found with id "${serverId}". Was the request matched with \`isCallbackRequest()\`?`
       );
     }
 
-    const serverId = matchingServer.id;
-    const code = url.searchParams.get("code");
-    const state = url.searchParams.get("state");
-    const error = url.searchParams.get("error");
-    const errorDescription = url.searchParams.get("error_description");
+    if (this.mcpConnections[serverId] === undefined) {
+      throw new Error(`Could not find serverId: ${serverId}`);
+    }
+
+    const conn = this.mcpConnections[serverId];
+    if (!conn.options.transport.authProvider) {
+      throw new Error(
+        "Trying to finalize authentication for a server connection without an authProvider"
+      );
+    }
+
+    const authProvider = conn.options.transport.authProvider;
+    authProvider.serverId = serverId;
+
+    // Two-phase state validation: check first (non-destructive), consume later
+    // This prevents DoS attacks where attacker consumes valid state before legitimate callback
+    const stateValidation = await authProvider.checkState(state);
+    if (!stateValidation.valid) {
+      throw new Error(`Invalid state: ${stateValidation.error}`);
+    }
 
-    // Handle OAuth error responses from the provider
     if (error) {
       return {
         serverId,
@@ -686,25 +721,14 @@ export class MCPClientManager {
     if (!code) {
       throw new Error("Unauthorized: no code provided");
     }
-    if (!state) {
-      throw new Error("Unauthorized: no state provided");
-    }
 
-    if (this.mcpConnections[serverId] === undefined) {
-      throw new Error(`Could not find serverId: ${serverId}`);
-    }
-
-    // If connection is already ready/connected, this is likely a duplicate callback
     if (
       this.mcpConnections[serverId].connectionState ===
         MCPConnectionState.READY ||
       this.mcpConnections[serverId].connectionState ===
         MCPConnectionState.CONNECTED
     ) {
-      // make sure auth_url is cleared
       this.clearServerAuthUrl(serverId);
-
-      // Already authenticated and ready, treat as success
       return {
         serverId,
         authSuccess: true
@@ -720,32 +744,20 @@ export class MCPClientManager {
       );
     }
 
-    const conn = this.mcpConnections[serverId];
-    if (!conn.options.transport.authProvider) {
-      throw new Error(
-        "Trying to finalize authentication for a server connection without an authProvider"
-      );
-    }
-
-    // Get clientId from auth provider (stored during redirectToAuthorization) or fallback to state for backward compatibility
-    const clientId = conn.options.transport.authProvider.clientId || state;
-
-    // Set the OAuth credentials
-    conn.options.transport.authProvider.clientId = clientId;
-    conn.options.transport.authProvider.serverId = serverId;
-
     try {
+      await authProvider.consumeState(state);
       await conn.completeAuthorization(code);
+      await authProvider.deleteCodeVerifier();
       this.clearServerAuthUrl(serverId);
       this._onServerStateChanged.fire();
 
       return {
         serverId,
         authSuccess: true
       };
-    } catch (error) {
+    } catch (authError) {
       const errorMessage =
-        error instanceof Error ? error.message : String(error);
+        authError instanceof Error ? authError.message : String(authError);
 
       this._onServerStateChanged.fire();
 

packages/agents/src/mcp/do-oauth-client-provider.ts

@@ -7,12 +7,25 @@ import type {
 } from "@modelcontextprotocol/sdk/shared/auth.js";
 import { nanoid } from "nanoid";
 
+const STATE_EXPIRATION_MS = 10 * 60 * 1000; // 10 minutes
+
+interface StoredState {
+  nonce: string;
+  serverId: string;
+  createdAt: number;
+}
+
 // A slight extension to the standard OAuthClientProvider interface because `redirectToAuthorization` doesn't give us the interface we need
 // This allows us to track authentication for a specific server and associated dynamic client registration
 export interface AgentsOAuthProvider extends OAuthClientProvider {
   authUrl: string | undefined;
   clientId: string | undefined;
   serverId: string | undefined;
+  checkState(
+    state: string
+  ): Promise<{ valid: boolean; serverId?: string; error?: string }>;
+  consumeState(state: string): Promise<void>;
+  deleteCodeVerifier(): Promise<void>;
 }
 
 export class DurableObjectOAuthClientProvider implements AgentsOAuthProvider {
@@ -48,7 +61,7 @@ export class DurableObjectOAuthClientProvider implements AgentsOAuthProvider {
   }
 
   get redirectUrl() {
-    return `${this.baseRedirectUrl}/${this.serverId}`;
+    return this.baseRedirectUrl;
   }
 
   get clientId() {
@@ -124,17 +137,92 @@ export class DurableObjectOAuthClientProvider implements AgentsOAuthProvider {
     return this._authUrl_;
   }
 
-  /**
-   * Because this operates on the server side (but we need browser auth), we send this url back to the user
-   * and require user interact to initiate the redirect flow
-   */
+  stateKey(nonce: string) {
+    return `/${this.clientName}/${this.serverId}/state/${nonce}`;
+  }
+
+  async state(): Promise<string> {
+    const nonce = nanoid();
+    const state = `${nonce}.${this.serverId}`;
+    const storedState: StoredState = {
+      nonce,
+      serverId: this.serverId,
+      createdAt: Date.now()
+    };
+    await this.storage.put(this.stateKey(nonce), storedState);
+    return state;
+  }
+
+  async checkState(
+    state: string
+  ): Promise<{ valid: boolean; serverId?: string; error?: string }> {
+    const parts = state.split(".");
+    if (parts.length !== 2) {
+      return { valid: false, error: "Invalid state format" };
+    }
+
+    const [nonce, serverId] = parts;
+    const key = this.stateKey(nonce);
+    const storedState = await this.storage.get<StoredState>(key);
+
+    if (!storedState) {
+      return { valid: false, error: "State not found or already used" };
+    }
+
+    if (storedState.serverId !== serverId) {
+      await this.storage.delete(key);
+      return { valid: false, error: "State serverId mismatch" };
+    }
+
+    const age = Date.now() - storedState.createdAt;
+    if (age > STATE_EXPIRATION_MS) {
+      await this.storage.delete(key);
+      return { valid: false, error: "State expired" };
+    }
+
+    return { valid: true, serverId };
+  }
+
+  async consumeState(state: string): Promise<void> {
+    const parts = state.split(".");
+    if (parts.length !== 2) {
+      // This should never happen since checkState validates format first.
+      // Log for debugging but don't throw - state consumption is best-effort.
+      console.warn(
+        `[OAuth] consumeState called with invalid state format: ${state.substring(0, 20)}...`
+      );
+      return;
+    }
+    const [nonce] = parts;
+    await this.storage.delete(this.stateKey(nonce));
+  }
+
   async redirectToAuthorization(authUrl: URL): Promise<void> {
-    // Generate secure random token for state parameter
-    const stateToken = nanoid();
-    authUrl.searchParams.set("state", stateToken);
     this._authUrl_ = authUrl.toString();
   }
 
+  async invalidateCredentials(
+    scope: "all" | "client" | "tokens" | "verifier"
+  ): Promise<void> {
+    if (!this._clientId_) return;
+
+    const deleteKeys: string[] = [];
+
+    if (scope === "all" || scope === "client") {
+      deleteKeys.push(this.clientInfoKey(this.clientId));
+    }
+    if (scope === "all" || scope === "tokens") {
+      deleteKeys.push(this.tokenKey(this.clientId));
+    }
+    if (scope === "all" || scope === "verifier") {
+      deleteKeys.push(this.codeVerifierKey(this.clientId));
+    }
+
+    if (deleteKeys.length > 0) {
+      await this.storage.delete(deleteKeys);
+    }
+  }
+
   codeVerifierKey(clientId: string) {
     return `${this.keyPrefix(clientId)}/code_verifier`;
   }
@@ -160,4 +248,8 @@ export class DurableObjectOAuthClientProvider implements AgentsOAuthProvider {
     }
     return codeVerifier;
   }
+
+  async deleteCodeVerifier(): Promise<void> {
+    await this.storage.delete(this.codeVerifierKey(this.clientId));
+  }
 }