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

package/src/useAgUiRuntime.thread-switch.test.tsx

@@ -3,10 +3,13 @@
 import { useState } from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { act, cleanup, renderHook, waitFor } from "@testing-library/react";
-import type { HttpAgent } from "@ag-ui/client";
-import type { ThreadMessage } from "@assistant-ui/core";
+import { HttpAgent } from "@ag-ui/client";
+import type { ThreadHistoryAdapter, ThreadMessage } from "@assistant-ui/core";
 import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
-import type { UseAgUiThreadListAdapter } from "./runtime/types";
+import type {
+  UseAgUiRuntimeOptions,
+  UseAgUiThreadListAdapter,
+} from "./runtime/types";
 import { useAgUiRuntime } from "./useAgUiRuntime";
 
 type ThreadLoad = Awaited<
@@ -40,17 +43,24 @@ function renderRuntime(
   load: (id: string) => Promise<ThreadLoad>,
   create: () => Promise<void> = async () => {},
   agentOverride?: HttpAgent,
+  options: Pick<
+    UseAgUiRuntimeOptions,
+    "onCancel" | "unstable_enableMessageQueue"
+  > & { history?: ThreadHistoryAdapter } = {},
 ) {
   const agent = {
     runAgent: vi.fn(),
     abortRun: vi.fn(),
   } as unknown as HttpAgent;
   const runtimeAgent = agentOverride ?? agent;
+  const { history, ...runtimeOptions } = options;
   return renderHook(() => {
     const [threadId, setThreadId] = useState("initial");
     return useAgUiRuntime({
       agent: runtimeAgent,
+      ...runtimeOptions,
       adapters: {
+        ...(history && { history }),
         threadList: {
           threadId,
           onSwitchToThread: (id) => {
@@ -73,6 +83,580 @@ afterEach(() => {
 });
 
 describe("useAgUiRuntime thread switching", () => {
+  it.each(["existing", "new"] as const)(
+    "persists a cancelled partial reply before switching to the %s thread",
+    async (destination) => {
+      const signals: AbortSignal[] = [];
+      const onAbort = vi.fn();
+      const appendHistory = vi.fn<ThreadHistoryAdapter["append"]>(
+        async () => {},
+      );
+      const history: ThreadHistoryAdapter = {
+        load: vi.fn(async () => ({ headId: null, messages: [] })),
+        append: appendHistory,
+      };
+      let subscriber!: Subscriber;
+      const runAgent = vi.fn(
+        async (
+          _input: unknown,
+          next: Subscriber,
+          options: { signal: AbortSignal },
+        ) => {
+          subscriber = next;
+          signals.push(options.signal);
+          await new Promise<void>((resolve) =>
+            options.signal.addEventListener(
+              "abort",
+              () => {
+                onAbort();
+                resolve();
+              },
+              { once: true },
+            ),
+          );
+        },
+      );
+      const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+      let adapterCalled = false;
+      let cancelledAtAdapter = false;
+      let partialAtAdapter = false;
+      const inspectAtAdapter = () => {
+        adapterCalled = true;
+        cancelledAtAdapter = signals[0]?.aborted ?? false;
+        partialAtAdapter = appendHistory.mock.calls.some(
+          ([entry]) =>
+            entry.message.role === "assistant" &&
+            entry.message.status?.type === "incomplete" &&
+            entry.message.content.some(
+              (part) => part.type === "text" && part.text === "partial",
+            ),
+        );
+      };
+      const { result } = renderRuntime(
+        async () => {
+          inspectAtAdapter();
+          return { messages: [message("loaded")] };
+        },
+        async () => {
+          inspectAtAdapter();
+        },
+        agent,
+        { history },
+      );
+
+      await waitFor(() =>
+        expect(result.current.thread.getState().isLoading).toBe(false),
+      );
+
+      act(() => {
+        void result.current.thread.append("old prompt");
+      });
+      await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+      act(() => {
+        subscriber.onTextMessageStartEvent?.({
+          event: { type: "TEXT_MESSAGE_START", messageId: "old-output" },
+        });
+        subscriber.onTextMessageContentEvent?.({
+          event: {
+            type: "TEXT_MESSAGE_CONTENT",
+            messageId: "old-output",
+            delta: "partial",
+          },
+        });
+      });
+      expect(result.current.thread.getState().messages.length).toBeGreaterThan(
+        0,
+      );
+
+      let switching!: Promise<void>;
+      act(() => {
+        if (destination === "existing") {
+          switching = result.current.threads.switchToThread("thread-a");
+        } else {
+          switching = result.current.threads.switchToNewThread();
+        }
+        expect(adapterCalled).toBe(true);
+        expect(partialAtAdapter).toBe(true);
+      });
+      await act(async () => {
+        await switching;
+      });
+
+      expect(cancelledAtAdapter).toBe(true);
+      expect(signals[0]?.aborted).toBe(true);
+      expect(onAbort).toHaveBeenCalledOnce();
+      expect(partialAtAdapter).toBe(true);
+      expect(
+        result.current.thread.getState().messages.map((m) => m.id),
+      ).toEqual(destination === "existing" ? ["loaded"] : []);
+      expect(runAgent).toHaveBeenCalledOnce();
+    },
+  );
+
+  it.each(["existing", "new"] as const)(
+    "aborts an HttpAgent's in-flight request when switching to the %s thread",
+    async (destination) => {
+      const requestSignals: AbortSignal[] = [];
+      const fetch = vi.fn((_url: string, init: RequestInit) => {
+        const signal = init.signal!;
+        requestSignals.push(signal);
+        return new Promise<Response>((_resolve, reject) =>
+          signal.addEventListener("abort", () => reject(signal.reason), {
+            once: true,
+          }),
+        );
+      });
+      const agent = new HttpAgent({ url: "https://agent.test/run", fetch });
+      const { result } = renderRuntime(
+        async () => ({ messages: [message("loaded")] }),
+        async () => {},
+        agent,
+      );
+
+      act(() => {
+        void result.current.thread.append("old prompt");
+      });
+      await waitFor(() => expect(requestSignals).toHaveLength(1));
+
+      await act(async () => {
+        await (destination === "existing"
+          ? result.current.threads.switchToThread("thread-a")
+          : result.current.threads.switchToNewThread());
+      });
+
+      expect(requestSignals[0]?.aborted).toBe(true);
+      expect(result.current.thread.getState().isRunning).toBe(false);
+      expect(fetch).toHaveBeenCalledOnce();
+    },
+  );
+
+  it("still switches when the agent's abortRun throws", async () => {
+    const signals: AbortSignal[] = [];
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _next: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        signals.push(options.signal);
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = {
+      runAgent,
+      abortRun: vi.fn(() => {
+        throw new Error("abort failed");
+      }),
+    } as unknown as HttpAgent;
+    const { result } = renderRuntime(
+      async () => ({ messages: [message("loaded")] }),
+      async () => {},
+      agent,
+    );
+
+    act(() => {
+      void result.current.thread.append("old prompt");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+
+    expect(signals[0]?.aborted).toBe(true);
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "loaded",
+    ]);
+  });
+
+  it("discards sends queued behind the active run when switching threads", async () => {
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    const { result } = renderRuntime(
+      async () => ({ messages: [message("loaded")] }),
+      undefined,
+      agent,
+      { unstable_enableMessageQueue: true },
+    );
+
+    act(() => {
+      void result.current.thread.append("active");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "queued" }],
+        parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+    expect(runAgent).toHaveBeenCalledOnce();
+
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "loaded",
+    ]);
+    expect(runAgent).toHaveBeenCalledOnce();
+  });
+
+  it("drops sends from onCancel and the caller during the switch", async () => {
+    const load = deferred<ThreadLoad>();
+    const signals: AbortSignal[] = [];
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        signals.push(options.signal);
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    let result!: ReturnType<typeof renderRuntime>["result"];
+    let cancelled = false;
+    ({ result } = renderRuntime(() => load.promise, undefined, agent, {
+      onCancel: () => {
+        if (cancelled) return;
+        cancelled = true;
+        void result.current.thread.append("started by onCancel");
+      },
+    }));
+
+    act(() => {
+      void result.current.thread.append("old prompt");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+
+    let switching!: Promise<void>;
+    act(() => {
+      switching = result.current.threads.switchToThread("thread-a");
+    });
+    act(() => {
+      void result.current.thread.append("during switch");
+    });
+    await act(async () => {
+      load.resolve({ messages: [message("loaded")] });
+      await switching;
+    });
+
+    expect(runAgent).toHaveBeenCalledOnce();
+    expect(signals.every((signal) => signal.aborted)).toBe(true);
+    expect(result.current.thread.getState().isRunning).toBe(false);
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "loaded",
+    ]);
+  });
+
+  it("drops every send started by onCancel while switching", async () => {
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    let result!: ReturnType<typeof renderRuntime>["result"];
+    const onCancel = vi.fn(() => {
+      void result.current.thread.append("from onCancel");
+    });
+    ({ result } = renderRuntime(
+      async () => ({ messages: [message("loaded")] }),
+      undefined,
+      agent,
+      { onCancel },
+    ));
+
+    act(() => {
+      void result.current.thread.append("old prompt");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+
+    expect(onCancel).toHaveBeenCalledOnce();
+    expect(runAgent).toHaveBeenCalledOnce();
+    expect(result.current.thread.getState().isRunning).toBe(false);
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "loaded",
+    ]);
+  });
+
+  it("drops a reload started by onCancel during the final reset", async () => {
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    let result!: ReturnType<typeof renderRuntime>["result"];
+    const onCancel = vi.fn(() => {
+      void result.current.thread.startRun({ parentId: null });
+    });
+    ({ result } = renderRuntime(
+      async () => ({ messages: [message("loaded")] }),
+      undefined,
+      agent,
+      { onCancel },
+    ));
+
+    act(() => {
+      void result.current.thread.append("old prompt");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+
+    expect(onCancel).toHaveBeenCalledOnce();
+    expect(runAgent).toHaveBeenCalledOnce();
+    expect(result.current.thread.getState().isRunning).toBe(false);
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "loaded",
+    ]);
+  });
+
+  it("lets a switch started by onCancel supersede the cancelled switch", async () => {
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    let result!: ReturnType<typeof renderRuntime>["result"];
+    let switchB!: Promise<void>;
+    ({ result } = renderRuntime(
+      async (id) => ({ messages: [message(id)] }),
+      undefined,
+      agent,
+      {
+        onCancel: () => {
+          switchB = result.current.threads.switchToThread("thread-b");
+        },
+      },
+    ));
+
+    act(() => {
+      void result.current.thread.append("old prompt");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    await act(async () => {
+      const switchA = result.current.threads.switchToThread("thread-a");
+      await Promise.all([switchA, switchB]);
+    });
+
+    expect(result.current.threads.getState().mainThreadId).toBe("thread-b");
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "thread-b",
+    ]);
+  });
+
+  it("resumes an in-flight run through history after the switch", async () => {
+    const resume = vi.fn(async function* () {
+      yield { content: [{ type: "text" as const, text: "resumed" }] };
+    });
+    const history: ThreadHistoryAdapter = {
+      load: vi.fn(async () => ({ headId: null, messages: [] })),
+      append: vi.fn(async () => {}),
+      resume,
+    };
+    const { result } = renderRuntime(
+      async () => ({ messages: [message("loaded")], unstable_resume: true }),
+      undefined,
+      undefined,
+      { history },
+    );
+
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+
+    await waitFor(() => expect(resume).toHaveBeenCalledOnce());
+  });
+
+  it("does not dispatch a queued send on the aborted run's idle edge", async () => {
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    const { result } = renderRuntime(
+      async () => ({ messages: [message("loaded")] }),
+      undefined,
+      agent,
+      { unstable_enableMessageQueue: true },
+    );
+
+    act(() => {
+      void result.current.thread.append("active");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "queued" }],
+        parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+    expect(runAgent).toHaveBeenCalledOnce();
+
+    let switching!: Promise<void>;
+    act(() => {
+      switching = result.current.threads.switchToThread("thread-a");
+    });
+    await act(async () => {
+      await switching;
+    });
+
+    expect(runAgent).toHaveBeenCalledOnce();
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "loaded",
+    ]);
+    act(() => {
+      void result.current.thread.append("new thread prompt");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+  });
+
+  it("keeps the queue usable after onCancel queues a send mid-switch", async () => {
+    const staleLoad = deferred<ThreadLoad>();
+    const runAgent = vi.fn(
+      async (
+        _input: unknown,
+        _subscriber: Subscriber,
+        options: { signal: AbortSignal },
+      ) => {
+        await new Promise<void>((resolve) =>
+          options.signal.addEventListener("abort", () => resolve(), {
+            once: true,
+          }),
+        );
+      },
+    );
+    const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+    let result!: ReturnType<typeof renderRuntime>["result"];
+    const onCancel = vi.fn(() => {
+      void result.current.thread.append("sent mid-switch");
+    });
+    ({ result } = renderRuntime(
+      (id) =>
+        id === "thread-a"
+          ? staleLoad.promise
+          : Promise.resolve({ messages: [message("thread-b")] }),
+      undefined,
+      agent,
+      { unstable_enableMessageQueue: true, onCancel },
+    ));
+
+    act(() => {
+      void result.current.thread.append("active");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    act(() => {
+      void result.current.threads.switchToThread("thread-a");
+    });
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-b");
+    });
+
+    expect(onCancel).toHaveBeenCalledOnce();
+    expect(runAgent).toHaveBeenCalledOnce();
+    expect(result.current.thread.getState().messages.map((m) => m.id)).toEqual([
+      "thread-b",
+    ]);
+    act(() => {
+      void result.current.thread.append("after the switch");
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+  });
+
+  it.each([false, true])(
+    "accepts a send after a newer switch completes with queue %s",
+    async (unstable_enableMessageQueue) => {
+      const staleLoad = deferred<ThreadLoad>();
+      const runAgent = vi.fn(async () => {});
+      const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+      const { result } = renderRuntime(
+        (id) =>
+          id === "thread-a"
+            ? staleLoad.promise
+            : Promise.resolve({ messages: [message("thread-b")] }),
+        undefined,
+        agent,
+        { unstable_enableMessageQueue },
+      );
+
+      act(() => {
+        void result.current.threads.switchToThread("thread-a");
+      });
+      await act(async () => {
+        await result.current.threads.switchToThread("thread-b");
+      });
+      expect(result.current.threads.getState().mainThreadId).toBe("thread-b");
+
+      act(() => {
+        void result.current.thread.append("send on B");
+      });
+      await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+    },
+  );
+
   it.each([undefined, { owner: "thread-a" }])(
     "ignores an older load and its resume request with state %j",
     async (state) => {