bees-budget

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { CodexLaunch, CodexSession } from '../types'
5
6const CONTINUE_LINE = 150_000
7const RETIRE_AT = 200_000
8const COLD_LIMIT = 100_000
9const CLAUDE_TTL_MS = 5 * 60_000
10const CODEX_TTL_MS = 30 * 60_000
11const SHOWN_FOR_MS = 60 * 60_000
12const REFRESH_MS = 15_000
13const HANDOVER = '[bees:handover]'
14const OPENCODE_URL = 'http://127.0.0.1:4096'
15
16const contexts = atom({ plugin: 'bees-budget', key: 'contexts' } as const, {})
17const codex = atom({ plugin: 'bees-budget', key: 'codex' } as const, [])
18const launches = atom({ plugin: 'bees-budget', key: 'launches' } as const, [])
19
20type Row = { vendor: string; model: string; label: string; tokens: number; at: number; ttlMs: number; isRunning: boolean; codex?: CodexSession }
21
22const contextTokens = (u: ModelUsage) =>
23  u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens + u.output_tokens
24
25const kilo = (n: number) => `${Math.round(n / 1000)}k`
26const isRetired = (tokens: number) => tokens >= RETIRE_AT
27const isCold = (r: Row, now: number) => !r.isRunning && now - r.at > r.ttlMs
28const isPastContinueLine = (tokens: number) => tokens >= CONTINUE_LINE
29const isBlocked = (r: Row, now: number) => isPastContinueLine(r.tokens) || (isCold(r, now) && r.tokens >= COLD_LIMIT)
30
31const refusal = (r: Row, now: number) => {
32  if (isRetired(r.tokens))
33    return `${r.label} is at ${kilo(r.tokens)} context, past the ${kilo(RETIRE_AT)} retire line: spawn a fresh bee and hand the resource over (bees reference/budgets.md). A retirement message may start with ${HANDOVER}.`
34  if (isPastContinueLine(r.tokens))
35    return `${r.label} is at ${kilo(r.tokens)} context, past the ${kilo(CONTINUE_LINE)} continue line: send the next round to a fresh bee (bees reference/budgets.md). A retirement message may start with ${HANDOVER}.`
36  if (isBlocked(r, now))
37    return `${r.label}'s cache is cold (idle ${Math.round((now - r.at) / 60_000)}m, clock ${r.ttlMs / 60_000}m) at ${kilo(r.tokens)} context, at or past ${kilo(COLD_LIMIT)}: spawn a fresh bee instead (bees reference/budgets.md).`
38  return undefined
39}
40
41const modelEffort = (model: string, effort: string | null | undefined) => (effort ? `${model}-${effort}` : model)
42
43const frontmatter = (text: string, key: string) => text.match(new RegExp(`^${key}:\\s*(\\S+)`, 'm'))?.[1]
44
45const agentModels = new Map<string, Promise<string>>()
46
47const readAgentModel = async ($: EngineInterface, type: string) => {
48  const text = await $.fs.read(`${await $.env.get('HOME')}/.claude/agents/${type}.md`).catch(() => '')
49  const model = (frontmatter(text, 'model') ?? '?').replace(/^claude-|-.*$/g, '')
50  return modelEffort(model, frontmatter(text, 'effort'))
51}
52
53const agentModel = ($: EngineInterface, type: string) => {
54  if (!agentModels.has(type)) agentModels.set(type, readAgentModel($, type))
55  return agentModels.get(type)!
56}
57
58const beeRows = async ($: EngineInterface): Promise<(Row & { id: string })[]> => {
59  const recorded = await read($, contexts)
60  const bees = (await $.agent.list()).flatMap(a => {
61    const bee = recorded[a.id]
62    return a.type.startsWith('bee-') && bee ? [{ ...a, ...bee }] : []
63  })
64  return Promise.all(bees.map(async a => ({
65    id: a.id, vendor: 'claude', model: await agentModel($, a.type), label: a.name ?? a.description,
66    tokens: a.tokens, at: a.at, ttlMs: CLAUDE_TTL_MS, isRunning: a.status === 'running',
67  })))
68}
69
70const assistant = `json_extract(data, '$.role') = 'assistant'`
71
72const codexQuery = (since: number, liveSince: number) => `
73  with touched as (select session_id, max(time_updated) as at from part group by session_id)
74  select s.id, s.title as label, s.directory, s.time_created as createdAt,
75    json_extract(l.data, '$.modelID') as model, json_extract(l.data, '$.variant') as effort,
76    coalesce(json_extract(m.data, '$.tokens.total'),
77      ifnull(json_extract(m.data, '$.tokens.input'), 0) + ifnull(json_extract(m.data, '$.tokens.output'), 0)
78      + ifnull(json_extract(m.data, '$.tokens.reasoning'), 0) + ifnull(json_extract(m.data, '$.tokens.cache.read'), 0)
79      + ifnull(json_extract(m.data, '$.tokens.cache.write'), 0)) as tokens,
80    coalesce(json_extract(m.data, '$.time.completed'), s.time_created) as at,
81    l.id is not null and json_extract(l.data, '$.time.completed') is null
82      and json_extract(l.data, '$.error') is null and t.at > ${liveSince} as isRunning
83  from session s join touched t on t.session_id = s.id
84  left join message m on m.id = (
85    select id from message where session_id = s.id and ${assistant}
86      and json_extract(data, '$.time.completed') is not null
87      and json_extract(data, '$.error') is null order by time_created desc limit 1)
88  left join message l on l.id = (
89    select id from message where session_id = s.id and ${assistant} order by time_created desc limit 1)
90  where s.parent_id is null and t.at > ${since} order by t.at desc limit 20`
91
92const queryOpencode = async <T,>($: EngineInterface, sql: string): Promise<T[]> => {
93  const data = (await $.env.get('XDG_DATA_HOME')) || `${await $.env.get('HOME')}/.local/share`
94  const ran = await $.process.run(['sqlite3', '-json', `${data}/opencode/opencode.db`, sql])
95  return ran.exitCode === 0 && ran.stdout.trim() ? JSON.parse(ran.stdout) : []
96}
97
98const readCodexSessions = async ($: EngineInterface, now: number): Promise<CodexSession[]> =>
99  (await queryOpencode<CodexSession>($, codexQuery(now - SHOWN_FOR_MS, now - CODEX_TTL_MS)))
100    .map(s => ({ ...s, isRunning: Boolean(s.isRunning) }))
101
102const isLaunchedHere = (s: CodexSession, ls: CodexLaunch[]) =>
103  ls.some(l => l.id === s.id || (l.title === s.label && s.createdAt >= l.since))
104
105const refresh = async ($: EngineInterface) => {
106  const sessions = await readCodexSessions($, await $.clock.now())
107  const { startedAt } = await $.session.usage()
108  const ls = await read($, launches)
109  await update($, codex, () => sessions.filter(s => isLaunchedHere(s, ls) && (s.isRunning || s.at >= startedAt)))
110  $.ui.invalidate('ui.render')
111}
112
113const unquote = (arg: string) => arg.replace(/^(['"])(.*)\1$/, '$2')
114const flag = (command: string, name: string) => {
115  const arg = command.match(new RegExp(`(?:^|\\s)${name}\\s+("[^"]*"|'[^']*'|\\S+)`))?.[1]
116  return arg && unquote(arg)
117}
118const basename = (path: string) => path.split('/').pop() ?? path
119
120const launchTitle = (command: string) => {
121  const tool = command.match(/opencode-(implement|review)|opencode run/)?.[0]
122  if (!tool) return undefined
123  const title = flag(command, '--title')
124  const out = flag(command, '-o')
125  if (title || !out) return title
126  return tool === 'opencode-review' ? `review ${basename(out)}` : basename(out)
127}
128
129const sessionIds = (text: string) => [...new Set(text.match(/\bses_[A-Za-z0-9]+/g) ?? [])]
130
131const recordLaunches = async ($: EngineInterface, command: string, output: string) => {
132  const since = await $.clock.now()
133  const title = launchTitle(command)
134  const found: CodexLaunch[] = [...sessionIds(`${command}\n${output}`).map(id => ({ id, since })), ...(title ? [{ title, since }] : [])]
135  if (found.length) await update($, launches, all => [...all, ...found])
136}
137
138const attach = async ($: EngineInterface, s: CodexSession) => {
139  const command = ['opencode', 'attach', OPENCODE_URL, '--dir', s.directory, '--session', s.id]
140  const socket = await $.env.get('KITTY_LISTEN_ON')
141  const window = await $.env.get('KITTY_WINDOW_ID')
142  if (!socket || !window) return $.ui.toast(command.join(' '))
143  const ran = await $.process.run(['kitty', '@', '--to', socket, 'launch', '--location=vsplit', `--next-to=id:${window}`,
144    '--cwd', s.directory, '--title', `codex ${s.label}`, ...command])
145  if (ran.exitCode !== 0) $.ui.toast(`kitty launch failed: ${ran.stderr.trim()}`)
146}
147
148const cacheText = (r: Row, now: number) =>
149  r.isRunning ? 'running' : isCold(r, now) ? 'cold' : `warm ${Math.ceil((r.ttlMs - (now - r.at)) / 60_000)}m`
150
151const statusColor = (r: Row, now: number) => (isBlocked(r, now) ? 'red' : r.isRunning ? 'green' : undefined)
152
153const tokenColor = (tokens: number) =>
154  isRetired(tokens) ? 'red' : isPastContinueLine(tokens) ? 'yellow' : undefined
155
156export const register: Register = on => {
157  on('session.start', async ($, e, next) => {
158    $.clock.every(REFRESH_MS, () => void refresh($))
159    void refresh($)
160    return next(e)
161  })
162
163  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
164    const ran = await next(e)
165    const output = ran.deny === undefined && !ran.isError ? `${ran.result.stdout}\n${ran.result.stderr}` : ''
166    await recordLaunches($, e.command, output)
167    return ran
168  })
169
170  on('turn.step', async function* ($, e, next) {
171    const step = yield* next(e)
172    const { agentId } = e
173    const { usage } = step
174    if (agentId && usage) {
175      const bee = { tokens: contextTokens(usage), at: await $.clock.now() }
176      await update($, contexts, all => ({ ...all, [agentId]: bee }))
177    }
178    return step
179  })
180
181  on('session.send', async ($, e, next) => {
182    if (e.text.startsWith(HANDOVER)) return next(e)
183    const to = e.to.replace(/ \[[^\]]*\]$/, '')
184    const bee = (await beeRows($)).find(r => r.id === to || r.label === to)
185    const reason = bee && refusal(bee, await $.clock.now())
186    return reason ? { isDelivered: false, reason } : next(e)
187  })
188
189  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
190    if (e.props.hasSurvey) return next(e)
191    const now = await $.clock.now()
192    const codexRows = (await read($, codex)).map(s => ({
193      ...s, vendor: 'codex', model: modelEffort(s.model?.split('-').pop() ?? '?', s.effort), ttlMs: CODEX_TTL_MS, codex: s,
194    }))
195    const rows = [...(await beeRows($)), ...codexRows].filter(r => r.isRunning || now - r.at < SHOWN_FOR_MS)
196    if (rows.length === 0) return next(e)
197
198    const { Box, Button, Text } = $.ui.resolve(e)
199    return (
200      <Box flexDirection="column">
201        {rows.map(r => (
202          <Box key={r.codex?.id ?? r.label} gap={2}>
203            <Text dimColor>{r.vendor} {r.model}</Text>
204            {r.codex
205              ? <Button key={r.codex.id} plain onPress={() => void attach($, r.codex!)}>{r.label}</Button>
206              : <Text wrap="truncate-end">{r.label}</Text>}
207            <Text color={tokenColor(r.tokens)}>{kilo(r.tokens)}</Text>
208            <Text dimColor={!statusColor(r, now)} color={statusColor(r, now)}>
209              {cacheText(r, now)}{isBlocked(r, now) ? ' · spawn fresh' : ''}
210            </Text>
211          </Box>
212        ))}
213      </Box>
214    )
215  })
216}
217

1export type BeeContext = { tokens: number; at: number }
2export type CodexSession = {
3  id: string; label: string; directory: string; createdAt: number; model: string | null; effort: string | null
4  tokens: number; at: number; isRunning: boolean
5}
6export type CodexLaunch = { id?: string; title?: string; since: number }
7
8declare module 'claude-code' {
9  interface PluginState {
10    'bees-budget': { contexts: Record<string, BeeContext>; codex: CodexSession[]; launches: CodexLaunch[] }
11  }
12}
13