Streaming and SSE

Copyexport class StreamingApi {
  private writer: WritableStreamDefaultWriter<Uint8Array>
  private encoder: TextEncoder
  private writable: WritableStream
  private abortSubscribers: (() => void | Promise<void>)[] = []
  responseReadable: ReadableStream
  /**
   * Whether the stream has been aborted.
   */
  aborted: boolean = false
  /**
   * Whether the stream has been closed normally.
   */
  closed: boolean = false
 
  constructor(writable: WritableStream, _readable: ReadableStream) {
    this.writable = writable
    this.writer = writable.getWriter()
    this.encoder = new TextEncoder()
 
    const reader = _readable.getReader()
 
    // in case the user disconnects, let the reader know to cancel
    // this in-turn results in responseReadable being closed
    // and writeSSE method no longer blocks indefinitely
    this.abortSubscribers.push(async () => {
      await reader.cancel()
    })
 
    this.responseReadable = new ReadableStream({
      async pull(controller) {
        const { done, value } = await reader.read()
        done ? controller.close() : controller.enqueue(value)
      },
      cancel: () => {
        if (!this.closed) {
          this.abort()
        }
      },
    })
  }

Copyasync pipe(body: ReadableStream) {
  this.writer.releaseLock()
  await body.pipeTo(this.writable, { preventClose: true })
  this.writer = this.writable.getWriter()
}

CopyonAbort(listener: () => void | Promise<void>) {
  this.abortSubscribers.push(listener)
}
 
/**
 * Abort the stream.
 * You can call this method when stream is aborted by external event.
 */
abort() {
  if (!this.aborted) {
    this.aborted = true
    this.abortSubscribers.forEach((subscriber) => subscriber())
  }
}

Copy// Until Bun v1.1.27, Bun didn't call cancel() on the ReadableStream for Response objects from Bun.serve()
if (isOldBunVersion()) {
  c.req.raw.signal.addEventListener('abort', () => {
    if (!stream.closed) {
      stream.abort()
    }
  })
}

Copyexport const stream = (
  c: Context,
  cb: (stream: StreamingApi) => Promise<void>,
  onError?: (e: Error, stream: StreamingApi) => Promise<void>
): Response => {
  const { readable, writable } = new TransformStream()
  const stream = new StreamingApi(writable, readable)

Copy// in bun, `c` is destroyed when the request is returned, so hold it until the end of streaming
contextStash.set(stream.responseReadable, c)
;(async () => {

Copyasync write(input: Uint8Array | string): Promise<StreamingApi> {
  try {
    if (typeof input === 'string') {
      input = this.encoder.encode(input)
    }
    await this.writer.write(input)
  } catch {
    // Do nothing. If you want to handle errors, create a stream by yourself.
  }
  return this
}
 
async writeln(input: string): Promise<StreamingApi> {
  await this.write(input + '\n')
  return this
}
 
sleep(ms: number): Promise<unknown> {
  return new Promise((res) => setTimeout(res, ms))
}
 
async close() {
  this.closed = true
  try {
    await this.writer.close()
  } catch {
    // Do nothing. If you want to handle errors, create a stream by yourself.
  }
}
 
async pipe(body: ReadableStream) {
  this.writer.releaseLock()
  await body.pipeTo(this.writable, { preventClose: true })
  this.writer = this.writable.getWriter()
}
 
onAbort(listener: () => void | Promise<void>) {
  this.abortSubscribers.push(listener)
}
 
/**
 * Abort the stream.
 * You can call this method when stream is aborted by external event.
 */
abort() {
  if (!this.aborted) {
    this.aborted = true
    this.abortSubscribers.forEach((subscriber) => subscriber())
  }
}

Copytry {
  if (typeof input === 'string') {
    input = this.encoder.encode(input)
  }
  await this.writer.write(input)
} catch {
  // Do nothing. If you want to handle errors, create a stream by yourself.
}

Copyexport const streamText = (
  c: Context,
  cb: (stream: StreamingApi) => Promise<void>,
  onError?: (e: Error, stream: StreamingApi) => Promise<void>
): Response => {
  c.header('Content-Type', TEXT_PLAIN)
  c.header('X-Content-Type-Options', 'nosniff')
  c.header('Transfer-Encoding', 'chunked')
  return stream(c, cb, onError)
}

Copyc.header('Content-Type', TEXT_PLAIN)
c.header('X-Content-Type-Options', 'nosniff')
c.header('Transfer-Encoding', 'chunked')

Copyexport const streamText = (
  c: Context,
  cb: (stream: StreamingApi) => Promise<void>,
  onError?: (e: Error, stream: StreamingApi) => Promise<void>
): Response => {

Copyconst contextStash: WeakMap<ReadableStream, Context> = new WeakMap<ReadableStream, Context>()

Copy// Until Bun v1.1.27, Bun didn't call cancel() on the ReadableStream for Response objects from Bun.serve()
if (isOldBunVersion()) {
  c.req.raw.signal.addEventListener('abort', () => {
    if (!stream.closed) {
      stream.abort()
    }
  })
}

Copyfor (const key of ['event', 'id', 'retry'] as (keyof SSEMessage)[]) {
  if (message[key] && /[\r\n]/.test(message[key] as string)) {
    throw new Error(`${key} must not contain "\\r" or "\\n"`)
  }
}