assistant-stream 0.3.13 → 0.3.15

package/src/resumable/createResumableAssistantStreamResponse.test.ts

@@ -0,0 +1,243 @@
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import {
+  RESUMABLE_STREAM_ID_HEADER,
+  createResumableAssistantStreamResponse,
+  createResumeAssistantStreamResponse,
+} from "./createResumableAssistantStreamResponse";
+import { createResumableStreamContext } from "./ResumableStreamContext";
+import { ResumableStreamError } from "./errors";
+import { createInMemoryResumableStreamStore } from "./stores/InMemoryResumableStreamStore";
+import { DataStreamDecoder } from "../core/serialization/data-stream/DataStream";
+import { AssistantTransportEncoder } from "../core/serialization/assistant-transport/AssistantTransport";
+import { AssistantMessageAccumulator } from "../core/accumulators/assistant-message-accumulator";
+import type { AssistantStreamChunk } from "../core/AssistantStreamChunk";
+
+async function decodeAssistantText(response: Response): Promise<string> {
+  const chunks: AssistantStreamChunk[] = [];
+  const stream = response.body!.pipeThrough(new DataStreamDecoder());
+  for await (const chunk of stream) chunks.push(chunk);
+
+  const acc = new AssistantMessageAccumulator();
+  const messageStream = new ReadableStream<AssistantStreamChunk>({
+    start(c) {
+      for (const ch of chunks) c.enqueue(ch);
+      c.close();
+    },
+  }).pipeThrough(acc);
+
+  const reader = messageStream.getReader();
+  let last:
+    | { parts: ReadonlyArray<{ type: string; text?: string }> }
+    | undefined;
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    last = value;
+  }
+  return (
+    last?.parts
+      .filter((p): p is { type: "text"; text: string } => p.type === "text")
+      .map((p) => p.text)
+      .join("") ?? ""
+  );
+}
+
+describe("createResumableAssistantStreamResponse", () => {
+  it("produces a response that decodes through DataStreamDecoder", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const response = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "s1",
+      callback: (controller) => {
+        controller.appendText("hello ");
+        controller.appendText("world");
+      },
+    });
+
+    expect(response.headers.get(RESUMABLE_STREAM_ID_HEADER)).toBe("s1");
+    expect(response.headers.get("content-type")).toContain("text/plain");
+    expect(await decodeAssistantText(response)).toBe("hello world");
+  });
+
+  it("resume endpoint replays the same payload byte-for-byte", async () => {
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({ store });
+
+    const first = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "s1",
+      callback: (controller) => {
+        controller.appendText("alpha");
+        const tool = controller.addToolCallPart({
+          toolName: "echo",
+          toolCallId: "t1",
+          args: { v: 1 },
+          response: { result: { ok: true } },
+        });
+        tool.close();
+      },
+    });
+    const firstBytes = await new Response(first.body).arrayBuffer();
+
+    const second = await createResumeAssistantStreamResponse({
+      context: ctx,
+      streamId: "s1",
+    });
+    const secondBytes = await new Response(second.body).arrayBuffer();
+
+    expect(new Uint8Array(secondBytes)).toEqual(new Uint8Array(firstBytes));
+    expect(second.headers.get(RESUMABLE_STREAM_ID_HEADER)).toBe("s1");
+  });
+
+  it("resume returns 404 when streamId is unknown", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const response = await createResumeAssistantStreamResponse({
+      context: ctx,
+      streamId: "missing",
+    });
+    expect(response.status).toBe(404);
+    const body = await response.json();
+    expect(body).toEqual({ error: "stream not found" });
+  });
+
+  it("custom encoder factory survives across producer + consumer", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const response = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "s1",
+      encoder: () => new AssistantTransportEncoder(),
+      callback: (controller) => controller.appendText("hi"),
+    });
+
+    expect(response.headers.get("content-type")).toBe("text/event-stream");
+    const text = await new Response(response.body).text();
+    expect(text).toContain("text-delta");
+    expect(text).toContain("[DONE]");
+  });
+
+  it("user headers override encoder defaults but not the stream id header", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    const response = await createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "s1",
+      headers: {
+        "Cache-Control": "private, max-age=0",
+        [RESUMABLE_STREAM_ID_HEADER]: "spoofed",
+      },
+      callback: (controller) => controller.appendText("hi"),
+    });
+    expect(response.headers.get("cache-control")).toBe("private, max-age=0");
+    expect(response.headers.get(RESUMABLE_STREAM_ID_HEADER)).toBe("s1");
+    await response.body!.cancel();
+  });
+
+  describe("producer crash", () => {
+    let suppressed: unknown[];
+    let original: NodeJS.UnhandledRejectionListener[];
+    const swallow: NodeJS.UnhandledRejectionListener = (reason) => {
+      suppressed.push(reason);
+    };
+
+    beforeEach(() => {
+      suppressed = [];
+      original = process.listeners("unhandledRejection");
+      for (const l of original) process.off("unhandledRejection", l);
+      process.on("unhandledRejection", swallow);
+    });
+
+    afterEach(async () => {
+      await new Promise((r) => setTimeout(r, 20));
+      process.off("unhandledRejection", swallow);
+      for (const l of original) process.on("unhandledRejection", l);
+    });
+
+    it("synchronous callback throw is encoded into the body and finalizes the stream", async () => {
+      const ctx = createResumableStreamContext({
+        store: createInMemoryResumableStreamStore(),
+      });
+      const response = await createResumableAssistantStreamResponse({
+        context: ctx,
+        streamId: "s1",
+        callback: () => {
+          throw new Error("boom");
+        },
+      });
+      expect(response.status).toBe(200);
+
+      const firstBytes = await new Response(response.body).arrayBuffer();
+      expect(new TextDecoder().decode(firstBytes)).toContain("Error: boom");
+
+      while ((await ctx.status("s1")) === "streaming") {
+        await new Promise((r) => setTimeout(r, 5));
+      }
+      expect(await ctx.status("s1")).toBe("done");
+
+      const replay = await ctx.resume("s1");
+      expect(replay).not.toBeNull();
+      const replayBytes = await new Response(replay).arrayBuffer();
+      expect(new Uint8Array(replayBytes)).toEqual(new Uint8Array(firstBytes));
+
+      while (suppressed.length === 0) {
+        await new Promise((r) => setTimeout(r, 5));
+      }
+      expect(suppressed.map((e) => (e as Error).message)).toContain("boom");
+    });
+  });
+
+  it("requireResume rejects with ResumableStreamError(missing) for unknown ids", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+    let captured: unknown;
+    try {
+      await ctx.requireResume("nonexistent");
+    } catch (e) {
+      captured = e;
+    }
+    expect(captured).toBeInstanceOf(ResumableStreamError);
+    expect((captured as ResumableStreamError).code).toBe("missing");
+  });
+
+  it("mid-stream consumer joins active production and receives every chunk in order", async () => {
+    const ctx = createResumableStreamContext({
+      store: createInMemoryResumableStreamStore(),
+    });
+
+    let emitted = 0;
+    const producerDone = createResumableAssistantStreamResponse({
+      context: ctx,
+      streamId: "s1",
+      callback: async (controller) => {
+        for (let i = 0; i < 5; i++) {
+          await new Promise((r) => setTimeout(r, 5));
+          controller.appendText(`chunk${i};`);
+          emitted += 1;
+        }
+      },
+    });
+
+    const producer = await producerDone;
+
+    while (emitted < 2) {
+      await new Promise((r) => setTimeout(r, 1));
+    }
+
+    const replay = await ctx.resume("s1");
+    expect(replay).not.toBeNull();
+
+    const [producerBody, replayText] = await Promise.all([
+      new Response(producer.body).arrayBuffer(),
+      decodeAssistantText(new Response(replay)),
+    ]);
+    expect(producerBody.byteLength).toBeGreaterThan(0);
+    expect(replayText).toBe("chunk0;chunk1;chunk2;chunk3;chunk4;");
+  });
+});

package/src/resumable/createResumableAssistantStreamResponse.ts

@@ -0,0 +1,80 @@
+import type { AssistantStreamEncoder } from "../core/AssistantStream";
+import {
+  createAssistantStream,
+  type AssistantStreamController,
+} from "../core/modules/assistant-stream";
+import { DataStreamEncoder } from "../core/serialization/data-stream/DataStream";
+import type { ResumableStreamContext } from "./ResumableStreamContext";
+
+export const RESUMABLE_STREAM_ID_HEADER = "x-resumable-stream-id";
+
+export type CreateResumableAssistantStreamResponseOptions = {
+  readonly context: ResumableStreamContext;
+  readonly streamId: string;
+  readonly callback: (
+    controller: AssistantStreamController,
+  ) => PromiseLike<void> | void;
+  /** Defaults to `DataStreamEncoder`. Also consulted for response headers. */
+  readonly encoder?: () => AssistantStreamEncoder;
+  readonly headers?: HeadersInit;
+};
+
+export async function createResumableAssistantStreamResponse(
+  options: CreateResumableAssistantStreamResponseOptions,
+): Promise<Response> {
+  const encoder = (options.encoder ?? (() => new DataStreamEncoder()))();
+
+  const stream = await options.context.run(options.streamId, () => {
+    const aStream = createAssistantStream(options.callback);
+    return aStream.pipeThrough(encoder);
+  });
+
+  return new Response(stream, {
+    headers: mergeHeaders(encoder.headers, options.headers, options.streamId),
+  });
+}
+
+export type CreateResumeAssistantStreamResponseOptions = {
+  readonly context: ResumableStreamContext;
+  readonly streamId: string;
+  /** Read for `headers` only; the encoder transform is not invoked on resume. */
+  readonly encoder?: () => AssistantStreamEncoder;
+  readonly headers?: HeadersInit;
+  /** Defaults to a 404 JSON response. */
+  readonly missingResponse?: () => Response;
+};
+
+export async function createResumeAssistantStreamResponse(
+  options: CreateResumeAssistantStreamResponseOptions,
+): Promise<Response> {
+  const stream = await options.context.resume(options.streamId);
+  if (!stream) {
+    return options.missingResponse?.() ?? defaultMissingResponse();
+  }
+  const encoder = (options.encoder ?? (() => new DataStreamEncoder()))();
+  return new Response(stream, {
+    headers: mergeHeaders(encoder.headers, options.headers, options.streamId),
+  });
+}
+
+function defaultMissingResponse(): Response {
+  return new Response(JSON.stringify({ error: "stream not found" }), {
+    status: 404,
+    headers: { "Content-Type": "application/json" },
+  });
+}
+
+function mergeHeaders(
+  encoderHeaders: Headers | undefined,
+  extra: HeadersInit | undefined,
+  streamId: string,
+): Headers {
+  const merged = new Headers(encoderHeaders ?? {});
+  if (extra) {
+    for (const [key, value] of new Headers(extra)) {
+      merged.set(key, value);
+    }
+  }
+  merged.set(RESUMABLE_STREAM_ID_HEADER, streamId);
+  return merged;
+}

package/src/resumable/errors.ts

@@ -0,0 +1,26 @@
+export type ResumableStreamErrorCode =
+  | "missing"
+  | "exists"
+  | "finalized"
+  | "invalid-id";
+
+export class ResumableStreamError extends Error {
+  readonly code: ResumableStreamErrorCode;
+
+  constructor(code: ResumableStreamErrorCode, message: string) {
+    super(message);
+    this.name = "ResumableStreamError";
+    this.code = code;
+  }
+}
+
+const STREAM_ID_PATTERN = /^[A-Za-z0-9_.:-]{1,256}$/;
+
+export function validateStreamId(streamId: string): void {
+  if (!STREAM_ID_PATTERN.test(streamId)) {
+    throw new ResumableStreamError(
+      "invalid-id",
+      `Invalid streamId: ${streamId} (must match ${STREAM_ID_PATTERN})`,
+    );
+  }
+}

package/src/resumable/index.ts

@@ -0,0 +1,36 @@
+export type {
+  ResumableStreamStore,
+  ResumableStreamRole,
+  ResumableStreamStatus,
+  ResumableStreamEntry,
+  ResumableStreamAcquireOptions,
+} from "./types";
+
+export {
+  ResumableStreamError,
+  type ResumableStreamErrorCode,
+} from "./errors";
+
+export {
+  createResumableStreamContext,
+  type ResumableStreamContext,
+  type ResumableStreamContextOptions,
+} from "./ResumableStreamContext";
+
+export {
+  createResumableAssistantStreamResponse,
+  createResumeAssistantStreamResponse,
+  RESUMABLE_STREAM_ID_HEADER,
+  type CreateResumableAssistantStreamResponseOptions,
+  type CreateResumeAssistantStreamResponseOptions,
+} from "./createResumableAssistantStreamResponse";
+
+export {
+  createInMemoryResumableStreamStore,
+  type InMemoryResumableStreamStoreOptions,
+} from "./stores/InMemoryResumableStreamStore";
+
+export type {
+  RedisLikeClient,
+  RedisResumableStreamStoreOptions,
+} from "./stores/redis-impl";

package/src/resumable/stores/InMemoryResumableStreamStore.test.ts

@@ -0,0 +1,285 @@
+import { describe, expect, it, vi } from "vitest";
+import { createInMemoryResumableStreamStore } from "./InMemoryResumableStreamStore";
+import type { ResumableStreamEntry } from "../types";
+
+const bytes = (s: string): Uint8Array => new TextEncoder().encode(s);
+
+const decode = (chunk: Uint8Array): string => new TextDecoder().decode(chunk);
+
+async function drain(
+  iter: AsyncIterable<ResumableStreamEntry>,
+): Promise<string[]> {
+  const out: string[] = [];
+  for await (const entry of iter) out.push(decode(entry.chunk));
+  return out;
+}
+
+describe("InMemoryResumableStreamStore", () => {
+  it("elects exactly one producer per stream id", async () => {
+    const store = createInMemoryResumableStreamStore();
+    const first = await store.acquire("a");
+    const second = await store.acquire("a");
+    const third = await store.acquire("a");
+    expect(first).toBe("producer");
+    expect(second).toBe("consumer");
+    expect(third).toBe("consumer");
+  });
+
+  it("isolates streams by id", async () => {
+    const store = createInMemoryResumableStreamStore();
+    expect(await store.acquire("a")).toBe("producer");
+    expect(await store.acquire("b")).toBe("producer");
+  });
+
+  it("replays buffered entries and tails new ones until finalize", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.append("a", bytes("hello "));
+    await store.append("a", bytes("world"));
+
+    const ac = new AbortController();
+    const collected: string[] = [];
+    const reading = (async () => {
+      for await (const entry of store.read("a", "", ac.signal)) {
+        collected.push(decode(entry.chunk));
+        if (collected.length === 3) {
+          await store.finalize("a", "done");
+        }
+      }
+    })();
+
+    await store.append("a", bytes("!"));
+    await reading;
+    expect(collected).toEqual(["hello ", "world", "!"]);
+  });
+
+  it("status transitions: missing → streaming → done", async () => {
+    const store = createInMemoryResumableStreamStore();
+    expect(await store.status("a")).toBe("missing");
+    await store.acquire("a");
+    expect(await store.status("a")).toBe("streaming");
+    await store.finalize("a", "done");
+    expect(await store.status("a")).toBe("done");
+  });
+
+  it("status reports error after error finalize", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.finalize("a", "error", "boom");
+    expect(await store.status("a")).toBe("error");
+  });
+
+  it("read throws on the next iteration after error finalize", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.append("a", bytes("partial"));
+    await store.finalize("a", "error", "boom");
+
+    const ac = new AbortController();
+    const seen: string[] = [];
+    await expect(async () => {
+      for await (const entry of store.read("a", "", ac.signal)) {
+        seen.push(decode(entry.chunk));
+      }
+    }).rejects.toThrow("boom");
+    expect(seen).toEqual(["partial"]);
+  });
+
+  it("a consumer joining after done replays everything", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.append("a", bytes("a"));
+    await store.append("a", bytes("b"));
+    await store.append("a", bytes("c"));
+    await store.finalize("a", "done");
+
+    const ac = new AbortController();
+    expect(await drain(store.read("a", "", ac.signal))).toEqual([
+      "a",
+      "b",
+      "c",
+    ]);
+  });
+
+  it("cursor advances and skips already-seen entries", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.append("a", bytes("1"));
+    await store.append("a", bytes("2"));
+    await store.append("a", bytes("3"));
+    await store.finalize("a", "done");
+
+    const ac = new AbortController();
+    const seen: { cursor: string; text: string }[] = [];
+    for await (const entry of store.read("a", "", ac.signal)) {
+      seen.push({ cursor: entry.cursor, text: decode(entry.chunk) });
+    }
+    expect(seen.map((s) => s.text)).toEqual(["1", "2", "3"]);
+
+    const afterFirst = seen[0]!.cursor;
+    expect(await drain(store.read("a", afterFirst, ac.signal))).toEqual([
+      "2",
+      "3",
+    ]);
+  });
+
+  it("aborting the read signal terminates without throwing", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    const ac = new AbortController();
+    const collected: string[] = [];
+    const reading = (async () => {
+      for await (const entry of store.read("a", "", ac.signal)) {
+        collected.push(decode(entry.chunk));
+      }
+    })();
+    await store.append("a", bytes("x"));
+    await new Promise((r) => setTimeout(r, 5));
+    ac.abort();
+    await reading;
+    expect(collected).toEqual(["x"]);
+  });
+
+  it("multiple consumers can read concurrently", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    const ac = new AbortController();
+    const a = drain(store.read("a", "", ac.signal));
+    const b = drain(store.read("a", "", ac.signal));
+    await store.append("a", bytes("x"));
+    await store.append("a", bytes("y"));
+    await store.finalize("a", "done");
+    expect(await a).toEqual(["x", "y"]);
+    expect(await b).toEqual(["x", "y"]);
+  });
+
+  it("delete prevents further reads and ends in-flight reads", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    const ac = new AbortController();
+    const reading = drain(store.read("a", "", ac.signal));
+    await store.append("a", bytes("x"));
+    await new Promise((r) => setTimeout(r, 0));
+    await store.delete("a");
+    expect(await reading).toEqual(["x"]);
+    expect(await store.status("a")).toBe("missing");
+  });
+
+  it("expired streams are evicted on next access", async () => {
+    let now = 1_000;
+    const store = createInMemoryResumableStreamStore({
+      defaultTtlMs: 100,
+      now: () => now,
+    });
+    await store.acquire("a");
+    await store.append("a", bytes("hi"));
+    expect(await store.status("a")).toBe("streaming");
+    now += 200;
+    expect(await store.status("a")).toBe("missing");
+  });
+
+  it("appending refreshes TTL", async () => {
+    let now = 1_000;
+    const store = createInMemoryResumableStreamStore({
+      defaultTtlMs: 100,
+      now: () => now,
+    });
+    await store.acquire("a");
+    now += 80;
+    await store.append("a", bytes("x"));
+    now += 80;
+    expect(await store.status("a")).toBe("streaming");
+  });
+
+  it("rejects append on finalized stream", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.finalize("a", "done");
+    await expect(store.append("a", bytes("late"))).rejects.toThrow(
+      /already finalized/,
+    );
+  });
+
+  it("rejects append on missing stream", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await expect(store.append("a", bytes("x"))).rejects.toThrow(
+      /Stream not found/,
+    );
+  });
+
+  it("finalize is idempotent", async () => {
+    const store = createInMemoryResumableStreamStore();
+    await store.acquire("a");
+    await store.finalize("a", "done");
+    await store.finalize("a", "done");
+    expect(await store.status("a")).toBe("done");
+  });
+
+  it("rejects append when chunk exceeds maxChunkBytes", async () => {
+    const store = createInMemoryResumableStreamStore({ maxChunkBytes: 4 });
+    await store.acquire("a");
+    await expect(store.append("a", bytes("hello"))).rejects.toThrow(
+      /Chunk exceeds maxChunkBytes: 5/,
+    );
+    await store.append("a", bytes("ok"));
+    expect(await store.status("a")).toBe("streaming");
+  });
+
+  it("rejects append when stream reaches maxEntriesPerStream", async () => {
+    const store = createInMemoryResumableStreamStore({
+      maxEntriesPerStream: 2,
+    });
+    await store.acquire("a");
+    await store.append("a", bytes("1"));
+    await store.append("a", bytes("2"));
+    await expect(store.append("a", bytes("3"))).rejects.toThrow(
+      /Stream exceeded maxEntriesPerStream: a/,
+    );
+  });
+
+  it("rejects acquire when active stream count exceeds maxStreams", async () => {
+    const store = createInMemoryResumableStreamStore({ maxStreams: 2 });
+    await store.acquire("a");
+    await store.acquire("b");
+    await expect(store.acquire("c")).rejects.toThrow(/maxStreams exceeded/);
+    expect(await store.acquire("a")).toBe("consumer");
+  });
+
+  it("gc sweeper evicts expired streams without explicit access", async () => {
+    vi.useFakeTimers();
+    try {
+      let now = 1_000;
+      const store = createInMemoryResumableStreamStore({
+        defaultTtlMs: 100,
+        gcIntervalMs: 50,
+        now: () => now,
+      });
+      await store.acquire("a");
+      await store.append("a", bytes("hi"));
+      now += 200;
+      vi.advanceTimersByTime(50);
+      expect(await store.status("a")).toBe("missing");
+      store.dispose();
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("dispose clears the gc interval", async () => {
+    vi.useFakeTimers();
+    try {
+      const clearSpy = vi.spyOn(globalThis, "clearInterval");
+      const store = createInMemoryResumableStreamStore({ gcIntervalMs: 50 });
+      store.dispose();
+      expect(clearSpy).toHaveBeenCalledTimes(1);
+      clearSpy.mockRestore();
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("dispose is a no-op when gcIntervalMs is undefined", () => {
+    const store = createInMemoryResumableStreamStore();
+    expect(() => store.dispose()).not.toThrow();
+  });
+});