@cloudflare/ai-chat 0.0.4 → 0.0.5

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/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([

package/src/tests/resumable-streaming.test.ts

@@ -2,6 +2,7 @@ import { env } from "cloudflare:test";
 import { describe, it, expect } from "vitest";
 import { MessageType, type OutgoingMessage } from "../types";
 import { connectChatWS, isUseChatResponseMessage } from "./test-utils";
+import { getAgentByName } from "agents";
 
 function isStreamResumingMessage(
   m: unknown
@@ -37,9 +38,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       const streamId = await agentStub.testStartStream("req-123");
       expect(streamId).toBeDefined();
@@ -59,9 +58,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       const streamId = await agentStub.testStartStream("req-456");
 
@@ -98,9 +95,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       const streamId = await agentStub.testStartStream("req-789");
 
@@ -127,9 +122,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       const streamId = await agentStub.testStartStream("req-error");
 
@@ -156,9 +149,7 @@ describe("Resumable Streaming", () => {
       );
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
       const streamId = await agentStub.testStartStream("req-resume");
       await agentStub.testStoreStreamChunk(
         streamId,
@@ -193,9 +184,7 @@ describe("Resumable Streaming", () => {
       );
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
       const streamId = await agentStub.testStartStream("req-ack");
       await agentStub.testStoreStreamChunk(
         streamId,
@@ -237,6 +226,82 @@ describe("Resumable Streaming", () => {
       ws2.close();
     });
 
+    it("does not deliver live chunks before ACK to resuming connections", async () => {
+      const room = crypto.randomUUID();
+
+      // First connection - start a stream
+      const { ws: ws1 } = await connectChatWS(
+        `/agents/test-chat-agent/${room}`
+      );
+      const messages1 = collectMessages(ws1);
+      await new Promise((r) => setTimeout(r, 50));
+
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
+      const streamId = await agentStub.testStartStream("req-live");
+
+      // Second connection - will be notified to resume
+      const { ws: ws2 } = await connectChatWS(
+        `/agents/test-chat-agent/${room}`
+      );
+      const messages2 = collectMessages(ws2);
+
+      await new Promise((r) => setTimeout(r, 100));
+
+      // Broadcast a live chunk while ws2 is pending resume (no ACK yet)
+      await agentStub.testBroadcastLiveChunk(
+        "req-live",
+        streamId,
+        '{"type":"text-delta","id":"0","delta":"A"}'
+      );
+
+      await new Promise((r) => setTimeout(r, 100));
+
+      // ws2 should NOT receive live chunks before ACK
+      const preAckChunks = messages2.filter(isUseChatResponseMessage);
+      expect(preAckChunks.length).toBe(0);
+
+      // ws1 should receive the live chunk
+      const ws1Chunks = messages1.filter(isUseChatResponseMessage);
+      expect(ws1Chunks.length).toBe(1);
+      expect(ws1Chunks[0].body).toBe(
+        '{"type":"text-delta","id":"0","delta":"A"}'
+      );
+
+      // Send ACK to resume
+      ws2.send(
+        JSON.stringify({
+          type: MessageType.CF_AGENT_STREAM_RESUME_ACK,
+          id: "req-live"
+        })
+      );
+
+      await new Promise((r) => setTimeout(r, 100));
+
+      // After ACK, ws2 should receive the replayed chunk
+      const postAckChunks = messages2.filter(isUseChatResponseMessage);
+      expect(postAckChunks.length).toBeGreaterThanOrEqual(1);
+      expect(postAckChunks[0].body).toBe(
+        '{"type":"text-delta","id":"0","delta":"A"}'
+      );
+
+      // Live chunks after ACK should be delivered
+      await agentStub.testBroadcastLiveChunk(
+        "req-live",
+        streamId,
+        '{"type":"text-delta","id":"0","delta":"B"}'
+      );
+
+      await new Promise((r) => setTimeout(r, 100));
+
+      const finalChunks = messages2.filter(isUseChatResponseMessage);
+      expect(finalChunks.some((m) => m.body?.includes('"delta":"B"'))).toBe(
+        true
+      );
+
+      ws1.close();
+      ws2.close();
+    });
+
     it("ignores ACK with wrong request ID", async () => {
       const room = crypto.randomUUID();
 
@@ -246,9 +311,7 @@ describe("Resumable Streaming", () => {
       );
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
       const streamId = await agentStub.testStartStream("req-correct");
       await agentStub.testStoreStreamChunk(
         streamId,
@@ -292,9 +355,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       // Insert a stale stream (6 minutes old)
       const staleStreamId = "stale-stream-123";
@@ -327,9 +388,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       // Insert a fresh stream (1 minute old)
       const freshStreamId = "fresh-stream-456";
@@ -363,9 +422,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       // Create a stream with chunks
       const streamId = await agentStub.testStartStream("req-clear");
@@ -405,9 +462,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       // Start first stream and add chunks without explicit flush
       const stream1 = await agentStub.testStartStream("req-1");
@@ -439,9 +494,7 @@ describe("Resumable Streaming", () => {
 
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
 
       const streamId = await agentStub.testStartStream("req-flush");
       await agentStub.testStoreStreamChunk(
@@ -470,9 +523,7 @@ describe("Resumable Streaming", () => {
       );
       await new Promise((r) => setTimeout(r, 50));
 
-      const agentStub = env.TestChatAgent.get(
-        env.TestChatAgent.idFromName(room)
-      );
+      const agentStub = await getAgentByName(env.TestChatAgent, room);
       const streamId = await agentStub.testStartStream("req-done");
       await agentStub.testStoreStreamChunk(
         streamId,

package/src/tests/worker.ts

@@ -1,6 +1,7 @@
 import { AIChatAgent } from "../";
 import type { UIMessage as ChatMessage } from "ai";
 import { callable, getCurrentAgent, routeAgentRequest } from "agents";
+import { MessageType, type OutgoingMessage } from "../types";
 
 // Type helper for tool call parts - extracts from ChatMessage parts
 type TestToolCallPart = Extract<
@@ -148,6 +149,29 @@ export class TestChatAgent extends AIChatAgent<Env> {
     this._storeStreamChunk(streamId, body);
   }
 
+  @callable()
+  testBroadcastLiveChunk(
+    requestId: string,
+    streamId: string,
+    body: string
+  ): void {
+    this._storeStreamChunk(streamId, body);
+    const message: OutgoingMessage = {
+      body,
+      done: false,
+      id: requestId,
+      type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
+    };
+    (
+      this as unknown as {
+        _broadcastChatMessage: (
+          msg: OutgoingMessage,
+          exclude?: string[]
+        ) => void;
+      }
+    )._broadcastChatMessage(message);
+  }
+
   @callable()
   testFlushChunkBuffer(): void {
     this._flushChunkBuffer();