@assistant-ui/react-ag-ui 0.0.71 → 0.1.1

package/src/runtime/{AgUiThreadRuntimeCore.test.ts → AgUiThreadController.test.ts}

@@ -6,7 +6,7 @@ import type {
   ThreadHistoryAdapter,
   ThreadMessage,
 } from "@assistant-ui/core";
-import { AgUiThreadRuntimeCore } from "./AgUiThreadRuntimeCore";
+import { AgUiThreadController } from "./AgUiThreadController";
 import { makeLogger } from "./logger";
 
 const flush = () => new Promise((resolve) => setTimeout(resolve, 0));
@@ -15,7 +15,7 @@ function createCore(
   history?: ThreadHistoryAdapter,
   agent = {} as AbstractAgent,
 ) {
-  const core = new AgUiThreadRuntimeCore({
+  const core = new AgUiThreadController({
     agent,
     logger: makeLogger(),
     showThinking: true,
@@ -141,7 +141,7 @@ function deferredHistory(messages: ThreadMessage[] = [storedUser]) {
   };
 }
 
-describe("AgUiThreadRuntimeCore late history loading", () => {
+describe("AgUiThreadController late history loading", () => {
   it("drops a queued send when loading is cancelled", async () => {
     const { history, finish } = deferredHistory();
     const runAgent = vi.fn(async () => {});
@@ -479,7 +479,7 @@ describe("AgUiThreadRuntimeCore late history loading", () => {
   });
 });
 
-describe("AgUiThreadRuntimeCore steerAway parent selection", () => {
+describe("AgUiThreadController steerAway parent selection", () => {
   it("starts a root branch for an explicit null parent", async () => {
     const { core } = createCore();
     await core.append(userMessage("old"));
@@ -497,7 +497,7 @@ describe("AgUiThreadRuntimeCore steerAway parent selection", () => {
   });
 });
 
-describe("AgUiThreadRuntimeCore activity deltas", () => {
+describe("AgUiThreadController activity deltas", () => {
   it("applies a delta to the snapshot the run streamed", async () => {
     type Subscriber = Record<string, ((payload?: unknown) => void) | undefined>;
     const runAgent = vi.fn(async (_input: unknown, subscriber: Subscriber) => {
@@ -521,7 +521,7 @@ describe("AgUiThreadRuntimeCore activity deltas", () => {
       subscriber.onRunFinalized?.();
     });
     const agent = { runAgent, abortRun: vi.fn() } as unknown as AbstractAgent;
-    const core = new AgUiThreadRuntimeCore({
+    const core = new AgUiThreadController({
       agent,
       logger: makeLogger(),
       showThinking: true,
@@ -538,3 +538,111 @@ describe("AgUiThreadRuntimeCore activity deltas", () => {
     });
   });
 });
+
+describe("AgUiThreadController setState", () => {
+  it("keeps a cancel the updater makes while a send waits for history", async () => {
+    const { history, finish } = deferredHistory();
+    const runAgent = vi.fn(async () => {});
+    const { core } = createCore(history, {
+      runAgent,
+      abortRun: vi.fn(),
+    } as unknown as AbstractAgent);
+    const loading = core.__internal_load();
+    const sending = core.append({
+      ...userMessage("cancelled"),
+      startRun: true,
+    });
+
+    core.setState(() => {
+      void core.cancel();
+      return { count: 1 };
+    });
+    finish();
+    await Promise.all([loading, sending]);
+
+    expect(core.getState()).toEqual({ count: 1 });
+    expect(core.getMessages().map((message) => message.id)).toEqual([
+      "restored",
+    ]);
+    expect(runAgent).not.toHaveBeenCalled();
+  });
+});
+
+describe("AgUiThreadController resuming a disconnected run", () => {
+  type Subscriber = {
+    onTextMessageStartEvent?: (payload: { event: unknown }) => void;
+    onTextMessageContentEvent?: (payload: { event: unknown }) => void;
+    onTextMessageEndEvent?: (payload: { event: unknown }) => void;
+    onRunFinishedEvent?: (payload: { event: unknown }) => void;
+  };
+  const sendText = (
+    subscriber: Subscriber,
+    messageId: string,
+    text: string,
+    end = true,
+  ) => {
+    subscriber.onTextMessageStartEvent?.({
+      event: { type: "TEXT_MESSAGE_START", messageId },
+    });
+    subscriber.onTextMessageContentEvent?.({
+      event: { type: "TEXT_MESSAGE_CONTENT", messageId, delta: text },
+    });
+    if (end) {
+      subscriber.onTextMessageEndEvent?.({
+        event: { type: "TEXT_MESSAGE_END", messageId },
+      });
+    }
+  };
+
+  it.each(["user", "assistant"] as const)(
+    "preserves replayed IDs and content when resume uses the %s parent",
+    async (parent) => {
+      const runAgent = vi.fn(
+        async (input: { runId: string }, subscriber: Subscriber) => {
+          if (runAgent.mock.calls.length === 1) {
+            sendText(subscriber, "assistant-a", "first", false);
+            throw new Error("stream disconnected");
+          }
+          sendText(subscriber, "assistant-a", "first and more");
+          sendText(subscriber, "assistant-b", "second answer");
+          subscriber.onRunFinishedEvent?.({
+            event: { type: "RUN_FINISHED", runId: input.runId },
+          });
+        },
+      );
+      const { core } = createCore(undefined, {
+        runAgent,
+        abortRun: vi.fn(),
+      } as unknown as AbstractAgent);
+
+      await expect(
+        core.append({ ...userMessage("question"), startRun: true }),
+      ).rejects.toThrow("stream disconnected");
+      const userId = core.getMessages()[0]!.id;
+      expect(core.getMessages()[1]).toMatchObject({
+        id: "assistant-a",
+        content: [{ type: "text", text: "first" }],
+      });
+
+      await core.resume({
+        parentId: parent === "user" ? userId : "assistant-a",
+        sourceId: null,
+        runConfig: {},
+      });
+
+      expect(runAgent).toHaveBeenCalledTimes(2);
+      expect(core.getMessages().map((message) => message.id)).toEqual([
+        userId,
+        "assistant-a",
+        "assistant-b",
+      ]);
+      expect(core.getMessages().slice(1)).toMatchObject([
+        { content: [{ type: "text", text: "first and more" }] },
+        { content: [{ type: "text", text: "second answer" }] },
+      ]);
+      expect(
+        core.getMessageRepository().messages.map(({ parentId }) => parentId),
+      ).toEqual([null, userId, "assistant-a"]);
+    },
+  );
+});