context-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Limit, Segment, Snapshot } from '../types'
5
6const snapshot = atom({ plugin: 'context-bar', key: 'snapshot' } as const, null)
7const isShown = atom({ plugin: 'context-bar', key: 'isShown' } as const, false)
8
9// /context's row names, mapped to the short labels and colors the bar uses
10const CATEGORIES: { match: RegExp; label: string; color: string }[] = [
11  { match: /^system prompt/i, label: 'System Prompt', color: '#6f8fdc' },
12  { match: /^(system )?tools/i, label: 'Tools', color: '#5bb8b8' },
13  { match: /^mcp/i, label: 'MCP Tools', color: '#9a7cf2' },
14  { match: /agent/i, label: 'Agents', color: '#86c77a' },
15  { match: /memory/i, label: 'Memory Files', color: '#e9c46a' },
16  { match: /skill/i, label: 'Skills', color: '#ee9bbb' },
17  { match: /^messages/i, label: 'Messages', color: '#e07f5a' },
18]
19const OTHER_COLOR = '#8a8f98'
20const FREE_COLOR = '#2c313a'
21
22export function fmt(n: number): string {
23  if (n < 1000) return String(Math.round(n))
24  if (n < 10_000) return (n / 1000).toFixed(1).replace(/\.0$/, '') + 'k'
25  if (n < 1_000_000) return Math.round(n / 1000) + 'k'
26  return (n / 1_000_000).toFixed(1).replace(/\.0$/, '') + 'M'
27}
28
29export function pct(tokens: number, window: number): number {
30  if (tokens <= 0 || window <= 0) return 0
31  return Math.max(1, Math.round((tokens / window) * 100))
32}
33
34// 各単語の先頭を大文字にする（未知のカテゴリ名用）
35export function titleCase(name: string): string {
36  return name.replace(/\b[a-z]/g, c => c.toUpperCase())
37}
38
39// 使用率に応じたバッジの背景色
40export function badgeColor(percent: number): 'error' | 'warning' | 'success' {
41  return percent >= 80 ? 'error' : percent >= 50 ? 'warning' : 'success'
42}
43
44// プラン上限の使用率（five_hour: セッション、seven_day: 週間）。サブスク以外や未計測なら null
45export function limitUsage(
46  rateLimits: readonly { kind: string; percentUsed: number; resetsAt?: string }[],
47  kind: string,
48): Limit | null {
49  const limit = rateLimits.find(r => r.kind === kind)
50  return limit ? { percent: Math.round(limit.percentUsed), resetsAt: limit.resetsAt ?? null } : null
51}
52
53// 保存済みの値は、リセット時刻を過ぎていなければまだ有効
54export function stillValid(limit: Limit | null | undefined, now: number): Limit | null {
55  if (!limit?.resetsAt) return null
56  const at = Date.parse(limit.resetsAt)
57  return Number.isNaN(at) || at <= now ? null : limit
58}
59
60// 最新の値があれば保存して使い、なければ（リロード直後や fork 直後など）前回保存した値で補う
61async function resolveLimits(
62  $: EngineInterface,
63  rateLimits: readonly { kind: string; percentUsed: number; resetsAt?: string }[],
64): Promise<{ usage: Limit | null; weekly: Limit | null }> {
65  const fresh = { usage: limitUsage(rateLimits, 'five_hour'), weekly: limitUsage(rateLimits, 'seven_day') }
66  if (fresh.usage || fresh.weekly) {
67    await $.store.set('limits', fresh)
68    return fresh
69  }
70  const saved = (await $.store.get('limits')) as { usage?: Limit | null; weekly?: Limit | null } | undefined
71  const now = await $.clock.now()
72  return { usage: stillValid(saved?.usage, now), weekly: stillValid(saved?.weekly, now) }
73}
74
75const WEEKDAYS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat']
76
77// リセット時刻をローカル時間で表示する（週間枠は曜日付き）
78export function fmtReset(iso: string | null, withDay: boolean): string {
79  if (!iso) return ''
80  const d = new Date(iso)
81  if (Number.isNaN(d.getTime())) return ''
82  const time = `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}`
83  return withDay ? `${WEEKDAYS[d.getDay()]} ${time}` : time
84}
85
86async function refresh($: EngineInterface): Promise<void> {
87  try {
88    const usage = await $.session.usage({ breakdown: 'summary' })
89    const b = usage.context.breakdown
90    if (!b) return
91    const window = usage.context.window || b.rawMaxTokens
92    const segments: Segment[] = []
93    let free = 0
94    for (const row of b.categories) {
95      if (row.kind === 'free') {
96        free += row.tokens
97        continue
98      }
99      if (row.kind !== 'used') continue
100      const known = CATEGORIES.find(c => c.match.test(row.name))
101      const label = known?.label ?? titleCase(row.name)
102      const existing = segments.find(s => s.label === label)
103      if (existing) existing.tokens += row.tokens
104      else segments.push({ label, tokens: row.tokens, color: known?.color ?? OTHER_COLOR })
105    }
106    // Keep the bar's order stable: the known categories first, in their order
107    segments.sort((a, z) => rank(a.label) - rank(z.label))
108    const used = segments.reduce((sum, s) => sum + s.tokens, 0)
109    const limits = await resolveLimits($, usage.rateLimits)
110    const next: Snapshot = {
111      used,
112      window,
113      compactAt: b.isAutoCompactEnabled ? (b.autoCompactThreshold ?? null) : null,
114      percent: pct(used, window),
115      segments,
116      free,
117      ...limits,
118    }
119    await update($, snapshot, () => next)
120  } catch {
121    // No session bound yet, or the breakdown failed: keep the last snapshot
122  }
123}
124
125function rank(label: string): number {
126  const i = CATEGORIES.findIndex(c => c.label === label)
127  return i === -1 ? CATEGORIES.length : i
128}
129
130export const register: Register = on => {
131  on('session.start', async ($, e, next) => {
132    await $.command.register({
133      name: 'context-bar',
134      description: 'Toggle the context-window bar above the prompt',
135      argumentHint: '[on|off]',
136      immediate: true,
137    })
138    const saved = await $.store.get('isShown')
139    if (saved === true) {
140      await update($, isShown, () => true)
141      $.clock.after(1500, () => void refresh($))
142    }
143    return next(e)
144  })
145
146  on('command.run', { command: 'context-bar' }, async ($, e) => {
147    const arg = (e.args ?? '').trim().toLowerCase()
148    const current = await read($, isShown)
149    const show = arg === 'on' ? true : arg === 'off' ? false : !current
150    await update($, isShown, () => show)
151    await $.store.set('isShown', show)
152    if (show) await refresh($)
153    return { text: show ? 'Context bar shown.' : 'Context bar hidden.' }
154  })
155
156  // Refresh after each main-loop turn, and after /clear, /resume, /compact
157  on('turn.complete', async ($, e, next) => {
158    const result = await next(e)
159    if (!e.agentId) void refresh($)
160    return result
161  })
162  // 使用率が動いたらバッジだけ更新する
163  on('session.measure', async ($, e, next) => {
164    const result = await next(e)
165    if (e.changed.includes('rateLimits')) {
166      const limits = await resolveLimits($, e.rateLimits)
167      await update($, snapshot, s => (s ? { ...s, ...limits } : s))
168    }
169    return result
170  }).catch(($, e, next) => next(e))
171  on('classic.SessionStart', async ($, e, next) => {
172    const result = await next(e)
173    $.clock.after(500, () => void refresh($))
174    return result
175  }).catch(($, e, next) => next(e))
176  on('classic.PostCompact', async ($, e, next) => {
177    const result = await next(e)
178    void refresh($)
179    return result
180  }).catch(($, e, next) => next(e))
181
182  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
183    if (e.props.hasSurvey || !(await read($, isShown))) return next(e)
184    const snap = await read($, snapshot)
185    const { Box, Text } = $.ui.resolve(e)
186    const width = Math.max(20, e.props.bodyColumns - 2)
187    // 後続のフック（進捗バーなど）の描画も残し、自分の枠の下に積む
188    const rest = await next(e)
189
190    if (!snap) {
191      return (
192        <Box flexDirection="column">
193          <Box flexDirection="column" borderStyle="round" borderDimColor paddingX={1}>
194            <Text>
195              <Text color="#e07f5a">◆ </Text>
196              <Text bold>Context</Text>
197              <Text dimColor> measuring… (fills in after the first reply)</Text>
198            </Text>
199          </Box>
200          {rest}
201        </Box>
202      )
203    }
204
205    const badge = badgeColor(snap.percent)
206    const summary =
207      `${fmt(snap.used)} of ${fmt(snap.window)}` +
208      (snap.compactAt ? ` · compacts at ${fmt(snap.compactAt)}` : '')
209
210    const header = (
211      <Box flexDirection="row" justifyContent="space-between">
212        <Text>
213          <Text color="#e07f5a">◆ </Text>
214          <Text bold>Context</Text>
215          {snap.usage && (
216            <Text>
217              <Text dimColor>{'   Usage '}</Text>
218              <Text bold backgroundColor={badgeColor(snap.usage.percent)} color="inverseText">{` ${snap.usage.percent}% `}</Text>
219              {snap.usage.resetsAt && <Text dimColor>{` ↻ ${fmtReset(snap.usage.resetsAt, false)}`}</Text>}
220            </Text>
221          )}
222          {snap.weekly && (
223            <Text>
224              <Text dimColor>{'   Week '}</Text>
225              <Text bold backgroundColor={badgeColor(snap.weekly.percent)} color="inverseText">{` ${snap.weekly.percent}% `}</Text>
226              {snap.weekly.resetsAt && <Text dimColor>{` ↻ ${fmtReset(snap.weekly.resetsAt, true)}`}</Text>}
227            </Text>
228          )}
229        </Text>
230        <Text>
231          <Text bold>{fmt(snap.used)}</Text>
232          <Text dimColor>{summary.slice(fmt(snap.used).length)}{'  '}</Text>
233          <Text bold backgroundColor={badge} color="inverseText">{` ${snap.percent}% `}</Text>
234          <Text> </Text>
235        </Text>
236      </Box>
237    )
238
239    const legend = (
240      <Box flexDirection="row" flexWrap="wrap" columnGap={2}>
241        {snap.segments.map(s => (
242          <Text>
243            <Text color={s.color}>■ </Text>
244            <Text>{s.label} </Text>
245            <Text bold>{fmt(s.tokens)}</Text>
246            <Text dimColor> {pct(s.tokens, snap.window)}%</Text>
247          </Text>
248        ))}
249        <Text>
250          <Text dimColor>■ Free </Text>
251          <Text>{fmt(snap.free)}</Text>
252        </Text>
253      </Box>
254    )
255
256    const bar = e.surface === 'terminal' ? terminalBar($, e, snap, width) : svgBar($, e, snap)
257
258    return (
259      <Box flexDirection="column">
260        <Box flexDirection="column" borderStyle="round" borderDimColor paddingX={1}>
261          {header}
262          {bar}
263          {legend}
264        </Box>
265        {rest}
266      </Box>
267    )
268  })
269}
270
271// A row of block characters, one cell per slice of the window
272function terminalBar($: EngineInterface, e: Parameters<EngineInterface['ui']['resolve']>[0], snap: Snapshot, width: number) {
273  const { Text } = $.ui.resolve(e)
274  const cells = Math.max(10, width - 4)
275  const marker = snap.compactAt ? Math.min(cells - 1, Math.floor((snap.compactAt / snap.window) * cells)) : -1
276  const parts: { text: string; color?: string; dim?: boolean }[] = []
277  let used = 0
278  for (const s of snap.segments) {
279    if (s.tokens <= 0) continue
280    const n = Math.max(1, Math.round((s.tokens / snap.window) * cells))
281    parts.push({ text: '█'.repeat(Math.min(n, cells - used)), color: s.color })
282    used += n
283    if (used >= cells) break
284  }
285  const rest = Math.max(0, cells - used)
286  if (rest > 0) {
287    const before = marker > used ? marker - used : rest
288    parts.push({ text: '░'.repeat(Math.min(before, rest)), dim: true })
289    if (marker >= used && marker < cells) {
290      parts.push({ text: '│', color: 'warning' })
291      parts.push({ text: '░'.repeat(Math.max(0, rest - before - 1)), dim: true })
292    }
293  }
294  return (
295    <Text wrap="truncate-end">
296      {parts.map(p => (
297        <Text color={p.color} dimColor={p.dim}>{p.text}</Text>
298      ))}
299    </Text>
300  )
301}
302
303// A rounded SVG bar for the Desktop app
304function svgBar($: EngineInterface, e: Parameters<EngineInterface['ui']['resolve']>[0], snap: Snapshot) {
305  if (e.surface !== 'desktop' && e.surface !== 'mobile') return terminalBar($, e, snap, 60)
306  const { Svg } = $.ui.resolve(e)
307  const W = 1000
308  const H = 22
309  let x = 0
310  const rects: string[] = []
311  for (const s of snap.segments) {
312    if (s.tokens <= 0) continue
313    const w = Math.max(3, (s.tokens / snap.window) * W)
314    rects.push(`<rect x="${x.toFixed(1)}" y="0" width="${w.toFixed(1)}" height="${H}" fill="${s.color}"/>`)
315    x += w
316  }
317  const marker = snap.compactAt
318    ? `<rect x="${((snap.compactAt / snap.window) * W - 2).toFixed(1)}" y="0" width="3" height="${H}" fill="#e9c46a"/>`
319    : ''
320  const source =
321    `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 ${W} ${H}" preserveAspectRatio="none">` +
322    `<defs><clipPath id="r"><rect width="${W}" height="${H}" rx="4"/></clipPath></defs>` +
323    `<g clip-path="url(#r)"><rect width="${W}" height="${H}" fill="${FREE_COLOR}"/>${rects.join('')}${marker}</g></svg>`
324  return (
325    <Svg
326      source={source}
327      alt={`Context ${snap.percent}% used: ${snap.segments.map(s => `${s.label} ${fmt(s.tokens)}`).join(', ')}`}
328      height={14}
329    />
330  )
331}
332

1export type Segment = { label: string; tokens: number; color: string }
2
3/** One plan limit window: percent used (0–100) and when it resets (ISO 8601) */
4export type Limit = { percent: number; resetsAt: string | null }
5
6export type Snapshot = {
7  used: number
8  window: number
9  compactAt: number | null
10  percent: number
11  segments: Segment[]
12  free: number
13  /** Plan usage of the five-hour session window, 0–100; null off a subscription or before the first reading */
14  usage: Limit | null
15  /** Plan usage of the seven-day window; null when not reported */
16  weekly: Limit | null
17}
18
19declare module 'claude-code' {
20  interface PluginState {
21    'context-bar': {
22      snapshot: Snapshot | null
23      isShown: boolean
24    }
25  }
26}
27