@assistant-ui/react-langchain 0.0.15 → 0.0.17

package/src/useStreamRuntime.ts

@@ -1,98 +1,40 @@
 /// <reference types="@assistant-ui/core/store" />
 "use client";
 
-import { useMemo, useRef, useState } from "react";
-import type {
-  AppendMessage,
-  AttachmentAdapter,
-  DictationAdapter,
-  ExternalStoreSharedOptions,
-  FeedbackAdapter,
-  RealtimeVoiceAdapter,
-  RemoteThreadListAdapter,
-  SpeechSynthesisAdapter,
-  ToolExecutionStatus,
+import { useEffect, useMemo, useRef, useState } from "react";
+import type { AppendMessage, ToolExecutionStatus } from "@assistant-ui/core";
+import {
+  generateId,
+  getExternalStoreMessages,
+  pickExternalStoreSharedOptions,
 } from "@assistant-ui/core";
-import { pickExternalStoreSharedOptions } from "@assistant-ui/core";
+import type { ThreadMessage } from "@assistant-ui/core";
 import {
   useCloudThreadListAdapter,
   useExternalStoreRuntime,
   useExternalMessageConverter,
   useRemoteThreadListRuntime,
 } from "@assistant-ui/core/react";
-import { useAui, useAuiState } from "@assistant-ui/store";
-import type { AssistantCloud } from "assistant-cloud";
+import { 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 {
-  useStream,
-  type UseStreamOptions,
-  type AssembledToolCall,
-} from "@langchain/react";
-import type { LangChainBaseMessage, LangChainToolCall } from "./types";
-import { convertLangChainBaseMessage, getMessageType } from "./convertMessages";
-
-const symbolLangChainRuntimeExtras = Symbol("langchain-runtime-extras");
-
-const EMPTY_TOOL_CALLS: readonly AssembledToolCall[] = [];
-
-type LangChainRuntimeExtras = {
-  [symbolLangChainRuntimeExtras]: true;
-  interrupt: { value?: unknown } | undefined;
-  interrupts: readonly { value?: unknown }[];
-  toolCalls: readonly AssembledToolCall[];
-  error: unknown;
-  submit: (
-    values: Record<string, unknown> | null | undefined,
-    options?: Record<string, unknown>,
-  ) => Promise<void>;
-  values: Record<string, unknown>;
-  messagesKey: string;
-};
-
-const asLangChainRuntimeExtras = (extras: unknown): LangChainRuntimeExtras => {
-  if (
-    typeof extras !== "object" ||
-    extras == null ||
-    !(symbolLangChainRuntimeExtras in extras)
-  )
-    throw new Error(
-      "This method can only be called when you are using useStreamRuntime",
-    );
-  return extras as LangChainRuntimeExtras;
-};
+  convertLangChainBaseMessage,
+  getMessageContent,
+  getMessageType,
+} from "./convertMessages";
+import { foldUIUpdates, mergeUIMessages } from "./uiMessages";
+import { langChainExtras } from "./runtimeExtras";
+import { resolveForkCheckpoint } from "./resolveForkCheckpoint";
+import { useLangChainStreamingTiming } from "./streamingTiming";
 
-type LangChainRuntimeExtraOptions = ExternalStoreSharedOptions & {
-  cloud?: AssistantCloud | undefined;
-  adapters?:
-    | {
-        attachments?: AttachmentAdapter | undefined;
-        speech?: SpeechSynthesisAdapter | undefined;
-        dictation?: DictationAdapter | undefined;
-        voice?: RealtimeVoiceAdapter | undefined;
-        feedback?: FeedbackAdapter | undefined;
-      }
-    | undefined;
-  /**
-   * When the user sends a new message while previous tool calls are
-   * still pending, automatically submit `tool` messages that cancel
-   * them so the agent's tool-call accounting stays consistent.
-   * Defaults to `true`.
-   */
-  autoCancelPendingToolCalls?: boolean | undefined;
-  /**
-   * Routes the Cancel button's click to `useStream().stop()`. On by
-   * default. Pass `false` to disable the Cancel button.
-   */
-  unstable_allowCancellation?: boolean | undefined;
-  /**
-   * Custom `RemoteThreadListAdapter`. When provided, replaces the
-   * cloud-backed thread list adapter.
-   */
-  unstable_threadListAdapter?: RemoteThreadListAdapter | undefined;
-  /** Custom thread-creation hook, forwarded to the cloud adapter. */
-  create?: (() => Promise<{ externalId: string | undefined }>) | undefined;
-  /** Custom thread-deletion hook, forwarded to the cloud adapter. */
-  delete?: ((threadId: string) => Promise<void>) | undefined;
-};
+const UI_CUSTOM_CHANNELS: readonly Channel[] = ["custom"];
 
 export const runConfigToSubmitOptions = (
   runConfig: AppendMessage["runConfig"],
@@ -101,6 +43,30 @@ export const runConfigToSubmitOptions = (
     ? { config: { configurable: runConfig.custom } }
     : undefined;
 
+/**
+ * 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[] => {
@@ -116,65 +82,29 @@ const getPendingToolCalls = (
   return [...pending.values()];
 };
 
-// Distribute the intersection through the union arms of `UseStreamOptions`
-// (`AgentServerOptions | CustomAdapterOptions`). Writing `UseStreamOptions & X`
-// directly collapses arm tracking, so `Omit<…, "cloud">` and the like would
-// produce a flattened structural type that no longer matches either arm.
-export type UseStreamRuntimeOptions = UseStreamOptions extends infer O
-  ? O extends UseStreamOptions
-    ? O & LangChainRuntimeExtraOptions
-    : never
-  : never;
-
-const getMessageContent = (msg: AppendMessage) => {
-  const allContent = [
-    ...msg.content,
-    ...(msg.attachments?.flatMap((a) => a.content) ?? []),
-  ];
-
-  const hasNonText = allContent.some(
-    (part) => part.type === "file" || part.type === "image",
-  );
-  const hasText = allContent.some((part) => part.type === "text");
-  if (hasNonText && !hasText) {
-    allContent.unshift({ type: "text", text: " " });
-  }
-
-  const content = allContent.map((part) => {
-    const type = part.type;
-    switch (type) {
-      case "text":
-        return { type: "text" as const, text: part.text };
-      case "image":
-        return { type: "image_url" as const, image_url: { url: part.image } };
-      case "file":
-        return {
-          type: "file" as const,
-          data: part.data,
-          mime_type: part.mimeType,
-          metadata: { filename: part.filename ?? "file" },
-          source_type: "base64" as const,
-        };
-      case "tool-call":
-        throw new Error("Tool call appends are not supported.");
-      default: {
-        const _exhaustiveCheck:
-          | "reasoning"
-          | "source"
-          | "audio"
-          | "data"
-          | "generative-ui" = type;
-        throw new Error(
-          `Unsupported append message part type: ${_exhaustiveCheck}`,
-        );
-      }
-    }
-  });
-
-  if (content.length === 1 && content[0]?.type === "text") {
-    return content[0].text ?? "";
+const toStagedHumanMessage = (
+  msg: AppendMessage,
+  id = generateId(),
+): LangChainBaseMessage & { id: string } => ({
+  id,
+  _getType: () => "human",
+  content: getMessageContent(msg),
+});
+
+const truncateLangChainBaseMessages = (
+  threadMessages: readonly ThreadMessage[],
+  parentId: string | null,
+): LangChainBaseMessage[] => {
+  if (parentId === null) return [];
+  const parentIndex = threadMessages.findIndex((m) => m.id === parentId);
+  if (parentIndex === -1) return [];
+  const truncated: LangChainBaseMessage[] = [];
+  for (let i = 0; i <= parentIndex && i < threadMessages.length; i++) {
+    truncated.push(
+      ...getExternalStoreMessages<LangChainBaseMessage>(threadMessages[i]!),
+    );
   }
-  return content;
+  return truncated;
 };
 
 type DistributiveOmit<T, K extends keyof any> = T extends unknown
@@ -190,6 +120,7 @@ const useStreamThreadRuntime = (
   const { adapters, autoCancelPendingToolCalls, unstable_allowCancellation } =
     options;
   const messagesKey = options.messagesKey ?? "messages";
+  const uiStateKey = options.uiStateKey ?? "ui";
 
   const externalId = useAuiState((s) => s.threadListItem.externalId) as
     | string
@@ -200,6 +131,9 @@ const useStreamThreadRuntime = (
   options.threadId = externalId;
 
   const stream = useStream(options);
+  const [stagedMessages, setStagedMessages] = useState<
+    LangChainBaseMessage[] | null
+  >(null);
 
   const [toolStatuses, setToolStatuses] = useState<
     Record<string, ToolExecutionStatus>
@@ -209,35 +143,143 @@ const useStreamThreadRuntime = (
   );
   const effectiveIsRunning = stream.isLoading || hasExecutingTools;
 
+  const uiStateValue = stream.values[uiStateKey];
+
+  const customEvents = useChannel(stream, UI_CUSTOM_CHANNELS);
+  const liveUiMessages = useMemo(
+    () => foldUIUpdates(customEvents),
+    [customEvents],
+  );
+
+  const mergedUiMessages = useMemo(
+    () => mergeUIMessages(liveUiMessages, uiStateValue),
+    [liveUiMessages, uiStateValue],
+  );
+
+  const visibleMessages =
+    stagedMessages ?? (stream.messages as LangChainBaseMessage[]);
+
+  const messageTiming = useLangChainStreamingTiming(
+    visibleMessages,
+    effectiveIsRunning,
+  );
+
+  const convertWithUI = useMemo<
+    useExternalMessageConverter.Callback<LangChainBaseMessage>
+  >(() => {
+    const uiMessagesByParent = groupUIMessagesByParent(mergedUiMessages);
+    return (message, metadata) =>
+      convertLangChainBaseMessage(message, {
+        ...metadata,
+        uiMessagesByParent,
+        messageTiming,
+      });
+  }, [mergedUiMessages, messageTiming]);
+
   const threadMessages = useExternalMessageConverter({
-    callback: convertLangChainBaseMessage,
-    messages: stream.messages as LangChainBaseMessage[],
+    callback: convertWithUI,
+    messages: visibleMessages,
     isRunning: effectiveIsRunning,
   });
 
   const streamRef = useRef(stream);
   streamRef.current = stream;
 
+  const visibleMessagesRef = useRef(visibleMessages);
+  visibleMessagesRef.current = visibleMessages;
+
+  const threadMessagesRef = useRef(threadMessages);
+  threadMessagesRef.current = threadMessages;
+
+  const stagedMessagesRef = useRef(
+    new Map<
+      string,
+      {
+        message: LangChainBaseMessage & { id: string };
+        runConfig: AppendMessage["runConfig"];
+      }
+    >(),
+  );
+  const stagedBaseMessagesRef = useRef<LangChainBaseMessage[] | null>(null);
+
+  useEffect(() => {
+    if (stagedMessagesRef.current.size === 0) return;
+
+    // Staged edits must keep their truncated base while stream updates arrive before promotion.
+    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);
+    }
+
+    if (remainingStagedMessages.length === 0) {
+      stagedBaseMessagesRef.current = null;
+      visibleMessagesRef.current = baseMessages;
+      setStagedMessages(null);
+      return;
+    }
+
+    const nextMessages = [...baseMessages, ...remainingStagedMessages];
+    visibleMessagesRef.current = nextMessages;
+    setStagedMessages(nextMessages);
+  }, [stream.messages]);
+
+  const getStagedRun = (parentId: string | null) => {
+    if (!parentId || !stagedMessagesRef.current.has(parentId)) 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);
+      }
+      if (message.id === parentId) break;
+    }
+
+    return {
+      messages: staged,
+      runConfig: stagedMessagesRef.current.get(parentId)!.runConfig,
+    };
+  };
+
+  const stageUserMessage = (msg: AppendMessage) => {
+    const stagedMessage = toStagedHumanMessage(msg);
+    stagedMessagesRef.current.set(stagedMessage.id, {
+      message: stagedMessage,
+      runConfig: msg.runConfig,
+    });
+    const nextMessages = [...visibleMessagesRef.current, stagedMessage];
+    visibleMessagesRef.current = nextMessages;
+    setStagedMessages(nextMessages);
+  };
+
   const extras = useMemo(
-    (): LangChainRuntimeExtras => ({
-      [symbolLangChainRuntimeExtras]: true,
-      interrupt: stream.interrupt,
-      interrupts: stream.interrupts,
-      toolCalls: stream.toolCalls,
-      error: stream.error,
-      submit: stream.submit,
-      values: stream.values,
-      messagesKey,
-    }),
-    [
-      stream.interrupt,
-      stream.interrupts,
-      stream.toolCalls,
-      stream.error,
-      stream.submit,
-      stream.values,
-      messagesKey,
-    ],
+    () =>
+      langChainExtras.provide({
+        interrupt: stream.interrupt,
+        interrupts: stream.interrupts,
+        toolCalls: stream.toolCalls,
+        subagents: stream.subagents,
+        subgraphs: stream.subgraphs,
+        stream,
+        error: stream.error,
+        submit: stream.submit,
+        respond: stream.respond,
+        respondAll: stream.respondAll,
+        values: stream.values,
+        messagesKey,
+      }),
+    [stream, messagesKey],
   );
 
   const runtime = useExternalStoreRuntime({
@@ -250,6 +292,11 @@ const useStreamThreadRuntime = (
     unstable_enableToolInvocations: true,
     setToolStatuses,
     onNew: async (msg) => {
+      if (!(msg.startRun ?? msg.role === "user")) {
+        stageUserMessage(msg);
+        return;
+      }
+
       const content = getMessageContent(msg);
       const cancellations =
         autoCancelPendingToolCalls !== false
@@ -288,6 +335,96 @@ const useStreamThreadRuntime = (
         ],
       });
     },
+    onReload: async (parentId, config) => {
+      const stagedRun = getStagedRun(parentId);
+      if (stagedRun) {
+        const promotedIds = new Set<string>();
+        for (const message of stagedRun.messages) {
+          if (!message.id) continue;
+          promotedIds.add(message.id);
+          stagedMessagesRef.current.delete(message.id);
+        }
+        stagedBaseMessagesRef.current = null;
+        if (stagedMessagesRef.current.size > 0) {
+          const nextMessages = visibleMessagesRef.current.filter(
+            (message) => !message.id || !promotedIds.has(message.id),
+          );
+          visibleMessagesRef.current = nextMessages;
+          setStagedMessages(nextMessages);
+        } else {
+          setStagedMessages(null);
+        }
+        await stream.submit(
+          {
+            [messagesKey]: stagedRun.messages.map((message) => ({
+              id: message.id,
+              type: "human",
+              content: message.content,
+            })),
+          },
+          runConfigToSubmitOptions(config.runConfig ?? stagedRun.runConfig),
+        );
+        return;
+      }
+
+      const threadId = externalId;
+      if (!threadId || parentId == null) return;
+      const s = streamRef.current;
+      const checkpointId = await resolveForkCheckpoint(
+        s.client,
+        threadId,
+        s.messages as readonly LangChainBaseMessage[],
+        parentId,
+        config.sourceId,
+        s[STREAM_CONTROLLER]?.messageMetadataStore?.getSnapshot?.(),
+        messagesKey,
+      );
+      if (!checkpointId) return;
+      await s.submit(null, {
+        forkFrom: checkpointId,
+        ...runConfigToSubmitOptions(config.runConfig),
+      });
+    },
+    onEdit: async (message) => {
+      if (!(message.startRun ?? message.role === "user")) {
+        const truncated = truncateLangChainBaseMessages(
+          threadMessagesRef.current,
+          message.parentId,
+        );
+        const stagedMessage = toStagedHumanMessage(message);
+        stagedMessagesRef.current.set(stagedMessage.id, {
+          message: stagedMessage,
+          runConfig: message.runConfig,
+        });
+        stagedBaseMessagesRef.current = truncated;
+        const nextMessages = [...truncated, stagedMessage];
+        visibleMessagesRef.current = nextMessages;
+        setStagedMessages(nextMessages);
+        return;
+      }
+
+      const threadId = externalId;
+      if (!threadId) return;
+      const s = streamRef.current;
+      const checkpointId = await resolveForkCheckpoint(
+        s.client,
+        threadId,
+        s.messages as readonly LangChainBaseMessage[],
+        message.parentId,
+        message.sourceId,
+        s[STREAM_CONTROLLER]?.messageMetadataStore?.getSnapshot?.(),
+        messagesKey,
+      );
+      if (!checkpointId) return;
+      const content = getMessageContent(message);
+      await s.submit(
+        { [messagesKey]: [{ type: "human", content }] },
+        {
+          forkFrom: checkpointId,
+          ...runConfigToSubmitOptions(message.runConfig),
+        },
+      );
+    },
     onCancel:
       unstable_allowCancellation !== false
         ? async () => {
@@ -329,6 +466,7 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
     unstable_threadListAdapter,
     create,
     delete: deleteFn,
+    onThreadIdChange,
     ...options
   } = rawOptions;
 
@@ -348,111 +486,6 @@ export const useStreamRuntime = (rawOptions: UseStreamRuntimeOptions) => {
     },
     adapter,
     allowNesting: true,
+    onThreadIdChange,
   });
 };
-
-/**
- * Read the current LangGraph interrupt state from the runtime extras.
- */
-export const useLangChainInterruptState = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return undefined;
-    return asLangChainRuntimeExtras(extras).interrupt;
-  });
-};
-
-/** Read the last run/hydration error from the runtime extras. */
-export const useLangChainError = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return undefined;
-    return asLangChainRuntimeExtras(extras).error;
-  });
-};
-
-/**
- * Read the root tool calls assembled by `useStream` from the `tools`
- * channel. Defaults to an empty array, so consumers can `.map` without
- * a guard. Useful for rendering pending/streamed tool calls and
- * approval UIs.
- */
-export const useLangChainToolCalls = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return EMPTY_TOOL_CALLS;
-    return asLangChainRuntimeExtras(extras).toolCalls ?? EMPTY_TOOL_CALLS;
-  });
-};
-
-/**
- * Returns a function to submit raw state updates to the LangGraph agent,
- * bypassing the normal message flow. Useful for sending interrupt resume
- * commands.
- */
-export const useLangChainSubmit = () => {
-  const aui = useAui();
-  return (
-    values: Record<string, unknown> | null | undefined,
-    options?: Record<string, unknown>,
-  ) => {
-    const extras = aui.thread().getState().extras;
-    const { submit } = asLangChainRuntimeExtras(extras);
-    return submit(values, options);
-  };
-};
-
-/**
- * Submit a list of LangChain-shaped messages on the current thread.
- * Parity helper for migrating from `useLangGraphSend`. Routes to
- * `useStream().submit({ [messagesKey]: messages }, options)`.
- */
-export const useLangChainSend = () => {
-  const aui = useAui();
-  return (
-    messages: readonly LangChainBaseMessage[],
-    options?: Record<string, unknown>,
-  ) => {
-    const { submit, messagesKey } = asLangChainRuntimeExtras(
-      aui.thread().getState().extras,
-    );
-    return submit({ [messagesKey]: messages }, options);
-  };
-};
-
-/**
- * Submit a `useStream` command (e.g. interrupt resume). Parity helper
- * for migrating from `useLangGraphSendCommand`. Note that v1's command
- * shape (`{ resume?, goto?, update? }`) differs from the legacy
- * `{ resume: string }` form — to carry a payload, use the input or
- * `stream.respond` instead.
- */
-export const useLangChainSendCommand = () => {
-  const submit = useLangChainSubmit();
-  return (command: Record<string, unknown>) => submit(null, { command });
-};
-
-/**
- * Read a custom LangGraph state key from the current thread. Mirrors
- * `useStream().values[key]` from `@langchain/react` and updates when the
- * stream emits new state.
- *
- * @example
- * ```tsx
- * const todos = useLangChainState<Todo[]>("todos");
- * const files = useLangChainState<Record<string, string>>("files", {});
- * ```
- */
-export function useLangChainState<T>(key: string): T | undefined;
-export function useLangChainState<T>(key: string, defaultValue: T): T;
-export function useLangChainState<T>(
-  key: string,
-  defaultValue?: T,
-): T | undefined {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return defaultValue;
-    const value = asLangChainRuntimeExtras(extras).values[key] as T | undefined;
-    return value !== undefined ? value : defaultValue;
-  });
-}