fix(mcp): harden OAuth callback recovery

.changeset/mcp-oauth-callback-robustness.md

@@ -0,0 +1,9 @@
+---
+"agents": patch
+---
+
+Harden the MCP client OAuth callback state machine against stray callbacks and stale auth URLs.
+
+Previously, a stray or invalid GET to the OAuth callback URL carrying a well-formed but unverifiable state nonce — with either an `error` or a `code` param — flipped an in-flight `authenticating` connection to `failed` and cleared its `auth_url`, after which the user's genuine callback was rejected and the authorization could never complete for the lifetime of the Durable Object. `handleCallbackRequest` now verifies the state nonce via `authProvider.checkState()` before failing the connection: callbacks whose nonce cannot be verified are logged and answered with an error without touching the connection state, and a genuine callback can still complete a connection that was spuriously moved to `failed`.
+
+Separately, `addMcpServer` on an `authenticating` connection kept returning an auth URL whose embedded state nonce had expired (the nonce TTL is 10 minutes) — both the persisted URL after a hibernation and the connection's live in-memory URL when the flow simply sat idle — so an OAuth flow not completed within the TTL became unrecoverable. An auth URL is now only returned while its nonce is still redeemable; otherwise `addMcpServer` reconnects and mints a fresh one.

packages/agents/src/index.ts

@@ -12264,12 +12264,18 @@ export class Agent<
     if (existingServer && this.mcp.mcpConnections[existingServer.id]) {
       const conn = this.mcp.mcpConnections[existingServer.id];
       if (conn.connectionState === MCPConnectionState.AUTHENTICATING) {
-        const liveAuthUrl = conn.options.transport.authProvider?.authUrl;
+        const authProvider = conn.options.transport.authProvider;
         const authUrl =
-          liveAuthUrl ||
-          (this._isAbsoluteHttpUrl(existingServer.auth_url)
-            ? existingServer.auth_url
-            : undefined);
+          (await this._redeemableAuthUrl(
+            existingServer.id,
+            authProvider?.authUrl,
+            authProvider
+          )) ||
+          (await this._redeemableAuthUrl(
+            existingServer.id,
+            existingServer.auth_url,
+            authProvider
+          ));
         if (authUrl) {
           return {
             id: existingServer.id,
@@ -12533,6 +12539,33 @@ export class Agent<
     }
   }
 
+  /**
+   * An auth URL embeds a one-time state nonce with a limited TTL, so both a
+   * connection's live in-memory URL and a row restored after hibernation can
+   * go stale before the flow completes. Serve a URL only while its nonce is
+   * still redeemable; otherwise the caller reconnects and mints a fresh one.
+   */
+  private async _redeemableAuthUrl(
+    serverId: string,
+    authUrl: string | null | undefined,
+    authProvider: AgentMcpOAuthProvider | undefined
+  ): Promise<string | undefined> {
+    if (!this._isAbsoluteHttpUrl(authUrl) || !authProvider) {
+      return undefined;
+    }
+    const state = new URL(authUrl).searchParams.get("state");
+    if (!state) {
+      return undefined;
+    }
+    authProvider.serverId = serverId;
+    try {
+      const stateValidation = await authProvider.checkState(state);
+      return stateValidation.valid ? authUrl : undefined;
+    } catch {
+      return undefined;
+    }
+  }
+
   async removeMcpServer(id: string) {
     await this.mcp.removeServer(id);
   }

packages/agents/src/mcp/client.ts

@@ -731,6 +731,25 @@ export class MCPClientManager {
     return callback();
   }
 
+  /**
+   * Verify that a callback's state nonce is still redeemable for this server.
+   * The provider may clean up expired or mismatched state while checking it;
+   * valid state remains available until `consumeState` runs.
+   */
+  private async isRedeemableOAuthState(
+    serverId: string,
+    authProvider: AgentMcpOAuthProvider,
+    state: string
+  ): Promise<boolean> {
+    authProvider.serverId = serverId;
+    try {
+      const stateValidation = await authProvider.checkState(state);
+      return stateValidation.valid;
+    } catch {
+      return false;
+    }
+  }
+
   private async consumeStaleOAuthState(
     serverId: string,
     authProvider: AgentMcpOAuthProvider,
@@ -1499,6 +1518,28 @@ export class MCPClientManager {
           }
           return this.oauthCallbackSuccess(validation.serverId, conn);
         }
+        // Only a callback carrying a genuine state nonce may fail the
+        // connection; a stray or invalid callback must not alter the
+        // connection state machine.
+        const authProvider = conn.options.transport.authProvider;
+        if (
+          validation.state &&
+          authProvider &&
+          !(await this.isRedeemableOAuthState(
+            validation.serverId,
+            authProvider,
+            validation.state
+          ))
+        ) {
+          console.warn(
+            `[MCPClientManager] Ignoring OAuth callback with unverified state for server "${validation.serverId}": ${validation.error}`
+          );
+          return {
+            serverId: validation.serverId,
+            authSuccess: false,
+            authError: validation.error
+          };
+        }
         return this.failConnection(validation.serverId, validation.error);
       }
 
@@ -1530,7 +1571,16 @@ export class MCPClientManager {
           await this.consumeStaleOAuthState(serverId, authProvider, state);
           return this.oauthCallbackSuccess(serverId, conn);
         }
-        throw new Error(stateValidation.error || "Invalid state");
+        // Same rule as the invalid branch above: a callback whose nonce
+        // cannot be verified must not alter the connection state machine.
+        console.warn(
+          `[MCPClientManager] Ignoring OAuth callback with unverified state for server "${serverId}": ${stateValidation.error ?? "Invalid state"}`
+        );
+        return {
+          serverId,
+          authSuccess: false,
+          authError: stateValidation.error || "Invalid state"
+        };
       }
 
       // A stale popup can complete after another callback already exchanged tokens.
@@ -1540,12 +1590,10 @@ export class MCPClientManager {
         return this.oauthCallbackSuccess(serverId, conn);
       }
 
-      if (conn.connectionState !== MCPConnectionState.AUTHENTICATING) {
-        throw new Error(
-          `Failed to authenticate: the client is in "${conn.connectionState}" state, expected "authenticating"`
-        );
-      }
-
+      // The auth-accepted states returned above. The two remaining states are
+      // AUTHENTICATING and FAILED; both may complete once checkState has
+      // proved the callback genuine. Accepting FAILED lets a valid callback
+      // recover a flow that an earlier stray callback disrupted.
       conn.connectionState = MCPConnectionState.CONNECTING;
       await authProvider.consumeState(state);
       await this.completeAuthorizationAndCleanupVerifier(

packages/agents/src/tests/agents/oauth.ts

@@ -136,6 +136,9 @@ export class TestOAuthAgent extends Agent {
               const [nonce] = parts;
               self.mockStateStorage.delete(nonce);
             },
+            async redirectToAuthorization(authUrl: URL): Promise<void> {
+              this.authUrl = authUrl.toString();
+            },
             async deleteCodeVerifier(): Promise<void> {
               // No-op for tests
             }
@@ -155,6 +158,41 @@ export class TestOAuthAgent extends Agent {
     this.mockStateStorage.set(nonce, { serverId, createdAt: Date.now() });
   }
 
+  // Sets the connection's live in-memory auth URL the same way a real OAuth
+  // redirect does, so tests can control how old its embedded nonce is.
+  async setLiveAuthUrlForTest(
+    serverId: string,
+    authUrl: string
+  ): Promise<void> {
+    const authProvider =
+      this.mcp.mcpConnections[serverId]?.options.transport.authProvider;
+    if (!authProvider) {
+      throw new Error(`Test error: OAuth provider ${serverId} not found`);
+    }
+    await authProvider.redirectToAuthorization(new URL(authUrl));
+  }
+
+  // Seeds a persisted OAuth state row under the same storage key the agent's
+  // real DurableObjectOAuthClientProvider reads, so tests can control how old
+  // the nonce embedded in a persisted auth_url is.
+  async seedPersistedOAuthState(
+    serverId: string,
+    nonce: string,
+    ageMs = 0
+  ): Promise<void> {
+    const provider = new DurableObjectOAuthClientProvider(
+      this.ctx.storage,
+      this.name,
+      "http://example.com/oauth/callback"
+    );
+    provider.serverId = serverId;
+    await this.ctx.storage.put(provider.stateKey(nonce), {
+      nonce,
+      serverId,
+      createdAt: Date.now() - ageMs
+    });
+  }
+
   setupMockMcpConnection(
     serverId: string,
     serverName: string,

packages/agents/src/tests/mcp/client-manager.test.ts

@@ -679,7 +679,7 @@ describe("MCPClientManager OAuth Integration", () => {
       expect(result.serverId).toBe(serverId);
     });
 
-    it("should fail connection when callback received for connection in failed state", async () => {
+    it("should complete a genuine callback received for a connection in failed state", async () => {
       const serverId = "test-server";
       const callbackUrl = "http://localhost:3000/callback";
       const stateStorage = createMockStateStorage();
@@ -706,7 +706,15 @@ describe("MCPClientManager OAuth Integration", () => {
 
       connection.init = vi.fn().mockResolvedValue(undefined);
       connection.client.close = vi.fn().mockResolvedValue(undefined);
+      // A stray or invalid callback may have spuriously failed the
+      // connection; a callback carrying a genuine nonce still completes.
       connection.connectionState = "failed";
+      connection.connectionError = "spurious failure";
+      const completeAuthSpy = vi
+        .spyOn(connection, "completeAuthorization")
+        .mockImplementation(async () => {
+          connection.connectionState = "connecting";
+        });
 
       manager.mcpConnections[serverId] = connection;
 
@@ -716,10 +724,11 @@ describe("MCPClientManager OAuth Integration", () => {
       );
 
       const result = await manager.handleCallbackRequest(callbackRequest);
-      expect(result.authSuccess).toBe(false);
-      expect(result.authError).toBe(
-        'Failed to authenticate: the client is in "failed" state, expected "authenticating"'
-      );
+      expect(result.authSuccess).toBe(true);
+      expect(completeAuthSpy).toHaveBeenCalledWith("test", {
+        alreadyAccepted: true
+      });
+      expect(connection.connectionError).toBe(null);
     });
 
     it("should recognize custom callback paths that do not contain '/callback'", async () => {
@@ -759,6 +768,104 @@ describe("MCPClientManager OAuth Integration", () => {
     });
   });
 
+  describe("OAuth Callback Robustness", () => {
+    const serverId = "test-server";
+    const callbackUrl = "http://localhost:3000/callback";
+    const authUrl = "https://auth.example.com/authorize";
+
+    function setupAuthenticatingConnection(
+      stateStorage: ReturnType<typeof createMockStateStorage>
+    ) {
+      saveServerToMock({
+        id: serverId,
+        name: "Test Server",
+        server_url: "http://test.com",
+        callback_url: callbackUrl,
+        client_id: null,
+        auth_url: authUrl,
+        server_options: null
+      });
+
+      const mockAuthProvider = createMockAuthProvider(stateStorage);
+      const connection = new MCPClientConnection(
+        new URL("http://test.com"),
+        { name: "test-client", version: "1.0.0" },
+        {
+          transport: { type: "auto", authProvider: mockAuthProvider },
+          client: {}
+        }
+      );
+      connection.init = vi.fn().mockResolvedValue(undefined);
+      connection.client.close = vi.fn().mockResolvedValue(undefined);
+      connection.connectionState = "authenticating";
+      manager.mcpConnections[serverId] = connection;
+      return connection;
+    }
+
+    it("should ignore a stray error callback whose state nonce was never issued", async () => {
+      const stateStorage = createMockStateStorage();
+      const connection = setupAuthenticatingConnection(stateStorage);
+
+      // A genuine flow is in flight when a stray callback arrives carrying a
+      // well-formed nonce that was never issued.
+      stateStorage.createState(serverId);
+      const strayState = `${nanoid()}.${serverId}`;
+      const result = await manager.handleCallbackRequest(
+        new Request(`${callbackUrl}?error=access_denied&state=${strayState}`)
+      );
+
+      expect(result.authSuccess).toBe(false);
+      expect(result.authError).toBe("access_denied");
+      expect(connection.connectionState).toBe("authenticating");
+      expect(mockStorageData.get(serverId)?.auth_url).toBe(authUrl);
+    });
+
+    it("should ignore a stray code callback whose state nonce was never issued", async () => {
+      const stateStorage = createMockStateStorage();
+      const connection = setupAuthenticatingConnection(stateStorage);
+
+      // Same stray-callback shape, but carrying a code param instead of an
+      // error param so it reaches the valid-callback path.
+      stateStorage.createState(serverId);
+      const strayState = `${nanoid()}.${serverId}`;
+      const result = await manager.handleCallbackRequest(
+        new Request(`${callbackUrl}?code=x&state=${strayState}`)
+      );
+
+      expect(result.authSuccess).toBe(false);
+      expect(result.authError).toBe("State not found or already used");
+      expect(connection.connectionState).toBe("authenticating");
+      expect(mockStorageData.get(serverId)?.auth_url).toBe(authUrl);
+    });
+
+    it("should complete a genuine callback after a stray code callback arrived first", async () => {
+      const stateStorage = createMockStateStorage();
+      const connection = setupAuthenticatingConnection(stateStorage);
+      const completeAuthSpy = vi
+        .spyOn(connection, "completeAuthorization")
+        .mockImplementation(async () => {
+          connection.connectionState = "connecting";
+        });
+
+      const state = stateStorage.createState(serverId);
+
+      const strayResult = await manager.handleCallbackRequest(
+        new Request(`${callbackUrl}?code=x&state=${nanoid()}.${serverId}`)
+      );
+      expect(strayResult.authSuccess).toBe(false);
+      expect(connection.connectionState).toBe("authenticating");
+
+      const result = await manager.handleCallbackRequest(
+        new Request(`${callbackUrl}?code=auth-code&state=${state}`)
+      );
+
+      expect(result.authSuccess).toBe(true);
+      expect(completeAuthSpy).toHaveBeenCalledWith("auth-code", {
+        alreadyAccepted: true
+      });
+    });
+  });
+
   describe("OAuth Security", () => {
     it("should clear auth_url but preserve callback_url after successful authentication", async () => {
       const serverId = "test-server";
@@ -857,6 +964,7 @@ describe("MCPClientManager OAuth Integration", () => {
       const result2 = await manager.handleCallbackRequest(callbackRequest2);
       expect(result2.authSuccess).toBe(false);
       expect(result2.authError).toBe("State not found or already used");
+      expect(connection.connectionState).toBe("authenticating");
     });
 
     it("should reject expired state (10 minute TTL)", async () => {
@@ -896,6 +1004,7 @@ describe("MCPClientManager OAuth Integration", () => {
       const result = await manager.handleCallbackRequest(callbackRequest);
       expect(result.authSuccess).toBe(false);
       expect(result.authError).toBe("State expired");
+      expect(connection.connectionState).toBe("authenticating");
     });
 
     it("should only match callbacks with valid state for existing servers", async () => {