nexus/frontend/src/services/sseHub.ts

import { createParser, type EventSourceMessage } from 'eventsource-parser'
import { apiFetch } from './api'

export type SseConnectionState = 'connecting' | 'open' | 'reconnecting' | 'closed' | 'unsupported' | 'error'

export interface ParsedSseEvent<T = unknown> {
  id: string | null
  event: string
  data: T
}

export interface SseSubscriptionOptions {
  lastEventId?: string | null
  onStateChange?: (state: SseConnectionState) => void
}

type Subscriber = {
  onEvent: (event: ParsedSseEvent) => void
  onStateChange?: (state: SseConnectionState) => void
}

type Connection = {
  url: string
  subscribers: Set<Subscriber>
  controller: AbortController | null
  loop: Promise<void> | null
  lastEventId: string | null
  retryMs: number
  stopped: boolean
}

const MAX_BUFFER_SIZE = 256 * 1024
const MAX_BACKOFF_MS = 30_000
const HEARTBEAT_TIMEOUT_MS = 45_000

function delay(ms: number, signal: AbortSignal): Promise<void> {
  return new Promise((resolve, reject) => {
    if (signal.aborted) {
      reject(signal.reason)
      return
    }
    const timer = globalThis.setTimeout(resolve, ms)
    signal.addEventListener('abort', () => {
      globalThis.clearTimeout(timer)
      reject(signal.reason)
    }, { once: true })
  })
}

function parseData(raw: string): unknown {
  try {
    return JSON.parse(raw)
  } catch {
    return raw
  }
}

export class AuthenticatedSseHub {
  private readonly connections = new Map<string, Connection>()

  subscribe(
    url: string,
    onEvent: (event: ParsedSseEvent) => void,
    options: SseSubscriptionOptions = {},
  ): () => void {
    let connection = this.connections.get(url)
    if (!connection) {
      connection = {
        url,
        subscribers: new Set(),
        controller: null,
        loop: null,
        lastEventId: options.lastEventId ?? null,
        retryMs: 1_000,
        stopped: false,
      }
      this.connections.set(url, connection)
    } else if (!connection.lastEventId && options.lastEventId) {
      connection.lastEventId = options.lastEventId
    }

    const subscriber: Subscriber = { onEvent, onStateChange: options.onStateChange }
    connection.subscribers.add(subscriber)
    connection.stopped = false
    if (!connection.loop) connection.loop = this.run(connection)

    return () => {
      connection?.subscribers.delete(subscriber)
      if (connection && connection.subscribers.size === 0) {
        connection.stopped = true
        connection.controller?.abort('no-subscribers')
        this.connections.delete(url)
      }
    }
  }

  closeAll(): void {
    for (const connection of this.connections.values()) {
      connection.stopped = true
      connection.controller?.abort('hub-closed')
      this.notify(connection, 'closed')
    }
    this.connections.clear()
  }

  private notify(connection: Connection, state: SseConnectionState): void {
    for (const subscriber of connection.subscribers) subscriber.onStateChange?.(state)
  }

  private dispatch(connection: Connection, event: ParsedSseEvent): void {
    queueMicrotask(() => {
      for (const subscriber of connection.subscribers) subscriber.onEvent(event)
    })
  }

  private async waitBeforeReconnect(
    connection: Connection,
    state: SseConnectionState,
  ): Promise<void> {
    if (connection.stopped || connection.subscribers.size === 0) return

    this.notify(connection, state)
    const controller = new AbortController()
    connection.controller = controller
    const jitter = Math.round(connection.retryMs * (0.8 + Math.random() * 0.4))
    await delay(jitter, controller.signal).catch(() => undefined)
    if (!connection.stopped) {
      connection.retryMs = Math.min(connection.retryMs * 2, MAX_BACKOFF_MS)
    }
  }

  private async run(connection: Connection): Promise<void> {
    try {
      while (!connection.stopped && connection.subscribers.size > 0) {
        const controller = new AbortController()
        connection.controller = controller
        this.notify(connection, connection.lastEventId ? 'reconnecting' : 'connecting')

        try {
          const headers = new Headers({
            Accept: 'text/event-stream',
            'Cache-Control': 'no-cache',
          })
          if (connection.lastEventId) headers.set('Last-Event-ID', connection.lastEventId)

          const response = await apiFetch(connection.url, { headers, signal: controller.signal })
          if (response.status === 404 || response.status === 501) {
            this.notify(connection, 'unsupported')
            await delay(60_000, controller.signal)
            continue
          }
          if (!response.ok || !response.body) {
            throw new Error(`SSE unavailable: HTTP ${response.status}`)
          }

          let parserFailure: Error | null = null
          let serverDirectedRetry = false
          const parser = createParser({
            maxBufferSize: MAX_BUFFER_SIZE,
            onRetry: retry => {
              serverDirectedRetry = true
              connection.retryMs = Math.min(Math.max(retry, 500), MAX_BACKOFF_MS)
            },
            onError: error => {
              if (error.type === 'max-buffer-size-exceeded') {
                parserFailure = error
                controller.abort(error)
              }
            },
            onEvent: (message: EventSourceMessage) => {
              if (!serverDirectedRetry) connection.retryMs = 1_000
              if (message.id) connection.lastEventId = message.id
              this.dispatch(connection, {
                id: message.id ?? null,
                event: message.event || 'message',
                data: parseData(message.data),
              })
            },
          })

          this.notify(connection, 'open')
          const reader = response.body.getReader()
          const decoder = new TextDecoder()
          let heartbeatTimer: ReturnType<typeof globalThis.setTimeout> | null = null
          const armHeartbeatTimeout = () => {
            if (heartbeatTimer !== null) globalThis.clearTimeout(heartbeatTimer)
            heartbeatTimer = globalThis.setTimeout(() => {
              controller.abort(new Error('SSE heartbeat timeout'))
            }, HEARTBEAT_TIMEOUT_MS)
          }
          armHeartbeatTimeout()
          try {
            while (!controller.signal.aborted) {
              const { value, done } = await reader.read()
              if (done) break
              armHeartbeatTimeout()
              parser.feed(decoder.decode(value, { stream: true }))
              if (parserFailure) throw parserFailure
            }
            const tail = decoder.decode()
            if (tail) parser.feed(tail)
            parser.reset({ consume: true })
          } finally {
            if (heartbeatTimer !== null) globalThis.clearTimeout(heartbeatTimer)
            reader.releaseLock()
          }

          // A clean EOF is still a disconnected live stream. Back off before
          // reconnecting so a server that closes immediately cannot trigger a
          // tight request loop.
          if (!controller.signal.aborted) {
            await this.waitBeforeReconnect(connection, 'reconnecting')
          }
        } catch (error) {
          if (connection.stopped || controller.signal.reason === 'no-subscribers' || controller.signal.reason === 'hub-closed') break
          await this.waitBeforeReconnect(connection, 'error')
        }
      }
    } finally {
      connection.controller = null
      connection.loop = null
      if (!connection.stopped && connection.subscribers.size > 0) {
        connection.loop = this.run(connection)
      }
    }
  }
}

export const sseHub = new AuthenticatedSseHub()