next-steps

1// Next Steps: after each turn, 2 or 3 likely next prompts above the input.
2// In an empty prompt, 1 to 3 drafts one (nothing sends on its own) and 0 dismisses.
3// On the desktop they are buttons instead: a click drafts one, the last button dismisses.
4import { atom, read, update } from 'claude-code'
5import type { EngineInterface, Register } from 'claude-code'
6
7import type { NextSteps } from '../types'
8import { bandsFile, parseOrder, stackBands } from './bands'
9
10const MODEL = 'haiku'
11const MAX_ITEMS = 3
12const MAX_CHARS = 60 // what the model is asked for
13const MAX_KEPT = 120 // a longer line is dropped, never cut: a cut draft would send half a prompt
14const MIN_ANSWER = 20 // a reply shorter than this is not worth a model call
15
16// Held by the host, so the list survives a hot reload of this file.
17const steps = atom({ plugin: 'next-steps', key: 'steps' } as const, null as NextSteps | null)
18// True while the list is on screen. where-am-i reads it and leaves out its own "next" meanwhile.
19const active = atom({ plugin: 'next-steps', key: 'active' } as const, false)
20
21let language = ''   // userConfig.language: empty keeps the original English
22
23let bandOrder: Promise<string[]> | undefined // read once per session: restart to apply
24
25export const register: Register = (on, options) => {
26  language = String((options as Record<string, unknown>)?.language ?? '').trim()
27  // The turn a suggestion may still be shown for. A submit or a new turn moves it on.
28  let latest = ''
29  let submits = 0
30  // A survey answers bare digits in the band: leave them to it.
31  let hasSurvey = false
32
33  on('prompt.submit', async ($, e, next) => {
34    latest = `submit-${++submits}`
35    await setSteps($, () => null)
36    return next(e)
37  })
38
39  on('turn.start', async ($, e, next) => {
40    latest = e.turnId
41    await setSteps($, () => null)
42    return next(e)
43  })
44
45  on('turn.complete', async ($, e, next) => {
46    const r = await next(e)
47    if (e.agentId) return r // a subagent's turn, not the person's
48    latest = e.turnId // also for a turn that raised no turn.start, and after a hot reload
49    if (e.reason !== 'answer' || e.answer.trim().length < MIN_ANSWER) return r
50    void suggest($, e.turnId, e.answer, () => latest).catch(() => {}) // in the background, so the turn ends at once
51    return r
52  })
53
54  // A digit typed into an empty prompt picks a suggestion instead of landing as text.
55  // A paste carries no key, so a pasted "1" lands as typed.
56  on('prompt.edit', async ($, e, next) => {
57    if (hasSurvey || !e.key || e.key.ctrl || e.key.meta || e.text !== '' || e.start !== 0 || e.end !== 0 || !/^[0-9]$/.test(e.inputText)) return next(e)
58    const s = await read($, steps)
59    if (!s) return next(e)
60    if (e.inputText === '0') {
61      await setSteps($, () => null)
62      return { text: '', cursor: 0 } // consumed: the box stays empty
63    }
64    const pick = s.items[Number(e.inputText) - 1]
65    if (!pick) return next(e)
66    return next({ ...e, inputText: pick }) // a draft: the person edits it or presses Enter
67  })
68
69  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
70    const rest = await next(e) // what other mods and Claude Code draw here stays
71    hasSurvey = e.props.hasSurvey
72    const s = await read($, steps)
73    if (e.props.hasSurvey || e.props.isWorking || !s || s.items.length === 0) return rest
74
75    const { Box, Button, Text } = $.ui.resolve(e)
76    const order = await (bandOrder ??= readBandOrder($).catch(() => []))
77
78    if (e.surface === 'vscode') {
79      const suggestions = (
80        <Box flexDirection="column" gap={1} alignItems="flex-start" paddingBottom={1}>
81          {s.items.map((text, i) => (
82            <Button
83              key={`next-step-${i + 1}`}
84              label={`${i + 1}. ${text}`}
85              variant="secondary"
86              onPress={async () => {
87                await $.prompt.fill({ text })
88                await setSteps($, () => null)
89              }}
90            />
91          ))}
92          <Button key="next-steps-dismiss" label="Dismiss" dimColor onPress={() => setSteps($, () => null)} />
93        </Box>
94      )
95      return stackBands(rest, 'next-steps', suggestions, 'above', order)
96    }
97
98    if (e.surface !== 'terminal') {
99      const chips = (
100        <Box flexDirection="row" flexWrap="wrap" gap={1} alignItems="center" paddingBottom={1}>
101          {s.items.map((text, i) => (
102            // The desktop draws Buttons natively and drops text color inside them, so the number sits outside.
103            <Box key={`next-step-row-${i + 1}`} gap={1} alignItems="center">
104              <Text color="suggestion" bold>{`${i + 1}`}</Text>
105              <Button
106                key={`next-step-${i + 1}`}
107                label={text}
108                variant="secondary"
109                onPress={async () => {
110                  await $.prompt.fill({ text })
111                  await setSteps($, () => null)
112                }}
113              />
114            </Box>
115          ))}
116          <Button key="next-steps-dismiss" label="✕" dimColor onPress={() => setSteps($, () => null)} />
117        </Box>
118      )
119      return stackBands(rest, 'next-steps', chips, 'above', order)
120    }
121
122    const sep = <Text dimColor>{'  ·  '}</Text>
123    const item = (text: string, i: number) => [
124      <Text color="cyan" bold>{`${i + 1} `}</Text>,
125      <Text>{text}</Text>,
126    ]
127    const dismiss = [<Text color="cyan" bold>{'0 '}</Text>, <Text dimColor>dismiss</Text>]
128
129    const lines =
130      // One per row needs a row each; a short band gets the one row, cut at its end.
131      fitsOnOneLine(s.items, e.props.bodyColumns) || e.props.maxRows < s.items.length + 2 ? (
132        <Text wrap="truncate-end">
133          <Text dimColor>{'next:  '}</Text>
134          {s.items.flatMap((t, i) => [...item(t, i), sep])}
135          {dismiss}
136        </Text>
137      ) : (
138        <Box flexDirection="column">
139          {s.items.map((t, i) => (
140            <Text wrap="truncate-end">
141              <Text dimColor>{i === 0 ? 'next:  ' : '       '}</Text>
142              {item(t, i)}
143              {i === s.items.length - 1 && sep}
144              {i === s.items.length - 1 && dismiss}
145            </Text>
146          ))}
147        </Box>
148      )
149
150    return stackBands(rest, 'next-steps', <Box paddingX={1}>{lines}</Box>, 'above', order) // above what other mods draw, unless the band order says otherwise
151  })
152}
153
154// The band order file, as the claude-code-mods checklist saves it; [] on any error.
155async function readBandOrder($: EngineInterface): Promise<string[]> {
156  const file = bandsFile(await $.env.get('CLAUDE_CONFIG_DIR'), await $.env.get('HOME'))
157  return file && (await $.fs.exists(file)) ? parseOrder(await $.fs.read(file)) : []
158}
159
160// Forgets the cached band order, for tests.
161export function resetBandOrder() {
162  bandOrder = undefined
163}
164
165// Writes the list, then whether it shows, for where-am-i.
166async function setSteps($: EngineInterface, fn: (before: NextSteps | null) => NextSteps | null) {
167  const after = await update($, steps, fn)
168  const showing = after !== null && after.items.length > 0
169  if ((await read($, active)) !== showing) await update($, active, () => showing)
170}
171
172// The whole list on one row, or one suggestion per row when it would not fit.
173export function fitsOnOneLine(items: string[], columns: number) {
174  const width = 'next:  '.length + items.reduce((n, t) => n + 2 + displayWidth(t) + 5, 0) + '0 dismiss'.length
175  return width <= columns - 2 // the band's padding
176}
177
178// Columns a line takes in the terminal: CJK, Hangul, full-width forms and emoji take two,
179// joiners, variation selectors and combining marks none. Where it guesses, it guesses wide:
180// too wide only means one suggestion per row.
181export function displayWidth(text: string) {
182  let width = 0
183  for (const ch of text) {
184    const c = ch.codePointAt(0) ?? 0
185    const zero =
186      (c >= 0x0300 && c <= 0x036f) || (c >= 0x200b && c <= 0x200f) || (c >= 0xfe00 && c <= 0xfe0f) || c === 0x20e3
187    const wide =
188      (c >= 0x1100 && c <= 0x115f) || (c >= 0x231a && c <= 0x231b) || (c >= 0x23e9 && c <= 0x23f3) ||
189      (c >= 0x25fd && c <= 0x25fe) || (c >= 0x2600 && c <= 0x27bf) || (c >= 0x2b1b && c <= 0x2b1c) ||
190      c === 0x2b50 || c === 0x2b55 || (c >= 0x2e80 && c <= 0xa4cf) || (c >= 0xac00 && c <= 0xd7a3) ||
191      (c >= 0xf900 && c <= 0xfaff) || (c >= 0xfe30 && c <= 0xfe4f) || (c >= 0xff00 && c <= 0xff60) ||
192      (c >= 0xffe0 && c <= 0xffe6) || (c >= 0x1f000 && c <= 0x1faff) || (c >= 0x20000 && c <= 0x3fffd)
193    width += zero ? 0 : wide ? 2 : 1
194  }
195  return width
196}
197
198async function suggest($: EngineInterface, turnId: string, answer: string, current: () => string) {
199  // Wait for background agents: the main agent picks up again once they finish.
200  const agents = await $.agent.list().catch(() => [])
201  if (agents.some(a => a.status === 'running')) return
202
203  const all = (await $.session.messages().catch(() => [])).filter(m => m.text.trim() !== '')
204  if (all.at(-1)?.role === 'assistant') all.pop() // the latest reply goes in on its own below
205  const messages = all.slice(-6)
206  const r = await $.model.complete({
207    model: MODEL,
208    maxTokens: 200,
209    timeoutMs: 20_000,
210    system:
211      'You predict what a person will type next to their coding assistant, after reading a transcript of their session. ' +
212      `Give 2 or 3 short, concrete follow-up prompts, each an instruction in the imperative, under ${MAX_CHARS} characters, ` +
213      'specific to what just happened (name the file, test or feature). Plain words, no em dashes, no numbering, no quotes. ' +
214      'One per line and nothing else. If nothing sensible follows, reply with the single word NONE.\n' +
215      'Example reply:\nRun the login tests again\nAdd the same check to signup.ts\nOpen a draft PR' +
216      (language
217        ? `\nWrite every line in ${language}. Keep file names, commands and code identifiers as they are.`
218        : ''),
219    prompt: [
220      `<transcript>\n${messages.map(m => `[${m.role === 'user' ? 'person' : 'assistant'}] ${m.text.slice(0, 800)}`).join('\n')}\n</transcript>`,
221      `<latest_reply>\n${answer.slice(0, 2000)}\n</latest_reply>`,
222      'Write the follow-up prompts the person is likely to send next: the lines only, not a reply to the transcript.',
223    ].join('\n\n'),
224  })
225  if (!r.isAnswered) return
226  const items = parseSteps(r.text)
227  // Checked inside the write, so a submit or a new turn that lands meanwhile always wins.
228  await setSteps($, before => (current() !== turnId ? before : items.length > 0 ? { turnId, items } : null))
229}
230
231// The model's lines, cleaned: no fences, bullets, numbers or em dashes; at most three.
232// A preamble, a "nothing to add" or a line too long to draft whole is dropped.
233export function parseSteps(text: string): string[] {
234  const body = text.replace(/```[a-z]*\n?/gi, '').trim()
235  let lines = body.split('\n')
236  if (body.startsWith('[') || body.startsWith('{')) {
237    try {
238      const o = JSON.parse(body) as unknown
239      const arr = Array.isArray(o) ? o : Object.values(o as Record<string, unknown>).find(Array.isArray)
240      lines = Array.isArray(arr) ? arr.filter((x): x is string => typeof x === 'string') : []
241    } catch {
242      // not JSON after all: read it as lines
243    }
244  }
245  const out: string[] = []
246  for (const line of lines) {
247    const clean = line
248      .replace(/\*\*/g, '')
249      .replace(/^\s*(?:[-*•]|\d+[.):])\s*/, '')
250      .replace(/^["'`]+|["'`,]+$/g, '')
251      .replace(/\s*\u2014\s*/g, ', ')
252      .replace(/\.$/, '')
253      .trim()
254    if (clean === '' || clean.endsWith(':') || clean.length > MAX_KEPT) continue
255    if (/^(none|no |nothing|here are|here's|sure)/i.test(clean) || /^[[\]{}]$/.test(clean)) continue
256    if (!out.some(o => o.toLowerCase() === clean.toLowerCase())) out.push(clean)
257    if (out.length === MAX_ITEMS) break
258  }
259  return out
260}
261

1// Band order: every band mod ships this same file (canonical copy: shared/bands.ts in the repo).
2//   The order file, <configDir>/claude-code-mods/bands.json, lists mod names top to bottom
3//   (the last sits right above the prompt). Each band mod wraps its band in a Box keyed
4//   `band:<name>` and returns one column keyed `bands`; the next band mod unpacks that
5//   column, adds its own band and sorts again, so the last one to draw sorts them all.
6//   Edit shared/bands.ts, then copy it over every mod's hooks/bands.ts (a repo test checks).
7//   Reading the file needs $, which Claude Code follows only within one file, so each mod
8//   reads it in its own register file with bandsFile and parseOrder, once per session.
9import type { RenderElement, RenderNode } from 'claude-code'
10
11const COLUMN = 'bands'
12const PREFIX = 'band:'
13
14// Where the checklist saves the order: CLAUDE_CONFIG_DIR, or ~/.claude. Undefined with neither.
15export function bandsFile(configDir: string | undefined, home: string | undefined): string | undefined {
16  const dir = configDir ? configDir.replace(/\/+$/, '') : home ? `${home}/.claude` : ''
17  return dir ? `${dir}/claude-code-mods/bands.json` : undefined
18}
19
20// The names in an order file; [] when it is no JSON or not { order: string[] }.
21export function parseOrder(text: string): string[] {
22  try {
23    const order = (JSON.parse(text) as { order?: unknown } | null)?.order
24    return Array.isArray(order) && order.every(n => typeof n === 'string') ? order : []
25  } catch {
26    return []
27  }
28}
29
30/**
31 * One column of bands: `mine` wrapped in a Box keyed `band:<name>`, placed
32 * `at` (above or below) what `next(e)` returned, then sorted by `order`.
33 * With an order, items that are no band come first, then bands not in the
34 * order (in their current order), then the listed ones in file order. With
35 * no order nothing moves, so the screen looks as it did before.
36 */
37export function stackBands(rest: RenderElement | null | undefined, name: string, mine: RenderNode, at: 'above' | 'below', order: string[]): RenderElement {
38  const own: RenderElement = { type: 'Box', props: { key: PREFIX + name, flexDirection: 'column' }, children: [mine] }
39  const others = unpack(rest)
40  const items = at === 'above' ? [own, ...others] : [...others, own]
41  return { type: 'Box', props: { key: COLUMN, flexDirection: 'column' }, children: order.length ? sorted(items, order) : items }
42}
43
44// Another band mod's column gives up its items; anything else is one item.
45function unpack(rest: RenderElement | null | undefined): RenderNode[] {
46  if (!rest) return []
47  if (!isColumn(rest)) return [rest]
48  const kids = rest.children ?? []
49  const owner = rest.part?.plugin
50  // The engine stamped the column with its owner: that stamp belongs on the owner's own band.
51  return kids.map(k => (owner && typeof k !== 'string' && bandName(k) === owner && !('part' in k && k.part) ? { ...k, part: rest.part } : k))
52}
53
54function isColumn(el: RenderElement): el is RenderElement & { type: 'Box' } {
55  return el.type === 'Box' && el.props?.key === COLUMN
56}
57
58function bandName(node: RenderNode): string | undefined {
59  if (typeof node === 'string' || node.type !== 'Box') return undefined
60  const key = node.props?.key
61  return typeof key === 'string' && key.startsWith(PREFIX) ? key.slice(PREFIX.length) : undefined
62}
63
64function sorted(items: RenderNode[], order: string[]): RenderNode[] {
65  const rank = (n: RenderNode) => {
66    const name = bandName(n)
67    if (name === undefined) return -2 // not a band: first
68    const i = order.indexOf(name)
69    return i < 0 ? -1 : i // a band not in the order: on top of the listed ones
70  }
71  return items.map((n, i) => ({ n, i, r: rank(n) })).sort((a, b) => a.r - b.r || a.i - b.i).map(x => x.n)
72}
73

1export type NextSteps = { turnId: string; items: string[] }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'next-steps': { steps: NextSteps | null; active: boolean }
6  }
7}
8