msg-timestamps

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Stamps } from '../types'
5
6// Reply text -> epoch ms when its first chunk arrived. The render props carry
7// no message id, so the text is the only key.
8// ponytail: text key, two identical blocks share one time; switch to an id if the props ever carry one
9const stamps = atom({ plugin: 'msg-timestamps', key: 'stamps' } as const, {} as Stamps)
10// System UTC offset in minutes, read from `date +%z` so the system timezone wins.
11const offsetMin = atom({ plugin: 'msg-timestamps', key: 'offsetMin' } as const, null as number | null)
12
13const MAX_STAMPS = 500
14
15export const key = (text: string) => text.trim()
16// Prompts get their own keys, so a prompt and a reply with the same text keep their own times.
17export const userKey = (text: string) => `user:${key(text)}`
18
19const trim = (merged: Stamps) => {
20  const keys = Object.keys(merged)
21  for (const k of keys.slice(0, Math.max(0, keys.length - MAX_STAMPS))) delete merged[k]
22  return merged
23}
24
25export const parseOffset = (z: string) => {
26  const m = /^([+-])(\d\d)(\d\d)$/.exec(z.trim())
27  return m ? (m[1] === '-' ? -1 : 1) * (Number(m[2]) * 60 + Number(m[3])) : 0
28}
29
30export const formatTime = (ms: number, offset: number) => {
31  const d = new Date(ms + offset * 60_000)
32  const pad = (n: number) => String(n).padStart(2, '0')
33  return `${pad(d.getUTCHours())}:${pad(d.getUTCMinutes())}:${pad(d.getUTCSeconds())}`
34}
35
36export const register: Register = on => {
37  on('session.start', async ($, e, next) => {
38    const { exitCode, stdout } = await $.process.run(['date', '+%z'])
39    if (exitCode === 0) await update($, offsetMin, () => parseOffset(stdout))
40    return next(e)
41  })
42
43  // Stamped before `next`, so the time is in state before the prompt's row is drawn.
44  on('prompt.submit', async ($, e, next) => {
45    const now = await $.clock.now()
46    await update($, stamps, old => trim({ ...old, [userKey(e.text)]: now }))
47    return next(e)
48  })
49
50  on('ui.render', { component: 'UserMessage' }, async ($, e, next) => {
51    const t = (await read($, stamps))[userKey(e.props.text)]
52    if (t === undefined) return next(e)
53
54    // A user row's text is drawn as plain text, not markdown, so no bold here.
55    const offset = (await read($, offsetMin)) ?? 0
56    return next({ ...e, props: { ...e.props, text: `[${formatTime(t, offset)}] ${e.props.text}` } })
57  })
58
59  on('turn.step', async function* ($, e, next) {
60    // Subagent replies are not drawn as AssistantMessage rows in the main transcript.
61    if (e.agentId) return yield* next(e)
62
63    const started = new Map<number, number>()
64    const texts = new Map<number, string>()
65    // The terminal draws a finished block once and never again, so the stamp
66    // must be in state before the engine sees the block end, not after the stream.
67    const flush = async () => {
68      if (started.size === 0) return
69      const done = [...started].map(([i, t]) => [key(texts.get(i) ?? ''), t] as const)
70      started.clear()
71      texts.clear()
72      await update($, stamps, old => trim({ ...old, ...Object.fromEntries(done) }))
73    }
74
75    const stream = next(e)
76    for await (const chunk of stream) {
77      if (chunk.kind === 'text') {
78        if (!started.has(chunk.index)) started.set(chunk.index, await $.clock.now())
79        texts.set(chunk.index, (texts.get(chunk.index) ?? '') + chunk.text)
80      } else {
81        await flush()
82      }
83      yield chunk
84    }
85    await flush()
86    return stream.result
87  })
88
89  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
90    const t = (await read($, stamps))[key(e.props.text)]
91    if (t === undefined) return next(e)
92
93    const offset = (await read($, offsetMin)) ?? 0
94    return next({ ...e, props: { ...e.props, text: `**[${formatTime(t, offset)}]** ${e.props.text}` } })
95  })
96}
97

1export type Stamps = Record<string, number>
2
3declare module 'claude-code' {
4  interface PluginState {
5    'msg-timestamps': { stamps: Stamps; offsetMin: number | null }
6  }
7}
8