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

/**
 * Pure reducer for evaluation-run session state.
 *
 * Extracted from `useRunSession` so the event-handling logic is:
 *   - a pure function (no refs, no API calls, no WebSocket manipulation)
 *   - trivially unit-testable
 *   - free of the stale-closure hazards that bit the previous
 *     `useCallback` + `useState` implementation.
 *
 * Side effects (closing the WebSocket, refreshing the Run object) are
 * handled by the hook via a `useEffect` that watches `completed`.
 */

import type { CaseSnapshot } from '../api'
import type {
  CaseError,
  CaseState,
  RunError,
  RunProgress,
  RuleResultState,
  TurnState,
} from './useRunSession'

export interface SessionState {
  cases: CaseState[]
  progress: RunProgress | null
  scenarioSnapshot: Record<string, CaseSnapshot>
  isLive: boolean
  completed: boolean
  finalStatus: 'completed' | 'failed' | 'cancelled' | null
  errorInfo: RunError | null
}

export const initialSessionState: SessionState = {
  cases: [],
  progress: null,
  scenarioSnapshot: {},
  isLive: false,
  completed: false,
  finalStatus: null,
  errorInfo: null,
}

export type SessionAction =
  | { type: 'RESET' }
  | { type: 'SET_LIVE'; live: boolean }
  | { type: 'SET_SNAPSHOT'; snapshot: Record<string, CaseSnapshot> }
  | { type: 'LOAD_HISTORY'; cases: CaseState[]; progress: RunProgress; snapshot: Record<string, CaseSnapshot> }
  | { type: 'WS_EVENT'; event: WsEvent }
  | { type: 'FINALIZE_FROM_RUN'; status: string; summary: any }
  | { type: 'SET_RUN'; run: { id: string } | null }

/**
 * Shape of a WebSocket event payload emitted by the backend. Typed loosely
 * because the backend sends a different shape per event name.
 */
export interface WsEvent {
  event: string
  [key: string]: any
}

// ── helpers ──────────────────────────────────────────────────────────────

function findOrCreateCase(cases: CaseState[], caseId: string, isDynamic = false): {
  cases: CaseState[]
  index: number
} {
  const idx = cases.findIndex((c) => c.caseId === caseId)
  if (idx !== -1) return { cases, index: idx }
  const newCase: CaseState = {
    caseId,
    status: 'waiting',
    turns: [],
    ruleResults: [],
    errors: [],
    isDynamic,
  }
  return { cases: [...cases, newCase], index: cases.length }
}

function updateCase(cases: CaseState[], index: number, patch: Partial<CaseState>): CaseState[] {
  const next = [...cases]
  next[index] = { ...next[index], ...patch }
  return next
}

function upsertTurn(cases: CaseState[], caseIndex: number, turn: TurnState): CaseState[] {
  const cs = cases[caseIndex]
  const turns = [...cs.turns]
  const existing = turns.findIndex((t) => t.roundIndex === turn.roundIndex)
  if (existing >= 0) turns[existing] = turn
  else turns.push(turn)
  return updateCase(cases, caseIndex, { turns, status: 'running' })
}

// ── reducer ──────────────────────────────────────────────────────────────

export function sessionReducer(state: SessionState, action: SessionAction): SessionState {
  switch (action.type) {
    case 'RESET':
      return initialSessionState

    case 'SET_LIVE':
      return { ...state, isLive: action.live }

    case 'SET_SNAPSHOT':
      return { ...state, scenarioSnapshot: action.snapshot }

    case 'LOAD_HISTORY':
      return {
        ...state,
        cases: action.cases,
        progress: action.progress,
        scenarioSnapshot: action.snapshot,
      }

    case 'SET_RUN':
      // The Run object itself lives in a separate useState; this action
      // only exists to piggy-back a reset of transient session state when
      // the user selects a different run.
      return state

    case 'FINALIZE_FROM_RUN': {
      const { status, summary } = action
      if (status === 'completed') {
        return { ...state, completed: true, isLive: false, finalStatus: 'completed', errorInfo: null }
      }
      if (status === 'failed') {
        const err = summary?.error
        let final: 'failed' | 'cancelled' = 'failed'
        let errInfo: RunError | null = null
        if (err && typeof err === 'object' && err.code === 'cancelled_by_user') {
          final = 'cancelled'
          errInfo = err
        } else if (typeof err === 'string') {
          if (err === 'cancelled by user') {
            final = 'cancelled'
            errInfo = { code: 'cancelled_by_user', message: '评测已手动停止' }
          } else {
            errInfo = { message: err }
          }
        } else if (err && typeof err === 'object') {
          errInfo = err
        }
        return { ...state, completed: true, isLive: false, finalStatus: final, errorInfo: errInfo }
      }
      return state
    }

    case 'WS_EVENT': {
      const ev = action.event
      switch (ev.event) {
        case 'case_start': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id)
          return {
            ...state,
            cases: updateCase(cases, index, { status: 'running' }),
            progress: {
              done: state.progress?.done ?? 0,
              total: ev.total ?? state.progress?.total ?? 0,
              currentIndex: ev.index ?? 0,
            },
          }
        }

        case 'case_end': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id)
          return {
            ...state,
            cases: updateCase(cases, index, { status: 'done', passed: !!ev.passed }),
            progress: {
              done: (state.progress?.done ?? 0) + 1,
              total: ev.total ?? state.progress?.total ?? 0,
              currentIndex: ev.index ?? state.progress?.currentIndex ?? 0,
            },
          }
        }

        case 'messages_generated': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id, true)
          return {
            ...state,
            cases: updateCase(cases, index, {
              generatedMessages: ev.messages,
              isDynamic: true,
            }),
          }
        }

        case 'turn_start': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id)
          const turn: TurnState = {
            roundIndex: ev.round,
            message: ev.message,
            pending: true,
          }
          return { ...state, cases: upsertTurn(cases, index, turn) }
        }

        case 'turn_end': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id)
          const existing = cases[index].turns.findIndex((t) => t.roundIndex === ev.round)
          const merged: TurnState = {
            roundIndex: ev.round,
            message: existing >= 0 ? cases[index].turns[existing].message : '',
            reply: ev.reply_text,
            latency_ms: ev.latency_ms,
            pending: false,
          }
          return { ...state, cases: upsertTurn(cases, index, merged) }
        }

        case 'turn_error': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id ?? 'unknown')
          const cs = cases[index]
          const errors: CaseError[] = [...cs.errors, { round: ev.round, message: ev.error ?? '未知错误' }]
          const turns = [...cs.turns]
          const existing = turns.findIndex((t) => t.roundIndex === ev.round)
          if (existing >= 0) {
            turns[existing] = { ...turns[existing], pending: false, error: ev.error }
          }
          return { ...state, cases: updateCase(cases, index, { errors, turns }) }
        }

        case 'rule_result': {
          const { cases, index } = findOrCreateCase(state.cases, ev.case_id)
          const cs = cases[index]
          const ruleResults: RuleResultState[] = [
            ...cs.ruleResults,
            {
              rule_type: ev.rule_type,
              passed: !!ev.passed,
              score: ev.score,
              reason: ev.reason,
            },
          ]
          return { ...state, cases: updateCase(cases, index, { ruleResults }) }
        }

        case 'error': {
          if (ev.case_id) {
            const { cases, index } = findOrCreateCase(state.cases, ev.case_id)
            const cs = cases[index]
            const errors: CaseError[] = [...cs.errors, { message: ev.error ?? '未知错误' }]
            return { ...state, cases: updateCase(cases, index, { errors }) }
          }
          return {
            ...state,
            errorInfo: { message: ev.error ?? '未知错误' },
          }
        }

        case 'run_completed': {
          const status: string = ev.status ?? 'completed'
          const reason: string | undefined = ev.reason
          const errorPayload = ev.error

          let final: 'completed' | 'failed' | 'cancelled' = 'completed'
          let errInfo: RunError | null = null
          if (status === 'completed') {
            final = 'completed'
          } else if (reason === 'cancelled') {
            final = 'cancelled'
          } else {
            final = 'failed'
          }

          if (typeof errorPayload === 'string') {
            errInfo = { message: errorPayload }
          } else if (errorPayload && typeof errorPayload === 'object') {
            errInfo = { code: errorPayload.code, message: errorPayload.message ?? '评测失败' }
          }

          // Mark any still-running cases as done so the UI doesn't show
          // spinners forever after the run ends.
          const cases = state.cases.map((c) =>
            c.status === 'running' ? { ...c, status: 'done' as const } : c,
          )

          return {
            ...state,
            cases,
            completed: true,
            isLive: false,
            finalStatus: final,
            errorInfo: errInfo,
          }
        }

        default:
          return state
      }
    }

    default:
      return state
  }
}