task-progress

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, Timer } from 'claude-code'
3
4import type { Run, SubAgent } from '../types'
5
6const run = atom({ plugin: 'task-progress', key: 'run' } as const, null)
7const agents = atom({ plugin: 'task-progress', key: 'agents' } as const, [])
8const clock = atom({ plugin: 'task-progress', key: 'now' } as const, 0)
9
10const LIVE = ['pending', 'running', 'waiting']
11
12const bar = (frac: number, width = 20) => {
13  const n = Math.round(Math.min(Math.max(frac, 0), 1) * width)
14  return '█'.repeat(n) + '░'.repeat(width - n)
15}
16
17const fmt = (ms: number) => {
18  const s = Math.max(0, Math.round(ms / 1000))
19  return s >= 60 ? `${Math.floor(s / 60)}m${s % 60}s` : `${s}s`
20}
21
22const MIN_TURN_MS = 15_000
23
24const ema = (avg: number | undefined, ms: number) => (avg ? avg * 0.7 + ms * 0.3 : ms)
25
26// Active time: wall clock minus time spent waiting on the user.
27export const active = (r: Run, now: number) =>
28  now - r.startedAt - (r.pausedMs ?? 0) - (r.pausedAt === undefined ? 0 : now - r.pausedAt)
29
30// Todo progress when the agent keeps a list; else elapsed against the average length.
31// Step pace is measured from when the list appeared, so exploration before it isn't spread over the steps.
32// ponytail: linear pace per step, weight later steps if the tail still runs long
33export const estimate = (r: Run, now: number) => {
34  const elapsed = active(r, now)
35  if (r.total > 0 && r.done > 0) {
36    const start = r.todoAt ?? 0
37    const perStep = (elapsed - start) / r.done
38    return { frac: r.done / r.total, elapsed, totalMs: elapsed + perStep * (r.total - r.done), isTodo: true }
39  }
40  if (r.total > 0) return { frac: 0, elapsed, totalMs: undefined, isTodo: true }
41  if (r.avgMs) return { frac: Math.min(elapsed / r.avgMs, 0.95), elapsed, totalMs: r.avgMs, isTodo: false }
42  return { frac: 0, elapsed, totalMs: undefined, isTodo: false }
43}
44
45const line = (r: Run, now: number) => {
46  const { frac, elapsed, totalMs, isTodo } = estimate(r, now)
47  const steps = isTodo ? `  ${r.done}/${r.total} 步` : ''
48  const time = totalMs === undefined
49    ? `已用 ${fmt(elapsed)} · 預估中…`
50    : elapsed > totalMs
51      ? `已用 ${fmt(elapsed)} · 已超出平均 ${fmt(elapsed - totalMs)}`
52      : `已用 ${fmt(elapsed)} · 剩餘 ~${fmt(totalMs - elapsed)} / 預估 ${fmt(totalMs)}`
53  return `${bar(frac)} ${Math.round(frac * 100)}%${steps}  ${time}`
54}
55
56// Applies a todo change to the main run, or to the subagent whose loop made the call.
57// The first change that gives the list steps stamps todoAt.
58const onTodos = async ($: EngineInterface, agentId: string | undefined, fn: (r: Run) => Partial<Run>) => {
59  const now = await $.clock.now()
60  const apply = <T extends Run>(r: T): T => {
61    const next = { ...r, ...fn(r) }
62    return next.todoAt === undefined && next.total > 0 ? { ...next, todoAt: active(next, now) } : next
63  }
64  return agentId === undefined
65    ? update($, run, r => (r ? apply(r) : r))
66    : update($, agents, list => list.map(a => (a.id === agentId ? apply(a) : a)))
67}
68
69// Tools that always wait on the user; any other main-loop call waits when its check says ask.
70const ASKS_USER = ['AskUserQuestion', 'ExitPlanMode']
71
72const pause = async ($: EngineInterface, id: string | undefined) => {
73  const now = await $.clock.now()
74  await update($, run, r => (r && r.pausedAt === undefined ? { ...r, pausedAt: now, pauseId: id } : r))
75}
76
77const resume = async ($: EngineInterface, id: string | undefined) => {
78  const now = await $.clock.now()
79  await update($, run, r =>
80    r && r.pausedAt !== undefined && r.pauseId === id
81      ? { ...r, pausedMs: (r.pausedMs ?? 0) + now - r.pausedAt, pausedAt: undefined, pauseId: undefined }
82      : r,
83  )
84}
85
86let tick: Timer | undefined
87
88function ensureTick($: EngineInterface) {
89  tick ??= $.clock.every(1000, async () => {
90    const now = await $.clock.now()
91    await update($, clock, () => now)
92
93    const tracked = await read($, agents)
94    if (tracked.length > 0) {
95      const live = new Set((await $.agent.list()).filter(a => LIVE.includes(a.status)).map(a => a.id))
96      const ended = tracked.filter(a => !live.has(a.id))
97      for (const a of ended) {
98        const key = `avgAgentMs:${a.type}`
99        await $.store.set(key, ema((await $.store.get(key)) as number | undefined, now - a.startedAt))
100      }
101      if (ended.length > 0) await update($, agents, list => list.filter(a => live.has(a.id)))
102    }
103
104    if ((await read($, run)) === null && (await read($, agents)).length === 0) {
105      tick?.cancel()
106      tick = undefined
107    }
108  })
109}
110
111export const register: Register = on => {
112  on('prompt.submit', async ($, e, next) => {
113    const now = await $.clock.now()
114    const avgKey = `avgTurnMs:${await $.session.root()}`
115    const avgMs = ((await $.store.get(avgKey)) ?? (await $.store.get('avgTurnMs'))) as number | undefined
116    await update($, clock, () => now)
117    await update($, run, () => ({ startedAt: now, done: 0, total: 0, avgMs, avgKey }))
118    ensureTick($)
119    return next(e)
120  }).catch(($, e, next) => next(e))
121
122  on('agent.spawn', async ($, e, next) => {
123    const res = await next(e)
124    if (res.deny === undefined && res.agentId) {
125      const id = res.agentId
126      const now = await $.clock.now()
127      const avgMs = (await $.store.get(`avgAgentMs:${e.subagentType}`)) as number | undefined
128      const label = e.description.length > 24 ? `${e.description.slice(0, 23)}…` : e.description
129      const a: SubAgent = { id, type: e.subagentType, label, startedAt: now, done: 0, total: 0, avgMs }
130      await update($, agents, list => [...list.filter(x => x.id !== id), a])
131      ensureTick($)
132    }
133    return res
134  }).catch(($, e, next) => next(e))
135
136  // ponytail: an asked call's pause also covers the tool's own run after approval; fine unless approved calls run long
137  on('tool.check', async ($, e, next) => {
138    const res = await next(e)
139    if (e.agentId === undefined && e.tool_use_id && res.decision === 'ask') await pause($, e.tool_use_id)
140    return res
141  }).catch(($, e, next) => next(e))
142
143  on('tool.call', async ($, e, next) => {
144    if (e.agentId === undefined && ASKS_USER.includes(e.tool)) await pause($, e.tool_use_id)
145    try {
146      return await next(e)
147    } finally {
148      if (e.agentId === undefined) await resume($, e.tool_use_id).catch(() => {})
149    }
150  }).catch(($, e, next) => next(e))
151
152  on('tool.call', { tool: 'TodoWrite' }, async ($, e, next) => {
153    const done = e.todos.filter(t => t.status === 'completed').length
154    await onTodos($, e.agentId, () => ({ done, total: e.todos.length }))
155    return next(e)
156  }).catch(($, e, next) => next(e))
157
158  on('tool.call', { tool: 'TaskCreate' }, async ($, e, next) => {
159    await onTodos($, e.agentId, r => ({ total: r.total + 1 }))
160    return next(e)
161  }).catch(($, e, next) => next(e))
162
163  on('tool.call', { tool: 'TaskUpdate' }, async ($, e, next) => {
164    if (e.status === 'completed') await onTodos($, e.agentId, r => ({ done: r.done + 1 }))
165    if (e.status === 'deleted') await onTodos($, e.agentId, r => ({ total: r.total - 1 }))
166    return next(e)
167  }).catch(($, e, next) => next(e))
168
169  // Subagent loops raise turn.complete too; only the main loop ends the run.
170  // Short or interrupted turns would drag the average away from real tasks, so they don't count.
171  on('turn.complete', async ($, e, next) => {
172    if (e.agentId !== undefined) return next(e)
173    const r = await read($, run)
174    if (r && !e.isAborted) {
175      const ms = active(r, await $.clock.now())
176      if (ms >= MIN_TURN_MS) {
177        if (r.avgKey) await $.store.set(r.avgKey, ema(r.avgMs, ms))
178        await $.store.set('avgTurnMs', ema((await $.store.get('avgTurnMs')) as number | undefined, ms))
179      }
180    }
181    await update($, run, () => null)
182    return next(e)
183  })
184
185  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
186    const r = await read($, run)
187    const subs = await read($, agents)
188    const showMain = e.props.isWorking && r !== null
189    if (e.props.hasSurvey || (!showMain && subs.length === 0)) return next(e)
190
191    const below = await next(e)
192    const now = await read($, clock)
193    const { Box, Text } = $.ui.resolve(e)
194
195    return (
196      <Box flexDirection="column">
197        {showMain && r ? <Text color="cyan">任務 {line(r, now)}</Text> : null}
198        {subs.map(a => (
199          <Text key={a.id} color="magenta">
200            {'  ↳ '}{a.label} <Text dimColor>({a.type})</Text> {line(a, now)}
201          </Text>
202        ))}
203        {below}
204      </Box>
205    )
206  })
207}
208

1export type Run = {
2  startedAt: number
3  done: number
4  total: number
5  avgMs?: number
6  // Time spent waiting on the user (permission dialogs, questions), excluded from estimates.
7  pausedMs?: number
8  pausedAt?: number
9  pauseId?: string
10  // Active time when the todo list first appeared; step pace is measured from here.
11  todoAt?: number
12  avgKey?: string
13}
14
15export type SubAgent = Run & { id: string; type: string; label: string }
16
17declare module 'claude-code' {
18  interface PluginState {
19    'task-progress': { run: Run | null; agents: SubAgent[]; now: number }
20  }
21}
22