token-weather

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionUsage } from 'claude-code'
3
4import type { Info, Reading } from '../types'
5import {
6  bar,
7  branchFromHead,
8  duration,
9  forecast,
10  limitLabel,
11  loadColor,
12  modelLabel,
13  shortTokens,
14  signedTokens,
15  sparkline,
16} from './forecast'
17
18const HISTORY = 12
19const TICK_MS = 30_000
20
21const readings = atom({ plugin: 'token-weather', key: 'readings' } as const, [])
22const isHidden = atom({ plugin: 'token-weather', key: 'isHidden' } as const, false)
23const info = atom({ plugin: 'token-weather', key: 'info' } as const, null)
24const now = atom({ plugin: 'token-weather', key: 'now' } as const, 0)
25
26async function readBranch($: EngineInterface, root: string) {
27  const head = await $.fs.read(`${root}/.git/HEAD`).catch(() => '')
28  return branchFromHead(head)
29}
30
31async function refresh($: EngineInterface, usage: SessionUsage) {
32  const repo = await $.session.repo()
33  const root = repo?.root ?? (await $.session.root())
34  const next: Info = {
35    model: await $.session.model(),
36    project: root.split('/').filter(Boolean).pop() ?? root,
37    branch: repo === null ? null : await readBranch($, repo.root),
38    startedAt: usage.startedAt,
39    limits: usage.rateLimits.map(l => ({ kind: l.kind, percent: l.percentUsed, resetsAt: l.resetsAt })),
40    costUsd: usage.cost?.usd,
41  }
42  const t = await $.clock.now()
43  await update($, info, () => next)
44  await update($, now, () => t)
45}
46
47async function record($: EngineInterface, usage: SessionUsage) {
48  const { tokens, percent, window } = usage.context
49  if (tokens === undefined || percent === undefined) return
50
51  const reading: Reading = { percent, tokens, window }
52  await update($, readings, list => {
53    const last = list[list.length - 1]
54    const isSame = last !== undefined && last.tokens === tokens
55    return isSame ? list : [...list, reading].slice(-HISTORY)
56  })
57}
58
59export const register: Register = on => {
60  on('session.start', async ($, e, next) => {
61    await $.command.register({
62      name: 'token-weather',
63      description: '切換 context 天氣預報（on / off / status）',
64    })
65    await refresh($, await $.session.usage())
66    $.clock.every(TICK_MS, async () => {
67      const t = await $.clock.now()
68      await update($, now, () => t)
69    })
70
71    return next(e)
72  })
73
74  on('session.measure', async ($, e, next) => {
75    const usage = await $.session.usage()
76    if (e.changed.includes('context')) await record($, usage)
77    await refresh($, usage)
78
79    return next(e)
80  })
81
82  on('turn.complete', async ($, e, next) => {
83    const usage = await $.session.usage()
84    await record($, usage)
85    await refresh($, usage)
86
87    return next(e)
88  })
89
90  on('command.run', { command: 'token-weather' }, async ($, e) => {
91    const arg = e.args.trim().toLowerCase()
92
93    if (arg === 'status') {
94      const list = await read($, readings)
95      const last = list[list.length - 1]
96      const shown = (await read($, isHidden)) ? '隱藏中' : '顯示中'
97
98      if (last === undefined) {
99        return { text: `Token Weather ${shown}，尚無資料（完成一輪對話後才會有）。` }
100      }
101
102      const w = forecast(last.percent)
103
104      return {
105        text: `Token Weather ${shown}：${w.icon} ${w.label} ${last.percent}%（${shortTokens(last.tokens)}/${shortTokens(last.window)}）`,
106      }
107    }
108
109    const hide = arg === 'off' ? true : arg === 'on' ? false : !(await read($, isHidden))
110    await update($, isHidden, () => hide)
111
112    return { text: hide ? 'Token Weather 已關閉。' : 'Token Weather 已開啟。' }
113  })
114
115  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
116    if (e.props.hasSurvey || (await read($, isHidden))) {
117      return next(e)
118    }
119
120    const list = await read($, readings)
121    const meta = await read($, info)
122    const t = await read($, now)
123    const last = list[list.length - 1]
124    // 還沒有任何資料（例如剛重開）時也畫等待提示，不然看起來像 mod 沒載入
125    const { Box, Text } = $.ui.resolve(e)
126    const prev = list[list.length - 2]
127    const delta = last !== undefined && prev !== undefined ? last.tokens - prev.tokens : undefined
128    const w = last === undefined ? undefined : forecast(last.percent)
129
130    return (
131      <Box flexDirection="column">
132        <Box>
133          {last === undefined || w === undefined ? (
134            <Text dimColor>☁ 等待第一個回應…</Text>
135          ) : (
136            <Box>
137              <Text color={w.color} bold>
138                {w.icon} {w.label} {last.percent}%
139              </Text>
140              <Text dimColor>
141                {'  '}
142                {shortTokens(last.tokens)}/{shortTokens(last.window)}
143                {'  '}
144              </Text>
145              <Text color={w.color}>{sparkline(list.map(r => r.percent))}</Text>
146              {delta !== undefined && <Text dimColor>{'  '}上一輪 {signedTokens(delta)}</Text>}
147            </Box>
148          )}
149        </Box>
150        {meta !== null && (
151          <Box>
152            <Text wrap="truncate-end">
153              <Text color="cyan">[{modelLabel(meta.model)}]</Text>
154              <Text dimColor> │ </Text>
155              <Text>{meta.project}</Text>
156              {meta.branch !== null && <Text color="magenta"> ({meta.branch})</Text>}
157              {meta.limits.map(l => (
158                <Text key={l.kind}>
159                  <Text dimColor> │ {limitLabel(l.kind)} </Text>
160                  <Text color={loadColor(l.percent)}>
161                    {bar(l.percent, 8)} {l.percent}%
162                  </Text>
163                  {l.resetsAt !== undefined && (
164                    <Text dimColor> ({duration(Date.parse(l.resetsAt) - t)}後重置)</Text>
165                  )}
166                </Text>
167              ))}
168              <Text dimColor> │ ⏱ {duration(t - meta.startedAt)}</Text>
169              {meta.costUsd !== undefined && <Text dimColor> │ ${meta.costUsd.toFixed(2)}</Text>}
170            </Text>
171          </Box>
172        )}
173      </Box>
174    )
175  })
176}
177

1export type Weather = { icon: string; label: string; color: string }
2
3const BARS = '▁▂▃▄▅▆▇█'
4
5export const forecast = (percent: number): Weather => {
6  if (percent >= 90) return { icon: '↯', label: '即將壓縮', color: 'red' }
7  if (percent >= 75) return { icon: '☇', label: '暴風雨', color: 'magenta' }
8  if (percent >= 50) return { icon: '☂', label: '陣雨', color: 'yellow' }
9  if (percent >= 25) return { icon: '☁', label: '多雲', color: 'cyan' }
10  return { icon: '☀', label: '晴', color: 'green' }
11}
12
13export const sparkline = (percents: readonly number[]): string =>
14  percents
15    .map(p => BARS[Math.min(BARS.length - 1, Math.floor((Math.max(0, p) / 100) * BARS.length))])
16    .join('')
17
18export const shortTokens = (n: number): string => {
19  if (n >= 1_000_000) return `${+(n / 1_000_000).toFixed(1)}M`
20  if (n >= 1_000) return `${+(n / 1_000).toFixed(1)}k`
21  return `${n}`
22}
23
24export const signedTokens = (n: number): string =>
25  n >= 0 ? `+${shortTokens(n)}` : `-${shortTokens(-n)}`
26
27export const bar = (percent: number, width = 10): string => {
28  const filled = Math.round((Math.min(100, Math.max(0, percent)) / 100) * width)
29  return '█'.repeat(filled) + '░'.repeat(width - filled)
30}
31
32export const duration = (ms: number): string => {
33  const m = Math.max(0, Math.floor(ms / 60_000))
34  if (m < 60) return `${m}m`
35  const h = Math.floor(m / 60)
36  if (h < 24) return `${h}h ${m % 60}m`
37  return `${Math.floor(h / 24)}d ${h % 24}h`
38}
39
40/** 依用量決定顏色：低於 50% 綠、低於 80% 黃，其餘紅。 */
41export const loadColor = (percent: number): string =>
42  percent >= 80 ? 'red' : percent >= 50 ? 'yellow' : 'green'
43
44export const limitLabel = (kind: string): string =>
45  kind === 'five_hour' ? '5h' : kind === 'seven_day' ? '7d' : kind
46
47/** 從 .git/HEAD 的內容取出分支名稱；detached HEAD 回前 7 碼。 */
48export const branchFromHead = (head: string): string | null => {
49  const text = head.trim()
50  const ref = /^ref: refs\/heads\/(.+)$/.exec(text)
51  if (ref) return ref[1] ?? null
52  return /^[0-9a-f]{7,}$/.test(text) ? text.slice(0, 7) : null
53}
54
55/** 把模型 ID 轉成好讀的名稱：claude-opus-5-5 → Opus 5.5；認不得就原樣顯示。 */
56export const modelLabel = (id: string): string => {
57  const m = /^claude-([a-z]+)-(\d+)-(\d+)(?!\d)/.exec(id)
58  if (m === null) return id
59  const [, family = '', major, minor] = m
60  const name = family.charAt(0).toUpperCase() + family.slice(1)
61  return `${name} ${major}.${minor}${id.includes('[1m]') ? ' (1M)' : ''}`
62}
63

1export type Reading = { percent: number; tokens: number; window: number }
2
3export type Limit = { kind: string; percent: number; resetsAt?: string }
4
5export type Info = {
6  model: string
7  project: string
8  branch: string | null
9  startedAt: number
10  limits: Limit[]
11  costUsd?: number
12}
13
14declare module 'claude-code' {
15  interface PluginState {
16    'token-weather': {
17      readings: Reading[]
18      isHidden: boolean
19      info: Info | null
20      now: number
21    }
22  }
23}
24