Fix client tool schemas lost after DO restart

.changeset/fix-client-tools-reconnect.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/ai-chat": patch
+---
+
+Fix client tool schemas lost after DO restart by re-sending them with CF_AGENT_TOOL_RESULT

packages/ai-chat/src/index.ts

@@ -482,7 +482,8 @@ export class AIChatAgent<
 
         // Handle client-side tool result
         if (data.type === MessageType.CF_AGENT_TOOL_RESULT) {
-          const { toolCallId, toolName, output, autoContinue } = data;
+          const { toolCallId, toolName, output, autoContinue, clientTools } =
+            data;
 
           // Apply the tool result
           this._applyToolResult(toolCallId, toolName, output).then(
@@ -533,7 +534,7 @@ export class AIChatAgent<
                           },
                           {
                             abortSignal,
-                            clientTools: this._lastClientTools
+                            clientTools: clientTools ?? this._lastClientTools
                           }
                         );
 

packages/ai-chat/src/react.tsx

@@ -777,14 +777,16 @@ export function useAgentChat<
 
             if (toolResults.length > 0) {
               // Send tool results to server first (server is source of truth)
+              const clientToolSchemas = extractClientToolSchemas(currentTools);
               for (const result of toolResults) {
                 agentRef.current.send(
                   JSON.stringify({
                     type: MessageType.CF_AGENT_TOOL_RESULT,
                     toolCallId: result.toolCallId,
                     toolName: result.toolName,
                     output: result.output,
-                    autoContinue: autoContinueAfterToolResult
+                    autoContinue: autoContinueAfterToolResult,
+                    clientTools: clientToolSchemas
                   })
                 );
               }
@@ -838,7 +840,8 @@ export function useAgentChat<
           toolCallId,
           toolName,
           output,
-          autoContinue: autoContinueAfterToolResult
+          autoContinue: autoContinueAfterToolResult,
+          clientTools: extractClientToolSchemas(toolsRef.current)
         })
       );
 
@@ -1267,7 +1270,8 @@ export function useAgentChat<
           toolCallId,
           toolName,
           output,
-          autoContinue: autoContinueAfterToolResult
+          autoContinue: autoContinueAfterToolResult,
+          clientTools: extractClientToolSchemas(toolsRef.current)
         })
       );
 

packages/ai-chat/src/tests/client-tools-reconnect.test.ts

@@ -0,0 +1,143 @@
+import { env } from "cloudflare:test";
+import { describe, it, expect } from "vitest";
+import { MessageType } from "../types";
+import type { UIMessage as ChatMessage } from "ai";
+import { connectChatWS } from "./test-utils";
+import { getAgentByName } from "agents";
+
+describe("Client tools after reconnect", () => {
+  it("should use client tools from CF_AGENT_TOOL_RESULT for continuation", async () => {
+    const room = crypto.randomUUID();
+
+    // Step 1: Set up a conversation with a pending tool call (simulates state before refresh)
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
+
+    const userMessage: ChatMessage = {
+      id: "msg1",
+      role: "user",
+      parts: [{ type: "text", text: "Change the background" }]
+    };
+
+    const toolCallId = "call_reconnect_test";
+    const assistantMessage: ChatMessage = {
+      id: "assistant-1",
+      role: "assistant",
+      parts: [
+        {
+          type: "tool-changeBackgroundColor",
+          toolCallId,
+          state: "input-available",
+          input: { color: "blue" }
+        }
+      ] as ChatMessage["parts"]
+    };
+
+    // Persist messages directly (simulates state loaded from SQLite after DO restart)
+    await agentStub.persistMessages([userMessage, assistantMessage]);
+
+    // Step 2: Connect (simulates reconnect after refresh)
+    // Note: We intentionally do NOT send a CF_AGENT_USE_CHAT_REQUEST first,
+    // so _lastClientTools is never set — this simulates DO restart.
+    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
+
+    // Wait for connection to be established
+    await new Promise((resolve) => setTimeout(resolve, 200));
+
+    // Clear any captured context from connection setup
+    await agentStub.clearCapturedContext();
+
+    // Step 3: Send tool result WITH clientTools (simulates client approval after reconnect)
+    const clientTools = [
+      {
+        name: "changeBackgroundColor",
+        description: "Changes the background color",
+        parameters: {
+          type: "object",
+          properties: { color: { type: "string" } }
+        }
+      },
+      {
+        name: "changeTextColor",
+        description: "Changes the text color",
+        parameters: {
+          type: "object",
+          properties: { color: { type: "string" } }
+        }
+      }
+    ];
+
+    ws.send(
+      JSON.stringify({
+        type: MessageType.CF_AGENT_TOOL_RESULT,
+        toolCallId,
+        toolName: "changeBackgroundColor",
+        output: { success: true },
+        autoContinue: true,
+        clientTools
+      })
+    );
+
+    // Wait for tool result to be applied + 500ms stream wait + continuation
+    await new Promise((resolve) => setTimeout(resolve, 1500));
+
+    // Step 4: Verify continuation received client tools
+    const capturedClientTools = await agentStub.getCapturedClientTools();
+    expect(capturedClientTools).toBeDefined();
+    expect(capturedClientTools).toHaveLength(2);
+    expect(capturedClientTools![0].name).toBe("changeBackgroundColor");
+    expect(capturedClientTools![1].name).toBe("changeTextColor");
+
+    ws.close();
+  });
+
+  it("should work without clientTools in CF_AGENT_TOOL_RESULT (backwards compat)", async () => {
+    const room = crypto.randomUUID();
+
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
+
+    // Persist a conversation with a pending tool call
+    await agentStub.persistMessages([
+      {
+        id: "msg1",
+        role: "user",
+        parts: [{ type: "text", text: "Do something" }]
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-testTool",
+            toolCallId: "call_compat_test",
+            state: "input-available",
+            input: {}
+          }
+        ] as ChatMessage["parts"]
+      }
+    ]);
+
+    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
+    await new Promise((resolve) => setTimeout(resolve, 200));
+    await agentStub.clearCapturedContext();
+
+    // Send tool result WITHOUT clientTools (old client behavior)
+    ws.send(
+      JSON.stringify({
+        type: MessageType.CF_AGENT_TOOL_RESULT,
+        toolCallId: "call_compat_test",
+        toolName: "testTool",
+        output: { success: true },
+        autoContinue: true
+        // No clientTools field
+      })
+    );
+
+    await new Promise((resolve) => setTimeout(resolve, 1500));
+
+    // Continuation should still fire, just without client tools
+    const capturedClientTools = await agentStub.getCapturedClientTools();
+    expect(capturedClientTools).toBeUndefined();
+
+    ws.close();
+  });
+});

packages/ai-chat/src/tests/worker.ts

@@ -6,6 +6,7 @@ import type {
 } from "ai";
 import { callable, getCurrentAgent, routeAgentRequest } from "agents";
 import { MessageType, type OutgoingMessage } from "../types";
+import type { ClientToolSchema } from "../";
 
 // Type helper for tool call parts - extracts from ChatMessage parts
 type TestToolCallPart = Extract<
@@ -34,7 +35,7 @@ export class TestChatAgent extends AIChatAgent<Env> {
   // Store captured body from onChatMessage options for testing
   private _capturedBody: Record<string, unknown> | undefined = undefined;
   // Store captured clientTools from onChatMessage options for testing
-  private _capturedClientTools: unknown[] | undefined = undefined;
+  private _capturedClientTools: ClientToolSchema[] | undefined = undefined;
 
   async onChatMessage(
     _onFinish: StreamTextOnFinishCallback<ToolSet>,
@@ -108,7 +109,7 @@ export class TestChatAgent extends AIChatAgent<Env> {
   }
 
   @callable()
-  getCapturedClientTools(): unknown[] | undefined {
+  getCapturedClientTools(): ClientToolSchema[] | undefined {
     return this._capturedClientTools;
   }
 

packages/ai-chat/src/types.ts

@@ -121,6 +121,12 @@ export type IncomingMessage<ChatMessage extends UIMessage = UIMessage> =
       output: unknown;
       /** Whether server should auto-continue the conversation after applying result */
       autoContinue?: boolean;
+      /** Client tool schemas for continuation (client is source of truth) */
+      clientTools?: Array<{
+        name: string;
+        description?: string;
+        parameters?: Record<string, unknown>;
+      }>;
     }
   | {
       /** Client sends tool approval response to server (for tools with needsApproval) */