@cloudflare/ai-chat 0.0.5 → 0.0.7

package/src/index.ts

@@ -535,6 +535,13 @@ export class AIChatAgent<
           );
           return;
         }
+
+        // Handle client-side tool approval response
+        if (data.type === MessageType.CF_AGENT_TOOL_APPROVAL) {
+          const { toolCallId, approved } = data;
+          this._applyToolApproval(toolCallId, approved);
+          return;
+        }
       }
 
       // Forward unhandled messages to consumer's onMessage
@@ -869,9 +876,9 @@ export class AIChatAgent<
    * @returns Response to send to the client or undefined
    */
   async onChatMessage(
-    // biome-ignore lint/correctness/noUnusedFunctionParameters: overridden later
+    // oxlint-disable-next-line eslint(no-unused-vars) -- params used by subclass overrides
     onFinish: StreamTextOnFinishCallback<ToolSet>,
-    // biome-ignore lint/correctness/noUnusedFunctionParameters: overridden later
+    // oxlint-disable-next-line eslint(no-unused-vars) -- params used by subclass overrides
     options?: OnChatMessageOptions
   ): Promise<Response | undefined> {
     throw new Error(
@@ -1008,7 +1015,9 @@ export class AIChatAgent<
       if (
         "toolCallId" in part &&
         "state" in part &&
-        part.state === "output-available"
+        (part.state === "output-available" ||
+          part.state === "approval-responded" ||
+          part.state === "approval-requested")
       ) {
         const toolCallId = part.toolCallId as string;
 
@@ -1970,6 +1979,144 @@ export class AIChatAgent<
     }
   }
 
+  /**
+   * Applies a tool approval response from the client, updating the persisted message.
+   * This is called when the client sends CF_AGENT_TOOL_APPROVAL for tools with needsApproval.
+   * Updates the tool part state from input-available/approval-requested to approval-responded.
+   *
+   * @param toolCallId - The tool call ID this approval is for
+   * @param approved - Whether the tool execution was approved
+   * @returns true if the approval was applied, false if the message was not found
+   */
+  private async _applyToolApproval(
+    toolCallId: string,
+    approved: boolean
+  ): Promise<boolean> {
+    // Find the message with this tool call.
+    // We check two locations:
+    // 1. _streamingMessage: in-memory message being actively built during AI response
+    //    (not yet persisted to SQLite or available in this.messages)
+    // 2. this.messages: persisted messages loaded from SQLite database
+    //
+    // The user can approve before streaming finishes (e.g., approval UI appears
+    // while AI is still generating text), so we must check _streamingMessage first.
+
+    let message: ChatMessage | undefined;
+
+    // Check streaming message first (in-memory, not yet persisted)
+    if (this._streamingMessage) {
+      for (const part of this._streamingMessage.parts) {
+        if ("toolCallId" in part && part.toolCallId === toolCallId) {
+          message = this._streamingMessage;
+          break;
+        }
+      }
+    }
+
+    // If not found in streaming message, check persisted messages (in SQLite).
+    // Retry with backoff in case streaming completes and persists between attempts.
+    if (!message) {
+      for (let attempt = 0; attempt < 10; attempt++) {
+        message = this._findMessageByToolCallId(toolCallId);
+        if (message) break;
+        await new Promise((resolve) => setTimeout(resolve, 100));
+      }
+    }
+
+    if (!message) {
+      console.warn(
+        `[AIChatAgent] _applyToolApproval: Could not find message with toolCallId ${toolCallId} after retries`
+      );
+      return false;
+    }
+
+    // Check if this is the streaming message (not yet persisted)
+    const isStreamingMessage = message === this._streamingMessage;
+
+    // Update the tool part with the approval
+    let updated = false;
+    if (isStreamingMessage) {
+      // Update in place - the message will be persisted when streaming completes
+      for (const part of message.parts) {
+        if (
+          "toolCallId" in part &&
+          part.toolCallId === toolCallId &&
+          "state" in part &&
+          (part.state === "input-available" ||
+            part.state === "approval-requested")
+        ) {
+          (part as { state: string; approval?: { approved: boolean } }).state =
+            "approval-responded";
+          (
+            part as { state: string; approval?: { approved: boolean } }
+          ).approval = { approved };
+          updated = true;
+          break;
+        }
+      }
+    } else {
+      // For persisted messages, create updated parts
+      const updatedParts = message.parts.map((part) => {
+        if (
+          "toolCallId" in part &&
+          part.toolCallId === toolCallId &&
+          "state" in part &&
+          (part.state === "input-available" ||
+            part.state === "approval-requested")
+        ) {
+          updated = true;
+          return {
+            ...part,
+            state: "approval-responded" as const,
+            approval: { approved }
+          };
+        }
+        return part;
+      }) as ChatMessage["parts"];
+
+      if (updated) {
+        // Create the updated message and strip OpenAI item IDs
+        const updatedMessage: ChatMessage = this._sanitizeMessageForPersistence(
+          {
+            ...message,
+            parts: updatedParts
+          }
+        );
+
+        // Persist the updated message
+        this.sql`
+          update cf_ai_chat_agent_messages 
+          set message = ${JSON.stringify(updatedMessage)}
+          where id = ${message.id}
+        `;
+
+        // Reload messages to update in-memory state
+        const persisted = this._loadMessagesFromDb();
+        this.messages = autoTransformMessages(persisted);
+      }
+    }
+
+    if (!updated) {
+      console.warn(
+        `[AIChatAgent] _applyToolApproval: Tool part with toolCallId ${toolCallId} not in input-available or approval-requested state`
+      );
+      return false;
+    }
+
+    // Broadcast the update to all clients (only for persisted messages)
+    if (!isStreamingMessage) {
+      const broadcastMessage = this._findMessageByToolCallId(toolCallId);
+      if (broadcastMessage) {
+        this._broadcastChatMessage({
+          type: MessageType.CF_AGENT_MESSAGE_UPDATED,
+          message: broadcastMessage
+        });
+      }
+    }
+
+    return true;
+  }
+
   private async _reply(
     id: string,
     response: Response,

package/src/react-tests/vitest.config.ts

@@ -2,6 +2,7 @@ import { defineConfig } from "vitest/config";
 
 export default defineConfig({
   test: {
+    name: "react",
     browser: {
       enabled: true,
       instances: [

package/src/react.tsx

@@ -354,10 +354,12 @@ export function useAgentChat<
   onToolCallRef.current = onToolCall;
 
   const agentUrl = new URL(
-    `${// @ts-expect-error we're using a protected _url property that includes query params
-    ((agent._url as string | null) || agent._pkurl)
-      ?.replace("ws://", "http://")
-      .replace("wss://", "https://")}`
+    `${
+      // @ts-expect-error we're using a protected _url property that includes query params
+      ((agent._url as string | null) || agent._pkurl)
+        ?.replace("ws://", "http://")
+        .replace("wss://", "https://")
+    }`
   );
 
   agentUrl.searchParams.delete("_pk");
@@ -835,6 +837,20 @@ export function useAgentChat<
     [autoContinueAfterToolResult]
   );
 
+  // Helper function to send tool approval to server
+  const sendToolApprovalToServer = useCallback(
+    (toolCallId: string, approved: boolean) => {
+      agentRef.current.send(
+        JSON.stringify({
+          type: MessageType.CF_AGENT_TOOL_APPROVAL,
+          toolCallId,
+          approved
+        })
+      );
+    },
+    []
+  );
+
   // Effect for new onToolCall callback pattern (v6 style)
   // This fires when there are tool calls that need client-side handling
   useEffect(() => {
@@ -1279,6 +1295,45 @@ export function useAgentChat<
       // If autoContinueAfterToolResult is true, server handles continuation
     };
 
+  // Wrapper that sends tool approval to server before updating local state.
+  // This prevents duplicate messages by ensuring server updates the message
+  // in place with the existing ID, rather than relying on ID resolution
+  // when sendMessage() is called later.
+  const addToolApprovalResponseAndNotifyServer: typeof useChatHelpers.addToolApprovalResponse =
+    (args) => {
+      const { id: approvalId, approved } = args;
+
+      // Find the toolCallId from the approval ID
+      // The approval ID is stored on the tool part's approval.id field
+      let toolCallId: string | undefined;
+      for (const msg of messagesRef.current) {
+        for (const part of msg.parts) {
+          if (
+            "toolCallId" in part &&
+            "approval" in part &&
+            (part.approval as { id?: string })?.id === approvalId
+          ) {
+            toolCallId = part.toolCallId as string;
+            break;
+          }
+        }
+        if (toolCallId) break;
+      }
+
+      if (toolCallId) {
+        // Send approval to server first (server updates message in place)
+        sendToolApprovalToServer(toolCallId, approved);
+      } else {
+        console.warn(
+          `[useAgentChat] addToolApprovalResponse: Could not find toolCallId for approval ID "${approvalId}". ` +
+            "Server will not be notified, which may cause duplicate messages."
+        );
+      }
+
+      // Call AI SDK's addToolApprovalResponse for local state update
+      useChatHelpers.addToolApprovalResponse(args);
+    };
+
   // Fix for issue #728: Merge client-side tool results with messages
   // so tool parts show output-available immediately after execution
   const messagesWithToolResults = useMemo(() => {
@@ -1382,6 +1437,12 @@ export function useAgentChat<
      * @deprecated Use `addToolOutput` instead.
      */
     addToolResult: addToolResultAndSendMessage,
+    /**
+     * Respond to a tool approval request. Use this for tools with `needsApproval`.
+     * This wrapper notifies the server before updating local state, preventing
+     * duplicate messages when sendMessage() is called afterward.
+     */
+    addToolApprovalResponse: addToolApprovalResponseAndNotifyServer,
     clearHistory: () => {
       useChatHelpers.setMessages([]);
       setClientToolResults(new Map());

package/src/tests/client-tool-duplicate-message.test.ts

@@ -542,3 +542,316 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.close();
   });
 });
+
+describe("Tool approval (needsApproval) duplicate message prevention", () => {
+  it("CF_AGENT_TOOL_APPROVAL updates existing message in place", async () => {
+    const room = crypto.randomUUID();
+    const ctx = createExecutionContext();
+    const req = new Request(
+      `http://example.com/agents/test-chat-agent/${room}`,
+      { headers: { Upgrade: "websocket" } }
+    );
+    const res = await worker.fetch(req, env, ctx);
+    expect(res.status).toBe(101);
+    const ws = res.webSocket as WebSocket;
+    ws.accept();
+    await ctx.waitUntil(Promise.resolve());
+
+    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const toolCallId = "call_approval_test";
+
+    // Persist assistant message with tool in input-available state (needs approval)
+    await agentStub.persistMessages([
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Execute tool" }]
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-testTool",
+            toolCallId,
+            state: "input-available",
+            input: { param: "value" }
+          }
+        ] as ChatMessage["parts"]
+      }
+    ]);
+
+    // Send CF_AGENT_TOOL_APPROVAL via WebSocket
+    ws.send(
+      JSON.stringify({
+        type: "cf_agent_tool_approval",
+        toolCallId,
+        approved: true
+      })
+    );
+
+    await new Promise((resolve) => setTimeout(resolve, 200));
+
+    const messages = (await agentStub.getPersistedMessages()) as ChatMessage[];
+    const assistantMessages = messages.filter((m) => m.role === "assistant");
+
+    // Should have exactly 1 assistant message (updated in place, not duplicated)
+    expect(assistantMessages.length).toBe(1);
+
+    const assistantMsg = assistantMessages[0];
+    // Message ID should be preserved
+    expect(assistantMsg.id).toBe("assistant-1");
+
+    // Tool state should be updated to approval-responded
+    const toolPart = assistantMsg.parts[0] as {
+      state: string;
+      approval?: { approved: boolean };
+    };
+    expect(toolPart.state).toBe("approval-responded");
+    expect(toolPart.approval).toEqual({ approved: true });
+
+    ws.close();
+  });
+
+  it("CF_AGENT_TOOL_APPROVAL handles rejection (approved: false)", async () => {
+    const room = crypto.randomUUID();
+    const ctx = createExecutionContext();
+    const req = new Request(
+      `http://example.com/agents/test-chat-agent/${room}`,
+      { headers: { Upgrade: "websocket" } }
+    );
+    const res = await worker.fetch(req, env, ctx);
+    expect(res.status).toBe(101);
+    const ws = res.webSocket as WebSocket;
+    ws.accept();
+    await ctx.waitUntil(Promise.resolve());
+
+    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const toolCallId = "call_rejection_test";
+
+    // Persist assistant message with tool in input-available state
+    await agentStub.persistMessages([
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Execute tool" }]
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-testTool",
+            toolCallId,
+            state: "input-available",
+            input: { param: "value" }
+          }
+        ] as ChatMessage["parts"]
+      }
+    ]);
+
+    // Send CF_AGENT_TOOL_APPROVAL with approved: false (rejection)
+    ws.send(
+      JSON.stringify({
+        type: "cf_agent_tool_approval",
+        toolCallId,
+        approved: false
+      })
+    );
+
+    await new Promise((resolve) => setTimeout(resolve, 200));
+
+    const messages = (await agentStub.getPersistedMessages()) as ChatMessage[];
+    const assistantMessages = messages.filter((m) => m.role === "assistant");
+
+    expect(assistantMessages.length).toBe(1);
+
+    const toolPart = assistantMessages[0].parts[0] as {
+      state: string;
+      approval?: { approved: boolean };
+    };
+    expect(toolPart.state).toBe("approval-responded");
+    expect(toolPart.approval).toEqual({ approved: false });
+
+    ws.close();
+  });
+
+  it("CF_AGENT_TOOL_APPROVAL updates tool in approval-requested state", async () => {
+    const room = crypto.randomUUID();
+    const ctx = createExecutionContext();
+    const req = new Request(
+      `http://example.com/agents/test-chat-agent/${room}`,
+      { headers: { Upgrade: "websocket" } }
+    );
+    const res = await worker.fetch(req, env, ctx);
+    expect(res.status).toBe(101);
+    const ws = res.webSocket as WebSocket;
+    ws.accept();
+    await ctx.waitUntil(Promise.resolve());
+
+    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const toolCallId = "call_approval_requested_test";
+
+    // Persist assistant message with tool in approval-requested state
+    await agentStub.persistMessages([
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Execute tool" }]
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-testTool",
+            toolCallId,
+            state: "approval-requested",
+            input: { param: "value" },
+            approval: { id: "approval-123" }
+          }
+        ] as ChatMessage["parts"]
+      }
+    ]);
+
+    // Send CF_AGENT_TOOL_APPROVAL
+    ws.send(
+      JSON.stringify({
+        type: "cf_agent_tool_approval",
+        toolCallId,
+        approved: true
+      })
+    );
+
+    await new Promise((resolve) => setTimeout(resolve, 200));
+
+    const messages = (await agentStub.getPersistedMessages()) as ChatMessage[];
+    const assistantMessages = messages.filter((m) => m.role === "assistant");
+
+    expect(assistantMessages.length).toBe(1);
+
+    const toolPart = assistantMessages[0].parts[0] as {
+      state: string;
+      approval?: { approved: boolean };
+    };
+    expect(toolPart.state).toBe("approval-responded");
+    expect(toolPart.approval).toEqual({ approved: true });
+
+    ws.close();
+  });
+
+  it("CF_AGENT_TOOL_APPROVAL with non-existent toolCallId logs warning", async () => {
+    const room = crypto.randomUUID();
+    const ctx = createExecutionContext();
+    const req = new Request(
+      `http://example.com/agents/test-chat-agent/${room}`,
+      { headers: { Upgrade: "websocket" } }
+    );
+    const res = await worker.fetch(req, env, ctx);
+    expect(res.status).toBe(101);
+    const ws = res.webSocket as WebSocket;
+    ws.accept();
+    await ctx.waitUntil(Promise.resolve());
+
+    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+
+    // Persist a message without any tool calls
+    await agentStub.persistMessages([
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Hello" }]
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "Hi there!" }] as ChatMessage["parts"]
+      }
+    ]);
+
+    // Send CF_AGENT_TOOL_APPROVAL for non-existent tool
+    ws.send(
+      JSON.stringify({
+        type: "cf_agent_tool_approval",
+        toolCallId: "non_existent_tool_call",
+        approved: true
+      })
+    );
+
+    await new Promise((resolve) => setTimeout(resolve, 1200)); // Wait for retries (10 * 100ms + buffer)
+
+    const messages = (await agentStub.getPersistedMessages()) as ChatMessage[];
+
+    // Messages should remain unchanged (no crash, graceful handling)
+    expect(messages.length).toBe(2);
+    const assistantMsg = messages.find((m) => m.role === "assistant");
+    expect(assistantMsg?.parts[0]).toEqual({ type: "text", text: "Hi there!" });
+
+    ws.close();
+  });
+
+  it("CF_AGENT_TOOL_APPROVAL does not update tool already in output-available state", async () => {
+    const room = crypto.randomUUID();
+    const ctx = createExecutionContext();
+    const req = new Request(
+      `http://example.com/agents/test-chat-agent/${room}`,
+      { headers: { Upgrade: "websocket" } }
+    );
+    const res = await worker.fetch(req, env, ctx);
+    expect(res.status).toBe(101);
+    const ws = res.webSocket as WebSocket;
+    ws.accept();
+    await ctx.waitUntil(Promise.resolve());
+
+    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const toolCallId = "call_already_completed";
+
+    // Persist assistant message with tool already in output-available state
+    await agentStub.persistMessages([
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Execute tool" }]
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-testTool",
+            toolCallId,
+            state: "output-available",
+            input: { param: "value" },
+            output: { result: "done" }
+          }
+        ] as ChatMessage["parts"]
+      }
+    ]);
+
+    // Send CF_AGENT_TOOL_APPROVAL for tool that's already completed
+    ws.send(
+      JSON.stringify({
+        type: "cf_agent_tool_approval",
+        toolCallId,
+        approved: true
+      })
+    );
+
+    await new Promise((resolve) => setTimeout(resolve, 200));
+
+    const messages = (await agentStub.getPersistedMessages()) as ChatMessage[];
+    const assistantMessages = messages.filter((m) => m.role === "assistant");
+
+    expect(assistantMessages.length).toBe(1);
+
+    // State should remain output-available (not changed to approval-responded)
+    const toolPart = assistantMessages[0].parts[0] as {
+      state: string;
+      output?: unknown;
+    };
+    expect(toolPart.state).toBe("output-available");
+    expect(toolPart.output).toEqual({ result: "done" });
+
+    ws.close();
+  });
+});

package/src/tests/vitest.config.ts

@@ -2,6 +2,7 @@ import { defineWorkersConfig } from "@cloudflare/vitest-pool-workers/config";
 
 export default defineWorkersConfig({
   test: {
+    name: "workers",
     deps: {
       optimizer: {
         ssr: {

package/src/tests/wrangler.jsonc

@@ -1,5 +1,5 @@
 {
-  "compatibility_date": "2025-04-17",
+  "compatibility_date": "2026-01-28",
   "compatibility_flags": [
     "nodejs_compat",
     // adding these flags since the vitest runner needs them

package/src/types.ts

@@ -18,7 +18,9 @@ export enum MessageType {
   /** Client sends tool result to server (for client-side tools) */
   CF_AGENT_TOOL_RESULT = "cf_agent_tool_result",
   /** Server notifies client that a message was updated (e.g., tool result applied) */
-  CF_AGENT_MESSAGE_UPDATED = "cf_agent_message_updated"
+  CF_AGENT_MESSAGE_UPDATED = "cf_agent_message_updated",
+  /** Client sends tool approval response to server (for tools with needsApproval) */
+  CF_AGENT_TOOL_APPROVAL = "cf_agent_tool_approval"
 }
 
 /**
@@ -119,4 +121,12 @@ export type IncomingMessage<ChatMessage extends UIMessage = UIMessage> =
       output: unknown;
       /** Whether server should auto-continue the conversation after applying result */
       autoContinue?: boolean;
+    }
+  | {
+      /** Client sends tool approval response to server (for tools with needsApproval) */
+      type: MessageType.CF_AGENT_TOOL_APPROVAL;
+      /** The tool call ID this approval is for */
+      toolCallId: string;
+      /** Whether the tool execution was approved */
+      approved: boolean;
     };

package/vitest.config.ts

@@ -0,0 +1,7 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+  test: {
+    projects: ["src/tests/vitest.config.ts", "src/react-tests/vitest.config.ts"]
+  }
+});