@assistant-ui/react-ag-ui 0.0.50 → 0.0.52

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,40 @@
1
+ import { describe, expect, expectTypeOf, it } from "vitest";
2
+ import { ToolSchema, type RunAgentParameters } from "@ag-ui/client";
3
+ import { toAgUiTools } from "./conversions";
4
+
5
+ describe("toAgUiTools", () => {
6
+ it("normalizes a missing description to an empty string", () => {
7
+ const tools = toAgUiTools({
8
+ described: { description: "Has one", parameters: { type: "object" } },
9
+ undescribed: { parameters: { type: "object" } },
10
+ });
11
+
12
+ expect(tools).toEqual([
13
+ {
14
+ name: "described",
15
+ description: "Has one",
16
+ parameters: { type: "object" },
17
+ },
18
+ { name: "undescribed", description: "", parameters: { type: "object" } },
19
+ ]);
20
+ });
21
+
22
+ it("serializes a description-less tool that upstream still accepts", () => {
23
+ const [tool] = toAgUiTools({ undescribed: { parameters: {} } });
24
+
25
+ expect(ToolSchema.safeParse(JSON.parse(JSON.stringify(tool))).success).toBe(
26
+ true,
27
+ );
28
+ });
29
+
30
+ it("produces tools the upstream run parameters accept", () => {
31
+ const tools = toAgUiTools({
32
+ search: { description: "Search", parameters: { type: "object" } },
33
+ });
34
+
35
+ expectTypeOf(tools).toExtend<NonNullable<RunAgentParameters["tools"]>>();
36
+
37
+ const parameters: RunAgentParameters = { runId: "run-1", tools };
38
+ expect(parameters.tools).toHaveLength(1);
39
+ });
40
+ });
@@ -1,6 +1,6 @@
1
1
  "use client";
2
2
 
3
- import type { InputContent } from "@ag-ui/client";
3
+ import type { InputContent, RunAgentParameters } from "@ag-ui/client";
4
4
  import type {
5
5
  ThreadMessageLike as CoreThreadMessageLike,
6
6
  ToolCallMessagePartMcpMetadata,
@@ -915,11 +915,7 @@ export function toAgUiMessages(
915
915
  return converted;
916
916
  }
917
917
 
918
- type AgUiTool = {
919
- name: string;
920
- description: string | undefined;
921
- parameters: unknown;
922
- };
918
+ type AgUiTool = NonNullable<RunAgentParameters["tools"]>[number];
923
919
 
924
920
  export function toAgUiTools(
925
921
  tools: Record<string, Tool> | undefined,
@@ -929,7 +925,7 @@ export function toAgUiTools(
929
925
  const toolsSchema = toToolsJSONSchema(tools);
930
926
  return Object.entries(toolsSchema).map(([name, tool]) => ({
931
927
  name,
932
- description: tool.description,
928
+ description: tool.description ?? "",
933
929
  parameters: tool.parameters,
934
930
  }));
935
931
  }
@@ -71,6 +71,13 @@ export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
71
71
  * Defaults to `true`.
72
72
  */
73
73
  autoCancelPendingToolCalls?: boolean | undefined;
74
+ /**
75
+ * Buffer a message sent while a run is in flight and send it once the run
76
+ * settles, exposing it on `composer.queue` for `ComposerPrimitive.Queue`.
77
+ * The runtime owns the queue lifecycle because flushing needs the agent's
78
+ * send path and the run's own busy and idle edges.
79
+ */
80
+ unstable_enableMessageQueue?: boolean | undefined;
74
81
  onError?: (e: Error) => void;
75
82
  onCancel?: () => void;
76
83
  adapters?: UseAgUiRuntimeAdapters;
@@ -0,0 +1,352 @@
1
+ // @vitest-environment jsdom
2
+
3
+ import { describe, it, expect, vi, afterEach } from "vitest";
4
+ import {
5
+ act,
6
+ cleanup,
7
+ render,
8
+ renderHook,
9
+ screen,
10
+ waitFor,
11
+ } from "@testing-library/react";
12
+ import type { AssistantRuntime } from "@assistant-ui/core";
13
+ import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
14
+ import { useAuiState } from "@assistant-ui/store";
15
+ import type { HttpAgent } from "@ag-ui/client";
16
+ import { useAgUiRuntime } from "./useAgUiRuntime";
17
+
18
+ const gatedAgent = () => {
19
+ let release!: () => void;
20
+ const gate = new Promise<void>((resolve) => {
21
+ release = resolve;
22
+ });
23
+ const runAgent = vi.fn(
24
+ async (_input: unknown, subscriber: { onRunFinalized?: () => void }) => {
25
+ if (runAgent.mock.calls.length === 1) await gate;
26
+ subscriber.onRunFinalized?.();
27
+ },
28
+ );
29
+ return {
30
+ agent: { runAgent } as unknown as HttpAgent,
31
+ runAgent,
32
+ release: () => release(),
33
+ };
34
+ };
35
+
36
+ // Asserts through a subscribed consumer rather than the runtime object, which
37
+ // stays correct on read even when the store identity never moves.
38
+ const QueuedPrompts = () => {
39
+ const prompts = useAuiState((s) =>
40
+ s.composer.queue.map((item) => item.prompt).join(","),
41
+ );
42
+ return <div data-testid="queued">{prompts}</div>;
43
+ };
44
+
45
+ const mount = (runtime: AssistantRuntime) =>
46
+ render(
47
+ <AssistantRuntimeProvider runtime={runtime}>
48
+ <QueuedPrompts />
49
+ </AssistantRuntimeProvider>,
50
+ );
51
+
52
+ afterEach(() => {
53
+ cleanup();
54
+ });
55
+
56
+ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
57
+ it("buffers a send during a run and flushes it once the run settles", async () => {
58
+ const { agent, runAgent, release } = gatedAgent();
59
+
60
+ const { result } = renderHook(() =>
61
+ useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
62
+ );
63
+ mount(result.current);
64
+
65
+ await act(async () => {
66
+ await result.current.thread.append({
67
+ role: "user",
68
+ content: [{ type: "text", text: "first" }],
69
+ });
70
+ });
71
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
72
+ expect(result.current.thread.getState().capabilities.queue).toBe(true);
73
+
74
+ await act(async () => {
75
+ await result.current.thread.append({
76
+ role: "user",
77
+ content: [{ type: "text", text: "second" }],
78
+ parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
79
+ });
80
+ });
81
+
82
+ // buffered rather than starting a second run, and visible to subscribers
83
+ expect(runAgent).toHaveBeenCalledTimes(1);
84
+ await waitFor(() =>
85
+ expect(screen.getByTestId("queued").textContent).toBe("second"),
86
+ );
87
+
88
+ await act(async () => {
89
+ release();
90
+ });
91
+
92
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
93
+ await waitFor(() =>
94
+ expect(screen.getByTestId("queued").textContent).toBe(""),
95
+ );
96
+ });
97
+
98
+ it("keeps accepting sends after a queued run fails", async () => {
99
+ const runAgent = vi.fn(
100
+ async (_input: unknown, subscriber: { onRunFinalized?: () => void }) => {
101
+ if (runAgent.mock.calls.length === 1) throw new Error("boom");
102
+ subscriber.onRunFinalized?.();
103
+ },
104
+ );
105
+ const agent = { runAgent } as unknown as HttpAgent;
106
+
107
+ const { result } = renderHook(() =>
108
+ useAgUiRuntime({
109
+ agent,
110
+ unstable_enableMessageQueue: true,
111
+ onError: () => {},
112
+ }),
113
+ );
114
+ mount(result.current);
115
+
116
+ await act(async () => {
117
+ await result.current.thread.append({
118
+ role: "user",
119
+ content: [{ type: "text", text: "first" }],
120
+ });
121
+ });
122
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
123
+
124
+ // the failed run must not leave the queue permanently busy
125
+ await act(async () => {
126
+ await result.current.thread.append({
127
+ role: "user",
128
+ content: [{ type: "text", text: "second" }],
129
+ parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
130
+ });
131
+ });
132
+
133
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
134
+ await waitFor(() =>
135
+ expect(screen.getByTestId("queued").textContent).toBe(""),
136
+ );
137
+ });
138
+
139
+ it("advances one item at a time when a queued dispatch fails", async () => {
140
+ const gate = (() => {
141
+ let release!: () => void;
142
+ const promise = new Promise<void>((r) => {
143
+ release = r;
144
+ });
145
+ return { promise, release: () => release() };
146
+ })();
147
+ // every run is held, so a second concurrent dispatch is observable as a
148
+ // call count rather than being hidden by sequential draining
149
+ const held: Array<() => void> = [];
150
+ const runAgent = vi.fn(async () => {
151
+ if (runAgent.mock.calls.length === 1) {
152
+ await gate.promise;
153
+ throw new Error("boom");
154
+ }
155
+ await new Promise<void>((resolve) => held.push(resolve));
156
+ });
157
+ const agent = { runAgent } as unknown as HttpAgent;
158
+
159
+ const { result } = renderHook(() =>
160
+ useAgUiRuntime({
161
+ agent,
162
+ unstable_enableMessageQueue: true,
163
+ onError: () => {},
164
+ }),
165
+ );
166
+ mount(result.current);
167
+
168
+ const send = async (text: string) => {
169
+ await act(async () => {
170
+ await result.current.thread.append({
171
+ role: "user",
172
+ content: [{ type: "text", text }],
173
+ parentId:
174
+ result.current.thread.getState().messages.at(-1)?.id ?? null,
175
+ });
176
+ });
177
+ };
178
+
179
+ await send("first");
180
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
181
+
182
+ // two behind the gated run, so a double release would drain both
183
+ await send("second");
184
+ await send("third");
185
+ await waitFor(() =>
186
+ expect(screen.getByTestId("queued").textContent).toBe("second,third"),
187
+ );
188
+
189
+ await act(async () => {
190
+ gate.release();
191
+ });
192
+
193
+ // exactly one item advances; a double release would have drained both
194
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
195
+ await waitFor(() =>
196
+ expect(screen.getByTestId("queued").textContent).toBe("third"),
197
+ );
198
+ expect(runAgent).toHaveBeenCalledTimes(2);
199
+ });
200
+
201
+ it("holds a queued message when the run settles on an interrupt", async () => {
202
+ const gate = (() => {
203
+ let release!: () => void;
204
+ const promise = new Promise<void>((r) => {
205
+ release = r;
206
+ });
207
+ return { promise, release: () => release() };
208
+ })();
209
+ const runAgent = vi.fn(
210
+ async (
211
+ _input: unknown,
212
+ subscriber: {
213
+ onRunFinishedEvent?: (e: unknown) => void;
214
+ onRunFinalized?: () => void;
215
+ },
216
+ ) => {
217
+ if (runAgent.mock.calls.length === 1) {
218
+ await gate.promise;
219
+ subscriber.onRunFinishedEvent?.({
220
+ event: {
221
+ type: "RUN_FINISHED",
222
+ runId: "run",
223
+ outcome: {
224
+ type: "interrupt",
225
+ interrupts: [{ id: "int-1", reason: "tool_call" }],
226
+ },
227
+ },
228
+ });
229
+ }
230
+ subscriber.onRunFinalized?.();
231
+ },
232
+ );
233
+ const agent = { runAgent } as unknown as HttpAgent;
234
+
235
+ const { result } = renderHook(() =>
236
+ useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
237
+ );
238
+ mount(result.current);
239
+
240
+ await act(async () => {
241
+ await result.current.thread.append({
242
+ role: "user",
243
+ content: [{ type: "text", text: "first" }],
244
+ });
245
+ });
246
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
247
+
248
+ await act(async () => {
249
+ await result.current.thread.append({
250
+ role: "user",
251
+ content: [{ type: "text", text: "second" }],
252
+ parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
253
+ });
254
+ });
255
+ await waitFor(() =>
256
+ expect(screen.getByTestId("queued").textContent).toBe("second"),
257
+ );
258
+
259
+ await act(async () => {
260
+ gate.release();
261
+ });
262
+
263
+ // append refuses a new run until the interrupt is answered, so the queued
264
+ // message must stay visible rather than being dropped into that refusal
265
+ await waitFor(() =>
266
+ expect(
267
+ result.current.thread.getState().messages.at(-1)?.status?.type,
268
+ ).toBe("requires-action"),
269
+ );
270
+ expect(runAgent).toHaveBeenCalledTimes(1);
271
+ expect(screen.getByTestId("queued").textContent).toBe("second");
272
+ });
273
+
274
+ it("does not wedge on an append that starts no run", async () => {
275
+ const { agent, runAgent } = gatedAgent();
276
+
277
+ const { result } = renderHook(() =>
278
+ useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
279
+ );
280
+ mount(result.current);
281
+
282
+ // startRun false never reaches the agent, so it produces no falling edge
283
+ await act(async () => {
284
+ await result.current.thread.append({
285
+ role: "user",
286
+ content: [{ type: "text", text: "staged" }],
287
+ startRun: false,
288
+ });
289
+ });
290
+ expect(runAgent).toHaveBeenCalledTimes(0);
291
+
292
+ await act(async () => {
293
+ await result.current.thread.append({
294
+ role: "user",
295
+ content: [{ type: "text", text: "real" }],
296
+ parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
297
+ });
298
+ });
299
+
300
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
301
+ await waitFor(() =>
302
+ expect(screen.getByTestId("queued").textContent).toBe(""),
303
+ );
304
+ });
305
+
306
+ it("clears queued items when the flag is turned off", async () => {
307
+ const { agent, runAgent } = gatedAgent();
308
+
309
+ const { result, rerender } = renderHook(
310
+ ({ enabled }: { enabled: boolean }) =>
311
+ useAgUiRuntime({ agent, unstable_enableMessageQueue: enabled }),
312
+ { initialProps: { enabled: true } },
313
+ );
314
+ mount(result.current);
315
+
316
+ await act(async () => {
317
+ await result.current.thread.append({
318
+ role: "user",
319
+ content: [{ type: "text", text: "first" }],
320
+ });
321
+ });
322
+ await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
323
+
324
+ await act(async () => {
325
+ await result.current.thread.append({
326
+ role: "user",
327
+ content: [{ type: "text", text: "second" }],
328
+ parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
329
+ });
330
+ });
331
+ await waitFor(() =>
332
+ expect(screen.getByTestId("queued").textContent).toBe("second"),
333
+ );
334
+
335
+ await act(async () => {
336
+ rerender({ enabled: false });
337
+ });
338
+
339
+ expect(result.current.thread.getState().capabilities.queue).toBe(false);
340
+ await waitFor(() =>
341
+ expect(screen.getByTestId("queued").textContent).toBe(""),
342
+ );
343
+ });
344
+
345
+ it("leaves the queue capability off when the flag is not set", () => {
346
+ const { agent } = gatedAgent();
347
+
348
+ const { result } = renderHook(() => useAgUiRuntime({ agent }));
349
+
350
+ expect(result.current.thread.getState().capabilities.queue).toBe(false);
351
+ });
352
+ });
@@ -1,12 +1,23 @@
1
1
  "use client";
2
2
 
3
- import { useCallback, useEffect, useMemo, useRef, useState } from "react";
3
+ import {
4
+ useCallback,
5
+ useEffect,
6
+ useMemo,
7
+ useRef,
8
+ useState,
9
+ useSyncExternalStore,
10
+ } from "react";
4
11
  import {
5
12
  useExternalStoreRuntime,
6
13
  useExternalStoreSharedOptions,
7
14
  useRuntimeAdapters,
8
15
  } from "@assistant-ui/core/react";
9
- import type { ToolExecutionStatus } from "@assistant-ui/core";
16
+ import { createMessageQueue } from "@assistant-ui/core";
17
+ import type {
18
+ MessageQueueController,
19
+ ToolExecutionStatus,
20
+ } from "@assistant-ui/core";
10
21
  import type {
11
22
  AssistantRuntime,
12
23
  AppendMessage,
@@ -22,8 +33,11 @@ import type {
22
33
  } from "./runtime/types";
23
34
  import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
24
35
  import { agUiExtras } from "./agUiExtras";
36
+ import type { QueueItemState } from "@assistant-ui/core/store";
25
37
 
26
38
  const EMPTY_INTERRUPTS: readonly AgUiInterrupt[] = [];
39
+ const EMPTY_QUEUE_ITEMS: readonly QueueItemState[] = Object.freeze([]);
40
+ const subscribeNoop = () => () => {};
27
41
 
28
42
  export type AgUiAssistantRuntime = AssistantRuntime & {
29
43
  /**
@@ -84,6 +98,77 @@ export function useAgUiRuntime(
84
98
  const hasExecutingTools = Object.values(toolStatuses).some(
85
99
  (s) => s?.type === "executing",
86
100
  );
101
+ const isRunning = core.isRunning() || hasExecutingTools;
102
+
103
+ // The driver sends through the agent core rather than the runtime, because
104
+ // the runtime's append routes every tail append back into this queue.
105
+ const queueRef = useRef<MessageQueueController | null>(null);
106
+ // Counts rendered busy edges, whether a run start or an interrupt arriving.
107
+ // A dispatch whose count never moves was never observed as busy, so no
108
+ // falling edge is coming to release the queue.
109
+ const busyEdgesRef = useRef(0);
110
+ if (options.unstable_enableMessageQueue && !queueRef.current) {
111
+ queueRef.current = createMessageQueue({
112
+ run: (message) => {
113
+ const controller = queueRef.current;
114
+ const edgesAtDispatch = busyEdgesRef.current;
115
+ // The queue drops the item before dispatching and stays busy until an
116
+ // idle edge. An append observed as busy is released by that falling
117
+ // edge, and releasing again here would advance the queue twice; one
118
+ // that never runs (a pre-run rejection, or startRun false) has to be
119
+ // released here or every later send buffers forever. An interrupt is
120
+ // rechecked because a run that resolves without yielding to React
121
+ // reaches here before the effect below has seen either edge.
122
+ const releaseIfNoRun = () => {
123
+ if (
124
+ queueRef.current !== controller ||
125
+ busyEdgesRef.current !== edgesAtDispatch ||
126
+ core.getPendingInterrupts() !== null
127
+ )
128
+ return;
129
+ controller?.notifyIdle();
130
+ };
131
+ void core.append(message).then(releaseIfNoRun, (e: unknown) => {
132
+ releaseIfNoRun();
133
+ logger.error?.("[agui] queued message failed to send", e);
134
+ });
135
+ },
136
+ });
137
+ } else if (!options.unstable_enableMessageQueue && queueRef.current) {
138
+ queueRef.current.clear();
139
+ queueRef.current = null;
140
+ }
141
+ const queueController = options.unstable_enableMessageQueue
142
+ ? queueRef.current
143
+ : null;
144
+
145
+ // Feeds the store memo below: the runtime core skips an adapter whose
146
+ // identity is unchanged, so queue items have to move the store reference or
147
+ // subscribers never see composer.queue change.
148
+ const queueItems = useSyncExternalStore(
149
+ queueController?.subscribe ?? subscribeNoop,
150
+ () => queueController?.adapter.items ?? EMPTY_QUEUE_ITEMS,
151
+ () => EMPTY_QUEUE_ITEMS,
152
+ );
153
+ const steerQueueItems = useSyncExternalStore(
154
+ queueController?.subscribe ?? subscribeNoop,
155
+ () => queueController?.adapter.steerItems ?? EMPTY_QUEUE_ITEMS,
156
+ () => EMPTY_QUEUE_ITEMS,
157
+ );
158
+
159
+ // Holds the queue while client-side tools are still executing, and while an
160
+ // interrupt is unanswered: `append` refuses a new run until it is resolved,
161
+ // and the queue drops an item before dispatching it, so draining into that
162
+ // refusal would lose the message instead of keeping it visible.
163
+ const queueBusy = isRunning || core.getPendingInterrupts() !== null;
164
+ useEffect(() => {
165
+ if (queueBusy) {
166
+ busyEdgesRef.current++;
167
+ queueController?.notifyBusy();
168
+ } else {
169
+ queueController?.notifyIdle();
170
+ }
171
+ }, [queueBusy, queueController]);
87
172
 
88
173
  const threadList = useMemo(() => {
89
174
  if (!threadListAdapter) return undefined;
@@ -143,7 +228,7 @@ export function useAgUiRuntime(
143
228
  isLoading: core.isLoading,
144
229
  messageRepository: core.getMessageRepository(),
145
230
  state: core.getState(),
146
- isRunning: core.isRunning() || hasExecutingTools,
231
+ isRunning,
147
232
  extras: agUiExtras.provide({
148
233
  interrupts:
149
234
  core.getPendingInterrupts()?.interrupts ?? EMPTY_INTERRUPTS,
@@ -157,10 +242,16 @@ export function useAgUiRuntime(
157
242
  unstable_enableToolInvocations: true,
158
243
  setToolStatuses,
159
244
  onNew: (message: AppendMessage) => core.append(message),
160
- onEdit: (message: AppendMessage) => core.edit(message),
161
- onReload: (parentId: string | null, config: { runConfig?: any }) =>
162
- core.reload(parentId, config),
245
+ onEdit: (message: AppendMessage) => {
246
+ queueController?.clear();
247
+ return core.edit(message);
248
+ },
249
+ onReload: (parentId: string | null, config: { runConfig?: any }) => {
250
+ queueController?.clear();
251
+ return core.reload(parentId, config);
252
+ },
163
253
  onCancel: async () => {
254
+ queueController?.clear();
164
255
  core.cancel();
165
256
  },
166
257
  onAddToolResult: (options) => core.addToolResult(options),
@@ -172,11 +263,21 @@ export function useAgUiRuntime(
172
263
  onLoadExternalState: (state: ReadonlyJSONValue) =>
173
264
  core.loadExternalState(state),
174
265
  adapters: adapterAdapters,
266
+ ...(queueController && { queue: queueController.adapter }),
175
267
  } satisfies ExternalStoreAdapter<ThreadMessage>;
176
268
  },
177
269
  // _version is intentionally included to trigger re-computation when core state changes via notifyUpdate
178
270
  // toolInvocations intentionally excluded: abort/resume use refs internally and work with stale captures
179
- [adapterAdapters, core, _version, hasExecutingTools, shared],
271
+ [
272
+ adapterAdapters,
273
+ core,
274
+ _version,
275
+ isRunning,
276
+ queueController,
277
+ queueItems,
278
+ steerQueueItems,
279
+ shared,
280
+ ],
180
281
  );
181
282
 
182
283
  const baseRuntime = useExternalStoreRuntime(store);