branch:
use-agent-chat.test.tsx
55852 bytesRaw
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<typeof useAgent>;
}
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 <div data-testid="messages">{JSON.stringify(chat.messages)}</div>;
};
const suspenseRendered = vi.fn();
const SuspenseObserver = () => {
suspenseRendered();
return "Suspended";
};
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"))
.toHaveTextContent(JSON.stringify(testMessages));
expect(getInitialMessages).toHaveBeenCalledTimes(1);
expect(suspenseRendered).toHaveBeenCalled();
suspenseRendered.mockClear();
await screen.rerender(<TestComponent />);
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<typeof useAgent>;
}) => {
const chat = useAgentChat({
agent,
getInitialMessages
});
return <div data-testid="messages">{JSON.stringify(chat.messages)}</div>;
};
const suspenseRendered = vi.fn();
const SuspenseObserver = () => {
suspenseRendered();
return "Suspended";
};
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"))
.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(<TestComponent agent={agentB} />);
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<UIMessage>
): 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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
const screen = await act(() =>
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
})
);
// 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<UIMessage>
): Promise<PrepareSendMessagesRequestResult> => {
// 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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
const screen = await act(() =>
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
})
);
// 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<string, AITool<unknown, unknown>> = {
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
const screen = await act(() =>
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
})
);
// 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<string, AITool> = {
showAlert: {
description: "Shows an alert",
execute: async () => ({ shown: true })
}
};
const prepareSendMessagesRequest = vi.fn(
(
_options: PrepareSendMessagesRequestOptions<UIMessage>
): 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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
const screen = await act(() =>
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
})
);
// 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<string, AITool> = {
showAlert: {
description: "Shows an alert",
parameters: {
type: "object",
properties: { message: { type: "string" } }
},
execute: mockExecute
}
};
const TestComponent = () => {
const chat = useAgentChat({
agent,
getInitialMessages: null,
tools
});
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
const screen = await act(() =>
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
})
);
// 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 (
<div>
<div data-testid="messages-count">{chat.messages.length}</div>
<div data-testid="tool-state">{toolState}</div>
</div>
);
};
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");
// 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<typeof useAgentChat> | 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 (
<div>
<div data-testid="messages-count">{chat.messages.length}</div>
<div data-testid="tool-parts-count">{toolPartsCount}</div>
<div data-testid="tool-state">{toolState}</div>
</div>
);
};
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"))
.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<typeof useAgentChat> | null = null;
const TestComponent = () => {
const chat = useAgentChat({
agent,
getInitialMessages: () => Promise.resolve(initialMessages)
});
chatInstance = chat;
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
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"))
.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<typeof useAgentChat> | null = null;
const TestComponent = () => {
const chat = useAgentChat({
agent,
getInitialMessages: null,
messages: [] as UIMessage[]
});
chatInstance = chat;
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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<typeof useAgentChat> | null = null;
const TestComponent = () => {
const chat = useAgentChat({
agent,
getInitialMessages: () => Promise.resolve(initialMessages)
});
chatInstance = chat;
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
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"))
.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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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<typeof useAgentChat> | null = null;
const TestComponent = () => {
const chat = useAgentChat({
agent,
getInitialMessages: () => Promise.resolve(initialMessages)
// No explicit autoContinueAfterToolResult — should default to true
});
chatInstance = chat;
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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<typeof useAgentChat> | null = null;
const TestComponent = () => {
renderCount++;
const chat = useAgentChat({
agent,
getInitialMessages: null,
messages: [] as UIMessage[]
});
chatInstance = chat;
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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<string, unknown>).addEventListener =
target.addEventListener.bind(target);
(agent as unknown as Record<string, unknown>).removeEventListener =
target.removeEventListener.bind(target);
let renderCount = 0;
const TestComponent = () => {
renderCount++;
const chat = useAgentChat({
agent,
getInitialMessages: null,
messages: []
});
return <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 <div data-testid="messages-count">{chat.messages.length}</div>;
};
await act(async () => {
render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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<typeof useAgentChat> | null = null;
const TestComponent = ({
agent
}: {
agent: ReturnType<typeof useAgent>;
}) => {
const chat = useAgentChat({
agent,
getInitialMessages: null,
messages: [] as UIMessage[]
});
chatInstance = chat;
return <div data-testid="status">{chat.status}</div>;
};
const screen = await act(async () => {
const screen = render(<TestComponent agent={agentOld} />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
await sleep(10);
return screen;
});
// Switch to the new agent
await act(async () => {
screen.rerender(<TestComponent agent={agentNew} />);
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<string, unknown>) => 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<string, unknown>) => 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<string, unknown>).addEventListener =
target.addEventListener.bind(target);
(agent as unknown as Record<string, unknown>).removeEventListener =
target.removeEventListener.bind(target);
return { agent, target, sentMessages };
}
function dispatch(target: EventTarget, data: Record<string, unknown>) {
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 (
<div>
<div data-testid="count">{chat.messages.length}</div>
<div data-testid="text">{textPart?.text ?? ""}</div>
<div data-testid="status">{chat.status}</div>
</div>
);
};
const screen = await act(async () => {
const screen = render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 (
<div>
<div data-testid="count">{chat.messages.length}</div>
<div data-testid="text">{textPart?.text ?? ""}</div>
<div data-testid="status">{chat.status}</div>
</div>
);
};
const screen = await act(async () => {
const screen = render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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 (
<div>
<div data-testid="count">{chat.messages.length}</div>
<div data-testid="text">{textPart?.text ?? ""}</div>
</div>
);
};
const screen = await act(async () => {
const screen = render(<TestComponent />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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<string, unknown>).addEventListener =
target.addEventListener.bind(target);
(agent as unknown as Record<string, unknown>).removeEventListener =
target.removeEventListener.bind(target);
return { agent, target };
}
function dispatch(target: EventTarget, data: Record<string, unknown>) {
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<typeof useAgent> }) {
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 (
<div>
<div data-testid="assistant-count">{assistantMessages.length}</div>
<div data-testid="assistant-ids">
{assistantMessages.map((message) => message.id).join(",")}
</div>
<div data-testid="tool-parts-count">{toolPartsCount}</div>
<div data-testid="text">{textPart?.text ?? ""}</div>
</div>
);
}
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(<TestComponent agent={agent} />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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(<TestComponent agent={agent} />, {
wrapper: ({ children }) => (
<StrictMode>
<Suspense fallback="Loading...">{children}</Suspense>
</StrictMode>
)
});
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");
});
});