jr-quota

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PluginOptions, Register } from 'claude-code'
3
4import type { JrQuotaBalance } from '../types'
5import { gauge, gradientText } from './jr/draw'
6import { GRADIENTS } from './jr/gradient'
7import { failureText, PROVIDERS } from './providers'
8
9const view = atom({ plugin: 'jr-quota', key: 'view' } as const, null)
10const bandHidden = atom({ plugin: 'jr-quota', key: 'isBandHidden' } as const, false)
11const band = { isLoaded: false }
12
13/** `hide`/`off`, `show`/`on`, or a bare command, which toggles; anything else is not about the band. */
14const bandVerb = (args: string) => {
15  const arg = args.trim().toLowerCase()
16  if (arg === 'hide' || arg === 'off') return 'hide' as const
17  if (arg === 'show' || arg === 'on') return 'show' as const
18  return arg === '' ? ('toggle' as const) : null
19}
20
21/** Hidden or shown is one choice for every session, kept across restarts. */
22async function isBandHidden($: EngineInterface) {
23  if (!band.isLoaded) {
24    band.isLoaded = true
25    if ((await $.store.get('isBandHidden')) === true) await update($, bandHidden, () => true)
26  }
27  return read($, bandHidden)
28}
29
30async function setBand($: EngineInterface, verb: 'hide' | 'show' | 'toggle') {
31  const hide = verb === 'toggle' ? !(await isBandHidden($)) : verb === 'hide'
32  await $.store.set('isBandHidden', hide)
33  await update($, bandHidden, () => hide)
34  return { text: hide ? 'jr quota · hidden (/jrquota to bring it back)' : 'jr quota · shown' }
35}
36const PANE = 'jr-quota'
37const EVERY_MS = 10 * 60_000
38
39const live: { isPolling: boolean; options: PluginOptions } = { isPolling: false, options: {} }
40
41/** The key in the environment, read by literal name so the mod's reads can be listed. */
42async function envKey($: EngineInterface, id: string) {
43  if (id === 'openrouter') return $.env.get('OPENROUTER_API_KEY')
44  if (id === 'firecrawl') return $.env.get('FIRECRAWL_API_KEY')
45  if (id === 'deepseek') return $.env.get('DEEPSEEK_API_KEY')
46  return undefined
47}
48
49async function check($: EngineInterface) {
50  const low = { usd: Number(live.options.lowUsd ?? 5), credits: Number(live.options.lowCredits ?? 500) }
51  const balances: JrQuotaBalance[] = []
52  for (const p of PROVIDERS) {
53    const fromOption = String(live.options[p.option] ?? '').trim()
54    const key = fromOption || (await envKey($, p.id)) || ''
55    if (!key) continue
56    try {
57      const res = await $.http.fetch(p.url, { headers: { Authorization: `Bearer ${key}`, Accept: 'application/json' } })
58      balances.push(
59        res.ok
60          ? { id: p.id, label: p.label, ...p.parse(JSON.parse(res.text), low) }
61          : { id: p.id, label: p.label, text: failureText(res.status), isLow: false, ratio: null, error: true },
62      )
63    } catch {
64      balances.push({ id: p.id, label: p.label, text: 'unreachable', isLow: false, ratio: null, error: true })
65    }
66  }
67  const checkedAt = await $.clock.now()
68  const before = await read($, view)
69  await update($, view, () => ({ balances, checkedAt }))
70  for (const b of balances) {
71    const was = before?.balances.find(x => x.id === b.id)
72    if (b.isLow && !was?.isLow) $.ui.toast(`jr quota · ${b.label} is low: ${b.text}`)
73  }
74}
75
76async function ensurePolling($: EngineInterface) {
77  if (live.isPolling) return
78  live.isPolling = true
79  await check($)
80  $.clock.every(EVERY_MS, () => {
81    void check($)
82  })
83}
84
85export const register: Register = (on, options) => {
86  live.options = options
87
88  on('session.start', async ($, e, next) => {
89    const started = await next(e)
90    await $.command.register({ name: 'jrquota', description: 'Show or hide the low-balance band; /jrquota details shows every balance' })
91    await ensurePolling($)
92    return started
93  })
94
95  on('command.run', { command: 'jrquota' }, async ($, e) => {
96    const verb = bandVerb(e.args)
97    if (verb) return setBand($, verb)
98    live.isPolling ? await check($) : await ensurePolling($)
99    const v = await read($, view)
100    if (!v || v.balances.length === 0) {
101      return { text: 'jr quota · no API key found. Set one in /config (jr-quota) or export OPENROUTER_API_KEY, FIRECRAWL_API_KEY or DEEPSEEK_API_KEY.' }
102    }
103    await $.ui.open({ id: PANE, title: 'jr quota', rows: v.balances.length + 3 })
104    return { text: 'jr quota opened.' }
105  })
106
107  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
108    const { Box, Text } = $.ui.resolve(e)
109    const v = await read($, view)
110    if (!v || v.balances.length === 0) return <Text dimColor>No API key configured.</Text>
111    return (
112      <Box flexDirection="column">
113        {gradientText(Text, 'jr quota', GRADIENTS.brand)}
114        {v.balances.map(b => (
115          <Box>
116            <Text color={b.error ? '#6b7089' : b.isLow ? '#ef4565' : '#2cb67d'}>{'● '}</Text>
117            <Text bold>{b.label.padEnd(12)}</Text>
118            {b.ratio !== null ? gauge(Text, b.ratio, 16, GRADIENTS.money) : <Text>{' '.repeat(16)}</Text>}
119            <Text color={b.isLow ? '#ef4565' : undefined} dimColor={b.error}>{`  ${b.text}`}</Text>
120          </Box>
121        ))}
122        <Text dimColor>{`checked ${new Date(v.checkedAt).toLocaleTimeString('en-GB', { hour: '2-digit', minute: '2-digit' })} · every 10 min`}</Text>
123      </Box>
124    )
125  })
126
127  // The band shows only what needs attention: a low balance.
128  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
129    const below = await next(e)
130    const v = await read($, view)
131    if (!live.isPolling) {
132      $.clock.after(0, () => {
133        void ensurePolling($)
134      })
135    }
136    // Never write state while drawing: the saved choice is loaded just after.
137    if (!band.isLoaded) {
138      $.clock.after(0, () => {
139        void isBandHidden($)
140      })
141    }
142    if (await read($, bandHidden)) return below
143    const lows = v?.balances.filter(b => b.isLow) ?? []
144    if (lows.length === 0 || e.props.hasSurvey) return below
145    const { Box, Text } = $.ui.resolve(e)
146    return (
147      <Box flexDirection="column">
148        <Text wrap="truncate-end">
149          <Text color="#ef4565">{'▲ '}</Text>
150          <Text>{lows.map(b => `${b.label} ${b.text}`).join(' · ')}</Text>
151          <Text dimColor>{' · /jrquota details'}</Text>
152        </Text>
153        {below}
154      </Box>
155    )
156  })
157}
158

1// Drawing helpers shared by every jr mod. Synced into each plugin's hooks/jr/.
2import type { Elements } from 'claude-code'
3
4import { bar, paint } from './gradient'
5
6type TextElement = Elements['terminal']['Text']
7
8/** `text` in a gradient, one color per character, bold by default. */
9export const gradientText = (Text: TextElement, text: string, stops: readonly string[], bold = true) => (
10  <Text bold={bold}>
11    {paint(text, stops).map(({ char, color }) => (
12      <Text color={color}>{char}</Text>
13    ))}
14  </Text>
15)
16
17/** A gauge: filled cells in the gradient, empty cells dim. */
18export const gauge = (Text: TextElement, ratio: number, width: number, stops: readonly string[]) => (
19  <Text>
20    {bar(ratio, width, stops).map(cell =>
21      cell.empty ? (
22        <Text color={cell.color} dimColor>
23          {cell.char}
24        </Text>
25      ) : (
26        <Text color={cell.color}>{cell.char}</Text>
27      ),
28    )}
29  </Text>
30)
31
32/** The jr tag that opens every band line: a small gradient label. */
33export const tag = (Text: TextElement, label: string, stops: readonly string[]) =>
34  gradientText(Text, `${label} `, stops)
35

1// Shared look of every jr mod: one palette, gradients drawn one cell at a time.
2// Synced into each plugin's hooks/jr/ by scripts/sync-shared.sh: edit it here only.
3
4export type Rgb = readonly [number, number, number]
5
6const hexToRgb = (hex: string): Rgb => {
7  const n = Number.parseInt(hex.replace('#', ''), 16)
8  return [(n >> 16) & 255, (n >> 8) & 255, n & 255]
9}
10
11const rgbToHex = ([r, g, b]: Rgb) =>
12  `#${[r, g, b].map(v => Math.round(v).toString(16).padStart(2, '0')).join('')}`
13
14export const GRADIENTS = {
15  // The jr signature: violet to pink to amber.
16  brand: ['#7f5af0', '#ff5ca8', '#ffb86b'],
17  // Calm to alarm, for gauges where full is bad.
18  heat: ['#2cb67d', '#f4d35e', '#ff8c42', '#ef4565'],
19  // Cool, for neutral figures.
20  ice: ['#3da9fc', '#7f5af0'],
21  // Money and growth.
22  money: ['#2cb67d', '#a8e6cf'],
23} as const satisfies Record<string, readonly string[]>
24
25export type GradientName = keyof typeof GRADIENTS
26
27/** The color at `t` (0..1) along the stops. */
28export const colorAt = (stops: readonly string[], t: number): string => {
29  if (stops.length === 1) return stops[0]!
30  const clamped = Math.min(1, Math.max(0, t))
31  const span = clamped * (stops.length - 1)
32  const i = Math.min(stops.length - 2, Math.floor(span))
33  const a = hexToRgb(stops[i]!)
34  const b = hexToRgb(stops[i + 1]!)
35  const local = span - i
36  return rgbToHex([0, 1, 2].map(k => a[k]! + (b[k]! - a[k]!) * local) as unknown as Rgb)
37}
38
39/** One color per character of `text`, spread across the stops. */
40export const paint = (text: string, stops: readonly string[]) => {
41  const chars = [...text]
42  return chars.map((char, i) => ({ char, color: colorAt(stops, chars.length > 1 ? i / (chars.length - 1) : 0) }))
43}
44
45const EIGHTHS = ['', '▏', '▎', '▍', '▌', '▋', '▊', '▉']
46
47/**
48 * A bar `width` cells wide filled to `ratio` (0..1), with eighth-cell precision.
49 * Each filled cell takes the gradient color of its position along the whole bar,
50 * so a nearly full heat bar ends red and a low one stays green.
51 */
52export const bar = (ratio: number, width: number, stops: readonly string[]) => {
53  const filled = Math.min(1, Math.max(0, ratio)) * width
54  const full = Math.floor(filled)
55  const partial = EIGHTHS[Math.round((filled - full) * 7)] ?? ''
56  const cells: { char: string; color: string; empty?: true }[] = []
57  for (let i = 0; i < width; i++) {
58    const color = colorAt(stops, width > 1 ? i / (width - 1) : 0)
59    if (i < full) cells.push({ char: '█', color })
60    else if (i === full && partial) cells.push({ char: partial, color })
61    else cells.push({ char: '·', color, empty: true })
62  }
63  return cells
64}
65
66export const formatTokens = (n: number) =>
67  n >= 1_000_000 ? `${(n / 1_000_000).toFixed(1)}M` : n >= 10_000 ? `${Math.round(n / 1000)}k` : n >= 1000 ? `${(n / 1000).toFixed(1)}k` : `${n}`
68
69export const formatDuration = (ms: number) => {
70  const minutes = Math.max(0, Math.round(ms / 60_000))
71  const h = Math.floor(minutes / 60)
72  const m = minutes % 60
73  return h > 0 ? `${h}h${String(m).padStart(2, '0')}` : `${m}min`
74}
75

1import type { JrQuotaBalance } from '../types'
2
3export type Provider = {
4  id: string
5  label: string
6  /** userConfig field holding the key, and the environment variable read when it is empty. */
7  option: string
8  env: string
9  url: string
10  /** Turns the response body into a balance, given the low thresholds. */
11  parse: (json: unknown, low: { usd: number; credits: number }) => Omit<JrQuotaBalance, 'id' | 'label'>
12}
13
14const num = (v: unknown) => (typeof v === 'number' ? v : typeof v === 'string' ? Number.parseFloat(v) : Number.NaN)
15
16export const PROVIDERS: Provider[] = [
17  {
18    id: 'openrouter',
19    label: 'OpenRouter',
20    option: 'openrouterKey',
21    env: 'OPENROUTER_API_KEY',
22    url: 'https://openrouter.ai/api/v1/credits',
23    parse: (json, low) => {
24      const data = (json as { data?: { total_credits?: unknown; total_usage?: unknown } }).data ?? {}
25      const total = num(data.total_credits)
26      const left = total - num(data.total_usage)
27      if (!Number.isFinite(left)) return { text: 'unexpected answer', isLow: false, ratio: null, error: true }
28      return { text: `$${left.toFixed(2)} left`, isLow: left < low.usd, ratio: total > 0 ? left / total : null, error: false }
29    },
30  },
31  {
32    id: 'firecrawl',
33    label: 'Firecrawl',
34    option: 'firecrawlKey',
35    env: 'FIRECRAWL_API_KEY',
36    url: 'https://api.firecrawl.dev/v2/team/credit-usage',
37    parse: (json, low) => {
38      const data = (json as { data?: Record<string, unknown> }).data ?? {}
39      const left = num(data.remainingCredits ?? data.remaining_credits)
40      const plan = num(data.planCredits ?? data.plan_credits)
41      if (!Number.isFinite(left)) return { text: 'unexpected answer', isLow: false, ratio: null, error: true }
42      return {
43        text: `${Math.round(left).toLocaleString('en-US')} credits left`,
44        isLow: left < low.credits,
45        ratio: Number.isFinite(plan) && plan > 0 ? Math.min(1, left / plan) : null,
46        error: false,
47      }
48    },
49  },
50  {
51    id: 'deepseek',
52    label: 'DeepSeek',
53    option: 'deepseekKey',
54    env: 'DEEPSEEK_API_KEY',
55    url: 'https://api.deepseek.com/user/balance',
56    parse: (json, low) => {
57      const info = (json as { balance_infos?: { currency?: string; total_balance?: unknown }[] }).balance_infos?.[0]
58      const left = num(info?.total_balance)
59      if (!info || !Number.isFinite(left)) return { text: 'unexpected answer', isLow: false, ratio: null, error: true }
60      return { text: `${left.toFixed(2)} ${info.currency ?? ''} left`.trim(), isLow: left < low.usd, ratio: null, error: false }
61    },
62  },
63]
64
65/** What a failed request means for the person. */
66export const failureText = (status: number) =>
67  status === 401 ? 'key refused' : status === 403 ? 'key lacks access (OpenRouter needs a management key)' : `HTTP ${status}`
68

1export type JrQuotaBalance = {
2  id: string
3  label: string
4  text: string
5  isLow: boolean
6  /** Share left of the plan or of what was bought, 0..1, when the provider says. */
7  ratio: number | null
8  error: boolean
9}
10
11export type JrQuotaView = { balances: JrQuotaBalance[]; checkedAt: number }
12
13declare module 'claude-code' {
14  interface PluginState {
15    'jr-quota': { view: JrQuotaView | null; isBandHidden: boolean }
16  }
17}
18