@assistant-ui/react-langchain 0.0.32 → 0.0.34

package/src/subagentMessagesProjection.ts

@@ -0,0 +1,50 @@
+import type { BaseMessage } from "@langchain/core/messages";
+import {
+  messagesProjection,
+  type ProjectionSpec,
+} from "@langchain/langgraph-sdk/stream";
+
+type MessagesProjection = ProjectionSpec<BaseMessage[]>;
+
+type ProjectionThread = Parameters<MessagesProjection["open"]>[0]["thread"];
+
+const exactDepthThread = (thread: ProjectionThread): ProjectionThread =>
+  new Proxy(thread, {
+    get(target, property) {
+      if (property === "subscribe") {
+        return (...[params]: Parameters<ProjectionThread["subscribe"]>) =>
+          target.subscribe({ ...params, depth: 0 });
+      }
+      const value = Reflect.get(target, property, target);
+      return typeof value === "function" ? value.bind(target) : value;
+    },
+  });
+
+/**
+ * The SDK's messages projection subscribes at the protocol's default depth
+ * of 1 and applies every delivered event, so a nested subagent's `values`
+ * snapshots rebuild its parent's store from the child's state while the child
+ * runs. The projection exposes no depth option, so this spec opens it against
+ * a thread whose subscription is pinned to depth 0: the namespace the SDK
+ * sets stays and only events at exactly that namespace are delivered, the
+ * rule the SDK's root projection already applies. The client unions every
+ * subscription's depth into the server filter, so a depth 0 subscription
+ * narrows nothing for other consumers.
+ */
+export const subagentMessagesProjection = (
+  namespace: readonly string[],
+): MessagesProjection => {
+  const projection = messagesProjection(namespace);
+  return {
+    key: `exact|${projection.key}`,
+    namespace: projection.namespace,
+    initial: projection.initial,
+    open({ thread, store, rootBus }) {
+      return projection.open({
+        thread: exactDepthThread(thread),
+        store,
+        rootBus,
+      });
+    },
+  };
+};

package/src/tests/settleOutsideAct.ts

@@ -0,0 +1,15 @@
+import { act, waitFor } from "@testing-library/react";
+
+// React 18's act holds renders until its callback settles, so a task that waits on a render starts inside act and settles outside it, where React can commit what the task waits on.
+export async function settleOutsideAct<T>(task: () => Promise<T>): Promise<T> {
+  let pending!: Promise<T>;
+  await act(async () => {
+    pending = task();
+    pending.catch(() => {});
+  });
+  await waitFor(() => pending.then(settled, settled));
+  await act(async () => {});
+  return pending;
+}
+
+const settled = () => {};

package/src/uiMessages.test.ts

@@ -2,10 +2,12 @@ import { describe, expect, it } from "vitest";
 import {
   applyUIUpdate,
   createUIFoldMemo,
+  createUISnapshotMemo,
   extractUIUpdate,
   foldUIUpdates,
   isUIUpdate,
   mergeUIMessages,
+  reconcileUISnapshot,
 } from "./uiMessages";
 import type { UIMessage } from "./types";
 
@@ -138,6 +140,110 @@ describe("extractUIUpdate", () => {
   });
 });
 
+describe("reconcileUISnapshot", () => {
+  const snapshot = (...entries: UIMessage[]) =>
+    entries.map((entry) => structuredClone(entry));
+
+  it("returns the previous list for an equal copy of the snapshot", () => {
+    const memo = createUISnapshotMemo();
+    const first = reconcileUISnapshot(
+      snapshot(ui("a", { x: 1 }), ui("b", { y: [1, 2] })),
+      memo,
+    );
+
+    const second = reconcileUISnapshot(
+      snapshot(ui("a", { x: 1 }), ui("b", { y: [1, 2] })),
+      memo,
+    );
+
+    expect(second).toBe(first);
+    expect(second).toEqual([ui("a", { x: 1 }), ui("b", { y: [1, 2] })]);
+  });
+
+  it("returns the same list for the same snapshot reference", () => {
+    const memo = createUISnapshotMemo();
+    const value = snapshot(ui("a"));
+    const first = reconcileUISnapshot(value, memo);
+
+    expect(reconcileUISnapshot(value, memo)).toBe(first);
+  });
+
+  it("replaces only the entry that changed", () => {
+    const memo = createUISnapshotMemo();
+    const [a, b] = reconcileUISnapshot(
+      snapshot(ui("a", { x: 1 }), ui("b", { y: 1 })),
+      memo,
+    );
+
+    const result = reconcileUISnapshot(
+      snapshot(ui("a", { x: 1 }), ui("b", { y: 2 })),
+      memo,
+    );
+
+    expect(result[0]).toBe(a);
+    expect(result[1]).not.toBe(b);
+    expect(result[1]).toEqual(ui("b", { y: 2 }));
+  });
+
+  it("keeps surviving entries across an addition and a removal", () => {
+    const memo = createUISnapshotMemo();
+    const [a, b] = reconcileUISnapshot(snapshot(ui("a"), ui("b")), memo);
+
+    const added = reconcileUISnapshot(
+      snapshot(ui("a"), ui("b"), ui("c")),
+      memo,
+    );
+    expect(added[0]).toBe(a);
+    expect(added[1]).toBe(b);
+    expect(added).toHaveLength(3);
+
+    const removed = reconcileUISnapshot(snapshot(ui("b"), ui("c")), memo);
+    expect(removed[0]).toBe(b);
+    expect(removed[1]).toBe(added[2]);
+    expect(removed).toHaveLength(2);
+  });
+
+  it("returns a new list when equal entries change position", () => {
+    const memo = createUISnapshotMemo();
+    const [a, b] = reconcileUISnapshot(snapshot(ui("a"), ui("b")), memo);
+
+    const result = reconcileUISnapshot(snapshot(ui("b"), ui("a")), memo);
+
+    expect(result).toEqual([ui("b"), ui("a")]);
+    expect(result[0]).toBe(b);
+    expect(result[1]).toBe(a);
+  });
+
+  it("treats an entry with a changed name or metadata as new", () => {
+    const memo = createUISnapshotMemo();
+    const [a] = reconcileUISnapshot(
+      snapshot({ ...ui("a"), metadata: { message_id: "m1" } }),
+      memo,
+    );
+
+    const renamed = reconcileUISnapshot(
+      snapshot({ ...ui("a"), name: "table", metadata: { message_id: "m1" } }),
+      memo,
+    );
+    expect(renamed[0]).not.toBe(a);
+
+    const moved = reconcileUISnapshot(
+      snapshot({ ...ui("a"), name: "table", metadata: { message_id: "m2" } }),
+      memo,
+    );
+    expect(moved[0]).not.toBe(renamed[0]);
+  });
+
+  it("returns an empty list for a non-array snapshot", () => {
+    const memo = createUISnapshotMemo();
+    expect(reconcileUISnapshot(undefined, memo)).toEqual([]);
+    reconcileUISnapshot(snapshot(ui("a")), memo);
+
+    expect(reconcileUISnapshot(null, memo)).toEqual([]);
+    expect(reconcileUISnapshot({ id: "a" }, memo)).toEqual([]);
+  });
+});
+
 describe("mergeUIMessages", () => {
   it("returns the snapshot when there are no live messages", () => {
     expect(mergeUIMessages([], [ui("a")])).toEqual([ui("a")]);

package/src/uiMessages.ts

@@ -1,4 +1,5 @@
 import type { Channel } from "@langchain/react";
+import { isJSONValueEqual } from "@assistant-ui/core/internal";
 import type { RemoveUIMessage, UIMessage } from "./types";
 
 /** Channels the generative-UI fold reads, at the root and per subagent. */
@@ -103,6 +104,48 @@ export const foldUIUpdates = (
   return acc;
 };
 
+export type UISnapshotMemo = {
+  snapshot: unknown;
+  entries: readonly UIMessage[];
+};
+
+export const createUISnapshotMemo = (): UISnapshotMemo => ({
+  snapshot: undefined,
+  entries: [],
+});
+
+/**
+ * Recovers entry identity across `values` events. The SDK rebuilds the
+ * `values` object from every snapshot and reconciles only the messages slot by
+ * id, so an unchanged UI list arrives as a new array of new objects on every
+ * superstep. An entry structurally equal to the previous entry with its id is
+ * replaced by that previous object, and a list whose entries all survive in
+ * place is replaced by the previous list.
+ */
+export const reconcileUISnapshot = (
+  snapshot: unknown,
+  memo: UISnapshotMemo,
+): readonly UIMessage[] => {
+  if (snapshot === memo.snapshot) return memo.entries;
+  memo.snapshot = snapshot;
+  if (!Array.isArray(snapshot)) {
+    if (memo.entries.length > 0) memo.entries = [];
+    return memo.entries;
+  }
+  const previous = memo.entries;
+  const previousById = new Map(previous.map((ui) => [ui.id, ui]));
+  let same = snapshot.length === previous.length;
+  const entries = (snapshot as UIMessage[]).map((ui, index) => {
+    const before = previousById.get(ui.id);
+    const entry =
+      before !== undefined && isJSONValueEqual(before, ui) ? before : ui;
+    if (entry !== previous[index]) same = false;
+    return entry;
+  });
+  if (!same) memo.entries = entries;
+  return memo.entries;
+};
+
 /**
  * Merges live-streamed UI with the state snapshot. The snapshot is
  * authoritative by id: once a UI lands in graph state it supersedes its live

package/src/useStreamRuntime.test.tsx

@@ -9,7 +9,7 @@ import type {
   RemoteThreadListAdapter,
 } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
-import type { LangChainBaseMessage } from "./types";
+import type { LangChainBaseMessage, LangChainToolCall } from "./types";
 import { startTransition, Suspense, type ReactNode } from "react";
 import {
   useLangChainRespond,
@@ -32,6 +32,7 @@ vi.mock("@langchain/react", () => ({
 }));
 
 import { useStreamRuntime } from "./useStreamRuntime";
+import { settleOutsideAct } from "./tests/settleOutsideAct";
 
 type MockStream = {
   messages: LangChainBaseMessage[];
@@ -194,18 +195,18 @@ describe("useStreamRuntime thread options", () => {
 
     const view = render(<TestRuntime />);
 
-    await act(async () => {
-      await capture.runtime!.threads.switchToThread("thread-a");
-    });
+    await settleOutsideAct(() =>
+      capture.runtime!.threads.switchToThread("thread-a"),
+    );
 
     const threadAOptions = mockUseStream.mock.calls
       .map(([options]) => options as { threadId?: string | null })
       .findLast((options) => options.threadId === "thread-a");
     expect(threadAOptions).toBeDefined();
 
-    await act(async () => {
-      await capture.runtime!.threads.switchToThread("thread-b");
-    });
+    await settleOutsideAct(() =>
+      capture.runtime!.threads.switchToThread("thread-b"),
+    );
 
     const threadBOptions = mockUseStream.mock.calls
       .map(([options]) => options as { threadId?: string | null })
@@ -917,6 +918,45 @@ describe("useStreamRuntime staged messages", () => {
   });
 });
 
+describe("useStreamRuntime pending tool call cancellation", () => {
+  it("cancels only the pending tool calls that carry an id", async () => {
+    const stream = createMockStream([
+      message("u1", "human", "look it up"),
+      {
+        id: "a1",
+        _getType: () => "ai",
+        content: "",
+        tool_calls: [
+          { name: "lookup", args: {} } as LangChainToolCall,
+          { id: "call-1", name: "search", args: {} },
+        ],
+      },
+    ]);
+    const { auiResult } = renderAui(stream);
+
+    await act(async () => {
+      await auiResult.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "never mind" }],
+      });
+    });
+
+    expect(stream.submit).toHaveBeenCalledTimes(1);
+    expect(stream.submit.mock.calls[0]![0]).toEqual({
+      messages: [
+        {
+          type: "tool",
+          name: "search",
+          tool_call_id: "call-1",
+          content: JSON.stringify({ cancelled: true }),
+          status: "error",
+        },
+        { id: expect.any(String), type: "human", content: "never mind" },
+      ],
+    });
+  });
+});
+
 describe("useStreamRuntime committed refs", () => {
   it("submits through the committed stream after an abandoned render", async () => {
     const streamA = createMockStream();
@@ -1072,6 +1112,106 @@ describe("useStreamRuntime subagent transcripts", () => {
     expect(nestedTranscript()).toBe(rendered);
   });
 
+  it("keeps messages and transcripts across equal copies of the UI state", async () => {
+    const stream = createMockStream([
+      message("human-1", "human", "delegate"),
+      {
+        id: "root-ai",
+        _getType: () => "ai",
+        content: "",
+        tool_calls: [{ id: "task-one", name: "task", args: {} }],
+      },
+    ]);
+    const transcript = [message("nested-ai", "ai", "nested answer")];
+    stream.subagents = new Map([
+      [
+        "task-one",
+        {
+          id: "task-one",
+          namespace: ["tools:task-one"],
+          status: "running",
+          parentId: null,
+          depth: 1,
+          startedAt: new Date(1_000),
+          completedAt: null,
+        },
+      ],
+    ]);
+    stream[streamController]!.registry.acquire.mockReturnValue({
+      store: { getSnapshot: () => transcript, subscribe: () => () => {} },
+      release: vi.fn(),
+    });
+    const uiState = (points: number[]) => [
+      {
+        type: "ui",
+        id: "ui-root",
+        name: "chart",
+        props: { points },
+        metadata: { message_id: "root-ai" },
+      },
+      {
+        type: "ui",
+        id: "ui-nested",
+        name: "chart",
+        props: { points },
+        metadata: { message_id: "nested-ai" },
+      },
+    ];
+    stream.values = { ui: uiState([1, 2]) };
+    const { auiResult, rerender } = renderAui(stream);
+    const nestedTranscript = () => {
+      const { messages } = auiResult.current.thread.getState();
+      for (const threadMessage of messages) {
+        for (const part of threadMessage.content) {
+          if (part.type === "tool-call" && part.toolCallId === "task-one")
+            return part.messages;
+        }
+      }
+      return undefined;
+    };
+
+    await waitFor(() =>
+      expect(nestedTranscript()?.[0]?.content).toMatchObject([
+        { type: "text", text: "nested answer" },
+        { type: "data", name: "chart", data: { points: [1, 2] } },
+      ]),
+    );
+    const [human, ai] = auiResult.current.thread.getState().messages;
+    expect(ai?.content).toMatchObject([
+      { type: "tool-call", toolCallId: "task-one" },
+      { type: "data", name: "chart", data: { points: [1, 2] } },
+    ]);
+    const rendered = nestedTranscript();
+
+    for (let i = 0; i < 3; i++) {
+      stream.values = { ui: uiState([1, 2]) };
+      await act(async () => {
+        rerender();
+      });
+    }
+
+    const messages = auiResult.current.thread.getState().messages;
+    expect(messages[0]).toBe(human);
+    expect(messages[1]).toBe(ai);
+    expect(nestedTranscript()).toBe(rendered);
+
+    stream.values = { ui: uiState([1, 2, 3]) };
+    await act(async () => {
+      rerender();
+    });
+
+    expect(
+      auiResult.current.thread.getState().messages[1]?.content,
+    ).toMatchObject([
+      { type: "tool-call", toolCallId: "task-one" },
+      { type: "data", name: "chart", data: { points: [1, 2, 3] } },
+    ]);
+    expect(nestedTranscript()?.[0]?.content).toMatchObject([
+      { type: "text", text: "nested answer" },
+      { type: "data", name: "chart", data: { points: [1, 2, 3] } },
+    ]);
+  });
+
   it("keeps messages and transcripts when custom events carry no UI update", async () => {
     const stream = createMockStream([
       message("human-1", "human", "delegate"),