@assistant-ui/react 0.14.12 → 0.14.14

package/src/legacy-runtime/runtime-cores/assistant-transport/replayBoundaryStream.ts

@@ -0,0 +1,146 @@
+"use client";
+
+import { useCallback, useEffect, useRef, useState } from "react";
+
+export const REPLAY_CONTENT_LENGTH_HEADER = "Aui-Replay-Content-Length";
+
+type ReplayBoundaryStreamOptions = {
+  setReplaying: (value: boolean) => void;
+  waitForRender: () => Promise<void>;
+};
+
+export const useReplayRenderWait = () => {
+  const [renderTicket, setRenderTicket] = useState(0);
+  const mountedRef = useRef(true);
+  const nextTicketRef = useRef(0);
+  const waitersRef = useRef<Array<{ ticket: number; resolve: () => void }>>([]);
+
+  const resolveWaiters = useCallback((committedTicket?: number) => {
+    const pendingWaiters = [];
+
+    for (const waiter of waitersRef.current) {
+      if (committedTicket === undefined || waiter.ticket <= committedTicket) {
+        waiter.resolve();
+      } else {
+        pendingWaiters.push(waiter);
+      }
+    }
+
+    waitersRef.current = pendingWaiters;
+  }, []);
+
+  useEffect(() => {
+    mountedRef.current = true;
+    resolveWaiters(renderTicket);
+  }, [renderTicket, resolveWaiters]);
+
+  useEffect(
+    () => () => {
+      mountedRef.current = false;
+      resolveWaiters();
+    },
+    [resolveWaiters],
+  );
+
+  return useCallback(
+    () =>
+      new Promise<void>((resolve) => {
+        setTimeout(() => {
+          if (!mountedRef.current) {
+            resolve();
+            return;
+          }
+
+          const ticket = nextTicketRef.current + 1;
+          nextTicketRef.current = ticket;
+          waitersRef.current.push({ ticket, resolve });
+          setRenderTicket(ticket);
+        }, 0);
+      }),
+    [],
+  );
+};
+
+const parseReplayContentLength = (headers: Headers): number => {
+  const raw = headers.get(REPLAY_CONTENT_LENGTH_HEADER);
+  if (raw == null) return 0;
+
+  const boundary = Number(raw);
+  return Number.isSafeInteger(boundary) && boundary > 0 ? boundary : 0;
+};
+
+// Gates replay bytes until isReplaying:true renders, then releases live bytes after isReplaying:false renders.
+export const createReplayBoundaryStream = async (
+  response: Response,
+  {
+    setReplaying,
+    waitForRender: waitForReplayRender,
+  }: ReplayBoundaryStreamOptions,
+): Promise<ReadableStream<Uint8Array>> => {
+  const body = response.body as ReadableStream<Uint8Array>;
+  const replayContentLength = parseReplayContentLength(response.headers);
+
+  if (replayContentLength <= 0) {
+    return body;
+  }
+
+  setReplaying(true);
+  await waitForReplayRender();
+
+  const reader = body.getReader();
+  let bytesForwarded = 0;
+  let replayFinished = false;
+
+  const finishReplay = async () => {
+    if (replayFinished) return;
+    replayFinished = true;
+
+    // Let replay bytes drain before rendering live mode, then render live mode before releasing live bytes.
+    await waitForReplayRender();
+    setReplaying(false);
+    await waitForReplayRender();
+  };
+
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      const { done, value } = await reader.read();
+
+      if (done) {
+        await finishReplay();
+        controller.close();
+        return;
+      }
+
+      if (replayFinished) {
+        controller.enqueue(value);
+        return;
+      }
+
+      const nextBytesForwarded = bytesForwarded + value.byteLength;
+
+      if (nextBytesForwarded < replayContentLength) {
+        bytesForwarded = nextBytesForwarded;
+        controller.enqueue(value);
+        return;
+      }
+
+      if (nextBytesForwarded === replayContentLength) {
+        controller.enqueue(value);
+        await finishReplay();
+        return;
+      }
+
+      const replayBytesInChunk = replayContentLength - bytesForwarded;
+
+      controller.enqueue(value.subarray(0, replayBytesInChunk));
+      await finishReplay();
+      controller.enqueue(value.subarray(replayBytesInChunk));
+    },
+    async cancel(reason) {
+      const wasFinished = replayFinished;
+      replayFinished = true;
+      if (!wasFinished) setReplaying(false);
+      await reader.cancel(reason);
+    },
+  });
+};

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime.ts

@@ -27,6 +27,10 @@ import type {
   SendCommandsRequestBody,
 } from "./types";
 import { useCommandQueue } from "./commandQueue";
+import {
+  createReplayBoundaryStream,
+  useReplayRenderWait,
+} from "./replayBoundaryStream";
 import { useRunManager } from "./runManager";
 import { useConvertedState } from "./useConvertedState";
 import type { ToolExecutionStatus } from "@assistant-ui/core";
@@ -116,6 +120,8 @@ const useAssistantTransportThreadRuntime = <T>(
   const agentStateRef = useRef(options.initialState);
   const [, rerender] = useState(0);
   const resumeFlagRef = useRef(false);
+  const [isReplaying, setIsReplaying] = useState(false);
+  const waitForReplayRender = useReplayRenderWait();
   const parentIdRef = useRef<string | null | undefined>(undefined);
   const commandQueue = useCommandQueue({
     onQueue: () => runManager.schedule(),
@@ -127,6 +133,7 @@ const useAssistantTransportThreadRuntime = <T>(
     onRun: async (signal: AbortSignal) => {
       const isResume = resumeFlagRef.current;
       resumeFlagRef.current = false;
+      setIsReplaying(false);
       const commands: QueuedCommand[] = isResume ? [] : commandQueue.flush();
       if (commands.length === 0 && !isResume)
         throw new Error("No commands to send");
@@ -182,6 +189,11 @@ const useAssistantTransportThreadRuntime = <T>(
         throw new Error("Response body is null");
       }
 
+      const body = await createReplayBoundaryStream(response, {
+        setReplaying: setIsReplaying,
+        waitForRender: waitForReplayRender,
+      });
+
       // Select decoder based on protocol option
       const protocol = options.protocol ?? "data-stream";
       const decoder =
@@ -190,7 +202,7 @@ const useAssistantTransportThreadRuntime = <T>(
           : new DataStreamDecoder();
 
       let err: string | undefined;
-      const stream = response.body.pipeThrough(decoder).pipeThrough(
+      const stream = body.pipeThrough(decoder).pipeThrough(
         new AssistantMessageAccumulator({
           initialMessage: createInitialMessage({
             unstable_state:
@@ -223,6 +235,7 @@ const useAssistantTransportThreadRuntime = <T>(
     },
     onFinish: options.onFinish,
     onCancel: () => {
+      setIsReplaying(false);
       const cmds = [
         ...commandQueue.state.inTransit,
         ...commandQueue.state.queued,
@@ -239,6 +252,7 @@ const useAssistantTransportThreadRuntime = <T>(
       });
     },
     onError: async (error) => {
+      setIsReplaying(false);
       const inTransitCmds = [...commandQueue.state.inTransit];
       const queuedCmds = [...commandQueue.state.queued];
 
@@ -288,6 +302,7 @@ const useAssistantTransportThreadRuntime = <T>(
     messages: converted.messages,
     state: converted.state,
     isRunning: converted.isRunning,
+    isLoading: isReplaying,
     adapters: options.adapters,
     unstable_enableToolInvocations: true,
     setToolStatuses,

package/src/unstable/useMentionAdapter.ts

@@ -45,7 +45,7 @@ export type Unstable_UseMentionAdapterOptions = {
   /** Categorized mentions for drill-down navigation. */
   readonly categories?: readonly Unstable_MentionCategory[];
   /**
-   * How tools registered via `useAssistantTool` integrate.
+   * How tools registered in model context integrate.
    * - `false`: exclude.
    * - `true`: include (default when no `items`/`categories`; as a category
    *   if `categories` is set, flat otherwise).
@@ -75,7 +75,7 @@ export type Unstable_MentionDirective = {
  * @deprecated Under active development and might change without notice.
  *
  * Creates a spreadable `{ adapter, directive }` bundle for `@` mentions.
- * Supports tools registered via `useAssistantTool`, explicit items, or both —
+ * Supports tools registered in model context, explicit items, or both —
  * flat or categorized.
  *
  * @example