eagerly load this.mcp.jsonSchema in ai chat agent

.changeset/beige-melons-brake.md

@@ -0,0 +1,12 @@
+---
+"@cloudflare/ai-chat": patch
+"agents": patch
+---
+
+fix: eagerly load this.mcp.jsonSchema in ai chat agent
+
+fixes https://github.com/cloudflare/agents/issues/718
+
+Because of hibernation, our preloaded this.mcp.jsonSchema can be lost when the DO wakes up again. We should probably move this helper into anotherl ibrary. Until then, a patch is to eagerly load it when an ai chat agent starts up.
+
+We also throw the error only if there's an active mcp connection.

packages/agents/src/mcp/client.ts

@@ -912,10 +912,12 @@ export class MCPClientManager {
    * @returns a set of tools that you can use with the AI SDK
    */
   getAITools(): ToolSet {
-    if (!this.jsonSchema) {
-      throw new Error("jsonSchema not initialized.");
+    if (Object.keys(this.mcpConnections).length > 0) {
+      // delay checking for .jsonSchema until we actually need it
+      if (!this.jsonSchema) {
+        throw new Error("invariant: mcp.jsonSchema not initialized.");
+      }
     }
-
     // Warn if tools are being read from non-ready connections
     for (const [id, conn] of Object.entries(this.mcpConnections)) {
       if (

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

@@ -1587,9 +1587,6 @@ describe("MCPClientManager OAuth Integration", () => {
       const name = "Test MCP Server";
       const callbackUrl = "http://localhost:3000/callback";
 
-      // Initialize jsonSchema (required for getAITools)
-      await manager.ensureJsonSchema();
-
       // Register server
       manager.registerServer(id, {
         url,
@@ -1637,6 +1634,13 @@ describe("MCPClientManager OAuth Integration", () => {
       expect(conn.connectionState).toBe("ready");
       expect(conn.tools).toHaveLength(1);
 
+      expect(() => manager.getAITools()).toThrow(
+        "invariant: mcp.jsonSchema not initialized."
+      );
+
+      // Initialize jsonSchema (required for getAITools)
+      await manager.ensureJsonSchema();
+
       // Get AI tools
       const tools = manager.getAITools();
 
@@ -1764,26 +1768,6 @@ describe("MCPClientManager OAuth Integration", () => {
         undefined
       );
     });
-
-    it("should throw error if jsonSchema not initialized", () => {
-      // Create a new manager without initializing jsonSchema
-      const mockStorage = {
-        sql: {
-          exec: <T extends Record<string, SqlStorageValue>>() =>
-            ([] as T[])[Symbol.iterator]()
-        },
-        get: async () => undefined,
-        put: async () => {}
-      } as unknown as DurableObjectStorage;
-
-      const newManager = new MCPClientManager("test-client", "1.0.0", {
-        storage: mockStorage
-      });
-
-      expect(() => newManager.getAITools()).toThrow(
-        "jsonSchema not initialized."
-      );
-    });
   });
 
   describe("clearAuthUrl()", () => {

packages/ai-chat/src/index.ts

@@ -515,6 +515,12 @@ export class AIChatAgent<
       // Forward unhandled messages to consumer's onMessage
       return _onMessage(connection, message);
     };
+
+    // since we're in a module that we know using ai-sdk, let's immediately ensure the .jsonSchema function is loaded
+    // todo: for some reason, wrapping this in a blockConcurrencyWhile is causing issues with the test runner
+    this.mcp.ensureJsonSchema().catch((error) => {
+      console.error("Error ensuring jsonSchema:", error);
+    });
   }
 
   /**