@assistant-ui/react-langchain 0.0.30 → 0.0.32

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,196 @@ 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,
+          startedAt: new Date(1_000),
+          completedAt: null,
+        },
+      ],
+    ]);
+    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);
+  });
+
+  it("keeps messages and transcripts when custom events carry no UI update", 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 chartEvent = (
+      props: Record<string, unknown>,
+      metadata: Record<string, unknown> = {},
+    ) => ({
+      method: "custom",
+      params: {
+        namespace: ["tools:task-one"],
+        data: {
+          type: "ui",
+          id: "ui-1",
+          name: "chart",
+          props,
+          metadata: { message_id: "nested-ai", ...metadata },
+        },
+      },
+    });
+    const events = [
+      chartEvent({ points: [1, 2] }),
+      chartEvent({ label: "first" }, { merge: true }),
+    ];
+    mockUseChannel.mockReturnValue(events);
+    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], label: "first" },
+        },
+      ]),
+    );
+    const [human, ai] = auiResult.current.thread.getState().messages;
+    const rendered = nestedTranscript();
+
+    const withProgress = [
+      ...events,
+      {
+        method: "custom",
+        params: { namespace: [], data: { name: "progress", payload: 1 } },
+      },
+    ];
+    mockUseChannel.mockReturnValue(withProgress);
+    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);
+
+    mockUseChannel.mockReturnValue([
+      ...withProgress,
+      chartEvent({ label: "second" }, { merge: true }),
+    ]);
+    await act(async () => {
+      rerender();
+    });
+
+    expect(nestedTranscript()?.[0]?.content).toMatchObject([
+      { type: "text", text: "nested answer" },
+      {
+        type: "data",
+        name: "chart",
+        data: { points: [1, 2], label: "second" },
+      },
+    ]);
+  });
+});

package/src/useStreamRuntime.ts

@@ -29,7 +29,6 @@ import {
 } from "@assistant-ui/core/react";
 import { useAui, useAuiState } from "@assistant-ui/store";
 import { STREAM_CONTROLLER, useChannel, useStream } from "@langchain/react";
-import type { Channel } from "@langchain/react";
 import type {
   LangChainBaseMessage,
   LangChainToolCall,
@@ -43,12 +42,21 @@ import {
   getMessageContent,
   getMessageType,
 } from "./convertMessages";
-import { foldUIUpdates, mergeUIMessages } from "./uiMessages";
+import {
+  attachSubagentTranscripts,
+  createAttachMemo,
+} from "./attachSubagentTranscripts";
+import { useSubagentTranscripts } from "./useSubagentTranscripts";
+import {
+  createUIFoldMemo,
+  foldUIUpdates,
+  mergeUIMessages,
+  UI_CUSTOM_CHANNELS,
+} from "./uiMessages";
 import { langChainExtras } from "./runtimeExtras";
 import { resolveForkCheckpoint } from "./resolveForkCheckpoint";
 import { useLangChainStreamingTiming } from "./streamingTiming";
-
-const UI_CUSTOM_CHANNELS: readonly Channel[] = ["custom"];
+import { LANGCHAIN_SDK } from "./sdkIdentity";
 
 export const runConfigToSubmitOptions = (
   runConfig: AppendMessage["runConfig"],
@@ -161,9 +169,10 @@ const useStreamThreadRuntime = (
   const uiStateValue = stream.values[uiStateKey];
 
   const customEvents = useChannel(stream, UI_CUSTOM_CHANNELS);
+  const [uiFoldMemo] = useState(createUIFoldMemo);
   const liveUiMessages = useMemo(
-    () => foldUIUpdates(customEvents),
-    [customEvents],
+    () => foldUIUpdates(customEvents, uiFoldMemo),
+    [customEvents, uiFoldMemo],
   );
 
   const mergedUiMessages = useMemo(
@@ -171,6 +180,11 @@ const useStreamThreadRuntime = (
     [liveUiMessages, uiStateValue],
   );
 
+  const uiMessagesByParent = useMemo(
+    () => groupUIMessagesByParent<UIMessage>(mergedUiMessages),
+    [mergedUiMessages],
+  );
+
   const visibleMessages =
     stagedMessages ?? (stream.messages as LangChainBaseMessage[]);
 
@@ -179,24 +193,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 +292,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 +435,7 @@ const useStreamThreadRuntime = (
     ...pickExternalStoreSharedOptions(options),
     isRunning: stream.isLoading,
     isLoading: stream.isThreadLoading,
-    messages: threadMessages,
+    messages: messagesWithTranscripts,
     adapters,
     extras,
     unstable_enableToolInvocations: true,
@@ -631,6 +654,7 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
 
   const aui = useAui();
   const cloudAdapter = useCloudThreadListAdapter({
+    sdk: LANGCHAIN_SDK,
     cloud,
     create: createCloudThreadListAdapterCreateFallback(
       create,