@assistant-ui/react-langchain 0.0.31 → 0.0.33

package/src/convertMessages.ts

@@ -6,6 +6,7 @@ import type { ReadonlyJSONObject } from "assistant-stream/utils";
 import {
   convertLangChainContentBlock,
   getCustomMetadata,
+  getMessageModality,
   uiMessageToDataPart,
   withAudioTranscript,
 } from "./converter";
@@ -55,6 +56,35 @@ const contentBlocks = (content: unknown): readonly LangChainContentBlock[] => {
   return [];
 };
 
+const normalizeToolCallArgs = (args: unknown): ReadonlyJSONObject => {
+  if (typeof args !== "object" || args === null || Array.isArray(args)) {
+    return {};
+  }
+
+  try {
+    const prototype = Object.getPrototypeOf(args);
+    return prototype === Object.prototype || prototype === null
+      ? (args as ReadonlyJSONObject)
+      : {};
+  } catch {
+    return {};
+  }
+};
+
+const toolCallArgs = (
+  value: unknown,
+): { args: ReadonlyJSONObject; argsText: string } => {
+  const args = normalizeToolCallArgs(value);
+  try {
+    const argsText = JSON.stringify(args);
+    return typeof argsText === "string"
+      ? { args, argsText }
+      : { args: {}, argsText: "{}" };
+  } catch {
+    return { args: {}, argsText: "{}" };
+  }
+};
+
 const contentToParts = (content: unknown) => {
   if (typeof content === "string")
     return [{ type: "text" as const, text: content }];
@@ -89,25 +119,31 @@ export const convertLangChainBaseMessage = (
         },
       };
 
-    case "human":
+    case "human": {
+      const modality = getMessageModality(message.additional_kwargs);
       return {
         role: "user",
         id: message.id,
         content: contentToParts(message.content),
         metadata: {
           custom: getCustomMetadata(message.additional_kwargs),
+          ...(modality && { modality }),
         },
       };
+    }
 
     case "ai": {
       const toolCallParts =
-        message.tool_calls?.map((tc) => ({
-          type: "tool-call" as const,
-          toolCallId: tc.id,
-          toolName: tc.name,
-          args: tc.args as ReadonlyJSONObject,
-          argsText: JSON.stringify(tc.args),
-        })) ?? [];
+        message.tool_calls?.map((tc) => {
+          const { args, argsText } = toolCallArgs(tc.args);
+          return {
+            type: "tool-call" as const,
+            toolCallId: tc.id,
+            toolName: tc.name,
+            args,
+            argsText,
+          };
+        }) ?? [];
 
       const assistantStatus =
         typeof message.status === "object" ? message.status : undefined;
@@ -120,6 +156,7 @@ export const convertLangChainBaseMessage = (
           : undefined) ?? [];
 
       const timing = metadata.messageTiming?.[message.id ?? ""];
+      const modality = getMessageModality(message.additional_kwargs);
 
       return {
         role: "assistant",
@@ -135,6 +172,7 @@ export const convertLangChainBaseMessage = (
         metadata: {
           custom: getCustomMetadata(message.additional_kwargs),
           ...(timing && { timing }),
+          ...(modality && { modality }),
         },
         ...(assistantStatus && { status: assistantStatus }),
       };

package/src/converter.ts

@@ -1,6 +1,7 @@
 import type {
   AppendMessage,
   DataMessagePart,
+  MessageModality,
   ThreadAssistantMessage,
   ThreadUserMessage,
 } from "@assistant-ui/core";
@@ -146,6 +147,11 @@ export const getCustomMetadata = (
 ): Record<string, unknown> =>
   (additionalKwargs?.metadata as Record<string, unknown>) ?? {};
 
+export const getMessageModality = (
+  additionalKwargs: Record<string, unknown> | undefined,
+): MessageModality | undefined =>
+  additionalKwargs?.modality === "voice" ? "voice" : undefined;
+
 export const uiMessageToDataPart = <
   TUIMessage extends { name: string; props: Record<string, unknown> },
 >(

package/src/subagentMessagesProjection.test.ts

@@ -0,0 +1,175 @@
+import { describe, expect, it, vi } from "vitest";
+import type { BaseMessage } from "@langchain/core/messages";
+import {
+  matchesSubscription,
+  SubscriptionHandle,
+} from "@langchain/langgraph-sdk/client";
+import {
+  type Event,
+  messagesProjection,
+  type ProjectionSpec,
+  StreamStore,
+} from "@langchain/langgraph-sdk/stream";
+import { subagentMessagesProjection } from "./subagentMessagesProjection";
+
+const PARENT = ["tools:parent"];
+const CHILD = ["tools:parent", "tools:child"];
+
+const human = (id: string, content: string) => ({
+  id,
+  type: "human",
+  content,
+});
+const ai = (id: string, content: string, tool_calls: unknown[] = []) => ({
+  id,
+  type: "ai",
+  content,
+  tool_calls,
+});
+const tool = (id: string, content: string, tool_call_id: string) => ({
+  id,
+  type: "tool",
+  content,
+  tool_call_id,
+});
+
+const values = (namespace: string[], messages: unknown[]) =>
+  ({
+    method: "values",
+    params: { namespace, data: { messages } },
+  }) as unknown as Event;
+
+const messagesEvent = (namespace: string[], data: Record<string, unknown>) =>
+  ({
+    method: "messages",
+    params: { namespace, node: "model", data },
+  }) as unknown as Event;
+
+/**
+ * Feeds the projection through the per-subscription matcher the SDK client
+ * applies before it hands an event to a subscription.
+ */
+const openProjection = (spec: ProjectionSpec<BaseMessage[]>) => {
+  let handle: SubscriptionHandle<Event> | undefined;
+  const unsubscribe = vi.fn(async () => {});
+  const subscribe = vi.fn(
+    async (params: SubscriptionHandle<Event>["params"]) => {
+      handle = new SubscriptionHandle<Event>(
+        "subscription",
+        params,
+        unsubscribe,
+      );
+      return handle;
+    },
+  );
+  const store = new StreamStore<BaseMessage[]>(spec.initial);
+  const runtime = spec.open({
+    thread: { subscribe } as never,
+    store,
+    rootBus: {
+      channels: [
+        "values",
+        "checkpoints",
+        "lifecycle",
+        "input",
+        "messages",
+        "tools",
+      ],
+      subscribe: () => () => {},
+    },
+  });
+  const push = (event: Event) => {
+    if (matchesSubscription(event, handle!.params)) handle!.push(event);
+  };
+  const ids = () => store.getSnapshot().map((message) => message.id);
+  return { subscribe, unsubscribe, runtime, push, ids };
+};
+
+const parentTurn = [
+  human("parent-human", "research"),
+  ai("parent-ai", "", [
+    { id: "call-child", name: "task", args: { subagent_type: "worker" } },
+  ]),
+];
+
+const childRun = [
+  values(CHILD, [human("child-human", "sub task")]),
+  messagesEvent(CHILD, { event: "message-start", id: "child-ai", role: "ai" }),
+  messagesEvent(CHILD, {
+    event: "content-block-start",
+    index: 0,
+    content: { type: "text", text: "hello from the child" },
+  }),
+  messagesEvent(CHILD, { event: "message-finish" }),
+  values(CHILD, [
+    human("child-human", "sub task"),
+    ai("child-ai", "hello from the child"),
+  ]),
+];
+
+describe("subagentMessagesProjection", () => {
+  it("keeps the SDK projection's identity and pins its subscription to depth 0", async () => {
+    const spec = subagentMessagesProjection(PARENT);
+    expect(spec.namespace).toEqual(PARENT);
+    expect(spec.initial).toEqual([]);
+    expect(spec.key).not.toBe(messagesProjection(PARENT).key);
+
+    const { subscribe, runtime } = openProjection(spec);
+    await vi.waitFor(() => expect(subscribe).toHaveBeenCalledOnce());
+    expect(subscribe).toHaveBeenCalledWith({
+      channels: ["messages", "values"],
+      namespaces: [PARENT],
+      depth: 0,
+    });
+    await runtime.dispose();
+  });
+
+  it("ignores a nested subagent's events while the parent's tool call runs", async () => {
+    const { subscribe, runtime, push, ids } = openProjection(
+      subagentMessagesProjection(PARENT),
+    );
+    await vi.waitFor(() => expect(subscribe).toHaveBeenCalledOnce());
+
+    push(values(PARENT, parentTurn));
+    await vi.waitFor(() =>
+      expect(ids()).toEqual(["parent-human", "parent-ai"]),
+    );
+
+    for (const event of childRun) push(event);
+    push(
+      messagesEvent(PARENT, {
+        event: "message-start",
+        id: "parent-sentinel",
+        role: "ai",
+      }),
+    );
+    push(messagesEvent(PARENT, { event: "message-finish" }));
+    await vi.waitFor(() => expect(ids()).toContain("parent-sentinel"));
+    expect(ids()).toEqual(["parent-human", "parent-ai", "parent-sentinel"]);
+
+    push(
+      values(PARENT, [
+        ...parentTurn,
+        tool("parent-tool", "done", "call-child"),
+      ]),
+    );
+    await vi.waitFor(() =>
+      expect(ids()).toEqual([
+        "parent-human",
+        "parent-ai",
+        "parent-tool",
+        "parent-sentinel",
+      ]),
+    );
+    await runtime.dispose();
+  });
+
+  it("unsubscribes the underlying subscription on dispose", async () => {
+    const { subscribe, unsubscribe, runtime } = openProjection(
+      subagentMessagesProjection(PARENT),
+    );
+    await vi.waitFor(() => expect(subscribe).toHaveBeenCalledOnce());
+    await runtime.dispose();
+    await vi.waitFor(() => expect(unsubscribe).toHaveBeenCalledOnce());
+  });
+});

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/uiMessages.test.ts

@@ -1,10 +1,13 @@
 import { describe, expect, it } from "vitest";
 import {
   applyUIUpdate,
+  createUIFoldMemo,
+  createUISnapshotMemo,
   extractUIUpdate,
   foldUIUpdates,
   isUIUpdate,
   mergeUIMessages,
+  reconcileUISnapshot,
 } from "./uiMessages";
 import type { UIMessage } from "./types";
 
@@ -137,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")]);
@@ -187,4 +294,46 @@ describe("foldUIUpdates", () => {
   it("returns an empty list for no events", () => {
     expect(foldUIUpdates([])).toEqual([]);
   });
+
+  const merge = (id: string, props: Record<string, unknown>) => ({
+    ...ui(id, props),
+    metadata: { merge: true },
+  });
+
+  it("keeps the previous list when appended events carry no UI update", () => {
+    const memo = createUIFoldMemo();
+    const events = [evt(ui("a", { x: 1 })), evt(merge("a", { y: 2 }))];
+    const folded = foldUIUpdates(events, memo);
+
+    expect(foldUIUpdates([...events, evt({ progress: 1 })], memo)).toBe(folded);
+  });
+
+  it("keeps entries that appended updates do not touch", () => {
+    const memo = createUIFoldMemo();
+    const events = [evt(ui("a", { x: 1 })), evt(merge("a", { y: 2 }))];
+    const [merged] = foldUIUpdates(events, memo);
+
+    const result = foldUIUpdates([...events, evt(ui("b"))], memo);
+
+    expect(result).toEqual([merge("a", { x: 1, y: 2 }), ui("b")]);
+    expect(result[0]).toBe(merged);
+  });
+
+  it("keeps folded entries when the buffer drops its oldest events", () => {
+    const memo = createUIFoldMemo();
+    const progress = evt({ progress: 1 });
+    foldUIUpdates([evt(ui("a", { x: 1 })), progress], memo);
+
+    expect(foldUIUpdates([progress, evt(merge("a", { y: 2 }))], memo)).toEqual([
+      merge("a", { x: 1, y: 2 }),
+    ]);
+  });
+
+  it("folds from scratch when the buffer is replaced", () => {
+    const memo = createUIFoldMemo();
+    foldUIUpdates([evt(ui("a"))], memo);
+
+    expect(foldUIUpdates([evt(ui("b"))], memo)).toEqual([ui("b")]);
+    expect(foldUIUpdates([], memo)).toEqual([]);
+  });
 });

package/src/uiMessages.ts

@@ -1,5 +1,10 @@
+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. */
+export const UI_CUSTOM_CHANNELS: readonly Channel[] = ["custom"];
+
 export type UIUpdate = UIMessage | RemoveUIMessage;
 
 export const isUIUpdate = (
@@ -62,15 +67,85 @@ export const extractUIUpdate = (
   return undefined;
 };
 
-export const foldUIUpdates = (events: readonly unknown[]): UIMessage[] => {
-  let acc: UIMessage[] = [];
-  for (const event of events) {
-    const update = extractUIUpdate(event);
+export type UIFoldMemo = {
+  events: readonly unknown[];
+  messages: UIMessage[];
+};
+
+export const createUIFoldMemo = (): UIFoldMemo => ({
+  events: [],
+  messages: [],
+});
+
+/**
+ * Folds `custom`-channel events into UI messages, continuing the fold held in
+ * `memo`. The channel buffer holds each event object once, appends new events
+ * and drops its oldest ones once full, so only events after the previously
+ * folded last event are applied, and entries no later event touches keep
+ * their identity. A buffer that no longer contains that event was replaced
+ * and is folded from scratch.
+ */
+export const foldUIUpdates = (
+  events: readonly unknown[],
+  memo: UIFoldMemo = createUIFoldMemo(),
+): UIMessage[] => {
+  const previous = memo.events;
+  const resumeAt =
+    previous.length === 0
+      ? 0
+      : events.lastIndexOf(previous.at(-1), previous.length - 1) + 1;
+  let acc = previous.length > 0 && resumeAt === 0 ? [] : memo.messages;
+  for (let i = resumeAt; i < events.length; i++) {
+    const update = extractUIUpdate(events[i]);
     if (update) acc = applyUIUpdate(acc, update);
   }
+  memo.events = events;
+  memo.messages = acc;
   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