app/frontend/src/components/Flow/useFlowSocket.ts

import type { QueryClient } from "@tanstack/react-query"
import { useQueryClient } from "@tanstack/react-query"
import { useEffect } from "react"

import { OpenAPI } from "@/client"
import { dashboardKeys, panelKeys } from "@/components/Dashboard/queries"
import { healthKeys } from "@/components/Health/queries"
import { runKeys } from "@/components/Runs/queries"
import { connectionStore } from "@/lib/connectionStore"
import { parseMediaFrame } from "@/lib/media"
import { apiToken } from "@/lib/portal"
import { type LogLine, liveStore, type ValueSource } from "./liveStore"
import { flowKeys } from "./queries"

const RECONNECT_MIN = 1000
const RECONNECT_MAX = 30000
/**
 * Consecutive token rejections before the session is given up on.
 *
 * A 1008 means the server would not take the token. Under a portal that is
 * routine — the handoff token lasts twelve hours against a local session's
 * eight days — so it has to be recoverable rather than fatal: the token is read
 * again for every attempt, so a refreshed one is picked up. A session that has
 * genuinely been revoked fails all three and falls through to the caller's auth
 * handler instead of retrying forever.
 */
const AUTH_ATTEMPTS = 3

type FlowEvent =
  | {
      type: "snapshot"
      values: Record<string, { value: unknown; ts: number | null }>
      nodes: { id: string; status: string; error?: string | null }[]
      paused?: string[]
      logs?: LogLine[]
      /** Emission counts per qualified node. Absent on older instances. */
      emits?: Record<string, number>
    }
  | {
      type: "message_value"
      name: string
      value: unknown
      ts: number
      source?: ValueSource
    }
  | { type: "node_started"; node: string }
  | {
      type: "node_queued"
      flow?: string
      node: string
      run?: string
      detail?: string
      ts?: number
    }
  | {
      type: "node_executed"
      flow?: string
      node: string
      outputs: number
      duration_ms?: number
    }
  | {
      type: "node_error"
      flow?: string
      node: string
      error: string
      ts?: number
    }
  | { type: "node_status"; node: string; status: string; error?: string | null }
  | { type: "node_error_acknowledged"; node: string; ts?: number }
  | ({ type: "node_log" } & LogLine)
  | { type: "flow_paused"; flow: string; paused: boolean }
  | {
      type: "node_health"
      flow?: string
      node: string
      health: "ok" | "down" | "unknown"
      detail?: string | null
      ts?: number
    }
  | { type: "flow_quarantined"; flow: string; error?: string; ts?: number }
  | { type: "engine_degraded"; reason?: string; ts?: number }
  | { type: "engine_fatal"; reason?: string; ts?: number }
  | {
      type: "cascade_dropped"
      flow?: string
      node?: string
      deliveries?: number
      ts?: number
    }
  | {
      type: "queue_unavailable"
      flow?: string
      node?: string
      error?: string
      ts?: number
    }
  | {
      type: "pipeline_rebuilt"
      nodes: { id: string; status: string; error?: string | null }[]
      paused?: string[]
    }
  | { type: "dashboard_changed"; dashboard?: string; ts?: number }
  | {
      type: "run_started" | "run_finished"
      flow: string
      run: string
      status?: string
      group?: string
      ts?: number
    }
  | {
      type: "run_metric"
      flow: string
      run: string
      names: string[]
      ts?: number
    }

function socketUrl(): string {
  const base = String(OpenAPI.BASE || window.location.origin)
  // Concatenated rather than resolved: an absolute path as the second argument
  // to `new URL` discards the base's own path, which under a portal
  // (`https://host/i/{id}`) would aim the socket at the wrong place entirely.
  const url = new URL(`${base.replace(/\/$/, "")}/api/v1/flows/ws`)
  url.protocol = url.protocol === "https:" ? "wss:" : "ws:"
  // Browsers cannot set headers on a websocket handshake, so the token rides
  // in the query string. Read per attempt, so a reconnect after the old one
  // expired carries whatever is current rather than what we started with.
  url.searchParams.set("token", apiToken())
  return url.toString()
}

/*
 * One socket for the page, owned here rather than by a component.
 *
 * Several places want live events at once — the app shell, the brain graph on
 * Home, the editor — and React mounts and unmounts them in an order none of
 * them controls: passive effects run children before parents, and StrictMode
 * runs mount → unmount → mount. Keeping the socket inside whichever hook
 * instance ran first made that one its owner, so leaving Home closed the socket
 * while the shell kept the count above zero and nothing ever reopened it. The
 * connection lives out here now with a plain refcount; a component only says
 * whether it is still watching.
 */
let socket: WebSocket | null = null
let timer: ReturnType<typeof setTimeout> | null = null
let retry = RECONNECT_MIN
/** How many mounted components want the socket open. */
let watchers = 0
/** Consecutive 1008s, cleared by a handshake the server accepts. */
let rejected = 0
let client: QueryClient | null = null
const authHandlers = new Set<() => void>()

function schedule() {
  if (timer) return
  // Jittered, because every client of an engine that restarted is counting
  // the same backoff from the same moment: without it they all come back
  // together, and keep coming back together.
  timer = setTimeout(
    () => {
      timer = null
      connect()
    },
    retry * (0.5 + Math.random()),
  )
  retry = Math.min(retry * 2, RECONNECT_MAX)
}

/**
 * Say which messages this page wants frames for.
 *
 * Bytes are the one thing the socket does not send unasked: a camera is
 * hundreds of kilobytes a second and most tabs are drawing no media at all. So
 * a Media tile or a thumbnail registers its message, and this tells the engine
 * — again on every reconnect, since the new socket knows nothing.
 */
function sendWanted() {
  if (socket?.readyState !== WebSocket.OPEN) return
  socket.send(JSON.stringify({ type: "media", names: liveStore.wantedNames() }))
}

liveStore.onWanted(sendWanted)

function connect() {
  if (watchers === 0 || socket || timer) return
  const ws = new WebSocket(socketUrl())
  ws.binaryType = "arraybuffer"
  socket = ws

  ws.onopen = () => {
    retry = RECONNECT_MIN
    rejected = 0
    liveStore.setConnected(true)
    connectionStore.setSocketOpen(true)
    // Whatever happened while the socket was down was missed, so nothing
    // held in cache can be trusted to still be current. Scoped to what this
    // socket actually feeds: an unqualified invalidation refetches every
    // query the page holds, and the usual reason the socket dropped is the
    // engine restarting — so every open tab and every wall panel did that at
    // once, at the moment it was least able to answer.
    for (const queryKey of [
      flowKeys.all,
      dashboardKeys.all,
      panelKeys.all,
      runKeys.all,
      healthKeys.all,
    ]) {
      client?.invalidateQueries({ queryKey })
    }
    sendWanted()
  }

  const handle = (message: FlowEvent) => {
    switch (message.type) {
      case "snapshot":
        liveStore.setValues(message.values)
        liveStore.setStatuses(message.nodes ?? [])
        liveStore.setPausedFlows(message.paused ?? [])
        liveStore.setLogs(message.logs ?? [])
        // Missing against an instance older than this bundle; the graph
        // then starts from zero the way it always did.
        liveStore.setEmits(message.emits ?? {})
        break
      case "message_value":
        liveStore.setValue(message.name, {
          value: message.value,
          ts: message.ts,
          source: message.source,
        })
        break
      case "node_queued":
        // Waiting for a machine, which looks exactly like hung from outside.
        // `node_started` overwrites this, so nothing has to clear it.
        liveStore.setStatus(message.node, {
          status: "queued",
          detail: message.detail,
        })
        break
      case "node_started":
        liveStore.setStatus(message.node, { status: "running" })
        break
      case "node_executed":
        liveStore.setStatus(message.node, { status: "success" })
        if (message.outputs > 0) liveStore.recordEmit(message.node)
        break
      case "node_error":
        liveStore.setStatus(message.node, {
          status: "error",
          error: message.error,
        })
        liveStore.recordEngineEvent({
          type: message.type,
          flow: message.flow,
          node: message.node,
          detail: message.error,
          ts: message.ts ?? Date.now() / 1000,
        })
        break
      case "node_health":
        // The canvas draws node health from the flow detail's `issues`, which
        // the server derives — so the screen only moved on mount, navigation
        // or a rebuild, never when health actually flipped. The store had a
        // health map of its own and nothing ever read it.
        client?.invalidateQueries({
          queryKey: message.flow ? flowKeys.detail(message.flow) : flowKeys.all,
        })
        if (message.health === "down") {
          liveStore.recordEngineEvent({
            type: message.type,
            flow: message.flow,
            node: message.node,
            detail: message.detail ?? "Reported itself down.",
            ts: message.ts ?? Date.now() / 1000,
          })
        }
        break
      case "flow_quarantined":
      case "engine_degraded":
      case "engine_fatal":
      case "cascade_dropped":
      case "queue_unavailable":
        liveStore.recordEngineEvent({
          type: message.type,
          flow: "flow" in message ? message.flow : undefined,
          node: "node" in message ? message.node : undefined,
          detail:
            ("error" in message ? message.error : undefined) ??
            ("reason" in message ? message.reason : undefined) ??
            "",
          ts: message.ts ?? Date.now() / 1000,
        })
        break
      case "node_status":
        liveStore.setStatus(message.node, {
          status: message.status as "active" | "error",
          error: message.error,
        })
        break
      case "node_error_acknowledged":
        // Someone dismissed it, here or in another browser. The engine has
        // already forgotten it, so drop the marker rather than posting back.
        liveStore.clearFailure(message.node)
        break
      case "node_log":
        liveStore.appendLog(message)
        break
      case "flow_paused":
        liveStore.setPaused(message.flow, message.paused)
        break
      case "pipeline_rebuilt":
        liveStore.setStatuses(message.nodes)
        liveStore.setPausedFlows(message.paused ?? [])
        // Someone published or started a flow, here or in another tab: the
        // markers on the flow chips are stale until the list is refetched.
        client?.invalidateQueries({ queryKey: flowKeys.all })
        break
      case "dashboard_changed":
        // Someone published, or changed which dashboards a panel shows. Refetch
        // rather than reload: a wall screen must not blank or sign in again.
        // `exact` matters — the list key is a prefix of every detail key,
        // including the draft an editor may have open.
        client?.invalidateQueries({ queryKey: dashboardKeys.all, exact: true })
        client?.invalidateQueries({
          queryKey: message.dashboard
            ? dashboardKeys.detail(message.dashboard)
            : panelKeys.all,
        })
        break
      case "run_started":
      case "run_finished":
        // One invalidation covers the lot: the list, the flow counts, the run
        // being watched, and any chart drawing a run's curve. A run is not a
        // live value, so nothing here goes through the live store.
        client?.invalidateQueries({ queryKey: runKeys.all })
        break
      case "run_metric":
        // A batch of readings was written. Only what draws them is refetched —
        // the run's own detail (its metric names and curve hang below that
        // key) and any comparison — so a training run reporting every second
        // does not re-read the list behind it.
        client?.invalidateQueries({ queryKey: runKeys.detail(message.run) })
        client?.invalidateQueries({ queryKey: runKeys.compares })
        break
    }
  }

  ws.onmessage = (event) => {
    // A frame that is not JSON, or one this bundle cannot read, costs the
    // frame rather than the connection: an exception thrown here escapes into
    // `window.onerror` and leaves whatever it had already applied behind.
    try {
      if (event.data instanceof ArrayBuffer) {
        // Media, sent in front of the value that names it — so the tile has
        // the bytes by the time it hears the message changed.
        const frame = parseMediaFrame(event.data)
        if (frame) liveStore.setBytes(frame.header.digest, frame.bytes)
        return
      }
      const payload = JSON.parse(event.data)
      if (payload?.type === "batch") {
        // A cascade publishes a dozen events at once and the engine coalesces
        // them into one frame. An instance older than this bundle sends
        // them one at a time, which is the branch below.
        for (const message of payload.events ?? []) handle(message)
      } else {
        handle(payload)
      }
    } catch (error) {
      console.warn("Dropped an unreadable socket frame", error)
    }
  }

  ws.onclose = (event) => {
    // A socket we already dropped: its close says nothing about the connection
    // we want now.
    if (socket !== ws) return
    socket = null
    liveStore.setConnected(false)
    connectionStore.setSocketOpen(false)
    if (watchers === 0) return
    if (event.code === 1008) {
      rejected += 1
      if (rejected >= AUTH_ATTEMPTS) {
        for (const handler of authHandlers) handler()
        return
      }
    }
    // 1013 is the portal saying the instance is not attached — the one
    // close code that means "offline" rather than "the socket dropped".
    if (event.code === 1013) {
      connectionStore.setOffline(null)
    }
    schedule()
  }
}

function release(onAuthFailure?: () => void) {
  watchers -= 1
  if (onAuthFailure) authHandlers.delete(onAuthFailure)
  if (watchers > 0) return
  if (timer) {
    clearTimeout(timer)
    timer = null
  }
  const ws = socket
  socket = null
  ws?.close()
  liveStore.setConnected(false)
  // Its `onclose` is dropped by the guard above, so say so here instead.
  connectionStore.setSocketOpen(false)
}

/**
 * Keeps one socket open while anything on the page wants live events.
 *
 * @param onAuthFailure called once the server has refused the token often
 *   enough to mean the session is gone, so the caller can send the user back to
 *   the login screen.
 */
export function useFlowSocket(onAuthFailure?: () => void): void {
  const queryClient = useQueryClient()

  useEffect(() => {
    client = queryClient
    if (onAuthFailure) authHandlers.add(onAuthFailure)
    watchers += 1
    connect()
    return () => release(onAuthFailure)
  }, [onAuthFailure, queryClient])
}