Skip to content

Commit c557334

Browse files
MarioCadenasMarioCadenas
committed
feat(appkit-ui): useAgentChat React hook wrapping connectSSE
Reviewer agentic finding #7 flagged ~20 lines of SSE parsing boilerplate in `template/client/src/pages/agents/AgentChat.tsx` that every scaffolded app would copy verbatim. The reviewer suggested extracting it; the cleaner shape is a thin React hook around the existing `connectSSE` utility in `@databricks/appkit-ui/js`, which already owns the buffer cap, abort signal composition, frame parsing, and retry/backoff logic. useAgentChat (new, in @databricks/appkit-ui/react) Single-stream chat primitive for agents-plugin chat endpoints. Wraps connectSSE and adds the stateful glue every chat UI ends up writing: - `content`: accumulates `response.output_text.delta` deltas - `events`: chronological log of every parsed event for replay / inspection - `threadId`: captured from the first `appkit.metadata` event and automatically forwarded on the next `send()` so the server reuses the same thread - `isStreaming` / `error` - `send(message)`: POSTs to `/api/agents/chat` (configurable via `endpoint`), streams the SSE response, parses each JSON-encoded data line, aborts any prior in-flight stream - `reset()`: drops state and aborts the active stream The hook stays narrow on purpose: one stream at a time, no multi-turn message-history ownership (callers compose that). Tool calls, approval gates, and any non-text event are surfaced through `onEvent` for the caller to render however it wants. `maxRetries: 0` is hard-wired into the connectSSE call — chat turns aren't idempotent and re-sending the payload on transient failure would either duplicate the user message or depend on server-side Last-Event-ID resumption (the agents plugin's StreamManager supports it, but failure-mode auditing is easier with retries explicitly off and the responsibility on the caller to re-send). Template migration `template/client/src/pages/agents/AgentChat.tsx` drops the inline fetch/decoder/SSE-buffer/parser block (~22 lines) and uses `useAgentChat({ agent, onEvent })` instead. Tool-call rows are pushed into the local `messages` array via the `onEvent` callback; the streaming assistant text mirrors from the hook's `content` field into the pending assistant row. Dev-playground's `useAgentStream` is left in place for now — it carries inspector + action-dispatcher wiring that doesn't fit on this hook. Migrating it onto `useAgentChat` is a follow-up that also covers the related Agentic finding #21 (duplicate SSE parsing across consumers). Tests: 12 new cases on useAgentChat covering posting payload shape, custom endpoint, delta accumulation, threadId capture + reuse, onEvent dispatch (including handlers that throw), malformed payload skipping, isStreaming lifecycle, reset(), send-while-streaming abort, and onError surfacing. 2310 tests passing across the workspace. Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
1 parent d82c518 commit c557334

4 files changed

Lines changed: 662 additions & 118 deletions

File tree

‎packages/appkit-ui/src/react/hooks/__tests__/use-agent-chat.test.ts‎

Lines changed: 325 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,325 @@
1+
import { act, renderHook, waitFor } from "@testing-library/react";
2+
import { afterEach, describe, expect, test, vi } from "vitest";
3+
4+
let capturedCallbacks: {
5+
onMessage?: (msg: { data: string }) => Promise<void>;
6+
onError?: (err: Error) => void;
7+
signal?: AbortSignal;
8+
url?: string;
9+
payload?: unknown;
10+
maxRetries?: number;
11+
} = {};
12+
13+
let resolveStream: (() => void) | null = null;
14+
let rejectStream: ((err: Error) => void) | null = null;
15+
16+
const mockConnectSSE = vi.fn().mockImplementation((opts: any) => {
17+
capturedCallbacks = {
18+
onMessage: opts.onMessage,
19+
onError: opts.onError,
20+
signal: opts.signal,
21+
url: opts.url,
22+
payload: opts.payload,
23+
maxRetries: opts.maxRetries,
24+
};
25+
return new Promise<void>((resolve, reject) => {
26+
resolveStream = resolve;
27+
rejectStream = reject;
28+
});
29+
});
30+
31+
vi.mock("@/js", () => ({
32+
connectSSE: (...args: unknown[]) => mockConnectSSE(...args),
33+
}));
34+
35+
import { useAgentChat } from "../use-agent-chat";
36+
37+
async function emit(data: string) {
38+
// Allow microtasks to settle before pushing the next message.
39+
await capturedCallbacks.onMessage?.({ data });
40+
}
41+
42+
describe("useAgentChat", () => {
43+
afterEach(() => {
44+
capturedCallbacks = {};
45+
resolveStream = null;
46+
rejectStream = null;
47+
vi.clearAllMocks();
48+
});
49+
50+
test("initial state is idle", () => {
51+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
52+
53+
expect(result.current.content).toBe("");
54+
expect(result.current.events).toEqual([]);
55+
expect(result.current.threadId).toBeNull();
56+
expect(result.current.isStreaming).toBe(false);
57+
expect(result.current.error).toBeNull();
58+
expect(typeof result.current.send).toBe("function");
59+
expect(typeof result.current.reset).toBe("function");
60+
});
61+
62+
test("send() posts to /api/agents/chat with the agent name and message", async () => {
63+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
64+
65+
act(() => {
66+
void result.current.send("hello");
67+
});
68+
69+
await waitFor(() => expect(mockConnectSSE).toHaveBeenCalled());
70+
71+
expect(capturedCallbacks.url).toBe("/api/agents/chat");
72+
expect(capturedCallbacks.payload).toEqual({
73+
message: "hello",
74+
agent: "helper",
75+
});
76+
// Chat turns are not safely retryable — assert we explicitly opt out.
77+
expect(capturedCallbacks.maxRetries).toBe(0);
78+
});
79+
80+
test("custom endpoint is forwarded to connectSSE", async () => {
81+
const { result } = renderHook(() =>
82+
useAgentChat({ agent: "helper", endpoint: "/v2/chat" }),
83+
);
84+
85+
act(() => {
86+
void result.current.send("hi");
87+
});
88+
89+
await waitFor(() => expect(mockConnectSSE).toHaveBeenCalled());
90+
expect(capturedCallbacks.url).toBe("/v2/chat");
91+
});
92+
93+
test("accumulates response.output_text.delta into content", async () => {
94+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
95+
96+
act(() => {
97+
void result.current.send("hi");
98+
});
99+
100+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
101+
102+
await act(async () => {
103+
await emit(
104+
JSON.stringify({
105+
type: "response.output_text.delta",
106+
delta: "Hello, ",
107+
}),
108+
);
109+
await emit(
110+
JSON.stringify({ type: "response.output_text.delta", delta: "world" }),
111+
);
112+
});
113+
114+
expect(result.current.content).toBe("Hello, world");
115+
});
116+
117+
test("captures threadId from appkit.metadata and reuses it on next send()", async () => {
118+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
119+
120+
act(() => {
121+
void result.current.send("first");
122+
});
123+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
124+
125+
await act(async () => {
126+
await emit(
127+
JSON.stringify({
128+
type: "appkit.metadata",
129+
data: { threadId: "t-123" },
130+
}),
131+
);
132+
});
133+
134+
expect(result.current.threadId).toBe("t-123");
135+
136+
// End the first stream so the next send() opens a new SSE.
137+
await act(async () => {
138+
resolveStream?.();
139+
await new Promise((r) => setTimeout(r, 0));
140+
});
141+
142+
mockConnectSSE.mockClear();
143+
act(() => {
144+
void result.current.send("second");
145+
});
146+
await waitFor(() => expect(mockConnectSSE).toHaveBeenCalled());
147+
148+
expect(capturedCallbacks.payload).toEqual({
149+
message: "second",
150+
agent: "helper",
151+
threadId: "t-123",
152+
});
153+
});
154+
155+
test("onEvent is invoked for every parsed event", async () => {
156+
const onEvent = vi.fn();
157+
const { result } = renderHook(() =>
158+
useAgentChat({ agent: "helper", onEvent }),
159+
);
160+
161+
act(() => {
162+
void result.current.send("hi");
163+
});
164+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
165+
166+
await act(async () => {
167+
await emit(
168+
JSON.stringify({ type: "response.output_text.delta", delta: "a" }),
169+
);
170+
await emit(
171+
JSON.stringify({
172+
type: "response.output_item.added",
173+
item: { type: "function_call", name: "get_weather", arguments: "{}" },
174+
}),
175+
);
176+
});
177+
178+
expect(onEvent).toHaveBeenCalledTimes(2);
179+
expect(onEvent).toHaveBeenNthCalledWith(
180+
1,
181+
expect.objectContaining({
182+
type: "response.output_text.delta",
183+
delta: "a",
184+
}),
185+
);
186+
expect(onEvent).toHaveBeenNthCalledWith(
187+
2,
188+
expect.objectContaining({
189+
type: "response.output_item.added",
190+
item: expect.objectContaining({ name: "get_weather" }),
191+
}),
192+
);
193+
});
194+
195+
test("throwing onEvent handler does not break the stream", async () => {
196+
const onEvent = vi.fn(() => {
197+
throw new Error("handler bug");
198+
});
199+
const { result } = renderHook(() =>
200+
useAgentChat({ agent: "helper", onEvent }),
201+
);
202+
203+
act(() => {
204+
void result.current.send("hi");
205+
});
206+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
207+
208+
await act(async () => {
209+
await emit(
210+
JSON.stringify({ type: "response.output_text.delta", delta: "x" }),
211+
);
212+
});
213+
214+
// Despite onEvent throwing, content still accumulated.
215+
expect(result.current.content).toBe("x");
216+
});
217+
218+
test("malformed event payloads are skipped silently", async () => {
219+
const onEvent = vi.fn();
220+
const { result } = renderHook(() =>
221+
useAgentChat({ agent: "helper", onEvent }),
222+
);
223+
224+
act(() => {
225+
void result.current.send("hi");
226+
});
227+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
228+
229+
await act(async () => {
230+
await emit("not-json");
231+
await emit("[DONE]");
232+
await emit("");
233+
await emit(
234+
JSON.stringify({ type: "response.output_text.delta", delta: "ok" }),
235+
);
236+
});
237+
238+
expect(result.current.content).toBe("ok");
239+
expect(onEvent).toHaveBeenCalledTimes(1);
240+
});
241+
242+
test("isStreaming toggles around the connectSSE lifecycle", async () => {
243+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
244+
245+
expect(result.current.isStreaming).toBe(false);
246+
247+
act(() => {
248+
void result.current.send("hi");
249+
});
250+
await waitFor(() => expect(result.current.isStreaming).toBe(true));
251+
252+
await act(async () => {
253+
resolveStream?.();
254+
await new Promise((r) => setTimeout(r, 0));
255+
});
256+
await waitFor(() => expect(result.current.isStreaming).toBe(false));
257+
});
258+
259+
test("reset() clears content, events, threadId, and aborts in-flight stream", async () => {
260+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
261+
262+
act(() => {
263+
void result.current.send("hi");
264+
});
265+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
266+
267+
await act(async () => {
268+
await emit(
269+
JSON.stringify({ type: "appkit.metadata", data: { threadId: "t-1" } }),
270+
);
271+
await emit(
272+
JSON.stringify({ type: "response.output_text.delta", delta: "x" }),
273+
);
274+
});
275+
276+
expect(result.current.threadId).toBe("t-1");
277+
expect(result.current.content).toBe("x");
278+
279+
const signal = capturedCallbacks.signal;
280+
expect(signal?.aborted).toBe(false);
281+
282+
act(() => {
283+
result.current.reset();
284+
});
285+
286+
expect(signal?.aborted).toBe(true);
287+
expect(result.current.content).toBe("");
288+
expect(result.current.events).toEqual([]);
289+
expect(result.current.threadId).toBeNull();
290+
expect(result.current.isStreaming).toBe(false);
291+
});
292+
293+
test("send() while a previous stream is in flight aborts the previous one", async () => {
294+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
295+
296+
act(() => {
297+
void result.current.send("first");
298+
});
299+
await waitFor(() => expect(capturedCallbacks.onMessage).toBeDefined());
300+
const firstSignal = capturedCallbacks.signal;
301+
302+
act(() => {
303+
void result.current.send("second");
304+
});
305+
expect(firstSignal?.aborted).toBe(true);
306+
});
307+
308+
test("onError surfaces a string error message", async () => {
309+
const { result } = renderHook(() => useAgentChat({ agent: "helper" }));
310+
311+
act(() => {
312+
void result.current.send("hi");
313+
});
314+
await waitFor(() => expect(capturedCallbacks.onError).toBeDefined());
315+
316+
await act(async () => {
317+
capturedCallbacks.onError?.(new Error("upstream 500"));
318+
resolveStream?.();
319+
await new Promise((r) => setTimeout(r, 0));
320+
});
321+
322+
expect(result.current.error).toBe("upstream 500");
323+
expect(result.current.isStreaming).toBe(false);
324+
});
325+
});

‎packages/appkit-ui/src/react/hooks/index.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,12 @@ export type {
1313
UseAnalyticsQueryOptions,
1414
UseAnalyticsQueryResult,
1515
} from "./types";
16+
export {
17+
type AgentChatEvent,
18+
type UseAgentChatOptions,
19+
type UseAgentChatResult,
20+
useAgentChat,
21+
} from "./use-agent-chat";
1622
export { useAnalyticsQuery } from "./use-analytics-query";
1723
export {
1824
type UseChartDataOptions,

0 commit comments

Comments
 (0)