@assistant-ui/react-ag-ui 0.1.0 → 0.1.2

package/src/runtime/AgUiThreadController.test.ts

@@ -539,6 +539,46 @@ describe("AgUiThreadController activity deltas", () => {
   });
 });
 
+describe("AgUiThreadController run state", () => {
+  type Subscriber = {
+    onStateSnapshotEvent?: (payload: {
+      event: { type: "STATE_SNAPSHOT"; snapshot: { count: number } };
+    }) => void;
+  };
+
+  it("sends an empty state when no snapshot exists and preserves null", async () => {
+    const runAgent = vi.fn(
+      async (_input: { state: unknown }, subscriber: Subscriber) => {
+        if (runAgent.mock.calls.length === 1) {
+          subscriber.onStateSnapshotEvent?.({
+            event: { type: "STATE_SNAPSHOT", snapshot: { count: 1 } },
+          });
+        }
+      },
+    );
+    const agent = {
+      runAgent,
+      abortRun: vi.fn(),
+    } as unknown as AbstractAgent;
+    const { core } = createCore(undefined, agent);
+
+    await core.append({ ...userMessage("initial"), startRun: true });
+    expect(runAgent.mock.calls[0]?.[0].state).toEqual({});
+    expect(agent.state).toEqual({});
+
+    core.resetState();
+    expect(core.getState()).toBeUndefined();
+    await core.append({ ...userMessage("after reset"), startRun: true });
+    expect(runAgent.mock.calls[1]?.[0].state).toEqual({});
+    expect(agent.state).toEqual({});
+
+    core.setState(null);
+    await core.append({ ...userMessage("explicit null"), startRun: true });
+    expect(runAgent.mock.calls[2]?.[0].state).toBeNull();
+    expect(agent.state).toBeNull();
+  });
+});
+
 describe("AgUiThreadController setState", () => {
   it("keeps a cancel the updater makes while a send waits for history", async () => {
     const { history, finish } = deferredHistory();
@@ -567,3 +607,82 @@ describe("AgUiThreadController setState", () => {
     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"]);
+    },
+  );
+});