AgentEvalTool/frontend/web/src/hooks/useRunSession.ts

import { useCallback, useEffect, useReducer, useRef, useState } from 'react'
import { runsApi, type CaseSnapshot, type Run } from '../api'
import {
  initialSessionState,
  sessionReducer,
  type WsEvent,
} from './sessionReducer'

// WebSocket reconnect config
const WS_MAX_RETRIES = 5
const WS_BASE_DELAY_MS = 1000  // 1s, 2s, 4s, 8s, 16s (capped at 30s)
const WS_MAX_DELAY_MS = 30000

export interface TurnState {
  roundIndex: number
  message: string
  reply?: string
  latency_ms?: number | null
  pending?: boolean
  error?: string
}

export interface CaseError {
  round?: number
  message: string
}

export interface RuleResultState {
  rule_type: string
  passed: boolean
  score?: number | null
  reason?: string | null
}

export interface CaseState {
  caseId: string
  status: 'waiting' | 'running' | 'done'
  turns: TurnState[]
  ruleResults: RuleResultState[]
  generatedMessages?: string[]
  errors: CaseError[]
  isDynamic?: boolean
  passed?: boolean
}

export interface RunProgress {
  done: number
  total: number
  currentIndex: number
}

export interface RunError {
  code?: string
  message: string
}

export interface RunSession {
  run: Run | null
  cases: CaseState[]
  progress: RunProgress | null
  scenarioSnapshot: Record<string, CaseSnapshot>
  isLive: boolean
  completed: boolean
  finalStatus: 'completed' | 'failed' | 'cancelled' | null
  errorInfo: RunError | null
  select: (run: Run | null, opts?: { live?: boolean }) => void
  cancel: () => Promise<void>
  refreshRun: () => Promise<void>
}

/**
 * Manages the lifecycle of one evaluation run: WebSocket events, historical
 * log loading, and polling. State transitions are driven by a pure reducer
 * (see `sessionReducer.ts`) so the logic is testable and free of the
 * stale-closure hazards that bit the previous useState-based implementation.
 */
export function useRunSession(): RunSession {
  // Run object is managed separately because it's driven by REST, not by
  // the WebSocket event stream.
  const [run, setRun] = useState<Run | null>(null)
  const [state, dispatch] = useReducer(sessionReducer, initialSessionState)

  const wsRef = useRef<WebSocket | null>(null)
  const pollRef = useRef<number | null>(null)
  const selectedIdRef = useRef<string | null>(null)
  const reconnectTimerRef = useRef<number | null>(null)
  const reconnectCountRef = useRef<number>(0)

  const clearPolling = () => {
    if (pollRef.current) {
      window.clearInterval(pollRef.current)
      pollRef.current = null
    }
  }

  const clearReconnect = () => {
    if (reconnectTimerRef.current) {
      window.clearTimeout(reconnectTimerRef.current)
      reconnectTimerRef.current = null
    }
  }

  const closeWs = () => {
    clearReconnect()
    if (wsRef.current) {
      try { wsRef.current.close() } catch { /* noop */ }
      wsRef.current = null
    }
  }

  // Side effect of `completed` transitioning to true: close the WS and
  // refresh the Run object so the UI shows the final status/summary.
  useEffect(() => {
    if (!state.completed) return
    closeWs()
    const id = selectedIdRef.current
    if (!id) return
    runsApi.get(id).then((res) => setRun(res.data)).catch(() => { /* noop */ })
  }, [state.completed])

  const loadHistoricalLogs = useCallback(async (runId: string) => {
    try {
      const res = await runsApi.logs(runId)
      const data = res.data
      const snapshot = data.scenario_snapshot ?? {}

      const caseMap = new Map<string, CaseState>()
      const ensure = (caseId: string): CaseState => {
        let cs = caseMap.get(caseId)
        if (!cs) {
          const snap = snapshot[caseId]
          cs = {
            caseId,
            status: 'done',
            turns: [],
            ruleResults: [],
            errors: [],
            isDynamic: snap?.type === 'dynamic',
          }
          caseMap.set(caseId, cs)
        }
        return cs
      }

      for (const t of data.turns) {
        const cs = ensure(t.case_id)
        cs.turns.push({
          roundIndex: t.round_index,
          message: t.sent_text,
          reply: t.reply_text,
          latency_ms: t.latency_ms,
        })
      }
      for (const r of data.results) {
        const cs = ensure(r.case_id)
        cs.ruleResults.push({
          rule_type: r.rule_type,
          passed: r.passed,
          score: r.score,
          reason: r.reason,
        })
      }
      const verdicts = data.case_verdicts ?? {}
      for (const cs of caseMap.values()) {
        cs.turns.sort((a, b) => a.roundIndex - b.roundIndex)
        cs.passed = verdicts[cs.caseId]?.passed ?? false
      }
      const cases = Array.from(caseMap.values())
      const total = cases.length

      dispatch({ type: 'LOAD_HISTORY', cases, progress: { done: total, total, currentIndex: total }, snapshot })
    } catch { /* noop */ }
  }, [])

  const select = useCallback((r: Run | null, opts?: { live?: boolean }) => {
    closeWs()
    clearPolling()
    dispatch({ type: 'RESET' })
    setRun(r)
    selectedIdRef.current = r?.id ?? null
    reconnectCountRef.current = 0

    if (!r) return

    const live = opts?.live ?? (r.status === 'running' || r.status === 'pending')
    dispatch({ type: 'SET_LIVE', live })

    if (live) {
      connectWs(r.id)
    } else {
      loadHistoricalLogs(r.id)
      dispatch({ type: 'FINALIZE_FROM_RUN', status: r.status, summary: r.summary })
    }
  }, [loadHistoricalLogs])  // eslint-disable-line react-hooks/exhaustive-deps

  function connectWs(runId: string) {
    if (selectedIdRef.current !== runId) return

    const proto = window.location.protocol === 'https:' ? 'wss' : 'ws'
    const ws = new WebSocket(`${proto}://${window.location.host}/ws/runs/${runId}`)
    wsRef.current = ws

    ws.onmessage = (msg) => {
      try {
        const ev = JSON.parse(msg.data) as WsEvent
        dispatch({ type: 'WS_EVENT', event: ev })
      } catch { /* noop */ }
    }

    ws.onclose = (e) => {
      if (selectedIdRef.current !== runId) return

      // Normal close (code 1000) or run already completed → finalize
      if (e.code === 1000 || e.wasClean) {
        dispatch({ type: 'SET_LIVE', live: false })
        runsApi.get(runId).then((res) => {
          setRun(res.data)
          dispatch({ type: 'FINALIZE_FROM_RUN', status: res.data.status, summary: res.data.summary })
        }).catch(() => { /* noop */ })
        return
      }

      // Abnormal close → attempt exponential backoff reconnect
      const retries = reconnectCountRef.current
      if (retries >= WS_MAX_RETRIES) {
        dispatch({ type: 'SET_LIVE', live: false })
        runsApi.get(runId).then((res) => {
          setRun(res.data)
          dispatch({ type: 'FINALIZE_FROM_RUN', status: res.data.status, summary: res.data.summary })
        }).catch(() => { /* noop */ })
        return
      }

      const delay = Math.min(WS_BASE_DELAY_MS * Math.pow(2, retries), WS_MAX_DELAY_MS)
      reconnectCountRef.current = retries + 1
      reconnectTimerRef.current = window.setTimeout(() => connectWs(runId), delay)
    }
  }

  const cancel = useCallback(async () => {
    if (!run) return
    await runsApi.cancel(run.id)
  }, [run])

  const refreshRun = useCallback(async () => {
    if (!run) return
    try {
      const res = await runsApi.get(run.id)
      setRun(res.data)
    } catch { /* noop */ }
  }, [run])

  // Poll during live runs to keep the Run object fresh (status/summary/completed_at)
  useEffect(() => {
    if (!state.isLive || !run) {
      clearPolling()
      return
    }
    pollRef.current = window.setInterval(() => {
      runsApi.get(run.id).then((res) => {
        if (selectedIdRef.current === run.id) setRun(res.data)
      }).catch(() => { /* noop */ })
    }, 3000)
    return () => clearPolling()
  }, [state.isLive, run?.id])

  useEffect(() => () => { closeWs(); clearPolling(); clearReconnect() }, [])

  return {
    run,
    cases: state.cases,
    progress: state.progress,
    scenarioSnapshot: state.scenarioSnapshot,
    isLive: state.isLive,
    completed: state.completed,
    finalStatus: state.finalStatus,
    errorInfo: state.errorInfo,
    select,
    cancel,
    refreshRun,
  }
}