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

package/src/runtime/AgUiThreadController.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi } from "vitest";
-import type { AbstractAgent } from "@ag-ui/client";
+import { HttpAgent, type AbstractAgent } from "@ag-ui/client";
 import type {
   AppendMessage,
   ThreadAssistantMessage,
@@ -539,6 +539,178 @@ describe("AgUiThreadController activity deltas", () => {
   });
 });
 
+describe("AgUiThreadController run state", () => {
+  const createStateAgent = (sentStates: unknown[]) =>
+    new HttpAgent({
+      url: "http://agent.test/agent",
+      initialState: { fromInitialState: true },
+      fetch: async (_url, init) => {
+        const body = JSON.parse(String(init?.body));
+        sentStates.push(body.state);
+        const events = [
+          { type: "RUN_STARTED", threadId: body.threadId, runId: body.runId },
+          { type: "STATE_SNAPSHOT", snapshot: { a: 1 } },
+          { type: "RUN_FINISHED", threadId: body.threadId, runId: body.runId },
+        ];
+        return new Response(
+          events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""),
+          { headers: { "content-type": "text/event-stream" } },
+        );
+      },
+    });
+
+  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();
+  });
+
+  it("isolates mounted threads sharing one agent through runs and setState", async () => {
+    const sentStates: unknown[] = [];
+    const agent = createStateAgent(sentStates);
+    const { core: threadA } = createCore(undefined, agent);
+    const { core: threadB } = createCore(undefined, agent);
+
+    await threadA.append({ ...userMessage("A"), startRun: true });
+    expect(threadA.getState()).toEqual({ a: 1 });
+    expect.soft(threadB.getState()).toEqual({ fromInitialState: true });
+
+    threadA.setState({ x: 1 });
+    await flush();
+    expect(threadA.getState()).toEqual({ x: 1 });
+    expect.soft(threadB.getState()).toEqual({ fromInitialState: true });
+
+    await threadB.append({ ...userMessage("B"), startRun: true });
+    expect
+      .soft(sentStates)
+      .toEqual([{ fromInitialState: true }, { fromInitialState: true }]);
+    expect.soft(threadA.getState()).toEqual({ x: 1 });
+  });
+
+  it("starts a later thread from the shared agent's initial state after a run", async () => {
+    const sentStates: unknown[] = [];
+    const agent = createStateAgent(sentStates);
+    const { core: threadA } = createCore(undefined, agent);
+
+    await threadA.append({ ...userMessage("A"), startRun: true });
+    expect(agent.state).toEqual({ a: 1 });
+
+    const { core: threadB } = createCore(undefined, agent);
+    expect.soft(threadB.getState()).toEqual({ fromInitialState: true });
+    await threadB.append({ ...userMessage("B"), startRun: true });
+    expect(sentStates).toEqual([
+      { fromInitialState: true },
+      { fromInitialState: true },
+    ]);
+  });
+
+  it("copies the initial state before the agent state is mutated", () => {
+    const agent = new HttpAgent({
+      url: "http://agent.test/agent",
+      initialState: { nested: { count: 0 } },
+    });
+    const { core } = createCore(undefined, agent);
+
+    agent.state.nested.count = 1;
+    expect(core.getState()).toEqual({ nested: { count: 0 } });
+
+    const { core: laterThread } = createCore(undefined, agent);
+    expect(laterThread.getState()).toEqual({ nested: { count: 0 } });
+  });
+
+  it("records an undefined initial state without adopting later agent state", () => {
+    const agent = {} as AbstractAgent;
+    createCore(undefined, agent);
+    agent.state = { fromPreviousThread: true };
+
+    const { core } = createCore(undefined, agent);
+    expect(core.getState()).toBeUndefined();
+  });
+
+  it("leaves the state undefined for an agent without initialState", () => {
+    const agent = new HttpAgent({ url: "http://agent.test/agent" });
+    const { core } = createCore(undefined, agent);
+
+    expect(agent.state).toEqual({});
+    expect(core.getState()).toBeUndefined();
+
+    const swapped = new HttpAgent({ url: "http://agent.test/other" });
+    core.updateOptions({
+      agent: swapped,
+      logger: makeLogger(),
+      showThinking: true,
+    });
+    expect(core.getState()).toBeUndefined();
+  });
+
+  it("seeds a new agent's initial state on swap when the thread has no snapshot", () => {
+    const { core } = createCore();
+    const agent = new HttpAgent({
+      url: "http://agent.test/agent",
+      initialState: { a: 1 },
+    });
+    const listener = vi.fn();
+    core.getSnapshot();
+    core.subscribe(listener);
+
+    core.updateOptions({ agent, logger: makeLogger(), showThinking: true });
+
+    expect(agent.state).toEqual({ a: 1 });
+    expect(core.getState()).toEqual({ a: 1 });
+    expect(core.getSnapshot().state).toEqual({ a: 1 });
+    expect(listener).toHaveBeenCalledOnce();
+  });
+
+  it("preserves the thread snapshot on swap and records the new agent's initial state", async () => {
+    const { core } = createCore();
+    core.setState({ fromThread: true });
+    const sentStates: unknown[] = [];
+    const agent = createStateAgent(sentStates);
+
+    core.updateOptions({ agent, logger: makeLogger(), showThinking: true });
+
+    expect(agent.state).toEqual({ fromInitialState: true });
+    expect(core.getState()).toEqual({ fromThread: true });
+    await core.append({ ...userMessage("after swap"), startRun: true });
+    expect(sentStates).toEqual([{ fromThread: true }]);
+
+    const { core: laterThread } = createCore(undefined, agent);
+    expect(laterThread.getState()).toEqual({ fromInitialState: true });
+  });
+});
+
 describe("AgUiThreadController setState", () => {
   it("keeps a cancel the updater makes while a send waits for history", async () => {
     const { history, finish } = deferredHistory();