fix: MCP tools intermittently unavailable after hibernation (#928)

.changeset/fix-mcp-tools-race-condition.md

@@ -0,0 +1,10 @@
+---
+"agents": patch
+"@cloudflare/ai-chat": patch
+---
+
+Fix race condition where MCP tools are intermittently unavailable in onChatMessage after hibernation.
+
+**`agents`**: Added `MCPClientManager.waitForConnections(options?)` which awaits all in-flight connection and discovery operations. Accepts an optional `{ timeout }` in milliseconds. Background restore promises from `restoreConnectionsFromStorage()` are now tracked so callers can wait for them to settle.
+
+**`@cloudflare/ai-chat`**: Added `waitForMcpConnections` opt-in config on `AIChatAgent`. Set to `true` to wait indefinitely, or `{ timeout: 10_000 }` to cap the wait. Default is `false` (non-blocking, preserving existing behavior). For lower-level control, call `this.mcp.waitForConnections()` directly in your `onChatMessage`.

examples/ai-chat/src/server.ts

@@ -1,5 +1,5 @@
 import { createWorkersAI } from "workers-ai-provider";
-import { routeAgentRequest } from "agents";
+import { routeAgentRequest, callable } from "agents";
 import { AIChatAgent, type OnChatMessageOptions } from "@cloudflare/ai-chat";
 import {
   streamText,
@@ -23,7 +23,39 @@ export class ChatAgent extends AIChatAgent {
   // Keep the last 200 messages in SQLite storage
   maxPersistedMessages = 200;
 
+  // Wait for MCP connections to restore after hibernation before processing messages
+  waitForMcpConnections = true;
+
+  onStart() {
+    // Configure OAuth popup behavior for MCP servers that require authentication
+    this.mcp.configureOAuthCallback({
+      customHandler: (result) => {
+        if (result.authSuccess) {
+          return new Response("<script>window.close();</script>", {
+            headers: { "content-type": "text/html" },
+            status: 200
+          });
+        }
+        return new Response(
+          `Authentication Failed: ${result.authError || "Unknown error"}`,
+          { headers: { "content-type": "text/plain" }, status: 400 }
+        );
+      }
+    });
+  }
+
+  @callable()
+  async addServer(name: string, url: string, host: string) {
+    return await this.addMcpServer(name, url, { callbackHost: host });
+  }
+
+  @callable()
+  async removeServer(serverId: string) {
+    await this.removeMcpServer(serverId);
+  }
+
   async onChatMessage(_onFinish: unknown, options?: OnChatMessageOptions) {
+    const mcpTools = this.mcp.getAITools();
     const workersai = createWorkersAI({ binding: this.env.AI });
 
     const result = streamText({
@@ -39,6 +71,9 @@ export class ChatAgent extends AIChatAgent {
         reasoning: "before-last-message"
       }),
       tools: {
+        // MCP tools from connected servers
+        ...mcpTools,
+
         // Server-side tool: executes automatically
         getWeather: tool({
           description: "Get the current weather for a city",

packages/agents/src/mcp/client.ts

@@ -192,6 +192,7 @@ export class MCPClientManager {
     callbackUrl: string
   ) => AgentMcpOAuthProvider;
   private _isRestored = false;
+  private _pendingConnections = new Map<string, Promise<void>>();
 
   /** @internal Protected for testing purposes. */
   protected readonly _onObservabilityEvent =
@@ -454,12 +455,61 @@ export class MCPClientManager {
       }
 
       // Start connection in background (don't await) to avoid blocking the DO
-      this._restoreServer(server.id, parsedOptions?.retry);
+      this._trackConnection(
+        server.id,
+        this._restoreServer(server.id, parsedOptions?.retry)
+      );
     }
 
     this._isRestored = true;
   }
 
+  /**
+   * Track a pending connection promise for a server.
+   * The promise is removed from the map when it settles.
+   */
+  private _trackConnection(serverId: string, promise: Promise<void>): void {
+    const tracked = promise.finally(() => {
+      // Only delete if it's still the same promise (not replaced by a newer one)
+      if (this._pendingConnections.get(serverId) === tracked) {
+        this._pendingConnections.delete(serverId);
+      }
+    });
+    this._pendingConnections.set(serverId, tracked);
+  }
+
+  /**
+   * Wait for all in-flight connection and discovery operations to settle.
+   * This is useful when you need MCP tools to be available before proceeding,
+   * e.g. before calling getAITools() after the agent wakes from hibernation.
+   *
+   * Returns once every pending connection has either connected and discovered,
+   * failed, or timed out. Never rejects.
+   *
+   * @param options.timeout - Maximum time in milliseconds to wait.
+   *   `0` returns immediately without waiting.
+   *   `undefined` (default) waits indefinitely.
+   */
+  async waitForConnections(options?: { timeout?: number }): Promise<void> {
+    if (this._pendingConnections.size === 0) {
+      return;
+    }
+    if (options?.timeout != null && options.timeout <= 0) {
+      return;
+    }
+    const settled = Promise.allSettled(this._pendingConnections.values());
+    if (options?.timeout != null && options.timeout > 0) {
+      let timerId: ReturnType<typeof setTimeout>;
+      const timer = new Promise<void>((resolve) => {
+        timerId = setTimeout(resolve, options.timeout);
+      });
+      await Promise.race([settled, timer]);
+      clearTimeout(timerId!);
+    } else {
+      await settled;
+    }
+  }
+
   /**
    * Internal method to restore a single server connection and discovery
    */
@@ -1024,11 +1074,19 @@ export class MCPClientManager {
   }
 
   /**
-   * Establish connection in the background after OAuth completion
-   * This method connects to the server and discovers its capabilities
+   * Establish connection in the background after OAuth completion.
+   * This method connects to the server and discovers its capabilities.
+   * The connection is automatically tracked so that `waitForConnections()`
+   * will include it.
    * @param serverId The server ID to establish connection for
    */
   async establishConnection(serverId: string): Promise<void> {
+    const promise = this._doEstablishConnection(serverId);
+    this._trackConnection(serverId, promise);
+    return promise;
+  }
+
+  private async _doEstablishConnection(serverId: string): Promise<void> {
     const conn = this.mcpConnections[serverId];
     if (!conn) {
       this._onObservabilityEvent.fire({
@@ -1221,6 +1279,9 @@ export class MCPClientManager {
   async closeAllConnections() {
     const ids = Object.keys(this.mcpConnections);
 
+    // Clear all pending connection tracking
+    this._pendingConnections.clear();
+
     // Cancel all in-flight discoveries
     for (const id of ids) {
       this.mcpConnections[id].cancelDiscovery();
@@ -1252,6 +1313,9 @@ export class MCPClientManager {
     // Cancel any in-flight discovery
     this.mcpConnections[id].cancelDiscovery();
 
+    // Remove from pending so waitForConnections() doesn't block on a closed server
+    this._pendingConnections.delete(id);
+
     await this.mcpConnections[id].client.close();
     delete this.mcpConnections[id];
 

packages/agents/src/tests/agents/index.ts

@@ -33,3 +33,4 @@ export {
   TestSessionAgentNoMicroCompaction,
   TestSessionAgentCustomRules
 } from "./session";
+export { TestWaitConnectionsAgent } from "./wait-connections";

packages/agents/src/tests/agents/wait-connections.ts

@@ -0,0 +1,167 @@
+import { Agent } from "../../index.ts";
+
+/**
+ * Test Agent that exposes waitForConnections() for E2E testing.
+ * Simulates the full hibernation → restore → wait → getAITools flow.
+ */
+export class TestWaitConnectionsAgent extends Agent<Record<string, unknown>> {
+  observability = undefined;
+
+  async onRequest(_request: Request): Promise<Response> {
+    return new Response("TestWaitConnectionsAgent");
+  }
+
+  /**
+   * Insert an MCP server row into SQLite (simulates pre-hibernation state).
+   */
+  insertMcpServer(
+    serverId: string,
+    name: string,
+    serverUrl: string,
+    callbackUrl: string,
+    authUrl: string | null
+  ): void {
+    this.sql`
+      INSERT OR REPLACE INTO cf_agents_mcp_servers (
+        id, name, server_url, client_id, auth_url, callback_url, server_options
+      ) VALUES (
+        ${serverId},
+        ${name},
+        ${serverUrl},
+        ${null},
+        ${authUrl},
+        ${callbackUrl},
+        ${null}
+      )
+    `;
+  }
+
+  /**
+   * Reset the restored flag so restoreConnectionsFromStorage can run again.
+   */
+  resetRestoredFlag(): void {
+    // @ts-expect-error - accessing private property for testing
+    this.mcp._isRestored = false;
+    // Clear existing connections
+    for (const id of Object.keys(this.mcp.mcpConnections)) {
+      delete this.mcp.mcpConnections[id];
+    }
+  }
+
+  /**
+   * Trigger restore and then immediately wait for all connections to settle.
+   * Returns info about the resulting connection states.
+   */
+  async restoreAndWait(timeout?: number): Promise<{
+    connectionIds: string[];
+    connectionStates: Record<string, string>;
+  }> {
+    await this.mcp.restoreConnectionsFromStorage(this.name);
+    await this.mcp.waitForConnections(
+      timeout != null ? { timeout } : undefined
+    );
+
+    const connectionStates: Record<string, string> = {};
+    for (const [id, conn] of Object.entries(this.mcp.mcpConnections)) {
+      connectionStates[id] = conn.connectionState;
+    }
+
+    return {
+      connectionIds: Object.keys(this.mcp.mcpConnections),
+      connectionStates
+    };
+  }
+
+  /**
+   * Trigger restore WITHOUT waiting, then immediately check states.
+   * This simulates the race condition (old behavior).
+   */
+  async restoreWithoutWait(): Promise<{
+    connectionIds: string[];
+    connectionStates: Record<string, string>;
+  }> {
+    await this.mcp.restoreConnectionsFromStorage(this.name);
+    // No waitForConnections — check states immediately
+
+    const connectionStates: Record<string, string> = {};
+    for (const [id, conn] of Object.entries(this.mcp.mcpConnections)) {
+      connectionStates[id] = conn.connectionState;
+    }
+
+    return {
+      connectionIds: Object.keys(this.mcp.mcpConnections),
+      connectionStates
+    };
+  }
+
+  /**
+   * Check if waitForConnections resolves when there are no pending connections.
+   */
+  async waitWithNoPending(): Promise<boolean> {
+    const start = Date.now();
+    await this.mcp.waitForConnections();
+    const elapsed = Date.now() - start;
+    return elapsed < 100; // Should be near-instant
+  }
+
+  /**
+   * Simulate a full hibernation round-trip:
+   *   onStart() (triggers restoreConnectionsFromStorage internally)
+   *   → waitForConnections()
+   *   → check connection states
+   *
+   * This tests the real lifecycle path, not the decomposed methods.
+   */
+  async hibernationRoundTrip(timeout?: number): Promise<{
+    connectionIds: string[];
+    connectionStates: Record<string, string>;
+  }> {
+    // onStart() calls restoreConnectionsFromStorage which fires
+    // background connections via _trackConnection
+    await this.onStart();
+
+    // This is what a consumer would call in onMessage / onChatMessage
+    await this.mcp.waitForConnections(
+      timeout != null ? { timeout } : undefined
+    );
+
+    const connectionStates: Record<string, string> = {};
+    for (const [id, conn] of Object.entries(this.mcp.mcpConnections)) {
+      connectionStates[id] = conn.connectionState;
+    }
+
+    return {
+      connectionIds: Object.keys(this.mcp.mcpConnections),
+      connectionStates
+    };
+  }
+
+  /**
+   * Simulate hibernation round-trip WITHOUT waiting — demonstrates the race.
+   */
+  async hibernationRoundTripNoWait(): Promise<{
+    connectionIds: string[];
+    connectionStates: Record<string, string>;
+  }> {
+    await this.onStart();
+
+    // Check states immediately — no waitForConnections
+    const connectionStates: Record<string, string> = {};
+    for (const [id, conn] of Object.entries(this.mcp.mcpConnections)) {
+      connectionStates[id] = conn.connectionState;
+    }
+
+    return {
+      connectionIds: Object.keys(this.mcp.mcpConnections),
+      connectionStates
+    };
+  }
+
+  hasMcpConnection(serverId: string): boolean {
+    return !!this.mcp.mcpConnections[serverId];
+  }
+
+  getConnectionState(serverId: string): string | null {
+    return this.mcp.mcpConnections[serverId]?.connectionState ?? null;
+  }
+}