@assistant-ui/react-langchain 0.0.30 → 0.0.32

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react-langchain",
-  "version": "0.0.30",
+  "version": "0.0.32",
   "description": "LangChain useStream adapter for assistant-ui",
   "keywords": [
     "langchain",
@@ -40,17 +40,21 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@assistant-ui/core": "^0.3.18",
-    "@assistant-ui/store": "^0.3.13",
+    "@assistant-ui/core": "^0.3.20",
+    "@assistant-ui/store": "^0.3.14",
     "assistant-cloud": "*",
-    "assistant-stream": "^0.3.42"
+    "assistant-stream": "^0.3.44"
   },
   "peerDependencies": {
-    "@langchain/react": "^1.0.2",
+    "@langchain/langgraph-sdk": "^1.9.20",
+    "@langchain/react": "^1.0.20",
     "@types/react": "*",
     "react": "^18 || ^19"
   },
   "peerDependenciesMeta": {
+    "@langchain/langgraph-sdk": {
+      "optional": true
+    },
     "@langchain/react": {
       "optional": true
     },
@@ -59,15 +63,15 @@
     }
   },
   "devDependencies": {
-    "@assistant-ui/x-buildutils": "0.0.26",
-    "@langchain/core": "^1.2.9",
-    "@langchain/langgraph-sdk": "^1.10.2",
-    "@langchain/react": "^1.0.35",
+    "@assistant-ui/x-buildutils": "0.0.28",
+    "@langchain/core": "^1.2.11",
+    "@langchain/langgraph-sdk": "^1.11.0",
+    "@langchain/react": "^1.1.0",
     "@testing-library/react": "^16.3.3",
-    "@types/react": "^19.2.18",
+    "@types/react": "^19.3.0",
     "jsdom": "^30.0.1",
-    "react": "^19.2.8",
-    "vitest": "^4.1.11"
+    "react": "^19.3.0",
+    "vitest": "^5.0.1"
   },
   "publishConfig": {
     "access": "public",
@@ -85,6 +89,7 @@
   "scripts": {
     "build": "aui-build",
     "test": "vitest run",
+    "typecheck": "tsc --noEmit",
     "test:watch": "vitest"
   }
 }

package/src/attachSubagentTranscripts.test.ts

@@ -0,0 +1,198 @@
+import {
+  bindExternalStoreMessage,
+  getExternalStoreMessages,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+import { describe, expect, it } from "vitest";
+import {
+  attachSubagentTranscripts,
+  createAttachMemo,
+  type SubagentTranscript,
+} from "./attachSubagentTranscripts";
+
+const assistantMessage = (id: string, toolCallId?: string): ThreadMessage =>
+  ({
+    id,
+    createdAt: new Date(),
+    role: "assistant",
+    status: { type: "complete", reason: "unknown" },
+    content: toolCallId
+      ? [
+          {
+            type: "tool-call",
+            toolCallId,
+            toolName: "task",
+            args: {},
+            argsText: "{}",
+          },
+        ]
+      : [{ type: "text", text: id }],
+    metadata: {
+      unstable_state: null,
+      unstable_annotations: [],
+      unstable_data: [],
+      steps: [],
+      custom: {},
+    },
+  }) as ThreadMessage;
+
+const transcript = (
+  id: string,
+  timing: SubagentTranscript["timing"] = {
+    startedAt: 1_000,
+    completedAt: 3_500,
+  },
+): SubagentTranscript => ({ messages: [assistantMessage(id)], timing });
+
+describe("attachSubagentTranscripts", () => {
+  it("returns the input array when no tool call has a transcript", () => {
+    const messages = [assistantMessage("message", "unknown")];
+
+    expect(
+      attachSubagentTranscripts(
+        messages,
+        new Map([["task", transcript("transcript")]]),
+        createAttachMemo(),
+      ),
+    ).toBe(messages);
+  });
+
+  it("attaches a transcript and preserves untouched message identities", () => {
+    const matching = assistantMessage("matching", "task");
+    const untouched = assistantMessage("untouched");
+    const messages = [matching, untouched];
+    const childTranscript = transcript("child");
+
+    const attached = attachSubagentTranscripts(
+      messages,
+      new Map([["task", childTranscript]]),
+      createAttachMemo(),
+    );
+
+    expect(attached).not.toBe(messages);
+    expect(attached[0]).not.toBe(matching);
+    expect(attached[1]).toBe(untouched);
+    const part = attached[0]!.content[0]!;
+    expect(part.type).toBe("tool-call");
+    if (part.type !== "tool-call") throw new Error("expected a tool call");
+    expect(part.messages).toBe(childTranscript.messages);
+    expect(part.timing).toBe(childTranscript.timing);
+  });
+
+  it("reuses the attached array when transcript identities are unchanged", () => {
+    const messages = [assistantMessage("matching", "task")];
+    const childTranscript = transcript("child");
+    const memo = createAttachMemo();
+
+    const first = attachSubagentTranscripts(
+      messages,
+      new Map([["task", childTranscript]]),
+      memo,
+    );
+    const second = attachSubagentTranscripts(
+      messages,
+      new Map([["task", childTranscript]]),
+      memo,
+    );
+
+    expect(second).toBe(first);
+  });
+
+  it("rebuilds only messages whose transcript changed", () => {
+    const messages = [
+      assistantMessage("first", "task-one"),
+      assistantMessage("second", "task-two"),
+    ];
+    const firstTranscript = transcript("first-transcript");
+    const secondTranscript = transcript("second-transcript");
+    const memo = createAttachMemo();
+    const first = attachSubagentTranscripts(
+      messages,
+      new Map([
+        ["task-one", firstTranscript],
+        ["task-two", secondTranscript],
+      ]),
+      memo,
+    );
+    const next = attachSubagentTranscripts(
+      messages,
+      new Map([
+        ["task-one", transcript("updated-transcript")],
+        ["task-two", secondTranscript],
+      ]),
+      memo,
+    );
+
+    expect(next).not.toBe(first);
+    expect(next[0]).not.toBe(first[0]);
+    expect(next[1]).toBe(first[1]);
+  });
+
+  it("rebuilds when the same transcript moves to another tool call", () => {
+    const first = assistantMessage("f", "call-a").content[0]!;
+    const second = assistantMessage("s", "call-b").content[0]!;
+    const message = {
+      ...assistantMessage("pair"),
+      content: [first, second],
+    } as ThreadMessage;
+    const nested = transcript("nested");
+    const memo = createAttachMemo();
+
+    const [before] = attachSubagentTranscripts(
+      [message],
+      new Map([["call-a", nested]]),
+      memo,
+    );
+    const [after] = attachSubagentTranscripts(
+      [message],
+      new Map([["call-b", nested]]),
+      memo,
+    );
+
+    expect(before?.content[0]).toMatchObject({ messages: nested.messages });
+    expect(after).not.toBe(before);
+    expect(after?.content[0]).toBe(first);
+    expect(after?.content[1]).toMatchObject({ messages: nested.messages });
+  });
+
+  it("keeps untouched parts of a partially matched message by reference", () => {
+    const text = { type: "text" as const, text: "hello" };
+    const matched = assistantMessage("m", "call-matched").content[0]!;
+    const unmatched = assistantMessage("u", "call-unmatched").content[0]!;
+    const message = {
+      ...assistantMessage("mixed"),
+      content: [text, matched, unmatched],
+    } as ThreadMessage;
+    const nested = transcript("nested");
+
+    const [attached] = attachSubagentTranscripts(
+      [message],
+      new Map([["call-matched", nested]]),
+      createAttachMemo(),
+    );
+
+    expect(attached?.content[0]).toBe(text);
+    expect(attached?.content[1]).toMatchObject({ messages: nested.messages });
+    expect(attached?.content[2]).toBe(unmatched);
+  });
+
+  it("preserves message and array external-store bindings", () => {
+    const message = assistantMessage("matching", "task");
+    const messages = [message];
+    const originalMessage = { id: "original-message" };
+    const originalMessages = [{ id: "original-messages" }];
+    bindExternalStoreMessage(message, originalMessage);
+    bindExternalStoreMessage(messages, originalMessages);
+
+    const attached = attachSubagentTranscripts(
+      messages,
+      new Map([["task", transcript("child")]]),
+      createAttachMemo(),
+    );
+
+    expect(getExternalStoreMessages(attached[0]!)).toEqual([originalMessage]);
+    expect(getExternalStoreMessages({ messages: attached })).toBe(
+      originalMessages,
+    );
+  });
+});

package/src/attachSubagentTranscripts.ts

@@ -0,0 +1,91 @@
+import {
+  bindExternalStoreMessage,
+  getExternalStoreMessages,
+  type ThreadMessage,
+  type ToolCallTiming,
+} from "@assistant-ui/core";
+
+export type SubagentTranscript = {
+  readonly messages: readonly ThreadMessage[];
+  /**
+   * Only present when the client watched the task run: the discovery
+   * timestamps are stamped when the task call is first seen, so a thread
+   * seeded from a checkpoint would report a duration of zero.
+   */
+  readonly timing?: ToolCallTiming;
+};
+
+type AttachedTranscript = readonly [string, SubagentTranscript];
+
+type AttachedMessage = {
+  transcripts: readonly AttachedTranscript[];
+  attached: ThreadMessage;
+};
+
+export type AttachMemo = {
+  messages: WeakMap<ThreadMessage, AttachedMessage>;
+  results: WeakMap<readonly ThreadMessage[], readonly ThreadMessage[]>;
+};
+
+export const createAttachMemo = (): AttachMemo => ({
+  messages: new WeakMap(),
+  results: new WeakMap(),
+});
+
+const sameTranscripts = (
+  a: readonly AttachedTranscript[],
+  b: readonly AttachedTranscript[],
+) =>
+  a.length === b.length &&
+  a.every(
+    ([toolCallId, transcript], index) =>
+      toolCallId === b[index]![0] && transcript === b[index]![1],
+  );
+
+export const attachSubagentTranscripts = (
+  messages: readonly ThreadMessage[],
+  transcripts: ReadonlyMap<string, SubagentTranscript>,
+  memo: AttachMemo,
+): readonly ThreadMessage[] => {
+  let attachedAny = false;
+  const next = messages.map((message) => {
+    const attachedTranscripts = message.content.flatMap((part) =>
+      part.type === "tool-call" && transcripts.has(part.toolCallId)
+        ? [[part.toolCallId, transcripts.get(part.toolCallId)!] as const]
+        : [],
+    );
+    if (attachedTranscripts.length === 0) return message;
+    attachedAny = true;
+
+    const cached = memo.messages.get(message);
+    if (cached && sameTranscripts(cached.transcripts, attachedTranscripts))
+      return cached.attached;
+
+    const attached = {
+      ...message,
+      content: message.content.map((part) => {
+        const transcript =
+          part.type === "tool-call"
+            ? transcripts.get(part.toolCallId)
+            : undefined;
+        return transcript
+          ? {
+              ...part,
+              messages: transcript.messages,
+              ...(transcript.timing && { timing: transcript.timing }),
+            }
+          : part;
+      }),
+    } as ThreadMessage;
+    memo.messages.set(message, { transcripts: attachedTranscripts, attached });
+    return attached;
+  });
+
+  if (!attachedAny) return messages;
+  const previous = memo.results.get(messages);
+  if (previous && previous.every((message, index) => message === next[index]))
+    return previous;
+  bindExternalStoreMessage(next, getExternalStoreMessages({ messages }));
+  memo.results.set(messages, next);
+  return next;
+};

package/src/sdkIdentity.ts

@@ -0,0 +1,9 @@
+import type { SdkIdentity } from "assistant-cloud";
+
+export const LANGCHAIN_SDK: SdkIdentity = {
+  name: "@assistant-ui/react-langchain",
+  version:
+    typeof __AUI_PACKAGE_VERSION__ === "string"
+      ? __AUI_PACKAGE_VERSION__
+      : "0.0.0",
+};

package/src/uiMessages.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it } from "vitest";
 import {
   applyUIUpdate,
+  createUIFoldMemo,
   extractUIUpdate,
   foldUIUpdates,
   isUIUpdate,
@@ -187,4 +188,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,9 @@
+import type { Channel } from "@langchain/react";
 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,12 +66,40 @@ 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;
 };