session-gauges

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { BackgroundTask, ContextFill } from '../types'
5
6const context = atom({ plugin: 'session-gauges', key: 'context' } as const, null)
7const tasks = atom({ plugin: 'session-gauges', key: 'tasks' } as const, {})
8
9const BAR_CELLS = 20
10const FINISHED = new Set(['completed', 'failed', 'killed', 'stopped', 'cancelled'])
11const NOTIFICATION = /<task-notification>([\s\S]*?)<\/task-notification>/g
12
13type Usage = { tokens?: number; window: number; percent?: number }
14
15// The engine reports the fill once a response of the live window carried one.
16const setContext = async ($: EngineInterface, usage: Usage) => {
17  if (usage.tokens === undefined || usage.percent === undefined) return
18  const fill: ContextFill = { tokens: usage.tokens, window: usage.window, percent: usage.percent }
19  await update($, context, () => fill)
20}
21
22const addTask = ($: EngineInterface, id: string | undefined, task: BackgroundTask) =>
23  id ? update($, tasks, all => ({ ...all, [id]: task })) : undefined
24
25const dropTasks = ($: EngineInterface, ids: string[]) =>
26  ids.length === 0
27    ? undefined
28    : update($, tasks, all => {
29        const kept = { ...all }
30        for (const id of ids) delete kept[id]
31        return kept
32      })
33
34// A task notification names the task and how it ended; ids of finished ones.
35export const finishedIds = (text: string): string[] => {
36  const ids: string[] = []
37  for (const [, body = ''] of text.matchAll(NOTIFICATION)) {
38    const id = /<task-id>\s*([^<\s]+)\s*<\/task-id>/.exec(body)?.[1]
39    const status = /<status>\s*([^<\s]+)\s*<\/status>/.exec(body)?.[1] ?? 'completed'
40    if (id && FINISHED.has(status.toLowerCase())) ids.push(id)
41  }
42  return ids
43}
44
45const clip = (s: string, n = 40) => (s.length > n ? `${s.slice(0, n - 1)}…` : s)
46
47export const formatTokens = (n: number) =>
48  n >= 1_000_000 ? `${(n / 1_000_000).toFixed(1)}M` : n >= 1000 ? `${Math.round(n / 1000)}k` : `${n}`
49
50export const bar = (percent: number) => {
51  const filled = Math.min(BAR_CELLS, Math.round((Math.max(0, percent) / 100) * BAR_CELLS))
52  return { filled: '█'.repeat(filled), empty: '░'.repeat(BAR_CELLS - filled) }
53}
54
55const fillColor = (percent: number) => (percent >= 80 ? 'error' : percent >= 60 ? 'warning' : 'success')
56
57// Every gauge is display only: if its bookkeeping throws, the event goes on as
58// if the mod were not there (each gating hook's .catch passes it through).
59export const register: Register = on => {
60  on('session.start', async ($, e, next) => {
61    const started = await next(e)
62    await setContext($, (await $.session.usage()).context)
63    return started
64  })
65
66  on('session.measure', async ($, e, next) => {
67    await setContext($, e.context)
68    return next(e)
69  })
70
71  // Launches: each tool that can leave work running reports the task's id.
72  on('tool.call', async ($, e, next) => {
73    const ran = await next(e)
74    if (ran.deny !== undefined || ran.isError) return ran
75
76    if (e.tool === 'Bash' || e.tool === 'PowerShell') {
77      const result = ran.result as { backgroundTaskId?: string } | undefined
78      await addTask($, result?.backgroundTaskId, { kind: 'shell', label: clip(e.description ?? e.command) })
79    } else if (e.tool === 'Agent') {
80      const result = ran.result as { status?: string; agentId?: string; taskId?: string } | undefined
81      if (result?.status === 'async_launched') {
82        await addTask($, result.agentId, { kind: 'agent', label: clip(e.description) })
83      } else if (result?.status === 'remote_launched') {
84        await addTask($, result.taskId, { kind: 'remote', label: clip(e.description) })
85      }
86    } else if (e.tool === 'Monitor' || e.tool === 'Workflow') {
87      const result = ran.result as { taskId?: string } | undefined
88      await addTask($, result?.taskId, { kind: e.tool.toLowerCase(), label: e.tool })
89    } else if (e.tool === 'TaskStop') {
90      await dropTasks($, [e.task_id ?? e.shell_id].filter((id): id is string => !!id))
91    }
92    return ran
93  }).catch(($, e, next) => next(e))
94
95  // Completions: the task notification lands in the conversation as a row.
96  on('session.append', async ($, e, next) => {
97    const text = e.message.content
98      .map(block => (block.type === 'text' ? block.text : ''))
99      .join('\n')
100    if (text.includes('<task-notification>')) await dropTasks($, finishedIds(text))
101    return next(e)
102  }).catch(($, e, next) => next(e))
103
104  // At each stop the engine lists what is still in flight: take it as truth,
105  // which also picks up a command the person backgrounded with ctrl+b.
106  on('classic.Stop', async ($, e, next) => {
107    const live = e.background_tasks
108    if (live) {
109      await update($, tasks, all => {
110        const inFlight: Record<string, BackgroundTask> = {}
111        for (const t of live) {
112          inFlight[t.id] = all[t.id] ?? { kind: t.type, label: clip(t.description || t.command || t.type) }
113        }
114        return inFlight
115      })
116    }
117    return next(e)
118  }).catch(($, e, next) => next(e))
119
120  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
121    if (e.props.hasSurvey) return next(e)
122
123    const fill = await read($, context)
124    const running = Object.values(await read($, tasks))
125    if (fill === null && running.length === 0) return next(e)
126
127    const { Box, Button, Text } = $.ui.resolve(e)
128    const cells = fill ? bar(fill.percent) : null
129
130    // The desktop app's own view tool opens its background-task pane.
131    const openTasks = async () => {
132      const shown = await $.tool.call({ tool: 'mcp__ccd_view__show_pane', pane: 'tasks' } as never)
133      if (shown.deny !== undefined || shown.isError) $.ui.toast('Could not open the tasks pane')
134    }
135
136    return (
137      <Box flexDirection="row" gap={1}>
138        {fill && cells ? (
139          <Text key="context">
140            <Text color={fillColor(fill.percent)}>{cells.filled}</Text>
141            <Text dimColor>{cells.empty}</Text>
142            <Text> {fill.percent}% </Text>
143            <Text dimColor>
144              {formatTokens(fill.tokens)} / {formatTokens(fill.window)}
145            </Text>
146          </Text>
147        ) : null}
148        {running.length > 0 ? (
149          <Button key="tasks" plain onPress={openTasks}>
150            <Text color="warning">
151              {fill ? '· ' : ''}⏳ {running.length} background {running.length === 1 ? 'task' : 'tasks'}:{' '}
152              {running.map(t => `${t.kind} (${t.label})`).join(', ')}
153            </Text>
154          </Button>
155        ) : null}
156      </Box>
157    )
158  })
159}
160

1export type ContextFill = { tokens: number; window: number; percent: number }
2
3export type BackgroundTask = { kind: string; label: string }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'session-gauges': {
8      context: ContextFill | null
9      tasks: Record<string, BackgroundTask>
10    }
11  }
12}
13