assumption-check

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Assumption, Item } from '../types'
5import { ASK, CHANGING, RETRY, TIMEOUT_MS, correctionPrompt, costLine, lastTurn, parseAssumptions } from './assume'
6import { STORE_KEY, storedSwitch, switchText, switchWord } from './toggle'
7
8const PANE = 'assumption-check'
9// The switch as this session read it (toggle.ts). Off, every hook passes its event on unchanged and no turn is
10// checked. The mod starts off; `/assumption-check on` turns it on for every session.
11const enabled = atom({ plugin: 'assumption-check', key: 'enabled' } as const, false)
12const items = atom({ plugin: 'assumption-check', key: 'items' } as const, [] as Item[])
13const checking = atom({ plugin: 'assumption-check', key: 'checking' } as const, false)
14const notice = atom({ plugin: 'assumption-check', key: 'notice' } as const, null as string | null)
15const lastCost = atom({ plugin: 'assumption-check', key: 'lastCost' } as const, null as string | null)
16const isFilled = atom({ plugin: 'assumption-check', key: 'isFilled' } as const, false)
17const round = atom({ plugin: 'assumption-check', key: 'round' } as const, 0)
18
19// This turn as the hooks saw it. `changed`: calls of the main loop that edited a file or ran a command.
20// `isBusy`: a turn of the main loop runs, so nothing is submitted. `notes`: the review's notes by item, kept
21// here rather than in the state so that typing does not redraw the box it types into.
22type Session = { changed: number; isBusy: boolean; notes: string[] }
23
24// The check: one Haiku call over the last turn, and one retry if the answer is not the JSON asked for. The list
25// is the newest checked turn's alone, so the previous one is cleared first.
26async function collect($: any, s: Session): Promise<void> {
27  await update($, checking, () => true)
28  await update($, notice, () => null)
29  await update($, items, () => [])
30  await update($, isFilled, () => false)
31  s.notes = []
32  try {
33    const turn = lastTurn(await $.session.messages())
34    const used = { input_tokens: 0, output_tokens: 0 }
35    let found: Assumption[] | undefined
36    let calls = 0
37    let failed: string | null = null
38    for (const ask of [ASK, `${ASK}\n\n${RETRY}`]) {
39      calls++
40      const answer = await $.model.complete({ model: 'haiku', prompt: `${turn}\n\n---\n${ask}`, maxTokens: 1500, timeoutMs: TIMEOUT_MS })
41      if (!answer.isAnswered) {
42        failed = answer.reason === 'aborted' ? `no answer within ${TIMEOUT_MS / 1000} s` : `the model call failed (${answer.reason})`
43        break
44      }
45      used.input_tokens += answer.usage.input_tokens
46      used.output_tokens += answer.usage.output_tokens
47      found = parseAssumptions(answer.text)
48      if (found !== undefined) break
49    }
50    await update($, lastCost, () => costLine(used, calls))
51    if (found === undefined) {
52      await update($, notice, () => failed ?? "could not read Haiku's answer as a list")
53      return
54    }
55    const list = found
56    await update($, items, () => list.map(a => ({ ...a, mark: null, note: '' })))
57    // A new list gets new keys, so that no note box keeps the text typed into the previous list's.
58    await update($, round, r => r + 1)
59    await update($, isFilled, () => false)
60  } finally {
61    await update($, checking, () => false)
62  }
63}
64
65async function openReview($: any): Promise<void> {
66  // Panes draw only in the terminal and the desktop app; elsewhere the command's line is the answer.
67  await $.ui.open({ id: PANE, title: 'Assumptions', focus: true, closeOnEscape: true }).catch(() => undefined)
68}
69
70async function mark($: any, n: number, value: 'right' | 'wrong'): Promise<void> {
71  await update($, items, list => list.map((i, k) => (k === n ? { ...i, mark: i.mark === value ? null : value } : i)))
72}
73
74function noteOn(s: Session, n: number, text: string): void {
75  s.notes[n] = text
76}
77
78// First press: the corrections go into the prompt box, after any draft there. Second press: they are submitted,
79// but never while Claude works (then Enter in the box is the author's to press).
80async function sendCorrections($: any, s: Session): Promise<void> {
81  if (await read($, isFilled)) {
82    if (s.isBusy) {
83      await update($, notice, () => 'Claude is working: press Enter in the prompt box when it is done')
84      return
85    }
86    const box = (await $.prompt.read()).text
87    if (box.trim() === '') {
88      await update($, notice, () => 'the prompt box is empty: nothing to send')
89      return
90    }
91    await $.prompt.fill({ text: '', mode: 'replace' })
92    await $.prompt.submit({ text: box, asUser: true })
93    await update($, items, () => [])
94    await update($, isFilled, () => false)
95    await update($, notice, () => 'corrections sent')
96    return
97  }
98  const prompt = correctionPrompt((await read($, items)).map((i, k) => ({ ...i, note: s.notes[k] ?? '' })))
99  if (prompt === undefined) {
100    await update($, notice, () => 'mark at least one assumption Wrong first')
101    return
102  }
103  const box = (await $.prompt.read()).text
104  const filled = await $.prompt.fill({ text: box.trim() === '' ? prompt : `${box.replace(/\s+$/, '')}\n\n${prompt}`, mode: 'replace' })
105  await update($, isFilled, () => filled.isFilled)
106  await update($, notice, () => (filled.isFilled ? 'the corrections are in the prompt box: press Enter there, or Send here' : 'the prompt box could not take them now'))
107}
108
109export const register: Register = on => {
110  const s: Session = { changed: 0, isBusy: false, notes: [] }
111
112  on('session.start', async ($, e, next) => {
113    // Registered even when the mod is off, so that it can be turned on.
114    await $.command
115      .register({ name: 'assumption-check', description: "Switch the assumption check on or off, show the last check's cost, or open the review", argumentHint: '[on|off|status|review]', immediate: true })
116      .catch(() => undefined)
117    const isOn = storedSwitch(await $.store.get(STORE_KEY).catch(() => undefined), false)
118    await update($, enabled, () => isOn)
119    return next(e)
120  })
121
122  on('command.run', { command: 'assumption-check' }, async ($, e) => {
123    const isReview = (e.args ?? '').trim().toLowerCase() === 'review'
124    const word = isReview ? 'status' : switchWord(e.args)
125    if (word === undefined) return { text: 'use /assumption-check on, off, status or review' }
126    if (word === 'status') {
127      if (!(await read($, enabled))) return { text: switchText('assumption-check', false) }
128      if (isReview) await openReview($)
129      const n = (await read($, items)).length
130      const cost = await read($, lastCost)
131      return { text: `${switchText('assumption-check', true)} · ${n} to check${cost === null ? ' · no turn checked yet' : ` · last check: ${cost}`}` }
132    }
133    const isOn = word === 'on'
134    await $.store.set(STORE_KEY, isOn)
135    await update($, enabled, () => isOn)
136    return { text: switchText('assumption-check', isOn) }
137  })
138
139  on('turn.start', async ($, e, next) => {
140    if (!(await read($, enabled))) return next(e)
141    s.changed = 0
142    s.isBusy = true
143    await update($, notice, () => null)
144    return next(e)
145  })
146
147  // An observer: counts the main loop's calls that changed something; the result goes back as it came.
148  on('tool.call', async ($, e, next) => {
149    if (!(await read($, enabled))) return next(e)
150    const ran = await next(e)
151    if (e.agentId === undefined && CHANGING.includes(String(e.tool)) && ran.deny === undefined) s.changed++
152    return ran
153  }).catch(($, e, next) => next(e))
154
155  // At most one check per turn, and none for a turn that only read or was interrupted. It runs after the turn
156  // has ended, so the turn never waits for it.
157  on('turn.complete', async ($, e, next) => {
158    if (!(await read($, enabled)) || e.agentId !== undefined) return next(e)
159    const done = await next(e)
160    s.isBusy = false
161    const isChecked = s.changed > 0 && !e.isAborted
162    s.changed = 0
163    if (isChecked) $.clock.after(0, () => void collect($, s).catch(() => undefined))
164    return done
165  }).catch(($, e, next) => next(e))
166
167  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
168    const below = await next(e)
169    if (!(await read($, enabled)) || e.props.hasSurvey) return below
170    const list = await read($, items)
171    const busy = await read($, checking)
172    const said = await read($, notice)
173    if (!busy && list.length === 0 && said === null) return below
174    const { Box, Button, Text } = $.ui.resolve(e)
175    const open = list.filter(i => i.mark === null).length
176    return (
177      <Box flexDirection="column">
178        <Box>
179          <Text dimColor wrap="truncate">
180            {busy ? 'assumptions ▸ checking the last turn… ' : list.length > 0 ? `assumptions ▸ ${open} to check ` : `assumptions ▸ ${said} `}
181          </Text>
182          {list.length > 0 && <Button key="review" label="Review" hotkey="a" onPress={() => openReview($)} />}
183        </Box>
184        {below}
185      </Box>
186    )
187  })
188
189  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
190    const ui = $.ui.resolve(e)
191    const { Box, Button, Text } = ui
192    // Mobile draws no text field: there the review marks and sends, without notes.
193    const Input = 'Input' in ui ? ui.Input : undefined
194    if (!(await read($, enabled))) return <Text dimColor>{switchText('assumption-check', false)}</Text>
195    const list = await read($, items)
196    const said = await read($, notice)
197    const filled = await read($, isFilled)
198    const r = await read($, round)
199    return (
200      <Box flexDirection="column">
201        <Text bold>{`Assumptions in the last turn · ${list.length}`}</Text>
202        {/* At the top, so that a list taller than the pane never hides it. */}
203        {list.length > 0 && (
204          <Box>
205            <Button key="send" label={filled ? 'Send' : 'Send corrections'} onPress={() => sendCorrections($, s)} />
206            {said !== null && <Text dimColor wrap="truncate">{` ${said}`}</Text>}
207          </Box>
208        )}
209        {list.length === 0 && said !== null && <Text dimColor>{said}</Text>}
210        {list.length === 0 && <Text dimColor>Nothing to check. After a turn that edits files or runs commands, the list appears here.</Text>}
211        {list.map((item, n) => (
212          <Box key={`item-${r}-${n}`} flexDirection="column">
213            <Text {...(item.mark === 'wrong' ? { color: 'error' } : item.mark === 'right' ? { color: 'success' } : {})}>
214              {`${n + 1}. ${item.claim}${item.mark !== null ? ` [${item.mark}]` : ''}`}
215            </Text>
216            {(item.where !== '' || item.why !== '') && <Text dimColor wrap="truncate">{`   ${[item.where, item.why].filter(Boolean).join(' · ')}`}</Text>}
217            <Box>
218              <Button key={`right-${r}-${n}`} label="Right" onPress={() => mark($, n, 'right')} />
219              <Button key={`wrong-${r}-${n}`} label="Wrong" onPress={() => mark($, n, 'wrong')} />
220              {Input !== undefined && (
221                <Input
222                  key={`note-${r}-${n}`}
223                  label=" Note: "
224                  placeholder="what it should be"
225                  submitLabel="keep"
226                  onInput={(value: string) => noteOn(s, n, value)}
227                  onSubmit={(value: string) => noteOn(s, n, value)}
228                />
229              )}
230            </Box>
231          </Box>
232        ))}
233      </Box>
234    )
235  })
236}
237

1// assumption-check's pure parts: which turns are checked, what Haiku reads and is asked, how its answer is read,
2// and the prompt the author's corrections become.
3import type { Assumption, Item } from '../types'
4
5// A turn that changed something: it edited a file or ran a command. A turn that only read is not checked.
6export const CHANGING: readonly string[] = ['Edit', 'Write', 'MultiEdit', 'NotebookEdit', 'Bash', 'PowerShell']
7export const MAX_ITEMS = 8
8export const TIMEOUT_MS = 60_000
9const INPUT_CHARS = 1500
10const TURN_CHARS = 30_000
11
12export const ASK =
13  'List the assumptions made in the last turn that the user did not state: choices of format, library, behaviour, naming, defaults, error handling, time zone or scope. Answer with JSON only, no prose: [{"claim": "...", "where": "file or command", "why": "why it was chosen"}], at most 8, most consequential first; [] if there are none.'
14export const RETRY = 'Your last answer was not valid JSON. Answer again with the JSON array only.'
15
16type Message = { role: string; text: string; toolUses?: readonly { tool: string; input: unknown }[] }
17
18// The last turn as text: from the author's last prompt to the end, each tool call with its input cut short.
19// A long turn keeps its end, where the answer and the last edits are.
20export function lastTurn(messages: readonly Message[]): string {
21  let start = messages.length - 1
22  while (start > 0 && !(messages[start]!.role === 'user' && messages[start]!.text !== '')) start--
23  const text = messages
24    .slice(Math.max(0, start))
25    .map(m => {
26      const uses = (m.toolUses ?? []).map(u => `[${u.tool}] ${JSON.stringify(u.input).slice(0, INPUT_CHARS)}`).join('\n')
27      return `${m.role.toUpperCase()}: ${m.text}${uses ? `\n${uses}` : ''}`
28    })
29    .join('\n\n')
30  return text.length > TURN_CHARS ? `…${text.slice(-TURN_CHARS)}` : text
31}
32
33// Haiku's answer as a list, or undefined when it is not the JSON asked for. Prose or a code fence around the
34// array is tolerated; items without a claim are dropped; at most 8 are kept.
35export function parseAssumptions(reply: string): Assumption[] | undefined {
36  const from = reply.indexOf('[')
37  const to = reply.lastIndexOf(']')
38  if (from < 0 || to < from) return undefined
39  let value: unknown
40  try {
41    value = JSON.parse(reply.slice(from, to + 1))
42  } catch {
43    return undefined
44  }
45  if (!Array.isArray(value)) return undefined
46  const text = (v: unknown) => (typeof v === 'string' ? v.replace(/\s+/g, ' ').trim() : '')
47  return value
48    .filter((v): v is Record<string, unknown> => v !== null && typeof v === 'object' && text((v as any).claim) !== '')
49    .slice(0, MAX_ITEMS)
50    .map(v => ({ claim: text(v.claim), where: text(v.where), why: text(v.why) }))
51}
52
53// The prompt the review becomes: the wrong ones with the author's notes, to be changed, and the right ones, to be
54// kept. Undefined when nothing was marked wrong.
55export function correctionPrompt(items: readonly Item[]): string | undefined {
56  const wrong = items.filter(i => i.mark === 'wrong')
57  if (wrong.length === 0) return undefined
58  const right = items.filter(i => i.mark === 'right')
59  const what = (i: Item) => `${i.claim}${i.where !== '' ? ` (${i.where})` : ''}`
60  // "You assumed" keeps the claim from reading as the request itself.
61  const fix = (i: Item, n: number) => `${n + 1}. You assumed: ${what(i)}. That is wrong${i.note.trim() !== '' ? `: ${i.note.trim()}` : '.'}`
62  const parts = [`Some choices in your last turn were wrong. Change only these:\n${wrong.map(fix).join('\n')}`]
63  if (right.length > 0) parts.push(`These were right; keep them as they are:\n${right.map((i, n) => `${n + 1}. ${what(i)}`).join('\n')}`)
64  parts.push('Leave everything else unchanged.')
65  return parts.join('\n\n')
66}
67
68export type Usage = { input_tokens: number; output_tokens: number }
69
70const thousands = (n: number) => String(n).replace(/\B(?=(\d{3})+(?!\d))/g, ',')
71
72export function costLine(usage: Usage, calls: number): string {
73  return `${thousands(usage.input_tokens)} in · ${thousands(usage.output_tokens)} out tokens (haiku${calls > 1 ? `, ${calls} calls` : ''})`
74}
75

1// The mod's on/off switch, `/<mod> on|off|status`. The value is kept in $.store under STORE_KEY: one store per
2// plugin, shared by every session and folder, so the switch holds across restarts. A session reads it as it
3// starts and when its own command changes it; a session already open elsewhere picks a change up when it next
4// starts. Pure helpers only: `claude plugin validate` follows `$` only within the file that uses it.
5export const STORE_KEY = 'enabled'
6
7export type SwitchWord = 'on' | 'off' | 'status'
8
9// The command's argument as a switch word: '' and 'status' ask, 'on' and 'off' set; anything else is the mod's own.
10export function switchWord(args: string | undefined): SwitchWord | undefined {
11  const word = (args ?? '').trim().toLowerCase()
12  if (word === '' || word === 'status') return 'status'
13  return word === 'on' || word === 'off' ? word : undefined
14}
15
16// The stored value, or the mod's default when it was never set (or holds something else).
17export function storedSwitch(value: unknown, byDefault: boolean): boolean {
18  return typeof value === 'boolean' ? value : byDefault
19}
20
21export function switchText(mod: string, isOn: boolean): string {
22  return isOn ? `on (/${mod} off turns it off)` : `off (/${mod} on turns it on)`
23}
24

1// One guess Claude made that the author did not state: what, where, and why Claude chose it.
2export type Assumption = { claim: string; where: string; why: string }
3
4// An assumption as the review holds it: the author's mark, and the note the corrections carry.
5export type Item = Assumption & { mark: 'right' | 'wrong' | null; note: string }
6
7declare module 'claude-code' {
8  interface PluginState {
9    // enabled: the mod's on/off switch as this session read it from the store, or as /assumption-check on|off set it.
10    // items: the last checked turn's assumptions, for this session only. checking: Haiku is reading the last turn.
11    // notice: what the last action did, or why the check failed. lastCost: the tokens the last check took.
12    // isFilled: the corrections are in the prompt box, so the review's button now sends them. round: how many lists
13    // this session has had, which keys the review's elements.
14    'assumption-check': { enabled: boolean; items: Item[]; checking: boolean; notice: string | null; lastCost: string | null; isFilled: boolean; round: number }
15  }
16}
17