125 lines
4.6 KiB
TypeScript
125 lines
4.6 KiB
TypeScript
'use client'
|
|
import { useEffect, useReducer } from 'react'
|
|
import { createSSEConnection } from '@/lib/sse'
|
|
import { getRun, getRunStreamUrl } from '@/lib/api-client'
|
|
import { AGENT_STEPS } from '@/lib/types/run'
|
|
import type { AgentStep } from '@/lib/types/run'
|
|
import type { ChiefAnalystReport } from '@/lib/types/agents'
|
|
import type { RunStreamState, TokenCount } from '../types'
|
|
|
|
const zeroTokens = (): TokenCount => ({ in: 0, out: 0 })
|
|
|
|
const initialState: RunStreamState = {
|
|
status: 'connecting',
|
|
steps: Object.fromEntries(AGENT_STEPS.map((s) => [s, 'pending'])) as RunStreamState['steps'],
|
|
reports: Object.fromEntries(AGENT_STEPS.map((s) => [s, []])) as unknown as RunStreamState['reports'],
|
|
tokensByStep: Object.fromEntries(AGENT_STEPS.map((s) => [s, zeroTokens()])) as RunStreamState['tokensByStep'],
|
|
tokensTotal: zeroTokens(),
|
|
verdict: null,
|
|
error: null,
|
|
chiefAnalystReport: null,
|
|
}
|
|
|
|
type Action =
|
|
| { type: 'AGENT_START'; step: AgentStep; turn: number }
|
|
| { type: 'AGENT_COMPLETE'; step: AgentStep; turn: number; report: string; tokens_in?: number; tokens_out?: number }
|
|
| { type: 'RUN_COMPLETE'; decision: string }
|
|
| { type: 'RUN_ERROR'; message: string }
|
|
| { type: 'CONNECTED' }
|
|
|
|
function reducer(state: RunStreamState, action: Action): RunStreamState {
|
|
switch (action.type) {
|
|
case 'CONNECTED':
|
|
return { ...state, status: 'running' }
|
|
|
|
case 'AGENT_START':
|
|
if (state.steps[action.step] !== 'pending') return state
|
|
return { ...state, steps: { ...state.steps, [action.step]: 'running' } }
|
|
|
|
case 'AGENT_COMPLETE': {
|
|
const dIn = action.tokens_in ?? 0
|
|
const dOut = action.tokens_out ?? 0
|
|
const prev = state.tokensByStep[action.step] ?? zeroTokens()
|
|
|
|
let chiefAnalystReport: ChiefAnalystReport | null = state.chiefAnalystReport
|
|
if (action.step === 'chief_analyst') {
|
|
try { chiefAnalystReport = JSON.parse(action.report) as ChiefAnalystReport }
|
|
catch { chiefAnalystReport = null }
|
|
}
|
|
|
|
return {
|
|
...state,
|
|
steps: { ...state.steps, [action.step]: 'done' },
|
|
reports: {
|
|
...state.reports,
|
|
[action.step]: [...(state.reports[action.step] ?? []), action.report],
|
|
},
|
|
tokensByStep: {
|
|
...state.tokensByStep,
|
|
[action.step]: { in: prev.in + dIn, out: prev.out + dOut },
|
|
},
|
|
tokensTotal: {
|
|
in: state.tokensTotal.in + dIn,
|
|
out: state.tokensTotal.out + dOut,
|
|
},
|
|
chiefAnalystReport,
|
|
}
|
|
}
|
|
|
|
case 'RUN_COMPLETE':
|
|
return { ...state, status: 'complete', verdict: action.decision as RunStreamState['verdict'] }
|
|
|
|
case 'RUN_ERROR':
|
|
return { ...state, status: 'error', error: action.message }
|
|
|
|
default:
|
|
return state
|
|
}
|
|
}
|
|
|
|
export function useRunStream(runId: string): RunStreamState {
|
|
const [state, dispatch] = useReducer(reducer, initialState)
|
|
|
|
useEffect(() => {
|
|
let close: (() => void) | undefined
|
|
let aborted = false
|
|
|
|
getRun(runId).then((run) => {
|
|
if (aborted) return
|
|
|
|
if (run.status === 'complete' && run.reports) {
|
|
dispatch({ type: 'CONNECTED' })
|
|
for (const [key, report] of Object.entries(run.reports)) {
|
|
const lastColon = key.lastIndexOf(':')
|
|
const step = key.slice(0, lastColon) as AgentStep
|
|
const turn = parseInt(key.slice(lastColon + 1), 10)
|
|
const tok = run.token_usage?.[key] ?? { tokens_in: 0, tokens_out: 0 }
|
|
dispatch({ type: 'AGENT_START', step, turn })
|
|
dispatch({ type: 'AGENT_COMPLETE', step, turn, report,
|
|
tokens_in: tok.tokens_in, tokens_out: tok.tokens_out })
|
|
}
|
|
dispatch({ type: 'RUN_COMPLETE', decision: run.decision ?? 'HOLD' })
|
|
return
|
|
}
|
|
|
|
const url = getRunStreamUrl(runId)
|
|
close = createSSEConnection(url, {
|
|
onOpen: () => dispatch({ type: 'CONNECTED' }),
|
|
onAgentStart: ({ step, turn }) =>
|
|
dispatch({ type: 'AGENT_START', step: step as AgentStep, turn }),
|
|
onAgentComplete: ({ step, turn, report, tokens_in, tokens_out }) =>
|
|
dispatch({ type: 'AGENT_COMPLETE', step: step as AgentStep, turn, report,
|
|
tokens_in, tokens_out }),
|
|
onRunComplete: ({ decision }) => dispatch({ type: 'RUN_COMPLETE', decision }),
|
|
onRunError: ({ message }) => dispatch({ type: 'RUN_ERROR', message }),
|
|
})
|
|
}).catch(() => {
|
|
if (!aborted) dispatch({ type: 'RUN_ERROR', message: 'Failed to load run' })
|
|
})
|
|
|
|
return () => { aborted = true; close?.() }
|
|
}, [runId])
|
|
|
|
return state
|
|
}
|