token-weather

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Mood } from '../types'
5
6const history = atom({ plugin: 'token-weather', key: 'history' } as const, [])
7const contextWindow = atom({ plugin: 'token-weather', key: 'window' } as const, 0)
8const mood = atom({ plugin: 'token-weather', key: 'mood' } as const, null)
9const prompts = atom({ plugin: 'token-weather', key: 'prompts' } as const, [])
10const isCompacted = atom({ plugin: 'token-weather', key: 'isCompacted' } as const, false)
11
12const BARS = '▁▂▃▄▅▆▇█'
13const SPAN = 12
14
15// U+FE0F asks for the color emoji form. Without it a terminal draws some of these from the text font
16// and borrows the rest from the emoji font, so they come out at different sizes and widths.
17const EMOJI = '\u{FE0F}'
18
19const FORECASTS = [
20  { below: 25, icon: `☀${EMOJI}`, word: 'Clear', color: '#FFA500' },
21  { below: 50, icon: `☁${EMOJI}`, word: 'Cloudy', color: 'cyan' },
22  { below: 75, icon: `☂${EMOJI}`, word: 'Showers', color: 'blue' },
23  { below: 90, icon: '☇', word: 'Storm', color: 'magenta' },
24  { below: Infinity, icon: '↯', word: 'Compact soon', color: 'red' },
25]
26
27const MOODS: Record<Mood, { icon: string; word: string; color: string }> = {
28  upbeat: { icon: `😊${EMOJI}`, word: 'Upbeat', color: 'green' },
29  curious: { icon: `🧐${EMOJI}`, word: 'Curious', color: 'cyan' },
30  neutral: { icon: `😐${EMOJI}`, word: 'Neutral', color: 'gray' },
31  confused: { icon: `🤔${EMOJI}`, word: 'Confused', color: '#FFA500' },
32  frustrated: { icon: `😤${EMOJI}`, word: 'Frustrated', color: 'red' },
33}
34
35const RECENT_PROMPTS = 6
36
37const MOOD_SYSTEM = `You read the latest exchanges between a developer and their coding assistant and judge how the developer feels about how the collaboration is going.
38Answer with exactly one word: upbeat, neutral, curious, confused or frustrated.
39- upbeat: pleased, thanking, excited, things are landing.
40- neutral: plainly directing work, no strong signal.
41- curious: exploring ideas, asking open questions, wanting to learn.
42- confused: unsure what happened, asking for clarification, contradicting themselves.
43- frustrated: repeating or correcting the same thing, terse, annoyed, the assistant missed the point.`
44
45function short(n: number) {
46  const [value, unit] = n >= 1e6 ? [n / 1e6, 'M'] : [n / 1e3, 'k']
47  return `${value.toFixed(1).replace(/\.0$/, '')}${unit}`
48}
49
50function sparkline(samples: number[]) {
51  const top = Math.max(...samples)
52  return samples
53    .map(n => BARS[top === 0 ? 0 : Math.min(BARS.length - 1, Math.floor((n / top) * BARS.length))])
54    .join('')
55}
56
57async function sample($: EngineInterface) {
58  const { context } = await $.session.usage()
59  if (context.tokens === undefined) return
60  await update($, contextWindow, () => context.window)
61  await update($, history, h => [...h, context.tokens!].slice(-SPAN))
62  await update($, isCompacted, () => false)
63}
64
65async function readMood($: EngineInterface, prompts: string[], answer: string) {
66  const transcript = prompts.map(p => `Developer: ${p.slice(0, 600)}`).join('\n\n')
67  const r = await $.model.complete({
68    model: 'haiku',
69    system: MOOD_SYSTEM,
70    prompt: `${transcript}\n\nAssistant's last reply: ${answer.slice(0, 800)}`,
71    maxTokens: 10,
72    effort: 'low',
73    timeoutMs: 8_000,
74  })
75  if (!r.isAnswered) return
76  const word = r.text
77    .trim()
78    .toLowerCase()
79    .replace(/[^a-z]/g, '')
80  if (word in MOODS) await update($, mood, () => word as Mood)
81}
82
83export const register: Register = on => {
84  on('prompt.submit', async ($, e, next) => {
85    await update($, prompts, p => [...p, e.text].slice(-RECENT_PROMPTS))
86    return next(e)
87  })
88
89  on('session.start', async ($, e, next) => {
90    const result = await next(e)
91    if ((await read($, history)).length === 0) await sample($)
92    return result
93  })
94
95  on('turn.complete', async ($, e, next) => {
96    const result = await next(e)
97    if (e.agentId !== undefined) return result
98    await sample($)
99    const recent = await read($, prompts)
100    if (recent.length > 0) await readMood($, recent, e.answer)
101    return result
102  })
103
104  // Right after a compaction the engine has no real context size until the next API response.
105  on('session.compact', async ($, e, next) => {
106    const result = await next(e)
107    if (e.agentId === undefined && e.trigger !== 'precompute' && result.messages) {
108      await update($, isCompacted, () => true)
109    }
110    return result
111  })
112
113  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
114    const samples = await read($, history)
115    const window = await read($, contextWindow)
116    const reading = await read($, mood)
117    const compacted = await read($, isCompacted)
118    if (e.props.hasSurvey || samples.length === 0 || window === 0) return next(e)
119
120    const tokens = samples[samples.length - 1]
121    const percent = Math.round((tokens / window) * 100)
122    const forecast = FORECASTS.find(f => percent < f.below)!
123    const delta = samples.length > 1 ? tokens - samples[samples.length - 2] : undefined
124    const { Box, Text } = $.ui.resolve(e)
125
126    const feeling = reading && MOODS[reading]
127
128    return (
129      <Box>
130        {compacted ? (
131          <Box>
132            <Text color="green" bold>
133              ↺ Compacted
134            </Text>
135            <Text dimColor>{'  new size after your next message  '}</Text>
136          </Box>
137        ) : (
138          <Box>
139            <Text color={forecast.color} bold>
140              {forecast.icon} {forecast.word}
141            </Text>
142            <Text>
143              {'  '}
144              {percent}% · {short(tokens)} / {short(window)}
145              {'  '}
146            </Text>
147          </Box>
148        )}
149        <Text color={compacted ? 'gray' : forecast.color}>{sparkline(samples)}</Text>
150        {!compacted && delta !== undefined && (
151          <Text dimColor>
152            {'  '}
153            {delta >= 0 ? '▲ +' : '▼ -'}
154            {short(Math.abs(delta))} last turn
155          </Text>
156        )}
157        {feeling && (
158          <Box>
159            <Text dimColor>{'  │  '}</Text>
160            <Text color={feeling.color} bold>
161              {feeling.icon} {feeling.word}
162            </Text>
163          </Box>
164        )}
165      </Box>
166    )
167  })
168}
169

1export type Mood = 'upbeat' | 'neutral' | 'curious' | 'confused' | 'frustrated'
2
3declare module 'claude-code' {
4  interface PluginState {
5    'token-weather': {
6      history: number[]
7      window: number
8      mood: Mood | null
9      prompts: string[]
10      isCompacted: boolean
11    }
12  }
13}
14