@assistant-ui/react-langchain 0.0.32 → 0.0.34

package/src/useSubagentTranscripts.fast-refresh.test.tsx

@@ -0,0 +1,152 @@
+// @vitest-environment jsdom
+
+import { act, useMemo } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+import type { UIMessage } from "./types";
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+const { streamController } = vi.hoisted(() => ({
+  streamController: Symbol("STREAM_CONTROLLER"),
+}));
+
+vi.mock("@langchain/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@langchain/react")>()),
+  STREAM_CONTROLLER: streamController,
+}));
+
+let renderer: RendererInternals | undefined;
+const fiberRoots = new Set<unknown>();
+vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+  supportsFiber: true,
+  inject: (internals: RendererInternals) => {
+    renderer = internals;
+    return 1;
+  },
+  onScheduleFiberRoot: () => {},
+  onCommitFiberRoot: (_id: number, root: unknown) => fiberRoots.add(root),
+  onCommitFiberUnmount: () => {},
+});
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+const { useSubagentTranscripts } = await import("./useSubagentTranscripts");
+
+afterEach(() => {
+  cleanup();
+  renderer?.setRefreshHandler(() => undefined);
+  fiberRoots.clear();
+});
+afterAll(() => vi.unstubAllGlobals());
+
+it("keeps a live subagent projection and namespace request through Fast Refresh, then disposes on unmount", async () => {
+  const listeners = new Set<() => void>();
+  let messages = [
+    { id: "message-one", _getType: () => "ai", content: "partial" },
+  ];
+  const store = {
+    getSnapshot: () => messages,
+    subscribe: (listener: () => void) => {
+      listeners.add(listener);
+      return () => listeners.delete(listener);
+    },
+  };
+  const release = vi.fn();
+  const acquire = vi.fn(() => ({ store, release }));
+  let settleFirst!: () => void;
+  const pending = new Promise<void>((resolve) => {
+    settleFirst = resolve;
+  });
+  const resolveSubagentNamespace = vi
+    .fn<() => Promise<void>>()
+    .mockReturnValueOnce(pending)
+    .mockResolvedValue(undefined);
+  const snapshot = (status: "running" | "complete") => ({
+    id: "task-one",
+    namespace: ["tools:task-one"],
+    status,
+    parentId: null,
+    depth: 1,
+    startedAt: new Date(1_000),
+    completedAt: status === "complete" ? new Date(4_000) : null,
+  });
+  const stream = {
+    subagents: new Map([["task-one", snapshot("running")]]),
+    [streamController]: {
+      registry: { acquire },
+      resolveSubagentNamespace,
+    },
+  };
+  const uiMaps = new Set<Map<string, UIMessage[]>>();
+  let transcript: ReturnType<typeof useSubagentTranscripts> | undefined;
+  let rendered: string | undefined;
+  const host = (name: string) => () => {
+    const uiMessagesByParent = useMemo(
+      () => new Map<string, UIMessage[]>(),
+      [],
+    );
+    uiMaps.add(uiMessagesByParent);
+    rendered = name;
+    transcript = useSubagentTranscripts(stream as never, uiMessagesByParent);
+    return null;
+  };
+  const Before = host("before");
+  const After = host("after");
+  const view = render(<Before />);
+
+  await waitFor(() => expect(transcript?.has("task-one")).toBe(true));
+  expect(resolveSubagentNamespace).toHaveBeenCalledOnce();
+  expect(listeners.size).toBe(1);
+  await act(async () => {
+    messages = [
+      { id: "message-one", _getType: () => "ai", content: "partial answer" },
+    ];
+    for (const listener of listeners) listener();
+  });
+
+  const family: Family = { current: After };
+  renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const root of fiberRoots) {
+      renderer!.scheduleRefresh(root, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+
+  expect(rendered).toBe("after");
+  expect(uiMaps.size).toBeGreaterThan(1);
+  expect(release).not.toHaveBeenCalled();
+  expect(acquire).toHaveBeenCalledOnce();
+  expect(listeners.size).toBe(1);
+  expect(resolveSubagentNamespace).toHaveBeenCalledOnce();
+
+  stream.subagents = new Map([["task-one", snapshot("complete")]]);
+  view.rerender(<After />);
+  await waitFor(() =>
+    expect(
+      transcript?.get("task-one")?.messages[0]?.metadata?.timing,
+    ).toMatchObject({
+      totalChunks: 2,
+    }),
+  );
+  await act(async () => settleFirst());
+  await waitFor(() =>
+    expect(resolveSubagentNamespace).toHaveBeenCalledTimes(2),
+  );
+
+  view.unmount();
+  await act(async () => {});
+  expect(release).toHaveBeenCalledOnce();
+  expect(listeners.size).toBe(0);
+});

package/src/useSubagentTranscripts.test.tsx

@@ -185,6 +185,7 @@ describe("useSubagentTranscripts", () => {
     expect(stream.acquire).toHaveBeenCalledTimes(2);
     expect(stream.resolveSubagentNamespace).not.toHaveBeenCalled();
     hook.unmount();
+    await act(async () => {});
     expect(stream.releases.get("tools:one")).toHaveBeenCalledOnce();
     expect(stream.releases.get("tools:two")).toHaveBeenCalledOnce();
   });
@@ -773,6 +774,7 @@ describe("useSubagentTranscripts", () => {
     expect(channelSpecs[0]?.namespace).toEqual(["tools:parent", "tools:child"]);
 
     hook.unmount();
+    await act(async () => {});
     expect(
       stream.uiReleases.get("tools:parent/tools:child"),
     ).toHaveBeenCalledOnce();

package/src/useSubagentTranscripts.ts

@@ -14,13 +14,10 @@ import {
   type ExternalMessageConversionCache,
   type useExternalMessageConverter,
 } from "@assistant-ui/core/react";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 import { STREAM_CONTROLLER, type AnyStream } from "@langchain/react";
 import type { BaseMessage } from "@langchain/core/messages";
-import {
-  channelProjection,
-  messagesProjection,
-  type Event,
-} from "@langchain/langgraph-sdk/stream";
+import { channelProjection, type Event } from "@langchain/langgraph-sdk/stream";
 import type { SubagentDiscoverySnapshot } from "@langchain/react";
 import {
   attachSubagentTranscripts,
@@ -31,6 +28,7 @@ import {
 import { convertLangChainBaseMessage } from "./convertMessages";
 import { groupUIMessagesByParent } from "./converter";
 import { langChainStreamingTimingAccessors } from "./streamingTiming";
+import { subagentMessagesProjection } from "./subagentMessagesProjection";
 import type { LangChainBaseMessage, UIMessage } from "./types";
 import {
   createUIFoldMemo,
@@ -301,7 +299,7 @@ const createSubagentTranscriptSource = (): SubagentTranscriptSource => {
         }
         if (source.resources.has(snapshot.id)) continue;
         const acquired = controller.registry.acquire(
-          messagesProjection(snapshot.namespace),
+          subagentMessagesProjection(snapshot.namespace),
         );
         const acquiredUI =
           snapshot.namespace.length > ROOT_UI_CHANNEL_DEPTH
@@ -516,7 +514,7 @@ export const useSubagentTranscripts = (
     source.reconcile(controller, stream.subagents, uiMessagesByParent);
   }, [controller, source, stream.subagents, uiMessagesByParent]);
 
-  useEffect(() => () => source.dispose(), [source]);
+  useReplaySafeEffect(() => () => source.dispose(), [source]);
 
   return useSyncExternalStore(
     source.subscribe,