@assistant-ui/react-langchain 0.0.27 → 0.0.29

package/src/streamingTiming.ts

@@ -1,76 +1,13 @@
 "use client";
 
 import type { MessageTiming } from "@assistant-ui/core";
-import {
-  useStreamingTiming,
-  type StreamingTimingAccessors,
-} from "@assistant-ui/core/react";
-import type { LangChainBaseMessage, LangChainContentBlock } from "./types";
+import { useStreamingTiming } from "@assistant-ui/core/react";
+import { createLangChainStreamingTimingAccessors } from "./converter";
 import { getMessageType } from "./convertMessages";
+import type { LangChainBaseMessage } from "./types";
 
-const findAiMessage = (
-  messages: readonly LangChainBaseMessage[],
-  messageId: string,
-): LangChainBaseMessage | undefined =>
-  messages.find((m) => getMessageType(m) === "ai" && m.id === messageId);
-
-const reasoningTextLength = (part: {
-  readonly summary?: ReadonlyArray<{ readonly text?: string }>;
-  readonly reasoning?: string;
-}): number => {
-  if (part.summary && part.summary.length > 0)
-    return part.summary.map((s) => s?.text ?? "").join("\n\n\n").length;
-  return part.reasoning?.length ?? 0;
-};
-
-const getTextLength = (
-  messages: readonly LangChainBaseMessage[],
-  messageId: string,
-): number => {
-  const m = findAiMessage(messages, messageId);
-  if (!m) return 0;
-  const content = m.content;
-  if (typeof content === "string") return content.length;
-  if (!Array.isArray(content)) return 0;
-  let len = 0;
-  for (const part of content as readonly LangChainContentBlock[]) {
-    switch (part.type) {
-      case "text":
-      case "text_delta":
-        if (typeof part.text === "string") len += part.text.length;
-        break;
-      case "thinking":
-        if (typeof part.thinking === "string") len += part.thinking.length;
-        break;
-      case "reasoning":
-        len += reasoningTextLength(part);
-        break;
-    }
-  }
-  return len;
-};
-
-const getToolCallCount = (
-  messages: readonly LangChainBaseMessage[],
-  messageId: string,
-): number => findAiMessage(messages, messageId)?.tool_calls?.length ?? 0;
-
-const getAssistantMessageId = (
-  messages: readonly LangChainBaseMessage[],
-): string | undefined => {
-  for (let i = messages.length - 1; i >= 0; i--) {
-    const m = messages[i];
-    if (m && getMessageType(m) === "ai" && m.id) return m.id;
-  }
-  return undefined;
-};
-
-export const langChainStreamingTimingAccessors: StreamingTimingAccessors<LangChainBaseMessage> =
-  {
-    getAssistantMessageId,
-    getTextLength,
-    getToolCallCount,
-  };
+export const langChainStreamingTimingAccessors =
+  createLangChainStreamingTimingAccessors<LangChainBaseMessage>(getMessageType);
 
 /**
  * Tracks per-message streaming timing for LangChain messages. Delegates to

package/src/types.ts

@@ -17,45 +17,7 @@ import type {
   SubgraphDiscoverySnapshot,
 } from "@langchain/react";
 
-/** Known content block types from @langchain/core messages. */
-export type LangChainContentBlock =
-  | { type: "text"; text: string }
-  | { type: "text_delta"; text: string }
-  | { type: "image_url"; image_url: string | { url?: string } }
-  | { type: "thinking"; thinking: string }
-  | {
-      type: "reasoning";
-      summary?: Array<{ type: "summary_text"; text?: string }>;
-      reasoning?: string;
-    }
-  | {
-      type: "file";
-      data: string;
-      mime_type: string;
-      source_type?: "base64";
-      metadata?: { filename?: string };
-    }
-  | {
-      type: "file";
-      url: string;
-      mime_type?: string;
-      source_type: "url";
-      metadata?: { filename?: string };
-    }
-  | {
-      type: "file";
-      id: string;
-      mime_type?: string;
-      source_type: "id";
-      metadata?: { filename?: string };
-    }
-  | {
-      type: "audio";
-      data: string;
-      mime_type: string;
-      source_type: "base64";
-    }
-  | { type: "tool_use" | "input_json_delta" };
+export type { LangChainContentBlock } from "./converter";
 
 export type LangChainToolCall = {
   id: string;

package/src/useStreamRuntime.test.tsx

@@ -10,7 +10,7 @@ import type {
 } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
 import type { LangChainBaseMessage } from "./types";
-import type { ReactNode } from "react";
+import { startTransition, Suspense, type ReactNode } from "react";
 import {
   useLangChainRespond,
   useLangChainRespondAll,
@@ -898,3 +898,76 @@ describe("useStreamRuntime staged messages", () => {
     });
   });
 });
+
+describe("useStreamRuntime committed refs", () => {
+  it("submits through the committed stream after an abandoned render", async () => {
+    const streamA = createMockStream();
+    const streamB = createMockStream();
+    const adapter = makeThreadListAdapter();
+    mockUseStream.mockImplementation((options: { apiUrl: string }) =>
+      options.apiUrl === "/api/b" ? streamB : streamA,
+    );
+    const host = renderHook(() =>
+      useStreamRuntime({
+        apiUrl: "/api/a",
+        unstable_threadListAdapter: adapter,
+      } as never),
+    );
+
+    const pending = new Promise<never>(() => {});
+    let blocked = false;
+    const interruptedRender = vi.fn();
+    const Blocker = () => {
+      if (blocked) {
+        interruptedRender();
+        throw pending;
+      }
+      return null;
+    };
+
+    const capture: { runtime: AssistantRuntime | null } = { runtime: null };
+    const Nested = ({ apiUrl }: { apiUrl: string }) => {
+      capture.runtime = useStreamRuntime({
+        apiUrl,
+        unstable_threadListAdapter: adapter,
+      } as never);
+      return null;
+    };
+    const Tree = ({ apiUrl }: { apiUrl: string }) => (
+      <AssistantRuntimeProvider runtime={host.result.current}>
+        <Suspense fallback={null}>
+          <Nested apiUrl={apiUrl} />
+          <Blocker />
+        </Suspense>
+      </AssistantRuntimeProvider>
+    );
+
+    const view = render(<Tree apiUrl="/api/a" />);
+    expect(capture.runtime).not.toBeNull();
+
+    act(() => {
+      blocked = true;
+      startTransition(() => view.rerender(<Tree apiUrl="/api/b" />));
+    });
+    expect(interruptedRender).toHaveBeenCalled();
+
+    await act(async () => {
+      await capture.runtime!.thread.append("hello");
+    });
+
+    expect(streamA.submit).toHaveBeenCalledOnce();
+    expect(streamB.submit).not.toHaveBeenCalled();
+
+    await act(async () => {
+      blocked = false;
+      view.rerender(<Tree apiUrl="/api/b" />);
+    });
+    await act(async () => {
+      await capture.runtime!.thread.append("second");
+    });
+
+    expect(streamB.submit).toHaveBeenCalledOnce();
+    view.unmount();
+    host.unmount();
+  });
+});

package/src/useStreamRuntime.ts

@@ -1,7 +1,14 @@
 /// <reference types="@assistant-ui/core/store" />
 "use client";
 
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 import type { AppendMessage, ToolExecutionStatus } from "@assistant-ui/core";
 import {
   generateId,
@@ -9,6 +16,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 +36,8 @@ import type {
   UIMessage,
   UseStreamRuntimeOptions,
 } from "./types";
+import { groupUIMessagesByParent } from "./converter";
+export { groupUIMessagesByParent } from "./converter";
 import {
   convertLangChainBaseMessage,
   getMessageContent,
@@ -47,44 +61,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,
@@ -191,7 +182,8 @@ const useStreamThreadRuntime = (
   const convertWithUI = useMemo<
     useExternalMessageConverter.Callback<LangChainBaseMessage>
   >(() => {
-    const uiMessagesByParent = groupUIMessagesByParent(mergedUiMessages);
+    const uiMessagesByParent =
+      groupUIMessagesByParent<UIMessage>(mergedUiMessages);
     return (message, metadata) =>
       convertLangChainBaseMessage(message, {
         ...metadata,
@@ -207,7 +199,9 @@ const useStreamThreadRuntime = (
   });
 
   const streamRef = useRef(stream);
-  streamRef.current = stream;
+  useInsertionEffect(() => {
+    streamRef.current = stream;
+  }, [stream]);
 
   const activeRunConfigRef = useRef<
     NormalizedRunConfigOptions["config"] | undefined
@@ -271,10 +265,14 @@ const useStreamThreadRuntime = (
   }, [stream.messages]);
 
   const visibleMessagesRef = useRef(visibleMessages);
-  visibleMessagesRef.current = visibleMessages;
+  useInsertionEffect(() => {
+    visibleMessagesRef.current = visibleMessages;
+  }, [visibleMessages]);
 
   const threadMessagesRef = useRef(threadMessages);
-  threadMessagesRef.current = threadMessages;
+  useInsertionEffect(() => {
+    threadMessagesRef.current = threadMessages;
+  }, [threadMessages]);
 
   const stagedMessagesRef = useRef(
     new Map<
@@ -409,7 +407,7 @@ const useStreamThreadRuntime = (
 
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),
-    isRunning: effectiveIsRunning,
+    isRunning: stream.isLoading,
     isLoading: stream.isThreadLoading,
     messages: threadMessages,
     adapters,
@@ -430,13 +428,7 @@ const useStreamThreadRuntime = (
         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
@@ -634,19 +626,20 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
     ...options
   } = rawOptions;
 
-  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;
 
   return useRemoteThreadListRuntime({
     runtimeHook: function RuntimeHook() {
-      return useStreamThreadRuntime(optionsRef.current);
+      return useStreamThreadRuntime(options);
     },
     adapter,
     allowNesting: true,