+
+
+ Total
+
+
+ {formatTokens(total)}
+
+
+
+
+
+
+
+ Input ↑
+
+
+ {formatTokens(tokensTotal.in)}
+
+
+
+
+
+ Output ↓
+
+
+ {formatTokens(tokensTotal.out)}
+
+
+
+ )
+}
diff --git a/ui/features/run-detail/hooks/useRunStream.ts b/ui/features/run-detail/hooks/useRunStream.ts
index 999311f2..d4394163 100644
--- a/ui/features/run-detail/hooks/useRunStream.ts
+++ b/ui/features/run-detail/hooks/useRunStream.ts
@@ -1,24 +1,28 @@
'use client'
import { useEffect, useReducer } from 'react'
import { createSSEConnection } from '@/lib/sse'
-import { getRunStreamUrl } from '@/lib/api-client'
+import { getRun, getRunStreamUrl } from '@/lib/api-client'
import { AGENT_STEPS } from '@/lib/types/run'
import type { AgentStep } from '@/lib/types/run'
-import type { RunStreamState } from '../types'
+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 RunStreamState['reports'],
+ steps: Object.fromEntries(AGENT_STEPS.map((s) => [s, 'pending'])) as RunStreamState['steps'],
+ reports: Object.fromEntries(AGENT_STEPS.map((s) => [s, []])) as RunStreamState['reports'],
+ tokensByStep: Object.fromEntries(AGENT_STEPS.map((s) => [s, zeroTokens()])) as RunStreamState['tokensByStep'],
+ tokensTotal: zeroTokens(),
verdict: null,
error: null,
}
type Action =
- | { type: 'AGENT_START'; step: AgentStep; turn: number }
- | { type: 'AGENT_COMPLETE'; step: AgentStep; turn: number; report: string }
- | { type: 'RUN_COMPLETE'; decision: string }
- | { type: 'RUN_ERROR'; message: string }
+ | { 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 {
@@ -27,19 +31,30 @@ function reducer(state: RunStreamState, action: Action): RunStreamState {
return { ...state, status: 'running' }
case 'AGENT_START':
- // Only transition to 'running' on first turn (don't regress from 'done')
if (state.steps[action.step] !== 'pending') return state
return { ...state, steps: { ...state.steps, [action.step]: 'running' } }
- case 'AGENT_COMPLETE':
+ case 'AGENT_COMPLETE': {
+ const dIn = action.tokens_in ?? 0
+ const dOut = action.tokens_out ?? 0
+ const prev = state.tokensByStep[action.step] ?? zeroTokens()
return {
...state,
- steps: { ...state.steps, [action.step]: 'done' },
+ 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,
+ },
}
+ }
case 'RUN_COMPLETE':
return { ...state, status: 'complete', verdict: action.decision as RunStreamState['verdict'] }
@@ -56,17 +71,43 @@ export function useRunStream(runId: string): RunStreamState {
const [state, dispatch] = useReducer(reducer, initialState)
useEffect(() => {
- const url = getRunStreamUrl(runId)
- const close = createSSEConnection(url, {
- onOpen: () => dispatch({ type: 'CONNECTED' }),
- onAgentStart: ({ step, turn }) =>
- dispatch({ type: 'AGENT_START', step: step as AgentStep, turn }),
- onAgentComplete: ({ step, turn, report }) =>
- dispatch({ type: 'AGENT_COMPLETE', step: step as AgentStep, turn, report }),
- onRunComplete: ({ decision }) => dispatch({ type: 'RUN_COMPLETE', decision }),
- onRunError: ({ message }) => dispatch({ type: 'RUN_ERROR', message }),
+ 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 close
+
+ return () => { aborted = true; close?.() }
}, [runId])
return state
diff --git a/ui/features/run-detail/types.ts b/ui/features/run-detail/types.ts
index cd099543..75dff0b5 100644
--- a/ui/features/run-detail/types.ts
+++ b/ui/features/run-detail/types.ts
@@ -1,10 +1,14 @@
import type { AgentStep, RunStatus } from '@/lib/types/run'
import type { Decision, StepStatus } from '@/lib/types/agents'
+export type TokenCount = { in: number; out: number }
+
export type RunStreamState = {
status: RunStatus | 'connecting'
steps: Record