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

package/src/runtime/adapter/subscriber.ts

@@ -22,6 +22,7 @@ type Subscriber = {
   onReasoningMessageStartEvent?: (payload: { event: unknown }) => void;
   onReasoningMessageContentEvent?: (payload: { event: unknown }) => void;
   onReasoningMessageEndEvent?: (payload: { event: unknown }) => void;
+  onReasoningEncryptedValueEvent?: (payload: { event: unknown }) => void;
   onToolCallStartEvent?: (payload: { event: unknown }) => void;
   onToolCallArgsEvent?: (payload: { event: unknown }) => void;
   onToolCallEndEvent?: (payload: { event: unknown }) => void;
@@ -124,6 +125,8 @@ export const createAgUiSubscriber = (
       dispatchIfValid(event, "REASONING_MESSAGE_CONTENT"),
     onReasoningMessageEndEvent: ({ event }) =>
       dispatchIfValid(event, "REASONING_MESSAGE_END"),
+    onReasoningEncryptedValueEvent: ({ event }) =>
+      dispatchIfValid(event, "REASONING_ENCRYPTED_VALUE"),
     onToolCallStartEvent: ({ event }) =>
       dispatchIfValid(event, "TOOL_CALL_START"),
     onToolCallArgsEvent: ({ event }) =>

package/src/runtime/event-parser.ts

@@ -171,6 +171,19 @@ export const parseAgUiEvent = (
         { type: "REASONING_MESSAGE_END" as const },
         { messageId: getString("messageId") },
       );
+    case "REASONING_ENCRYPTED_VALUE": {
+      const entityId = getString("entityId");
+      const encryptedValue = getString("encryptedValue");
+      const subtype = getString("subtype");
+      if (!entityId || !encryptedValue) return null;
+      if (subtype !== "message" && subtype !== "tool-call") return null;
+      return {
+        type: "REASONING_ENCRYPTED_VALUE" as const,
+        subtype,
+        entityId,
+        encryptedValue,
+      };
+    }
     case "REASONING_END":
       return withOptional(
         { type: "REASONING_END" as const },

package/src/runtime/types.ts

@@ -149,6 +149,12 @@ export type AgUiEvent =
   | { type: "REASONING_MESSAGE_START"; messageId?: string }
   | { type: "REASONING_MESSAGE_CONTENT"; messageId?: string; delta: string }
   | { type: "REASONING_MESSAGE_END"; messageId?: string }
+  | {
+      type: "REASONING_ENCRYPTED_VALUE";
+      subtype: "message" | "tool-call";
+      entityId: string;
+      encryptedValue: string;
+    }
   | { type: "REASONING_END"; messageId?: string }
   | {
       type: "TOOL_CALL_START";

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,7 +251,7 @@ export function useAgUiRuntime(
           return core.reload(parentId, config);
         },
         onCancel: async () => {
-          queueController?.clear();
+          queueController?.notifyCancelled();
           core.cancel();
         },
         onAddToolResult: (options) => core.addToolResult(options),