@assistant-ui/react-langchain 0.0.30 → 0.0.31

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react-langchain",
-  "version": "0.0.30",
+  "version": "0.0.31",
   "description": "LangChain useStream adapter for assistant-ui",
   "keywords": [
     "langchain",
@@ -40,17 +40,21 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@assistant-ui/core": "^0.3.18",
+    "@assistant-ui/core": "^0.3.19",
     "@assistant-ui/store": "^0.3.13",
     "assistant-cloud": "*",
-    "assistant-stream": "^0.3.42"
+    "assistant-stream": "^0.3.43"
   },
   "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,7 +63,7 @@
     }
   },
   "devDependencies": {
-    "@assistant-ui/x-buildutils": "0.0.26",
+    "@assistant-ui/x-buildutils": "0.0.27",
     "@langchain/core": "^1.2.9",
     "@langchain/langgraph-sdk": "^1.10.2",
     "@langchain/react": "^1.0.35",
@@ -67,7 +71,7 @@
     "@types/react": "^19.2.18",
     "jsdom": "^30.0.1",
     "react": "^19.2.8",
-    "vitest": "^4.1.11"
+    "vitest": "^5.0.0"
   },
   "publishConfig": {
     "access": "public",

package/src/attachSubagentTranscripts.test.ts

@@ -0,0 +1,190 @@
+import {
+  bindExternalStoreMessage,
+  getExternalStoreMessages,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+import { describe, expect, it } from "vitest";
+import {
+  attachSubagentTranscripts,
+  createAttachMemo,
+} 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) => [assistantMessage(id)];
+
+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);
+  });
+
+  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 });
+    expect(after).not.toBe(before);
+    expect(after?.content[0]).toBe(first);
+    expect(after?.content[1]).toMatchObject({ messages: nested });
+  });
+
+  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 });
+    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,72 @@
+import {
+  bindExternalStoreMessage,
+  getExternalStoreMessages,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+
+type AttachedTranscript = readonly [string, readonly ThreadMessage[]];
+
+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, readonly ThreadMessage[]>,
+  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) =>
+        part.type === "tool-call" && transcripts.has(part.toolCallId)
+          ? { ...part, messages: transcripts.get(part.toolCallId)! }
+          : 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/useStreamRuntime.test.tsx

@@ -1,7 +1,7 @@
 // @vitest-environment jsdom
 
 import { act, render, renderHook, waitFor } from "@testing-library/react";
-import { describe, expect, it, vi } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
 import type {
   AssistantRuntime,
@@ -20,7 +20,7 @@ import {
 } from "./hooks";
 
 const { mockUseChannel, mockUseStream, streamController } = vi.hoisted(() => ({
-  mockUseChannel: vi.fn(() => []),
+  mockUseChannel: vi.fn((): unknown[] => []),
   mockUseStream: vi.fn(),
   streamController: Symbol("STREAM_CONTROLLER"),
 }));
@@ -40,7 +40,7 @@ type MockStream = {
   values: Record<string, unknown>;
   interrupts: unknown[];
   toolCalls: unknown[];
-  subagents: unknown[];
+  subagents: ReadonlyMap<string, unknown>;
   subgraphs: unknown[];
   error: unknown;
   submit: ReturnType<typeof vi.fn>;
@@ -53,6 +53,10 @@ type MockStream = {
     messageMetadataStore: {
       getSnapshot: ReturnType<typeof vi.fn>;
     };
+    resolveSubagentNamespace: ReturnType<typeof vi.fn>;
+    registry: {
+      acquire: ReturnType<typeof vi.fn>;
+    };
   };
 };
 
@@ -75,7 +79,7 @@ const createMockStream = (
   values: {},
   interrupts: [],
   toolCalls: [],
-  subagents: [],
+  subagents: new Map(),
   subgraphs: [],
   error: undefined,
   submit: vi.fn(async () => {}),
@@ -88,6 +92,10 @@ const createMockStream = (
     messageMetadataStore: {
       getSnapshot: vi.fn(),
     },
+    resolveSubagentNamespace: vi.fn(async () => {}),
+    registry: {
+      acquire: vi.fn(),
+    },
   },
 });
 
@@ -177,7 +185,11 @@ describe("useStreamRuntime thread options", () => {
         unstable_threadListAdapter: threadListAdapter,
       } as never);
       capture.runtime = runtime;
-      return <AssistantRuntimeProvider runtime={runtime} />;
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          {null}
+        </AssistantRuntimeProvider>
+      );
     };
 
     const view = render(<TestRuntime />);
@@ -286,7 +298,11 @@ describe("useStreamRuntime thread options", () => {
     const TestRuntime = () => {
       const runtime = useStreamRuntime({ apiUrl: "/api" } as never);
       capture.runtime = runtime;
-      return <AssistantRuntimeProvider runtime={runtime} />;
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          {null}
+        </AssistantRuntimeProvider>
+      );
     };
     const view = render(<TestRuntime />);
     await waitFor(() => expect(capture.runtime).not.toBeNull());
@@ -721,7 +737,9 @@ describe("useStreamRuntime run configuration", () => {
     });
 
     await act(async () => {
-      await capture.send!([{ type: "human", content: "next" }]);
+      await capture.send!([
+        { type: "human", content: "next" } as unknown as LangChainBaseMessage,
+      ]);
     });
 
     expect(stream.submit).toHaveBeenLastCalledWith(
@@ -971,3 +989,84 @@ describe("useStreamRuntime committed refs", () => {
     host.unmount();
   });
 });
+
+describe("useStreamRuntime subagent transcripts", () => {
+  afterEach(() => {
+    mockUseChannel.mockReset();
+  });
+
+  it("renders live UI messages inside the transcript of the subagent that pushed them", 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,
+        },
+      ],
+    ]);
+    stream[streamController]!.registry.acquire.mockReturnValue({
+      store: { getSnapshot: () => transcript, subscribe: () => () => {} },
+      release: vi.fn(),
+    });
+    const uiEvent = {
+      method: "custom",
+      params: {
+        namespace: ["tools:task-one"],
+        data: {
+          type: "ui",
+          id: "ui-1",
+          name: "chart",
+          props: { points: [1, 2] },
+          metadata: { message_id: "nested-ai" },
+        },
+      },
+    };
+    mockUseChannel.mockReturnValue([uiEvent]);
+    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 rendered = nestedTranscript();
+
+    mockUseChannel.mockReturnValue([
+      uiEvent,
+      {
+        method: "custom",
+        params: { namespace: [], data: { name: "progress", payload: 1 } },
+      },
+    ]);
+    await act(async () => {
+      rerender();
+    });
+
+    expect(nestedTranscript()).toBe(rendered);
+  });
+});

package/src/useStreamRuntime.ts

@@ -43,10 +43,16 @@ import {
   getMessageContent,
   getMessageType,
 } from "./convertMessages";
+import {
+  attachSubagentTranscripts,
+  createAttachMemo,
+} from "./attachSubagentTranscripts";
+import { useSubagentTranscripts } from "./useSubagentTranscripts";
 import { foldUIUpdates, mergeUIMessages } from "./uiMessages";
 import { langChainExtras } from "./runtimeExtras";
 import { resolveForkCheckpoint } from "./resolveForkCheckpoint";
 import { useLangChainStreamingTiming } from "./streamingTiming";
+import { LANGCHAIN_SDK } from "./sdkIdentity";
 
 const UI_CUSTOM_CHANNELS: readonly Channel[] = ["custom"];
 
@@ -171,6 +177,11 @@ const useStreamThreadRuntime = (
     [liveUiMessages, uiStateValue],
   );
 
+  const uiMessagesByParent = useMemo(
+    () => groupUIMessagesByParent<UIMessage>(mergedUiMessages),
+    [mergedUiMessages],
+  );
+
   const visibleMessages =
     stagedMessages ?? (stream.messages as LangChainBaseMessage[]);
 
@@ -179,24 +190,33 @@ const useStreamThreadRuntime = (
     effectiveIsRunning,
   );
 
+  const subagentTranscripts = useSubagentTranscripts(
+    stream,
+    uiMessagesByParent,
+  );
+
   const convertWithUI = useMemo<
     useExternalMessageConverter.Callback<LangChainBaseMessage>
-  >(() => {
-    const uiMessagesByParent =
-      groupUIMessagesByParent<UIMessage>(mergedUiMessages);
-    return (message, metadata) =>
+  >(
+    () => (message, metadata) =>
       convertLangChainBaseMessage(message, {
         ...metadata,
         uiMessagesByParent,
         messageTiming,
-      });
-  }, [mergedUiMessages, messageTiming]);
+      }),
+    [uiMessagesByParent, messageTiming],
+  );
 
   const threadMessages = useExternalMessageConverter({
     callback: convertWithUI,
     messages: visibleMessages,
     isRunning: effectiveIsRunning,
   });
+  const [memo] = useState(createAttachMemo);
+  const messagesWithTranscripts = useMemo(
+    () => attachSubagentTranscripts(threadMessages, subagentTranscripts, memo),
+    [threadMessages, subagentTranscripts, memo],
+  );
 
   const streamRef = useRef(stream);
   useInsertionEffect(() => {
@@ -269,10 +289,10 @@ const useStreamThreadRuntime = (
     visibleMessagesRef.current = visibleMessages;
   }, [visibleMessages]);
 
-  const threadMessagesRef = useRef(threadMessages);
+  const threadMessagesRef = useRef(messagesWithTranscripts);
   useInsertionEffect(() => {
-    threadMessagesRef.current = threadMessages;
-  }, [threadMessages]);
+    threadMessagesRef.current = messagesWithTranscripts;
+  }, [messagesWithTranscripts]);
 
   const stagedMessagesRef = useRef(
     new Map<
@@ -412,7 +432,7 @@ const useStreamThreadRuntime = (
     ...pickExternalStoreSharedOptions(options),
     isRunning: stream.isLoading,
     isLoading: stream.isThreadLoading,
-    messages: threadMessages,
+    messages: messagesWithTranscripts,
     adapters,
     extras,
     unstable_enableToolInvocations: true,
@@ -631,6 +651,7 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
 
   const aui = useAui();
   const cloudAdapter = useCloudThreadListAdapter({
+    sdk: LANGCHAIN_SDK,
     cloud,
     create: createCloudThreadListAdapterCreateFallback(
       create,