Time Estimate

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, TurnCompleteInput } from 'claude-code'
3
4import type { Job, TrackedAgent } from '../types'
5import { reconcile } from './agents'
6import { MAX_HISTORY, NO_TASKS, applyTool, describe, describeCountdown, estimate, formatDuration, nextCountdown } from './estimate'
7
8const job = atom({ plugin: 'time-estimate', key: 'job' } as const, null)
9const now = atom({ plugin: 'time-estimate', key: 'now' } as const, 0)
10const tasks = atom({ plugin: 'time-estimate', key: 'tasks' } as const, NO_TASKS)
11const agents = atom({ plugin: 'time-estimate', key: 'agents' } as const, [])
12const starts = atom({ plugin: 'time-estimate', key: 'starts' } as const, {})
13
14const HISTORY_KEY = 'durations'
15const AGENT_HISTORY_KEY = 'agentDurations'
16const AGENT_MIN_SAMPLE = 1
17const TRACKED_TOOLS = ['TaskCreate', 'TaskUpdate', 'TodoWrite']
18
19const asHistory = (stored: unknown): number[] =>
20  Array.isArray(stored) ? stored.filter((ms): ms is number => typeof ms === 'number') : []
21
22const asAgentHistory = (stored: unknown): Record<string, number[]> =>
23  typeof stored === 'object' && stored !== null && !Array.isArray(stored)
24    ? Object.fromEntries(Object.entries(stored).map(([label, list]) => [label, asHistory(list)]))
25    : {}
26
27let history: number[] = []
28let agentHistory: Record<string, number[]> = {}
29
30const completeMain = async ($: EngineInterface, e: TurnCompleteInput) => {
31  const current = await read($, job)
32  if (current === null || current.turnId !== e.turnId) return
33
34  await update($, job, () => null)
35  if (e.reason !== 'answer') return
36  history = [...history, e.durationMs].slice(-MAX_HISTORY)
37  await $.store.set(HISTORY_KEY, history)
38}
39
40const completeAgent = async ($: EngineInterface, e: TurnCompleteInput, agentId: string) => {
41  const tracked = (await read($, agents)).find(agent => agent.id === agentId)
42  const isTimed = e.reason === 'answer' && tracked !== undefined && tracked.isSeenFromStart
43  if (!isTimed) return
44
45  const past = agentHistory[tracked.label] ?? []
46  agentHistory = { ...agentHistory, [tracked.label]: [...past, e.durationMs].slice(-MAX_HISTORY) }
47  await $.store.set(AGENT_HISTORY_KEY, agentHistory)
48}
49
50const refreshCountdown = async ($: EngineInterface, at: number) => {
51  const progress = await read($, tasks)
52  await update($, job, current => {
53    if (current === null) return null
54    const guess = estimate(Math.max(0, at - current.startedAt), progress, current.doneAtStart, history)
55
56    return { ...current, countdown: nextCountdown(current.countdown, at, guess, progress.done) }
57  })
58}
59
60const agentLine = (agent: TrackedAgent, at: number): string => {
61  const elapsed = Math.max(0, at - agent.startedAt)
62  const children = { total: agent.childIds.length, done: agent.doneChildIds.length }
63  const guess = estimate(elapsed, children, 0, agentHistory[agent.label] ?? [], AGENT_MIN_SAMPLE)
64  const since = agent.isSeenFromStart ? '' : '+'
65  const childText = children.total > 0 ? ` · sub-agents ${children.done}/${children.total}` : ''
66
67  return `${formatDuration(elapsed)}${since}${childText} · ${describe(guess, 'runs', AGENT_MIN_SAMPLE)}`
68}
69
70export const register: Register = on => {
71  on('session.start', async ($, e, next) => {
72    history = asHistory(await $.store.get(HISTORY_KEY))
73    agentHistory = asAgentHistory(await $.store.get(AGENT_HISTORY_KEY))
74    $.clock.every(1000, async () => {
75      const at = await $.clock.now()
76      const listed = await $.agent.list()
77      const knownStarts = await read($, starts)
78      await update($, agents, tracked => reconcile(tracked, listed, knownStarts, at))
79      const isBusy = (await read($, agents)).length > 0 || (await read($, job)) !== null
80      if (isBusy) await update($, now, () => at)
81      await refreshCountdown($, at)
82    })
83
84    return next(e)
85  })
86
87  on('turn.start', async ($, e, next) => {
88    const startedAt = await $.clock.now()
89    const agentId = e.agentId
90    if (agentId !== undefined) {
91      await update($, starts, known => (agentId in known ? known : { ...known, [agentId]: startedAt }))
92      return next(e)
93    }
94
95    const current: Job = { turnId: e.turnId, startedAt, doneAtStart: (await read($, tasks)).done }
96    await update($, job, () => current)
97    await update($, now, () => startedAt)
98    await refreshCountdown($, startedAt)
99
100    return next(e)
101  })
102
103  on('tool.call', async ($, e, next) => {
104    const ran = await next(e)
105    const isCounted = TRACKED_TOOLS.includes(String(e.tool)) && ran.deny === undefined && ran.isError !== true
106    if (isCounted) await update($, tasks, current => applyTool(current, String(e.tool), e))
107
108    return ran
109  }).catch(($, e, next) => next(e))
110
111  on('turn.complete', async ($, e, next) => {
112    await (e.agentId === undefined ? completeMain($, e) : completeAgent($, e, e.agentId))
113
114    return next(e)
115  })
116
117  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
118    const current = await read($, job)
119    const running = await read($, agents)
120    const main = e.props.isWorking ? current : null
121    if (e.props.hasSurvey || (main === null && running.length === 0)) return next(e)
122
123    const { Box, Text } = $.ui.resolve(e)
124    const at = await read($, now)
125    const progress = await read($, tasks)
126    const room = Math.max(1, e.props.maxRows - 3 - (main === null ? 0 : 1))
127    const mainElapsed = main === null ? 0 : Math.max(0, at - main.startedAt)
128
129    return (
130      <Box key="time-estimate" flexDirection="column" marginTop={1} marginLeft={4} paddingX={2} borderStyle="round" borderColor="claude">
131        {main !== null && (
132          <Box key="main">
133            <Text bold color="claude">
134              ⏳ {describeCountdown(main.countdown, at, estimate(mainElapsed, progress, main.doneAtStart, history))}
135            </Text>
136          </Box>
137        )}
138        {running.slice(0, room).map(agent => (
139          <Box key={`agent-${agent.id}`}>
140            <Text bold color="claude">⚙ {agent.label}</Text>
141            <Text dimColor> · {agentLine(agent, at)}</Text>
142          </Box>
143        ))}
144      </Box>
145    )
146  })
147}
148

1import type { AgentInfo } from 'claude-code'
2
3import type { TrackedAgent } from '../types'
4
5const ACTIVE = ['pending', 'running', 'waiting']
6
7type Listed = Pick<AgentInfo, 'id' | 'description' | 'type' | 'status' | 'parentId' | 'name'>
8
9const union = (a: string[], b: string[]): string[] => [...new Set([...a, ...b])]
10
11const labelOf = (agent: Listed): string => agent.description || agent.name || agent.type
12
13export const reconcile = (
14  tracked: TrackedAgent[],
15  listed: Listed[],
16  starts: Record<string, number>,
17  at: number,
18): TrackedAgent[] =>
19  listed
20    .filter(agent => agent.parentId === undefined && ACTIVE.includes(agent.status))
21    .map(agent => {
22      const known = tracked.find(one => one.id === agent.id)
23      const children = listed.filter(child => child.parentId === agent.id)
24      const doneChildren = children.filter(child => !ACTIVE.includes(child.status))
25
26      return {
27        id: agent.id,
28        label: labelOf(agent),
29        startedAt: starts[agent.id] ?? known?.startedAt ?? at,
30        isSeenFromStart: starts[agent.id] !== undefined,
31        childIds: union(known?.childIds ?? [], children.map(child => child.id)),
32        doneChildIds: union(known?.doneChildIds ?? [], doneChildren.map(child => child.id)),
33      }
34    })
35

1import type { Countdown, Tasks } from '../types'
2
3export const MIN_SAMPLE = 5
4export const MAX_HISTORY = 50
5
6export type Estimate =
7  | { kind: 'tasks'; remainingMs: number }
8  | { kind: 'history'; remainingMs: number; sample: number }
9  | { kind: 'beyond'; sample: number }
10  | { kind: 'learning'; sample: number }
11
12export const NO_TASKS: Tasks = { total: 0, done: 0, doneIds: [] }
13
14const median = (values: number[]): number => {
15  const sorted = [...values].sort((a, b) => a - b)
16  const mid = Math.floor(sorted.length / 2)
17  const upper = sorted[mid] ?? 0
18
19  return sorted.length % 2 === 1 ? upper : ((sorted[mid - 1] ?? upper) + upper) / 2
20}
21
22// Conditioned on the turn having already run `elapsedMs`: only past turns that
23// outlasted it say anything about how much longer this one will take.
24const fromHistory = (elapsedMs: number, history: number[], minSample: number): Estimate => {
25  if (history.length < minSample) return { kind: 'learning', sample: history.length }
26
27  const longer = history.filter(ms => ms > elapsedMs)
28
29  return longer.length === 0
30    ? { kind: 'beyond', sample: history.length }
31    : { kind: 'history', remainingMs: median(longer) - elapsedMs, sample: history.length }
32}
33
34export type Progress = Pick<Tasks, 'total' | 'done'>
35
36export const estimate = (
37  elapsedMs: number,
38  tasks: Progress,
39  doneAtStart: number,
40  history: number[],
41  minSample = MIN_SAMPLE,
42): Estimate => {
43  const doneThisJob = tasks.done - doneAtStart
44  const left = tasks.total - tasks.done
45
46  return doneThisJob > 0 && left > 0
47    ? { kind: 'tasks', remainingMs: (elapsedMs / doneThisJob) * left }
48    : fromHistory(elapsedMs, history, minSample)
49}
50
51type TaskUpdateInput = { taskId: string; status?: string }
52type TodoWriteInput = { todos: ReadonlyArray<{ status: string }> }
53
54const applyUpdate = (tasks: Tasks, { taskId, status }: TaskUpdateInput): Tasks => {
55  const wasDone = tasks.doneIds.includes(taskId)
56
57  switch (true) {
58    case status === 'completed' && !wasDone:
59      return { ...tasks, done: tasks.done + 1, doneIds: [...tasks.doneIds, taskId] }
60    case status === 'deleted':
61      return {
62        total: Math.max(0, tasks.total - 1),
63        done: wasDone ? tasks.done - 1 : tasks.done,
64        doneIds: tasks.doneIds.filter(id => id !== taskId),
65      }
66    default:
67      return tasks
68  }
69}
70
71const fromTodos = ({ todos }: TodoWriteInput): Tasks => ({
72  total: todos.length,
73  done: todos.filter(todo => todo.status === 'completed').length,
74  doneIds: [],
75})
76
77export const applyTool = (tasks: Tasks, tool: string, input: unknown): Tasks => {
78  switch (tool) {
79    case 'TaskCreate':
80      return { ...tasks, total: tasks.total + 1 }
81    case 'TaskUpdate':
82      return applyUpdate(tasks, input as TaskUpdateInput)
83    case 'TodoWrite':
84      return fromTodos(input as TodoWriteInput)
85    default:
86      return tasks
87  }
88}
89
90export const formatDuration = (ms: number): string => {
91  const seconds = Math.max(0, Math.round(ms / 1000))
92  const h = Math.floor(seconds / 3600)
93  const m = Math.floor((seconds % 3600) / 60)
94  const s = seconds % 60
95  const pad = (n: number) => String(n).padStart(2, '0')
96
97  switch (true) {
98    case h > 0:
99      return `${h}h ${pad(m)}m`
100    case m > 0:
101      return `${m}m ${pad(s)}s`
102    default:
103      return `${s}s`
104  }
105}
106
107export const describe = (e: Estimate, noun = 'turns', minSample = MIN_SAMPLE): string => {
108  switch (e.kind) {
109    case 'tasks':
110      return `~${formatDuration(e.remainingMs)} left`
111    case 'history':
112      return `usually ~${formatDuration(e.remainingMs)} more (${e.sample} past ${noun})`
113    case 'beyond':
114      return `longer than your last ${e.sample} ${noun}`
115    case 'learning':
116      return `learning your pace (${e.sample}/${minSample} ${noun})`
117  }
118}
119
120// The end time is held between ticks so the display counts down one second per second;
121// it moves only when the estimate's footing changes (a task finished) or the time runs out.
122export const nextCountdown = (
123  prev: Countdown | undefined,
124  at: number,
125  guess: Estimate,
126  done: number,
127): Countdown | undefined => {
128  if (!('remainingMs' in guess)) return undefined
129  const basis = `${guess.kind}:${done}`
130  const isStale = prev === undefined || prev.basis !== basis || prev.endsAt <= at
131
132  return isStale ? { endsAt: at + guess.remainingMs, basis } : prev
133}
134
135export const describeCountdown = (countdown: Countdown | undefined, at: number, guess: Estimate): string => {
136  switch (true) {
137    case countdown !== undefined:
138      return `${formatDuration(countdown.endsAt - at)} left`
139    case guess.kind === 'learning':
140      return 'estimating…'
141    default:
142      return 'longer than usual'
143  }
144}
145

1export type Countdown = { endsAt: number; basis: string }
2
3export type Job = { turnId: string; startedAt: number; doneAtStart: number; countdown?: Countdown }
4
5export type Tasks = { total: number; done: number; doneIds: string[] }
6
7export type TrackedAgent = {
8  id: string
9  label: string
10  startedAt: number
11  isSeenFromStart: boolean
12  childIds: string[]
13  doneChildIds: string[]
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    'time-estimate': {
19      job: Job | null
20      now: number
21      tasks: Tasks
22      agents: TrackedAgent[]
23      starts: Record<string, number>
24    }
25  }
26}
27