qtalker---/frontend/src/lib/ws.ts

// Reconnecting WebSocket client for /ws/session with a typed event emitter.
// Exponential backoff on unexpected close; surfaces connection status so the
// UI can show "the connection to the other side is unstable".

import type { ClientFrame, ServerFrame } from './types'

export type VeilConnectionState =
  | 'connecting'
  | 'open'
  | 'unstable' // reconnecting after a drop
  | 'closed'

type FrameHandler = (frame: ServerFrame) => void
type StateHandler = (state: VeilConnectionState) => void

export type VeilSocketOptions = {
  url?: string
  /** Base backoff delay (ms); doubles each retry up to maxBackoffMs. */
  baseBackoffMs?: number
  maxBackoffMs?: number
  /** Ping interval to keep the veil warm (ms). 0 disables. */
  pingIntervalMs?: number
  /** Injectable WebSocket constructor (tests). */
  socketFactory?: (url: string) => WebSocket
}

export function defaultSessionUrl(): string {
  const proto = location.protocol === 'https:' ? 'wss://' : 'ws://'
  return `${proto}${location.host}/ws/session`
}

export class VeilSocket {
  private ws: WebSocket | null = null
  private frameHandlers = new Set<FrameHandler>()
  private stateHandlers = new Set<StateHandler>()
  private readonly url: string
  private readonly baseBackoffMs: number
  private readonly maxBackoffMs: number
  private readonly pingIntervalMs: number
  private readonly socketFactory: (url: string) => WebSocket
  private attempts = 0
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null
  private pingTimer: ReturnType<typeof setInterval> | null = null
  private deliberatelyClosed = false
  private outbox: string[] = []
  private _state: VeilConnectionState = 'closed'

  constructor(opts: VeilSocketOptions = {}) {
    this.url = opts.url ?? defaultSessionUrl()
    this.baseBackoffMs = opts.baseBackoffMs ?? 800
    this.maxBackoffMs = opts.maxBackoffMs ?? 15000
    this.pingIntervalMs = opts.pingIntervalMs ?? 25000
    this.socketFactory = opts.socketFactory ?? ((url: string) => new WebSocket(url))
  }

  get state(): VeilConnectionState {
    return this._state
  }

  onFrame(handler: FrameHandler): () => void {
    this.frameHandlers.add(handler)
    return () => this.frameHandlers.delete(handler)
  }

  onState(handler: StateHandler): () => void {
    this.stateHandlers.add(handler)
    return () => this.stateHandlers.delete(handler)
  }

  connect(): void {
    this.deliberatelyClosed = false
    this.openSocket()
  }

  close(): void {
    this.deliberatelyClosed = true
    this.clearTimers()
    this.ws?.close()
    this.ws = null
    this.setState('closed')
  }

  /** Send a frame; queued while connecting, dropped while closed/unstable. */
  send(frame: ClientFrame): void {
    const raw = JSON.stringify(frame)
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(raw)
    } else if (this._state === 'connecting') {
      if (this.outbox.length < 64) this.outbox.push(raw)
    }
  }

  private setState(state: VeilConnectionState): void {
    if (this._state === state) return
    this._state = state
    this.stateHandlers.forEach((h) => h(state))
  }

  private clearTimers(): void {
    if (this.reconnectTimer !== null) {
      clearTimeout(this.reconnectTimer)
      this.reconnectTimer = null
    }
    if (this.pingTimer !== null) {
      clearInterval(this.pingTimer)
      this.pingTimer = null
    }
  }

  private openSocket(): void {
    this.clearTimers()
    this.setState('connecting')
    let ws: WebSocket
    try {
      ws = this.socketFactory(this.url)
    } catch {
      this.scheduleReconnect()
      return
    }
    this.ws = ws

    ws.onopen = () => {
      this.attempts = 0
      this.setState('open')
      // Flush frames queued while connecting.
      for (const raw of this.outbox) ws.send(raw)
      this.outbox = []
      if (this.pingIntervalMs > 0) {
        this.pingTimer = setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'ping' }))
          }
        }, this.pingIntervalMs)
      }
    }

    ws.onmessage = (ev: MessageEvent) => {
      if (typeof ev.data !== 'string') return
      let frame: ServerFrame
      try {
        frame = JSON.parse(ev.data) as ServerFrame
      } catch {
        return // ignore malformed whispers
      }
      this.frameHandlers.forEach((h) => h(frame))
    }

    ws.onclose = () => {
      this.clearTimers()
      this.ws = null
      if (!this.deliberatelyClosed) this.scheduleReconnect()
      else this.setState('closed')
    }

    ws.onerror = () => {
      // onclose follows onerror; nothing extra to do here.
    }
  }

  private scheduleReconnect(): void {
    this.setState('unstable')
    const delay = Math.min(
      this.maxBackoffMs,
      this.baseBackoffMs * 2 ** this.attempts,
    )
    this.attempts++
    this.reconnectTimer = setTimeout(() => {
      if (!this.deliberatelyClosed) this.openSocket()
    }, delay)
  }
}