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

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