@assistant-ui/react-ag-ui 0.0.49 → 0.0.51

package/src/useAgUiRuntime.queue.test.tsx

@@ -0,0 +1,352 @@
+// @vitest-environment jsdom
+
+import { describe, it, expect, vi, afterEach } from "vitest";
+import {
+  act,
+  cleanup,
+  render,
+  renderHook,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import type { AssistantRuntime } from "@assistant-ui/core";
+import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
+import { useAuiState } from "@assistant-ui/store";
+import type { HttpAgent } from "@ag-ui/client";
+import { useAgUiRuntime } from "./useAgUiRuntime";
+
+const gatedAgent = () => {
+  let release!: () => void;
+  const gate = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  const runAgent = vi.fn(
+    async (_input: unknown, subscriber: { onRunFinalized?: () => void }) => {
+      if (runAgent.mock.calls.length === 1) await gate;
+      subscriber.onRunFinalized?.();
+    },
+  );
+  return {
+    agent: { runAgent } as unknown as HttpAgent,
+    runAgent,
+    release: () => release(),
+  };
+};
+
+// Asserts through a subscribed consumer rather than the runtime object, which
+// stays correct on read even when the store identity never moves.
+const QueuedPrompts = () => {
+  const prompts = useAuiState((s) =>
+    s.composer.queue.map((item) => item.prompt).join(","),
+  );
+  return <div data-testid="queued">{prompts}</div>;
+};
+
+const mount = (runtime: AssistantRuntime) =>
+  render(
+    <AssistantRuntimeProvider runtime={runtime}>
+      <QueuedPrompts />
+    </AssistantRuntimeProvider>,
+  );
+
+afterEach(() => {
+  cleanup();
+});
+
+describe("useAgUiRuntime unstable_enableMessageQueue", () => {
+  it("buffers a send during a run and flushes it once the run settles", 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));
+    expect(result.current.thread.getState().capabilities.queue).toBe(true);
+
+    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,
+      });
+    });
+
+    // buffered rather than starting a second run, and visible to subscribers
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("second"),
+    );
+
+    await act(async () => {
+      release();
+    });
+
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+    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 }) => {
+        if (runAgent.mock.calls.length === 1) throw new Error("boom");
+        subscriber.onRunFinalized?.();
+      },
+    );
+    const agent = { runAgent } as unknown as HttpAgent;
+
+    const { result } = renderHook(() =>
+      useAgUiRuntime({
+        agent,
+        unstable_enableMessageQueue: true,
+        onError: () => {},
+      }),
+    );
+    mount(result.current);
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+
+    // the failed run must not leave the queue permanently busy
+    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(runAgent).toHaveBeenCalledTimes(2));
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe(""),
+    );
+  });
+
+  it("advances one item at a time when a queued dispatch fails", async () => {
+    const gate = (() => {
+      let release!: () => void;
+      const promise = new Promise<void>((r) => {
+        release = r;
+      });
+      return { promise, release: () => release() };
+    })();
+    // every run is held, so a second concurrent dispatch is observable as a
+    // call count rather than being hidden by sequential draining
+    const held: Array<() => void> = [];
+    const runAgent = vi.fn(async () => {
+      if (runAgent.mock.calls.length === 1) {
+        await gate.promise;
+        throw new Error("boom");
+      }
+      await new Promise<void>((resolve) => held.push(resolve));
+    });
+    const agent = { runAgent } as unknown as HttpAgent;
+
+    const { result } = renderHook(() =>
+      useAgUiRuntime({
+        agent,
+        unstable_enableMessageQueue: true,
+        onError: () => {},
+      }),
+    );
+    mount(result.current);
+
+    const send = async (text: string) => {
+      await act(async () => {
+        await result.current.thread.append({
+          role: "user",
+          content: [{ type: "text", text }],
+          parentId:
+            result.current.thread.getState().messages.at(-1)?.id ?? null,
+        });
+      });
+    };
+
+    await send("first");
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+
+    // two behind the gated run, so a double release would drain both
+    await send("second");
+    await send("third");
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("second,third"),
+    );
+
+    await act(async () => {
+      gate.release();
+    });
+
+    // exactly one item advances; a double release would have drained both
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("third"),
+    );
+    expect(runAgent).toHaveBeenCalledTimes(2);
+  });
+
+  it("holds a queued message when the run settles on an interrupt", async () => {
+    const gate = (() => {
+      let release!: () => void;
+      const promise = new Promise<void>((r) => {
+        release = r;
+      });
+      return { promise, release: () => release() };
+    })();
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        subscriber: {
+          onRunFinishedEvent?: (e: unknown) => void;
+          onRunFinalized?: () => void;
+        },
+      ) => {
+        if (runAgent.mock.calls.length === 1) {
+          await gate.promise;
+          subscriber.onRunFinishedEvent?.({
+            event: {
+              type: "RUN_FINISHED",
+              runId: "run",
+              outcome: {
+                type: "interrupt",
+                interrupts: [{ id: "int-1", reason: "tool_call" }],
+              },
+            },
+          });
+        }
+        subscriber.onRunFinalized?.();
+      },
+    );
+    const agent = { runAgent } as unknown as HttpAgent;
+
+    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"),
+    );
+
+    await act(async () => {
+      gate.release();
+    });
+
+    // append refuses a new run until the interrupt is answered, so the queued
+    // message must stay visible rather than being dropped into that refusal
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status?.type,
+      ).toBe("requires-action"),
+    );
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    expect(screen.getByTestId("queued").textContent).toBe("second");
+  });
+
+  it("does not wedge on an append that starts no run", async () => {
+    const { agent, runAgent } = gatedAgent();
+
+    const { result } = renderHook(() =>
+      useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
+    );
+    mount(result.current);
+
+    // startRun false never reaches the agent, so it produces no falling edge
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "staged" }],
+        startRun: false,
+      });
+    });
+    expect(runAgent).toHaveBeenCalledTimes(0);
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "real" }],
+        parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe(""),
+    );
+  });
+
+  it("clears queued items when the flag is turned off", async () => {
+    const { agent, runAgent } = gatedAgent();
+
+    const { result, rerender } = renderHook(
+      ({ enabled }: { enabled: boolean }) =>
+        useAgUiRuntime({ agent, unstable_enableMessageQueue: enabled }),
+      { initialProps: { enabled: 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"),
+    );
+
+    await act(async () => {
+      rerender({ enabled: false });
+    });
+
+    expect(result.current.thread.getState().capabilities.queue).toBe(false);
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe(""),
+    );
+  });
+
+  it("leaves the queue capability off when the flag is not set", () => {
+    const { agent } = gatedAgent();
+
+    const { result } = renderHook(() => useAgUiRuntime({ agent }));
+
+    expect(result.current.thread.getState().capabilities.queue).toBe(false);
+  });
+});

package/src/useAgUiRuntime.ts

@@ -1,12 +1,23 @@
 "use client";
 
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+  useSyncExternalStore,
+} from "react";
 import {
   useExternalStoreRuntime,
   useExternalStoreSharedOptions,
   useRuntimeAdapters,
 } from "@assistant-ui/core/react";
-import type { ToolExecutionStatus } from "@assistant-ui/core";
+import { createMessageQueue } from "@assistant-ui/core";
+import type {
+  MessageQueueController,
+  ToolExecutionStatus,
+} from "@assistant-ui/core";
 import type {
   AssistantRuntime,
   AppendMessage,
@@ -22,8 +33,11 @@ import type {
 } from "./runtime/types";
 import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
 import { agUiExtras } from "./agUiExtras";
+import type { QueueItemState } from "@assistant-ui/core/store";
 
 const EMPTY_INTERRUPTS: readonly AgUiInterrupt[] = [];
+const EMPTY_QUEUE_ITEMS: readonly QueueItemState[] = Object.freeze([]);
+const subscribeNoop = () => () => {};
 
 export type AgUiAssistantRuntime = AssistantRuntime & {
   /**
@@ -84,6 +98,77 @@ export function useAgUiRuntime(
   const hasExecutingTools = Object.values(toolStatuses).some(
     (s) => s?.type === "executing",
   );
+  const isRunning = core.isRunning() || hasExecutingTools;
+
+  // The driver sends through the agent core rather than the runtime, because
+  // the runtime's append routes every tail append back into this queue.
+  const queueRef = useRef<MessageQueueController | null>(null);
+  // Counts rendered busy edges, whether a run start or an interrupt arriving.
+  // A dispatch whose count never moves was never observed as busy, so no
+  // falling edge is coming to release the queue.
+  const busyEdgesRef = useRef(0);
+  if (options.unstable_enableMessageQueue && !queueRef.current) {
+    queueRef.current = createMessageQueue({
+      run: (message) => {
+        const controller = queueRef.current;
+        const edgesAtDispatch = busyEdgesRef.current;
+        // The queue drops the item before dispatching and stays busy until an
+        // idle edge. An append observed as busy is released by that falling
+        // edge, and releasing again here would advance the queue twice; one
+        // that never runs (a pre-run rejection, or startRun false) has to be
+        // released here or every later send buffers forever. An interrupt is
+        // rechecked because a run that resolves without yielding to React
+        // reaches here before the effect below has seen either edge.
+        const releaseIfNoRun = () => {
+          if (
+            queueRef.current !== controller ||
+            busyEdgesRef.current !== edgesAtDispatch ||
+            core.getPendingInterrupts() !== null
+          )
+            return;
+          controller?.notifyIdle();
+        };
+        void core.append(message).then(releaseIfNoRun, (e: unknown) => {
+          releaseIfNoRun();
+          logger.error?.("[agui] queued message failed to send", e);
+        });
+      },
+    });
+  } else if (!options.unstable_enableMessageQueue && queueRef.current) {
+    queueRef.current.clear();
+    queueRef.current = null;
+  }
+  const queueController = options.unstable_enableMessageQueue
+    ? queueRef.current
+    : null;
+
+  // Feeds the store memo below: the runtime core skips an adapter whose
+  // identity is unchanged, so queue items have to move the store reference or
+  // subscribers never see composer.queue change.
+  const queueItems = useSyncExternalStore(
+    queueController?.subscribe ?? subscribeNoop,
+    () => queueController?.adapter.items ?? EMPTY_QUEUE_ITEMS,
+    () => EMPTY_QUEUE_ITEMS,
+  );
+  const steerQueueItems = useSyncExternalStore(
+    queueController?.subscribe ?? subscribeNoop,
+    () => queueController?.adapter.steerItems ?? EMPTY_QUEUE_ITEMS,
+    () => EMPTY_QUEUE_ITEMS,
+  );
+
+  // Holds the queue while client-side tools are still executing, and while an
+  // interrupt is unanswered: `append` refuses a new run until it is resolved,
+  // and the queue drops an item before dispatching it, so draining into that
+  // refusal would lose the message instead of keeping it visible.
+  const queueBusy = isRunning || core.getPendingInterrupts() !== null;
+  useEffect(() => {
+    if (queueBusy) {
+      busyEdgesRef.current++;
+      queueController?.notifyBusy();
+    } else {
+      queueController?.notifyIdle();
+    }
+  }, [queueBusy, queueController]);
 
   const threadList = useMemo(() => {
     if (!threadListAdapter) return undefined;
@@ -143,7 +228,7 @@ export function useAgUiRuntime(
         isLoading: core.isLoading,
         messageRepository: core.getMessageRepository(),
         state: core.getState(),
-        isRunning: core.isRunning() || hasExecutingTools,
+        isRunning,
         extras: agUiExtras.provide({
           interrupts:
             core.getPendingInterrupts()?.interrupts ?? EMPTY_INTERRUPTS,
@@ -157,10 +242,16 @@ export function useAgUiRuntime(
         unstable_enableToolInvocations: true,
         setToolStatuses,
         onNew: (message: AppendMessage) => core.append(message),
-        onEdit: (message: AppendMessage) => core.edit(message),
-        onReload: (parentId: string | null, config: { runConfig?: any }) =>
-          core.reload(parentId, config),
+        onEdit: (message: AppendMessage) => {
+          queueController?.clear();
+          return core.edit(message);
+        },
+        onReload: (parentId: string | null, config: { runConfig?: any }) => {
+          queueController?.clear();
+          return core.reload(parentId, config);
+        },
         onCancel: async () => {
+          queueController?.clear();
           core.cancel();
         },
         onAddToolResult: (options) => core.addToolResult(options),
@@ -172,11 +263,21 @@ export function useAgUiRuntime(
         onLoadExternalState: (state: ReadonlyJSONValue) =>
           core.loadExternalState(state),
         adapters: adapterAdapters,
+        ...(queueController && { queue: queueController.adapter }),
       } satisfies ExternalStoreAdapter<ThreadMessage>;
     },
     // _version is intentionally included to trigger re-computation when core state changes via notifyUpdate
     // toolInvocations intentionally excluded: abort/resume use refs internally and work with stale captures
-    [adapterAdapters, core, _version, hasExecutingTools, shared],
+    [
+      adapterAdapters,
+      core,
+      _version,
+      isRunning,
+      queueController,
+      queueItems,
+      steerQueueItems,
+      shared,
+    ],
   );
 
   const baseRuntime = useExternalStoreRuntime(store);