@assistant-ui/react-ag-ui 0.0.64 → 0.0.65

package/src/runtime/types.ts

@@ -180,16 +180,22 @@ export type AgUiEvent =
       delta: string;
       subagentRunId?: string;
     }
-  | { type: "THINKING_START"; title?: string }
-  | { type: "THINKING_TEXT_MESSAGE_START" }
+  | { type: "THINKING_START"; title?: string; timestamp?: number }
+  | { type: "THINKING_TEXT_MESSAGE_START"; timestamp?: number }
   | { type: "THINKING_TEXT_MESSAGE_CONTENT"; delta: string }
-  | { type: "THINKING_TEXT_MESSAGE_END" }
-  | { type: "THINKING_END" }
-  | { type: "REASONING_START"; messageId?: string; subagentRunId?: string }
+  | { type: "THINKING_TEXT_MESSAGE_END"; timestamp?: number }
+  | { type: "THINKING_END"; timestamp?: number }
+  | {
+      type: "REASONING_START";
+      messageId?: string;
+      subagentRunId?: string;
+      timestamp?: number;
+    }
   | {
       type: "REASONING_MESSAGE_START";
       messageId?: string;
       subagentRunId?: string;
+      timestamp?: number;
     }
   | {
       type: "REASONING_MESSAGE_CONTENT";
@@ -201,6 +207,7 @@ export type AgUiEvent =
       type: "REASONING_MESSAGE_END";
       messageId?: string;
       subagentRunId?: string;
+      timestamp?: number;
     }
   | {
       type: "REASONING_ENCRYPTED_VALUE";
@@ -209,7 +216,12 @@ export type AgUiEvent =
       encryptedValue: string;
       subagentRunId?: string;
     }
-  | { type: "REASONING_END"; messageId?: string; subagentRunId?: string }
+  | {
+      type: "REASONING_END";
+      messageId?: string;
+      subagentRunId?: string;
+      timestamp?: number;
+    }
   | {
       type: "TOOL_CALL_START";
       toolCallId: string;

package/src/useAgUiRuntime.queue.test.tsx

@@ -96,11 +96,67 @@ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
     );
   });
 
+  it("holds queued sends while sending is disabled", async () => {
+    const { agent, runAgent, release } = gatedAgent();
+
+    const { result, rerender } = renderHook(
+      ({ isSendDisabled }) =>
+        useAgUiRuntime({
+          agent,
+          unstable_enableMessageQueue: true,
+          isSendDisabled,
+        }),
+      { initialProps: { isSendDisabled: false } },
+    );
+    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"),
+    );
+    rerender({ isSendDisabled: true });
+
+    await act(async () => {
+      release();
+    });
+    await waitFor(() =>
+      expect(result.current.thread.getState().isRunning).toBe(false),
+    );
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    expect(screen.getByTestId("queued").textContent).toBe("second");
+
+    rerender({ isSendDisabled: false });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe(""),
+    );
+  });
+
   it("preserves queued sends when the active run is cancelled", async () => {
     const { agent, runAgent, release } = gatedAgent();
 
-    const { result } = renderHook(() =>
-      useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
+    const { result, rerender } = renderHook(
+      ({ isSendDisabled }) =>
+        useAgUiRuntime({
+          agent,
+          unstable_enableMessageQueue: true,
+          isSendDisabled,
+        }),
+      { initialProps: { isSendDisabled: false } },
     );
     mount(result.current);
 
@@ -135,6 +191,11 @@ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
       expect(screen.getByTestId("queued").textContent).toBe("second"),
     );
 
+    rerender({ isSendDisabled: true });
+    rerender({ isSendDisabled: false });
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    expect(screen.getByTestId("queued").textContent).toBe("second");
+
     await act(async () => {
       await result.current.thread.append({
         role: "user",

package/src/useAgUiRuntime.ts

@@ -175,6 +175,14 @@ export function useAgUiRuntime(
   // 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 (options.isSendDisabled === true) {
+      queueController?.hold();
+    } else {
+      queueController?.release();
+    }
+  }, [options.isSendDisabled, queueController]);
+
   useEffect(() => {
     if (!queueController) return;
     if (