@assistant-ui/react-langchain 0.0.31 → 0.0.33

package/src/useStreamRuntime.test.tsx

@@ -1015,6 +1015,8 @@ describe("useStreamRuntime subagent transcripts", () => {
           status: "running",
           parentId: null,
           depth: 1,
+          startedAt: new Date(1_000),
+          completedAt: null,
         },
       ],
     ]);
@@ -1069,4 +1071,214 @@ describe("useStreamRuntime subagent transcripts", () => {
 
     expect(nestedTranscript()).toBe(rendered);
   });
+
+  it("keeps messages and transcripts across equal copies of the UI state", 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 uiState = (points: number[]) => [
+      {
+        type: "ui",
+        id: "ui-root",
+        name: "chart",
+        props: { points },
+        metadata: { message_id: "root-ai" },
+      },
+      {
+        type: "ui",
+        id: "ui-nested",
+        name: "chart",
+        props: { points },
+        metadata: { message_id: "nested-ai" },
+      },
+    ];
+    stream.values = { ui: uiState([1, 2]) };
+    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 [human, ai] = auiResult.current.thread.getState().messages;
+    expect(ai?.content).toMatchObject([
+      { type: "tool-call", toolCallId: "task-one" },
+      { type: "data", name: "chart", data: { points: [1, 2] } },
+    ]);
+    const rendered = nestedTranscript();
+
+    for (let i = 0; i < 3; i++) {
+      stream.values = { ui: uiState([1, 2]) };
+      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);
+
+    stream.values = { ui: uiState([1, 2, 3]) };
+    await act(async () => {
+      rerender();
+    });
+
+    expect(
+      auiResult.current.thread.getState().messages[1]?.content,
+    ).toMatchObject([
+      { type: "tool-call", toolCallId: "task-one" },
+      { type: "data", name: "chart", data: { points: [1, 2, 3] } },
+    ]);
+    expect(nestedTranscript()?.[0]?.content).toMatchObject([
+      { type: "text", text: "nested answer" },
+      { type: "data", name: "chart", data: { points: [1, 2, 3] } },
+    ]);
+  });
+
+  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

@@ -1,4 +1,4 @@
-/// <reference types="@assistant-ui/core/store" />
+/// <reference types="@assistant-ui/core/store" preserve="true" />
 "use client";
 
 import {
@@ -19,6 +19,7 @@ import type { ThreadMessage } from "@assistant-ui/core";
 import {
   createCloudThreadListAdapterCreateFallback,
   createToolCallCancellationStub,
+  getThreadMessageText,
   scanPendingToolCalls,
 } from "@assistant-ui/core/internal";
 import {
@@ -29,14 +30,13 @@ 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,
   UIMessage,
   UseStreamRuntimeOptions,
 } from "./types";
-import { groupUIMessagesByParent } from "./converter";
+import { getMessageModality, groupUIMessagesByParent } from "./converter";
 export { groupUIMessagesByParent } from "./converter";
 import {
   convertLangChainBaseMessage,
@@ -48,14 +48,19 @@ import {
   createAttachMemo,
 } from "./attachSubagentTranscripts";
 import { useSubagentTranscripts } from "./useSubagentTranscripts";
-import { foldUIUpdates, mergeUIMessages } from "./uiMessages";
+import {
+  createUIFoldMemo,
+  createUISnapshotMemo,
+  foldUIUpdates,
+  mergeUIMessages,
+  reconcileUISnapshot,
+  UI_CUSTOM_CHANNELS,
+} 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"];
-
 export const runConfigToSubmitOptions = (
   runConfig: AppendMessage["runConfig"],
 ) =>
@@ -92,6 +97,15 @@ const toStagedHumanMessage = (
   content: getMessageContent(msg),
 });
 
+const toStagedMessageInput = (message: LangChainBaseMessage) => ({
+  id: message.id,
+  type: getMessageType(message) === "ai" ? ("ai" as const) : ("human" as const),
+  content: message.content,
+  ...(message.additional_kwargs && {
+    additional_kwargs: message.additional_kwargs,
+  }),
+});
+
 const humanContentText = (content: LangChainBaseMessage["content"]) => {
   if (typeof content === "string") return content;
   if (!Array.isArray(content)) return "";
@@ -164,12 +178,17 @@ const useStreamThreadRuntime = (
   );
   const effectiveIsRunning = stream.isLoading || hasExecutingTools;
 
-  const uiStateValue = stream.values[uiStateKey];
+  const [uiSnapshotMemo] = useState(createUISnapshotMemo);
+  const uiStateValue = reconcileUISnapshot(
+    stream.values[uiStateKey],
+    uiSnapshotMemo,
+  );
 
   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(
@@ -302,6 +321,7 @@ const useStreamThreadRuntime = (
         runConfig: AppendMessage["runConfig"];
         reconcileOnEcho: boolean;
         baseMessageCount: number;
+        transcriptStatus?: "unsent" | "sent";
       }
     >(),
   );
@@ -357,19 +377,121 @@ const useStreamThreadRuntime = (
   }, [stream.messages]);
 
   const getStagedRun = (parentId: string | null) => {
-    if (!parentId || !stagedMessagesRef.current.has(parentId)) return null;
+    const parent = parentId
+      ? stagedMessagesRef.current.get(parentId)
+      : undefined;
+    if (!parent || parent.transcriptStatus === "sent") return null;
 
     const staged: LangChainBaseMessage[] = [];
     for (const message of visibleMessagesRef.current) {
-      if (message.id && stagedMessagesRef.current.has(message.id)) {
-        staged.push(stagedMessagesRef.current.get(message.id)!.message);
+      const entry = message.id
+        ? stagedMessagesRef.current.get(message.id)
+        : undefined;
+      if (entry && entry.transcriptStatus !== "sent") {
+        staged.push(entry.message);
       }
       if (message.id === parentId) break;
     }
 
     return {
       messages: staged,
-      runConfig: stagedMessagesRef.current.get(parentId)!.runConfig,
+      runConfig: parent.runConfig,
+    };
+  };
+
+  const appendVoiceTranscript = (message: ThreadMessage) => {
+    const transcript = {
+      id: message.id,
+      _getType: () => (message.role === "assistant" ? "ai" : "human"),
+      content: getThreadMessageText(message),
+      ...(message.metadata.modality && {
+        additional_kwargs: { modality: message.metadata.modality },
+      }),
+    };
+    stagedMessagesRef.current.set(transcript.id, {
+      message: transcript,
+      runConfig: undefined,
+      reconcileOnEcho: false,
+      baseMessageCount: streamRef.current.messages.length,
+      transcriptStatus: "unsent",
+    });
+    const nextMessages = [...visibleMessagesRef.current, transcript];
+    visibleMessagesRef.current = nextMessages;
+    setStagedMessages(nextMessages);
+  };
+
+  const getUnsentTranscripts = () =>
+    visibleMessagesRef.current.filter(
+      (message) =>
+        message.id !== undefined &&
+        stagedMessagesRef.current.get(message.id)?.transcriptStatus ===
+          "unsent",
+    );
+
+  const setTranscriptStatus = (
+    messages: readonly LangChainBaseMessage[],
+    status: "unsent" | "sent",
+  ) => {
+    for (const message of messages) {
+      const staged = message.id
+        ? stagedMessagesRef.current.get(message.id)
+        : undefined;
+      if (staged?.transcriptStatus) staged.transcriptStatus = status;
+    }
+  };
+
+  // Reserving before the submit keeps an overlapping submit from carrying the
+  // same transcript; a failed submit hands it back to the next run.
+  const submitCarryingTranscripts = async (
+    transcripts: readonly LangChainBaseMessage[],
+    submit: () => Promise<void>,
+  ) => {
+    setTranscriptStatus(transcripts, "sent");
+    try {
+      await submit();
+    } catch (error) {
+      setTranscriptStatus(transcripts, "unsent");
+      throw error;
+    }
+  };
+
+  const dropTranscripts = (messages: readonly LangChainBaseMessage[]) => {
+    for (const message of messages) {
+      if (
+        message.id &&
+        stagedMessagesRef.current.get(message.id)?.transcriptStatus
+      )
+        removeStagedMessage(message.id);
+    }
+  };
+
+  const isTranscriptMessage = (message: LangChainBaseMessage) => {
+    const staged = message.id
+      ? stagedMessagesRef.current.get(message.id)
+      : undefined;
+    if (staged?.transcriptStatus === "unsent") return true;
+    const type = getMessageType(message);
+    return (
+      (type === "human" || type === "ai") &&
+      getMessageModality(message.additional_kwargs) !== undefined
+    );
+  };
+
+  // A transcript reaches the graph in the same input as the message after it, so
+  // no checkpoint ends at one; a fork starts before the trailing transcripts.
+  const planForkTranscripts = (parentId: string | null) => {
+    const visible = visibleMessagesRef.current;
+    const parentIndex =
+      parentId == null ? -1 : visible.findIndex((m) => m.id === parentId);
+    if (parentId != null && parentIndex === -1)
+      return { forkParentId: parentId, transcripts: [], truncated: [] };
+
+    let start = parentIndex + 1;
+    while (start > 0 && isTranscriptMessage(visible[start - 1]!)) start--;
+    return {
+      forkParentId: start > 0 ? (visible[start - 1]!.id ?? null) : null,
+      transcripts: visible.slice(start, parentIndex + 1),
+      truncated: visible.slice(parentIndex + 1),
     };
   };
 
@@ -459,27 +581,32 @@ const useStreamThreadRuntime = (
       // longer holds appends on that barrier.
       try {
         const { externalId } = await aui.threadListItem.initialize();
-        await streamRef.current.submit(
-          {
-            [messagesKey]: [
-              ...cancellations,
-              {
-                id: stagedMessageId,
-                type: "human",
-                content,
-              },
-            ],
-          },
-          {
-            ...runConfigToSubmitOptions(msg.runConfig),
-            ...(externalId != null ? { threadId: externalId } : {}),
-          },
+        const transcripts = getUnsentTranscripts();
+        await submitCarryingTranscripts(transcripts, () =>
+          streamRef.current.submit(
+            {
+              [messagesKey]: [
+                ...cancellations,
+                ...transcripts.map(toStagedMessageInput),
+                {
+                  id: stagedMessageId,
+                  type: "human",
+                  content,
+                },
+              ],
+            },
+            {
+              ...runConfigToSubmitOptions(msg.runConfig),
+              ...(externalId != null ? { threadId: externalId } : {}),
+            },
+          ),
         );
       } catch (error) {
         removeStagedMessage(stagedMessageId);
         throw error;
       }
     },
+    onVoiceTranscript: appendVoiceTranscript,
     onAddToolResult: async ({
       messageId,
       toolCallId,
@@ -510,9 +637,18 @@ const useStreamThreadRuntime = (
     onReload: async (parentId, config) => {
       const stagedRun = getStagedRun(parentId);
       if (stagedRun) {
+        if (
+          config.sourceId &&
+          stagedMessagesRef.current.get(config.sourceId)?.transcriptStatus
+        )
+          removeStagedMessage(config.sourceId);
         const promotedIds = new Set<string>();
         for (const message of stagedRun.messages) {
-          if (!message.id) continue;
+          if (
+            !message.id ||
+            stagedMessagesRef.current.get(message.id)?.transcriptStatus
+          )
+            continue;
           promotedIds.add(message.id);
           stagedMessagesRef.current.delete(message.id);
         }
@@ -528,15 +664,13 @@ const useStreamThreadRuntime = (
         }
         const runConfig = config.runConfig ?? stagedRun.runConfig;
         setActiveRunConfig(runConfig);
-        await stream.submit(
-          {
-            [messagesKey]: stagedRun.messages.map((message) => ({
-              id: message.id,
-              type: "human",
-              content: message.content,
-            })),
-          },
-          runConfigToSubmitOptions(runConfig),
+        await submitCarryingTranscripts(stagedRun.messages, () =>
+          stream.submit(
+            {
+              [messagesKey]: stagedRun.messages.map(toStagedMessageInput),
+            },
+            runConfigToSubmitOptions(runConfig),
+          ),
         );
         return;
       }
@@ -544,21 +678,30 @@ const useStreamThreadRuntime = (
       const threadId = externalId;
       if (!threadId || parentId == null) return;
       const s = streamRef.current;
+      const fork = planForkTranscripts(parentId);
       const checkpointId = await resolveForkCheckpoint(
         s.client,
         threadId,
         s.messages as readonly LangChainBaseMessage[],
-        parentId,
+        fork.forkParentId,
         config.sourceId,
         s[STREAM_CONTROLLER]?.messageMetadataStore?.getSnapshot?.(),
         messagesKey,
       );
       if (!checkpointId) return;
+      dropTranscripts(fork.truncated);
       setActiveRunConfig(config.runConfig);
-      await s.submit(null, {
-        forkFrom: checkpointId,
-        ...runConfigToSubmitOptions(config.runConfig),
-      });
+      await submitCarryingTranscripts(fork.transcripts, () =>
+        s.submit(
+          fork.transcripts.length > 0
+            ? { [messagesKey]: fork.transcripts.map(toStagedMessageInput) }
+            : null,
+          {
+            forkFrom: checkpointId,
+            ...runConfigToSubmitOptions(config.runConfig),
+          },
+        ),
+      );
     },
     onEdit: async (message) => {
       if (!(message.startRun ?? message.role === "user")) {
@@ -583,24 +726,33 @@ const useStreamThreadRuntime = (
       const threadId = externalId;
       if (!threadId) return;
       const s = streamRef.current;
+      const fork = planForkTranscripts(message.parentId);
       const checkpointId = await resolveForkCheckpoint(
         s.client,
         threadId,
         s.messages as readonly LangChainBaseMessage[],
-        message.parentId,
+        fork.forkParentId,
         message.sourceId,
         s[STREAM_CONTROLLER]?.messageMetadataStore?.getSnapshot?.(),
         messagesKey,
       );
       if (!checkpointId) return;
+      dropTranscripts(fork.truncated);
       const content = getMessageContent(message);
       setActiveRunConfig(message.runConfig);
-      await s.submit(
-        { [messagesKey]: [{ type: "human", content }] },
-        {
-          forkFrom: checkpointId,
-          ...runConfigToSubmitOptions(message.runConfig),
-        },
+      await submitCarryingTranscripts(fork.transcripts, () =>
+        s.submit(
+          {
+            [messagesKey]: [
+              ...fork.transcripts.map(toStagedMessageInput),
+              { type: "human", content },
+            ],
+          },
+          {
+            forkFrom: checkpointId,
+            ...runConfigToSubmitOptions(message.runConfig),
+          },
+        ),
       );
     },
     onCancel: