Brain Drain for Claude Code

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import * as agentBars from './agents.ts'
5import * as core from './core.ts'
6import * as notify from './notify.ts'
7import type { DrainAgent, DrainAgents, DrainMeta, DrainMiss, DrainNotify, DrainPrefs, DrainReading } from '../types'
8
9// Brain Drain, the mod: the data source for the VS Code extension (the bridge's
10// session file) and a coloured line under the prompt (or above it, by choice). Never calls a model,
11// never compacts, never rewrites the prompt. The one thing it hands the model is the self-notice,
12// off unless asked for: a line of facts when a threshold is crossed.
13
14const reading = atom({ plugin: 'brain-drain', key: 'reading' } as const, null)
15const prefs = atom({ plugin: 'brain-drain', key: 'prefs' } as const, null)
16const drafting = atom({ plugin: 'brain-drain', key: 'drafting' } as const, false)
17const turns = atom({ plugin: 'brain-drain', key: 'turns' } as const, null)
18const meta = atom({ plugin: 'brain-drain', key: 'meta' } as const, null)
19const agents = atom({ plugin: 'brain-drain', key: 'agents' } as const, null)
20
21const DEFAULT_PREFS: DrainPrefs = { style: core.DEFAULT_STYLE, place: 'below', warnSeconds: 60, sound: false, agents: false, notify: notify.DEFAULT_NOTIFY }
22const CHIME = 'assets/chime.wav'
23const MISS_RETENTION_S = 60 * 86400
24
25type $ = EngineInterface
26
27async function nowS($: $): Promise<number> {
28  return (await $.clock.now()) / 1000
29}
30
31/** The prefs in `$.store`, sanitised and merged over the defaults. */
32async function loadPrefs($: $): Promise<DrainPrefs> {
33  const stored = (await $.store.get('prefs')) as Partial<DrainPrefs> | undefined
34  const p: DrainPrefs = { ...DEFAULT_PREFS, ...(stored ?? {}) }
35  if (!core.isStyle(p.style)) p.style = core.DEFAULT_STYLE
36  if (p.place !== 'below' && p.place !== 'above') p.place = 'below'
37  p.agents = p.agents === true
38  p.notify = notify.sanitize(p.notify)
39  return p
40}
41
42/**
43 * The session's prefs. While `$.state` is blank (a /clear, not yet seeded again)
44 * the stored ones, so a `/drain` change made then keeps the other stored choices.
45 * Reads only: a timer or a render may call it.
46 */
47async function getPrefs($: $): Promise<DrainPrefs> {
48  return (await read($, prefs)) ?? (await loadPrefs($))
49}
50
51async function setPrefs($: $, change: Partial<DrainPrefs>): Promise<DrainPrefs> {
52  const next = { ...(await getPrefs($)), ...change }
53  await update($, prefs, () => next)
54  await $.store.set('prefs', next)
55  return next
56}
57
58async function getMeta($: $): Promise<DrainMeta> {
59  return (await read($, meta)) ?? {}
60}
61
62async function patchMeta($: $, change: Partial<DrainMeta>): Promise<void> {
63  await update($, meta, m => ({ ...(m ?? {}), ...change }))
64}
65
66async function patchReading($: $, fn: (r: DrainReading) => DrainReading): Promise<DrainReading> {
67  let out = core.emptyReading()
68  await update($, reading, r => {
69    out = fn(r ?? core.emptyReading())
70    return out
71  })
72  return out
73}
74
75async function readJson($: $, path: string): Promise<unknown> {
76  try {
77    return JSON.parse((await $.fs.read(path)) as string)
78  } catch {
79    return undefined
80  }
81}
82
83// ---------------------------------------------------------------- writing files
84
85// `$.fs.write` is not atomic, so a file is written beside its target and moved
86// into place. A failed move removes its temp file; `$.fs` cannot delete, so a
87// host with no commands at all leaves that one file, once per process.
88let canMove = true
89
90async function writeAtomic($: $, path: string, text: string): Promise<void> {
91  if (canMove) {
92    const tmp = `${path}.mod-${Math.random().toString(36).slice(2, 10)}.tmp`
93    await $.fs.write(tmp, text)
94    try {
95      const moved = await $.process.run(['mv', '-f', tmp, path], { timeoutMs: 5000 })
96      if (moved.exitCode === 0) return
97      await $.process.run(['rm', '-f', tmp], { timeoutMs: 5000 }).catch(() => undefined)
98    } catch {
99      // No host commands here (a host other than the CLI): write in place from now on.
100      canMove = false
101    }
102  }
103  await $.fs.write(path, text)
104}
105
106let writeChain: Promise<void> = Promise.resolve()
107
108/** Queues a write of this session's file; writes run one at a time, each from the latest reading. */
109function writeSessionFile($: $): Promise<void> {
110  writeChain = writeChain.then(() => doWriteSessionFile($)).catch(() => undefined)
111  return writeChain
112}
113
114async function doWriteSessionFile($: $): Promise<void> {
115  await syncSession($)
116  // Who holds the status line is checked before every write, not only at a
117  // turn's end: after Connect Status Line the bridge writes the extension's
118  // directory at once, and a long turn must not keep writing it too (#4983).
119  await refreshOutDir($)
120  const m = await getMeta($)
121  const r = await read($, reading)
122  // A `claude -p` run has no status line, so the bridge writes no file for it.
123  // Neither does the mod: in primary mode it would take over the VS Code gauge.
124  if (!m.interactive || !m.outDir || !r) return
125  const ids = { sessionId: await $.session.id(), cwd: await $.session.cwd(), projectDir: await $.session.root() }
126  const file = core.sessionFile(r, ids, await nowS($))
127  await writeAtomic($, `${m.outDir}/${ids.sessionId}.json`, JSON.stringify(file, null, 2))
128}
129
130/**
131 * session.start fires once per process, never on /clear or /resume. Claude
132 * Code 2.1.300 empties `$.state` on /clear, and `ensureSeeded` builds it again.
133 * It keeps `$.state` on a /resume in the same process, and so may another host
134 * on /clear; that lands here: when the id changes the reading starts again,
135 * from this mod's own file for that id (a /resume), else empty but for the
136 * account's TTL, the rate limits and the model. The 1-second tick notices the
137 * change (`followSession`), so it does not wait for a request.
138 */
139async function syncSession($: $): Promise<void> {
140  const id = await $.session.id()
141  const m = await getMeta($)
142  if (m.sessionId === id) return
143  // Transcripts are named after their session, side by side; the next classic event confirms the path.
144  // A SessionStart that came before the engine's id switched has already stored the new path: kept.
145  const oldName = `/${m.sessionId}.jsonl`
146  const transcriptPath = m.transcriptPath?.endsWith(`/${id}.jsonl`)
147    ? m.transcriptPath
148    : m.sessionId !== undefined && m.transcriptPath?.endsWith(oldName) ? `${m.transcriptPath.slice(0, -oldName.length)}/${id}.jsonl` : undefined
149  await patchMeta($, { sessionId: id, transcriptPath, alertedFor: undefined, pendingMiss: undefined, notified: undefined })
150  await clearAgents($)
151  const prev = await read($, reading)
152  if (m.sessionId === undefined || !prev) return
153  const own = m.outDir ? core.readingFromFile(await readJson($, `${m.outDir}/${id}.json`), prev.ttl, prev.ttlSource) : undefined
154  const usage = await $.session.usage()
155  const next = { ...(own ?? core.freshSession(prev)), rateLimits: core.limitsFrom(usage.rateLimits), costUsd: usage.cost?.usd }
156  await update($, reading, () => applyContext(next, usage.context))
157}
158
159// ---------------------------------------------------------------- the bridge handover
160
161async function chooseOutDir($: $, home: string): Promise<{ outDir: string; shadow: boolean }> {
162  let live = true
163  try {
164    // The merged settings: what this session actually runs, a project's override included.
165    live = core.bridgeLive(await $.settings.read())
166  } catch {
167    // Unreadable: assume the bridge is live and stay out of its directory.
168  }
169  return { outDir: core.outDirFor(home, live), shadow: live }
170}
171
172/** Re-reads the status line, so re-wiring it mid-session moves the writes at the next write. */
173async function refreshOutDir($: $): Promise<void> {
174  const m = await getMeta($)
175  if (!m.home) return
176  const where = await chooseOutDir($, m.home)
177  if (where.outDir !== m.outDir) await patchMeta($, where)
178}
179
180async function accountTtl($: $, m: DrainMeta): Promise<number | undefined> {
181  const paths = [`${m.outDir}/account_ttl.json`]
182  if (m.home && m.shadow) paths.push(`${m.home}/${core.PRIMARY_DIR}/account_ttl.json`)
183  for (const p of paths) {
184    const f = (await readJson($, p)) as { detected_ttl?: unknown } | undefined
185    if (f && typeof f.detected_ttl === 'number' && f.detected_ttl > 0) return f.detected_ttl
186  }
187  return undefined
188}
189
190// The extension prunes its own directory (session files after 36 h, temp files
191// after 1 h). Nothing prunes the shadow directory, so the mod does, alike.
192const SESSION_MAX_AGE_MS = 36 * 3600_000
193const TMP_MAX_AGE_MS = 3600_000
194const META_FILES = new Set(['account_ttl.json', 'config.json'])
195
196async function pruneShadow($: $, m: DrainMeta): Promise<void> {
197  if (!m.shadow || !m.interactive || !m.outDir || !canMove) return
198  const now = await $.clock.now()
199  const stale: string[] = []
200  let entries: { name: string; kind: string }[] = []
201  try {
202    entries = await $.fs.list(m.outDir)
203  } catch {
204    return
205  }
206  for (const entry of entries) {
207    if (entry.kind !== 'file' || META_FILES.has(entry.name)) continue
208    const maxAge = entry.name.endsWith('.tmp') ? TMP_MAX_AGE_MS : entry.name.endsWith('.json') ? SESSION_MAX_AGE_MS : undefined
209    if (maxAge === undefined) continue
210    const path = `${m.outDir}/${entry.name}`
211    try {
212      if ((await $.fs.stat(path)).mtimeMs < now - maxAge) stale.push(path)
213    } catch {
214      // gone already
215    }
216  }
217  if (stale.length > 0) await $.process.run(['rm', '-f', ...stale], { timeoutMs: 5000 }).catch(() => undefined)
218}
219
220// ---------------------------------------------------------------- exact TTL from the transcript
221
222async function tail($: $, path: string, bytes: number, size: number): Promise<string> {
223  if (canMove) {
224    try {
225      const r = await $.process.run(['tail', '-c', String(bytes), path], { timeoutMs: 5000 })
226      if (r.exitCode === 0) return r.stdout
227    } catch {
228      // fall through to a whole read
229    }
230  }
231  if (size > 4 * 1024 * 1024) throw new Error('transcript too large to read whole')
232  const text = (await $.fs.read(path)) as string
233  return text.slice(Math.max(0, text.length - bytes))
234}
235
236async function refreshTtl($: $): Promise<void> {
237  const m = await getMeta($)
238  const path = m.transcriptPath
239  if (!path || !m.outDir) return
240  // Headless with nothing drawing: the TTL would show nowhere, so spawn nothing for it.
241  if (!m.interactive && (await $.session.surfaces()).length === 0) return
242  let size = 0
243  try {
244    size = (await $.fs.stat(path)).size
245  } catch {
246    return
247  }
248  let found: { ttl: number; atS: number } | undefined
249  for (const window of core.TAIL_WINDOWS) {
250    const cut = size > window
251    try {
252      found = core.ttlFromTail(await tail($, path, window, size), cut)
253    } catch {
254      return
255    }
256    if (found || !cut) break
257  }
258  if (!found) return
259  const hit = found
260  await patchReading($, r => ({ ...r, ttl: hit.ttl, ttlSource: 'transcript' }))
261  const accountPath = `${m.outDir}/account_ttl.json`
262  const stored = (await readJson($, accountPath)) as { detected_ttl?: number; detected_at?: number } | undefined
263  if (m.interactive && core.shouldWriteAccountTtl(stored, hit.ttl, hit.atS)) {
264    await writeAtomic($, accountPath, JSON.stringify({ detected_ttl: hit.ttl, detected_at: hit.atS }))
265  }
266}
267
268// turn.start does not say whose turn it is; turn.step and turn.complete do.
269async function dropTurn($: $, turnId: string): Promise<void> {
270  if (!(await read($, turns))?.includes(turnId)) return
271  await update($, turns, t => (t ?? []).filter(id => id !== turnId))
272}
273
274// ---------------------------------------------------------------- subagent token bars (prototype)
275
276// Read on a timer, never while drawing: the rows draw from `agents` in $.state.
277const AGENT_REFRESH_MS = 5000
278/** Each `tail` answers at most 4 MiB; a long addition catches up over several reads. */
279const READS_PER_FILE = 8
280const WHOLE_READ_MAX = 16 * 1024 * 1024
281
282type Tracked = { acc: agentBars.Acc; label?: string }
283let agentDir: string | undefined
284let tracked = new Map<string, Tracked>()
285let agentBusy = false
286let lastAgentSig: string | undefined
287
288/** `<dir>/<sessionId>.jsonl` → `<dir>/<sessionId>/subagents`. */
289function subagentDir(transcriptPath: string): string | undefined {
290  return transcriptPath.endsWith('.jsonl') ? `${transcriptPath.slice(0, -'.jsonl'.length)}/subagents` : undefined
291}
292
293/** Agent transcripts: the directory's own, and a workflow's one level down in `workflows/<wf>/`. */
294async function agentFiles($: $, dir: string): Promise<{ path: string; size: number; mtimeMs: number }[]> {
295  const out: { path: string; size: number; mtimeMs: number }[] = []
296  const walk = async (d: string, depth: number): Promise<void> => {
297    let entries: { name: string; kind: string; size: number; mtimeMs: number }[]
298    try {
299      entries = await $.fs.list(d)
300    } catch {
301      return
302    }
303    for (const e of entries) {
304      if (e.kind === 'file' && e.name.startsWith('agent-') && e.name.endsWith('.jsonl')) out.push({ path: `${d}/${e.name}`, size: e.size, mtimeMs: e.mtimeMs })
305      else if (e.kind === 'dir' && (depth === 0 ? e.name === 'workflows' : depth === 1)) await walk(`${d}/${e.name}`, depth + 1)
306    }
307  }
308  await walk(dir, 0)
309  return out
310}
311
312/** Reads what the file gained since the last refresh. */
313async function readNew($: $, path: string, t: Tracked, size: number): Promise<void> {
314  if (size < t.acc.offset) t.acc = agentBars.newAcc() // replaced, not appended to
315  if (size === t.acc.offset) return
316  // The first read takes the file whole, one call and no process; a host command reads only what each later refresh adds.
317  if (canMove && (t.acc.offset > 0 || size > WHOLE_READ_MAX)) {
318    try {
319      for (let i = 0; i < READS_PER_FILE && t.acc.offset < size; i++) {
320        const before = t.acc.offset
321        const r = await $.process.run(['tail', '-c', `+${before + 1}`, path], { timeoutMs: 5000 })
322        if (r.exitCode !== 0) return
323        agentBars.feed(t.acc, r.stdout, r.isStdoutTruncated)
324        if (!r.isStdoutTruncated || t.acc.offset === before) return
325      }
326      return
327    } catch {
328      // No host commands here: read the file whole instead.
329    }
330  }
331  if (size > WHOLE_READ_MAX) return
332  const acc = agentBars.newAcc()
333  agentBars.feed(acc, (await $.fs.read(path)) as string)
334  t.acc = acc
335}
336
337async function clearAgents($: $): Promise<void> {
338  agentDir = undefined
339  tracked = new Map()
340  lastAgentSig = undefined
341  if ((await read($, agents)) !== null) await update($, agents, () => null)
342}
343
344async function setAgents($: $, next: DrainAgents): Promise<void> {
345  // The guess moves with the clock, so it is part of what counts as a change.
346  const now = await $.clock.now()
347  const sig = JSON.stringify(next) + next.agents.map(a => agentBars.activity(a, now) ?? '').join(',')
348  if (sig === lastAgentSig) return
349  lastAgentSig = sig
350  await update($, agents, () => next)
351  $.ui.invalidate('ui.render')
352}
353
354async function refreshAgents($: $): Promise<void> {
355  if (!(await getPrefs($)).agents) return
356  const m = await getMeta($)
357  const dir = m.transcriptPath ? subagentDir(m.transcriptPath) : undefined
358  if (dir !== agentDir) {
359    await clearAgents($)
360    agentDir = dir
361  }
362  if (!dir || (await $.session.surfaces()).length === 0) return
363  const files = await agentFiles($, dir)
364  const status = new Map<string, string>()
365  try {
366    for (const a of await $.agent.list()) status.set(a.id, a.status)
367  } catch {
368    // No agent list here: the file's last write is the only clue.
369  }
370  const list: DrainAgent[] = []
371  for (const f of files) {
372    const id = f.path.slice(f.path.lastIndexOf('/') + 1 + 'agent-'.length, -'.jsonl'.length)
373    let t = tracked.get(f.path)
374    if (!t) tracked.set(f.path, (t = { acc: agentBars.newAcc() }))
375    if (t.label === undefined) {
376      const sidecar = await readJson($, `${f.path.slice(0, -'.jsonl'.length)}.meta.json`)
377      if (sidecar !== undefined) t.label = agentBars.labelFrom(sidecar, id)
378    }
379    try {
380      await readNew($, f.path, t, f.size)
381    } catch {
382      // unreadable this time; the next refresh tries again
383    }
384    const a = t.acc
385    if (a.requests === 0) continue
386    const st = status.get(id) ?? status.get(`agent-${id}`)
387    list.push({
388      id, label: t.label ?? id, tokens: a.tokens, output: a.output, requests: a.requests,
389      model: a.model, effort: a.effort, mtimeMs: f.mtimeMs, ...(st === undefined ? {} : { status: st }),
390    })
391  }
392  await setAgents($, { dir, agents: list })
393}
394
395/** One refresh at a time; a tick that finds one running skips. */
396async function refreshAgentsOnce($: $): Promise<void> {
397  if (agentBusy) return
398  agentBusy = true
399  try {
400    await refreshAgents($)
401  } catch (err) {
402    $.ui.log(`brain-drain: agents: ${String(err)}`, { to: 'debug' })
403  } finally {
404    agentBusy = false
405  }
406}
407
408// ---------------------------------------------------------------- display
409
410let lastStatus: string | undefined
411
412async function currentView($: $): Promise<{ view: core.View; r: DrainReading } | undefined> {
413  const r = await read($, reading)
414  if (!r) return undefined
415  return { view: core.makeView(r, await nowS($), await read($, drafting)), r }
416}
417
418async function refreshStatusLine($: $): Promise<void> {
419  const p = await getPrefs($)
420  let text: string | undefined
421  if (p.style === 'line') {
422    const cv = await currentView($)
423    text = cv ? core.lineText(cv.view) : undefined
424  }
425  if (text === lastStatus) return
426  lastStatus = text
427  $.ui.status(text)
428}
429
430let following: Promise<void> | undefined
431
432/**
433 * An in-process /resume changes the session id and keeps `$.state`, and no
434 * event reaches the mod until the next request. Checked every second: one id
435 * read, and files only when the id changed.
436 */
437async function followSession($: $): Promise<void> {
438  if (following || seeding) return
439  const id = await $.session.id()
440  if ((await getMeta($)).sessionId === id || following) return
441  following = (async () => {
442    await ensureSeeded($)
443    await syncSession($)
444    await writeSessionFile($)
445    await refreshStatusLine($)
446  })()
447    .catch(err => $.ui.log(`brain-drain: following the session: ${String(err)}`, { to: 'debug' }))
448    .finally(() => {
449      following = undefined
450    })
451  await following
452}
453
454async function tick($: $): Promise<void> {
455  await followSession($)
456  const p = await getPrefs($)
457  const r = await read($, reading)
458  if (!r) return
459  const t = await nowS($)
460  const c = core.cacheState(r, t)
461  const isTicking = c.kind === 'live' || (c.kind === 'expired' && c.since <= 2)
462  if (isTicking && core.isBandStyle(p.style)) $.ui.invalidate('ui.render')
463  if (p.style === 'line') await refreshStatusLine($)
464
465  // The expiry alert: once per cache window, only while no turn runs.
466  if (c.kind !== 'live' || p.warnSeconds <= 0 || c.left > p.warnSeconds) return
467  if ((await read($, turns))?.length) return
468  const m = await getMeta($)
469  if (m.alertedFor === r.lastApiS) return
470  await patchMeta($, { alertedFor: r.lastApiS })
471  $.ui.toast(`Cache expires in ${c.left} s · the next message re-writes ~${core.fmtTokens(r.tokens ?? 0)} tokens`, { timeoutMs: 8000 })
472  if (p.sound) $.audio.play({ asset: CHIME }).catch(() => undefined)
473}
474
475// ---------------------------------------------------------------- the ledger
476
477async function misses($: $): Promise<DrainMiss[]> {
478  const out: DrainMiss[] = []
479  for (const key of await $.store.keys()) {
480    if (!key.startsWith('miss:')) continue
481    const v = (await $.store.get(key)) as DrainMiss | undefined
482    if (v && typeof v.at === 'number') out.push(v)
483  }
484  return out
485}
486
487async function pruneMisses($: $): Promise<void> {
488  const cutoff = (await nowS($)) - MISS_RETENTION_S
489  for (const key of await $.store.keys()) {
490    if (!key.startsWith('miss:')) continue
491    const atMs = Number(key.split(':')[1])
492    if (Number.isFinite(atMs) && atMs / 1000 < cutoff) await $.store.delete(key)
493  }
494}
495
496// ---------------------------------------------------------------- the self-notice
497
498async function notifySettings($: $): Promise<DrainNotify> {
499  return notify.effective(notify.sanitize((await getPrefs($)).notify), (await getMeta($)).notifyEnv)
500}
501
502let notifyChain: Promise<unknown> = Promise.resolve()
503
504/**
505 * The note to hand the model now, or undefined. Marks what it reports as fired
506 * before it returns, one call at a time, so two tool calls ending together
507 * cannot both carry it. Works whatever `meta.interactive` says: a `claude -p`
508 * run gets the note though it gets no session file.
509 */
510function takeNotice($: $): Promise<{ text: string; crossed: string[] } | undefined> {
511  const run = notifyChain.then(async () => {
512    // The first prompt after a /clear comes before any request: seed first, so
513    // BRAIN_DRAIN_NOTIFY and the stored thresholds count.
514    await ensureSeeded($).catch(() => undefined)
515    const n = await notifySettings($)
516    if (!n.on) return undefined
517    await syncSession($)
518    const usage = await $.session.usage()
519    const f = notify.figures(usage.context, core.limitsFrom(usage.rateLimits))
520    const before = (await getMeta($)).notified ?? []
521    const t = await nowS($)
522    const { crossed, fired } = notify.evaluate(f, n, before, t)
523    if (crossed.length > 0 || fired.length !== before.length) await patchMeta($, { notified: fired })
524    return crossed.length > 0 ? { text: notify.noticeText(f, n, crossed, t), crossed } : undefined
525  })
526  notifyChain = run.catch(() => undefined)
527  return run
528}
529
530/** A note that never reached the model (the prompt was dropped): its thresholds fire again next time. */
531async function unfire($: $, crossed: string[]): Promise<void> {
532  await update($, meta, m => ({ ...(m ?? {}), notified: (m?.notified ?? []).filter(k => !crossed.includes(k)) }))
533}
534
535const THRESHOLD_KIND: Record<string, 'context' | 'fiveHour'> = { context: 'context', '5h': 'fiveHour', '5-hour': 'fiveHour' }
536
537/** Strips one pair of surrounding quotes, as typed in `/drain notify text context 35 "stop"`. */
538function unquote(s: string): string {
539  const t = s.trim()
540  return t.length >= 2 && ((t[0] === '"' && t.endsWith('"')) || (t[0] === "'" && t.endsWith("'"))) ? t.slice(1, -1) : t
541}
542
543async function notifyCommand($: $, args: string): Promise<string> {
544  const words = args.trim().split(/\s+/).filter(w => w !== '')
545  const sub = (words[0] ?? '').toLowerCase()
546  const saved = notify.sanitize((await getPrefs($)).notify)
547  const m = await getMeta($)
548  const save = async (change: Partial<DrainNotify>) => {
549    await setPrefs($, { notify: { ...saved, ...change } })
550  }
551  const envNote = m.notifyEnv ? ` This run follows ${notify.ENV_NAME}, which overrides the saved settings until it ends.` : ''
552  if (sub === '') {
553    const n = await notifySettings($)
554    const usage = await $.session.usage()
555    const now = notify.factsLine(notify.figures(usage.context, core.limitsFrom(usage.rateLimits)), n, await nowS($))
556    return [...notify.describe(n, m.notifyEnv !== undefined, m.notifyEnvError, m.notified ?? [], now), 'Change with /drain notify on|off, context <n..>, 5h <n..>, week on|off, margin <n>, text <kind> [n] "<text>".'].join('\n')
557  }
558  if (sub === 'on' || sub === 'off') {
559    await save({ on: sub === 'on' })
560    return sub === 'on'
561      ? `Self-notice on: when this session crosses a threshold, the model gets one line of facts with the next prompt or tool result. /drain notify shows the thresholds.${envNote}`
562      : `Self-notice off.${envNote}`
563  }
564  const kind = THRESHOLD_KIND[sub]
565  if (kind) {
566    const rest = words.slice(1)
567    if (rest.length === 1 && rest[0]!.toLowerCase() === 'off') {
568      await save({ [kind]: [] })
569      return `${sub} thresholds off.${envNote}`
570    }
571    const nums = rest.map(w => Number(w.replace(/%$/, '')))
572    if (nums.length === 0 || nums.length > 10 || !nums.every(x => Number.isFinite(x) && x >= 0 && x <= 100)) {
573      return `Give up to ten percentages from 0 to 100, e.g. /drain notify ${sub} ${kind === 'context' ? '70 90' : '80 95'}, or /drain notify ${sub} off.`
574    }
575    const list = [...new Set(nums)].sort((a, b) => a - b)
576    await save({ [kind]: list })
577    return `${sub} thresholds: ${list.map(t => `${t}%`).join(', ')}.${saved.on ? '' : ' (The self-notice is off; /drain notify on turns it on.)'}${envNote}`
578  }
579  if (sub === 'week') {
580    const v = (words[1] ?? '').toLowerCase()
581    if (v !== 'on' && v !== 'off') return 'Use /drain notify week on or /drain notify week off.'
582    await save({ week: v === 'on' })
583    return v === 'on' ? `Weekly pace threshold on: fires when the week's use passes the pace line (plus ${saved.margin} points).${envNote}` : `Weekly pace threshold off.${envNote}`
584  }
585  if (sub === 'margin') {
586    const n = Number(words[1])
587    if (words.length !== 2 || !Number.isFinite(n) || Math.abs(n) > 100) return 'Give the margin in points, e.g. /drain notify margin 5.'
588    await save({ margin: n })
589    return `Pace line margin: ${n} points.${envNote}`
590  }
591  if (sub === 'text') {
592    // text context 35 "..." | text 5h 80 "..." | text week "..."; no text clears it.
593    const k = (words[1] ?? '').toLowerCase()
594    const tk = THRESHOLD_KIND[k]
595    let key: string
596    let restStart: number
597    if (k === 'week') {
598      key = 'week'
599      restStart = 2
600    } else if (tk && words[2] !== undefined && Number.isFinite(Number(words[2]))) {
601      key = `${tk === 'context' ? 'context' : '5h'}:${Number(words[2])}`
602      restStart = 3
603    } else {
604      return 'Use /drain notify text context 35 "your text", /drain notify text 5h 80 "your text" or /drain notify text week "your text". Leave the text out to clear it.'
605    }
606    const text = unquote(args.trim().split(/\s+/).slice(restStart).join(' '))
607    const next = { ...saved.text }
608    if (text === '') delete next[key]
609    else next[key] = text
610    await save({ text: next })
611    const label = key.replace(':', ' ')
612    if (text === '') return `Text for ${label} cleared.${envNote}`
613    const configured = key === 'week' ? saved.week : key.startsWith('context:') ? saved.context.includes(Number(key.slice(8))) : saved.fiveHour.includes(Number(key.slice(3)))
614    return `When ${label} fires, the note adds: ${text}${configured ? '' : ` (No ${label} threshold is set yet.)`}${envNote}`
615  }
616  return 'Unknown /drain notify option. Use /drain notify on|off, context <n..>, 5h <n..>, week on|off, margin <n>, text <kind> [n] "<text>", or /drain notify to show the settings.'
617}
618
619// ---------------------------------------------------------------- /drain
620
621const SAMPLE: DrainReading = {
622  ...core.emptyReading(3600, 'transcript'),
623  percent: 31,
624  tokens: 312000,
625  window: 1000000,
626  rateLimits: [
627    { kind: 'five_hour', pct: 72 },
628    { kind: 'seven_day', pct: 44 },
629  ],
630  model: 'claude-opus-5-5',
631  effort: 'xhigh',
632  lastApiS: 0,
633  observed: true,
634}
635
636async function runCommand($: $, args: string): Promise<string> {
637  const [first = '', ...rest] = args.trim().split(/\s+/)
638  let sub = first.toLowerCase()
639  let arg = rest.join(' ')
640  // `/drain bars` and `/drain styles bars` are shorthand for `/drain style bars`.
641  if (core.isStyle(sub)) [sub, arg] = ['style', sub]
642  else if (sub === 'styles' && arg !== '') sub = 'style'
643  switch (sub) {
644    case 'help':
645      return ['Brain Drain commands:', ...core.HELP].join('\n')
646    case 'styles': {
647      const live = await read($, reading)
648      const t = await nowS($)
649      const r = live && live.percent !== undefined ? live : { ...SAMPLE, lastApiS: t - 2 }
650      const p = await getPrefs($)
651      const lines = [live && live.percent !== undefined ? 'Styles, with this session\'s figures:' : 'Styles, with sample figures:']
652      for (const style of core.STYLES) {
653        lines.push('', `${style}${style === p.style ? '  (current)' : ''}`)
654        lines.push(...core.styleText(style, core.makeView(r, t, false)).split('\n').map(l => `  ${l}`))
655      }
656      lines.push('', 'While the cache has expired and the prompt box holds text, the band adds:')
657      lines.push(`  ${core.styleText('minimal', core.makeView({ ...r, lastApiS: t - r.ttl - 60 }, t, true)).split('\n')[1] ?? ''}`)
658      lines.push('', 'Switch with /drain <name>, e.g. /drain gauge.')
659      return lines.join('\n')
660    }
661    case 'style': {
662      const name = arg.toLowerCase()
663      if (!core.isStyle(name)) return `Unknown style "${arg}". Styles: ${core.STYLES.join(', ')}.`
664      await setPrefs($, { style: name })
665      await refreshStatusLine($)
666      $.ui.invalidate('ui.render')
667      return name === 'off'
668        ? 'Display off. The session file for the VS Code extension is still written.'
669        : name === 'line'
670          ? 'Style: line (plain text under the prompt; Claude Code prefixes it with the mod name).'
671          : `Style: ${name}.`
672    }
673    case 'place': {
674      if (arg !== 'below' && arg !== 'above') return 'Place: below (under the prompt, the default) or above (the band above the prompt).'
675      await setPrefs($, { place: arg })
676      $.ui.invalidate('ui.render')
677      return `Place: ${arg} the prompt.`
678    }
679    case 'warn': {
680      const n = Number(arg)
681      if (!Number.isInteger(n) || n < 0 || n > 3600) return 'Give whole seconds from 0 to 3600, e.g. /drain warn 90. 0 turns the alert off.'
682      await setPrefs($, { warnSeconds: n })
683      return n === 0 ? 'Expiry alert off.' : `Expiry alert ${n} s before the cache expires.`
684    }
685    case 'sound': {
686      const v = arg.toLowerCase()
687      if (v !== 'on' && v !== 'off') return 'Use /drain sound on or /drain sound off.'
688      await setPrefs($, { sound: v === 'on' })
689      if (v === 'on') {
690        try {
691          await $.audio.play({ asset: CHIME })
692        } catch (err) {
693          return `Sound on, but the sample did not play: ${String(err)}`
694        }
695        return 'Sound on (that was the chime).'
696      }
697      return 'Sound off.'
698    }
699    case 'agents': {
700      const v = arg.toLowerCase()
701      if (v === 'on' || v === 'off') {
702        await setPrefs($, { agents: v === 'on' })
703        if (v === 'on') await refreshAgentsOnce($)
704        else await clearAgents($)
705        $.ui.invalidate('ui.render')
706        return v === 'on'
707          ? 'Subagent bars on (prototype): one row per subagent of this session under the gauge. Tokens are each request\'s input + cache write + cache read + output; "running?" is a guess from the transcript\'s last write.'
708          : 'Subagent bars off.'
709      }
710      if (v !== '') return 'Use /drain agents on, /drain agents off, or /drain agents to list them.'
711      if (!(await getPrefs($)).agents) return 'Subagent bars are off. /drain agents on turns them on.'
712      await refreshAgentsOnce($)
713      return agentBars.agentText((await read($, agents))?.agents ?? [], await $.clock.now())
714    }
715    case 'notify':
716      return notifyCommand($, args.trim().slice(first.length))
717    case '':
718    case 'report': {
719      const p = await getPrefs($)
720      const m = await getMeta($)
721      const r = await read($, reading)
722      const dir = m.outDir && m.home ? m.outDir.replace(m.home, '~') : '(not writing: no HOME)'
723      const lines = ['Brain Drain · cache misses (a request after the TTL ran out that re-wrote the cache)']
724      lines.push(...core.ledgerReport(await misses($), await nowS($)))
725      lines.push('Cost is the session ledger\'s measured rise over the turn that opened cold; misses before a cost reading have none.')
726      lines.push('')
727      lines.push(`TTL ${r ? `${r.ttl >= 3600 ? '1h' : '5m'} (${r.ttlSource})` : 'unknown'} · style ${p.style} · alert ${p.warnSeconds > 0 ? `${p.warnSeconds} s` : 'off'} · sound ${p.sound ? 'on' : 'off'} · agents ${p.agents ? 'on' : 'off'} · notify ${(await notifySettings($)).on ? 'on' : 'off'}`)
728      if (!m.interactive) lines.push('Session file: not written (not an interactive session)')
729      else lines.push(`Session file: ${dir}${m.shadow ? ' (shadow copy: the bridge is still the status line, so its directory is left alone)' : ''}`)
730      lines.push('More: /drain help')
731      return lines.join('\n')
732    }
733    default:
734      return `Unknown subcommand "${first}".\n${core.HELP.join('\n')}`
735  }
736}
737
738// ---------------------------------------------------------------- setup
739
740/** Once per load of this code, from session.start. The command and the timers survive /clear. */
741async function setup($: $, interactive: boolean): Promise<void> {
742  await $.command.register({
743    name: 'drain',
744    description: 'Brain Drain: cache-miss report, display styles, expiry alert',
745    argumentHint: '[help|styles|<style>|place below|above|warn <s>|sound on|off|agents on|off|notify ...]',
746    immediate: true,
747  })
748  $.clock.every(1000, () => {
749    void tick($).catch(() => undefined)
750  })
751  $.clock.every(AGENT_REFRESH_MS, () => {
752    void refreshAgentsOnce($)
753  })
754  await seed($, () => seedSession($, interactive, false))
755}
756
757// One seeding at a time. A module variable survives /clear; a hot reload resets
758// it, and session.start seeds again after one.
759let seeding: Promise<void> | undefined
760/**
761 * What session.start said, kept for the seedings after a /clear. The process
762 * stays what it was: a session that did not report itself interactive (`-p`,
763 * Anthropic's VS Code chat panel) must not start writing a file after a /clear.
764 * Unknown (a hook ran before this load's session.start) counts as not
765 * interactive: no guess may make a chat-panel session write a file, and
766 * session.start seeds again with the true value right after, a hot reload included.
767 */
768let startInteractive: boolean | undefined
769
770async function seed($: $, run: () => Promise<void>): Promise<void> {
771  while (seeding) await seeding
772  seeding = run()
773    .catch(err => $.ui.log(`brain-drain: seeding failed: ${String(err)}`, { to: 'debug' }))
774    .finally(() => {
775      seeding = undefined
776    })
777  await seeding
778}
779
780/**
781 * Seeds `$.state` again when the host has emptied it: Claude Code 2.1.300
782 * empties it on /clear, under a new session id, and fires no session.start.
783 * (An in-process /resume keeps it; `followSession` handles that.) `interactive`
784 * is set by every seeding, so its absence marks a blank state. Called first
785 * thing by the hooks that run after a /clear, and by the tick.
786 */
787async function ensureSeeded($: $): Promise<void> {
788  while (seeding) await seeding
789  if ((await getMeta($)).interactive !== undefined) return
790  if (seeding) return ensureSeeded($)
791  await seed($, async () => {
792    // Re-checked inside the slot: another caller may have seeded meanwhile.
793    if ((await getMeta($)).interactive !== undefined) return
794    await seedSession($, startInteractive ?? false, true)
795  })
796}
797
798/**
799 * Puts into `$.state` what this session needs: prefs, where to write, the
800 * reading and the context window. `blank`: the host emptied `$.state`, so a
801 * reading found there is a stray made since and is built again; otherwise (a hot
802 * reload) a kept reading stays.
803 */
804async function seedSession($: $, interactive: boolean, blank: boolean): Promise<void> {
805  const p = await loadPrefs($)
806  await update($, prefs, () => p)
807  const env = notify.parseEnv(await $.env.get('BRAIN_DRAIN_NOTIFY'))
808  if (env.error) $.ui.log(`brain-drain: ${env.error}`, { to: 'debug' })
809
810  const home = await $.env.get('HOME')
811  const where = home ? await chooseOutDir($, home) : { outDir: undefined, shadow: false }
812  await patchMeta($, { interactive, notifyEnv: env.value, notifyEnvError: env.error, home, outDir: where.outDir, shadow: where.shadow })
813  await syncSession($)
814  const m = await getMeta($)
815
816  // A reload keeps the reading in $.state. A new process or an emptied state
817  // seeds it from this mod's own earlier file (a resume), else from the account's TTL.
818  if ((blank || !(await read($, reading))) && m.outDir) {
819    const acct = await accountTtl($, m)
820    const ttl = acct ?? core.DEFAULT_TTL
821    const source = acct ? 'account' : 'default'
822    const own = core.readingFromFile(await readJson($, `${m.outDir}/${await $.session.id()}.json`), ttl, source)
823    await update($, reading, () => own ?? core.emptyReading(ttl, source))
824  } else if (blank || !(await read($, reading))) {
825    await update($, reading, () => core.emptyReading())
826  }
827
828  const usage = await $.session.usage()
829  await patchReading($, r => applyContext({ ...r, rateLimits: core.limitsFrom(usage.rateLimits), costUsd: usage.cost?.usd ?? r.costUsd }, usage.context))
830  await pruneMisses($)
831  await pruneShadow($, m)
832  await writeSessionFile($)
833  lastStatus = undefined
834  await refreshStatusLine($)
835}
836
837function applyContext(r: DrainReading, c: { tokens?: number; window: number; percent?: number }): DrainReading {
838  if (c.tokens === undefined) {
839    // A fresh session has no fill yet. After a compact (or on a resume) the
840    // fill is not known: hold the last figure and flag it, as the bridge does.
841    const held = (r.percent ?? 0) > 0
842    return {
843      ...r,
844      window: c.window,
845      unknown: held,
846      percent: held ? r.percent : 0,
847      // A compact seen by this process starts the cache clock again.
848      lastApiS: held && r.observed ? undefined : r.lastApiS,
849      observed: held && r.observed ? false : r.observed,
850    }
851  }
852  return { ...r, tokens: c.tokens, window: c.window, percent: c.percent ?? r.percent, unknown: false }
853}
854
855// ---------------------------------------------------------------- hooks
856
857export const register: Register = on => {
858  on('session.start', async ($, e, next) => {
859    const result = await next(e)
860    try {
861      startInteractive = e.isInteractive
862      await setup($, e.isInteractive)
863    } catch (err) {
864      $.ui.log(`brain-drain: setup failed: ${String(err)}`, { to: 'debug' })
865    }
866    return result
867  })
868
869  // One model request. Main thread only: subagents keep caches of their own.
870  on('turn.step', async function* ($, e, next) {
871    if (e.agentId !== undefined) {
872      await dropTurn($, e.turnId).catch(() => undefined)
873      return yield* next(e)
874    }
875    await ensureSeeded($).catch(() => undefined)
876    await syncSession($).catch(() => undefined)
877    const startMs = await $.clock.now()
878    const costAtStart = (await read($, reading))?.costUsd
879    const result = yield* next(e)
880    try {
881      const usage = result.usage
882      if (usage) {
883        const endS = (await $.clock.now()) / 1000
884        const prev = (await read($, reading)) ?? core.emptyReading()
885        const startS = startMs / 1000
886        const kind = core.classifyStep(prev, startS, usage)
887        let upgrade = core.warmGapUpgrade(prev, startS, usage)
888        if (upgrade) {
889          // The earlier requests' rows are on disk and name the tier exactly. A cache
890          // something else kept warm must not pass for the 1-hour tier until Stop.
891          await refreshTtl($).catch(() => undefined)
892          if ((await read($, reading))?.ttlSource === 'transcript') upgrade = false
893        }
894        const r = await patchReading($, cur => ({
895          ...cur,
896          lastApiS: endS,
897          observed: true,
898          cacheRead: usage.cache_read_input_tokens,
899          cacheCreate: usage.cache_creation_input_tokens,
900          outputTokens: usage.output_tokens,
901          ...core.stepFill(usage, cur.window),
902          model: usage.model || e.model,
903          effort: e.effort === undefined ? undefined : String(e.effort),
904          ...(upgrade ? { ttl: 3600, ttlSource: 'warm-gap' as const } : {}),
905        }))
906        if (kind === 'miss' && prev.lastApiS !== undefined) {
907          const sessionId = await $.session.id()
908          const root = await $.session.root()
909          const key = core.missKey(startMs, sessionId)
910          const miss: DrainMiss = {
911            at: Math.floor(startS),
912            session: sessionId,
913            project: root.split('/').pop() ?? '',
914            tokens: usage.cache_creation_input_tokens,
915            gapS: Math.round(startS - prev.lastApiS),
916            ttl: prev.ttl,
917          }
918          // If the session ledger already moved for this request, the cost is known now.
919          if (costAtStart !== undefined && r.costUsd !== undefined && r.costUsd !== costAtStart) {
920            miss.costUsd = r.costUsd - costAtStart
921          } else if (costAtStart !== undefined) {
922            await patchMeta($, { pendingMiss: { key, costBefore: costAtStart } })
923          }
924          await $.store.set(key, miss)
925        }
926        // Queued, not awaited: the next request must not wait on a file move.
927        void writeSessionFile($)
928      }
929    } catch (err) {
930      $.ui.log(`brain-drain: turn.step: ${String(err)}`, { to: 'debug' })
931    }
932    return result
933  })
934
935  on('session.measure', async ($, e, next) => {
936    const result = await next(e)
937    try {
938      await ensureSeeded($)
939      await patchReading($, r => applyContext({ ...r, rateLimits: core.limitsFrom(e.rateLimits), costUsd: e.cost?.usd ?? r.costUsd }, e.context))
940      const m = await getMeta($)
941      if (m.pendingMiss && e.cost && e.changed.includes('cost')) {
942        const pending = m.pendingMiss
943        const stored = (await $.store.get(pending.key)) as DrainMiss | undefined
944        if (stored) await $.store.set(pending.key, { ...stored, costUsd: e.cost.usd - pending.costBefore })
945        await patchMeta($, { pendingMiss: undefined })
946      }
947      await writeSessionFile($)
948    } catch (err) {
949      $.ui.log(`brain-drain: session.measure: ${String(err)}`, { to: 'debug' })
950    }
951    return result
952  })
953
954  // The transcript's path rides on every classic event; the TTL is read once the turn's rows are on disk.
955  on('classic.SessionStart', async ($, e, next) => {
956    const result = await next(e)
957    try {
958      // A /clear or a /resume fires this, though not session.start.
959      const sid = await $.session.id()
960      if (e.session_id && e.session_id !== sid) $.ui.log(`brain-drain: SessionStart (${e.source}) names session ${e.session_id}, the engine says ${sid}`, { to: 'debug' })
961      await ensureSeeded($)
962      await syncSession($)
963      await refreshOutDir($)
964      if (e.transcript_path) await patchMeta($, { transcriptPath: e.transcript_path })
965      await refreshTtl($)
966      await writeSessionFile($)
967    } catch {
968      // the TTL falls back to the previous value or the account's
969    }
970    return result
971  })
972
973  on('classic.Stop', async ($, e, next) => {
974    const result = await next(e)
975    try {
976      await ensureSeeded($)
977      await refreshOutDir($)
978      if (e.transcript_path) await patchMeta($, { transcriptPath: e.transcript_path })
979      await refreshTtl($)
980      await writeSessionFile($)
981    } catch {
982      // as above
983    }
984    return result
985  })
986
987  // The alert waits while a main turn runs. A subagent's turn leaves the list at
988  // its first request, so a background agent never holds the alert off.
989  on('turn.start', async ($, e, next) => {
990    await update($, turns, t => [...(t ?? []), e.turnId])
991    return next(e)
992  })
993
994  on('turn.complete', async ($, e, next) => {
995    const result = await next(e)
996    await dropTurn($, e.turnId)
997    return result
998  })
999
1000  // Observe only: the box the editor computed goes back exactly as it came.
1001  let lastDrafting: boolean | undefined
1002  on('prompt.edit', async ($, e, next) => {
1003    const box = await next(e)
1004    const has = box.text.trim().length > 0
1005    if (has !== lastDrafting) {
1006      lastDrafting = has
1007      await update($, drafting, () => has).catch(() => undefined)
1008    }
1009    return box
1010  })
1011
1012  // The self-notice rides on the prompt (attached on the way down) and on a
1013  // main-loop tool result, so a long run with no new prompt still hears it.
1014  on('prompt.submit', async ($, e, next) => {
1015    if (lastDrafting !== false) {
1016      lastDrafting = false
1017      await update($, drafting, () => false).catch(() => undefined)
1018    }
1019    const note = await takeNotice($).catch(err => {
1020      $.ui.log(`brain-drain: notify: ${String(err)}`, { to: 'debug' })
1021      return undefined
1022    })
1023    if (!note) return next(e)
1024    const result = await next({ ...e, context: [...(e.context ?? []), note.text] })
1025    if (result.drop !== undefined) await unfire($, note.crossed).catch(() => undefined)
1026    return result
1027  })
1028
1029  // Subagents' calls (agentId set) and a plugin's own $.tool.call (no context reaches a model there) pass untouched.
1030  on('tool.call', async ($, e, next) => {
1031    const result = await next(e)
1032    if (e.agentId !== undefined || next.origin.plugin !== 'engine' || result.deny !== undefined) return result
1033    try {
1034      const note = await takeNotice($)
1035      if (note) return { ...result, context: [...(result.context ?? []), note.text] }
1036    } catch (err) {
1037      $.ui.log(`brain-drain: notify: ${String(err)}`, { to: 'debug' })
1038    }
1039    return result
1040  })
1041
1042  on('command.run', { command: 'drain' }, async ($, e) => {
1043    await ensureSeeded($).catch(() => undefined)
1044    return { text: await runCommand($, e.args) }
1045  })
1046
1047  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
1048    const others = await next(e)
1049    if (e.props.hasSurvey) return others
1050    const p = await read($, prefs)
1051    if ((p?.place ?? 'below') !== 'above') return others
1052    const rows = await gaugeRows($, p, e.props.bodyColumns).catch(() => undefined)
1053    if (!rows) return others
1054    const { Box, Text } = $.ui.resolve(e)
1055    return (
1056      <Box flexDirection="column">
1057        {drawRows(Text, rows)}
1058        {others}
1059      </Box>
1060    )
1061  })
1062
1063  // Below the prompt: the gauge takes the hint line's place and the engine's
1064  // own line (shortcuts, "esc to interrupt") is kept underneath it.
1065  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
1066    const engine = await next(e)
1067    const p = await read($, prefs)
1068    if ((p?.place ?? 'below') !== 'below') return engine
1069    const rows = await gaugeRows($, p, e.viewport?.columns).catch(() => undefined)
1070    if (!rows) return engine
1071    const { Box, Text } = $.ui.resolve(e)
1072    return (
1073      <Box flexDirection="column">
1074        {drawRows(Text, rows)}
1075        {engine}
1076      </Box>
1077    )
1078  })
1079}
1080
1081/** The rows to draw, or undefined to draw nothing. Callers swallow a throw (a state read refused mid-unload) and pass the other rows through. */
1082async function gaugeRows($: $, p: DrainPrefs | null, cols: number | undefined): Promise<core.Row[] | undefined> {
1083  const style = p?.style ?? core.DEFAULT_STYLE
1084  if (!core.isBandStyle(style)) return undefined
1085  const r = await read($, reading)
1086  if (!r) return undefined
1087  const view = core.makeView(r, await nowS($), await read($, drafting))
1088  const rows = core.renderRows(style, view, cols)
1089  if (p?.agents && rows.length > 0) {
1090    const a = await read($, agents)
1091    if (a && a.agents.length > 0) rows.push(...agentBars.agentRows(a.agents, await $.clock.now(), cols))
1092  }
1093  return rows
1094}
1095
1096// eslint-disable-next-line @typescript-eslint/no-explicit-any
1097function drawRows(Text: any, rows: core.Row[]) {
1098  return rows.map((row, i) => (
1099    <Text key={i === 0 ? 'gauge' : `row-${i}`} wrap="truncate-end">
1100      {row.map(seg => (
1101        <Text color={seg.color} dimColor={seg.dim} bold={seg.bold}>
1102          {seg.text}
1103        </Text>
1104      ))}
1105    </Text>
1106  ))
1107}
1108

1// Pure logic for the per-subagent token bars (prototype, `/drain agents on`): no `$`, so the tests call it directly.
2//
3// A session's subagents write `<transcript dir>/<sessionId>/subagents/agent-<id>.jsonl`
4// (workflow agents one level down, in `workflows/<wf>/`), append-only, beside an
5// `agent-<id>.meta.json` naming the agent. register.tsx reads the new bytes of each
6// file on a timer and feeds them here; the prompt draws from the stored result.
7import { bar, cellWidth, effortLabel, fmtTokens, plain, shortModel } from './core.ts'
8import type { Row, Seg } from './core.ts'
9import type { DrainAgent } from '../types'
10
11/** An agent counts as running, as a guess, while its file changed this recently. */
12export const RUNNING_GUESS_MS = 30_000
13export const MAX_ROWS = 5
14
15/**
16 * One transcript's running count. A response is written as one record per content
17 * block, each carrying the response's usage; the input fields repeat and
18 * `output_tokens` grows to its final figure on the last block. So each response is
19 * counted once, by `message.id`, with the largest output count seen.
20 */
21export type Acc = {
22  /** Bytes of the file consumed: always just past a newline. */
23  offset: number
24  /** Per response: [input + cache write + cache read, output] as last seen. */
25  seen: Record<string, [number, number]>
26  tokens: number
27  output: number
28  requests: number
29  model?: string
30  effort?: string
31}
32
33export function newAcc(): Acc {
34  return { offset: 0, seen: {}, tokens: 0, output: 0, requests: 0 }
35}
36
37/** UTF-8 bytes a string takes; the module has no TextEncoder to ask. */
38export function utf8Bytes(s: string): number {
39  let n = 0
40  for (let i = 0; i < s.length; i++) {
41    const c = s.charCodeAt(i)
42    if (c < 0x80) n += 1
43    else if (c < 0x800) n += 2
44    else if (c >= 0xd800 && c <= 0xdbff && i + 1 < s.length) {
45      n += 4
46      i++
47    } else n += 3
48  }
49  return n
50}
51
52const num = (v: unknown): number => (typeof v === 'number' && Number.isFinite(v) ? v : 0)
53
54function take(acc: Acc, line: string): void {
55  if (!line.includes('"usage"')) return
56  let row: unknown
57  try {
58    row = JSON.parse(line)
59  } catch {
60    return
61  }
62  if (typeof row !== 'object' || row === null) return
63  const r = row as { type?: unknown; message?: unknown; requestId?: unknown; uuid?: unknown; effort?: unknown }
64  if (r.type !== 'assistant' || typeof r.message !== 'object' || r.message === null) return
65  const m = r.message as { id?: unknown; model?: unknown; usage?: Record<string, unknown> }
66  const u = m.usage
67  if (typeof u !== 'object' || u === null || m.model === '<synthetic>') return
68  const input = num(u.input_tokens) + num(u.cache_creation_input_tokens) + num(u.cache_read_input_tokens)
69  let out = num(u.output_tokens)
70  const key = typeof m.id === 'string' ? m.id : typeof r.requestId === 'string' ? r.requestId : typeof r.uuid === 'string' ? r.uuid : undefined
71  const prev = key === undefined ? undefined : acc.seen[key]
72  if (prev) {
73    // The count only grows over a response's blocks; a lower one is never the final figure.
74    out = Math.max(out, prev[1])
75    acc.tokens -= prev[0] + prev[1]
76    acc.output -= prev[1]
77  } else acc.requests++
78  if (key !== undefined) acc.seen[key] = [input, out]
79  acc.tokens += input + out
80  acc.output += out
81  if (typeof m.model === 'string') acc.model = m.model
82  if (typeof r.effort === 'string') acc.effort = r.effort
83}
84
85/**
86 * Feeds text read from `acc.offset` onwards. Only whole lines are taken: a last
87 * line still being written stays unread, and the offset stops before it.
88 * `cut` says the read was cut short (stdout's cap): a chunk with no newline is
89 * then the middle of one huge line, passed over so the file cannot stall on it.
90 * A read that starts mid-line skips to the next line.
91 */
92export function feed(acc: Acc, text: string, cut = false): void {
93  const end = text.lastIndexOf('\n')
94  if (end < 0) {
95    if (cut) acc.offset += utf8Bytes(text)
96    return
97  }
98  const midLine = acc.offset > 0 && text[0] !== '{' && text[0] !== '\n'
99  let body = text.slice(0, end + 1)
100  acc.offset += utf8Bytes(body)
101  if (midLine) body = body.slice(body.indexOf('\n') + 1)
102  for (const line of body.split('\n')) if (line) take(acc, line)
103}
104
105/** The label a meta sidecar gives: the agent's name, else its type. */
106export function labelFrom(meta: unknown, id: string): string {
107  if (typeof meta === 'object' && meta !== null) {
108    const m = meta as { name?: unknown; agentType?: unknown }
109    if (typeof m.name === 'string' && m.name) return m.name
110    if (typeof m.agentType === 'string' && m.agentType) return m.agentType
111  }
112  return id
113}
114
115/** `running` from the engine's own agent list, `running?` guessed from the file's last write, else nothing. */
116export function activity(a: DrainAgent, nowMs: number): 'running' | 'running?' | undefined {
117  if (a.status === 'running' || a.status === 'pending' || a.status === 'waiting') return 'running'
118  if (a.status !== undefined) return undefined
119  return nowMs - a.mtimeMs <= RUNNING_GUESS_MS ? 'running?' : undefined
120}
121
122function fit(label: string, width: number): string {
123  const chars = [...label]
124  if (chars.length <= width) return label + ' '.repeat(width - chars.length)
125  return `${chars.slice(0, width - 1).join('')}…`
126}
127
128type Level = { cells: number; out: boolean; model: boolean; labelMax: number }
129const LEVELS: Level[] = [
130  { cells: 10, out: true, model: true, labelMax: 18 },
131  { cells: 10, out: false, model: true, labelMax: 18 },
132  { cells: 10, out: false, model: false, labelMax: 18 },
133  { cells: 5, out: false, model: false, labelMax: 12 },
134]
135
136/**
137 * The rows under the gauge: one per agent (active first, then the most recently
138 * written, at most MAX_ROWS), each bar scaled to the session's largest agent,
139 * then one total line. Empty when the session has no agents.
140 */
141export function agentRows(agents: DrainAgent[], nowMs: number, cols?: number): Row[] {
142  if (agents.length === 0) return []
143  const max = Math.max(1, ...agents.map(a => a.tokens))
144  const ranked = [...agents].sort((a, b) => {
145    const act = Number(activity(b, nowMs) !== undefined) - Number(activity(a, nowMs) !== undefined)
146    return act !== 0 ? act : b.mtimeMs - a.mtimeMs
147  })
148  const shown = ranked.slice(0, MAX_ROWS)
149  const total = agents.reduce((n, a) => n + a.tokens, 0)
150  const output = agents.reduce((n, a) => n + a.output, 0)
151  const more = agents.length - shown.length
152
153  const build = (lv: Level): Row[] => {
154    const width = Math.min(lv.labelMax, Math.max(...shown.map(a => [...a.label].length)))
155    const tokW = Math.max(...shown.map(a => fmtTokens(a.tokens).length))
156    const outW = Math.max(...shown.map(a => fmtTokens(a.output).length))
157    const rows: Row[] = shown.map(a => {
158      const row: Row = [{ text: `  ${fit(a.label, width)} `, dim: true }, ...bar(a.tokens / max, lv.cells, undefined)]
159      row.push({ text: ` ${fmtTokens(a.tokens).padStart(tokW)}` })
160      if (lv.out) row.push({ text: `  out ${fmtTokens(a.output).padEnd(outW)}`, dim: true })
161      const model = shortModel(a.model)
162      const effort = effortLabel(a.effort)
163      if (lv.model && model) row.push({ text: `  ${effort ? `${model}·${effort}` : model}`, dim: true })
164      const act = activity(a, nowMs)
165      if (act) row.push({ text: `  ${act}` })
166      return row
167    })
168    const sum: Seg[] = [{ text: `  ${agents.length} agent${agents.length === 1 ? '' : 's'} · ${fmtTokens(total)} tokens`, dim: true }]
169    if (lv.out) sum.push({ text: ` · ${fmtTokens(output)} out`, dim: true })
170    if (more > 0) sum.push({ text: ` · +${more} more`, dim: true })
171    rows.push(sum)
172    return rows
173  }
174  for (const lv of LEVELS) {
175    const rows = build(lv)
176    if (cols === undefined || rows.every(r => cellWidth(plain(r)) <= cols)) return rows
177  }
178  return build(LEVELS[LEVELS.length - 1]!)
179}
180
181/** The rows as plain text, for `/drain agents`. */
182export function agentText(agents: DrainAgent[], nowMs: number): string {
183  if (agents.length === 0) return 'No subagents in this session yet.'
184  return agentRows(agents, nowMs).map(plain).join('\n')
185}
186

1// Pure logic: no `$`, so the tests call it directly. register.tsx wires it to events.
2import type { DrainLimit, DrainMiss, DrainReading, DrainStyle } from '../types'
3
4export const STYLES: readonly DrainStyle[] = ['minimal', 'bars', 'gauge', 'compact', 'line', 'off']
5export const BAND_STYLES: readonly DrainStyle[] = ['minimal', 'bars', 'gauge', 'compact']
6export const DEFAULT_STYLE: DrainStyle = 'minimal'
7export const DEFAULT_TTL = 300
8
9// Theme keys, so the light, dark and colour-blind themes each pick a readable shade.
10export const OK = 'success'
11export const WARN = 'warning'
12export const BAD = 'error'
13
14// The extension's thresholds: context 60/80, cache 60 s/15 s, rate limits 70/90.
15export const CONTEXT_WARN = 60
16export const CONTEXT_DANGER = 80
17export const CACHE_WARN_S = 60
18export const CACHE_DANGER_S = 15
19export const LIMIT_WARN = 70
20export const LIMIT_DANGER = 90
21
22// A request counts as a cache miss only when it re-wrote at least this much.
23export const MIN_MISS_TOKENS = 2000
24
25export const TAIL_WINDOWS = [65536, 262144, 1048576] as const
26
27export function isStyle(s: string): s is DrainStyle {
28  return (STYLES as readonly string[]).includes(s)
29}
30
31export function isBandStyle(s: DrainStyle): boolean {
32  return BAND_STYLES.includes(s)
33}
34
35// ---------------------------------------------------------------- colours
36
37export function isImplausible(r: Pick<DrainReading, 'percent' | 'tokens' | 'window'>): boolean {
38  if (r.window !== undefined && r.window > 0 && r.tokens !== undefined && r.tokens > r.window) return true
39  return (r.percent ?? 0) > 100
40}
41
42export function contextColor(pct: number | undefined, unknown: boolean, implausible: boolean): string | undefined {
43  if (unknown || implausible || pct === undefined) return undefined
44  if (pct >= CONTEXT_DANGER) return BAD
45  if (pct >= CONTEXT_WARN) return WARN
46  return OK
47}
48
49export type CacheState =
50  | { kind: 'none' }
51  | { kind: 'live'; left: number }
52  | { kind: 'expired'; since: number }
53
54export function cacheState(r: Pick<DrainReading, 'lastApiS' | 'ttl'>, nowS: number): CacheState {
55  if (r.lastApiS === undefined || r.lastApiS <= 0) return { kind: 'none' }
56  const left = Math.floor(r.lastApiS + r.ttl - nowS)
57  if (left <= 0) return { kind: 'expired', since: -left }
58  return { kind: 'live', left }
59}
60
61export function cacheColor(c: CacheState, ctx: string | undefined): string | undefined {
62  if (c.kind === 'expired') return BAD
63  if (c.kind === 'none') return undefined
64  if (c.left < CACHE_DANGER_S) return BAD
65  if (c.left < CACHE_WARN_S) return WARN
66  return ctx
67}
68
69export function limitColor(pct: number): string | undefined {
70  if (pct >= LIMIT_DANGER) return BAD
71  if (pct >= LIMIT_WARN) return WARN
72  return undefined
73}
74
75// ---------------------------------------------------------------- formatting
76
77export function fmtCountdown(s: number): string {
78  const m = Math.floor(s / 60)
79  return `${m}:${String(s % 60).padStart(2, '0')}`
80}
81
82export function fmtTokens(n: number): string {
83  if (n >= 1e6) return `${(n / 1e6).toFixed(1).replace(/\.0$/, '')}M`
84  if (n >= 1000) return `${Math.round(n / 1000)}k`
85  return String(n)
86}
87
88export function fmtDuration(s: number): string {
89  const h = Math.floor(s / 3600)
90  const m = Math.floor((s % 3600) / 60)
91  return h > 0 ? `${h}h ${m}m` : `${m}m`
92}
93
94export function shortModel(id: string | undefined): string | undefined {
95  if (!id) return undefined
96  const m = id.toLowerCase()
97  for (const family of ['opus', 'sonnet', 'haiku', 'fable']) if (m.includes(family)) return family
98  return id
99}
100
101const EFFORT_LABELS: Record<string, string> = { low: 'low', medium: 'medium', high: 'high', xhigh: 'x-high', max: 'max' }
102
103export function effortLabel(e: string | undefined): string | undefined {
104  if (!e) return undefined
105  return EFFORT_LABELS[e] ?? e
106}
107
108/** Local wall-clock time, `18:20`. */
109export function clock(epochS: number): string {
110  const d = new Date(epochS * 1000)
111  return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}`
112}
113
114export function limitLabel(kind: string): string {
115  if (kind === 'five_hour') return '5h'
116  if (kind === 'seven_day') return 'wk'
117  if (kind === 'spend_limit') return '$'
118  return kind
119}
120
121// ---------------------------------------------------------------- the view
122
123export type View = {
124  ctx: { pct?: number; unknown: boolean; implausible: boolean; color?: string }
125  cache: CacheState & { color?: string; ttl: number }
126  /** `resetsAt` only for the 5-hour window, and only while it lies ahead. */
127  limits: { label: string; pct: number; color?: string; resetsAt?: number }[]
128  modelEffort?: string
129  /** Set while the prompt box holds text and the cache has expired. */
130  cold?: { tokens?: number }
131}
132
133export function makeView(r: DrainReading, nowS: number, drafting: boolean): View {
134  const implausible = isImplausible(r)
135  const ctxColor = contextColor(r.percent, r.unknown, implausible)
136  const cache = cacheState(r, nowS)
137  const model = shortModel(r.model)
138  const effort = effortLabel(r.effort)
139  return {
140    ctx: { pct: r.percent, unknown: r.unknown, implausible, color: ctxColor },
141    cache: { ...cache, color: cacheColor(cache, ctxColor), ttl: r.ttl },
142    limits: r.rateLimits.map(l => ({
143      label: limitLabel(l.kind),
144      pct: l.pct,
145      color: limitColor(l.pct),
146      resetsAt: l.kind === 'five_hour' && l.resetsAt !== undefined && l.resetsAt > nowS ? l.resetsAt : undefined,
147    })),
148    modelEffort: model === undefined ? undefined : effort === undefined ? model : `${model}·${effort}`,
149    cold: drafting && cache.kind === 'expired' ? { tokens: r.tokens } : undefined,
150  }
151}
152
153// ---------------------------------------------------------------- rows
154
155export type Seg = { text: string; color?: string; dim?: boolean; bold?: boolean }
156export type Row = Seg[]
157
158const WIDE = new Set(['🧠'])
159
160/** Terminal cells a string takes; the brain glyph is a wide character. */
161export function cellWidth(s: string): number {
162  let n = 0
163  for (const ch of s) n += WIDE.has(ch) ? 2 : 1
164  return n
165}
166
167export function plain(row: Row): string {
168  return row.map(s => s.text).join('')
169}
170
171function pctText(v: View): string {
172  if (v.ctx.pct === undefined) return '–%'
173  const n = Math.round(v.ctx.pct)
174  if (v.ctx.implausible) return `⚠ ${n}%`
175  return v.ctx.unknown ? `~${n}%` : `${n}%`
176}
177
178function cacheText(v: View): string | undefined {
179  if (v.cache.kind === 'none') return undefined
180  return v.cache.kind === 'expired' ? 'expired' : fmtCountdown(v.cache.left)
181}
182
183const SEP: Seg = { text: ' │ ', dim: true }
184
185function joinGroups(groups: Row[]): Row {
186  const out: Row = []
187  for (const g of groups) {
188    if (g.length === 0) continue
189    if (out.length > 0) out.push(SEP)
190    out.push(...g)
191  }
192  return out
193}
194
195function limitsGroup(v: View, glue: string, withPct: boolean, withReset = false): Row {
196  const g: Row = []
197  v.limits.forEach((l, i) => {
198    if (i > 0) g.push({ text: glue, dim: true })
199    g.push({ text: `${l.label} `, dim: true }, { text: `${Math.round(l.pct)}${withPct ? '%' : ''}`, color: l.color })
200    if (withReset && l.resetsAt !== undefined) g.push({ text: ` ↻${clock(l.resetsAt)}`, dim: true })
201  })
202  return g
203}
204
205export function bar(fraction: number, cells: number, color: string | undefined): Row {
206  const f = Math.max(0, Math.min(1, fraction))
207  const filled = Math.round(f * cells)
208  const row: Row = []
209  if (filled > 0) row.push({ text: '▰'.repeat(filled), color })
210  if (cells - filled > 0) row.push({ text: '▱'.repeat(cells - filled), dim: true })
211  return row
212}
213
214function minimalRow(v: View, withModel: boolean, withReset: boolean): Row {
215  const groups: Row[] = [[{ text: '🧠 ' }, { text: pctText(v), color: v.ctx.color, bold: true }]]
216  const c = cacheText(v)
217  if (c) groups.push([{ text: '◷ ', dim: true }, { text: c, color: v.cache.color }])
218  groups.push(limitsGroup(v, ' · ', true, withReset))
219  if (withModel && v.modelEffort) groups.push([{ text: v.modelEffort, dim: true }])
220  return joinGroups(groups)
221}
222
223function barsRow(v: View, withModel: boolean, wide: boolean): Row {
224  const ctxCells = wide ? 10 : 5
225  const groups: Row[] = [[
226    { text: 'ctx ', dim: true },
227    ...bar((v.ctx.pct ?? 0) / 100, ctxCells, v.ctx.color),
228    { text: ` ${pctText(v)}`, color: v.ctx.color, bold: true },
229  ]]
230  const c = cacheText(v)
231  if (c) {
232    const left = v.cache.kind === 'live' ? v.cache.left : 0
233    groups.push([
234      { text: 'cache ', dim: true },
235      ...bar(left / Math.max(left, v.cache.ttl, 1), ctxCells, v.cache.color),
236      { text: ` ${c}`, color: v.cache.color },
237    ])
238  }
239  const lim: Row = []
240  v.limits.forEach((l, i) => {
241    if (i > 0) lim.push({ text: '  ' })
242    lim.push({ text: `${l.label} `, dim: true }, ...bar(l.pct / 100, wide ? 5 : 3, l.color), {
243      text: ` ${Math.round(l.pct)}%`,
244      color: l.color,
245    })
246  })
247  groups.push(lim)
248  if (withModel && v.modelEffort) groups.push([{ text: v.modelEffort, dim: true }])
249  return joinGroups(groups)
250}
251
252function gaugeRow(v: View, withModel: boolean, track: number): Row {
253  const left = Math.max(0, Math.min(100, 100 - (v.ctx.pct ?? 0)))
254  const pos = Math.round((left / 100) * (track - 1))
255  const needle: Row = [{ text: 'E ', dim: true }]
256  if (pos > 0) needle.push({ text: '━'.repeat(pos), color: v.ctx.color })
257  needle.push({ text: '◆', color: v.ctx.color, bold: true })
258  if (track - 1 - pos > 0) needle.push({ text: '─'.repeat(track - 1 - pos), dim: true })
259  needle.push({ text: ' F ', dim: true })
260  const leftText = v.ctx.pct === undefined ? '–' : v.ctx.unknown ? `~${Math.round(left)}%` : `${Math.round(left)}%`
261  needle.push({ text: v.ctx.implausible ? '⚠ ?' : leftText, color: v.ctx.color, bold: true }, { text: ' left', dim: true })
262  const groups: Row[] = [needle]
263  const c = cacheText(v)
264  if (c) groups.push([{ text: '◷ ', dim: true }, { text: c, color: v.cache.color }])
265  groups.push(limitsGroup(v, ' · ', true))
266  if (withModel && v.modelEffort) groups.push([{ text: v.modelEffort, dim: true }])
267  return joinGroups(groups)
268}
269
270function compactRow(v: View, withLimits: boolean): Row {
271  const row: Row = [{ text: '🧠' }, { text: pctText(v), color: v.ctx.color, bold: true }]
272  const c = cacheText(v)
273  if (c) row.push({ text: ' ◷', dim: true }, { text: v.cache.kind === 'expired' ? 'exp' : c, color: v.cache.color })
274  if (withLimits) {
275    for (const l of v.limits) {
276      row.push({ text: ` ${l.label}`, dim: true }, { text: String(Math.round(l.pct)), color: l.color })
277    }
278  }
279  return row
280}
281
282function coldRows(v: View): Row[] {
283  if (!v.cold) return []
284  const amount = v.cold.tokens === undefined ? 'the whole context' : `~${fmtTokens(v.cold.tokens)} tokens`
285  return [
286    [{ text: '▲ cache expired', color: BAD, bold: true }, { text: ` · your next message re-writes ${amount}`, color: BAD }],
287    [{ text: '▲ cold', color: BAD, bold: true }, { text: ` · next msg re-writes ${v.cold.tokens === undefined ? 'all' : `~${fmtTokens(v.cold.tokens)}`}`, color: BAD }],
288  ]
289}
290
291/** The gauge row's candidates for a style, widest first. */
292function candidates(style: DrainStyle, v: View): Row[] {
293  switch (style) {
294    case 'minimal':
295      // The model goes first when space is short, then the 5-hour reset time.
296      return [minimalRow(v, true, true), minimalRow(v, false, true), minimalRow(v, false, false), compactRow(v, true), compactRow(v, false)]
297    case 'bars':
298      return [barsRow(v, true, true), barsRow(v, false, true), barsRow(v, false, false), compactRow(v, true), compactRow(v, false)]
299    case 'gauge':
300      return [gaugeRow(v, true, 16), gaugeRow(v, false, 16), gaugeRow(v, false, 8), compactRow(v, true), compactRow(v, false)]
301    case 'compact':
302      return [compactRow(v, true), compactRow(v, false)]
303    default:
304      return []
305  }
306}
307
308function fit(rows: Row[], cols: number | undefined): Row | undefined {
309  if (rows.length === 0) return undefined
310  if (cols === undefined) return rows[0]
311  return rows.find(r => cellWidth(plain(r)) <= cols) ?? rows[rows.length - 1]
312}
313
314/**
315 * The band's rows for a style: the gauge row, and under it the cold-cache
316 * warning while one applies. Narrower forms are picked when `cols` is tight.
317 */
318export function renderRows(style: DrainStyle, v: View, cols?: number): Row[] {
319  const main = fit(candidates(style, v), cols)
320  if (!main) return []
321  const out: Row[] = [main]
322  const cold = coldRows(v)
323  if (cold.length > 0) {
324    const pick = style === 'compact' ? cold[1] : fit(cold, cols)
325    if (pick) out.push(pick)
326  }
327  return out
328}
329
330/** The `line` style: one plain line for `$.ui.status`, which takes no colour. */
331export function lineText(v: View): string {
332  const parts = [`🧠 ${pctText(v)}`]
333  const c = cacheText(v)
334  if (c) parts.push(v.cache.kind === 'expired' ? 'cache expired' : `cache ${c}`)
335  for (const l of v.limits) parts.push(`${l.label} ${Math.round(l.pct)}%`)
336  if (v.modelEffort) parts.push(v.modelEffort)
337  let text = parts.join(' · ')
338  if (v.cold) {
339    const amount = v.cold.tokens === undefined ? 'the whole context' : `~${fmtTokens(v.cold.tokens)} tokens`
340    text += ` · next message re-writes ${amount}`
341  }
342  return text
343}
344
345/** Text form of a style, for `/drain styles`. */
346export function styleText(style: DrainStyle, v: View): string {
347  if (style === 'line') return lineText(v)
348  if (style === 'off') return '(nothing drawn; the session file is still written)'
349  return renderRows(style, v).map(plain).join('\n')
350}
351
352// ---------------------------------------------------------------- handover
353
354/** True while the effective status line is the bridge script, so the bridge owns the session file. */
355export function bridgeLive(settings: unknown): boolean {
356  if (typeof settings !== 'object' || settings === null) return false
357  const sl = (settings as { statusLine?: unknown }).statusLine
358  if (typeof sl !== 'object' || sl === null) return false
359  const cmd = (sl as { command?: unknown }).command
360  return typeof cmd === 'string' && cmd.includes('context-bridge.sh')
361}
362
363export const SHADOW_DIR = '.claude/fuel-gauge-mod'
364export const PRIMARY_DIR = '.claude/fuel-gauge'
365
366export function outDirFor(home: string, live: boolean): string {
367  return `${home}/${live ? SHADOW_DIR : PRIMARY_DIR}`
368}
369
370// ---------------------------------------------------------------- TTL from the transcript
371
372/**
373 * The cache TTL the most recent main-thread cache write used, read from the end
374 * of a transcript. `cutFirst` drops the first line, which a read from the
375 * middle of the file cuts in half. Mirrors the bridge's ttl_from_transcript.
376 */
377export function ttlFromTail(text: string, cutFirst: boolean): { ttl: number; atS: number } | undefined {
378  const lines = text.split('\n')
379  if (cutFirst && lines.length > 0) lines.shift()
380  for (let i = lines.length - 1; i >= 0; i--) {
381    const line = lines[i]
382    if (!line || !line.includes('ephemeral_')) continue
383    let row: unknown
384    try {
385      row = JSON.parse(line)
386    } catch {
387      continue
388    }
389    if (typeof row !== 'object' || row === null || Array.isArray(row)) continue
390    const r = row as { isSidechain?: unknown; message?: unknown; timestamp?: unknown }
391    if (r.isSidechain) continue
392    const msg = r.message as { usage?: { cache_creation?: Record<string, unknown> } } | undefined
393    const tiers = msg && typeof msg === 'object' ? msg.usage?.cache_creation : undefined
394    if (typeof tiers !== 'object' || tiers === null) continue
395    const at = typeof r.timestamp === 'string' ? Date.parse(r.timestamp) : NaN
396    const atS = Number.isFinite(at) ? Math.floor(at / 1000) : 0
397    if (Number(tiers.ephemeral_1h_input_tokens ?? 0) > 0) return { ttl: 3600, atS }
398    if (Number(tiers.ephemeral_5m_input_tokens ?? 0) > 0) return { ttl: 300, atS }
399  }
400  return undefined
401}
402
403/** The bridge's account_ttl.json rule: write only a new value observed no earlier than the stored one. */
404export function shouldWriteAccountTtl(
405  stored: { detected_ttl?: number; detected_at?: number } | undefined,
406  ttl: number,
407  atS: number,
408): boolean {
409  if (ttl <= 0 || atS <= 0) return false
410  const prevTtl = stored?.detected_ttl ?? 0
411  const prevAt = stored?.detected_at ?? 0
412  return ttl !== prevTtl && atS >= prevAt
413}
414
415// ---------------------------------------------------------------- misses
416
417export type StepUsage = { input_tokens: number; cache_read_input_tokens: number; cache_creation_input_tokens: number }
418type CacheUsage = Pick<StepUsage, 'cache_read_input_tokens' | 'cache_creation_input_tokens'>
419
420/**
421 * The context fill one request reports: the three input counters the status line
422 * sums, over the same raw window, rounded as the bridge's figure is. Set on every
423 * request so a long turn does not show its first request's fill until the turn
424 * ends; session.measure still has the last word at turn end (#4986).
425 */
426export function stepFill(usage: StepUsage, window: number | undefined): { tokens: number; percent?: number; unknown: false } {
427  const tokens = usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
428  if (window === undefined || window <= 0) return { tokens, unknown: false }
429  return { tokens, percent: Math.round(tokens / window * 100), unknown: false }
430}
431
432/**
433 * Classifies one main-thread request. `first`: nothing earlier was seen in this
434 * process (a new or resumed session), never a miss. `miss`: the gap since the
435 * last request outran the TTL and the request re-wrote more than it read.
436 */
437export function classifyStep(
438  prev: Pick<DrainReading, 'lastApiS' | 'observed' | 'ttl'>,
439  startS: number,
440  usage: CacheUsage,
441): 'first' | 'warm' | 'miss' {
442  if (!prev.observed || prev.lastApiS === undefined || prev.lastApiS <= 0) return 'first'
443  const gap = startS - prev.lastApiS
444  const wrote = usage.cache_creation_input_tokens
445  const readTokens = usage.cache_read_input_tokens
446  if (gap > prev.ttl && wrote >= MIN_MISS_TOKENS && readTokens < wrote) return 'miss'
447  return 'warm'
448}
449
450/**
451 * The bridge's fallback upgrade, made stricter: a cache that served most of a
452 * request after more than five minutes cannot be on the five-minute tier.
453 * Only when the transcript gave no answer.
454 */
455export function warmGapUpgrade(prev: Pick<DrainReading, 'lastApiS' | 'observed' | 'ttl' | 'ttlSource'>, startS: number, usage: CacheUsage): boolean {
456  if (prev.ttlSource === 'transcript' || prev.ttl >= 3600) return false
457  if (!prev.observed || prev.lastApiS === undefined) return false
458  return startS - prev.lastApiS > 310 && usage.cache_read_input_tokens > usage.cache_creation_input_tokens
459}
460
461export function missKey(atMs: number, session: string): string {
462  return `miss:${Math.round(atMs)}:${session.slice(0, 8)}`
463}
464
465function startOfLocalDayS(nowS: number): number {
466  const d = new Date(nowS * 1000)
467  d.setHours(0, 0, 0, 0)
468  return Math.floor(d.getTime() / 1000)
469}
470
471function tally(list: DrainMiss[]): string {
472  if (list.length === 0) return 'no misses'
473  const tokens = list.reduce((n, m) => n + m.tokens, 0)
474  const priced = list.filter(m => typeof m.costUsd === 'number')
475  const usd = priced.reduce((n, m) => n + (m.costUsd ?? 0), 0)
476  const cost = priced.length === 0 ? 'no measured cost' : `$${usd.toFixed(2)} measured (${priced.length} of ${list.length})`
477  return `${list.length} miss${list.length === 1 ? '' : 'es'} · ${fmtTokens(tokens)} tokens re-written · ${cost}`
478}
479
480export function ledgerReport(misses: DrainMiss[], nowS: number): string[] {
481  const today = startOfLocalDayS(nowS)
482  const week = nowS - 7 * 86400
483  const lines = [
484    `Today:        ${tally(misses.filter(m => m.at >= today))}`,
485    `Last 7 days:  ${tally(misses.filter(m => m.at >= week))}`,
486  ]
487  const last = [...misses].sort((a, b) => b.at - a.at)[0]
488  if (last) {
489    const when = new Date(last.at * 1000)
490    const stamp = `${when.toDateString().slice(4, 10)} ${String(when.getHours()).padStart(2, '0')}:${String(when.getMinutes()).padStart(2, '0')}`
491    lines.push(
492      `Last miss:    ${stamp} in ${last.project || 'a session'}, ${fmtTokens(last.tokens)} tokens after ${fmtDuration(last.gapS)} idle (TTL ${last.ttl >= 3600 ? '1h' : '5m'})`,
493    )
494  }
495  return lines
496}
497
498// ---------------------------------------------------------------- the session file
499
500export type Ids = { sessionId: string; cwd: string; projectDir: string }
501
502function limitFields(limits: DrainLimit[], kind: string): { pct: number; reset: number } {
503  const l = limits.find(x => x.kind === kind)
504  return { pct: l ? l.pct : -1, reset: l?.resetsAt ?? 0 }
505}
506
507/** The bridge's session file, field for field, plus `writer: "mod"`. */
508export function sessionFile(r: DrainReading, ids: Ids, nowS: number): Record<string, unknown> {
509  const used = r.percent ?? 0
510  const h5 = limitFields(r.rateLimits, 'five_hour')
511  const d7 = limitFields(r.rateLimits, 'seven_day')
512  return {
513    cwd: ids.cwd,
514    project_dir: ids.projectDir,
515    used_percentage: used,
516    remaining_percentage: Math.max(0, 100 - used),
517    context_unknown: r.unknown,
518    window_implausible: isImplausible(r),
519    total_input_tokens: r.tokens ?? 0,
520    total_output_tokens: r.outputTokens,
521    context_window_size: r.window ?? 200000,
522    model: r.model ?? 'unknown',
523    session_id: ids.sessionId,
524    timestamp: nowS,
525    cache_read_input_tokens: r.cacheRead,
526    cache_creation_input_tokens: r.cacheCreate,
527    detected_ttl: r.ttl,
528    last_api_call_ts: r.lastApiS ?? 0,
529    rate_limit_5h_pct: h5.pct,
530    rate_limit_5h_reset: h5.reset,
531    rate_limit_7d_pct: d7.pct,
532    rate_limit_7d_reset: d7.reset,
533    effort_level: r.effort ?? '',
534    writer: 'mod',
535  }
536}
537
538export function emptyReading(ttl = DEFAULT_TTL, source: DrainReading['ttlSource'] = 'default'): DrainReading {
539  return { unknown: false, rateLimits: [], observed: false, ttl, ttlSource: source, cacheRead: 0, cacheCreate: 0, outputTokens: 0 }
540}
541
542/** A new session in the same process (/clear): the account's TTL, limits and the model carry over, nothing else. */
543export function freshSession(prev: DrainReading): DrainReading {
544  return { ...emptyReading(prev.ttl, prev.ttlSource), rateLimits: prev.rateLimits, model: prev.model, effort: prev.effort, window: prev.window }
545}
546
547/** A reading restored from this mod's own earlier file (a resume or a fresh process). */
548export function readingFromFile(file: unknown, fallbackTtl: number, fallbackSource: DrainReading['ttlSource']): DrainReading | undefined {
549  if (typeof file !== 'object' || file === null) return undefined
550  const f = file as Record<string, unknown>
551  if (f.writer !== 'mod') return undefined
552  const num = (k: string): number | undefined => (typeof f[k] === 'number' ? (f[k] as number) : undefined)
553  const ttl = num('detected_ttl')
554  const r = emptyReading(ttl && ttl > 0 ? ttl : fallbackTtl, ttl && ttl > 0 ? 'previous' : fallbackSource)
555  const last = num('last_api_call_ts')
556  r.lastApiS = last && last > 0 ? last : undefined
557  r.percent = num('used_percentage')
558  r.window = num('context_window_size')
559  r.model = typeof f.model === 'string' && f.model !== 'unknown' ? f.model : undefined
560  r.effort = typeof f.effort_level === 'string' && f.effort_level !== '' ? f.effort_level : undefined
561  return r
562}
563
564export function limitsFrom(list: readonly { kind: string; percentUsed: number; resetsAt?: string }[]): DrainLimit[] {
565  return list.map(l => {
566    const t = l.resetsAt ? Date.parse(l.resetsAt) : NaN
567    return Number.isFinite(t) ? { kind: l.kind, pct: l.percentUsed, resetsAt: Math.floor(t / 1000) } : { kind: l.kind, pct: l.percentUsed }
568  })
569}
570
571export const HELP = [
572  '/drain                cache-miss report: today, last 7 days, tokens re-written, measured cost',
573  '/drain styles         preview every display style with the current figures',
574  '/drain <style>        switch style: minimal, bars, gauge, compact, line, off',
575  '                      (/drain style <name> and /drain styles <name> work too)',
576  '/drain place below|above  draw under the prompt (default) or in the band above it',
577  '/drain warn <seconds>  how long before expiry the alert fires (default 60, 0 turns it off)',
578  '/drain sound on|off   play a short chime with the alert (off by default)',
579  '/drain agents on|off  prototype: one token bar per subagent under the gauge (off by default)',
580  '/drain agents         the subagent rows as text',
581  '/drain notify on|off  self-notice: one line of facts to the model when a threshold is crossed (off by default)',
582  '/drain notify         the self-notice settings and the note as it would read now',
583  '/drain notify context 70 90 | 5h 80 95 | week on|off | margin 5   thresholds (each list replaces the last; off clears)',
584  '/drain notify text context 35 "your text"   your own words after the facts when that threshold fires',
585  '                      (text 5h 80 "...", text week "..."; no text clears it)',
586  '                      BRAIN_DRAIN_NOTIFY in the environment overrides these for one run (see the README)',
587  '/drain help           this list',
588]
589

1import type { DrainLimit, DrainNotify } from '../types'
2import { clock, fmtTokens } from './core.ts'
3
4// The self-notice: a one-line note of facts handed to the main session's model
5// when a figure crosses a threshold. Pure: register.tsx reads the figures, keeps
6// which thresholds have fired, and attaches the text.
7
8export const DEFAULT_NOTIFY: DrainNotify = { on: false, context: [70, 90], fiveHour: [80, 95], week: true, margin: 5, text: {} }
9
10export const ENV_NAME = 'BRAIN_DRAIN_NOTIFY'
11const MAX_LINES = 10
12const MAX_TEXT = 2000
13
14export type Figures = {
15  context?: { pct: number; tokens: number; window: number }
16  fiveHour?: { pct: number; resetsAt?: number }
17  week?: { pct: number; resetsAt?: number }
18}
19
20/** The figures as the engine reports them. A context with no token count (fresh, or just compacted) is not known. */
21export function figures(context: { tokens?: number; window: number; percent?: number }, limits: DrainLimit[]): Figures {
22  const f: Figures = {}
23  if (context.tokens !== undefined && context.window > 0) {
24    f.context = { pct: context.percent ?? (context.tokens / context.window) * 100, tokens: context.tokens, window: context.window }
25  }
26  const five = limits.find(l => l.kind === 'five_hour')
27  if (five) f.fiveHour = { pct: five.pct, resetsAt: five.resetsAt }
28  const week = limits.find(l => l.kind === 'seven_day')
29  if (week) f.week = { pct: week.pct, resetsAt: week.resetsAt }
30  return f
31}
32
33/** The share of the week gone since the weekly reset, as a percentage, plus the margin. */
34export function paceLine(resetsAtS: number, nowS: number, margin: number): number {
35  const daysSince = Math.min(7, Math.max(0, 7 - (resetsAtS - nowS) / 86400))
36  return (daysSince / 7) * 100 + margin
37}
38
39/** Each configured threshold with whether it stands crossed now; undefined where its figure is not known. */
40function states(f: Figures, n: DrainNotify, nowS: number): { key: string; crossed?: boolean }[] {
41  const out: { key: string; crossed?: boolean }[] = []
42  for (const t of n.context) out.push({ key: `context:${t}`, crossed: f.context ? f.context.pct >= t : undefined })
43  for (const t of n.fiveHour) out.push({ key: `5h:${t}`, crossed: f.fiveHour ? f.fiveHour.pct >= t : undefined })
44  if (n.week) {
45    const w = f.week
46    out.push({ key: 'week', crossed: w && w.resetsAt !== undefined ? w.pct > paceLine(w.resetsAt, nowS, n.margin) : undefined })
47  }
48  return out
49}
50
51/**
52 * Which thresholds are newly crossed, and the fired set after this reading:
53 * a crossed threshold fires once; one whose figure fell back below re-arms; one
54 * whose figure is not known keeps its state.
55 */
56export function evaluate(f: Figures, n: DrainNotify, fired: readonly string[], nowS: number): { crossed: string[]; fired: string[] } {
57  const was = new Set(fired)
58  const now = new Set(fired)
59  const crossed: string[] = []
60  for (const s of states(f, n, nowS)) {
61    if (s.crossed === undefined) continue
62    if (s.crossed && !was.has(s.key)) {
63      crossed.push(s.key)
64      now.add(s.key)
65    } else if (!s.crossed) {
66      now.delete(s.key)
67    }
68  }
69  return { crossed, fired: [...now] }
70}
71
72/** The facts, one line: `Brain Drain: context 71% (712k of 1M) · 5-hour window 82%, resets 18:20 · week 40% (pace line 43%)`. */
73export function factsLine(f: Figures, n: DrainNotify, nowS: number): string {
74  const parts: string[] = []
75  if (f.context) parts.push(`context ${Math.round(f.context.pct)}% (${fmtTokens(f.context.tokens)} of ${fmtTokens(f.context.window)})`)
76  if (f.fiveHour) parts.push(`5-hour window ${Math.round(f.fiveHour.pct)}%${f.fiveHour.resetsAt !== undefined ? `, resets ${clock(f.fiveHour.resetsAt)}` : ''}`)
77  if (f.week) {
78    const pace = f.week.resetsAt !== undefined ? ` (pace line ${Math.round(paceLine(f.week.resetsAt, nowS, n.margin))}%)` : ''
79    parts.push(`week ${Math.round(f.week.pct)}%${pace}`)
80  }
81  return `Brain Drain: ${parts.length > 0 ? parts.join(' · ') : 'no figures yet'}`
82}
83
84/** The note as the model receives it: the facts, then the user's own text for each threshold just crossed, one per line. */
85export function noticeText(f: Figures, n: DrainNotify, crossed: readonly string[], nowS: number): string {
86  const extra = crossed.map(k => n.text[k]).filter((t): t is string => typeof t === 'string' && t.trim() !== '')
87  return [factsLine(f, n, nowS), ...extra].join('\n')
88}
89
90// ---------------------------------------------------------------- settings
91
92function lines(v: unknown): number[] | undefined {
93  if (!Array.isArray(v) || v.length > MAX_LINES) return undefined
94  if (!v.every(x => typeof x === 'number' && Number.isFinite(x) && x >= 0 && x <= 100)) return undefined
95  return [...new Set(v as number[])].sort((a, b) => a - b)
96}
97
98/** A text key names a threshold kind and, for context and 5h, its value: `context:35`, `5h:80`, `week`. */
99export function isTextKey(k: string): boolean {
100  return k === 'week' || /^(context|5h):\d+(\.\d+)?$/.test(k)
101}
102
103function texts(v: unknown): Record<string, string> | undefined {
104  if (typeof v !== 'object' || v === null || Array.isArray(v)) return undefined
105  const out: Record<string, string> = {}
106  for (const [k, t] of Object.entries(v)) {
107    if (!isTextKey(k) || typeof t !== 'string' || t.length > MAX_TEXT) return undefined
108    if (t.trim() !== '') out[k] = t
109  }
110  return out
111}
112
113/** Saved prefs, field by field: a field of the wrong shape falls back to its default. */
114export function sanitize(v: unknown): DrainNotify {
115  const o = typeof v === 'object' && v !== null ? (v as Record<string, unknown>) : {}
116  return {
117    on: o.on === true,
118    context: lines(o.context) ?? DEFAULT_NOTIFY.context,
119    fiveHour: lines(o.fiveHour) ?? DEFAULT_NOTIFY.fiveHour,
120    week: typeof o.week === 'boolean' ? o.week : DEFAULT_NOTIFY.week,
121    margin: typeof o.margin === 'number' && Number.isFinite(o.margin) && Math.abs(o.margin) <= 100 ? o.margin : DEFAULT_NOTIFY.margin,
122    text: texts(o.text) ?? {},
123  }
124}
125
126/**
127 * `BRAIN_DRAIN_NOTIFY`, the per-run override for launchers: `on`, `off`, or a
128 * JSON object with any of `on`, `context`, `5h`, `week`, `margin`, `text`.
129 * Each key given replaces the saved value for this run; a JSON value turns the
130 * feature on unless it says `"on": false`. Anything else is an error, and the
131 * saved prefs stand.
132 */
133export function parseEnv(raw: string | undefined): { value?: Partial<DrainNotify>; error?: string } {
134  if (raw === undefined || raw.trim() === '') return {}
135  const s = raw.trim()
136  if (s === 'on' || s === 'off') return { value: { on: s === 'on' } }
137  let o: unknown
138  try {
139    o = JSON.parse(s)
140  } catch {
141    return { error: `${ENV_NAME} is neither on, off nor JSON` }
142  }
143  if (typeof o !== 'object' || o === null || Array.isArray(o)) return { error: `${ENV_NAME} must be a JSON object` }
144  const src = o as Record<string, unknown>
145  const out: Partial<DrainNotify> = { on: true }
146  for (const k of Object.keys(src)) {
147    if (!['on', 'context', '5h', 'week', 'margin', 'text'].includes(k)) return { error: `${ENV_NAME}: unknown key "${k}"` }
148  }
149  if ('on' in src) {
150    if (typeof src.on !== 'boolean') return { error: `${ENV_NAME}: "on" must be true or false` }
151    out.on = src.on
152  }
153  if ('context' in src) {
154    const l = lines(src.context)
155    if (!l) return { error: `${ENV_NAME}: "context" must be a list of numbers from 0 to 100` }
156    out.context = l
157  }
158  if ('5h' in src) {
159    const l = lines(src['5h'])
160    if (!l) return { error: `${ENV_NAME}: "5h" must be a list of numbers from 0 to 100` }
161    out.fiveHour = l
162  }
163  if ('week' in src) {
164    if (typeof src.week !== 'boolean') return { error: `${ENV_NAME}: "week" must be true or false` }
165    out.week = src.week
166  }
167  if ('margin' in src) {
168    if (typeof src.margin !== 'number' || !Number.isFinite(src.margin) || Math.abs(src.margin) > 100) return { error: `${ENV_NAME}: "margin" must be a number from -100 to 100` }
169    out.margin = src.margin
170  }
171  if ('text' in src) {
172    const t = texts(src.text)
173    if (!t) return { error: `${ENV_NAME}: "text" must map keys like "context:35", "5h:80" or "week" to strings` }
174    out.text = t
175  }
176  return { value: out }
177}
178
179/** The settings in force: the saved ones with the run's override laid over them. */
180export function effective(saved: DrainNotify, override: Partial<DrainNotify> | undefined): DrainNotify {
181  return override ? { ...saved, ...override } : saved
182}
183
184/** `/drain notify` with no arguments: the settings in force and where they come from. */
185export function describe(n: DrainNotify, fromEnv: boolean, envError: string | undefined, fired: readonly string[], now: string): string[] {
186  const out = [`Self-notice ${n.on ? 'on' : 'off'}${fromEnv ? ` (this run: ${ENV_NAME} overrides the saved settings)` : ''}`]
187  if (envError) out.push(`${envError}; the saved settings stand.`)
188  out.push(`  context  ${n.context.length > 0 ? n.context.map(t => `${t}%`).join(', ') : 'off'}`)
189  out.push(`  5h       ${n.fiveHour.length > 0 ? n.fiveHour.map(t => `${t}%`).join(', ') : 'off'}`)
190  out.push(`  week     ${n.week ? `above the pace line (days since the weekly reset / 7 x 100, plus ${n.margin} points)` : 'off'}`)
191  const keys = Object.keys(n.text)
192  if (keys.length > 0) {
193    out.push('  your text:')
194    for (const k of keys) out.push(`    ${k.replace(':', ' ')}: ${n.text[k]}`)
195  }
196  out.push(`  fired this session: ${fired.length > 0 ? fired.map(k => k.replace(':', ' ')).join(', ') : 'none'}`)
197  out.push(`The note reads, now: ${now}`)
198  return out
199}
200

1export type DrainStyle = 'minimal' | 'bars' | 'gauge' | 'compact' | 'line' | 'off'
2
3export type DrainLimit = { kind: string; pct: number; resetsAt?: number }
4
5export type DrainReading = {
6  percent?: number
7  tokens?: number
8  window?: number
9  /** Context fill not known since a compact or a resume: the last figure is held. */
10  unknown: boolean
11  rateLimits: DrainLimit[]
12  model?: string
13  effort?: string
14  /** Epoch seconds of the last main-thread API response. */
15  lastApiS?: number
16  /** True once this process saw a request of its own; miss detection needs it. */
17  observed: boolean
18  ttl: number
19  ttlSource: 'transcript' | 'previous' | 'account' | 'default' | 'warm-gap'
20  cacheRead: number
21  cacheCreate: number
22  outputTokens: number
23  costUsd?: number
24}
25
26export type DrainPlace = 'below' | 'above'
27
28export type DrainPrefs = {
29  style: DrainStyle
30  place: DrainPlace
31  warnSeconds: number
32  sound: boolean
33  /** Prototype: per-subagent token bars under the gauge. Off unless asked for. */
34  agents: boolean
35  /** The self-notice: a facts-only note to the main session's model when a threshold is crossed. Off unless asked for. */
36  notify: DrainNotify
37}
38
39export type DrainNotify = {
40  on: boolean
41  /** Context fill thresholds, percent. */
42  context: number[]
43  /** Five-hour window thresholds, percent. */
44  fiveHour: number[]
45  /** Weekly utilisation above the pace line: days since the weekly reset / 7 x 100, plus `margin` points. */
46  week: boolean
47  margin: number
48  /** The user's own text, appended after the facts when that threshold fires: `context:35`, `5h:80`, `week`. */
49  text: Record<string, string>
50}
51
52/** One subagent of this session, as its transcript adds up. */
53export type DrainAgent = {
54  /** The file's stem after `agent-`, as the engine names the agent. */
55  id: string
56  label: string
57  /** Every request's input + cache write + cache read + output, each response counted once. */
58  tokens: number
59  output: number
60  requests: number
61  model?: string
62  effort?: string
63  /** The transcript's last write, epoch milliseconds. */
64  mtimeMs: number
65  /** The engine's word for the agent (`$.agent.list()`), when it lists it. */
66  status?: string
67}
68
69export type DrainAgents = { dir: string; agents: DrainAgent[] }
70
71export type DrainMeta = {
72  /** False under `claude -p` and the SDK: no session file then, as the bridge writes none. */
73  interactive?: boolean
74  /** The session the reading belongs to; /clear, /resume and a fork change it. */
75  sessionId?: string
76  home?: string
77  outDir?: string
78  shadow?: boolean
79  transcriptPath?: string
80  alertedFor?: number
81  /** Self-notice thresholds fired in this session and not re-armed: `context:70`, `5h:80`, `week`. */
82  notified?: string[]
83  /** This run's `BRAIN_DRAIN_NOTIFY`, parsed once at setup. */
84  notifyEnv?: Partial<DrainNotify>
85  notifyEnvError?: string
86  pendingMiss?: { key: string; costBefore: number }
87}
88
89export type DrainMiss = {
90  at: number
91  session: string
92  project: string
93  tokens: number
94  gapS: number
95  ttl: number
96  costUsd?: number
97}
98
99declare module 'claude-code' {
100  interface PluginState {
101    'brain-drain': {
102      reading: DrainReading | null
103      prefs: DrainPrefs | null
104      drafting: boolean
105      /** Turns begun and not yet complete; a subagent's leaves at its first request. */
106      turns: string[] | null
107      meta: DrainMeta | null
108      agents: DrainAgents | null
109    }
110  }
111}
112