@assistant-ui/react 0.15.14 → 0.15.16

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

@@ -1,455 +0,0 @@
-// @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/legacy-runtime/runtime-cores/assistant-transport/replayBoundaryStream.ts

@@ -1,172 +0,0 @@
-"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);
-    },
-  });
-};