timekeeper

1import { atom, memberOf, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { RunningTool } from '../types'
5import {
6  backgroundIdOf,
7  formatClock,
8  formatElapsed,
9  labelOf,
10  localOffsetMinutes,
11  parseNotification,
12  textKey,
13} from './format'
14import type { TaskNotification } from './format'
15
16// A tool call shows in the band once it has run this long.
17const SHOW_AFTER_MS = 3_000
18// A turn at least this long ends with a toast.
19const TOAST_AFTER_MS = 30_000
20
21const STAMPS = { plugin: 'timekeeper', key: 'stamps' } as const
22const stamps = atom(STAMPS, null)
23const running = atom({ plugin: 'timekeeper', key: 'running' } as const, [])
24const now = atom({ plugin: 'timekeeper', key: 'now' } as const, 0)
25const lastRow = atom({ plugin: 'timekeeper', key: 'lastRow' } as const, null)
26
27const clock = (ms: number) => formatClock(ms, localOffsetMinutes(ms))
28// Where a turn's closing block is marked: under its row id and its text's key.
29const finalKey = (id: string) => `final:${id}`
30
31// Marks the main loop's latest reply block as the one that closed its turn.
32async function markFinal($: EngineInterface) {
33  const row = await read($, lastRow)
34
35  if (row !== null) {
36    const at = await $.clock.now()
37    await $.state.set({ ...STAMPS, id: finalKey(row.uuid) }, at)
38    await $.state.set({ ...STAMPS, id: finalKey(row.key) }, at)
39    await update($, lastRow, () => null)
40  }
41}
42
43// Takes a background task that ended out of the band and says how it ended.
44async function settle($: EngineInterface, note: TaskNotification) {
45  const ended = (await read($, running)).find(
46    t => t.isBackground && (t.id === note.taskId || t.callId === note.toolUseId),
47  )
48
49  if (ended === undefined) {
50    return
51  }
52
53  await update($, running, list => list.filter(t => t.id !== ended.id))
54
55  const at = await $.clock.now()
56  const what = ended.label === '' ? ended.tool : `${ended.tool} ${ended.label}`
57  $.ui.toast(`Background ${what}: ${note.status} after ${formatElapsed(at - ended.startedAt)}`, {
58    timeoutMs: 8000,
59  })
60}
61
62export const register: Register = on => {
63  on('session.start', async ($, e, next) => {
64    // Foreground calls left over from before a reload are no longer tracked;
65    // background tasks run on, and their notifications still end them.
66    await update($, running, list => list.filter(t => t.isBackground === true))
67    // A reload lands when a turn ends, and may be what saw that turn's end.
68    await markFinal($).catch(() => {})
69
70    $.clock.every(1000, async () => {
71      if ((await read($, running)).length > 0) {
72        const at = await $.clock.now()
73        await update($, now, () => at)
74      }
75    })
76
77    return next(e)
78  })
79
80  // Messages carry no time of their own: note when each row is stored, under
81  // its id and under its text.
82  on('session.append', async ($, e, next) => {
83    try {
84      const { type, isMeta, content } = e.message
85      const isMessage = (type === 'user' || type === 'assistant') && isMeta !== true
86      const texts = content.flatMap(block =>
87        block.type === 'text' && typeof block.text === 'string' ? [block.text] : [],
88      )
89
90      // A background task's end reaches the conversation as a notification row.
91      if (type !== 'assistant' && e.agentId === undefined) {
92        for (const text of texts) {
93          const note = parseNotification(text)
94
95          if (note !== undefined) {
96            await settle($, note)
97          }
98        }
99      }
100
101      if (isMessage && e.agentId === undefined) {
102        if (texts.length > 0) {
103          const at = await $.clock.now()
104          await $.state.set({ ...STAMPS, id: e.uuid }, at)
105
106          for (const text of texts) {
107            await $.state.set({ ...STAMPS, id: textKey(text) }, at)
108          }
109
110          const last = texts.at(-1)
111
112          if (type === 'assistant' && last !== undefined) {
113            await update($, lastRow, () => ({ uuid: e.uuid, key: textKey(last) }))
114          }
115        }
116      }
117    } catch {
118      // A row is stored whether or not its time was noted.
119    }
120
121    return next(e)
122  })
123
124  on('ui.render', { component: 'UserMessage' }, async ($, e, next) => {
125    const at =
126      (await read($, memberOf(stamps, e))) ??
127      (await $.state.get({ ...STAMPS, id: textKey(e.props.text) })).value
128
129    if (typeof at !== 'number' || e.props.text.trim() === '') {
130      return next(e)
131    }
132
133    return next({ ...e, props: { ...e.props, text: `${e.props.text}  ▶ ${clock(at)}` } })
134  })
135
136  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
137    const at =
138      (await read($, memberOf(stamps, e))) ??
139      (await $.state.get({ ...STAMPS, id: textKey(e.props.text) })).value
140
141    if (typeof at !== 'number' || e.props.text.trim() === '') {
142      return next(e)
143    }
144
145    // The block that closed its turn carries the stop mark.
146    const isFinal =
147      (await $.state.get({ ...STAMPS, id: finalKey(e.requestId) })).value != null ||
148      (await $.state.get({ ...STAMPS, id: finalKey(textKey(e.props.text)) })).value != null
149    const stamp = isFinal ? `■ ${clock(at)}` : clock(at)
150
151    return next({ ...e, props: { ...e.props, text: `${e.props.text}\n\n${stamp}` } })
152  })
153
154  on('tool.call', async ($, e, next) => {
155    try {
156      const startedAt = await $.clock.now()
157      const entry: RunningTool = {
158        id: e.tool_use_id,
159        callId: e.tool_use_id,
160        tool: e.tool,
161        label: labelOf(e),
162        startedAt,
163        isSubagent: e.agentId !== undefined,
164        isBackground: false,
165      }
166
167      await update($, running, list => [...list.filter(t => t.id !== entry.id), entry])
168      await update($, now, () => startedAt)
169    } catch {
170      // The call runs whether or not it is tracked.
171    }
172
173    // A call that handed its work to the background stays in the band under
174    // the task's id, until the task's notification ends it.
175    let taskId: string | undefined
176    let stopped: string | undefined
177
178    try {
179      const ran = await next(e)
180
181      if (ran.deny === undefined && ran.isError !== true) {
182        taskId = backgroundIdOf(e.tool, ran.result)
183        stopped = e.tool === 'TaskStop' ? (e.task_id ?? e.shell_id) : undefined
184      }
185
186      return ran
187    } finally {
188      await update($, running, list =>
189        list.flatMap(t => {
190          if (t.id === stopped) {
191            return []
192          }
193
194          if (t.id !== e.tool_use_id) {
195            return [t]
196          }
197
198          return taskId === undefined ? [] : [{ ...t, id: taskId, isBackground: true }]
199        }),
200      ).catch(() => {})
201    }
202  })
203
204  on('turn.complete', async ($, e, next) => {
205    if (e.agentId === undefined) {
206      try {
207        // The main loop's calls have all settled by now.
208        await update($, running, list => list.filter(t => t.isSubagent || t.isBackground))
209        await markFinal($)
210
211        if (e.durationMs >= TOAST_AFTER_MS && !e.isAborted) {
212          const at = await $.clock.now()
213          $.ui.toast(`Finished ${clock(at)}, took ${formatElapsed(e.durationMs)}`, { timeoutMs: 8000 })
214        }
215      } catch {
216        // The turn ends either way.
217      }
218    }
219
220    return next(e)
221  })
222
223  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
224    if (e.props.hasSurvey) {
225      return next(e)
226    }
227
228    const at = await read($, now)
229    const shown = (await read($, running)).filter(
230      t => t.isBackground || at - t.startedAt >= SHOW_AFTER_MS,
231    )
232
233    if (shown.length === 0) {
234      return next(e)
235    }
236
237    const { Box, Button, Text } = $.ui.resolve(e)
238
239    return (
240      <Box flexDirection="column">
241        {shown.map(t => (
242          <Text wrap="truncate-end">
243            <Text color="warning">{formatElapsed(at - t.startedAt)}</Text>
244            <Text bold> {t.tool}</Text>
245            <Text>{t.label === '' ? '' : ` ${t.label}`}</Text>
246            <Text dimColor>
247              {' '}
248              ({t.isBackground ? 'background, ' : ''}since {clock(t.startedAt)}
249              {t.isSubagent ? ', subagent' : ''})
250            </Text>
251          </Text>
252        ))}
253        {shown.some(t => t.isBackground) ? (
254          <Button
255            key="clear-background"
256            label="clear background rows"
257            plain
258            dimColor
259            onPress={() => update($, running, list => list.filter(t => !t.isBackground))}
260          />
261        ) : (
262          ''
263        )}
264      </Box>
265    )
266  })
267}
268

1const pad = (n: number) => String(n).padStart(2, '0')
2
3// Minutes east of UTC at that moment, as the environment's own clock has it.
4export const localOffsetMinutes = (ms: number) => -new Date(ms).getTimezoneOffset()
5
6export const formatClock = (ms: number, offsetMinutes: number) => {
7  const d = new Date(ms + offsetMinutes * 60_000)
8
9  return `${pad(d.getUTCHours())}:${pad(d.getUTCMinutes())}:${pad(d.getUTCSeconds())}`
10}
11
12export const formatElapsed = (ms: number) => {
13  const total = Math.max(0, Math.floor(ms / 1000))
14  const h = Math.floor(total / 3600)
15  const m = Math.floor((total % 3600) / 60)
16  const s = total % 60
17
18  if (h > 0) {
19    return `${h}h ${pad(m)}m ${pad(s)}s`
20  }
21
22  return m > 0 ? `${m}m ${pad(s)}s` : `${s}s`
23}
24
25// A short id for a message's text (FNV-1a), for rows whose drawn id is not
26// the id they were stored under.
27export const textKey = (text: string) => {
28  const trimmed = text.trim()
29  let hash = 0x811c9dc5
30
31  for (let i = 0; i < trimmed.length; i += 1) {
32    hash ^= trimmed.charCodeAt(i)
33    hash = Math.imul(hash, 0x01000193)
34  }
35
36  return `text:${trimmed.length}:${(hash >>> 0).toString(16)}`
37}
38
39const LABEL_KEYS = ['description', 'command', 'file_path', 'path', 'pattern', 'url', 'query', 'skill']
40
41export const labelOf = (input: Readonly<Record<string, unknown>>) => {
42  for (const key of LABEL_KEYS) {
43    const value = input[key]
44
45    if (typeof value === 'string' && value.trim() !== '') {
46      const line = value.trim().split('\n')[0] ?? ''
47
48      return line.length > 60 ? `${line.slice(0, 59)}…` : line
49    }
50  }
51
52  return ''
53}
54
55const stringAt = (record: Readonly<Record<string, unknown>>, key: string) => {
56  const value = record[key]
57
58  return typeof value === 'string' && value !== '' ? value : undefined
59}
60
61// The id of the background task a tool call left running, read from the
62// call's result; undefined when the call finished its work itself.
63export const backgroundIdOf = (tool: string, result: unknown) => {
64  if (typeof result !== 'object' || result === null) {
65    return undefined
66  }
67
68  const record = result as Readonly<Record<string, unknown>>
69
70  if (record.status === 'async_launched') {
71    return stringAt(record, 'agentId')
72  }
73
74  if (record.status === 'remote_launched' || tool === 'Monitor' || tool === 'Workflow') {
75    return stringAt(record, 'taskId')
76  }
77
78  return stringAt(record, 'backgroundTaskId')
79}
80
81export type TaskNotification = { taskId?: string; toolUseId?: string; status: string }
82
83const tagOf = (text: string, tag: string) =>
84  new RegExp(`<${tag}>\\s*([^<]*?)\\s*</${tag}>`).exec(text)?.[1] || undefined
85
86// What a background task's notification says of it; undefined for any other
87// text, and for a notification that reports no end (it names no status).
88export const parseNotification = (text: string): TaskNotification | undefined => {
89  if (!text.includes('<task-notification>')) {
90    return undefined
91  }
92
93  const status = tagOf(text, 'status')
94
95  if (status === undefined) {
96    return undefined
97  }
98
99  return { taskId: tagOf(text, 'task-id'), toolUseId: tagOf(text, 'tool-use-id'), status }
100}
101

1// The main loop's latest reply block: its row id and its text's key.
2export type LastRow = { uuid: string; key: string }
3
4// `id` is the call's tool_use_id while it runs in the foreground, and the
5// background task's id once the call handed its work to the background.
6export type RunningTool = {
7  id: string
8  callId: string
9  tool: string
10  label: string
11  startedAt: number
12  isSubagent: boolean
13  isBackground: boolean
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    timekeeper: {
19      stamps: StateFamily<number | null>
20      running: RunningTool[]
21      now: number
22      lastRow: LastRow | null
23    }
24  }
25}
26