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'
7
8const MODEL = 'haiku'
9const MAX_ITEMS = 3
10const MAX_CHARS = 60 // what the model is asked for
11const MAX_KEPT = 120 // a longer line is dropped, never cut: a cut draft would send half a prompt
12const MIN_ANSWER = 20 // a reply shorter than this is not worth a model call
13
14// Held by the host, so the list survives a hot reload of this file.
15const steps = atom({ plugin: 'next-steps', key: 'steps' } as const, null as NextSteps | null)
16// True while the list is on screen. where-am-i reads it and leaves out its own "next" meanwhile.
17const active = atom({ plugin: 'next-steps', key: 'active' } as const, false)
18
19export const register: Register = on => {
20  // The turn a suggestion may still be shown for. A submit or a new turn moves it on.
21  let latest = ''
22  let submits = 0
23  // A survey answers bare digits in the band: leave them to it.
24  let hasSurvey = false
25
26  on('prompt.submit', async ($, e, next) => {
27    latest = `submit-${++submits}`
28    await setSteps($, () => null)
29    return next(e)
30  })
31
32  on('turn.start', async ($, e, next) => {
33    latest = e.turnId
34    await setSteps($, () => null)
35    return next(e)
36  })
37
38  on('turn.complete', async ($, e, next) => {
39    const r = await next(e)
40    if (e.agentId) return r // a subagent's turn, not the person's
41    latest = e.turnId // also for a turn that raised no turn.start, and after a hot reload
42    if (e.reason !== 'answer' || e.answer.trim().length < MIN_ANSWER) return r
43    void suggest($, e.turnId, e.answer, () => latest).catch(() => {}) // in the background, so the turn ends at once
44    return r
45  })
46
47  // A digit typed into an empty prompt picks a suggestion instead of landing as text.
48  // A paste carries no key, so a pasted "1" lands as typed.
49  on('prompt.edit', async ($, e, next) => {
50    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)
51    const s = await read($, steps)
52    if (!s) return next(e)
53    if (e.inputText === '0') {
54      await setSteps($, () => null)
55      return { text: '', cursor: 0 } // consumed: the box stays empty
56    }
57    const pick = s.items[Number(e.inputText) - 1]
58    if (!pick) return next(e)
59    return next({ ...e, inputText: pick }) // a draft: the person edits it or presses Enter
60  })
61
62  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
63    const rest = await next(e) // what other mods and Claude Code draw here stays
64    hasSurvey = e.props.hasSurvey
65    const s = await read($, steps)
66    if (e.props.hasSurvey || e.props.isWorking || !s || s.items.length === 0) return rest
67
68    const { Box, Text } = $.ui.resolve(e)
69    const sep = <Text dimColor>{'  ·  '}</Text>
70    const item = (text: string, i: number) => [
71      <Text color="cyan" bold>{`${i + 1} `}</Text>,
72      <Text>{text}</Text>,
73    ]
74    const dismiss = [<Text color="cyan" bold>{'0 '}</Text>, <Text dimColor>dismiss</Text>]
75
76    const lines =
77      // One per row needs a row each; a short band gets the one row, cut at its end.
78      fitsOnOneLine(s.items, e.props.bodyColumns) || e.props.maxRows < s.items.length + 2 ? (
79        <Text wrap="truncate-end">
80          <Text dimColor>{'next:  '}</Text>
81          {s.items.flatMap((t, i) => [...item(t, i), sep])}
82          {dismiss}
83        </Text>
84      ) : (
85        <Box flexDirection="column">
86          {s.items.map((t, i) => (
87            <Text wrap="truncate-end">
88              <Text dimColor>{i === 0 ? 'next:  ' : '       '}</Text>
89              {item(t, i)}
90              {i === s.items.length - 1 && sep}
91              {i === s.items.length - 1 && dismiss}
92            </Text>
93          ))}
94        </Box>
95      )
96
97    return (
98      <Box flexDirection="column">
99        <Box paddingX={1}>{lines}</Box>
100        {rest}
101      </Box>
102    )
103  })
104}
105
106// Writes the list, then whether it shows, for where-am-i.
107async function setSteps($: EngineInterface, fn: (before: NextSteps | null) => NextSteps | null) {
108  const after = await update($, steps, fn)
109  const showing = after !== null && after.items.length > 0
110  if ((await read($, active)) !== showing) await update($, active, () => showing)
111}
112
113// The whole list on one row, or one suggestion per row when it would not fit.
114export function fitsOnOneLine(items: string[], columns: number) {
115  const width = 'next:  '.length + items.reduce((n, t) => n + 2 + t.length + 5, 0) + '0 dismiss'.length
116  return width <= columns - 2 // the band's padding
117}
118
119async function suggest($: EngineInterface, turnId: string, answer: string, current: () => string) {
120  // Wait for background agents: the main agent picks up again once they finish.
121  const agents = await $.agent.list().catch(() => [])
122  if (agents.some(a => a.status === 'running')) return
123
124  const all = (await $.session.messages().catch(() => [])).filter(m => m.text.trim() !== '')
125  if (all.at(-1)?.role === 'assistant') all.pop() // the latest reply goes in on its own below
126  const messages = all.slice(-6)
127  const r = await $.model.complete({
128    model: MODEL,
129    maxTokens: 200,
130    timeoutMs: 20_000,
131    system:
132      'You predict what a person will type next to their coding assistant, after reading a transcript of their session. ' +
133      `Give 2 or 3 short, concrete follow-up prompts, each an instruction in the imperative, under ${MAX_CHARS} characters, ` +
134      'specific to what just happened (name the file, test or feature). Plain words, no em dashes, no numbering, no quotes. ' +
135      'One per line and nothing else. If nothing sensible follows, reply with the single word NONE.\n' +
136      'Example reply:\nRun the login tests again\nAdd the same check to signup.ts\nOpen a draft PR',
137    prompt: [
138      `<transcript>\n${messages.map(m => `[${m.role === 'user' ? 'person' : 'assistant'}] ${m.text.slice(0, 800)}`).join('\n')}\n</transcript>`,
139      `<latest_reply>\n${answer.slice(0, 2000)}\n</latest_reply>`,
140      'Write the follow-up prompts the person is likely to send next: the lines only, not a reply to the transcript.',
141    ].join('\n\n'),
142  })
143  if (!r.isAnswered) return
144  const items = parseSteps(r.text)
145  // Checked inside the write, so a submit or a new turn that lands meanwhile always wins.
146  await setSteps($, before => (current() !== turnId ? before : items.length > 0 ? { turnId, items } : null))
147}
148
149// The model's lines, cleaned: no fences, bullets, numbers or em dashes; at most three.
150// A preamble, a "nothing to add" or a line too long to draft whole is dropped.
151export function parseSteps(text: string): string[] {
152  const body = text.replace(/```[a-z]*\n?/gi, '').trim()
153  let lines = body.split('\n')
154  if (body.startsWith('[') || body.startsWith('{')) {
155    try {
156      const o = JSON.parse(body) as unknown
157      const arr = Array.isArray(o) ? o : Object.values(o as Record<string, unknown>).find(Array.isArray)
158      lines = Array.isArray(arr) ? arr.filter((x): x is string => typeof x === 'string') : []
159    } catch {
160      // not JSON after all: read it as lines
161    }
162  }
163  const out: string[] = []
164  for (const line of lines) {
165    const clean = line
166      .replace(/\*\*/g, '')
167      .replace(/^\s*(?:[-*•]|\d+[.):])\s*/, '')
168      .replace(/^["'`]+|["'`,]+$/g, '')
169      .replace(/\s*\u2014\s*/g, ', ')
170      .replace(/\.$/, '')
171      .trim()
172    if (clean === '' || clean.endsWith(':') || clean.length > MAX_KEPT) continue
173    if (/^(none|no |nothing|here are|here's|sure)/i.test(clean) || /^[[\]{}]$/.test(clean)) continue
174    if (!out.some(o => o.toLowerCase() === clean.toLowerCase())) out.push(clean)
175    if (out.length === MAX_ITEMS) break
176  }
177  return out
178}
179

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