omp-quota

1import type { On } from 'claude-code'
2import { quotaModelOf } from './quota-model.ts'
3import type { QuotaView, WindowCell } from './quota-model.ts'
4import { readUsage, worsenedProviders } from './usage.ts'
5import type { OmpOutcome, UsageReading } from './usage.ts'
6
7const FETCH_ARGV = ['omp', 'usage', '--json']
8const INVALIDATE_ARGV = ['omp', 'usage', 'invalidate']
9const OMP_TIMEOUT_MS = 30_000
10const POLL_MS = 300_000
11const BAND_KEY = 'band'
12
13let view: QuotaView = { usage: null, failure: null, lastGoodAt: null }
14let bandOn = false
15let poll: { cancel(): void } | null = null
16let fetchesStarted = 0
17let lastPublished = 0
18
19async function runOmp($: any, home: string, argv: string[]): Promise<OmpOutcome> {
20  try {
21    // No shell, so no `~`; and env can only override the PI_CODING_AGENT_DIR that
22    // Claude Code's own settings pass down, never unset it.
23    const result = await $.process.run(argv, {
24      env: { PI_CODING_AGENT_DIR: `${home}/.omp/agent` },
25      timeoutMs: OMP_TIMEOUT_MS,
26    })
27    return { kind: 'exited', exitCode: result.exitCode, stdout: result.stdout }
28  } catch {
29    return { kind: 'rejected' }
30  }
31}
32
33async function fetchUsage($: any): Promise<UsageReading> {
34  const home = await $.env.get('HOME')
35  if (!home) return { ok: false, reason: 'HOME is unset' }
36  return readUsage(await runOmp($, home, FETCH_ARGV))
37}
38
39// A fetch that started before one already published would put older data back on screen.
40async function publish($: any, seq: number, reading: UsageReading) {
41  const now = await $.clock.now()
42  if (seq < lastPublished) return
43  lastPublished = seq
44  if (reading.ok) {
45    const worsened = worsenedProviders(view.usage, reading.usage)
46    view = { usage: reading.usage, failure: null, lastGoodAt: now }
47    if (worsened.length) $.ui.toast(`omp quota: ${worsened.map((w) => `${w.provider} now ${w.status}`).join(', ')}`)
48  } else {
49    view = { ...view, failure: reading.reason }
50  }
51  $.ui.invalidate('ui.render')
52}
53
54async function fetchAndPublish($: any) {
55  const seq = ++fetchesStarted
56  await publish($, seq, await fetchUsage($))
57}
58
59async function refresh($: any) {
60  const home = await $.env.get('HOME')
61  if (!home) {
62    await publish($, ++fetchesStarted, { ok: false, reason: 'HOME is unset' })
63    return { text: 'omp quota refresh failed: HOME is unset' }
64  }
65  const invalidated = await runOmp($, home, INVALIDATE_ARGV)
66  const seq = ++fetchesStarted
67  const reading = readUsage(await runOmp($, home, FETCH_ARGV))
68  await publish($, seq, reading)
69  const note = invalidated.kind === 'exited' && invalidated.exitCode === 0 ? '' : ' (cache not invalidated)'
70  return { text: reading.ok ? `omp quota refreshed${note}` : `omp quota refresh failed: ${reading.reason}${note}` }
71}
72
73async function readBand($: any) {
74  try {
75    bandOn = (await $.store.get(BAND_KEY)) === true
76  } catch {
77    bandOn = false
78  }
79}
80
81async function toggleBand($: any) {
82  bandOn = !bandOn
83  try {
84    await $.store.set(BAND_KEY, bandOn)
85  } catch {
86    // The band still toggles for this session; only the next one will not remember it.
87  }
88  $.ui.invalidate('ui.render')
89  return {}
90}
91
92// Every window cell is padded to the widest one in its column, so windows line up across
93// providers; a line's last cell is left unpadded so no line ends in spaces.
94async function renderBand($: any, e: any, beneath: unknown) {
95  const { Box, Text } = await $.ui.resolve(e)
96  const model = quotaModelOf(view, await $.clock.now())
97  const cells = model.providers.flatMap((p) => p.windows)
98  const nameWidth = Math.max(0, ...model.providers.map((p) => p.provider.length))
99  const windowWidth = Math.max(0, ...cells.map((c) => c.window.length))
100  const cellText = (c: WindowCell) =>
101    `${c.window.padEnd(windowWidth)} ${c.share.padStart(4)} ${c.resets}${c.status ? ` ${c.status}` : ''}`
102  const columnWidths: number[] = []
103  for (const { windows } of model.providers) {
104    windows.slice(0, -1).forEach((c, i) => (columnWidths[i] = Math.max(columnWidths[i] ?? 0, cellText(c).length)))
105  }
106  const cellSpans = (c: WindowCell, i: number, last: boolean) => {
107    const padded = last ? cellText(c) : cellText(c).padEnd(columnWidths[i]!)
108    const at = windowWidth + 1 + 4 - c.share.length
109    const share = c.shareColor ? Text({ color: c.shareColor, children: [c.share] }) : c.share
110    return ['  ', padded.slice(0, at), share, padded.slice(at + c.share.length)]
111  }
112  const lines = []
113  if (model.notice) lines.push(Text({ dimColor: true, wrap: 'truncate-end', children: [model.notice] }))
114  for (const { provider, windows } of model.providers) {
115    const spans = windows.flatMap((c, i) => cellSpans(c, i, i === windows.length - 1))
116    lines.push(Text({ wrap: 'truncate-end', children: [provider.padEnd(nameWidth), ...spans] }))
117  }
118  if (beneath) lines.push(beneath)
119  return Box({ flexDirection: 'column', children: lines })
120}
121
122export function register(on: On) {
123  on('session.start', async ($, e, next) => {
124    try {
125      await $.command.register({
126        name: 'quota',
127        description: 'Show omp provider quota',
128        argumentHint: '[refresh]',
129        immediate: true,
130      })
131    } catch {
132      // A refused /quota leaves the poll working.
133    }
134    await readBand($)
135    void fetchAndPublish($).catch(() => {})
136    // A reload re-fires session.start on this instance; a second timer would double the cadence.
137    poll?.cancel()
138    poll = $.clock.every(POLL_MS, () => {
139      void fetchAndPublish($).catch(() => {})
140    })
141    return next(e)
142  })
143
144  on('command.run', { command: 'quota' }, async ($, e) => {
145    const args = (e.args ?? '').trim()
146    if (args === '') return toggleBand($)
147    if (args === 'refresh') return refresh($)
148    return { text: 'usage: /quota [refresh]' }
149  })
150
151  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
152    if (!bandOn || e.props.hasSurvey) return next(e)
153    // The band is one instance: a tree that drops next(e) hides every other plugin's band.
154    return renderBand($, e, await next(e))
155  })
156}
157

1import type { LimitQuota, Usage } from './usage.ts'
2
3export type QuotaView = { usage: Usage | null; failure: string | null; lastGoodAt: number | null }
4
5export type WindowCell = { window: string; share: string; shareColor?: string; status: string | null; resets: string }
6
7export type ProviderSection = { provider: string; windows: WindowCell[] }
8
9export type QuotaModel = { notice: string | null; providers: ProviderSection[] }
10
11const MINUTE = 60_000
12const HOUR = 60 * MINUTE
13const DAY = 24 * HOUR
14
15const RED = '#e5484d'
16const ORANGE = '#f5a524'
17const GREEN = '#46a758'
18
19function percentOf(share: number | null): string {
20  return share === null ? '—' : `${Math.round(share * 100)}%`
21}
22
23export function shareColorOf(share: number | null): string | undefined {
24  if (share === null) return undefined
25  const percent = Math.round(share * 100)
26  return percent <= 30 ? RED : percent <= 60 ? ORANGE : GREEN
27}
28
29function durationOf(ms: number): string {
30  if (ms >= DAY) return `${Math.floor(ms / DAY)}d ${Math.floor((ms % DAY) / HOUR)}h`
31  if (ms >= HOUR) return `${Math.floor(ms / HOUR)}h ${Math.floor((ms % HOUR) / MINUTE)}m`
32  return `${Math.floor(ms / MINUTE)}m`
33}
34
35function resetsOf(resetsAt: number | null, nowMs: number): string {
36  if (resetsAt === null) return '—'
37  return resetsAt <= nowMs ? 'now' : durationOf(resetsAt - nowMs)
38}
39
40function leastIndexOf(limits: LimitQuota[]): number {
41  let least = 0
42  limits.forEach((limit, i) => {
43    const share = limits[least]!.share
44    if (limit.share !== null && (share === null || limit.share < share)) least = i
45  })
46  return least
47}
48
49function worstBadStatusOf(limits: LimitQuota[]): string | null {
50  if (limits.some((l) => l.status === 'exhausted')) return 'exhausted'
51  return limits.some((l) => l.status === 'warning') ? 'warning' : null
52}
53
54function byNullLast(a: number | null, b: number | null): number {
55  return (a ?? Infinity) - (b ?? Infinity)
56}
57
58function windowsOf(limits: LimitQuota[], nowMs: number): WindowCell[] {
59  const groups = new Map<string, LimitQuota[]>()
60  for (const limit of limits) {
61    const key = limit.windowLabel || limit.label
62    groups.set(key, [...(groups.get(key) ?? []), limit])
63  }
64  return [...groups]
65    .map(([window, members]) => ({ window, members, least: members[leastIndexOf(members)]! }))
66    .sort((a, b) => byNullLast(a.least.resetsAt, b.least.resetsAt))
67    .map(({ window, members, least }) => ({
68      window,
69      share: percentOf(least.share),
70      shareColor: shareColorOf(least.share),
71      status: worstBadStatusOf(members),
72      resets: resetsOf(least.resetsAt, nowMs),
73    }))
74}
75
76function sectionsOf(usage: Usage, nowMs: number): ProviderSection[] {
77  return usage.providers
78    .filter((provider) => provider.limits.length)
79    .sort((a, b) => byNullLast(a.share, b.share))
80    .map((provider) => ({ provider: provider.provider, windows: windowsOf(provider.limits, nowMs) }))
81}
82
83export function quotaModelOf(view: QuotaView, nowMs: number): QuotaModel {
84  if (!view.usage) {
85    return { notice: view.failure ? `Unavailable: ${view.failure}` : 'Fetching omp usage', providers: [] }
86  }
87  const notice = view.failure
88    ? `Stale: ${view.failure}; showing data from ${durationOf(nowMs - (view.lastGoodAt ?? nowMs))} ago`
89    : null
90  return { notice, providers: sectionsOf(view.usage, nowMs) }
91}
92

1export type QuotaStatus = 'ok' | 'warning' | 'exhausted'
2
3export type LimitQuota = {
4  id: string
5  label: string
6  windowLabel: string
7  resetsAt: number | null
8  share: number | null
9  status: QuotaStatus | null
10}
11
12export type ProviderQuota = {
13  provider: string
14  limits: LimitQuota[]
15  share: number | null
16  status: QuotaStatus | null
17}
18
19export type Usage = { providers: ProviderQuota[] }
20
21export type OmpOutcome = { kind: 'exited'; exitCode: number; stdout: string } | { kind: 'rejected' }
22
23export type UsageReading = { ok: true; usage: Usage } | { ok: false; reason: string }
24
25const STATUS_RANK: Record<QuotaStatus, number> = { ok: 0, warning: 1, exhausted: 2 }
26
27function isFiniteNumber(value: unknown): value is number {
28  return typeof value === 'number' && Number.isFinite(value)
29}
30
31function field(value: unknown, key: string): unknown {
32  return value !== null && typeof value === 'object' ? (value as Record<string, unknown>)[key] : undefined
33}
34
35export function remainingShareOf(amount: unknown): number | null {
36  const remaining = field(amount, 'remainingFraction')
37  const used = field(amount, 'usedFraction')
38  const share = isFiniteNumber(remaining) ? remaining : isFiniteNumber(used) ? 1 - used : null
39  return share === null ? null : Math.min(1, Math.max(0, share))
40}
41
42function statusOf(value: unknown): QuotaStatus | null {
43  return value === 'ok' || value === 'warning' || value === 'exhausted' ? value : null
44}
45
46function stringOf(value: unknown): string {
47  return typeof value === 'string' ? value : ''
48}
49
50function limitOf(raw: unknown): LimitQuota | null {
51  const id = field(raw, 'id')
52  if (typeof id !== 'string') return null
53  const window = field(raw, 'window')
54  const resetsAt = field(window, 'resetsAt')
55  return {
56    id,
57    label: stringOf(field(raw, 'label')),
58    windowLabel: stringOf(field(window, 'label')),
59    resetsAt: isFiniteNumber(resetsAt) ? resetsAt : null,
60    share: remainingShareOf(field(raw, 'amount')),
61    status: statusOf(field(raw, 'status')),
62  }
63}
64
65function providerOf(provider: string, limits: LimitQuota[]): ProviderQuota {
66  const shares = limits.map((l) => l.share).filter((s): s is number => s !== null)
67  const statuses = limits.map((l) => l.status).filter((s): s is QuotaStatus => s !== null)
68  return {
69    provider,
70    limits,
71    share: shares.length ? Math.min(...shares) : null,
72    status: statuses.length ? statuses.reduce((a, b) => (STATUS_RANK[b] > STATUS_RANK[a] ? b : a)) : null,
73  }
74}
75
76export function readUsage(outcome: OmpOutcome): UsageReading {
77  if (outcome.kind === 'rejected') return { ok: false, reason: 'omp did not answer' }
78  if (outcome.exitCode !== 0) return { ok: false, reason: `omp exited ${outcome.exitCode}` }
79  let parsed: unknown
80  try {
81    parsed = JSON.parse(outcome.stdout)
82  } catch {
83    return { ok: false, reason: 'omp output unreadable' }
84  }
85  const reports = field(parsed, 'reports')
86  if (!Array.isArray(reports)) return { ok: false, reason: 'omp output unreadable' }
87  if (reports.length === 0) return { ok: false, reason: 'omp reported no providers' }
88  const grouped = new Map<string, LimitQuota[]>()
89  for (const report of reports) {
90    const provider = stringOf(field(report, 'provider')) || '(unnamed)'
91    const rawLimits = field(report, 'limits')
92    const limits = (Array.isArray(rawLimits) ? rawLimits : []).map(limitOf).filter((l): l is LimitQuota => l !== null)
93    grouped.set(provider, [...(grouped.get(provider) ?? []), ...limits])
94  }
95  return { ok: true, usage: { providers: [...grouped].map(([provider, limits]) => providerOf(provider, limits)) } }
96}
97
98export type Worsened = { provider: string; status: 'warning' | 'exhausted' }
99
100export function worsenedProviders(previous: Usage | null, current: Usage): Worsened[] {
101  if (!previous) return []
102  const before = new Map(previous.providers.map((p) => [p.provider, p.status]))
103  return current.providers.flatMap((p) =>
104    before.get(p.provider) === 'ok' && (p.status === 'warning' || p.status === 'exhausted')
105      ? [{ provider: p.provider, status: p.status }]
106      : [],
107  )
108}
109