@assistant-ui/react-langchain 0.0.31 → 0.0.32

package/src/attachSubagentTranscripts.test.ts

@@ -7,6 +7,7 @@ import { describe, expect, it } from "vitest";
 import {
   attachSubagentTranscripts,
   createAttachMemo,
+  type SubagentTranscript,
 } from "./attachSubagentTranscripts";
 
 const assistantMessage = (id: string, toolCallId?: string): ThreadMessage =>
@@ -35,7 +36,13 @@ const assistantMessage = (id: string, toolCallId?: string): ThreadMessage =>
     },
   }) as ThreadMessage;
 
-const transcript = (id: string) => [assistantMessage(id)];
+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", () => {
@@ -68,7 +75,8 @@ describe("attachSubagentTranscripts", () => {
     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);
+    expect(part.messages).toBe(childTranscript.messages);
+    expect(part.timing).toBe(childTranscript.timing);
   });
 
   it("reuses the attached array when transcript identities are unchanged", () => {
@@ -141,10 +149,10 @@ describe("attachSubagentTranscripts", () => {
       memo,
     );
 
-    expect(before?.content[0]).toMatchObject({ messages: nested });
+    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 });
+    expect(after?.content[1]).toMatchObject({ messages: nested.messages });
   });
 
   it("keeps untouched parts of a partially matched message by reference", () => {
@@ -164,7 +172,7 @@ describe("attachSubagentTranscripts", () => {
     );
 
     expect(attached?.content[0]).toBe(text);
-    expect(attached?.content[1]).toMatchObject({ messages: nested });
+    expect(attached?.content[1]).toMatchObject({ messages: nested.messages });
     expect(attached?.content[2]).toBe(unmatched);
   });
 

package/src/attachSubagentTranscripts.ts

@@ -2,9 +2,20 @@ import {
   bindExternalStoreMessage,
   getExternalStoreMessages,
   type ThreadMessage,
+  type ToolCallTiming,
 } from "@assistant-ui/core";
 
-type AttachedTranscript = readonly [string, readonly ThreadMessage[]];
+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[];
@@ -33,7 +44,7 @@ const sameTranscripts = (
 
 export const attachSubagentTranscripts = (
   messages: readonly ThreadMessage[],
-  transcripts: ReadonlyMap<string, readonly ThreadMessage[]>,
+  transcripts: ReadonlyMap<string, SubagentTranscript>,
   memo: AttachMemo,
 ): readonly ThreadMessage[] => {
   let attachedAny = false;
@@ -52,11 +63,19 @@ export const attachSubagentTranscripts = (
 
     const attached = {
       ...message,
-      content: message.content.map((part) =>
-        part.type === "tool-call" && transcripts.has(part.toolCallId)
-          ? { ...part, messages: transcripts.get(part.toolCallId)! }
-          : part,
-      ),
+      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;

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;
 };
 

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,114 @@ describe("useStreamRuntime subagent transcripts", () => {
 
     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,
@@ -48,14 +47,17 @@ import {
   createAttachMemo,
 } from "./attachSubagentTranscripts";
 import { useSubagentTranscripts } from "./useSubagentTranscripts";
-import { foldUIUpdates, mergeUIMessages } from "./uiMessages";
+import {
+  createUIFoldMemo,
+  foldUIUpdates,
+  mergeUIMessages,
+  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"],
 ) =>
@@ -167,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(