agent-usage

1import type { EngineInterface, Register, SessionRateLimit } from 'claude-code'
2
3type SessionState = 'idle' | 'running' | 'waiting' | 'ended'
4
5// These tools block on the user until they return; everything else waits only while a permission dialog is open
6const QUESTION_TOOLS = new Set(['AskUserQuestion', 'ExitPlanMode'])
7const HEARTBEAT_MS = 60_000
8
9let turnActive = false
10let updatedAt = ''
11// tool_use_id -> tool name, for calls currently waiting on the user
12const waitingOn = new Map<string, string>()
13
14const currentState = (): SessionState => (waitingOn.size > 0 ? 'waiting' : turnActive ? 'running' : 'idle')
15
16const homeDir = async ($: EngineInterface) => (await $.env.get('USERPROFILE')) ?? (await $.env.get('HOME'))
17
18const writeUsage = async ($: EngineInterface, rateLimits: SessionRateLimit[]) => {
19  if (rateLimits.length === 0) return
20  const home = await homeDir($)
21  if (home === undefined) return
22  const windows = Object.fromEntries(
23    rateLimits.map(r => [r.kind, { percentUsed: r.percentUsed, resetsAt: r.resetsAt ?? null }]),
24  )
25  const updatedAt = new Date(await $.clock.now()).toISOString()
26  await $.fs.write(`${home}/.agent-usage/claude.json`, JSON.stringify({ updatedAt, ...windows }, null, 2))
27}
28
29// updatedAt marks the last state change; heartbeatAt alone moves on the timer so the reader can drop dead sessions
30const writeSession = async ($: EngineInterface, state: SessionState, sessionId?: string, touch = true) => {
31  const home = await homeDir($)
32  if (home === undefined) return
33  const id = sessionId ?? (await $.session.id())
34  const now = new Date(await $.clock.now()).toISOString()
35  if (touch || updatedAt === '') updatedAt = now
36  const detail = state === 'waiting' ? ([...waitingOn.values()].at(-1) ?? null) : null
37  const record = { sessionId: id, cwd: await $.session.cwd(), state, detail, updatedAt, heartbeatAt: now }
38  await $.fs.write(`${home}/.agent-usage/sessions/${id}.json`, JSON.stringify(record, null, 2))
39}
40
41const heartbeat = ($: EngineInterface) => writeSession($, currentState(), undefined, false)
42
43const stopWaiting = async ($: EngineInterface, ...keys: string[]) => {
44  const removed = keys.filter(k => waitingOn.delete(k))
45  if (removed.length > 0) await writeSession($, currentState())
46}
47
48export const register: Register = on => {
49  on('session.start', async ($, e, next) => {
50    const result = await next(e)
51    await writeUsage($, (await $.session.usage()).rateLimits)
52    // Not gated on isInteractive: the VS Code extension runs through the SDK and never got a heartbeat
53    await writeSession($, currentState())
54    $.clock.every(HEARTBEAT_MS, () => heartbeat($))
55    return result
56  })
57
58  on('session.measure', async ($, e, next) => {
59    if (e.changed.includes('rateLimits')) await writeUsage($, e.rateLimits)
60    return next(e)
61  })
62
63  on('turn.start', async ($, e, next) => {
64    turnActive = true
65    await writeSession($, currentState())
66    return next(e)
67  })
68
69  on('turn.complete', async ($, e, next) => {
70    if (e.agentId === undefined) {
71      turnActive = false
72      waitingOn.clear()
73      await writeSession($, currentState())
74    }
75    return next(e)
76  })
77
78  // tool.check answers 'ask' before the mode (e.g. auto) settles it, so only PermissionRequest means a dialog is shown.
79  // It carries no tool_use_id, hence the key by tool name, cleared when that tool's call resolves.
80  on('classic.PermissionRequest', async ($, e, next) => {
81    waitingOn.set(`perm:${e.tool_name}`, e.tool_name)
82    await writeSession($, currentState())
83    return next(e)
84  })
85
86  on('tool.call', async ($, e, next) => {
87    if (QUESTION_TOOLS.has(e.tool)) {
88      waitingOn.set(e.tool_use_id, e.tool)
89      await writeSession($, currentState())
90    }
91    try {
92      return await next(e)
93    } finally {
94      await stopWaiting($, e.tool_use_id, `perm:${e.tool}`)
95    }
96  })
97
98  on('session.end', async ($, e, next) => {
99    turnActive = false
100    waitingOn.clear()
101    await writeSession($, 'ended', e.sessionId)
102    return next(e)
103  })
104}
105