@cloudflare/ai-chat 0.0.4 → 0.0.6

package/src/react-tests/setup.ts

@@ -0,0 +1,3 @@
+// properly set up the act environment for react-tests
+// @ts-expect-error - react specific API
+globalThis.IS_REACT_ACT_ENVIRONMENT = true;

package/src/react-tests/use-agent-chat.test.tsx

@@ -10,6 +10,10 @@ import {
 } from "../react";
 import type { useAgent } from "agents/react";
 
+function sleep(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 function createAgent({ name, url }: { name: string; url: string }) {
   const target = new EventTarget();
   const baseAgent = {
@@ -63,15 +67,17 @@ describe("useAgentChat", () => {
       return "Suspended";
     };
 
-    const screen = await act(() =>
-      render(<TestComponent />, {
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
         wrapper: ({ children }) => (
           <StrictMode>
             <Suspense fallback={<SuspenseObserver />}>{children}</Suspense>
           </StrictMode>
         )
-      })
-    );
+      });
+      await sleep(10);
+      return screen;
+    });
 
     await expect
       .element(screen.getByTestId("messages"))
@@ -123,15 +129,18 @@ describe("useAgentChat", () => {
       return "Suspended";
     };
 
-    const screen = await act(() =>
-      render(<TestComponent agent={agentA} />, {
+    const screen = await act(async () => {
+      const screen = render(<TestComponent agent={agentA} />, {
         wrapper: ({ children }) => (
           <StrictMode>
             <Suspense fallback={<SuspenseObserver />}>{children}</Suspense>
           </StrictMode>
         )
-      })
-    );
+      });
+
+      await sleep(10);
+      return screen;
+    });
 
     await expect
       .element(screen.getByTestId("messages"))
@@ -145,7 +154,10 @@ describe("useAgentChat", () => {
 
     suspenseRendered.mockClear();
 
-    await act(() => screen.rerender(<TestComponent agent={agentB} />));
+    await act(async () => {
+      screen.rerender(<TestComponent agent={agentB} />);
+      await sleep(10);
+    });
 
     await expect
       .element(screen.getByTestId("messages"))
@@ -219,7 +231,7 @@ describe("useAgentChat", () => {
         _options: PrepareSendMessagesRequestOptions<UIMessage>
       ): Promise<PrepareSendMessagesRequestResult> => {
         // Simulate async operation like fetching tool definitions
-        await new Promise((resolve) => setTimeout(resolve, 10));
+        await sleep(10);
         return {
           body: {
             clientTools: [
@@ -475,26 +487,24 @@ describe("useAgentChat client-side tool execution (issue #728)", () => {
       );
     };
 
-    const screen = await act(() =>
-      render(<TestComponent />, {
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
         wrapper: ({ children }) => (
           <StrictMode>
             <Suspense fallback="Loading...">{children}</Suspense>
           </StrictMode>
         )
-      })
-    );
+      });
+      // The tool should have been automatically executed
+      await sleep(10);
+      return screen;
+    });
 
     // Wait for initial messages to load
     await expect
       .element(screen.getByTestId("messages-count"))
       .toHaveTextContent("2");
 
-    // The tool should have been automatically executed
-    await act(async () => {
-      await new Promise((resolve) => setTimeout(resolve, 100));
-    });
-
     // Verify the tool execute was called
     expect(mockExecute).toHaveBeenCalled();
 
@@ -574,15 +584,17 @@ describe("useAgentChat client-side tool execution (issue #728)", () => {
       );
     };
 
-    const screen = await act(() =>
-      render(<TestComponent />, {
+    const screen = await act(async () => {
+      const screen = render(<TestComponent />, {
         wrapper: ({ children }) => (
           <StrictMode>
             <Suspense fallback="Loading...">{children}</Suspense>
           </StrictMode>
         )
-      })
-    );
+      });
+      await sleep(10);
+      return screen;
+    });
 
     await expect
       .element(screen.getByTestId("messages-count"))

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

@@ -12,6 +12,7 @@ export default defineConfig({
       ],
       provider: "playwright"
     },
-    clearMocks: true
+    clearMocks: true,
+    setupFiles: ["./setup.ts"]
   }
 });

package/src/react.tsx

@@ -835,6 +835,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 +1293,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 +1435,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/chat-context.test.ts

@@ -3,6 +3,7 @@ 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("AIChatAgent Connection Context - Issue #711", () => {
   it("getCurrentAgent() should return connection in onChatMessage and nested async functions (tool execute)", async () => {
@@ -10,7 +11,7 @@ describe("AIChatAgent Connection Context - Issue #711", () => {
     const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
 
     // Get the agent stub to access captured context
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     // Clear any previous captured context
     await agentStub.clearCapturedContext();

package/src/tests/chat-persistence.test.ts

@@ -4,6 +4,7 @@ import worker from "./worker";
 import { MessageType } from "../types";
 import type { UIMessage as ChatMessage } from "ai";
 import { connectChatWS } from "./test-utils";
+import { getAgentByName } from "agents";
 
 // Type helper for tool call parts - extracts ToolUIPart from ChatMessage parts
 type TestToolCallPart = Extract<
@@ -198,7 +199,7 @@ describe("Chat Agent Persistence", () => {
 
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     await agentStub.testPersistToolCall("msg-tool-1", "getLocalTime");
 
@@ -256,7 +257,7 @@ describe("Chat Agent Persistence", () => {
 
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     const userMessage: ChatMessage = {
       id: "user-1",
@@ -362,7 +363,7 @@ describe("Chat Agent Persistence", () => {
 
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     const userMessage: ChatMessage = {
       id: "user-1",

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

@@ -1,4 +1,5 @@
 import { createExecutionContext, env } from "cloudflare:test";
+import { getAgentByName } from "agents";
 import { describe, it, expect } from "vitest";
 import worker from "./worker";
 import type { UIMessage as ChatMessage } from "ai";
@@ -17,7 +18,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
     const toolCallId = "call_merge_test";
 
     // Persist assistant message with tool in input-available state
@@ -91,7 +92,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
     const toolCallId = "call_tool_result_test";
 
     // Persist assistant message with tool in input-available state
@@ -164,7 +165,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
     const toolCallId = "call_tool_result_auto_continue";
 
     // Persist assistant message with tool in input-available state
@@ -239,7 +240,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     // Persist message with OpenAI itemId in providerMetadata (simulates OpenAI Responses API)
     await agentStub.persistMessages([
@@ -299,7 +300,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
     const toolCallId = "call_openai_strip_test";
 
     // Persist message with tool that has OpenAI itemId in callProviderMetadata
@@ -362,7 +363,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     // Persist message with other metadata alongside itemId
     await agentStub.persistMessages([
@@ -433,7 +434,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     // Persist message with empty reasoning part (simulates OpenAI Responses API)
     await agentStub.persistMessages([
@@ -487,7 +488,7 @@ describe("Client-side tool duplicate message prevention", () => {
     ws.accept();
     await ctx.waitUntil(Promise.resolve());
 
-    const agentStub = env.TestChatAgent.get(env.TestChatAgent.idFromName(room));
+    const agentStub = await getAgentByName(env.TestChatAgent, room);
 
     // Persist message with non-empty reasoning part
     await agentStub.persistMessages([
@@ -541,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();
+  });
+});