assistant-stream/src/resumable/ResumableStreamContext.ts

Streaming utilities for AI assistants

import { ResumableStreamError } from "./errors";
import { withPromiseOrValue } from "../core/utils/withPromiseOrValue";
import type {
  ResumableStreamRole,
  ResumableStreamAcquisition,
  ResumableStreamLease,
  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>;
}

const invokeObservabilityHook = <TArgs extends unknown[]>(
  name: string,
  hook: ((...args: TArgs) => void) | undefined,
  ...args: TArgs
): void => {
  if (!hook) return;
  void withPromiseOrValue(
    () => hook(...args),
    () => {},
    (error) => {
      console.error(`resumable stream ${name} hook failed:`, error);
    },
  );
};

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 acquisition:
        | ResumableStreamAcquisition
        | { role: ResumableStreamRole; lease?: undefined } = store.acquireLease
        ? await store.acquireLease(streamId, acquireOptions)
        : { role: await store.acquire(streamId, acquireOptions) };
      invokeObservabilityHook(
        "onAcquire",
        onAcquire,
        streamId,
        acquisition.role,
      );
      if (acquisition.role === "producer") {
        startProducerTask(
          store,
          streamId,
          makeStream,
          {
            waitUntil,
            onAppend,
            onFinalize,
            onError,
          },
          acquisition.lease,
        );
      }
      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,
  lease: ResumableStreamLease | undefined,
): void {
  const { waitUntil, onAppend, onFinalize, onError } = hooks;
  const task = (async () => {
    let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
    try {
      reader = makeStream().getReader();
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        await store.append(streamId, value, lease);
        invokeObservabilityHook(
          "onAppend",
          onAppend,
          streamId,
          value.byteLength,
        );
      }
      const finalized = await store.finalize(
        streamId,
        "done",
        undefined,
        lease,
      );
      if (finalized === false) {
        invokeObservabilityHook(
          "onError",
          onError,
          streamId,
          new ResumableStreamError(
            "missing",
            `Stream no longer owned by this producer: ${streamId}`,
          ),
        );
        return;
      }
      invokeObservabilityHook("onFinalize", onFinalize, streamId, "done");
    } catch (err) {
      invokeObservabilityHook("onError", onError, streamId, err);
      const message = err instanceof Error ? err.message : String(err);
      const cancellation = (async () => {
        try {
          await reader?.cancel(err);
        } catch (cancelErr) {
          console.error("resumable stream reader cancel failed:", cancelErr);
        }
      })();
      try {
        const finalized = await store.finalize(
          streamId,
          "error",
          message,
          lease,
        );
        if (finalized === false) return;
        invokeObservabilityHook(
          "onFinalize",
          onFinalize,
          streamId,
          "error",
          message,
        );
      } catch (finalizeErr) {
        console.error("resumable stream finalize failed:", finalizeErr);
      } finally {
        await cancellation;
      }
    } finally {
      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);
      }
    },
    async cancel() {
      ac.abort();
      try {
        await iterator?.return?.();
      } catch {}
    },
  });
}