assistant-stream 0.3.13 → 0.3.14

package/src/resumable/stores/InMemoryResumableStreamStore.ts

@@ -0,0 +1,237 @@
+import { DEFAULT_TTL_MS } from "../constants";
+import { ResumableStreamError, validateStreamId } from "../errors";
+import type {
+  ResumableStreamEntry,
+  ResumableStreamStatus,
+  ResumableStreamStore,
+} from "../types";
+
+type FinalizeMarker = { kind: "done" } | { kind: "error"; error: string };
+
+type StreamState = {
+  entries: ResumableStreamEntry[];
+  nextSeq: number;
+  expiresAt: number;
+  ttlMs: number;
+  final: FinalizeMarker | undefined;
+  waiters: Array<() => void>;
+};
+
+const cursorOf = (seq: number): string => seq.toString(36);
+const seqFromCursor = (cursor: string): number => {
+  if (cursor === "") return 0;
+  const parsed = Number.parseInt(cursor, 36);
+  return Number.isNaN(parsed) ? 0 : parsed;
+};
+
+export type InMemoryResumableStreamStoreOptions = {
+  readonly defaultTtlMs?: number;
+  readonly now?: () => number;
+  readonly maxChunkBytes?: number;
+  readonly maxEntriesPerStream?: number;
+  readonly maxStreams?: number;
+  readonly gcIntervalMs?: number;
+};
+
+export function createInMemoryResumableStreamStore(
+  options: InMemoryResumableStreamStoreOptions = {},
+): ResumableStreamStore & { dispose: () => void } {
+  const streams = new Map<string, StreamState>();
+  const defaultTtlMs = options.defaultTtlMs ?? DEFAULT_TTL_MS;
+  const now = options.now ?? Date.now;
+  const maxChunkBytes = options.maxChunkBytes;
+  const maxEntriesPerStream = options.maxEntriesPerStream;
+  const maxStreams = options.maxStreams;
+
+  const evictExpired = (): void => {
+    const t = now();
+    for (const [id, state] of streams) {
+      if (state.expiresAt > t) continue;
+      streams.delete(id);
+      state.final ??= { kind: "error", error: "Stream expired" };
+      notify(state);
+    }
+  };
+
+  const notify = (state: StreamState): void => {
+    const waiters = state.waiters;
+    state.waiters = [];
+    for (const wake of waiters) wake();
+  };
+
+  const findStartIndex = (state: StreamState, cursor: string): number => {
+    if (cursor === "") return 0;
+    const after = seqFromCursor(cursor);
+    let lo = 0;
+    let hi = state.entries.length;
+    while (lo < hi) {
+      const mid = (lo + hi) >>> 1;
+      const seq = seqFromCursor(state.entries[mid]!.cursor);
+      if (seq <= after) lo = mid + 1;
+      else hi = mid;
+    }
+    return lo;
+  };
+
+  const waitForUpdate = (
+    state: StreamState,
+    signal: AbortSignal,
+    wakeBy?: number,
+  ): Promise<void> =>
+    new Promise<void>((resolve) => {
+      let settled = false;
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const wake = () => {
+        if (settled) return;
+        settled = true;
+        if (timer !== undefined) clearTimeout(timer);
+        signal.removeEventListener("abort", wake);
+        const idx = state.waiters.indexOf(wake);
+        if (idx !== -1) state.waiters.splice(idx, 1);
+        resolve();
+      };
+      if (signal.aborted) {
+        wake();
+        return;
+      }
+      state.waiters.push(wake);
+      signal.addEventListener("abort", wake, { once: true });
+      if (wakeBy !== undefined) {
+        if (wakeBy > 0) {
+          timer = setTimeout(wake, wakeBy);
+        } else {
+          // already past the deadline; resolve so the caller can re-check
+          // expiration without waiting for an external notify.
+          wake();
+        }
+      }
+    });
+
+  const requireActive = (streamId: string): StreamState => {
+    evictExpired();
+    const state = streams.get(streamId);
+    if (!state) throw new Error(`Stream not found: ${streamId}`);
+    if (state.final) {
+      throw new ResumableStreamError(
+        "finalized",
+        `Stream already finalized: ${streamId}`,
+      );
+    }
+    return state;
+  };
+
+  const gcTimer =
+    options.gcIntervalMs !== undefined
+      ? setInterval(evictExpired, options.gcIntervalMs)
+      : undefined;
+  gcTimer?.unref?.();
+
+  return {
+    async acquire(streamId, acquireOptions) {
+      validateStreamId(streamId);
+      evictExpired();
+      const existing = streams.get(streamId);
+      if (existing) return "consumer";
+
+      if (maxStreams !== undefined && streams.size >= maxStreams) {
+        throw new Error("maxStreams exceeded");
+      }
+
+      const ttlMs = acquireOptions?.ttlMs ?? defaultTtlMs;
+      streams.set(streamId, {
+        entries: [],
+        nextSeq: 1,
+        expiresAt: now() + ttlMs,
+        ttlMs,
+        final: undefined,
+        waiters: [],
+      });
+      return "producer";
+    },
+
+    async append(streamId, chunk) {
+      validateStreamId(streamId);
+      if (maxChunkBytes !== undefined && chunk.byteLength > maxChunkBytes) {
+        throw new Error(`Chunk exceeds maxChunkBytes: ${chunk.byteLength}`);
+      }
+      const state = requireActive(streamId);
+      if (
+        maxEntriesPerStream !== undefined &&
+        state.entries.length >= maxEntriesPerStream
+      ) {
+        throw new Error(`Stream exceeded maxEntriesPerStream: ${streamId}`);
+      }
+      const seq = state.nextSeq;
+      state.nextSeq += 1;
+      state.entries.push({ cursor: cursorOf(seq), chunk });
+      state.expiresAt = now() + state.ttlMs;
+      notify(state);
+    },
+
+    async finalize(streamId, status, error) {
+      validateStreamId(streamId);
+      evictExpired();
+      const state = streams.get(streamId);
+      if (!state) throw new Error(`Stream not found: ${streamId}`);
+      if (state.final) return;
+      state.final =
+        status === "done"
+          ? { kind: "done" }
+          : { kind: "error", error: error ?? "Stream errored" };
+      state.expiresAt = now() + state.ttlMs;
+      notify(state);
+    },
+
+    async *read(streamId, cursor, signal) {
+      validateStreamId(streamId);
+      evictExpired();
+      const state = streams.get(streamId);
+      if (!state) throw new Error(`Stream not found: ${streamId}`);
+
+      let idx = findStartIndex(state, cursor);
+
+      while (true) {
+        if (signal.aborted) return;
+
+        while (idx < state.entries.length) {
+          if (signal.aborted) return;
+          yield state.entries[idx]!;
+          idx += 1;
+        }
+
+        if (state.final) {
+          if (state.final.kind === "error") {
+            throw new Error(state.final.error);
+          }
+          return;
+        }
+
+        const wakeBy = state.expiresAt - now();
+        await waitForUpdate(state, signal, wakeBy);
+        evictExpired();
+      }
+    },
+
+    async status(streamId): Promise<ResumableStreamStatus> {
+      validateStreamId(streamId);
+      evictExpired();
+      const state = streams.get(streamId);
+      if (!state) return "missing";
+      if (!state.final) return "streaming";
+      return state.final.kind === "error" ? "error" : "done";
+    },
+
+    async delete(streamId) {
+      validateStreamId(streamId);
+      const state = streams.get(streamId);
+      if (!state) return;
+      streams.delete(streamId);
+      state.final ??= { kind: "done" };
+      notify(state);
+    },
+
+    dispose() {
+      if (gcTimer !== undefined) clearInterval(gcTimer);
+    },
+  };
+}

package/src/resumable/stores/ioredis.ts

@@ -0,0 +1,123 @@
+import type {
+  ChainableCommander,
+  Cluster as IoRedisCluster,
+  Redis as IoRedis,
+} from "ioredis";
+import {
+  RedisResumableStreamStore,
+  type PipelineCommand,
+  type RedisLikeClient,
+  type RedisResumableStreamStoreOptions,
+} from "./redis-impl";
+import type { ResumableStreamStore } from "../types";
+
+export type IoRedisLike = IoRedis | IoRedisCluster;
+
+/**
+ * Resumable stream store backed by [`ioredis`](https://www.npmjs.com/package/ioredis)
+ * v5. Accepts a `Redis` or `Cluster` instance.
+ */
+export function createIoredisResumableStreamStore(
+  client: IoRedisLike,
+  options?: RedisResumableStreamStoreOptions,
+): ResumableStreamStore {
+  return new RedisResumableStreamStore(adapt(client), options);
+}
+
+function adapt(client: IoRedisLike): RedisLikeClient {
+  return {
+    async setNX(key, value, ttlSec) {
+      const result = await client.set(key, value, "EX", ttlSec, "NX");
+      return result === "OK";
+    },
+    async set(key, value, ttlSec) {
+      await client.set(key, value, "EX", ttlSec);
+    },
+    async get(key) {
+      return client.get(key);
+    },
+    async expire(key, ttlSec) {
+      await client.expire(key, ttlSec);
+    },
+    async exists(key) {
+      const result = await client.exists(key);
+      return result > 0;
+    },
+    async del(keys) {
+      if (keys.length === 0) return;
+      await client.del(...keys);
+    },
+    async xAdd(key, fields) {
+      const id = await client.xadd(key, "*", ...toFieldArgs(fields));
+      return id ?? "";
+    },
+    async xRange(key, start, end) {
+      const entries = await client.xrangeBuffer(key, start, end);
+      return entries.map(([idBuf, fieldArray]) => ({
+        id: idBuf.toString("utf8"),
+        fields: bufferFieldsToRecord(fieldArray),
+      }));
+    },
+    async pipeline(commands) {
+      if (commands.length === 0) return;
+      const pipe = client.pipeline();
+      for (const cmd of commands) {
+        applyPipelineCommand(pipe, cmd);
+      }
+      const results = (await pipe.exec()) ?? [];
+      for (const [err] of results) {
+        if (err) throw err;
+      }
+    },
+  };
+}
+
+function applyPipelineCommand(
+  pipe: ChainableCommander,
+  cmd: PipelineCommand,
+): void {
+  switch (cmd.type) {
+    case "xAdd":
+      pipe.xadd(cmd.key, "*", ...toFieldArgs(cmd.fields));
+      return;
+    case "expire":
+      pipe.expire(cmd.key, cmd.ttlSec);
+      return;
+    case "set":
+      pipe.set(cmd.key, cmd.value, "EX", cmd.ttlSec);
+      return;
+  }
+}
+
+function toFieldArgs(
+  fields: Record<string, string | Uint8Array>,
+): Array<string | Buffer> {
+  const args: Array<string | Buffer> = [];
+  for (const [k, v] of Object.entries(fields)) {
+    args.push(k, typeof v === "string" ? v : toBuffer(v));
+  }
+  return args;
+}
+
+function toBuffer(bytes: Uint8Array): Buffer {
+  if (Buffer.isBuffer(bytes)) return bytes;
+  return Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+}
+
+function bufferFieldsToRecord(
+  fields: Buffer[],
+): Record<string, string | Uint8Array> {
+  const out: Record<string, string | Uint8Array> = {};
+  for (let i = 0; i + 1 < fields.length; i += 2) {
+    const key = fields[i]?.toString("utf8");
+    const value = fields[i + 1];
+    if (key !== undefined && value !== undefined) {
+      out[key] = new Uint8Array(
+        value.buffer,
+        value.byteOffset,
+        value.byteLength,
+      );
+    }
+  }
+  return out;
+}

package/src/resumable/stores/redis-impl.ts

@@ -0,0 +1,304 @@
+import { DEFAULT_TTL_MS } from "../constants";
+import { ResumableStreamError, validateStreamId } from "../errors";
+import type {
+  ResumableStreamAcquireOptions,
+  ResumableStreamEntry,
+  ResumableStreamRole,
+  ResumableStreamStatus,
+  ResumableStreamStore,
+} from "../types";
+
+const DEFAULT_POLL_INTERVAL_MS = 100;
+const DEFAULT_KEY_PREFIX = "aui:resumable";
+
+const FIELD_CHUNK = "c";
+const FIELD_FIN = "fin";
+const FIELD_ERROR = "error";
+
+const FIN_DONE = "done";
+const FIN_ERROR = "error";
+
+const STREAM_START_ID = "0-0";
+
+export type PipelineCommand =
+  | {
+      readonly type: "xAdd";
+      readonly key: string;
+      readonly fields: Record<string, string | Uint8Array>;
+    }
+  | { readonly type: "expire"; readonly key: string; readonly ttlSec: number }
+  | {
+      readonly type: "set";
+      readonly key: string;
+      readonly value: string;
+      readonly ttlSec: number;
+    };
+
+/**
+ * Structural Redis-client interface. The bundled `redis` and `ioredis`
+ * adapters wrap their respective clients to satisfy it; custom or proxied
+ * clients can implement it directly.
+ */
+export interface RedisLikeClient {
+  setNX(key: string, value: string, ttlSec: number): Promise<boolean>;
+  set(key: string, value: string, ttlSec: number): Promise<void>;
+  get(key: string): Promise<string | null>;
+  expire(key: string, ttlSec: number): Promise<void>;
+  exists(key: string): Promise<boolean>;
+  del(keys: string[]): Promise<void>;
+  xAdd(
+    key: string,
+    fields: Record<string, string | Uint8Array>,
+  ): Promise<string>;
+  xRange(
+    key: string,
+    start: string,
+    end: string,
+  ): Promise<
+    Array<{ id: string; fields: Record<string, string | Uint8Array> }>
+  >;
+  /** Executes the commands as a single pipeline batch (one round trip). */
+  pipeline(commands: readonly PipelineCommand[]): Promise<void>;
+}
+
+export type RedisResumableStreamStoreOptions = {
+  readonly keyPrefix?: string;
+  readonly defaultTtlMs?: number;
+  /** Defaults to 100ms. Lower values reduce read latency, raise traffic. */
+  readonly pollIntervalMs?: number;
+  readonly maxChunkBytes?: number;
+};
+
+export class RedisResumableStreamStore implements ResumableStreamStore {
+  private readonly client: RedisLikeClient;
+  private readonly keyPrefix: string;
+  private readonly defaultTtlMs: number;
+  private readonly pollIntervalMs: number;
+  private readonly maxChunkBytes: number | undefined;
+
+  constructor(
+    client: RedisLikeClient,
+    options: RedisResumableStreamStoreOptions = {},
+  ) {
+    this.client = client;
+    this.keyPrefix = options.keyPrefix ?? DEFAULT_KEY_PREFIX;
+    this.defaultTtlMs = options.defaultTtlMs ?? DEFAULT_TTL_MS;
+    this.pollIntervalMs = options.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS;
+    this.maxChunkBytes = options.maxChunkBytes;
+  }
+
+  async acquire(
+    streamId: string,
+    options?: ResumableStreamAcquireOptions,
+  ): Promise<ResumableStreamRole> {
+    validateStreamId(streamId);
+    const ttlSec = msToSec(options?.ttlMs ?? this.defaultTtlMs);
+    const meta = JSON.stringify({ status: "streaming", ttlSec });
+    const acquired = await this.client.setNX(
+      this.metaKey(streamId),
+      meta,
+      ttlSec,
+    );
+    if (acquired) {
+      // a prior producer's data key may outlive its expired meta key.
+      await this.client.del([this.dataKey(streamId)]);
+      return "producer";
+    }
+    return "consumer";
+  }
+
+  async append(streamId: string, chunk: Uint8Array): Promise<void> {
+    validateStreamId(streamId);
+    if (
+      this.maxChunkBytes !== undefined &&
+      chunk.byteLength > this.maxChunkBytes
+    ) {
+      throw new Error(
+        `Chunk exceeds maxChunkBytes (${chunk.byteLength} > ${this.maxChunkBytes})`,
+      );
+    }
+    const dataKey = this.dataKey(streamId);
+    const metaKey = this.metaKey(streamId);
+    const meta = await this.readMeta(streamId);
+    if (!meta) {
+      throw new Error(`Stream not found: ${streamId}`);
+    }
+    if (meta.status !== "streaming") {
+      throw new ResumableStreamError(
+        "finalized",
+        `Stream already finalized: ${streamId}`,
+      );
+    }
+    const ttlSec = meta.ttlSec ?? msToSec(this.defaultTtlMs);
+    await this.client.pipeline([
+      { type: "xAdd", key: dataKey, fields: { [FIELD_CHUNK]: chunk } },
+      { type: "expire", key: dataKey, ttlSec },
+      { type: "expire", key: metaKey, ttlSec },
+    ]);
+  }
+
+  async finalize(
+    streamId: string,
+    status: "done" | "error",
+    error?: string,
+  ): Promise<void> {
+    validateStreamId(streamId);
+    const dataKey = this.dataKey(streamId);
+    const metaKey = this.metaKey(streamId);
+    const existing = await this.readMeta(streamId);
+    if (!existing) {
+      throw new Error(`Stream not found: ${streamId}`);
+    }
+    // a second finalize must not append a duplicate FIN entry.
+    if (existing.status !== "streaming") return;
+    const ttlSec = existing.ttlSec ?? msToSec(this.defaultTtlMs);
+    const meta = JSON.stringify(
+      status === "error"
+        ? { status: "error", error: error ?? "Stream errored", ttlSec }
+        : { status: "done", ttlSec },
+    );
+    const fields: Record<string, string> = {
+      [FIELD_FIN]: status === "error" ? FIN_ERROR : FIN_DONE,
+    };
+    if (status === "error") {
+      fields[FIELD_ERROR] = error ?? "Stream errored";
+    }
+    await this.client.pipeline([
+      { type: "set", key: metaKey, value: meta, ttlSec },
+      { type: "xAdd", key: dataKey, fields },
+      { type: "expire", key: dataKey, ttlSec },
+    ]);
+  }
+
+  async *read(
+    streamId: string,
+    cursor: string,
+    signal: AbortSignal,
+  ): AsyncIterable<ResumableStreamEntry> {
+    validateStreamId(streamId);
+    const dataKey = this.dataKey(streamId);
+    const metaKey = this.metaKey(streamId);
+    const initialMeta = await this.client.get(metaKey);
+    if (initialMeta === null) {
+      throw new Error(`Stream not found: ${streamId}`);
+    }
+
+    let lastId = cursor === "" ? STREAM_START_ID : cursor;
+
+    while (true) {
+      if (signal.aborted) return;
+
+      const start = lastId === STREAM_START_ID ? "-" : `(${lastId}`;
+      const entries = await this.client.xRange(dataKey, start, "+");
+
+      for (const entry of entries) {
+        if (signal.aborted) return;
+        lastId = entry.id;
+
+        const fin = readString(entry.fields[FIELD_FIN]);
+        if (fin === FIN_DONE) return;
+        if (fin === FIN_ERROR) {
+          throw new Error(
+            readString(entry.fields[FIELD_ERROR]) ?? "Stream errored",
+          );
+        }
+
+        const raw = entry.fields[FIELD_CHUNK];
+        if (raw === undefined) continue;
+        yield { cursor: entry.id, chunk: toBytes(raw) };
+      }
+
+      if (entries.length > 0) continue;
+
+      const stillExists = await this.client.exists(metaKey);
+      if (!stillExists) return;
+
+      await sleep(this.pollIntervalMs, signal);
+    }
+  }
+
+  async status(streamId: string): Promise<ResumableStreamStatus> {
+    validateStreamId(streamId);
+    const meta = await this.client.get(this.metaKey(streamId));
+    if (meta === null) return "missing";
+    const parsed = parseMeta(meta);
+    if (parsed?.status === "streaming") return "streaming";
+    if (parsed?.status === "done") return "done";
+    if (parsed?.status === "error") return "error";
+    return "missing";
+  }
+
+  async delete(streamId: string): Promise<void> {
+    validateStreamId(streamId);
+    await this.client.del([this.metaKey(streamId), this.dataKey(streamId)]);
+  }
+
+  private async readMeta(streamId: string): Promise<ParsedMeta | undefined> {
+    const raw = await this.client.get(this.metaKey(streamId));
+    if (raw === null) return undefined;
+    return parseMeta(raw);
+  }
+
+  // {streamId} is a Redis Cluster hash tag so both keys live on the same
+  // shard; multi-key DEL and same-stream pipelines stay single-slot.
+  private metaKey(streamId: string): string {
+    return `${this.keyPrefix}:{${streamId}}:meta`;
+  }
+
+  private dataKey(streamId: string): string {
+    return `${this.keyPrefix}:{${streamId}}:data`;
+  }
+}
+
+type ParsedMeta = {
+  status?: string;
+  error?: string;
+  ttlSec?: number;
+};
+
+function parseMeta(value: string): ParsedMeta | undefined {
+  try {
+    const parsed = JSON.parse(value) as ParsedMeta;
+    return parsed && typeof parsed === "object" ? parsed : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function msToSec(ms: number): number {
+  return Math.max(1, Math.ceil(ms / 1000));
+}
+
+function sleep(ms: number, signal: AbortSignal): Promise<void> {
+  return new Promise<void>((resolve) => {
+    if (signal.aborted) {
+      resolve();
+      return;
+    }
+    const timer = setTimeout(() => {
+      signal.removeEventListener("abort", onAbort);
+      resolve();
+    }, ms);
+    const onAbort = () => {
+      clearTimeout(timer);
+      resolve();
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+  });
+}
+
+const SHARED_DECODER = new TextDecoder();
+const SHARED_ENCODER = new TextEncoder();
+
+function readString(
+  value: string | Uint8Array | undefined,
+): string | undefined {
+  if (value === undefined) return undefined;
+  if (typeof value === "string") return value;
+  return SHARED_DECODER.decode(value);
+}
+
+function toBytes(value: string | Uint8Array): Uint8Array {
+  if (value instanceof Uint8Array) return value;
+  return SHARED_ENCODER.encode(value);
+}