tori/web/src/ws.ts

export interface WsPlanUpdate {
  type: 'PlanUpdate'
  workflow_id: string
  steps: { order: number; description: string; command: string }[]
}

export interface WsStepStatusUpdate {
  type: 'StepStatusUpdate'
  step_id: string
  status: string
  output: string
}

export interface WsWorkflowStatusUpdate {
  type: 'WorkflowStatusUpdate'
  workflow_id: string
  status: string
}

export interface WsError {
  type: 'Error'
  message: string
}

export type WsMessage = WsPlanUpdate | WsStepStatusUpdate | WsWorkflowStatusUpdate | WsError

export type WsHandler = (msg: WsMessage) => void

export function connectWs(projectId: string, onMessage: WsHandler): { close: () => void } {
  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:'
  const url = `${proto}//${location.host}/ws/${projectId}`
  let ws: WebSocket | null = null
  let reconnectTimer: ReturnType<typeof setTimeout> | null = null
  let closed = false

  function connect() {
    if (closed) return
    ws = new WebSocket(url)

    ws.onmessage = (e) => {
      try {
        const msg: WsMessage = JSON.parse(e.data)
        onMessage(msg)
      } catch {
        // ignore malformed messages
      }
    }

    ws.onclose = () => {
      if (!closed) {
        reconnectTimer = setTimeout(connect, 2000)
      }
    }

    ws.onerror = () => {
      ws?.close()
    }
  }

  connect()

  return {
    close() {
      closed = true
      if (reconnectTimer) clearTimeout(reconnectTimer)
      ws?.close()
    },
  }
}