54 lines · 2.1 KB
const SSE_HEARTBEAT_MS = 15_000;
const MAX_SUBSCRIBERS = 100;

export type StreamedLogRow = { seq: number; step: number; text: string };
type Writer = WritableStreamDefaultWriter<Uint8Array>;

/** The server-sent event streams of a Runner Durable Object's log viewers. */
export class LogSubscribers {
  private readonly writers = new Set<Writer>();
  private readonly encoder = new TextEncoder();

  private frame(payload: Record<string, unknown>, id?: number): Uint8Array {
    return this.encoder.encode(`${id === undefined ? "" : `id: ${id}\n`}data: ${JSON.stringify(payload)}\n\n`);
  }

  broadcast(payload: Record<string, unknown>, id?: number): void {
    const frame = this.frame(payload, id);
    for (const writer of this.writers) writer.write(frame).catch(() => this.writers.delete(writer));
  }

  /** Sends a final event and ends every stream. */
  finish(payload: Record<string, unknown>): void {
    this.broadcast(payload);
    for (const writer of this.writers) writer.close().catch(() => undefined);
    this.writers.clear();
  }

  /**
   * A new stream: the backlog, then either `final` (the job is over, so the stream ends at once) or live
   * events until `finish`, with keepalives so idle proxies don't drop it.
   */
  open(backlog: StreamedLogRow[], final: Record<string, unknown> | null): Response {
    if (this.writers.size >= MAX_SUBSCRIBERS) return new Response("Too many log viewers", { status: 429 });
    const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
    const writer = writable.getWriter();
    for (const row of backlog) writer.write(this.frame({ type: "log", ...row }, row.seq));
    if (final) {
      writer.write(this.frame(final));
      writer.close();
    } else {