@assistant-ui/react-langchain 0.0.29 → 0.0.31

package/src/converter.ts

@@ -61,7 +61,10 @@ export const convertLangChainContentBlock = (
   switch (type) {
     case "text":
     case "text_delta":
-      return { type: "text" as const, text: part.text };
+      return {
+        type: "text" as const,
+        text: typeof part.text === "string" ? part.text : "",
+      };
     case "image_url": {
       const image =
         typeof part.image_url === "string"
@@ -98,15 +101,13 @@ export const convertLangChainContentBlock = (
       };
     }
     case "thinking":
-      return { type: "reasoning" as const, text: part.thinking };
-    case "reasoning":
-      return {
-        type: "reasoning" as const,
-        text:
-          part.summary && part.summary.length > 0
-            ? part.summary.map((s) => s?.text ?? "").join("\n\n\n")
-            : (part.reasoning ?? ""),
-      };
+      return hasVisibleText(part.thinking)
+        ? { type: "reasoning" as const, text: part.thinking }
+        : null;
+    case "reasoning": {
+      const text = getReasoningText(part);
+      return text ? { type: "reasoning" as const, text } : null;
+    }
     case "tool_use":
     case "input_json_delta":
       return null;
@@ -262,13 +263,14 @@ export const getMessageContent = (msg: AppendMessage) => {
   return content;
 };
 
-const reasoningTextLength = (part: {
+const getReasoningText = (part: {
   readonly summary?: ReadonlyArray<{ readonly text?: string }>;
   readonly reasoning?: string;
-}): number => {
-  if (part.summary && part.summary.length > 0)
-    return part.summary.map((s) => s?.text ?? "").join("\n\n\n").length;
-  return part.reasoning?.length ?? 0;
+}): string => {
+  const summary = part.summary?.map((s) => s?.text ?? "").join("\n\n\n") ?? "";
+  if (hasVisibleText(summary)) return summary;
+  const reasoning = part.reasoning ?? "";
+  return hasVisibleText(reasoning) ? reasoning : "";
 };
 
 export const createLangChainStreamingTimingAccessors = <
@@ -299,16 +301,17 @@ export const createLangChainStreamingTimingAccessors = <
     if (!Array.isArray(content)) return 0;
     let len = 0;
     for (const part of content as readonly LangChainContentBlock[]) {
+      if (typeof part !== "object" || part === null) continue;
       switch (part.type) {
         case "text":
         case "text_delta":
           if (typeof part.text === "string") len += part.text.length;
           break;
         case "thinking":
-          if (typeof part.thinking === "string") len += part.thinking.length;
+          if (hasVisibleText(part.thinking)) len += part.thinking.length;
           break;
         case "reasoning":
-          len += reasoningTextLength(part);
+          len += getReasoningText(part).length;
           break;
       }
     }

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/streamingTiming.test.tsx

@@ -34,4 +34,54 @@ describe("useLangChainStreamingTiming", () => {
       Math.ceil("deduced".length / 4),
     );
   });
+
+  it("counts a thinking block only when the converter renders it", () => {
+    const tokenCount = (thinking: string) => {
+      const messages: LangChainBaseMessage[] = [
+        {
+          id: "msg-1",
+          _getType: () => "ai",
+          content: [{ type: "thinking", thinking }],
+        },
+      ];
+      const { result, rerender } = renderHook(
+        ({ msgs, running }) => useLangChainStreamingTiming(msgs, running),
+        { initialProps: { msgs: messages, running: true } },
+      );
+      act(() => {
+        rerender({ msgs: messages, running: false });
+      });
+      return result.current["msg-1"]?.tokenCount;
+    };
+
+    expect(tokenCount("deduced")).toBe(Math.ceil("deduced".length / 4));
+    expect(tokenCount("   ")).toBeUndefined();
+  });
+
+  it("counts the summary text the converter renders, not the reasoning it shadows", () => {
+    const messages: LangChainBaseMessage[] = [
+      {
+        id: "msg-1",
+        _getType: () => "ai",
+        content: [
+          {
+            type: "reasoning",
+            reasoning: "partial thinking",
+            summary: [{ type: "summary_text", text: "first summary" }],
+          },
+        ],
+      },
+    ];
+    const { result, rerender } = renderHook(
+      ({ running }) => useLangChainStreamingTiming(messages, running),
+      { initialProps: { running: true } },
+    );
+    act(() => {
+      rerender({ running: false });
+    });
+
+    expect(result.current["msg-1"]?.tokenCount).toBe(
+      Math.ceil("first summary".length / 4),
+    );
+  });
 });

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<
@@ -324,6 +344,9 @@ const useStreamThreadRuntime = (
     if (remainingStagedMessages.length === 0) {
       stagedBaseMessagesRef.current = null;
       visibleMessagesRef.current = baseMessages;
+      // Reconciling against the upstream stream mutates the staged refs above,
+      // which cannot happen during render.
+      // eslint-disable-next-line react-hooks/set-state-in-effect
       setStagedMessages(null);
       return;
     }
@@ -409,7 +432,7 @@ const useStreamThreadRuntime = (
     ...pickExternalStoreSharedOptions(options),
     isRunning: stream.isLoading,
     isLoading: stream.isThreadLoading,
-    messages: threadMessages,
+    messages: messagesWithTranscripts,
     adapters,
     extras,
     unstable_enableToolInvocations: true,
@@ -628,6 +651,7 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
 
   const aui = useAui();
   const cloudAdapter = useCloudThreadListAdapter({
+    sdk: LANGCHAIN_SDK,
     cloud,
     create: createCloudThreadListAdapterCreateFallback(
       create,