@assistant-ui/react 0.14.13 → 0.14.15

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

@@ -0,0 +1,426 @@
+// @vitest-environment jsdom
+
+import {
+  AssistantMessageAccumulator,
+  DataStreamDecoder,
+} from "assistant-stream";
+import { act, renderHook } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import {
+  createReplayBoundaryStream,
+  REPLAY_CONTENT_LENGTH_HEADER,
+  useReplayRenderWait,
+} from "./replayBoundaryStream";
+
+const encoder = new TextEncoder();
+const decoder = new TextDecoder();
+
+const createBody = (chunks: readonly string[]) =>
+  new ReadableStream<Uint8Array>({
+    start(controller) {
+      for (const chunk of chunks) {
+        controller.enqueue(encoder.encode(chunk));
+      }
+      controller.close();
+    },
+  });
+
+const createResponse = (
+  chunks: readonly string[],
+  replayContentLength?: number | string,
+) =>
+  new Response(createBody(chunks), {
+    headers:
+      replayContentLength === undefined
+        ? undefined
+        : { [REPLAY_CONTENT_LENGTH_HEADER]: String(replayContentLength) },
+  });
+
+const createRenderWait = () => {
+  const pending: Array<() => void> = [];
+  const waitForRender = vi.fn(
+    () =>
+      new Promise<void>((resolve) => {
+        pending.push(resolve);
+      }),
+  );
+
+  const releaseNext = async () => {
+    for (let i = 0; pending.length === 0 && i < 10; i++) {
+      await Promise.resolve();
+    }
+
+    const resolve = pending.shift();
+    expect(resolve).toBeDefined();
+    resolve!();
+    await Promise.resolve();
+  };
+
+  return { waitForRender, releaseNext };
+};
+
+const readText = async (stream: ReadableStream<Uint8Array>) => {
+  const reader = stream.getReader();
+  const chunks: string[] = [];
+
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    chunks.push(decoder.decode(value, { stream: true }));
+  }
+  chunks.push(decoder.decode());
+
+  return chunks.join("");
+};
+
+describe("useReplayRenderWait", () => {
+  it("resolves after its own render ticket commits", async () => {
+    vi.useFakeTimers();
+
+    try {
+      const { result } = renderHook(() => useReplayRenderWait());
+
+      let resolved = false;
+      const wait = result.current().then(() => {
+        resolved = true;
+      });
+
+      await Promise.resolve();
+      expect(resolved).toBe(false);
+
+      await act(async () => {
+        vi.runOnlyPendingTimers();
+      });
+      await wait;
+
+      expect(resolved).toBe(true);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+});
+
+describe("createReplayBoundaryStream", () => {
+  it("short-circuits responses without a valid replay content length", async () => {
+    const setReplaying = vi.fn();
+    const waitForRender = vi.fn();
+
+    for (const replayContentLength of [undefined, "abc", "3.5", "-1"]) {
+      setReplaying.mockClear();
+      waitForRender.mockClear();
+
+      const body = await createReplayBoundaryStream(
+        createResponse(["live"], replayContentLength),
+        {
+          setReplaying,
+          waitForRender,
+        },
+      );
+
+      expect(await readText(body)).toBe("live");
+      expect(setReplaying).not.toHaveBeenCalled();
+      expect(waitForRender).not.toHaveBeenCalled();
+    }
+  });
+
+  it("pauses at the replay boundary before releasing live bytes", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    const replayPrefix = '0:"hi"\nb:{"toolCallId":"call-1","toolName":"te';
+    const liveSuffix = 'st"}\n';
+    const replayContentLength = encoder.encode(replayPrefix).byteLength;
+
+    const streamPromise = createReplayBoundaryStream(
+      createResponse([replayPrefix + liveSuffix], replayContentLength),
+      { setReplaying, waitForRender },
+    );
+
+    expect(setReplaying).toHaveBeenCalledWith(true);
+    expect(waitForRender).toHaveBeenCalledTimes(1);
+
+    await releaseNext();
+    const stream = await streamPromise;
+    const reader = stream.getReader();
+
+    await expect(reader.read()).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(replayPrefix),
+    });
+    expect(waitForRender).toHaveBeenCalledTimes(2);
+    expect(setReplaying).toHaveBeenCalledTimes(1);
+
+    let liveReadResolved = false;
+    const liveRead = reader.read().then((read) => {
+      liveReadResolved = true;
+      return read;
+    });
+    await Promise.resolve();
+    expect(liveReadResolved).toBe(false);
+
+    await releaseNext();
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+    expect(waitForRender).toHaveBeenCalledTimes(3);
+    await Promise.resolve();
+    expect(liveReadResolved).toBe(false);
+
+    await releaseNext();
+    await expect(liveRead).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(liveSuffix),
+    });
+  });
+
+  it("pauses when a chunk ends exactly at the replay boundary", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    const replayPrefix = "replay";
+    const liveSuffix = "live";
+    const replayContentLength = encoder.encode(replayPrefix).byteLength;
+
+    const streamPromise = createReplayBoundaryStream(
+      createResponse([replayPrefix, liveSuffix], replayContentLength),
+      { setReplaying, waitForRender },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+    const reader = stream.getReader();
+
+    await expect(reader.read()).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(replayPrefix),
+    });
+    expect(setReplaying).toHaveBeenCalledTimes(1);
+
+    let liveReadResolved = false;
+    const liveRead = reader.read().then((read) => {
+      liveReadResolved = true;
+      return read;
+    });
+    await Promise.resolve();
+    expect(liveReadResolved).toBe(false);
+
+    await releaseNext();
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+    await releaseNext();
+
+    await expect(liveRead).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(liveSuffix),
+    });
+  });
+
+  it("accumulates replay bytes across chunks before splitting live bytes", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    const firstReplayChunk = "part1";
+    const secondReplayChunk = "part2";
+    const firstLiveChunk = "live1";
+    const secondLiveChunk = "live2";
+    const replayContentLength = encoder.encode(
+      firstReplayChunk + secondReplayChunk,
+    ).byteLength;
+
+    const streamPromise = createReplayBoundaryStream(
+      createResponse(
+        [firstReplayChunk, secondReplayChunk + firstLiveChunk, secondLiveChunk],
+        replayContentLength,
+      ),
+      { setReplaying, waitForRender },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+    const reader = stream.getReader();
+
+    await expect(reader.read()).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(firstReplayChunk),
+    });
+    await expect(reader.read()).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(secondReplayChunk),
+    });
+    expect(setReplaying).toHaveBeenCalledTimes(1);
+
+    let liveReadResolved = false;
+    const liveRead = reader.read().then((read) => {
+      liveReadResolved = true;
+      return read;
+    });
+    await Promise.resolve();
+    expect(liveReadResolved).toBe(false);
+
+    await releaseNext();
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+    await releaseNext();
+
+    await expect(liveRead).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(firstLiveChunk),
+    });
+    await expect(reader.read()).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(secondLiveChunk),
+    });
+  });
+
+  it("clears replaying when the stream ends before the boundary", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    const streamPromise = createReplayBoundaryStream(
+      createResponse(["hi"], 10),
+      {
+        setReplaying,
+        waitForRender,
+      },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+    const text = readText(stream);
+    await releaseNext();
+    await releaseNext();
+
+    await expect(text).resolves.toBe("hi");
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+  });
+
+  it("clears replaying when the gated stream is cancelled", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    let cancelled = false;
+    const body = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(encoder.encode("hi"));
+      },
+      cancel() {
+        cancelled = true;
+      },
+    });
+    const streamPromise = createReplayBoundaryStream(
+      new Response(body, { headers: { [REPLAY_CONTENT_LENGTH_HEADER]: "10" } }),
+      { setReplaying, waitForRender },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+    await stream.cancel("done");
+
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+    expect(cancelled).toBe(true);
+  });
+
+  it("does not wait for replay completion after cancellation unblocks a read", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    let cancelled = false;
+    const body = new ReadableStream<Uint8Array>({
+      cancel() {
+        cancelled = true;
+      },
+    });
+    const streamPromise = createReplayBoundaryStream(
+      new Response(body, { headers: { [REPLAY_CONTENT_LENGTH_HEADER]: "10" } }),
+      { setReplaying, waitForRender },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+    const reader = stream.getReader();
+    const read = reader.read();
+
+    await Promise.resolve();
+    expect(waitForRender).toHaveBeenCalledTimes(1);
+
+    await reader.cancel("done");
+    await expect(read).resolves.toMatchObject({ done: true });
+
+    expect(waitForRender).toHaveBeenCalledTimes(1);
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+    expect(cancelled).toBe(true);
+  });
+
+  it("does not clear replaying twice when cancelled during replay completion", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    const replayStr = "replay";
+    const replayContentLength = encoder.encode(replayStr).byteLength;
+
+    const streamPromise = createReplayBoundaryStream(
+      createResponse([replayStr], replayContentLength),
+      { setReplaying, waitForRender },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+    const reader = stream.getReader();
+
+    await expect(reader.read()).resolves.toMatchObject({
+      done: false,
+      value: encoder.encode(replayStr),
+    });
+    expect(waitForRender).toHaveBeenCalledTimes(2);
+
+    const cancel = reader.cancel("done");
+    await Promise.resolve();
+    expect(setReplaying).toHaveBeenCalledTimes(1);
+
+    await releaseNext();
+    expect(setReplaying).toHaveBeenCalledTimes(2);
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+
+    await releaseNext();
+    await cancel;
+    expect(setReplaying).toHaveBeenCalledTimes(2);
+  });
+
+  it("lets data-stream parsing commit replayed text before live tool calls", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
+    const replayPrefix = '0:"hi"\nb:{"toolCallId":"call-1","toolName":"te';
+    const liveSuffix = 'st"}\n';
+    const replayContentLength = encoder.encode(replayPrefix).byteLength;
+
+    const streamPromise = createReplayBoundaryStream(
+      createResponse([replayPrefix + liveSuffix], replayContentLength),
+      { setReplaying, waitForRender },
+    );
+
+    await releaseNext();
+    const messages = (await streamPromise)
+      .pipeThrough(new DataStreamDecoder())
+      .pipeThrough(new AssistantMessageAccumulator({ throttle: true }));
+    const reader = messages.getReader();
+
+    let sawReplayedText = false;
+    while (!sawReplayedText) {
+      const replayedMessage = await reader.read();
+      expect(replayedMessage.done).toBe(false);
+      expect(
+        replayedMessage.value?.parts.some((part) => part.type === "tool-call"),
+      ).toBe(false);
+      sawReplayedText =
+        replayedMessage.value?.parts.some(
+          (part) => part.type === "text" && part.text === "hi",
+        ) ?? false;
+    }
+    expect(setReplaying).toHaveBeenCalledTimes(1);
+
+    const liveMessagePromise = reader.read();
+    await releaseNext();
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
+    await releaseNext();
+
+    let liveMessage = await liveMessagePromise;
+    while (
+      !liveMessage.done &&
+      !liveMessage.value?.parts.some(
+        (part) => part.type === "tool-call" && part.toolName === "test",
+      )
+    ) {
+      liveMessage = await reader.read();
+    }
+    expect(liveMessage.done).toBe(false);
+  });
+});

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,