import { StrictMode, Suspense, act } from "react"; import { describe, expect, it, vi } from "vitest"; import { render } from "vitest-browser-react"; import type { UIMessage } from "ai"; import { useAgentChat, type PrepareSendMessagesRequestOptions, type PrepareSendMessagesRequestResult, type AITool } from "../react"; import type { useAgent } from "agents/react"; function sleep(ms: number) { return new Promise((resolve) => setTimeout(resolve, ms)); } function createAgent({ name, url, send }: { name: string; url: string; send?: (data: string) => void; }) { const target = new EventTarget(); const baseAgent = { _pkurl: url, _pk: name, // Use name as pk to distinguish agents _url: null as string | null, addEventListener: target.addEventListener.bind(target), agent: "Chat", close: () => {}, id: "fake-agent", name, removeEventListener: target.removeEventListener.bind(target), send: send ?? (() => {}), dispatchEvent: target.dispatchEvent.bind(target) }; return baseAgent as unknown as ReturnType; } describe("useAgentChat", () => { it("should cache initial message responses across re-renders", async () => { const agent = createAgent({ name: "thread-alpha", url: "ws://localhost:3000/agents/chat/thread-alpha?_pk=abc" }); const testMessages = [ { id: "1", role: "user" as const, parts: [{ type: "text" as const, text: "Hi" }] }, { id: "2", role: "assistant" as const, parts: [{ type: "text" as const, text: "Hello" }] } ]; const getInitialMessages = vi.fn(() => Promise.resolve(testMessages)); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages }); return

{JSON.stringify(chat.messages)}

; }; const suspenseRendered = vi.fn(); const SuspenseObserver = () => { suspenseRendered(); return "Suspended"; }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( }>{children} ) }); await sleep(10); return screen; }); await expect .element(screen.getByTestId("messages")) .toHaveTextContent(JSON.stringify(testMessages)); expect(getInitialMessages).toHaveBeenCalledTimes(1); expect(suspenseRendered).toHaveBeenCalled(); suspenseRendered.mockClear(); await screen.rerender(); await expect .element(screen.getByTestId("messages")) .toHaveTextContent(JSON.stringify(testMessages)); expect(getInitialMessages).toHaveBeenCalledTimes(1); expect(suspenseRendered).not.toHaveBeenCalled(); }); it("should refetch initial messages when the agent name changes", async () => { const url = "ws://localhost:3000/agents/chat/thread-a?_pk=abc"; const agentA = createAgent({ name: "thread-a", url }); const agentB = createAgent({ name: "thread-b", url }); const getInitialMessages = vi.fn(async ({ name }: { name: string }) => [ { id: "1", role: "assistant" as const, parts: [{ type: "text" as const, text: `Hello from ${name}` }] } ]); const TestComponent = ({ agent }: { agent: ReturnType; }) => { const chat = useAgentChat({ agent, getInitialMessages }); return

{JSON.stringify(chat.messages)}

; }; const suspenseRendered = vi.fn(); const SuspenseObserver = () => { suspenseRendered(); return "Suspended"; }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( }>{children} ) }); await sleep(10); return screen; }); await expect .element(screen.getByTestId("messages")) .toHaveTextContent("Hello from thread-a"); expect(getInitialMessages).toHaveBeenCalledTimes(1); expect(getInitialMessages).toHaveBeenNthCalledWith( 1, expect.objectContaining({ name: "thread-a" }) ); suspenseRendered.mockClear(); await act(async () => { screen.rerender(); await sleep(10); }); await expect .element(screen.getByTestId("messages")) .toHaveTextContent("Hello from thread-b"); expect(getInitialMessages).toHaveBeenCalledTimes(2); expect(getInitialMessages).toHaveBeenNthCalledWith( 2, expect.objectContaining({ name: "thread-b" }) ); }); it("should accept prepareSendMessagesRequest option without errors", async () => { const agent = createAgent({ name: "thread-with-tools", url: "ws://localhost:3000/agents/chat/thread-with-tools?_pk=abc" }); const prepareSendMessagesRequest = vi.fn( ( _options: PrepareSendMessagesRequestOptions ): PrepareSendMessagesRequestResult => ({ body: { clientTools: [ { name: "showAlert", description: "Shows an alert to the user", parameters: { message: { type: "string" } } } ] }, headers: { "X-Client-Tool-Count": "1" } }) ); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, // Skip fetching initial messages prepareSendMessagesRequest }); return

{chat.messages.length}

; }; const screen = await act(() => render(, { wrapper: ({ children }) => ( {children} ) }) ); // Verify component renders without errors await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("0"); }); it("should handle async prepareSendMessagesRequest", async () => { const agent = createAgent({ name: "thread-async-prepare", url: "ws://localhost:3000/agents/chat/thread-async-prepare?_pk=abc" }); const prepareSendMessagesRequest = vi.fn( async ( _options: PrepareSendMessagesRequestOptions ): Promise => { // Simulate async operation like fetching tool definitions await sleep(10); return { body: { clientTools: [ { name: "navigateToPage", description: "Navigates to a page" } ] } }; } ); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, prepareSendMessagesRequest }); return

{chat.messages.length}

; }; const screen = await act(() => render(, { wrapper: ({ children }) => ( {children} ) }) ); // Verify component renders without errors await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("0"); }); it("should auto-extract schemas from tools with execute functions", async () => { const agent = createAgent({ name: "thread-client-tools", url: "ws://localhost:3000/agents/chat/thread-client-tools?_pk=abc" }); // Tools with execute functions have their schemas auto-extracted and sent to server const tools: Record> = { showAlert: { description: "Shows an alert dialog to the user", parameters: { type: "object", properties: { message: { type: "string", description: "The message to display" } }, required: ["message"] }, execute: async (input) => { // Client-side execution const { message } = input as { message: string }; return { shown: true, message }; } }, changeBackgroundColor: { description: "Changes the page background color", parameters: { type: "object", properties: { color: { type: "string" } } }, execute: async (input) => { const { color } = input as { color: string }; return { success: true, color }; } } }; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, tools }); return

{chat.messages.length}

; }; const screen = await act(() => render(, { wrapper: ({ children }) => ( {children} ) }) ); // Verify component renders without errors await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("0"); }); it("should combine auto-extracted tools with prepareSendMessagesRequest", async () => { const agent = createAgent({ name: "thread-combined", url: "ws://localhost:3000/agents/chat/thread-combined?_pk=abc" }); const tools: Record = { showAlert: { description: "Shows an alert", execute: async () => ({ shown: true }) } }; const prepareSendMessagesRequest = vi.fn( ( _options: PrepareSendMessagesRequestOptions ): PrepareSendMessagesRequestResult => ({ body: { customData: "extra-context", userTimezone: "America/New_York" }, headers: { "X-Custom-Header": "custom-value" } }) ); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, tools, prepareSendMessagesRequest }); return

{chat.messages.length}

; }; const screen = await act(() => render(, { wrapper: ({ children }) => ( {children} ) }) ); // Verify component renders without errors await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("0"); }); it("should work with tools that have execute functions for client-side execution", async () => { const agent = createAgent({ name: "thread-tools-execution", url: "ws://localhost:3000/agents/chat/thread-tools-execution?_pk=abc" }); const mockExecute = vi.fn().mockResolvedValue({ success: true }); // Single unified tools object - schema + execute in one place const tools: Record = { showAlert: { description: "Shows an alert", parameters: { type: "object", properties: { message: { type: "string" } } }, execute: mockExecute } }; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, tools }); return

{chat.messages.length}

; }; const screen = await act(() => render(, { wrapper: ({ children }) => ( {children} ) }) ); // Verify component renders without errors await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("0"); }); }); describe("useAgentChat client-side tool execution (issue #728)", () => { it("should update tool part state from input-available to output-available when addToolResult is called", async () => { const agent = createAgent({ name: "tool-state-test", url: "ws://localhost:3000/agents/chat/tool-state-test?_pk=abc" }); const mockExecute = vi.fn().mockResolvedValue({ location: "New York" }); // Initial messages with a tool call in input-available state const initialMessages: UIMessage[] = [ { id: "msg-1", role: "user", parts: [{ type: "text", text: "Where am I?" }] }, { id: "msg-2", role: "assistant", parts: [ { type: "tool-getLocation", toolCallId: "tool-call-1", state: "input-available", input: {} } ] } ]; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages), experimental_automaticToolResolution: true, tools: { getLocation: { execute: mockExecute } } }); // Find the tool part to check its state const assistantMsg = chat.messages.find((m) => m.role === "assistant"); const toolPart = assistantMsg?.parts.find( (p) => "toolCallId" in p && p.toolCallId === "tool-call-1" ); const toolState = toolPart && "state" in toolPart ? toolPart.state : "not-found"; return (

{chat.messages.length}

{toolState}

); }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); // 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"); // Verify the tool execute was called expect(mockExecute).toHaveBeenCalled(); // the tool part should be updated to output-available // in the SAME message (msg-2), not in a new message await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("2"); // Should still be 2 messages, not 3 // The tool state should be output-available after addToolResult await expect .element(screen.getByTestId("tool-state")) .toHaveTextContent("output-available"); }); it("should not create duplicate tool parts when client executes tool", async () => { const agent = createAgent({ name: "duplicate-test", url: "ws://localhost:3000/agents/chat/duplicate-test?_pk=abc" }); const mockExecute = vi.fn().mockResolvedValue({ confirmed: true }); const initialMessages: UIMessage[] = [ { id: "msg-1", role: "assistant", parts: [ { type: "text", text: "Should I proceed?" }, { type: "tool-askForConfirmation", toolCallId: "confirm-1", state: "input-available", input: { message: "Proceed with action?" } } ] } ]; let chatInstance: ReturnType | null = null; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages), tools: { askForConfirmation: { execute: mockExecute } } }); chatInstance = chat; // Count tool parts with this toolCallId const toolPartsCount = chat.messages.reduce((count, msg) => { return ( count + msg.parts.filter( (p) => "toolCallId" in p && p.toolCallId === "confirm-1" ).length ); }, 0); // Get the tool state const toolPart = chat.messages .flatMap((m) => m.parts) .find((p) => "toolCallId" in p && p.toolCallId === "confirm-1"); const toolState = toolPart && "state" in toolPart ? toolPart.state : "not-found"; return (

{chat.messages.length}

{toolPartsCount}

{toolState}

); }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("1"); // Manually trigger addToolResult to simulate user confirming await act(async () => { if (chatInstance) { await chatInstance.addToolResult({ tool: "askForConfirmation", toolCallId: "confirm-1", output: { confirmed: true } }); } }); // There should still be exactly ONE tool part with this toolCallId await expect .element(screen.getByTestId("tool-parts-count")) .toHaveTextContent("1"); // The tool state should be updated to output-available await expect .element(screen.getByTestId("tool-state")) .toHaveTextContent("output-available"); }); }); describe("useAgentChat setMessages", () => { it("should handle functional updater and sync resolved messages to server", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "set-messages-test", url: "ws://localhost:3000/agents/chat/set-messages-test?_pk=abc", send: (data: string) => sentMessages.push(data) }); const initialMessages: UIMessage[] = [ { id: "msg-1", role: "user", parts: [{ type: "text", text: "Hello" }] }, { id: "msg-2", role: "assistant", parts: [{ type: "text", text: "Hi there!" }] } ]; let chatInstance: ReturnType | null = null; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages) }); chatInstance = chat; return

{chat.messages.length}

; }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("2"); // Use functional updater to append a message const newMessage: UIMessage = { id: "msg-3", role: "user", parts: [{ type: "text", text: "Follow up" }] }; await act(async () => { chatInstance!.setMessages((prev) => [...prev, newMessage]); await sleep(10); }); await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("3"); // Verify the server received the RESOLVED messages (not empty array) const chatMessagesSent = sentMessages .map((m) => JSON.parse(m)) .filter((m) => m.type === "cf_agent_chat_messages"); expect(chatMessagesSent.length).toBeGreaterThan(0); const lastSent = chatMessagesSent[chatMessagesSent.length - 1]; // Should have the full 3 messages, NOT an empty array expect(lastSent.messages.length).toBe(3); expect(lastSent.messages[2].id).toBe("msg-3"); }); it("should handle array setMessages and sync to server", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "set-messages-array-test", url: "ws://localhost:3000/agents/chat/set-messages-array-test?_pk=abc", send: (data: string) => sentMessages.push(data) }); let chatInstance: ReturnType | null = null; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] as UIMessage[] }); chatInstance = chat; return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); // Set messages with an array directly const newMessages: UIMessage[] = [ { id: "arr-1", role: "user", parts: [{ type: "text", text: "Replaced" }] } ]; await act(async () => { chatInstance!.setMessages(newMessages); await sleep(10); }); // Verify the server received the array const chatMessagesSent = sentMessages .map((m) => JSON.parse(m)) .filter((m) => m.type === "cf_agent_chat_messages"); expect(chatMessagesSent.length).toBeGreaterThan(0); const lastSent = chatMessagesSent[chatMessagesSent.length - 1]; expect(lastSent.messages.length).toBe(1); expect(lastSent.messages[0].id).toBe("arr-1"); }); }); describe("useAgentChat clearHistory", () => { it("should clear local state and send CF_AGENT_CHAT_CLEAR to server", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "clear-test", url: "ws://localhost:3000/agents/chat/clear-test?_pk=abc", send: (data: string) => sentMessages.push(data) }); const initialMessages: UIMessage[] = [ { id: "clear-1", role: "user", parts: [{ type: "text", text: "Hello" }] } ]; let chatInstance: ReturnType | null = null; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages) }); chatInstance = chat; return

{chat.messages.length}

; }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("1"); await act(async () => { chatInstance!.clearHistory(); await sleep(10); }); await expect .element(screen.getByTestId("messages-count")) .toHaveTextContent("0"); // Verify CF_AGENT_CHAT_CLEAR was sent const clearMessages = sentMessages .map((m) => JSON.parse(m)) .filter((m) => m.type === "cf_agent_chat_clear"); expect(clearMessages.length).toBe(1); }); }); describe("useAgentChat autoContinueAfterToolResult default", () => { it("should send autoContinue: true by default with tool results", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "auto-continue-default", url: "ws://localhost:3000/agents/chat/auto-continue-default?_pk=abc", send: (data: string) => sentMessages.push(data) }); const initialMessages: UIMessage[] = [ { id: "msg-1", role: "assistant", parts: [ { type: "tool-getLocation", toolCallId: "tc-default-1", state: "input-available", input: {} } ] } ]; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages), // No explicit autoContinueAfterToolResult — should default to true onToolCall: ({ toolCall, addToolOutput }) => { addToolOutput({ toolCallId: toolCall.toolCallId, output: { lat: 51.5, lng: -0.1 } }); } }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(50); }); // Find the CF_AGENT_TOOL_RESULT message const toolResultMessages = sentMessages .map((m) => JSON.parse(m)) .filter((m) => m.type === "cf_agent_tool_result"); expect(toolResultMessages.length).toBeGreaterThanOrEqual(1); // Default should be autoContinue: true expect(toolResultMessages[0].autoContinue).toBe(true); }); it("should send autoContinue: false when explicitly disabled", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "auto-continue-disabled", url: "ws://localhost:3000/agents/chat/auto-continue-disabled?_pk=abc", send: (data: string) => sentMessages.push(data) }); const initialMessages: UIMessage[] = [ { id: "msg-1", role: "assistant", parts: [ { type: "tool-getLocation", toolCallId: "tc-disabled-1", state: "input-available", input: {} } ] } ]; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages), autoContinueAfterToolResult: false, // Explicitly disabled onToolCall: ({ toolCall, addToolOutput }) => { addToolOutput({ toolCallId: toolCall.toolCallId, output: { lat: 51.5, lng: -0.1 } }); } }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(50); }); const toolResultMessages = sentMessages .map((m) => JSON.parse(m)) .filter((m) => m.type === "cf_agent_tool_result"); expect(toolResultMessages.length).toBeGreaterThanOrEqual(1); expect(toolResultMessages[0].autoContinue).toBe(false); }); it("should send autoContinue: true by default with tool approvals", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "auto-continue-approval", url: "ws://localhost:3000/agents/chat/auto-continue-approval?_pk=abc", send: (data: string) => sentMessages.push(data) }); // Tool part must have approval.id so the wrapper can find the toolCallId const initialMessages: UIMessage[] = [ { id: "msg-1", role: "assistant", parts: [ { type: "tool-dangerousAction", toolCallId: "tc-approval-1", state: "approval-requested", input: { action: "delete" }, approval: { id: "approval-req-1" } } ] } ]; let chatInstance: ReturnType | null = null; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages) // No explicit autoContinueAfterToolResult — should default to true }); chatInstance = chat; return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(50); }); // Send approval via the hook using the approval request ID await act(async () => { if (chatInstance) { chatInstance.addToolApprovalResponse({ id: "approval-req-1", approved: true }); } await sleep(10); }); // Find the CF_AGENT_TOOL_APPROVAL message const approvalMessages = sentMessages .map((m) => JSON.parse(m)) .filter((m) => m.type === "cf_agent_tool_approval"); expect(approvalMessages.length).toBeGreaterThanOrEqual(1); expect(approvalMessages[0].autoContinue).toBe(true); expect(approvalMessages[0].approved).toBe(true); }); }); describe("useAgentChat onToolCall", () => { it("should fire onToolCall for input-available tool parts", async () => { const agent = createAgent({ name: "ontoolcall-test", url: "ws://localhost:3000/agents/chat/ontoolcall-test?_pk=abc", send: () => {} }); const toolCallReceived = vi.fn(); const initialMessages: UIMessage[] = [ { id: "msg-tool-1", role: "assistant", parts: [ { type: "tool-getLocation", toolCallId: "tc-1", state: "input-available", input: { query: "current" } } ] } ]; const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages), onToolCall: ({ toolCall, addToolOutput }) => { toolCallReceived(toolCall); addToolOutput({ toolCallId: toolCall.toolCallId, output: { lat: 40.7, lng: -74.0 } }); } }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(50); }); // onToolCall should have been called with the tool call details expect(toolCallReceived).toHaveBeenCalledWith( expect.objectContaining({ toolCallId: "tc-1", toolName: "getLocation", input: { query: "current" } }) ); }); }); describe("useAgentChat re-render stability", () => { it("should not cause infinite re-renders when idle", async () => { const agent = createAgent({ name: "rerender-idle", url: "ws://localhost:3000/agents/chat/rerender-idle?_pk=abc" }); let renderCount = 0; const TestComponent = () => { renderCount++; const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); // Capture render count after initial mount const afterMountCount = renderCount; // Wait to see if more renders happen (would indicate an infinite loop) await sleep(200); // In Strict Mode, React double-renders. After mount stabilizes, // there should be NO additional renders (no infinite loop). expect(renderCount).toBe(afterMountCount); }); it("should not re-render excessively when messages are set", async () => { const agent = createAgent({ name: "rerender-messages", url: "ws://localhost:3000/agents/chat/rerender-messages?_pk=abc" }); let renderCount = 0; let chatInstance: ReturnType | null = null; const TestComponent = () => { renderCount++; const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] as UIMessage[] }); chatInstance = chat; return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); const beforeSetMessages = renderCount; // Set messages await act(async () => { chatInstance!.setMessages([ { id: "msg-1", role: "user", parts: [{ type: "text", text: "Hello" }] } ]); await sleep(10); }); const afterSetMessages = renderCount; // Wait to see if renders stabilize await sleep(200); // Should have re-rendered for the setMessages call but then stopped. // Allow some re-renders (React batching, state updates) but not infinite. const rendersFromSetMessages = afterSetMessages - beforeSetMessages; expect(rendersFromSetMessages).toBeLessThan(10); // No additional renders after stabilizing expect(renderCount).toBe(afterSetMessages); }); it("should stabilize after receiving a broadcast message", async () => { const target = new EventTarget(); const agent = createAgent({ name: "rerender-broadcast", url: "ws://localhost:3000/agents/chat/rerender-broadcast?_pk=abc" }); // Override addEventListener/removeEventListener to use our target (agent as unknown as Record).addEventListener = target.addEventListener.bind(target); (agent as unknown as Record).removeEventListener = target.removeEventListener.bind(target); let renderCount = 0; const TestComponent = () => { renderCount++; const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); const beforeBroadcast = renderCount; // Simulate a server broadcast (CF_AGENT_CHAT_MESSAGES) await act(async () => { target.dispatchEvent( new MessageEvent("message", { data: JSON.stringify({ type: "cf_agent_chat_messages", messages: [ { id: "broadcast-1", role: "user", parts: [{ type: "text", text: "From other tab" }] } ] }) }) ); await sleep(10); }); const afterBroadcast = renderCount; // Wait for stabilization await sleep(200); // Should have re-rendered for the broadcast but then stopped const rendersFromBroadcast = afterBroadcast - beforeBroadcast; expect(rendersFromBroadcast).toBeGreaterThan(0); // Must have re-rendered expect(rendersFromBroadcast).toBeLessThan(10); // But not infinitely // No additional renders after stabilizing expect(renderCount).toBe(afterBroadcast); }); }); describe("useAgentChat body option", () => { it("should include static body fields in sent messages", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "body-static-test", url: "ws://localhost:3000/agents/chat/body-static-test?_pk=abc", send: (data: string) => sentMessages.push(data) }); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [], body: { timezone: "America/New_York", userId: "user-123" } }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); // The body fields should be included when the transport sends messages // We can verify by checking that the component rendered without errors // (the actual body merging is tested via the sent WS messages) expect(sentMessages).toBeDefined(); }); it("should include dynamic body fields from function", async () => { const sentMessages: string[] = []; let callCount = 0; const agent = createAgent({ name: "body-dynamic-test", url: "ws://localhost:3000/agents/chat/body-dynamic-test?_pk=abc", send: (data: string) => sentMessages.push(data) }); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [], body: () => { callCount++; return { timestamp: Date.now(), requestNumber: callCount }; } }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); // Component should render without errors with function body expect(callCount).toBeDefined(); }); it("should work alongside prepareSendMessagesRequest", async () => { const sentMessages: string[] = []; const agent = createAgent({ name: "body-combined-test", url: "ws://localhost:3000/agents/chat/body-combined-test?_pk=abc", send: (data: string) => sentMessages.push(data) }); const prepareSendMessagesRequest = vi.fn(() => ({ body: { fromPrepare: true } })); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [], body: { fromBody: true }, prepareSendMessagesRequest }); return

{chat.messages.length}

; }; await act(async () => { render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); }); // Both body and prepareSendMessagesRequest should coexist without errors expect(sentMessages).toBeDefined(); }); }); describe("useAgentChat stale agent ref (issue #929)", () => { it("should use the new agent's send method after agent switch, not the old one", async () => { const oldSend = vi.fn(); const newSend = vi.fn(); const agentOld = createAgent({ name: "thread-old", url: "ws://localhost:3000/agents/chat/thread-old?_pk=old", send: oldSend }); const agentNew = createAgent({ name: "thread-new", url: "ws://localhost:3000/agents/chat/thread-new?_pk=new", send: newSend }); let chatInstance: ReturnType | null = null; const TestComponent = ({ agent }: { agent: ReturnType; }) => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] as UIMessage[] }); chatInstance = chat; return

{chat.status}

; }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); // Switch to the new agent await act(async () => { screen.rerender(); await sleep(10); }); // Clear any sends that happened during setup (e.g., stream resume requests) oldSend.mockClear(); newSend.mockClear(); // Clear history triggers agent.send() — this should go to the NEW agent await act(async () => { chatInstance!.clearHistory(); await sleep(10); }); // The clear message should have been sent to the NEW agent, not the old one const newSendCalls = newSend.mock.calls .map((args) => JSON.parse(args[0] as string)) .filter((m: Record) => m.type === "cf_agent_chat_clear"); expect(newSendCalls.length).toBe(1); // The old agent should NOT have received the clear message const oldSendCalls = oldSend.mock.calls .map((args) => JSON.parse(args[0] as string)) .filter((m: Record) => m.type === "cf_agent_chat_clear"); expect(oldSendCalls.length).toBe(0); }); }); describe("useAgentChat stream resumption (issue #896)", () => { function createAgentWithTarget({ name, url }: { name: string; url: string }) { const target = new EventTarget(); const sentMessages: string[] = []; const agent = createAgent({ name, url, send: (data: string) => sentMessages.push(data) }); // Wire up the target so we can dispatch messages to the hook (agent as unknown as Record).addEventListener = target.addEventListener.bind(target); (agent as unknown as Record).removeEventListener = target.removeEventListener.bind(target); return { agent, target, sentMessages }; } function dispatch(target: EventTarget, data: Record) { target.dispatchEvent( new MessageEvent("message", { data: JSON.stringify(data) }) ); } it("should process resumed stream chunks progressively and update status", async () => { const { agent, target } = createAgentWithTarget({ name: "replay-complete-test", url: "ws://localhost:3000/agents/chat/replay-complete-test?_pk=abc" }); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] as UIMessage[] }); const assistantMsg = chat.messages.find( (m: UIMessage) => m.role === "assistant" ); const textPart = assistantMsg?.parts.find( (p: UIMessage["parts"][number]) => p.type === "text" ) as { text?: string } | undefined; return (

{chat.messages.length}

{textPart?.text ?? ""}

{chat.status}

); }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); // Initially no messages await expect.element(screen.getByTestId("count")).toHaveTextContent("0"); // Simulate server sending CF_AGENT_STREAM_RESUMING // The transport's reconnectToStream picks this up and returns a ReadableStream await act(async () => { dispatch(target, { type: "cf_agent_stream_resuming", id: "req-1" }); await sleep(10); }); // Simulate replay chunks — now processed progressively by useChat's pipeline await act(async () => { dispatch(target, { type: "cf_agent_use_chat_response", id: "req-1", body: '{"type":"text-start","id":"t1"}', done: false, replay: true }); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-1", body: '{"type":"text-delta","id":"t1","delta":"Hello world"}', done: false, replay: true }); await sleep(10); }); // Chunks are processed progressively by useChat — message appears immediately await expect.element(screen.getByTestId("count")).toHaveTextContent("1"); await expect .element(screen.getByTestId("text")) .toHaveTextContent("Hello world"); }); it("should flush and finalize after done:true for orphaned streams", async () => { const { agent, target } = createAgentWithTarget({ name: "orphaned-done-test", url: "ws://localhost:3000/agents/chat/orphaned-done-test?_pk=abc" }); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] as UIMessage[] }); const assistantMsg = chat.messages.find( (m: UIMessage) => m.role === "assistant" ); const textPart = assistantMsg?.parts.find( (p: UIMessage["parts"][number]) => p.type === "text" ) as { text?: string } | undefined; return (

{chat.messages.length}

{textPart?.text ?? ""}

{chat.status}

); }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); // Simulate resume + replay + done (orphaned stream path) await act(async () => { dispatch(target, { type: "cf_agent_stream_resuming", id: "req-orphaned" }); await sleep(5); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-orphaned", body: '{"type":"text-start","id":"t1"}', done: false, replay: true }); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-orphaned", body: '{"type":"text-delta","id":"t1","delta":"partial from hibernation"}', done: false, replay: true }); // done:true signals orphaned stream is finalized dispatch(target, { type: "cf_agent_use_chat_response", id: "req-orphaned", body: "", done: true, replay: true }); await sleep(10); }); // Message should be flushed with the accumulated text await expect.element(screen.getByTestId("count")).toHaveTextContent("1"); await expect .element(screen.getByTestId("text")) .toHaveTextContent("partial from hibernation"); }); it("should continue receiving live chunks after replayComplete", async () => { const { agent, target } = createAgentWithTarget({ name: "replay-then-live-test", url: "ws://localhost:3000/agents/chat/replay-then-live-test?_pk=abc" }); const TestComponent = () => { const chat = useAgentChat({ agent, getInitialMessages: null, messages: [] as UIMessage[] }); const assistantMsg = chat.messages.find( (m: UIMessage) => m.role === "assistant" ); const textPart = assistantMsg?.parts.find( (p: UIMessage["parts"][number]) => p.type === "text" ) as { text?: string } | undefined; return (

{chat.messages.length}

{textPart?.text ?? ""}

); }; const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); // Replay phase await act(async () => { dispatch(target, { type: "cf_agent_stream_resuming", id: "req-live" }); await sleep(5); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-live", body: '{"type":"text-start","id":"t1"}', done: false, replay: true }); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-live", body: '{"type":"text-delta","id":"t1","delta":"replayed-"}', done: false, replay: true }); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-live", body: "", done: false, replay: true, replayComplete: true }); await sleep(10); }); // After replay, message should show replayed text await expect.element(screen.getByTestId("count")).toHaveTextContent("1"); await expect .element(screen.getByTestId("text")) .toHaveTextContent("replayed-"); // Now simulate a live chunk arriving (no replay flag) await act(async () => { dispatch(target, { type: "cf_agent_use_chat_response", id: "req-live", body: '{"type":"text-delta","id":"t1","delta":"and live!"}', done: false }); await sleep(10); }); // The live chunk should append to the same message await expect.element(screen.getByTestId("count")).toHaveTextContent("1"); await expect .element(screen.getByTestId("text")) .toHaveTextContent("replayed-and live!"); }); }); describe("useAgentChat tool approval continuations (issue #1108)", () => { function createAgentWithTarget({ name, url }: { name: string; url: string }) { const target = new EventTarget(); const agent = createAgent({ name, url }); (agent as unknown as Record).addEventListener = target.addEventListener.bind(target); (agent as unknown as Record).removeEventListener = target.removeEventListener.bind(target); return { agent, target }; } function dispatch(target: EventTarget, data: Record) { target.dispatchEvent( new MessageEvent("message", { data: JSON.stringify(data) }) ); } const initialMessages: UIMessage[] = [ { id: "assistant-local", role: "assistant", parts: [ { type: "tool-dangerousAction", toolCallId: "tc-approval-1", state: "approval-responded", input: { action: "delete" }, approval: { id: "approval-req-1", approved: true } } ] } ]; function TestComponent({ agent }: { agent: ReturnType }) { const chat = useAgentChat({ agent, getInitialMessages: () => Promise.resolve(initialMessages) }); const assistantMessages = chat.messages.filter( (message) => message.role === "assistant" ); const textPart = assistantMessages .flatMap((message) => message.parts) .find((part) => part.type === "text") as { text?: string } | undefined; const toolPartsCount = assistantMessages.reduce((count, message) => { return ( count + message.parts.filter( (part) => "toolCallId" in part && part.toolCallId === "tc-approval-1" ).length ); }, 0); return (

{assistantMessages.length}

{assistantMessages.map((message) => message.id).join(",")}

{toolPartsCount}

{textPart?.text ?? ""}

); } it("keeps the existing assistant id for continuation start chunks", async () => { const { agent, target } = createAgentWithTarget({ name: "continuation-start-id", url: "ws://localhost:3000/agents/chat/continuation-start-id?_pk=abc" }); const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); await expect .element(screen.getByTestId("assistant-count")) .toHaveTextContent("1"); await act(async () => { dispatch(target, { type: "cf_agent_use_chat_response", id: "req-continuation-start", continuation: true, body: '{"type":"start","messageId":"assistant-stream"}', done: false }); await sleep(10); }); await expect .element(screen.getByTestId("assistant-count")) .toHaveTextContent("1"); await expect .element(screen.getByTestId("assistant-ids")) .toHaveTextContent("assistant-local"); await expect .element(screen.getByTestId("tool-parts-count")) .toHaveTextContent("1"); }); it("keeps merging continuations when broadcasts replace assistant ids mid-stream", async () => { const { agent, target } = createAgentWithTarget({ name: "continuation-remap-id", url: "ws://localhost:3000/agents/chat/continuation-remap-id?_pk=abc" }); const screen = await act(async () => { const screen = render(, { wrapper: ({ children }) => ( {children} ) }); await sleep(10); return screen; }); await act(async () => { dispatch(target, { type: "cf_agent_use_chat_response", id: "req-continuation-remap", continuation: true, body: '{"type":"start","messageId":"assistant-stream"}', done: false }); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-continuation-remap", continuation: true, body: '{"type":"text-start","id":"text-1"}', done: false }); await sleep(10); }); await act(async () => { dispatch(target, { type: "cf_agent_chat_messages", messages: [ { id: "assistant-server", role: "assistant", parts: [ { type: "tool-dangerousAction", toolCallId: "tc-approval-1", state: "approval-responded", input: { action: "delete" }, approval: { id: "approval-req-1", approved: true } } ] } ] }); dispatch(target, { type: "cf_agent_use_chat_response", id: "req-continuation-remap", continuation: true, body: '{"type":"text-delta","id":"text-1","delta":"done"}', done: false }); await sleep(10); }); await expect .element(screen.getByTestId("assistant-count")) .toHaveTextContent("1"); await expect .element(screen.getByTestId("assistant-ids")) .toHaveTextContent("assistant-server"); await expect .element(screen.getByTestId("tool-parts-count")) .toHaveTextContent("1"); await expect.element(screen.getByTestId("text")).toHaveTextContent("done"); }); });