assistant-stream 0.3.13 → 0.3.14

package/src/resumable/ResumableStreamContext.test.ts

@@ -0,0 +1,274 @@
+import { describe, expect, it } from "vitest";
+import { createResumableStreamContext } from "./ResumableStreamContext";
+import { ResumableStreamError } from "./errors";
+import { createInMemoryResumableStreamStore } from "./stores/InMemoryResumableStreamStore";
+
+const enc = new TextEncoder();
+const dec = new TextDecoder();
+
+const bytes = (s: string): Uint8Array => enc.encode(s);
+
+async function collect(stream: ReadableStream<Uint8Array>): Promise<string> {
+  const reader = stream.getReader();
+  let out = "";
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) return out;
+      out += dec.decode(value, { stream: true });
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+function makeStringStream(parts: string[]): ReadableStream<Uint8Array> {
+  return new ReadableStream<Uint8Array>({
+    async start(controller) {
+      for (const part of parts) {
+        controller.enqueue(bytes(part));
+        await Promise.resolve();
+      }
+      controller.close();
+    },
+  });
+}
+
+describe("createResumableStreamContext", () => {
+  it("producer caller receives full byte stream", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const stream = await ctx.run("a", () =>
+      makeStringStream(["hello ", "world"]),
+    );
+    expect(await collect(stream)).toBe("hello world");
+  });
+
+  it("second caller becomes consumer and receives identical bytes", async () => {
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({ store });
+
+    const producerStream = await ctx.run("a", () =>
+      makeStringStream(["one ", "two ", "three"]),
+    );
+    const consumerStream = await ctx.run("a", () =>
+      makeStringStream(["should-not-run"]),
+    );
+
+    const [a, b] = await Promise.all([
+      collect(producerStream),
+      collect(consumerStream),
+    ]);
+    expect(a).toBe("one two three");
+    expect(b).toBe("one two three");
+  });
+
+  it("late consumer after done replays via resume", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const producer = await ctx.run("a", () =>
+      makeStringStream(["alpha", "beta", "gamma"]),
+    );
+    expect(await collect(producer)).toBe("alphabetagamma");
+
+    const replay = await ctx.resume("a");
+    expect(replay).not.toBeNull();
+    expect(await collect(replay!)).toBe("alphabetagamma");
+  });
+
+  it("resume returns null for missing streams", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    expect(await ctx.resume("nope")).toBeNull();
+  });
+
+  it("requireResume throws ResumableStreamError for missing streams", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    await expect(ctx.requireResume("nope")).rejects.toBeInstanceOf(
+      ResumableStreamError,
+    );
+    await expect(ctx.requireResume("nope")).rejects.toMatchObject({
+      code: "missing",
+    });
+  });
+
+  it("requireResume returns the replay stream when it exists", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const producer = await ctx.run("a", () => makeStringStream(["hi"]));
+    expect(await collect(producer)).toBe("hi");
+
+    const replay = await ctx.requireResume("a");
+    expect(await collect(replay)).toBe("hi");
+  });
+
+  it("status tracks lifecycle", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    expect(await ctx.status("a")).toBe("missing");
+    const stream = await ctx.run("a", () => makeStringStream(["x"]));
+    await collect(stream);
+    expect(await ctx.status("a")).toBe("done");
+  });
+
+  it("delete removes stream state", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const stream = await ctx.run("a", () => makeStringStream(["x"]));
+    await collect(stream);
+    expect(await ctx.status("a")).toBe("done");
+    await ctx.delete("a");
+    expect(await ctx.status("a")).toBe("missing");
+  });
+
+  it("producer keeps writing after the local consumer cancels", async () => {
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({ store });
+
+    let producerEmitted = 0;
+    const slowStream = new ReadableStream<Uint8Array>({
+      async start(controller) {
+        for (let i = 0; i < 5; i++) {
+          await new Promise((r) => setTimeout(r, 5));
+          controller.enqueue(bytes(`chunk${i};`));
+          producerEmitted += 1;
+        }
+        controller.close();
+      },
+    });
+
+    const stream = await ctx.run("a", () => slowStream);
+    const reader = stream.getReader();
+    const first = await reader.read();
+    expect(first.done).toBe(false);
+    await reader.cancel();
+
+    while (producerEmitted < 5) {
+      await new Promise((r) => setTimeout(r, 5));
+    }
+    while ((await ctx.status("a")) === "streaming") {
+      await new Promise((r) => setTimeout(r, 5));
+    }
+
+    const replay = await ctx.resume("a");
+    expect(replay).not.toBeNull();
+    expect(await collect(replay!)).toBe("chunk0;chunk1;chunk2;chunk3;chunk4;");
+  });
+
+  it("propagates producer errors to consumers", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const failing = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(bytes("partial;"));
+        controller.error(new Error("oops"));
+      },
+    });
+    const stream = await ctx.run("a", () => failing);
+    await expect(collect(stream)).rejects.toThrow("oops");
+    expect(await ctx.status("a")).toBe("error");
+  });
+
+  it("waitUntil receives the producer task promise", async () => {
+    const promises: Promise<unknown>[] = [];
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+      waitUntil: (p) => promises.push(p),
+    });
+    const stream = await ctx.run("a", () => makeStringStream(["x"]));
+    expect(await collect(stream)).toBe("x");
+    expect(promises.length).toBe(1);
+    await Promise.all(promises);
+  });
+
+  it("onAcquire fires for both producer and consumer roles", async () => {
+    const calls: Array<{ id: string; role: string }> = [];
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+      onAcquire: (id, role) => calls.push({ id, role }),
+    });
+    const producer = await ctx.run("a", () => makeStringStream(["x"]));
+    const consumer = await ctx.run("a", () => makeStringStream(["unused"]));
+    await Promise.all([collect(producer), collect(consumer)]);
+    expect(calls).toEqual([
+      { id: "a", role: "producer" },
+      { id: "a", role: "consumer" },
+    ]);
+  });
+
+  it("onAppend fires per appended chunk with byteLength", async () => {
+    const calls: Array<{ id: string; byteLength: number }> = [];
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+      onAppend: (id, byteLength) => calls.push({ id, byteLength }),
+    });
+    const stream = await ctx.run("a", () => makeStringStream(["ab", "cde"]));
+    expect(await collect(stream)).toBe("abcde");
+    expect(calls).toEqual([
+      { id: "a", byteLength: 2 },
+      { id: "a", byteLength: 3 },
+    ]);
+  });
+
+  it("onFinalize fires on successful completion", async () => {
+    const calls: Array<{
+      id: string;
+      status: "done" | "error";
+      error: string | undefined;
+    }> = [];
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+      onFinalize: (id, status, error) => calls.push({ id, status, error }),
+    });
+    const stream = await ctx.run("a", () => makeStringStream(["x"]));
+    await collect(stream);
+    expect(calls).toEqual([{ id: "a", status: "done", error: undefined }]);
+  });
+
+  it("onFinalize fires with error status when producer fails", async () => {
+    const calls: Array<{
+      id: string;
+      status: "done" | "error";
+      error: string | undefined;
+    }> = [];
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+      onFinalize: (id, status, error) => calls.push({ id, status, error }),
+    });
+    const failing = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.error(new Error("boom"));
+      },
+    });
+    const stream = await ctx.run("a", () => failing);
+    await expect(collect(stream)).rejects.toThrow("boom");
+    expect(calls).toEqual([{ id: "a", status: "error", error: "boom" }]);
+  });
+
+  it("onError fires when the producer task throws", async () => {
+    const errors: Array<{ id: string; error: unknown }> = [];
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+      onError: (id, error) => errors.push({ id, error }),
+    });
+    const failing = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.error(new Error("boom"));
+      },
+    });
+    const stream = await ctx.run("a", () => failing);
+    await expect(collect(stream)).rejects.toThrow("boom");
+    expect(errors).toHaveLength(1);
+    expect(errors[0]!.id).toBe("a");
+    expect((errors[0]!.error as Error).message).toBe("boom");
+  });
+});

package/src/resumable/ResumableStreamContext.ts

@@ -0,0 +1,187 @@
+import { ResumableStreamError } from "./errors";
+import type {
+  ResumableStreamRole,
+  ResumableStreamStatus,
+  ResumableStreamStore,
+} from "./types";
+
+export type ResumableStreamContextOptions = {
+  readonly store: ResumableStreamStore;
+  readonly ttlMs?: number;
+  /**
+   * Required on serverless runtimes that terminate background work when the
+   * request handler returns (Vercel, Cloudflare). Pass `after` from
+   * `next/server` so the producer task outlives the response.
+   */
+  readonly waitUntil?: (promise: Promise<unknown>) => void;
+  readonly onAcquire?: (streamId: string, role: ResumableStreamRole) => void;
+  readonly onAppend?: (streamId: string, byteLength: number) => void;
+  readonly onFinalize?: (
+    streamId: string,
+    status: "done" | "error",
+    error?: string,
+  ) => void;
+  readonly onError?: (streamId: string, error: unknown) => void;
+};
+
+export interface ResumableStreamContext {
+  /** Producer or consumer entrypoint. Atomically elects the role. */
+  run(
+    streamId: string,
+    makeStream: () => ReadableStream<Uint8Array>,
+  ): Promise<ReadableStream<Uint8Array>>;
+
+  /** Returns `null` when the stream does not exist. */
+  resume(streamId: string): Promise<ReadableStream<Uint8Array> | null>;
+
+  /** Throws `ResumableStreamError("missing")` when the stream does not exist. */
+  requireResume(streamId: string): Promise<ReadableStream<Uint8Array>>;
+
+  status(streamId: string): Promise<ResumableStreamStatus>;
+
+  delete(streamId: string): Promise<void>;
+}
+
+export function createResumableStreamContext(
+  options: ResumableStreamContextOptions,
+): ResumableStreamContext {
+  const { store, waitUntil, onAcquire, onAppend, onFinalize, onError } =
+    options;
+  const acquireOptions =
+    options.ttlMs !== undefined ? { ttlMs: options.ttlMs } : undefined;
+
+  return {
+    async run(streamId, makeStream) {
+      const role = await store.acquire(streamId, acquireOptions);
+      onAcquire?.(streamId, role);
+      if (role === "producer") {
+        startProducerTask(store, streamId, makeStream, {
+          waitUntil,
+          onAppend,
+          onFinalize,
+          onError,
+        });
+      }
+      return readFromStore(store, streamId);
+    },
+
+    async resume(streamId) {
+      const status = await store.status(streamId);
+      if (status === "missing") return null;
+      return readFromStore(store, streamId);
+    },
+
+    async requireResume(streamId) {
+      const status = await store.status(streamId);
+      if (status === "missing") {
+        throw new ResumableStreamError(
+          "missing",
+          `resumable stream not found: ${streamId}`,
+        );
+      }
+      return readFromStore(store, streamId);
+    },
+
+    async status(streamId) {
+      return store.status(streamId);
+    },
+
+    async delete(streamId) {
+      await store.delete(streamId);
+    },
+  };
+}
+
+type ProducerHooks = {
+  readonly waitUntil: ((promise: Promise<unknown>) => void) | undefined;
+  readonly onAppend:
+    | ((streamId: string, byteLength: number) => void)
+    | undefined;
+  readonly onFinalize:
+    | ((streamId: string, status: "done" | "error", error?: string) => void)
+    | undefined;
+  readonly onError: ((streamId: string, error: unknown) => void) | undefined;
+};
+
+function startProducerTask(
+  store: ResumableStreamStore,
+  streamId: string,
+  makeStream: () => ReadableStream<Uint8Array>,
+  hooks: ProducerHooks,
+): void {
+  const { waitUntil, onAppend, onFinalize, onError } = hooks;
+  const task = (async () => {
+    let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
+    let cancelled = false;
+    try {
+      reader = makeStream().getReader();
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        await store.append(streamId, value);
+        onAppend?.(streamId, value.byteLength);
+      }
+      await store.finalize(streamId, "done");
+      onFinalize?.(streamId, "done");
+    } catch (err) {
+      cancelled = true;
+      onError?.(streamId, err);
+      const message = err instanceof Error ? err.message : String(err);
+      try {
+        await reader?.cancel(err);
+      } catch (cancelErr) {
+        console.error("resumable stream reader cancel failed:", cancelErr);
+      }
+      try {
+        await store.finalize(streamId, "error", message);
+        onFinalize?.(streamId, "error", message);
+      } catch (finalizeErr) {
+        console.error("resumable stream finalize failed:", finalizeErr);
+      }
+    } finally {
+      if (!cancelled) reader?.releaseLock();
+    }
+  })();
+
+  if (waitUntil) waitUntil(task);
+  task.catch((err) => {
+    console.error("resumable producer task failed:", err);
+  });
+}
+
+function readFromStore(
+  store: ResumableStreamStore,
+  streamId: string,
+): ReadableStream<Uint8Array> {
+  const ac = new AbortController();
+  let iterator: AsyncIterator<{ chunk: Uint8Array }> | undefined;
+
+  return new ReadableStream<Uint8Array>({
+    start() {
+      iterator = store.read(streamId, "", ac.signal)[Symbol.asyncIterator]();
+    },
+    async pull(controller) {
+      try {
+        if (!iterator) return;
+        const { done, value } = await iterator.next();
+        if (done) {
+          controller.close();
+          return;
+        }
+        controller.enqueue(value.chunk);
+      } catch (err) {
+        // the platform never calls cancel() on errored streams, so unwind
+        // the store iterator and abort the signal explicitly.
+        ac.abort();
+        try {
+          await iterator?.return?.();
+        } catch {}
+        controller.error(err);
+      }
+    },
+    cancel() {
+      ac.abort();
+      iterator?.return?.().catch(() => {});
+    },
+  });
+}

package/src/resumable/__tests__/integration.test.ts

@@ -0,0 +1,159 @@
+import { describe, expect, it } from "vitest";
+import {
+  RESUMABLE_STREAM_ID_HEADER,
+  createResumableAssistantStreamResponse,
+  createResumeAssistantStreamResponse,
+} from "../createResumableAssistantStreamResponse";
+import { createResumableStreamContext } from "../ResumableStreamContext";
+import { createInMemoryResumableStreamStore } from "../stores/InMemoryResumableStreamStore";
+import {
+  AssistantTransportDecoder,
+  AssistantTransportEncoder,
+} from "../../core/serialization/assistant-transport/AssistantTransport";
+import { DataStreamDecoder } from "../../core/serialization/data-stream/DataStream";
+import { AssistantMessageAccumulator } from "../../core/accumulators/assistant-message-accumulator";
+import type { AssistantStreamChunk } from "../../core/AssistantStreamChunk";
+import type { AssistantMessage } from "../../core/utils/types";
+
+async function decodeViaDataStream(body: ReadableStream<Uint8Array>) {
+  const stream = body.pipeThrough(new DataStreamDecoder());
+  const out: string[] = [];
+  for await (const chunk of stream) {
+    if (chunk.type === "text-delta") out.push(chunk.textDelta);
+  }
+  return out.join("");
+}
+
+async function accumulate(
+  body: ReadableStream<Uint8Array>,
+  decoder: TransformStream<Uint8Array, AssistantStreamChunk>,
+): Promise<AssistantMessage> {
+  const messages = body
+    .pipeThrough(decoder)
+    .pipeThrough(new AssistantMessageAccumulator());
+  let last: AssistantMessage | undefined;
+  for await (const msg of messages) last = msg;
+  if (!last) throw new Error("accumulator yielded no messages");
+  return last;
+}
+
+describe("resumable integration", () => {
+  it("ten concurrent consumers see identical bytes for a 50 chunk stream", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+
+    const producer = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "fanout",
+      callback: async (controller) => {
+        for (let i = 0; i < 50; i++) {
+          controller.appendText(`chunk-${i}`);
+          await Promise.resolve();
+        }
+      },
+    });
+
+    const consumers = await Promise.all(
+      Array.from({ length: 10 }, async () => {
+        const r = await ctx.resume("fanout");
+        if (!r) throw new Error("resume returned null");
+        return new Response(r).arrayBuffer();
+      }),
+    );
+    const producerBytes = await new Response(producer.body).arrayBuffer();
+
+    const expected = new Uint8Array(producerBytes);
+    for (const c of consumers) {
+      expect(new Uint8Array(c)).toEqual(expected);
+    }
+
+    const text = await decodeViaDataStream(
+      new Response(new Uint8Array(producerBytes)).body!,
+    );
+    expect(text).toBe(
+      Array.from({ length: 50 }, (_, i) => `chunk-${i}`).join(""),
+    );
+  });
+
+  it("non-ASCII payload round-trips byte-for-byte through resume", async () => {
+    const allCodePoints = Array.from({ length: 256 }, (_, i) =>
+      String.fromCodePoint(i),
+    ).join("");
+    const emoji = "hello 🌍 🚀 ✨";
+    const payload = allCodePoints + emoji;
+
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const producer = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "binary",
+      callback: (controller) => controller.appendText(payload),
+    });
+    const producerBytes = new Uint8Array(
+      await new Response(producer.body).arrayBuffer(),
+    );
+
+    const replay = await createResumeAssistantStreamResponse({
+      context: ctx,
+      streamId: "binary",
+    });
+    const replayBytes = new Uint8Array(
+      await new Response(replay.body).arrayBuffer(),
+    );
+    expect(replayBytes).toEqual(producerBytes);
+    expect(replay.headers.get(RESUMABLE_STREAM_ID_HEADER)).toBe("binary");
+
+    const decoded = await decodeViaDataStream(new Response(replayBytes).body!);
+    expect(decoded).toBe(payload);
+  });
+
+  it("AssistantTransportEncoder round-trips through resume into an identical message", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const producer = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "transport",
+      encoder: () => new AssistantTransportEncoder(),
+      callback: (controller) => {
+        controller.appendText("processing... ");
+        const tool = controller.addToolCallPart({
+          toolName: "lookup",
+          toolCallId: "tool-1",
+          args: { query: "weather" },
+          response: { result: { temperature: 72 } },
+        });
+        tool.close();
+        controller.appendText("done");
+      },
+    });
+    const producerMessage = await accumulate(
+      producer.body!,
+      new AssistantTransportDecoder(),
+    );
+
+    const replay = await createResumeAssistantStreamResponse({
+      context: ctx,
+      streamId: "transport",
+      encoder: () => new AssistantTransportEncoder(),
+    });
+    expect(replay.headers.get("content-type")).toBe("text/event-stream");
+    const replayMessage = await accumulate(
+      replay.body!,
+      new AssistantTransportDecoder(),
+    );
+
+    expect(replayMessage.parts).toEqual(producerMessage.parts);
+    expect(replayMessage.status).toEqual(producerMessage.status);
+    const toolPart = replayMessage.parts.find((p) => p.type === "tool-call");
+    expect(toolPart).toBeDefined();
+    expect(toolPart).toMatchObject({
+      toolName: "lookup",
+      toolCallId: "tool-1",
+      args: { query: "weather" },
+      result: { temperature: 72 },
+    });
+  });
+});

package/src/resumable/constants.ts

@@ -0,0 +1 @@
+export const DEFAULT_TTL_MS = 24 * 60 * 60 * 1000;