Resume a stream reliably

import { isCheckfuError, type Checkfu } from "@checkfu/sdk"

type FollowedEvent = {
  readonly seq: number
  readonly type: string
  readonly payload?: unknown
}

/** One transaction in your own database: both writes commit, or neither does. */
export interface CursorTransaction {
  /** Lock this Session's cursor until commit, creating its row if absent. */
  loadCursor(sessionId: string): Promise<number | undefined>
  saveCursor(sessionId: string, seq: number): Promise<void>
}

export interface CursorStore {
  transaction<T>(work: (tx: CursorTransaction) => Promise<T>): Promise<T>
}

const TERMINAL_SESSION_EVENTS = new Set([
  "session.status_completed",
  "session.status_failed",
  "session.status_canceled",
])

const stopReason = (payload: unknown) =>
  typeof payload === "object" && payload !== null && "stop_reason" in payload
    ? payload.stop_reason
    : undefined

export const stopsWorkflowFollow = (event: FollowedEvent) =>
  TERMINAL_SESSION_EVENTS.has(event.type) ||
  (event.type === "session.status_idle" &&
    stopReason(event.payload) !== "requires_action")

/** The stream closed without a settlement event; nothing failed, so follow again. */
export class StreamEndedError extends Error {}

/** Only an SDK iterator read can produce this reconnectable failure. */
export class RetryableStreamError extends Error {}

/** Follow again only when the stream read failed or ended before settlement. */
export const isRetryableStreamFailure = (error: unknown) =>
  error instanceof StreamEndedError ||
  error instanceof RetryableStreamError

export const followSession = async (
  checkfu: Pick<Checkfu, "sessions">,
  store: CursorStore,
  sessionId: string,
  apply: (tx: CursorTransaction, event: FollowedEvent) => Promise<void>,
): Promise<void> => {
  const start = (await store.transaction((tx) => tx.loadCursor(sessionId))) ?? 0
  const iterator = checkfu.sessions.events.stream(sessionId, {
    lastEventId: start,
  })[Symbol.asyncIterator]()
  let failed = false
  try {
    for (;;) {
      let next: Awaited<ReturnType<typeof iterator.next>>
      try {
        next = await iterator.next()
      } catch (error) {
        if (isCheckfuError(error) && error.retryable) {
          throw new RetryableStreamError("stream reconnect budget exhausted", { cause: error })
        }
        throw error
      }
      if (next.done) break
      const event = next.value
      await store.transaction(async (tx) => {
        // The cursor lock serializes concurrent followers before either applies an effect.
        if (event.seq <= ((await tx.loadCursor(sessionId)) ?? 0)) return
        await apply(tx, event)
        await tx.saveCursor(sessionId, event.seq)
      })
      if (stopsWorkflowFollow(event)) return
    }
  } catch (error) {
    failed = true
    throw error
  } finally {
    if (failed) {
      try {
        await iterator.return?.()
      } catch {
        // Preserve the primary failure if cleanup also fails.
      }
    } else {
      await iterator.return?.()
    }
  }
  throw new StreamEndedError("event stream ended before the Session settled")
}