@assistant-ui/react-ag-ui 0.0.53 → 0.0.56

package/src/useAgUiRuntime.queue.test.tsx

@@ -27,7 +27,7 @@ const gatedAgent = () => {
     },
   );
   return {
-    agent: { runAgent } as unknown as HttpAgent,
+    agent: { runAgent, abortRun: vi.fn() } as unknown as HttpAgent,
     runAgent,
     release: () => release(),
   };
@@ -95,6 +95,73 @@ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
     );
   });
 
+  it("preserves queued sends when the active run is cancelled", async () => {
+    const { agent, runAgent, release } = gatedAgent();
+
+    const { result } = renderHook(() =>
+      useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
+    );
+    mount(result.current);
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "second" }],
+        parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("second"),
+    );
+
+    act(() => {
+      result.current.thread.cancelRun();
+    });
+    await act(async () => {
+      release();
+    });
+
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("second"),
+    );
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "third" }],
+        parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(3));
+    expect(runAgent.mock.calls[1]?.[0]).toEqual(
+      expect.objectContaining({
+        messages: expect.arrayContaining([
+          expect.objectContaining({ role: "user", content: "second" }),
+        ]),
+      }),
+    );
+    expect(runAgent.mock.calls[2]?.[0]).toEqual(
+      expect.objectContaining({
+        messages: expect.arrayContaining([
+          expect.objectContaining({ role: "user", content: "third" }),
+        ]),
+      }),
+    );
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe(""),
+    );
+  });
+
   it("keeps accepting sends after a queued run fails", async () => {
     const runAgent = vi.fn(
       async (_input: unknown, subscriber: { onRunFinalized?: () => void }) => {

package/src/useAgUiRuntime.ts

@@ -251,10 +251,14 @@ export function useAgUiRuntime(
           return core.reload(parentId, config);
         },
         onCancel: async () => {
-          queueController?.clear();
+          queueController?.notifyCancelled();
           core.cancel();
         },
         onAddToolResult: (options) => core.addToolResult(options),
+        onRespondToToolApproval: (options) =>
+          core.respondToToolApproval(options).catch((error: unknown) => {
+            core.reportError(error);
+          }),
         onResume: (config) => core.resume(config),
         setMessages: (messages: readonly ThreadMessage[]) =>
           core.applyExternalMessages(messages),