session-models

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { AgentsSeen } from '../types'
5
6const agentsSeen = atom({ plugin: 'session-models', key: 'agents' } as const, {} as AgentsSeen)
7
8export type Thresholds = { liveAgents: number; context: number }
9
10// the crowd toast re-arms only once the count drops this far below the threshold,
11// and never fires twice within the quiet period
12const CROWD_HYSTERESIS = 2
13const CROWD_QUIET_MS = 10 * 60_000
14// a busy session steps many times a second across its agents: redraw at most this often
15const REFRESH_MIN_MS = 2_000
16// the surfaces that raise the AbovePrompt band; any other gets the line as a status
17const BAND_SURFACES: readonly string[] = ['terminal', 'desktop']
18
19// claude-opus-5-5 -> Opus; an id outside the known families is shown as-is
20export function family(model: string): string {
21  const m = /(opus|sonnet|haiku|fable)/i.exec(model)
22  if (m === null) return model.replace(/^claude-/, '')
23  const name = m[1]!.toLowerCase()
24  return name[0]!.toUpperCase() + name.slice(1)
25}
26
27// 955492 -> 955K, 1200000 -> 1.2M
28export function tokens(n: number): string {
29  return n >= 1_000_000 ? `${(n / 1_000_000).toFixed(1)}M` : `${Math.round(n / 1000)}K`
30}
31
32// the agents running now: the main thread always, plus every running subagent or teammate
33export function liveAgents(agents: AgentsSeen, running: ReadonlySet<string>): string[] {
34  return Object.keys(agents).filter(id => id === 'main' || running.has(id))
35}
36
37// the live agent carrying the most context, when any carries at least `threshold`
38export function biggestContext(agents: AgentsSeen, live: readonly string[], threshold: number) {
39  let top: { id: string; context: number } | undefined
40  for (const id of live) {
41    const context = agents[id]?.context ?? 0
42    if (context >= threshold && (top === undefined || context > top.context)) top = { id, context }
43  }
44  return top
45}
46
47// "model: Opus · agents: Opus (29, 13 live), Sonnet (4) · ⚠ 14 live · ⚠ ctx 955K"
48// the agents part appears once anything beyond the main thread has run;
49// an agent that used two families counts under each
50export function formatModels(
51  current: string,
52  agents: AgentsSeen,
53  running: ReadonlySet<string>,
54  limits: Thresholds,
55): string {
56  return modelParts(current, agents, running, limits).join(' · ')
57}
58
59// the line's parts, as formatModels joins them
60export function modelParts(
61  current: string,
62  agents: AgentsSeen,
63  running: ReadonlySet<string>,
64  limits: Thresholds,
65): string[] {
66  const parts = [`model: ${family(current)}`]
67  const live = new Set(liveAgents(agents, running))
68  if (Object.keys(agents).some(id => id !== 'main')) {
69    const counts = new Map<string, { total: number; live: number }>()
70    for (const [id, seen] of Object.entries(agents)) {
71      for (const f of new Set(seen.models.map(family))) {
72        const c = counts.get(f) ?? { total: 0, live: 0 }
73        counts.set(f, { total: c.total + 1, live: c.live + (live.has(id) ? 1 : 0) })
74      }
75    }
76    const list = [...counts]
77      .sort((a, b) => b[1].live - a[1].live || b[1].total - a[1].total)
78      .map(([f, c]) => (c.live > 0 ? `${f} (${c.total}, ${c.live} live)` : `${f} (${c.total})`))
79      .join(', ')
80    parts.push(`agents: ${list}`)
81  }
82  if (live.size >= limits.liveAgents) parts.push(`⚠ ${live.size} live`)
83  const top = biggestContext(agents, [...live], limits.context)
84  if (top !== undefined) parts.push(`⚠ ctx ${tokens(top.context)}`)
85  return parts
86}
87
88// the line as runs of text: each part marked ⚠ a warning run of its own, the other
89// parts and the separators merged into dim runs
90export type Run = { text: string; isWarning: boolean }
91export function runs(parts: readonly string[]): Run[] {
92  const out: Run[] = []
93  parts.forEach((part, i) => {
94    const pieces: Run[] = i > 0 ? [{ text: ' · ', isWarning: false }] : []
95    pieces.push({ text: part, isWarning: part.startsWith('⚠') })
96    for (const piece of pieces) {
97      const last = out[out.length - 1]
98      if (last !== undefined && !last.isWarning && !piece.isWarning) last.text += piece.text
99      else out.push({ ...piece })
100    }
101  })
102  return out
103}
104
105// whether no surface the session draws on raises a band (or nothing draws at all, as
106// under -p or an SDK host), so the line goes out as a status instead; a phone or IDE
107// attached beside a terminal goes without, so the terminal keeps its quiet band alone
108export function needsStatus(surfaces: readonly string[]): boolean {
109  return !surfaces.some(s => BAND_SURFACES.includes(s))
110}
111
112export type Crowd = { isArmed: boolean; lastToastAt: number }
113
114// whether this live count earns a crowd toast, and the alarm's next state
115export function crowdAlarm(crowd: Crowd, live: number, now: number, threshold: number) {
116  if (live < threshold - CROWD_HYSTERESIS) return { toast: false, crowd: { ...crowd, isArmed: true } }
117  if (live >= threshold && crowd.isArmed && now - crowd.lastToastAt >= CROWD_QUIET_MS) {
118    return { toast: true, crowd: { isArmed: false, lastToastAt: now } }
119  }
120  return { toast: false, crowd }
121}
122
123// an agent by the task it was given, falling back to its id when it has none
124export function agentLabel(id: string, description: string | undefined): string {
125  if (id === 'main') return 'The main thread'
126  return description !== undefined && description !== '' ? description : `Agent ${id.slice(0, 8)}`
127}
128
129// a configured number, or the default when unset or not a number, held within bounds
130export function setting(value: unknown, fallback: number, min: number, max: number): number {
131  const n = Number(value ?? fallback)
132  return Number.isFinite(n) ? Math.min(max, Math.max(min, n)) : fallback
133}
134
135// module state: a reload starts these over, so a toast may repeat once after one
136let crowd: Crowd = { isArmed: true, lastToastAt: -Infinity }
137const flaggedContext = new Set<string>()
138let lastRefreshAt = -Infinity
139// the line the band draws
140let line: string[] | undefined
141
142// the band draws the line where it is raised; the status carries it where it is not
143async function show($: EngineInterface, parts: string[]) {
144  if (line === undefined || line.join(' · ') !== parts.join(' · ')) {
145    line = parts
146    $.ui.invalidate('ui.render')
147  }
148  $.ui.status(needsStatus(await $.session.surfaces()) ? parts.join(' · ') : undefined)
149}
150
151async function refresh($: EngineInterface, limits: Thresholds) {
152  const now = await $.clock.now()
153  lastRefreshAt = now
154  const agents = await read($, agentsSeen)
155  const listed = await $.agent.list()
156  const running = new Set(listed.filter(a => a.status === 'running').map(a => a.id))
157  const descriptions = new Map(listed.map(a => [a.id, a.description]))
158  await show($, modelParts(await $.session.model(), agents, running, limits))
159
160  const live = liveAgents(agents, running)
161  const alarm = crowdAlarm(crowd, live.length, now, limits.liveAgents)
162  crowd = alarm.crowd
163  if (alarm.toast) $.ui.toast(`${live.length} agents are running at once in this session`)
164  for (const id of live) {
165    const context = agents[id]?.context ?? 0
166    if (context >= limits.context && !flaggedContext.has(id)) {
167      flaggedContext.add(id)
168      $.ui.toast(`${agentLabel(id, descriptions.get(id))} is re-reading ${tokens(context)} of context on every request`)
169    }
170  }
171}
172
173export const register: Register = (on, options) => {
174  const limits: Thresholds = {
175    liveAgents: setting(options.liveAgentsWarn, 6, 2, 100),
176    context: setting(options.contextWarnK, 300, 10, 10_000) * 1000,
177  }
178
179  on('session.start', async ($, e, next) => {
180    const result = await next(e)
181    await refresh($, limits)
182    // agents finish between steps: keep the live counts honest
183    $.clock.every(15_000, () => refresh($, limits).catch(() => undefined))
184    return result
185  })
186
187  on('session.end', { reason: 'clear' }, async ($, e, next) => {
188    await update($, agentsSeen, () => ({}))
189    flaggedContext.clear()
190    crowd = { isArmed: true, lastToastAt: -Infinity }
191    const result = await next(e)
192    await refresh($, limits)
193    return result
194  })
195
196  // every model request of the main loop, a subagent or an in-process teammate
197  on('turn.step', async function* ($, e, next) {
198    const result = yield* next(e)
199    const usage = result.usage
200    // a request that got no response used no model
201    if (usage === null) return result
202    // the model that actually answered (a fallback may differ from the one asked)
203    const model = usage.model
204    const id = e.agentId ?? 'main'
205    const context = usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
206    // a new agent or model redraws at once; otherwise the throttle decides
207    const isNew = !(await read($, agentsSeen))[id]?.models.includes(model)
208    await update($, agentsSeen, agents => {
209      const seen = agents[id] ?? { models: [], context: 0 }
210      return {
211        ...agents,
212        [id]: { models: seen.models.includes(model) ? seen.models : [...seen.models, model], context },
213      }
214    })
215    if (isNew || (await $.clock.now()) - lastRefreshAt >= REFRESH_MIN_MS) await refresh($, limits)
216    return result
217  })
218
219  on('classic.PostModelSwitch', async ($, e, next) => {
220    const result = await next(e)
221    await refresh($, limits)
222    return result
223  })
224
225  // a phone or an IDE joining or leaving changes whether the status is needed
226  on('session.attach', async ($, e, next) => {
227    const result = await next(e)
228    if (line !== undefined) await show($, line)
229    return result
230  })
231  on('session.detach', async ($, e, next) => {
232    const result = await next(e)
233    if (line !== undefined) await show($, line)
234    return result
235  })
236
237  // one dim line in the band above the prompt, under what the plugins beneath drew;
238  // only the parts marked ⚠ take the warning color
239  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
240    const below = await next(e)
241    if (e.props.hasSurvey || line === undefined) return below
242    const { Box, Text } = $.ui.resolve(e)
243    return (
244      <Box flexDirection="column">
245        {below}
246        <Text wrap="truncate-end">
247          {runs(line).map(r => (r.isWarning ? <Text color="warning">{r.text}</Text> : <Text dimColor>{r.text}</Text>))}
248        </Text>
249      </Box>
250    )
251  })
252}
253

1// what the mod has seen of one agent ('main', or a subagent or teammate id)
2export type AgentSeen = {
3  // every model id it has called
4  models: string[]
5  // tokens its latest request carried: input plus cache reads and writes
6  context: number
7}
8
9export type AgentsSeen = Record<string, AgentSeen>
10
11declare module 'claude-code' {
12  interface PluginState {
13    'session-models': { agents: AgentsSeen }
14  }
15}
16