thrift

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Coach, Diet, Route, Tier, Weather } from '../types'
5import {
6  applyDiet,
7  fmtTokens,
8  isPinnedSection,
9  isReset,
10  isSameModel,
11  measureSections,
12  modelForTier,
13  parseTier,
14  contextLine,
15  routeLine,
16  savedTokens,
17} from './logic'
18
19const PANE = 'thrift-diet'
20const MAX_POINTS = 30
21
22const route = atom({ plugin: 'thrift', key: 'route' } as const, { tier: null, model: null, isApplied: false, current: null } as Route)
23const coach = atom({ plugin: 'thrift', key: 'coach' } as const, { tokens: null, crossing: 0, dismissedCrossing: 0 } as Coach)
24const weather = atom({ plugin: 'thrift', key: 'weather' } as const, { points: [], limit: null, isCompaction: false } as Weather)
25const diet = atom({ plugin: 'thrift', key: 'diet' } as const, { cwd: '', dropped: [], sections: [] } as Diet)
26
27const RUBRIC =
28  'You route a coding assistant\'s prompts to a model. Classify the person\'s message by the difficulty of the work it asks for.\n' +
29  'trivial: a lookup, a one-line edit, a rename, a yes/no, a short factual or formatting question.\n' +
30  'normal: ordinary coding work: a bug fix or feature in a few files, explaining code, writing tests.\n' +
31  'hard: design across many files, subtle debugging, security or concurrency, large refactors, open-ended research.\n' +
32  'Reply with exactly one word: trivial, normal or hard.'
33
34async function classify($: EngineInterface, text: string): Promise<Tier | null> {
35  const r = await $.model.complete({ model: 'haiku', system: RUBRIC, prompt: text.slice(0, 1500), maxTokens: 8 }).catch(() => undefined)
36  return r && r.isAnswered ? parseTier(r.text) : null
37}
38
39async function setRoute($: EngineInterface, next: Route) {
40  await update($, route, () => next)
41}
42
43async function noteContext($: EngineInterface, tokens: number, threshold: number) {
44  const prev = await read($, coach)
45  const was = threshold > 0 && (prev.tokens ?? 0) >= threshold
46  const isOver = threshold > 0 && tokens >= threshold
47  const isNewCrossing = isOver && !was
48  await update($, coach, c => ({ ...c, tokens, crossing: isNewCrossing ? c.crossing + 1 : c.crossing }))
49  if (isNewCrossing) $.ui.toast(`thrift: context passed ${fmtTokens(threshold)} tokens. The band offers a handoff or a compaction.`)
50}
51
52// The engine compacts at autoCompactThreshold when it is on; otherwise the window is the limit.
53async function readLimit($: EngineInterface): Promise<{ limit: number | null; isCompaction: boolean }> {
54  try {
55    const { context } = await $.session.usage({ breakdown: 'summary' })
56    const b = context.breakdown
57    if (b?.isAutoCompactEnabled && b.autoCompactThreshold) return { limit: b.autoCompactThreshold, isCompaction: true }
58    return { limit: b?.rawMaxTokens ?? context.window, isCompaction: false }
59  } catch {
60    return { limit: null, isCompaction: false }
61  }
62}
63
64async function recordTurn($: EngineInterface) {
65  const { context } = await $.session.usage()
66  if (context.tokens === undefined) return
67  const tokens = context.tokens
68  const prev = await read($, weather)
69  const reset = isReset(prev.points[prev.points.length - 1], tokens)
70  const limit = reset || prev.limit === null ? await readLimit($) : { limit: prev.limit, isCompaction: prev.isCompaction }
71  await update($, weather, w => ({
72    ...w,
73    points: (reset ? [tokens] : [...w.points, tokens]).slice(-MAX_POINTS),
74    limit: limit.limit,
75    isCompaction: limit.isCompaction,
76  }))
77}
78
79async function dismiss($: EngineInterface) {
80  await update($, coach, c => ({ ...c, dismissedCrossing: c.crossing }))
81}
82
83async function handOff($: EngineInterface, text: string) {
84  await dismiss($)
85  await $.prompt.submit({ text })
86}
87
88async function compactNow($: EngineInterface) {
89  try {
90    const r = await $.session.compact()
91    if (r.skip !== undefined) return $.ui.toast(`thrift: compaction skipped: ${r.skip}`)
92    await update($, coach, c => ({ ...c, tokens: null, dismissedCrossing: c.crossing }))
93  } catch {
94    $.ui.toast('thrift: wait for the turn to end, then press Compact again.')
95  }
96}
97
98async function refreshSections($: EngineInterface) {
99  const composed = await $.prompt.compose()
100  const sections = measureSections(composed.sections)
101  await update($, diet, d => ({ ...d, sections }))
102}
103
104async function toggleSection($: EngineInterface, id: string) {
105  if (isPinnedSection(id)) return
106  const d = await update($, diet, cur => ({ ...cur, dropped: cur.dropped.includes(id) ? cur.dropped.filter(x => x !== id) : [...cur.dropped, id] }))
107  await $.store.set(`diet:${d.cwd}`, d.dropped)
108  $.ui.toast('thrift: takes effect on the next request (one prompt-cache miss).')
109}
110
111async function resetAll($: EngineInterface) {
112  await update($, coach, () => ({ tokens: null, crossing: 0, dismissedCrossing: 0 }))
113  await update($, weather, () => ({ points: [], limit: null, isCompaction: false }))
114  await update($, route, () => ({ tier: null, model: null, isApplied: false, current: null }))
115}
116
117export const register: Register = (on, options) => {
118  const mode = options.routerMode === 'off' || options.routerMode === 'auto' ? options.routerMode : 'suggest'
119  const threshold = typeof options.coachThreshold === 'number' ? options.coachThreshold : 150000
120  const handoffText = typeof options.handoffPrompt === 'string' && options.handoffPrompt.trim() ? options.handoffPrompt : 'Write a handoff note for the next session.'
121  const isWeatherOn = options.weather !== false
122  const isDietOn = options.diet !== false
123
124  // The model the running turn is routed to; set at prompt.submit, cleared when the turn completes.
125  let target: string | null = null
126
127  on('session.start', async ($, e, next) => {
128    const dropped = isDietOn ? await $.store.get(`diet:${e.cwd}`) : []
129    await update($, diet, d => ({ ...d, cwd: e.cwd, dropped: Array.isArray(dropped) ? dropped.map(String) : [] }))
130    if (isDietOn) await $.command.register({ name: 'diet', description: 'Choose system prompt sections to drop for this project' })
131    return next(e)
132  })
133
134  on('session.end', async ($, e, next) => {
135    if (e.reason === 'clear') {
136      target = null
137      await resetAll($)
138    }
139    return next(e)
140  })
141
142  // Router: a Haiku call sizes the prompt before the turn starts.
143  on('prompt.submit', async ($, e, next) => {
144    if (mode === 'off' || e.command || e.isMeta || e.turnId !== undefined || e.text.trim().length < 4) return next(e)
145    const tier = await classify($, e.text)
146    if (!tier) return next(e)
147    const wanted = modelForTier(options, tier)
148    const current = await $.session.model()
149    const isDifferent = !isSameModel(current, wanted)
150    target = mode === 'auto' && isDifferent ? wanted : null
151    await setRoute($, { tier, model: wanted, isApplied: target !== null, current })
152    return next(e)
153  }).catch(($, e, next) => next(e))
154
155  on('turn.step', async function* ($, e, next) {
156    if (target !== null && e.agentId === undefined && !isSameModel(e.model, target)) return yield* next({ ...e, model: target })
157    return yield* next(e)
158  })
159
160  on('turn.complete', async ($, e, next) => {
161    if (e.agentId === undefined) {
162      target = null
163      if (isWeatherOn) await recordTurn($).catch(() => undefined)
164    }
165    return next(e)
166  })
167
168  // Coach: the context size, as the engine measures it after each response.
169  on('session.measure', async ($, e, next) => {
170    if (e.context.tokens !== undefined && threshold > 0) await noteContext($, e.context.tokens, threshold)
171    return next(e)
172  })
173
174  // Diet: dropped sections leave the system prompt; pinned ones never do.
175  on('prompt.compose', async ($, e, next) => {
176    const composed = await next(e)
177    const { dropped } = await read($, diet)
178    if (!isDietOn || dropped.length === 0) return composed
179    const kept = applyDiet(composed.sections, dropped)
180    return kept.sections.length === composed.sections.length ? composed : { sections: kept.sections }
181  })
182
183  on('command.run', { command: 'diet' }, async $ => {
184    await refreshSections($)
185    await $.ui.open({ id: PANE, title: 'Instruction diet', focus: true, closeOnEscape: true })
186    return { text: 'Instruction diet opened.' }
187  })
188
189  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
190    const rest = await next(e)
191    if (e.props.hasSurvey) return rest
192
193    const r = await read($, route)
194    const c = await read($, coach)
195    const w = await read($, weather)
196    const d = await read($, diet)
197    const { Box, Text, Button } = $.ui.resolve(e)
198
199    const head: string[] = []
200    const routed = mode === 'off' ? undefined : routeLine(r, r.current !== null && r.model !== null && isSameModel(r.current, r.model))
201    if (routed) head.push(routed)
202    const saved = savedTokens(d.sections, d.dropped)
203    if (isDietOn && saved > 0) head.push(`diet saves ~${fmtTokens(saved)} per request`)
204
205    const line = isWeatherOn ? contextLine(w.points) : undefined
206    const isCoaching = threshold > 0 && c.tokens !== null && c.tokens >= threshold && c.crossing !== c.dismissedCrossing
207    const shown = line ?? (isCoaching ? `context ${fmtTokens(c.tokens ?? 0)}` : undefined)
208
209    if (head.length === 0 && !shown) return rest
210
211    return (
212      <Box flexDirection="column">
213        {head.length > 0 && <Text dimColor>thrift · {head.join(' · ')}</Text>}
214        {shown && (
215          <Box>
216            {isCoaching ? <Text color="yellow">{shown}, past {fmtTokens(threshold)}: </Text> : <Text dimColor>{shown}</Text>}
217            {isCoaching && <Button key="handoff" label="Hand off" variant="primary" onPress={() => handOff($, handoffText)} />}
218            {isCoaching && <Text> </Text>}
219            {isCoaching && <Button key="compact" label="Compact" onPress={() => compactNow($)} />}
220            {isCoaching && <Text> </Text>}
221            {isCoaching && <Button key="dismiss" label="Dismiss" dimColor onPress={() => dismiss($)} />}
222          </Box>
223        )}
224        {rest}
225      </Box>
226    )
227  })
228
229  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
230    const { Box, Text, Button } = $.ui.resolve(e)
231    const d = await read($, diet)
232    const total = d.sections.reduce((sum, s) => sum + s.tokens, 0)
233    const saved = savedTokens(d.sections, d.dropped)
234    return (
235      <Box flexDirection="column">
236        <Text bold>
237          System prompt: ~{fmtTokens(total)} tokens in {d.sections.length} sections; dropped here: ~{fmtTokens(saved)}
238        </Text>
239        <Text dimColor>Sizes are estimates (4 characters a token). Choices are kept per folder: {d.cwd || 'this project'}</Text>
240        {d.sections.map(s => {
241          const isDropped = d.dropped.includes(s.id)
242          return (
243            <Box key={s.id}>
244              <Text dimColor={isDropped}>{s.id.padEnd(24)}</Text>
245              <Text dimColor>{`~${fmtTokens(s.tokens)}`.padStart(7)} </Text>
246              {s.isPinned ? <Text dimColor>pinned</Text> : <Button key={`toggle-${s.id}`} label={isDropped ? 'Restore' : 'Drop'} onPress={() => toggleSection($, s.id)} />}
247            </Box>
248          )
249        })}
250      </Box>
251    )
252  })
253}
254

1import type { DietSection, Tier } from '../types'
2
3export const TIERS: readonly Tier[] = ['trivial', 'normal', 'hard']
4
5// Sections that carry the engine's tool and safety rules; dropping them breaks the session.
6const PINNED = new Set(['bare', 'lean_body', 'intro', 'system', 'tools', 'actions'])
7
8export const isPinnedSection = (id: string) => PINNED.has(id)
9
10// One token is about four characters of English or code.
11export const estimateTokens = (text: string) => Math.ceil(text.length / 4)
12
13export const fmtTokens = (n: number) => (n >= 10_000 ? `${Math.round(n / 1000)}k` : n >= 1000 ? `${(n / 1000).toFixed(1).replace(/\.0$/, '')}k` : String(Math.round(n)))
14
15export function parseTier(label: string | undefined): Tier | null {
16  const t = (label ?? '').trim().toLowerCase()
17  return TIERS.find(tier => t.startsWith(tier)) ?? null
18}
19
20export function modelForTier(options: Record<string, unknown>, tier: Tier): string {
21  const fallback = { trivial: 'haiku', normal: 'sonnet', hard: 'opus' }[tier]
22  const key = `${tier}Model`
23  const value = options[key]
24  return typeof value === 'string' && value.trim() ? value.trim() : fallback
25}
26
27// True when `current` (a full id such as claude-sonnet-x) already is the wanted alias or id.
28export const isSameModel = (current: string, wanted: string) => current.toLowerCase().includes(wanted.toLowerCase())
29
30// Least-squares slope over the last `window` points, in tokens per turn.
31export function slope(points: readonly number[], window = 6): number {
32  const p = points.slice(-window)
33  if (p.length < 2) return 0
34  const n = p.length
35  const meanX = (n - 1) / 2
36  const meanY = p.reduce((a, b) => a + b, 0) / n
37  let num = 0
38  let den = 0
39  p.forEach((y, x) => {
40    num += (x - meanX) * (y - meanY)
41    den += (x - meanX) ** 2
42  })
43  return num / den
44}
45
46// One line: the context size once and how fast it grows per turn (each turn re-sends all of it).
47export function contextLine(points: readonly number[]): string | undefined {
48  const last = points[points.length - 1]
49  if (last === undefined) return undefined
50  if (points.length < 2) return `context ${fmtTokens(last)}`
51  const s = slope(points)
52  return s > 0 ? `context ${fmtTokens(last)}, growing ~${fmtTokens(s)} per turn` : `context ${fmtTokens(last)}, steady`
53}
54
55const TIER_WORDS: Record<string, string> = { trivial: 'simple', normal: 'routine', hard: 'hard' }
56
57// The router's line, shown only when it would change something.
58export function routeLine(r: { tier: string | null; model: string | null; isApplied: boolean; current: string | null }, isSame: boolean): string | undefined {
59  if (!r.tier || !r.model || isSame) return undefined
60  const kind = TIER_WORDS[r.tier] ?? r.tier
61  return r.isApplied
62    ? `${kind} prompt: this turn runs on ${r.model}`
63    : `${kind} prompt: ${r.model} would do${r.current ? ` (you're on ${r.current})` : ''}; /model ${r.model} to switch`
64}
65
66// A compaction or /clear shows as a sharp fall; the trend starts over from there.
67export const isReset = (prev: number | undefined, next: number) => prev !== undefined && next < prev * 0.7
68
69export const measureSections = (sections: readonly { id: string; text: string }[]): DietSection[] =>
70  sections.map(s => ({ id: s.id, tokens: estimateTokens(s.text), isPinned: isPinnedSection(s.id) }))
71
72// Drops the named sections, never a pinned one; order is kept so shared sections stay ahead of session ones.
73export function applyDiet<T extends { id: string; text: string }>(sections: readonly T[], dropped: readonly string[]): { sections: T[]; saved: number } {
74  const gone = new Set(dropped.filter(id => !isPinnedSection(id)))
75  const kept = sections.filter(s => !gone.has(s.id))
76  const saved = sections.filter(s => gone.has(s.id)).reduce((sum, s) => sum + estimateTokens(s.text), 0)
77  return { sections: kept, saved }
78}
79
80export const savedTokens = (sections: readonly DietSection[], dropped: readonly string[]) =>
81  sections.filter(s => dropped.includes(s.id) && !s.isPinned).reduce((sum, s) => sum + s.tokens, 0)
82

1export type Tier = 'trivial' | 'normal' | 'hard'
2
3export type Route = {
4  tier: Tier | null
5  /** The model the tier maps to; null while nothing is classified. */
6  model: string | null
7  /** True when the turn runs on `model`; false when it is only a suggestion. */
8  isApplied: boolean
9  /** The session's model when the prompt was classified. */
10  current: string | null
11}
12
13export type Coach = {
14  tokens: number | null
15  /** Counts each time the context passes the threshold from below. */
16  crossing: number
17  dismissedCrossing: number
18}
19
20export type Weather = {
21  /** Context size at the end of each turn, oldest first. */
22  points: number[]
23  /** Tokens at which the engine compacts, or fills the window; null until read. */
24  limit: number | null
25  isCompaction: boolean
26}
27
28export type DietSection = { id: string; tokens: number; isPinned: boolean }
29
30export type Diet = {
31  cwd: string
32  dropped: string[]
33  sections: DietSection[]
34}
35
36declare module 'claude-code' {
37  interface PluginState {
38    thrift: { route: Route; coach: Coach; weather: Weather; diet: Diet }
39  }
40}
41