quota-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { QuotaStats, QuotaWindow } from '../types'
5
6const stats = atom({ plugin: 'quota-band', key: 'stats' } as const, {
7  turn: 0,
8  speed: null,
9  turnCost: null,
10  total: 0,
11  costAtStart: 0,
12  context: 0,
13} as QuotaStats)
14const windows = atom({ plugin: 'quota-band', key: 'windows' } as const, [])
15const cacheExpiresAt = atom({ plugin: 'quota-band', key: 'cacheExpiresAt' } as const, null)
16const now = atom({ plugin: 'quota-band', key: 'now' } as const, 0)
17// The auto-continue mod's state, read only: absent when it is not installed.
18const autoOn = atom({ plugin: 'auto-continue', key: 'isEnabled' } as const, null)
19const autoWait = atom({ plugin: 'auto-continue', key: 'wait' } as const, null)
20const autoScheduled = atom({ plugin: 'auto-continue', key: 'scheduled' } as const, null)
21
22const LABELS: Record<string, string> = { five_hour: '5h', seven_day: '周', spend_limit: '花费' }
23const BAR = 10
24const DIVIDER = 'cyan'
25const SPEED = '#8b8cf8'
26const MONEY = '#4caf6e'
27const CACHE = '#4fb3a9'
28const WARN = '#e2a94e'
29const ALERT = '#e0846b'
30
31// Output speed: ignore replies too short or too quick to time fairly.
32const MIN_TOKENS = 20
33const MIN_SECONDS = 0.5
34const MAX_SPEED = 1000
35
36// The prompt cache lives an hour after the last answer.
37const CACHE_TTL_MS = 60 * 60 * 1000
38const CACHE_WARN_S = 20 * 60
39const CACHE_ALERT_S = 10 * 60
40
41// Narrower than these, the row drops its bars, then its labels and total cost.
42const FULL_COLUMNS = 110
43const MID_COLUMNS = 80
44
45const level = (pct: number) => (pct >= 80 ? ALERT : pct >= 50 ? WARN : MONEY)
46
47const bar = (pct: number) => {
48  const full = Math.max(0, Math.min(BAR, Math.round((pct / 100) * BAR)))
49  return '█'.repeat(full) + '░'.repeat(BAR - full)
50}
51
52const pad = (n: number) => String(n).padStart(2, '0')
53// `14:30` today, `10/22 14:30` on another day.
54const clockTime = (ms: number, at: number) => {
55  const d = new Date(ms)
56  const hm = `${pad(d.getHours())}:${pad(d.getMinutes())}`
57  return d.toDateString() === new Date(at).toDateString() ? hm : `${d.getMonth() + 1}/${d.getDate()} ${hm}`
58}
59const usd = (n: number) => `$${n.toFixed(2)}`
60
61const resetIn = (resetsAt: string | undefined, at: number) => {
62  if (!resetsAt) return ''
63  const min = Math.max(0, Math.round((Date.parse(resetsAt) - at) / 60000))
64  if (Number.isNaN(min)) return ''
65  if (min < 60) return `${min}分钟后重置`
66  if (min < 24 * 60) return `${Math.floor(min / 60)}h${min % 60}m后重置`
67  return `${Math.round(min / 1440)}天后重置`
68}
69
70async function refresh($: EngineInterface) {
71  const { rateLimits, cost, context } = await $.session.usage()
72  const list: QuotaWindow[] = rateLimits.map(r => ({ kind: r.kind, percentUsed: r.percentUsed, resetsAt: r.resetsAt }))
73  await update($, windows, () => list)
74  const at = await $.clock.now()
75  await update($, now, () => at)
76  const total = cost?.usd ?? 0
77  const turn = await $.session.turns()
78  await update($, stats, s => ({
79    ...s,
80    turn,
81    total,
82    turnCost: s.turnCost === null ? null : total - s.costAtStart,
83    context: context.percent ?? s.context,
84  }))
85}
86
87export const register: Register = on => {
88  on('session.start', async ($, e, next) => {
89    const started = await next(e)
90    await refresh($)
91    $.clock.every(1000, async () => {
92      const at = await $.clock.now()
93      await update($, now, () => at)
94    })
95    $.clock.every(60_000, () => refresh($))
96    return started
97  })
98
99  on('turn.start', async ($, e, next) => {
100    if (e.agentId === undefined) {
101      const total = (await $.session.usage()).cost?.usd ?? 0
102      await update($, stats, s => ({ ...s, costAtStart: total, turnCost: 0 }))
103    }
104    return next(e)
105  })
106
107  on('turn.step', async function* ($, e, next) {
108    if (e.agentId !== undefined) return yield* next(e)
109
110    // Timed from the response's first chunk, so time-to-first-token is left out.
111    let startedAt: number | null = null
112    for await (const chunk of next(e)) {
113      if (startedAt === null) startedAt = await $.clock.now()
114      if (chunk.kind === 'stop' && chunk.usage !== null) {
115        const tokens = chunk.usage.output_tokens
116        const seconds = ((await $.clock.now()) - startedAt) / 1000
117        const tps = tokens / seconds
118        if (tokens >= MIN_TOKENS && seconds >= MIN_SECONDS && tps <= MAX_SPEED) {
119          await update($, stats, s => ({ ...s, speed: Math.round(tps) }))
120        }
121      }
122      yield chunk
123    }
124    await refresh($)
125  })
126
127  on('turn.complete', async ($, e, next) => {
128    const done = await next(e)
129    await refresh($)
130    if (e.agentId === undefined) {
131      const at = await $.clock.now()
132      await update($, cacheExpiresAt, () => at + CACHE_TTL_MS)
133    }
134    return done
135  })
136
137  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
138    // The desktop draws desktop-cockpit's band instead.
139    if (e.props.hasSurvey || e.surface === 'desktop') return next(e)
140    const s = await read($, stats)
141    const list = await read($, windows)
142    const expiresAt = await read($, cacheExpiresAt)
143    const at = await read($, now)
144    const { Box, Text } = $.ui.resolve(e)
145    const cols = e.props.bodyColumns
146    const isFull = cols >= FULL_COLUMNS
147    const isMid = cols >= MID_COLUMNS
148
149    // A bold line marking where the transcript ends and the tool area begins.
150    const divider = (
151      <Text key="divider" color={DIVIDER} wrap="truncate-end">
152        {'━'.repeat(Math.max(1, cols))}
153      </Text>
154    )
155
156    const left = expiresAt === null || at === 0 ? null : Math.max(0, Math.floor((expiresAt - at) / 1000))
157    const cacheColor =
158      e.props.isWorking || left === null ? CACHE : left <= CACHE_ALERT_S ? ALERT : left <= CACHE_WARN_S ? WARN : CACHE
159    const cacheText =
160      e.props.isWorking ? '…' : left === null ? '—' : left === 0 ? '已过期' : `${pad(Math.floor(left / 60))}:${pad(left % 60)}`
161
162    // auto-continue beside the quota: waiting for a reset, a set time, or just on.
163    const wait = await read($, autoWait)
164    const scheduled = await read($, autoScheduled)
165    const isAutoOn = (await read($, autoOn)) === true
166    const auto = wait ? (
167      <Box key="auto" flexDirection="row" columnGap={1} flexShrink={0}>
168        <Text color={WARN}>⏳</Text>
169        <Text bold color={WARN}>{`${clockTime(wait.wakeAt, at)}${isMid ? ' 自动续' : ''}`}</Text>
170      </Box>
171    ) : scheduled ? (
172      <Box key="auto" flexDirection="row" columnGap={1} flexShrink={0}>
173        <Text color={SPEED}>⏰</Text>
174        <Text bold color={SPEED}>{`${clockTime(scheduled.at, at)}${isMid ? ' 定时续' : ''}`}</Text>
175      </Box>
176    ) : isAutoOn ? (
177      <Box key="auto" flexDirection="row" flexShrink={0}>
178        <Text dimColor>{isMid ? '⟳ 自动续' : '⟳'}</Text>
179      </Box>
180    ) : null
181
182    const row = (
183      <Box key="row" flexDirection="row" columnGap={isMid ? 3 : 2} paddingX={1} flexWrap="nowrap" overflow="hidden">
184        <Box flexDirection="row" columnGap={1} flexShrink={0}>
185          <Text dimColor>{isMid ? 'Turn' : '#'}</Text>
186          <Text bold>{String(s.turn)}</Text>
187        </Box>
188        <Box flexDirection="row" columnGap={1} flexShrink={0}>
189          <Text color={SPEED}>⚡</Text>
190          <Text bold>{s.speed === null ? '—' : isMid ? `${s.speed} tok/s` : `${s.speed}/s`}</Text>
191        </Box>
192        <Box flexDirection="row" columnGap={1} flexShrink={0}>
193          <Text bold color={MONEY}>{s.turnCost === null ? '—' : usd(s.turnCost)}</Text>
194          {isMid ? <Text dimColor>|</Text> : null}
195          {isMid ? <Text bold>{usd(s.total)}</Text> : null}
196        </Box>
197        <Box flexDirection="row" columnGap={1} flexShrink={0}>
198          <Text color={cacheColor}>⧗</Text>
199          <Text bold color={cacheColor}>{cacheText}</Text>
200        </Box>
201        {list.map(w => (
202          <Box key={w.kind} flexDirection="row" columnGap={1} flexShrink={0}>
203            <Text dimColor>{LABELS[w.kind] ?? w.kind}</Text>
204            {isFull ? <Text color={level(w.percentUsed)}>{bar(w.percentUsed)}</Text> : null}
205            <Text bold color={level(w.percentUsed)}>{`${Math.round(w.percentUsed)}%`}</Text>
206            {isFull ? <Text dimColor>{resetIn(w.resetsAt, at)}</Text> : null}
207          </Box>
208        ))}
209        {auto}
210        <Box flexGrow={1} flexShrink={1} />
211        <Box flexDirection="row" columnGap={1} flexShrink={0}>
212          <Text dimColor>上下文</Text>
213          {isMid ? <Text color={level(s.context)}>{bar(s.context)}</Text> : null}
214          <Text bold color={level(s.context)}>{`${Math.round(s.context)}%`}</Text>
215        </Box>
216      </Box>
217    )
218
219    // Keep whatever the plugins beneath draw in the band under ours.
220    const below = await next(e)
221    return (
222      <Box flexDirection="column">
223        {divider}
224        {row}
225        {below}
226      </Box>
227    )
228  })
229}
230

1export type QuotaWindow = { kind: string; percentUsed: number; resetsAt?: string }
2
3export type QuotaStats = {
4  turn: number
5  speed: number | null
6  turnCost: number | null
7  total: number
8  costAtStart: number
9  context: number
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    // what the band reads of the auto-continue mod, when it is installed
15    'auto-continue': {
16      isEnabled: boolean | null
17      wait: { wakeAt: number; kind: string | null } | null
18      scheduled: { at: number; text: string } | null
19    }
20    'quota-band': {
21      stats: QuotaStats
22      windows: QuotaWindow[]
23      cacheExpiresAt: number | null
24      now: number
25    }
26  }
27}
28