next-steps

1// Next Steps: after each turn, 2 or 3 likely next prompts above the input.
2// In an empty prompt, the leader key and then 1 to 3 drafts one (nothing sends on its own),
3// the leader and 0 dismisses. A bare digit types as usual, so a stray 1 is never a pick.
4// Local fork of hamzafer/claude-code-mods next-steps 0.1.1: the leader key and the card are the changes.
5import { atom, read, update } from 'claude-code'
6import type { EngineInterface, Register } from 'claude-code'
7
8import type { NextSteps } from '../types'
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
15const LEADERS = [';', ',', '\\'] // what the person may set the leader to; ; by default
16const ROW_COLORS = ['#f472b6', '#a5b4fc', '#5eead4'] // suggestion 1, 2 and 3: bullet, key and compact chip
17const ROW_TINTS = ['#f9a8d4', '#c7d2fe', '#99f6e4'] // their text, a shade lighter (as where-am-i's card)
18const BORDER = '#fb7185'
19const TITLE: [string, string][] = [['Suggested ', '#fdba74'], ['next ', '#f472b6'], ['steps', '#a5b4fc']]
20const CARD_ROWS = 5 // the card's rows besides the suggestions: two borders, title, hint, dismiss
21const CARD_MIN_COLUMNS = 40 // narrower than this, or shorter than the card, the band draws one compact row
22const HEAD = 7 // '✦ next ' and '✦ pick ' in the compact row, and the indent under them
23
24// Held by the host, so the list survives a hot reload of this file.
25const steps = atom({ plugin: 'next-steps', key: 'steps' } as const, null as NextSteps | null)
26// True while the list is on screen. where-am-i reads it and leaves out its own "next" meanwhile.
27const active = atom({ plugin: 'next-steps', key: 'active' } as const, false)
28// True while the box holds the leader alone: the band then shows the digits to press.
29const armed = atom({ plugin: 'next-steps', key: 'armed' } as const, false)
30// The leader as /config set it, for where-am-i to draw the keys with.
31const leaderKey = atom({ plugin: 'next-steps', key: 'leader' } as const, ';')
32// where-am-i's value, read here: once it hosts the card, the suggestions are drawn inside it, never a card of their own.
33const whereHosts = { plugin: 'where-am-i', key: 'hosts' } as const
34
35export const register: Register = (on, options) => {
36  const leader = LEADERS.find(l => l === options.leader) ?? ';'
37  // The turn a suggestion may still be shown for. A submit or a new turn moves it on.
38  let latest = ''
39  let submits = 0
40  // A survey answers bare digits in the band: leave them to it.
41  let hasSurvey = false
42  // What `armed` was last written as, so a keystroke that changes nothing writes nothing.
43  const arm = { now: false }
44
45  on('session.start', async ($, e, next) => {
46    const r = await next(e)
47    await update($, leaderKey, () => leader).catch(() => {})
48    return r
49  })
50
51  on('prompt.submit', async ($, e, next) => {
52    latest = `submit-${++submits}`
53    await setSteps($, () => null)
54    await setArmed($, arm, false)
55    return next(e)
56  })
57
58  on('turn.start', async ($, e, next) => {
59    latest = e.turnId
60    await setSteps($, () => null)
61    return next(e)
62  })
63
64  on('turn.complete', async ($, e, next) => {
65    const r = await next(e)
66    if (e.agentId) return r // a subagent's turn, not the person's
67    latest = e.turnId // also for a turn that raised no turn.start, and after a hot reload
68    if (e.reason !== 'answer' || e.answer.trim().length < MIN_ANSWER) return r
69    void suggest($, e.turnId, e.answer, () => latest).catch(() => {}) // in the background, so the turn ends at once
70    return r
71  })
72
73  // The leader typed alone into an empty prompt lands as text; a digit right after it picks a
74  // suggestion in its place. A paste carries no key, so a pasted ";1" lands as typed.
75  on('prompt.edit', async ($, e, next) => {
76    const n = leader.length
77    const isPick = !hasSurvey && !!e.key && !e.key.ctrl && !e.key.meta && e.text === leader && e.start === n && e.end === n && /^[0-9]$/.test(e.inputText)
78    const s = isPick ? await read($, steps) : null
79    if (s && e.inputText === '0') {
80      await setSteps($, () => null)
81      await setArmed($, arm, false)
82      return { text: '', cursor: 0 } // consumed: the leader goes too, the box is empty again
83    }
84    const pick = s?.items[Number(e.inputText) - 1]
85    if (pick) {
86      await setArmed($, arm, false)
87      return { text: pick, cursor: pick.length } // the leader becomes the draft: the person edits it or presses Enter
88    }
89    const r = await next(e)
90    await setArmed($, arm, r.text === leader)
91    return r
92  })
93
94  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
95    const rest = await next(e) // what other mods and Claude Code draw here stays
96    hasSurvey = e.props.hasSurvey
97    const s = await read($, steps)
98    if (e.props.hasSurvey || e.props.isWorking || !s || s.items.length === 0) return rest
99    const { value: hosted = false } = await $.state.get(whereHosts)
100    if (hosted) return rest // where-am-i draws them in its card
101
102    const isArmed = await read($, armed)
103    const { Box, Text } = $.ui.resolve(e)
104    // With the leader typed, the keys drop it: only the digit is left to press.
105    const key = (n: number) => (isArmed ? `${n}` : `${leader}${n}`)
106
107    // The card: a bordered box, a title, a hint, one row per suggestion with its key on the right.
108    if (e.props.maxRows >= s.items.length + CARD_ROWS && e.props.bodyColumns >= CARD_MIN_COLUMNS) {
109      const digits = s.items.map((_, i) => `${i + 1}`)
110      const choices = digits.length > 1 ? `${digits.slice(0, -1).join(', ')} or ${digits.at(-1)}` : (digits[0] ?? '1')
111      const hint = isArmed
112        ? `press ${choices} to draft it · 0 dismisses`
113        : `${s.items.length} idea${s.items.length === 1 ? '' : 's'} · type ${leader} then a number`
114      const row = (bullet: JSX.Element, text: JSX.Element, keyText: JSX.Element) => (
115        <Box flexDirection="row">
116          <Box flexGrow={1} flexShrink={1}>
117            <Text wrap="truncate-end">
118              {bullet}
119              {text}
120            </Text>
121          </Box>
122          <Box flexShrink={0} marginLeft={2}>
123            {keyText}
124          </Box>
125        </Box>
126      )
127      return (
128        <Box flexDirection="column">
129          <Box flexDirection="column" borderStyle="round" borderColor={isArmed ? 'warning' : BORDER} paddingX={1} width={e.props.bodyColumns}>
130            <Text wrap="truncate-end">
131              <Text color="#f472b6">{'✧ '}</Text>
132              {TITLE.map(([word, color]) => (
133                <Text color={color} bold>
134                  {word}
135                </Text>
136              ))}
137            </Text>
138            <Text wrap="truncate-end" {...(isArmed ? { color: 'warning', bold: true } : { dimColor: true })}>
139              {hint}
140            </Text>
141            {s.items.map((t, i) =>
142              row(
143                <Text color={ROW_COLORS[i] ?? BORDER}>{'● '}</Text>,
144                <Text color={ROW_TINTS[i] ?? BORDER}>{t}</Text>,
145                <Text color={ROW_COLORS[i] ?? BORDER} bold>
146                  {key(i + 1)}
147                </Text>,
148              ),
149            )}
150            {row(<Text dimColor>{'○ '}</Text>, <Text dimColor>{'Dismiss'}</Text>, <Text dimColor>{key(0)}</Text>)}
151          </Box>
152          {rest}
153        </Box>
154      )
155    }
156
157    // The compact row, for a short or narrow band.
158    const head = isArmed ? <Text color="warning" bold>{'✦ pick '}</Text> : <Text color="claude" bold>{'✦ next '}</Text>
159    const sep = <Text>{'   '}</Text>
160    const item = (text: string, i: number) => [
161      <Text backgroundColor={ROW_COLORS[i] ?? BORDER} color="black" bold>{` ${key(i + 1)} `}</Text>,
162      <Text>{' '}</Text>,
163      <Text color={ROW_TINTS[i] ?? BORDER}>{text}</Text>,
164    ]
165    const dismiss = [<Text inverse dimColor>{` ${key(0)} `}</Text>, <Text dimColor>{' dismiss'}</Text>]
166
167    const lines =
168      // One per row needs a row each; a short band gets the one row, cut at its end.
169      fitsOnOneLine(s.items, e.props.bodyColumns, leader) || e.props.maxRows < s.items.length + 2 ? (
170        <Text wrap="truncate-end">
171          {head}
172          {s.items.flatMap((t, i) => [...item(t, i), sep])}
173          {dismiss}
174        </Text>
175      ) : (
176        <Box flexDirection="column">
177          {s.items.map((t, i) => (
178            <Text wrap="truncate-end">
179              {i === 0 ? head : <Text>{' '.repeat(HEAD)}</Text>}
180              {item(t, i)}
181              {i === s.items.length - 1 && sep}
182              {i === s.items.length - 1 && dismiss}
183            </Text>
184          ))}
185        </Box>
186      )
187
188    return (
189      <Box flexDirection="column">
190        <Box paddingX={1}>{lines}</Box>
191        {rest}
192      </Box>
193    )
194  })
195}
196
197// Writes the list, then whether it shows, for where-am-i.
198async function setSteps($: EngineInterface, fn: (before: NextSteps | null) => NextSteps | null) {
199  const after = await update($, steps, fn)
200  const showing = after !== null && after.items.length > 0
201  if ((await read($, active)) !== showing) await update($, active, () => showing)
202}
203
204// Whether the band holds the leader alone, written only when that changes.
205async function setArmed($: EngineInterface, arm: { now: boolean }, value: boolean) {
206  if (arm.now === value) return
207  arm.now = value
208  await update($, armed, () => value)
209}
210
211// The whole list on one row, or one suggestion per row when it would not fit.
212export function fitsOnOneLine(items: string[], columns: number, leader = ';') {
213  const chip = leader.length + 3 // ' ;1 '
214  const width = HEAD + items.reduce((n, t) => n + chip + 1 + t.length + 3, 0) + chip + ' dismiss'.length
215  return width <= columns - 2 // the band's padding
216}
217
218async function suggest($: EngineInterface, turnId: string, answer: string, current: () => string) {
219  // Wait for background agents: the main agent picks up again once they finish.
220  const agents = await $.agent.list().catch(() => [])
221  if (agents.some(a => a.status === 'running')) return
222
223  const all = (await $.session.messages().catch(() => [])).filter(m => m.text.trim() !== '')
224  if (all.at(-1)?.role === 'assistant') all.pop() // the latest reply goes in on its own below
225  const messages = all.slice(-6)
226  const r = await $.model.complete({
227    model: MODEL,
228    maxTokens: 200,
229    timeoutMs: 20_000,
230    system:
231      'You predict what a person will type next to their coding assistant, after reading a transcript of their session. ' +
232      `Give 2 or 3 short, concrete follow-up prompts, each an instruction in the imperative, under ${MAX_CHARS} characters, ` +
233      'specific to what just happened (name the file, test or feature). Plain words, no em dashes, no numbering, no quotes. ' +
234      'One per line and nothing else. If nothing sensible follows, reply with the single word NONE.\n' +
235      'Example reply:\nRun the login tests again\nAdd the same check to signup.ts\nOpen a draft PR',
236    prompt: [
237      `<transcript>\n${messages.map(m => `[${m.role === 'user' ? 'person' : 'assistant'}] ${m.text.slice(0, 800)}`).join('\n')}\n</transcript>`,
238      `<latest_reply>\n${answer.slice(0, 2000)}\n</latest_reply>`,
239      'Write the follow-up prompts the person is likely to send next: the lines only, not a reply to the transcript.',
240    ].join('\n\n'),
241  })
242  if (!r.isAnswered) return
243  const items = parseSteps(r.text)
244  // Checked inside the write, so a submit or a new turn that lands meanwhile always wins.
245  await setSteps($, before => (current() !== turnId ? before : items.length > 0 ? { turnId, items } : null))
246}
247
248// The model's lines, cleaned: no fences, bullets, numbers or em dashes; at most three.
249// A preamble, a "nothing to add" or a line too long to draft whole is dropped.
250export function parseSteps(text: string): string[] {
251  const body = text.replace(/```[a-z]*\n?/gi, '').trim()
252  let lines = body.split('\n')
253  if (body.startsWith('[') || body.startsWith('{')) {
254    try {
255      const o = JSON.parse(body) as unknown
256      const arr = Array.isArray(o) ? o : Object.values(o as Record<string, unknown>).find(Array.isArray)
257      lines = Array.isArray(arr) ? arr.filter((x): x is string => typeof x === 'string') : []
258    } catch {
259      // not JSON after all: read it as lines
260    }
261  }
262  const out: string[] = []
263  for (const line of lines) {
264    const clean = line
265      .replace(/\*\*/g, '')
266      .replace(/^\s*(?:[-*•]|\d+[.):])\s*/, '')
267      .replace(/^["'`]+|["'`,]+$/g, '')
268      .replace(/\s*\u2014\s*/g, ', ')
269      .replace(/\.$/, '')
270      .trim()
271    if (clean === '' || clean.endsWith(':') || clean.length > MAX_KEPT) continue
272    if (/^(none|no |nothing|here are|here's|sure)/i.test(clean) || /^[[\]{}]$/.test(clean)) continue
273    if (!out.some(o => o.toLowerCase() === clean.toLowerCase())) out.push(clean)
274    if (out.length === MAX_ITEMS) break
275  }
276  return out
277}
278

1export type NextSteps = { turnId: string; items: string[] }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'next-steps': { steps: NextSteps | null; active: boolean; armed: boolean; leader: string }
6    // Owned by the where-am-i mod, read here. Kept the same as its own contract.
7    'where-am-i': { recap: { goal: string; now: string; waiting: string; next: string } | null; live: string; hosts: boolean }
8  }
9}
10