next-steps

1import { atom, read, update } from 'claude-code'
2import type { Elements, EngineInterface, Register } from 'claude-code'
3
4import { buildPrompt, fileNameOf, fillPlan, parseSteps, shouldSuggest, SYSTEM } from './lib/steps'
5
6const stepsAtom = atom({ plugin: 'next-steps', key: 'steps' } as const, [] as string[])
7const mutedAtom = atom({ plugin: 'next-steps', key: 'isMuted' } as const, false)
8/** Indexes ticked for insertion, in the order they were ticked. */
9const pickedAtom = atom({ plugin: 'next-steps', key: 'picked' } as const, [] as number[])
10/** True when the pane could not be placed and the suggestions fall back into the band. */
11const inBandAtom = atom({ plugin: 'next-steps', key: 'inBand' } as const, false)
12
13const PANE = 'next-steps'
14const EDIT_TOOLS = new Set(['Edit', 'Write', 'NotebookEdit'])
15const HOTKEYS = ['1', '2', '3']
16
17type Ask = { readonly answer: string; readonly files: readonly string[]; readonly model: string }
18
19/** One small completion; a failure leaves the band empty, never an error. */
20async function suggest($: EngineInterface, ask: Ask): Promise<void> {
21  if (await read($, mutedAtom)) return
22  const project = fileNameOf(await $.session.root())
23  const reply = await $.model.complete({
24    model: ask.model,
25    system: SYSTEM,
26    prompt: buildPrompt({ answer: ask.answer, files: ask.files, project }),
27    maxTokens: 300,
28    effort: 'low',
29    timeoutMs: 20_000,
30  })
31  if (!reply.isAnswered) return
32  const steps = parseSteps(reply.text)
33  await update($, stepsAtom, () => steps)
34  await update($, pickedAtom, () => [])
35  if (steps.length === 0) return
36  const opened = await $.ui.open({ id: PANE, title: 'Next steps' }).catch(() => ({ isPlaced: false }))
37  await update($, inBandAtom, () => !opened.isPlaced)
38}
39
40async function lastAnswer($: EngineInterface): Promise<string> {
41  const messages = await $.session.messages()
42  if ('deny' in messages) return ''
43  return [...messages].reverse().find(m => m.role === 'assistant' && m.text.trim() !== '')?.text ?? ''
44}
45
46type Table = Pick<Elements['terminal'], 'Box' | 'Text' | 'Button'>
47
48/** Clears the suggestions and closes their pane. */
49async function clearSteps($: EngineInterface): Promise<void> {
50  await update($, stepsAtom, () => [])
51  await update($, pickedAtom, () => [])
52  await update($, inBandAtom, () => false)
53  await $.ui.close({ id: PANE }).catch(() => undefined)
54}
55
56/** Ticks a step on or off. Any number may be ticked before they go in together. */
57async function toggleStep($: EngineInterface, index: number): Promise<void> {
58  await update($, pickedAtom, list => (list.includes(index) ? list.filter(i => i !== index) : [...list, index]))
59}
60
61/** Ticks every step, keeping the order of any already ticked first. */
62async function pickAll($: EngineInterface): Promise<void> {
63  const steps = await read($, stepsAtom)
64  await update($, pickedAtom, list => [...list, ...steps.map((_, i) => i).filter(i => !list.includes(i))])
65}
66
67/** Puts every ticked step in the prompt box without sending it, keeping any text already there. */
68async function insertSteps($: EngineInterface): Promise<void> {
69  const steps = await read($, stepsAtom)
70  const picked = await read($, pickedAtom)
71  const text = picked
72    .map(i => steps[i] ?? '')
73    .filter(t => t !== '')
74    .join('\n\n')
75  if (text === '') return
76  const box = await $.prompt.read().catch(() => ({ text: '' }))
77  const plan = fillPlan(box.text, text)
78  const filled = await $.prompt.fill({ text: plan.text, mode: plan.mode }).catch(() => ({ isFilled: false }))
79  if (filled.isFilled) await clearSteps($)
80}
81
82/** The label of the insert action: it names how many steps are going in. */
83export const insertLabel = (count: number): string =>
84  count === 0 ? 'Tick steps to insert' : count === 1 ? 'Insert 1 step' : `Insert ${count} steps`
85
86/**
87 * The suggestions as a checklist: tick as many as you want, then insert them together.
88 * Ticked rows draw bright and unticked rows dim, and Insert is always on show,
89 * so the pane reads as multi-select before anything is pressed.
90 */
91function stepsTree($: EngineInterface, els: Table, steps: readonly string[], picked: readonly number[], columns: number) {
92  const { Box, Button, Text } = els
93  const width = Math.max(20, columns - 12)
94  const isClipped = (t: string): boolean => t.length > width
95  const clip = (t: string): string => (isClipped(t) ? `${t.slice(0, width - 1)}…` : t)
96  const count = picked.length
97  const isAllPicked = count === steps.length
98  return (
99    <Box key="steps" flexDirection="column">
100      <Text dimColor>Tick any number (click or {steps.length === 1 ? 'press 1' : `1-${steps.length}`}), then insert them together.</Text>
101      {steps.map((step, i) => {
102        const isPicked = picked.includes(i)
103        return (
104          <Box key={`row-${i}`} flexDirection="column" marginTop={1}>
105            <Button
106              key={`step-${i}`}
107              hotkey={HOTKEYS[i]}
108              plain
109              dimColor={!isPicked}
110              label={`${isPicked ? '[x]' : '[ ]'} ${clip(step)}`}
111              onPress={() => {
112                void toggleStep($, i)
113              }}
114            />
115            {isClipped(step) && (
116              <Text dimColor wrap="wrap">
117                {step}
118              </Text>
119            )}
120          </Box>
121        )
122      })}
123      <Box key="actions" flexDirection="row" gap={2} marginTop={1}>
124        <Button
125          key="insert"
126          hotkey="i"
127          variant="primary"
128          dimColor={count === 0}
129          label={insertLabel(count)}
130          onPress={() => {
131            void insertSteps($)
132          }}
133        />
134        {isAllPicked ? (
135          <Button
136            key="clear"
137            hotkey="c"
138            dimColor
139            label="Clear"
140            onPress={() => {
141              void update($, pickedAtom, () => [])
142            }}
143          />
144        ) : (
145          <Button
146            key="all"
147            hotkey="a"
148            dimColor
149            label="Select all"
150            onPress={() => {
151              void pickAll($)
152            }}
153          />
154        )}
155        <Button
156          key="dismiss"
157          role="dismiss"
158          plain
159          dimColor
160          label="Dismiss"
161          onPress={() => {
162            void clearSteps($)
163          }}
164        />
165      </Box>
166    </Box>
167  )
168}
169
170export const register: Register = (on, options) => {
171  const model = typeof options.model === 'string' && options.model !== '' ? options.model : 'haiku'
172  const minAnswerChars = typeof options.minAnswerChars === 'number' ? options.minAnswerChars : 300
173
174  // What this turn did, for the suggestion prompt. Reset on each prompt.
175  let toolCalls = 0
176  let files: string[] = []
177
178  on('session.start', async ($, e, next) => {
179    await $.command.register({
180      name: 'next',
181      description: 'Next steps: refresh the suggestions, or turn them on or off',
182      argumentHint: '[on|off]',
183    })
184    return next(e)
185  })
186
187  on('prompt.submit', async ($, e, next) => {
188    toolCalls = 0
189    files = []
190    await clearSteps($)
191    return next(e)
192  })
193
194  on('tool.call', ($, e, next) => {
195    if (e.agentId === undefined) {
196      toolCalls += 1
197      const path = 'file_path' in e && typeof e.file_path === 'string' ? e.file_path : ''
198      if (EDIT_TOOLS.has(String(e.tool)) && path !== '' && !files.includes(path)) files = [...files, path]
199    }
200    return next(e)
201  })
202
203  on('turn.complete', async ($, e, next) => {
204    const result = await next(e)
205    const isMainAnswer = e.agentId === undefined && e.reason === 'answer' && !e.isAborted
206    if (isMainAnswer && shouldSuggest({ answer: e.answer, toolCalls, minAnswerChars })) {
207      const ask: Ask = { answer: e.answer, files: files.map(fileNameOf), model }
208      // Off the turn's path: the turn ends now, the band fills a moment later.
209      $.clock.after(0, () => {
210        void suggest($, ask).catch(() => undefined)
211      })
212    }
213    return result
214  })
215
216  on('command.run', { command: 'next' }, async ($, e) => {
217    const arg = e.args.trim().toLowerCase()
218    if (arg === 'off') {
219      await update($, mutedAtom, () => true)
220      await clearSteps($)
221      return { text: 'Next-step suggestions are off for this session.' }
222    }
223    await update($, mutedAtom, () => false)
224    if (arg === 'on') return { text: 'Next-step suggestions are on.' }
225    const answer = await lastAnswer($)
226    if (answer === '') return { text: 'No answer yet to suggest from.' }
227    await suggest($, { answer, files: files.map(fileNameOf), model })
228    const steps = await read($, stepsAtom)
229    return { text: steps.length ? `Suggested:\n${steps.map((s, i) => `${i + 1}. ${s}`).join('\n')}` : 'Nothing to suggest.' }
230  })
231
232  // The suggestions get a pane of their own, so the band above the prompt — the
233  // token-ledger bar — is never replaced and never scrolls with them. Only a
234  // terminal too narrow for a pane falls back to the band.
235  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
236    const width = e.viewport?.columns ?? 80
237    return stepsTree($, $.ui.resolve(e), await read($, stepsAtom), await read($, pickedAtom), width)
238  })
239
240  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
241    const steps = await read($, stepsAtom)
242    const isInBand = await read($, inBandAtom)
243    if (e.props.hasSurvey || steps.length === 0 || !isInBand || (await read($, mutedAtom))) return next(e)
244    return stepsTree($, $.ui.resolve(e), steps, await read($, pickedAtom), e.props.bodyColumns ?? 80)
245  })
246}
247

1// Pure prompt building and reply parsing for next-step suggestions.
2
3import { redact } from './redact'
4
5export const MAX_STEPS = 3
6export const MAX_STEP_CHARS = 90
7/** How much of the answer's tail goes to the model: the ending says what is left. */
8export const ANSWER_TAIL_CHARS = 4_000
9
10export const SYSTEM = [
11  'You suggest what the user should ask Claude Code to do next.',
12  'Read the end of the assistant answer and the files it touched.',
13  'Reply with a JSON array of 1 to 3 strings and nothing else.',
14  'Each string is a prompt the user could send next: imperative, specific, under 80 characters.',
15  'Prefer verifying or finishing the work just done over new work.',
16  'If nothing useful follows, reply [].',
17].join(' ')
18
19export const shouldSuggest = (args: {
20  readonly answer: string
21  readonly toolCalls: number
22  readonly minAnswerChars: number
23}): boolean => args.answer.trim().length >= args.minAnswerChars || args.toolCalls >= 3
24
25export const buildPrompt = (args: {
26  readonly answer: string
27  readonly files: readonly string[]
28  readonly project: string
29}): string => {
30  // Personal data never leaves for the suggestion call. Redact the whole
31  // answer before cutting it, so no number is split across the cut, and
32  // redact file and project names too (`client-8001015009087.csv`).
33  const scrub = (text: string): string => redact(text, { salt: 'next-steps', redactAmounts: false }).text
34  const safe = scrub(args.answer).slice(-ANSWER_TAIL_CHARS)
35  const files = args.files.length ? scrub(args.files.slice(-15).join(', ')) : 'none'
36  return `Project: ${scrub(args.project)}\nFiles touched this turn: ${files}\n\nEnd of the assistant answer:\n"""\n${safe}\n"""`
37}
38
39const clean = (s: string): string => s.replace(/\s+/g, ' ').trim()
40
41/** Reads the model's reply as a JSON array, or failing that as a bulleted list. */
42export const parseSteps = (reply: string): string[] => {
43  const fromJson = (() => {
44    const match = reply.match(/\[[\s\S]*\]/)
45    if (!match) return undefined
46    try {
47      const parsed: unknown = JSON.parse(match[0])
48      return Array.isArray(parsed) ? parsed.filter((x): x is string => typeof x === 'string') : undefined
49    } catch {
50      return undefined
51    }
52  })()
53  const raw =
54    fromJson ??
55    reply
56      .split('\n')
57      .map(line => line.match(/^\s*(?:[-*•]|\d+[.)])\s+(.*)$/)?.[1])
58      .filter((x): x is string => x !== undefined)
59  const unique = [...new Set(raw.map(clean).filter(s => s.length > 0))]
60  return unique.slice(0, MAX_STEPS).map(s => (s.length > MAX_STEP_CHARS ? `${s.slice(0, MAX_STEP_CHARS - 1)}…` : s))
61}
62
63export const fileNameOf = (p: string): string => p.replace(/\\/g, '/').split('/').at(-1) ?? p
64
65/**
66 * How the ticked steps join the prompt box: never over the user's own text.
67 * `text` is already the ticked steps joined, so one call covers any number.
68 */
69export const fillPlan = (current: string, text: string): { readonly text: string; readonly mode: 'replace' | 'append' } =>
70  current.trim() === '' ? { text, mode: 'replace' } : { text: `${current.endsWith('\n') ? '' : '\n\n'}${text}`, mode: 'append' }
71

1// Pure detection and tokenising of South African personal information.
2// No engine calls here, so the tests cover it directly.
3
4export type Kind = 'EMAIL' | 'SA_ID' | 'CARD' | 'PHONE' | 'BANK_ACC' | 'PASSPORT' | 'AMOUNT' | 'MEDIA'
5
6export type RedactOptions = {
7  readonly salt: string
8  readonly redactAmounts: boolean
9}
10
11export type RedactResult = {
12  readonly text: string
13  readonly counts: Readonly<Partial<Record<Kind, number>>>
14}
15
16/** Tokens hold letters only, so no later digit pattern can match inside one. */
17export const TOKEN_PATTERN = /\[(?:EMAIL|SA_ID|CARD|PHONE|BANK_ACC|PASSPORT|AMOUNT):[a-p]{8}\]/
18
19export const luhnValid = (digits: string): boolean => {
20  if (!/^\d+$/.test(digits)) return false
21  let sum = 0
22  for (let i = 0; i < digits.length; i += 1) {
23    let d = Number(digits[digits.length - 1 - i])
24    if (i % 2 === 1) {
25      d *= 2
26      if (d > 9) d -= 9
27    }
28    sum += d
29  }
30  return sum % 10 === 0
31}
32
33const DAYS_IN_MONTH = [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
34
35/** 13 digits: YYMMDD, 4-digit sequence, citizenship 0/1/2, race digit, Luhn check. */
36export const isSaId = (digits: string): boolean => {
37  if (!/^\d{13}$/.test(digits)) return false
38  const month = Number(digits.slice(2, 4))
39  const day = Number(digits.slice(4, 6))
40  if (month < 1 || month > 12) return false
41  if (day < 1 || day > (DAYS_IN_MONTH[month - 1] ?? 0)) return false
42  if (!'012'.includes(digits[10] ?? '')) return false
43  return luhnValid(digits)
44}
45
46/** FNV-1a over salt + value, spelled in the letters a-p: stable per session, no digits. */
47export const tokenFor = (kind: Kind, value: string, salt: string): string => {
48  let hash = 0x811c9dc5
49  const input = `${salt}\u241f${kind}\u241f${value}`
50  for (let i = 0; i < input.length; i += 1) {
51    hash ^= input.charCodeAt(i)
52    hash = Math.imul(hash, 0x01000193) >>> 0
53  }
54  const letters = hash
55    .toString(16)
56    .padStart(8, '0')
57    .replace(/[0-9a-f]/g, c => String.fromCharCode(97 + parseInt(c, 16)))
58  return `[${kind}:${letters}]`
59}
60
61type Rule = {
62  readonly kind: Kind
63  readonly pattern: RegExp
64  /** Which capture group holds the sensitive part; 0 for the whole match. */
65  readonly group: number
66  readonly accept: (value: string) => boolean
67}
68
69const digitsOf = (s: string): string => s.replace(/\D/g, '')
70
71/** `logo@2x.png`, `pkg@1.2.3.tgz`: file names and package specs, not addresses. */
72const FILE_LIKE = /\.(?:png|jpe?g|gif|svg|webp|ico|tgz|zip|gz|js|mjs|cjs|ts|tsx|jsx|json|css|scss|map|lock|md|html?|py|rb|go|rs|java|php|vue|txt|log|csv|xml|ya?ml)$/i
73
74/** The same person spelled two ways ("082 123 4567", "+27821234567") gets one token. */
75const normalise = (kind: Kind, value: string): string => {
76  if (kind === 'EMAIL' || kind === 'PASSPORT') return value.toLowerCase()
77  const digits = digitsOf(value)
78  if (kind === 'PHONE') return digits.replace(/^(?:0027|27|0)/, '')
79  return digits
80}
81
82const RULES: readonly Rule[] = [
83  {
84    kind: 'EMAIL',
85    pattern: /(?<![\w.%+-])[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}(?![\w:@-])/g,
86    group: 0,
87    accept: v => !FILE_LIKE.test(v) && !/^git@/i.test(v),
88  },
89  {
90    kind: 'SA_ID',
91    pattern: /(?<!\d)(?<!\d\.)\d{6}[ -]?\d{4}[ -]?\d{3}(?!\d)(?!\.\d)/g,
92    group: 0,
93    accept: v => isSaId(digitsOf(v)),
94  },
95  {
96    kind: 'CARD',
97    pattern: /(?<!\d)(?<!\d\.)\d(?:[ -]?\d){12,18}(?!\d)(?!\.\d)/g,
98    group: 0,
99    accept: v => {
100      const d = digitsOf(v)
101      return d.length >= 13 && d.length <= 19 && /^[2-6]/.test(d) && luhnValid(d)
102    },
103  },
104  {
105    kind: 'PHONE',
106    pattern: /(?<![\w+])(?:\+27[ -]?(?:\(0\))?|0027|0)[ -.]?\(?[1-8]\d\)?(?:[ -.]?\d){7}(?!\w)(?!\.\d)/g,
107    group: 0,
108    accept: () => true,
109  },
110  {
111    kind: 'BANK_ACC',
112    pattern:
113      /\b(?:acc(?:ount)?|acct|rekening(?:nommer)?)(?:\s*(?:no|nr|number|num|#))?\.?\s*[:#]?\s*(\d(?:[ -]?\d){5,15})(?!\d)/gi,
114    group: 1,
115    accept: () => true,
116  },
117  {
118    kind: 'PASSPORT',
119    pattern: /\bpas(?:s)?(?:port|poort)(?:\s*(?:no|nr|number|#))?\.?\s*[:#]?\s*([A-Z]{0,2}\d{7,9})\b/gi,
120    group: 1,
121    accept: () => true,
122  },
123]
124
125const AMOUNT_RULE: Rule = {
126  kind: 'AMOUNT',
127  pattern: /(?<![A-Za-z])R\s?\d{1,3}(?:[ ,]\d{3})*(?:[.,]\d{2})?(?![\d])/g,
128  group: 0,
129  accept: () => true,
130}
131
132const applyRule = (
133  text: string,
134  rule: Rule,
135  salt: string,
136): { readonly text: string; readonly count: number } => {
137  let count = 0
138  const out = text.replace(rule.pattern, (match: string, ...rest: unknown[]) => {
139    const part = rule.group === 0 ? match : rest[rule.group - 1]
140    if (typeof part !== 'string' || !rule.accept(part)) return match
141    count += 1
142    const token = tokenFor(rule.kind, normalise(rule.kind, part), salt)
143    return rule.group === 0 ? token : match.replace(part, token)
144  })
145  return { text: out, count }
146}
147
148/** Returns a new string with every detected item swapped for its token. */
149export const redact = (text: string, options: RedactOptions): RedactResult => {
150  const rules = options.redactAmounts ? [...RULES, AMOUNT_RULE] : RULES
151  return rules.reduce<RedactResult>(
152    (acc, rule) => {
153      const step = applyRule(acc.text, rule, options.salt)
154      if (step.count === 0) return acc
155      return {
156        text: step.text,
157        counts: { ...acc.counts, [rule.kind]: (acc.counts[rule.kind] ?? 0) + step.count },
158      }
159    },
160    { text, counts: {} },
161  )
162}
163
164export const totalOf = (counts: Readonly<Partial<Record<Kind, number>>>): number =>
165  Object.values(counts).reduce<number>((a, b) => a + (b ?? 0), 0)
166
167export const mergeCounts = (
168  a: Readonly<Partial<Record<Kind, number>>>,
169  b: Readonly<Partial<Record<Kind, number>>>,
170): Partial<Record<Kind, number>> => {
171  const keys = new Set([...Object.keys(a), ...Object.keys(b)]) as Set<Kind>
172  return Object.fromEntries([...keys].map(k => [k, (a[k] ?? 0) + (b[k] ?? 0)]))
173}
174

1/** A next prompt the band offers. */
2export type NextStep = string
3
4declare module 'claude-code' {
5  interface PluginState {
6    'next-steps': {
7      /** The suggestions on offer; empty hides them. */
8      steps: NextStep[]
9      /** Indexes ticked for insertion, in the order they were ticked. */
10      picked: number[]
11      /** True when the pane could not be placed and the suggestions fall back into the band. */
12      inBand: boolean
13      /** Turned off for this session with `/next off`. */
14      isMuted: boolean
15    }
16  }
17}
18
19