@assistant-ui/react-langchain 0.0.26 → 0.0.28

package/src/useStreamRuntime.test.tsx

@@ -5,6 +5,7 @@ import { describe, expect, it, vi } from "vitest";
 import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
 import type {
   AssistantRuntime,
+  AppendMessage,
   RemoteThreadListAdapter,
 } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
@@ -156,10 +157,12 @@ const makeThreadListAdapter = (): RemoteThreadListAdapter => ({
 
 const deferred = <T,>() => {
   let resolve!: (value: T) => void;
-  const promise = new Promise<T>((res) => {
+  let reject!: (reason?: unknown) => void;
+  const promise = new Promise<T>((res, rej) => {
     resolve = res;
+    reject = rej;
   });
-  return { promise, resolve };
+  return { promise, resolve, reject };
 };
 
 describe("useStreamRuntime thread options", () => {
@@ -202,7 +205,7 @@ describe("useStreamRuntime thread options", () => {
     view.unmount();
   });
 
-  it("submits with the initialized thread id when initialization resolves before rerender", async () => {
+  it("renders before initialization and submits with the initialized thread id", async () => {
     const stream = createMockStream();
     mockUseStream.mockReturnValue(stream);
     const initialization = deferred<{
@@ -213,18 +216,29 @@ describe("useStreamRuntime thread options", () => {
     threadListAdapter.list = vi.fn(async () => ({ threads: [] }));
     threadListAdapter.initialize = vi.fn(() => initialization.promise);
 
-    const capture: { runtime: AssistantRuntime | null } = { runtime: null };
+    const capture: {
+      runtime: AssistantRuntime | null;
+      aui?: ReturnType<typeof useAui>;
+    } = { runtime: null };
+    const Capture = () => {
+      capture.aui = useAui();
+      return null;
+    };
     const TestRuntime = () => {
       const runtime = useStreamRuntime({
         apiUrl: "/api",
         unstable_threadListAdapter: threadListAdapter,
       } as never);
       capture.runtime = runtime;
-      return <AssistantRuntimeProvider runtime={runtime} />;
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Capture />
+        </AssistantRuntimeProvider>
+      );
     };
 
     const view = render(<TestRuntime />);
-    await waitFor(() => expect(capture.runtime).not.toBeNull());
+    await waitFor(() => expect(capture.aui).toBeDefined());
 
     await act(async () => {
       capture.runtime!.thread.append({
@@ -235,6 +249,7 @@ describe("useStreamRuntime thread options", () => {
     });
 
     expect(stream.submit).not.toHaveBeenCalled();
+    expect(getText(capture.aui!)).toEqual(["hello"]);
 
     await act(async () => {
       initialization.resolve({ remoteId: "thread-b", externalId: "thread-b" });
@@ -242,10 +257,25 @@ describe("useStreamRuntime thread options", () => {
 
     await waitFor(() =>
       expect(stream.submit).toHaveBeenCalledWith(
-        { messages: [{ type: "human", content: "hello" }] },
+        {
+          messages: [
+            expect.objectContaining({
+              id: expect.any(String),
+              type: "human",
+              content: "hello",
+            }),
+          ],
+        },
         { threadId: "thread-b" },
       ),
     );
+
+    stream.messages = [message("echo-hello", "human", "hello")];
+    view.rerender(<TestRuntime />);
+    await waitFor(() => {
+      expect(getText(capture.aui!)).toEqual(["hello"]);
+      expect(capture.aui!.thread.getState().messages).toHaveLength(1);
+    });
     view.unmount();
   });
 
@@ -280,6 +310,100 @@ describe("useStreamRuntime thread options", () => {
     }
     view.unmount();
   });
+
+  it.each(["initialization", "submit"] as const)(
+    "removes the staged message when %s fails",
+    async (failurePoint) => {
+      const stream = createMockStream();
+      mockUseStream.mockReturnValue(stream);
+      const initialization = deferred<{
+        remoteId: string;
+        externalId: string;
+      }>();
+      const threadListAdapter = makeThreadListAdapter();
+      threadListAdapter.list = vi.fn(async () => ({ threads: [] }));
+      threadListAdapter.initialize = vi.fn(() => initialization.promise);
+
+      const capture: {
+        runtime: AssistantRuntime | null;
+        aui?: ReturnType<typeof useAui>;
+      } = { runtime: null };
+      const Capture = () => {
+        capture.aui = useAui();
+        return null;
+      };
+      const TestRuntime = () => {
+        const runtime = useStreamRuntime({
+          apiUrl: "/api",
+          unstable_threadListAdapter: threadListAdapter,
+        } as never);
+        capture.runtime = runtime;
+        return (
+          <AssistantRuntimeProvider runtime={runtime}>
+            <Capture />
+          </AssistantRuntimeProvider>
+        );
+      };
+
+      const view = render(<TestRuntime />);
+      await waitFor(() => expect(capture.aui).toBeDefined());
+
+      if (failurePoint === "submit") {
+        stream.submit.mockRejectedValueOnce(new Error("submit failed"));
+      }
+
+      const core = (
+        capture.runtime!.thread as unknown as {
+          __internal_threadBinding: {
+            getState(): { append(message: AppendMessage): Promise<void> };
+          };
+        }
+      ).__internal_threadBinding.getState();
+      let appendPromise!: Promise<void>;
+      await act(async () => {
+        appendPromise = core.append({
+          role: "user",
+          content: [{ type: "text", text: "failed" }],
+          parentId: null,
+          sourceId: null,
+          runConfig: undefined,
+          attachments: [],
+          metadata: { custom: {} },
+          createdAt: new Date(0),
+        });
+        await Promise.resolve();
+      });
+      const appendResult = appendPromise.then(
+        () => undefined,
+        (error: unknown) => error,
+      );
+      expect(getText(capture.aui!)).toEqual(["failed"]);
+
+      await act(async () => {
+        if (failurePoint === "initialization") {
+          initialization.reject(new Error("initialize failed"));
+        } else {
+          initialization.resolve({
+            remoteId: "thread-failed",
+            externalId: "thread-failed",
+          });
+        }
+      });
+      await expect(appendResult).resolves.toMatchObject({
+        message:
+          failurePoint === "initialization"
+            ? "initialize failed"
+            : "submit failed",
+      });
+      await waitFor(() => expect(getText(capture.aui!)).toEqual([]));
+      if (failurePoint === "initialization") {
+        expect(stream.submit).not.toHaveBeenCalled();
+      } else {
+        expect(stream.submit).toHaveBeenCalledTimes(1);
+      }
+      view.unmount();
+    },
+  );
 });
 
 describe("useStreamRuntime run configuration", () => {
@@ -331,7 +455,9 @@ describe("useStreamRuntime run configuration", () => {
     expect(stream.submit).toHaveBeenNthCalledWith(
       1,
       {
-        messages: [{ type: "human", content: "hello" }],
+        messages: [
+          expect.objectContaining({ type: "human", content: "hello" }),
+        ],
       },
       config,
     );
@@ -659,6 +785,55 @@ describe("useStreamRuntime staged messages", () => {
     expect(stream.submit).not.toHaveBeenCalled();
   });
 
+  it("does not resurrect a staged draft that an edit already truncated", async () => {
+    const stream = createMockStream([
+      message("u1", "human", "first"),
+      message("a1", "ai", "first answer"),
+      message("u2", "human", "second"),
+    ]);
+    const { auiResult, rerender } = renderAui(stream);
+
+    await act(async () => {
+      auiResult.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "draft" }],
+        startRun: false,
+      });
+    });
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual([
+        "first",
+        "first answer",
+        "second",
+        "draft",
+      ]);
+    });
+
+    await act(async () => {
+      auiResult.current.thread.append({
+        role: "user",
+        parentId: "u1",
+        content: [{ type: "text", text: "edited" }],
+        startRun: false,
+      });
+    });
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual(["first", "edited"]);
+    });
+
+    stream.messages = [
+      message("u1", "human", "first"),
+      message("a1", "ai", "first answer from refresh"),
+      message("u2", "human", "second from refresh"),
+    ];
+    rerender();
+
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual(["first", "edited"]);
+    });
+    expect(stream.submit).not.toHaveBeenCalled();
+  });
+
   it("keeps later staged messages visible after promoting one staged parent", async () => {
     const stream = createMockStream([message("u1", "human", "earlier")]);
     const { auiResult, rerender } = renderAui(stream);

package/src/useStreamRuntime.ts

@@ -9,6 +9,11 @@ import {
   pickExternalStoreSharedOptions,
 } from "@assistant-ui/core";
 import type { ThreadMessage } from "@assistant-ui/core";
+import {
+  createCloudThreadListAdapterCreateFallback,
+  createToolCallCancellationStub,
+  scanPendingToolCalls,
+} from "@assistant-ui/core/internal";
 import {
   useCloudThreadListAdapter,
   useExternalStoreRuntime,
@@ -24,6 +29,8 @@ import type {
   UIMessage,
   UseStreamRuntimeOptions,
 } from "./types";
+import { groupUIMessagesByParent } from "./converter";
+export { groupUIMessagesByParent } from "./converter";
 import {
   convertLangChainBaseMessage,
   getMessageContent,
@@ -47,44 +54,21 @@ type NormalizedRunConfigOptions = NonNullable<
   ReturnType<typeof runConfigToSubmitOptions>
 >;
 
-/**
- * Group the graph's accumulated `UIMessage`s by the assistant message they
- * belong to. Non-array state and entries without a parent link are dropped.
- * The parent id comes from `metadata.message_id` (Python SDK) or
- * `metadata.id` (JS SDK).
- */
-export const groupUIMessagesByParent = (
-  value: unknown,
-): Map<string, UIMessage[]> => {
-  const map = new Map<string, UIMessage[]>();
-  if (!Array.isArray(value)) return map;
-  for (const ui of value as UIMessage[]) {
-    const parentId = ui.metadata?.message_id ?? ui.metadata?.id;
-    if (!parentId) continue;
-    const existing = map.get(parentId);
-    if (existing) {
-      existing.push(ui);
-    } else {
-      map.set(parentId, [ui]);
-    }
-  }
-  return map;
-};
-
 const getPendingToolCalls = (
   messages: readonly LangChainBaseMessage[],
-): LangChainToolCall[] => {
-  const pending = new Map<string, LangChainToolCall>();
-  for (const m of messages) {
-    const type = getMessageType(m);
-    if (type === "ai") {
-      for (const tc of m.tool_calls ?? []) pending.set(tc.id, tc);
-    } else if (type === "tool" && m.tool_call_id) {
-      pending.delete(m.tool_call_id);
-    }
-  }
-  return [...pending.values()];
-};
+): LangChainToolCall[] =>
+  scanPendingToolCalls(
+    messages,
+    (message) => {
+      const type = getMessageType(message);
+      if (type === "ai") return { toolCalls: message.tool_calls ?? [] };
+      if (type === "tool" && message.tool_call_id) {
+        return { toolCallId: message.tool_call_id };
+      }
+      return undefined;
+    },
+    (toolCall) => toolCall.id,
+  );
 
 const toStagedHumanMessage = (
   msg: AppendMessage,
@@ -95,6 +79,26 @@ const toStagedHumanMessage = (
   content: getMessageContent(msg),
 });
 
+const humanContentText = (content: LangChainBaseMessage["content"]) => {
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return "";
+  return content
+    .filter(
+      (part): part is { type: "text"; text: string } =>
+        typeof part === "object" &&
+        part !== null &&
+        part.type === "text" &&
+        typeof part.text === "string",
+    )
+    .map((part) => part.text)
+    .join("");
+};
+
+const hasSameMessageContent = (
+  a: LangChainBaseMessage,
+  b: LangChainBaseMessage,
+) => humanContentText(a.content) === humanContentText(b.content);
+
 const truncateLangChainBaseMessages = (
   threadMessages: readonly ThreadMessage[],
   parentId: string | null,
@@ -171,7 +175,8 @@ const useStreamThreadRuntime = (
   const convertWithUI = useMemo<
     useExternalMessageConverter.Callback<LangChainBaseMessage>
   >(() => {
-    const uiMessagesByParent = groupUIMessagesByParent(mergedUiMessages);
+    const uiMessagesByParent =
+      groupUIMessagesByParent<UIMessage>(mergedUiMessages);
     return (message, metadata) =>
       convertLangChainBaseMessage(message, {
         ...metadata,
@@ -262,11 +267,12 @@ const useStreamThreadRuntime = (
       {
         message: LangChainBaseMessage & { id: string };
         runConfig: AppendMessage["runConfig"];
+        reconcileOnEcho: boolean;
+        baseMessageCount: number;
       }
     >(),
   );
   const stagedBaseMessagesRef = useRef<LangChainBaseMessage[] | null>(null);
-
   useEffect(() => {
     if (stagedMessagesRef.current.size === 0) return;
 
@@ -274,18 +280,32 @@ const useStreamThreadRuntime = (
     const baseMessages =
       stagedBaseMessagesRef.current ??
       (stream.messages as LangChainBaseMessage[]);
-    const baseMessageIds = new Set(
-      baseMessages.flatMap((message) => (message.id ? [message.id] : [])),
-    );
     const remainingStagedMessages: LangChainBaseMessage[] = [];
-    const seenStagedIds = new Set<string>();
-    for (const message of visibleMessagesRef.current) {
-      if (!message.id || seenStagedIds.has(message.id)) continue;
-      if (baseMessageIds.has(message.id)) continue;
-      const staged = stagedMessagesRef.current.get(message.id);
-      if (!staged) continue;
-      remainingStagedMessages.push(staged.message);
-      seenStagedIds.add(message.id);
+    const matchedBaseMessageIndexes = new Set<number>();
+    const visibleStagedIds = new Set(
+      visibleMessagesRef.current.flatMap((m) => (m.id ? [m.id] : [])),
+    );
+    for (const [id, staged] of stagedMessagesRef.current) {
+      if (!visibleStagedIds.has(id)) continue;
+      const echoed = baseMessages.some((message, index) => {
+        if (matchedBaseMessageIndexes.has(index)) return false;
+        if (message.id === id) {
+          matchedBaseMessageIndexes.add(index);
+          return true;
+        }
+        if (
+          !staged.reconcileOnEcho ||
+          index < staged.baseMessageCount ||
+          getMessageType(message) !== "human" ||
+          !hasSameMessageContent(message, staged.message)
+        ) {
+          return false;
+        }
+        matchedBaseMessageIndexes.add(index);
+        return true;
+      });
+      if (echoed) stagedMessagesRef.current.delete(id);
+      else remainingStagedMessages.push(staged.message);
     }
 
     if (remainingStagedMessages.length === 0) {
@@ -317,15 +337,32 @@ const useStreamThreadRuntime = (
     };
   };
 
-  const stageUserMessage = (msg: AppendMessage) => {
+  const stageUserMessage = (msg: AppendMessage, reconcileOnEcho = false) => {
     const stagedMessage = toStagedHumanMessage(msg);
     stagedMessagesRef.current.set(stagedMessage.id, {
       message: stagedMessage,
       runConfig: msg.runConfig,
+      reconcileOnEcho,
+      baseMessageCount: streamRef.current.messages.length,
     });
     const nextMessages = [...visibleMessagesRef.current, stagedMessage];
     visibleMessagesRef.current = nextMessages;
     setStagedMessages(nextMessages);
+    return stagedMessage;
+  };
+
+  const removeStagedMessage = (id: string) => {
+    if (!stagedMessagesRef.current.delete(id)) return;
+    const nextMessages = visibleMessagesRef.current.filter(
+      (message) => message.id !== id,
+    );
+    visibleMessagesRef.current = nextMessages;
+    if (stagedMessagesRef.current.size === 0) {
+      stagedBaseMessagesRef.current = null;
+      setStagedMessages(null);
+    } else {
+      setStagedMessages(nextMessages);
+    }
   };
 
   const extras = useMemo(
@@ -357,7 +394,7 @@ const useStreamThreadRuntime = (
 
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),
-    isRunning: effectiveIsRunning,
+    isRunning: stream.isLoading,
     isLoading: stream.isThreadLoading,
     messages: threadMessages,
     adapters,
@@ -370,32 +407,42 @@ const useStreamThreadRuntime = (
         return;
       }
 
+      const stagedMessage = stageUserMessage(msg, true);
+      const stagedMessageId = stagedMessage.id;
       setActiveRunConfig(msg.runConfig);
       const content = getMessageContent(msg);
       const cancellations =
         autoCancelPendingToolCalls !== false
           ? getPendingToolCalls(
               streamRef.current.messages as readonly LangChainBaseMessage[],
-            ).map((t) => ({
-              type: "tool" as const,
-              name: t.name,
-              tool_call_id: t.id,
-              content: JSON.stringify({ cancelled: true }),
-              status: "error" as const,
-            }))
+            ).map(createToolCallCancellationStub)
           : [];
       // A null threadId is not a no-op for the SDK: it rebinds the controller
       // away from its self-created thread and forces a fresh one, so the
       // submit waits for initialization to produce an identity; core no
       // longer holds appends on that barrier.
-      const { externalId } = await aui.threadListItem.initialize();
-      await streamRef.current.submit(
-        { [messagesKey]: [...cancellations, { type: "human", content }] },
-        {
-          ...runConfigToSubmitOptions(msg.runConfig),
-          ...(externalId != null ? { threadId: externalId } : {}),
-        },
-      );
+      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 } : {}),
+          },
+        );
+      } catch (error) {
+        removeStagedMessage(stagedMessageId);
+        throw error;
+      }
     },
     onAddToolResult: async ({
       messageId,
@@ -487,6 +534,8 @@ const useStreamThreadRuntime = (
         stagedMessagesRef.current.set(stagedMessage.id, {
           message: stagedMessage,
           runConfig: message.runConfig,
+          reconcileOnEcho: false,
+          baseMessageCount: 0,
         });
         stagedBaseMessagesRef.current = truncated;
         const nextMessages = [...truncated, stagedMessage];
@@ -567,9 +616,13 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
   const optionsRef = useRef(options);
   optionsRef.current = options;
 
+  const aui = useAui();
   const cloudAdapter = useCloudThreadListAdapter({
     cloud,
-    create,
+    create: createCloudThreadListAdapterCreateFallback(
+      create,
+      aui.threadListItem,
+    ),
     delete: deleteFn,
   });
   const adapter = unstable_threadListAdapter ?? cloudAdapter;