@assistant-ui/core 0.3.13 → 0.3.15

package/src/react/runtimes/assistant-transport/replayBoundaryStream.test.ts

@@ -0,0 +1,455 @@
+// @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),
+    replayContentLength === undefined
+      ? undefined
+      : {
+          headers: {
+            [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 body = createBody(["hi"]);
+    const streamPromise = createReplayBoundaryStream(
+      new Response(body, {
+        headers: { [REPLAY_CONTENT_LENGTH_HEADER]: "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);
+    expect(body.locked).toBe(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);
+    expect(body.locked).toBe(false);
+  });
+
+  it("releases the response body reader after a read failure", async () => {
+    const { waitForRender, releaseNext } = createRenderWait();
+    const error = new Error("stream failed");
+    const body = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.error(error);
+      },
+    });
+    const streamPromise = createReplayBoundaryStream(
+      new Response(body, { headers: { [REPLAY_CONTENT_LENGTH_HEADER]: "10" } }),
+      { setReplaying: vi.fn(), waitForRender },
+    );
+
+    await releaseNext();
+    const stream = await streamPromise;
+
+    await expect(stream.getReader().read()).rejects.toBe(error);
+    expect(body.locked).toBe(false);
+  });
+
+  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/react/runtimes/assistant-transport/replayBoundaryStream.ts

@@ -0,0 +1,172 @@
+"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;
+  let readerCleanup: Promise<void> | undefined;
+
+  const releaseReader = () => {
+    if (readerCleanup) return readerCleanup;
+    reader.releaseLock();
+    readerCleanup = Promise.resolve();
+    return readerCleanup;
+  };
+
+  const cancelReader = (reason?: unknown) => {
+    if (readerCleanup) return readerCleanup;
+    readerCleanup = (async () => {
+      try {
+        await reader.cancel(reason);
+      } finally {
+        reader.releaseLock();
+      }
+    })();
+    return readerCleanup;
+  };
+
+  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) {
+      try {
+        const { done, value } = await reader.read();
+
+        if (done) {
+          await releaseReader();
+          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));
+      } catch (error) {
+        await cancelReader(error).catch(() => {});
+        throw error;
+      }
+    },
+    async cancel(reason) {
+      const wasFinished = replayFinished;
+      replayFinished = true;
+      if (!wasFinished) setReplaying(false);
+      await cancelReader(reason);
+    },
+  });
+};