Next Steps

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Suggestion } from '../types'
5
6const suggestions = atom({ plugin: 'next-steps', key: 'suggestions' } as const, [] as Suggestion[])
7// What the person last asked, so the suggestions follow from it.
8const lastPrompt = atom({ plugin: 'next-steps', key: 'lastPrompt' } as const, '')
9// True while the suggestions for the last reply are being written.
10const loading = atom({ plugin: 'next-steps', key: 'loading' } as const, false)
11// Which suggestions are ticked, by index; Send sends them together.
12const selected = atom({ plugin: 'next-steps', key: 'selected' } as const, [] as number[])
13// Claude's last answer, kept so Send can write the prompt with full context.
14const lastAnswer = atom({ plugin: 'next-steps', key: 'lastAnswer' } as const, '')
15// True while Send is writing the prompt.
16const composing = atom({ plugin: 'next-steps', key: 'composing' } as const, false)
17
18// Writes the message Send submits: a real, specific request built from the ticked
19// steps and the conversation, not a one-line restatement of the label.
20const COMPOSE_SYSTEM =
21  "You write the next message a user sends to an AI coding assistant. You are given the user's " +
22  "last request, the assistant's last answer, and one or more next steps the user picked. Write " +
23  'that message as the user, in the first person, ready to send. For each step: say exactly what to ' +
24  'do, name the specific files, functions, commands, URLs or values from the conversation it ' +
25  'involves, give any constraints or context the assistant needs, and say what done looks like. ' +
26  'With several steps, number them in the order given and say to finish each before the next. ' +
27  'Be concrete and concise: no greetings, no filler, no restating the whole conversation. Output ' +
28  'only the message.'
29
30const MAX = 6
31
32// The model that writes the message Send submits. 'haiku' is near-instant; 'sonnet' is
33// slower but more specific. Change this one line to switch.
34const SEND_MODEL = 'haiku'
35
36const SYSTEM =
37  'You suggest what a user might ask a coding assistant to do next. Reply with only a JSON array ' +
38  `of up to ${MAX} objects, each {"label": string}. ` +
39  '"label" is the button text: imperative, at most 5 words and under 30 characters, specific to this conversation. ' +
40  'Make the suggestions distinct from each other. ' +
41  'If nothing useful follows, reply [].'
42
43const clip = (s: string, max: number) => (s.length > max ? s.slice(0, max) + '…' : s)
44const words = (s: string, max: number) => s.split(/\s+/).slice(0, max).join(' ')
45
46// Button labels must stay short or the band overflows: at most 30 characters, cut at
47// a word boundary with an ellipsis.
48const LABEL_CHARS = 30
49function shortLabel(s: string): string {
50  const w = s.replace(/\s+/g, ' ').trim()
51  if (w.length <= LABEL_CHARS) return w
52  const cut = w.slice(0, LABEL_CHARS - 1)
53  const space = cut.lastIndexOf(' ')
54  return (space > 12 ? cut.slice(0, space) : cut).replace(/[\s,.;:-]+$/, '') + '…'
55}
56
57function parseList(text: string): Suggestion[] {
58  const m = text.match(/\[[\s\S]*\]/)
59  if (!m) return []
60  try {
61    const list = JSON.parse(m[0])
62    if (!Array.isArray(list)) return []
63    return list
64      .filter(x => x && typeof x.label === 'string' && x.label.trim() !== '')
65      .map(x => {
66        const label = shortLabel(x.label.trim())
67        const prompt = typeof x.prompt === 'string' && x.prompt.trim() !== '' ? x.prompt.trim() : label
68        return { label, prompt }
69      })
70      .slice(0, MAX)
71  } catch {
72    return []
73  }
74}
75
76// ---------- cache: reuse a model reply when the exact same request repeats ----------
77
78type CacheEntry = { text: string; at: number }
79
80// A small LRU with expiry. The key is the whole request (model, instructions, the full
81// prompt with its context, limits), so any change in context is a different key and a
82// stale reply can never come back.
83function makeCache(limit: number, ttlMs: number, now: () => number = () => Date.now()) {
84  const entries = new Map<string, CacheEntry>()
85  return {
86    get(key: string): string | undefined {
87      const hit = entries.get(key)
88      if (!hit) return undefined
89      if (now() - hit.at > ttlMs) {
90        entries.delete(key)
91        return undefined
92      }
93      entries.delete(key) // refresh its place as most recently used
94      entries.set(key, hit)
95      return hit.text
96    },
97    set(key: string, text: string) {
98      entries.delete(key)
99      entries.set(key, { text, at: now() })
100      while (entries.size > limit) entries.delete(entries.keys().next().value as string)
101    },
102    get size() {
103      return entries.size
104    },
105  }
106}
107
108type CompleteRequest = {
109  model: string
110  system: string
111  prompt: string
112  maxTokens: number
113  effort: 'low'
114  timeoutMs: number
115}
116
117// 40 replies, kept 30 minutes. Lives in memory: a reload starts it empty.
118const replyCache = makeCache(40, 30 * 60_000)
119const cacheStats = { hits: 0, modelCalls: 0 }
120
121const cacheKey = (r: CompleteRequest) =>
122  JSON.stringify([r.model, r.system, r.prompt, r.maxTokens, r.effort])
123
124// Answers from the cache when the identical request was answered before; otherwise asks
125// the model and keeps a good answer. Only answered replies are cached, never errors.
126async function completeCached($: EngineInterface, req: CompleteRequest): Promise<string | undefined> {
127  const key = cacheKey(req)
128  const cached = replyCache.get(key)
129  if (cached !== undefined) {
130    cacheStats.hits++
131    return cached
132  }
133  cacheStats.modelCalls++
134  const r = await $.model.complete(req)
135  if (!r.isAnswered || !r.text.trim()) return undefined
136  replyCache.set(key, r.text)
137  return r.text
138}
139
140async function suggest($: EngineInterface, answer: string) {
141  const [asked] = await Promise.all([
142    read($, lastPrompt),
143    update($, lastAnswer, () => answer),
144    update($, loading, () => true),
145  ])
146  // Labels only (the full prompt is written at Send), at low effort: a short, fast reply.
147  const reply = await completeCached($, {
148    model: 'haiku',
149    system: SYSTEM,
150    prompt: `User asked:\n${clip(asked, 3000)}\n\nAssistant replied:\n${clip(answer, 6000)}\n\nNext steps (JSON array):`,
151    maxTokens: 350,
152    effort: 'low',
153    timeoutMs: 15000,
154  })
155  await Promise.all([
156    update($, suggestions, () => (reply ? parseList(reply) : [])),
157    update($, selected, () => []),
158    update($, loading, () => false),
159  ])
160}
161
162async function toggle($: EngineInterface, i: number) {
163  await update($, selected, l => (l.includes(i) ? l.filter(x => x !== i) : [...l, i].sort((a, b) => a - b)))
164}
165
166// Sends the ticked steps as one well-written prompt, composed by SEND_MODEL with the
167// conversation as context; falls back to the plain labels if that fails.
168async function sendSelected($: EngineInterface) {
169  const [list, ticked, busy, asked, answer] = await Promise.all([
170    read($, suggestions),
171    read($, selected),
172    read($, composing),
173    read($, lastPrompt),
174    read($, lastAnswer),
175  ])
176  const picked = ticked.map(i => list[i]).filter(Boolean)
177  if (picked.length === 0 || busy) return
178  // Fallback if the writer fails: the labels as a plain list.
179  let text =
180    picked.length === 1
181      ? picked[0].label
182      : 'Please do the following, in this order:\n' + picked.map((s, i) => `${i + 1}. ${s.label}`).join('\n')
183  await update($, composing, () => true)
184  try {
185    const steps = picked.map((s, i) => `${i + 1}. ${s.label}`).join('\n')
186    const reply = await completeCached($, {
187      model: SEND_MODEL,
188      system: COMPOSE_SYSTEM,
189      prompt:
190        `User's last request:\n${clip(asked, 4000)}\n\n` +
191        `Assistant's last answer:\n${clip(answer, 10000)}\n\n` +
192        `Next steps the user picked, in order:\n${steps}\n\nThe message:`,
193      maxTokens: 900,
194      effort: 'low',
195      timeoutMs: 25000,
196    })
197    if (reply) text = reply.trim()
198  } catch {
199    // keep the fallback
200  }
201  await Promise.all([
202    update($, composing, () => false),
203    update($, suggestions, () => []),
204    update($, selected, () => []),
205  ])
206  await $.prompt.submit({ text, asUser: true })
207}
208
209
210// ---------- /mod-bench: timing measurements, run on demand ----------
211
212// Scratch values the storage benchmark writes, so it never touches real state.
213const benchA = atom({ plugin: 'next-steps', key: 'benchA' } as const, 0)
214const benchB = atom({ plugin: 'next-steps', key: 'benchB' } as const, 0)
215const benchC = atom({ plugin: 'next-steps', key: 'benchC' } as const, 0)
216
217// The module's own clock: no round-trip to the host, so short timings stay honest.
218const tick = () => (typeof performance !== 'undefined' ? performance.now() : Date.now())
219
220// How long each real message waited inside this mod before going out, newest last.
221const submitWaits: number[] = []
222
223type Stat = { runs: number; min: number; median: number; mean: number; max: number }
224function stat(xs: number[]): Stat {
225  const a = [...xs].sort((x, y) => x - y)
226  const r = (n: number) => Math.round(n * 100) / 100
227  return {
228    runs: a.length,
229    min: r(a[0] ?? 0),
230    median: r(a[Math.floor((a.length - 1) / 2)] ?? 0),
231    mean: r(a.reduce((x, y) => x + y, 0) / (a.length || 1)),
232    max: r(a[a.length - 1] ?? 0),
233  }
234}
235
236async function timeIt(n: number, fn: () => Promise<unknown>): Promise<Stat> {
237  const xs: number[] = []
238  for (let i = 0; i < n; i++) {
239    const t0 = tick()
240    await fn()
241    xs.push(tick() - t0)
242  }
243  return stat(xs)
244}
245
246async function runBench($: EngineInterface): Promise<string> {
247  const N = 20
248  const bump = (n: number) => n + 1
249  const writesSequential = await timeIt(N, async () => {
250    await update($, benchA, bump)
251    await update($, benchB, bump)
252    await update($, benchC, bump)
253  })
254  const writesBatched = await timeIt(N, () =>
255    Promise.all([update($, benchA, bump), update($, benchB, bump), update($, benchC, bump)]),
256  )
257  const readsSequential = await timeIt(N, async () => {
258    await read($, suggestions)
259    await read($, selected)
260    await read($, composing)
261    await read($, loading)
262  })
263  const readsBatched = await timeIt(N, () =>
264    Promise.all([read($, suggestions), read($, selected), read($, composing), read($, loading)]),
265  )
266
267  // Model calls on the current conversation: labels, then Send written by each model.
268  const [asked, answer, list] = await Promise.all([read($, lastPrompt), read($, lastAnswer), read($, suggestions)])
269  const labels = list.filter(x => x && typeof x === 'object' && x.label).slice(0, 2).map(x => x.label)
270  const steps = (labels.length ? labels : ['Summarize what changed in this session'])
271    .map((l, i) => `${i + 1}. ${l}`)
272    .join('\n')
273  const composeInput =
274    `User's last request:\n${clip(asked, 4000)}\n\n` +
275    `Assistant's last answer:\n${clip(answer, 10000)}\n\n` +
276    `Next steps the user picked, in order:\n${steps}\n\nThe message:`
277  const timed = async (model: string, system: string, prompt: string, maxTokens: number) => {
278    const t0 = tick()
279    const r = await $.model.complete({ model, system, prompt, maxTokens, effort: 'low', timeoutMs: 60000 })
280    return {
281      model,
282      ms: Math.round(tick() - t0),
283      outputTokens: r.usage.output_tokens,
284      text: r.isAnswered ? r.text : `(no reply: ${r.reason})`,
285    }
286  }
287  const labelsCall = await timed(
288    'haiku',
289    SYSTEM,
290    `User asked:\n${clip(asked, 3000)}\n\nAssistant replied:\n${clip(answer, 6000)}\n\nNext steps (JSON array):`,
291    350,
292  )
293  const sendHaiku = await timed('haiku', COMPOSE_SYSTEM, composeInput, 900)
294  const sendSonnet = await timed('sonnet', COMPOSE_SYSTEM, composeInput, 900)
295
296  const result = {
297    measuredAt: new Date().toISOString(),
298    storage: { writesSequential, writesBatched, readsSequential, readsBatched },
299    realMessagesWaitedInMod: stat(submitWaits),
300    cache: { entries: replyCache.size, hits: cacheStats.hits, modelCalls: cacheStats.modelCalls },
301    models: { labels: labelsCall, sendHaiku, sendSonnet },
302    stepsUsed: steps,
303  }
304  await $.fs.write(`${$.plugin.root}/bench.json`, JSON.stringify(result, null, 2))
305  const ms = (x: Stat) => `${x.median} ms median`
306  return [
307    'mod-bench done (full results in next-steps/bench.json):',
308    `  3 writes, one at a time: ${ms(writesSequential)}  ·  batched: ${ms(writesBatched)}`,
309    `  4 reads, one at a time: ${ms(readsSequential)}  ·  batched: ${ms(readsBatched)}`,
310    `  cache: ${cacheStats.hits} hits, ${cacheStats.modelCalls} model calls, ${replyCache.size} entries`,
311    `  Haiku labels: ${labelsCall.ms} ms  ·  Send with Haiku: ${sendHaiku.ms} ms  ·  Send with Sonnet: ${sendSonnet.ms} ms`,
312  ].join('\n')
313}
314
315export const register: Register = on => {
316  on('session.start', async ($, e, next) => {
317    const result = await next(e)
318    await $.command.register({ name: 'mod-bench', description: 'Time next-steps storage and model calls' })
319    return result
320  })
321
322  on('command.run', { command: 'mod-bench' }, async $ => ({ text: await runBench($) }))
323
324  on('prompt.submit', async ($, e, next) => {
325    const t0 = tick()
326    // A new request makes the old suggestions stale.
327    if (!e.turnId) {
328      // In the background: the message goes out without waiting on these.
329      void Promise.all([
330        update($, lastPrompt, () => e.text),
331        update($, suggestions, () => []),
332        update($, selected, () => []),
333      ])
334    }
335    submitWaits.push(tick() - t0)
336    if (submitWaits.length > 50) submitWaits.shift()
337    return next(e)
338  })
339
340  on('turn.complete', async ($, e, next) => {
341    const result = await next(e)
342    if (!e.agentId && e.reason === 'answer' && e.answer.trim()) {
343      void suggest($, e.answer).catch(() => update($, loading, () => false))
344    }
345    return result
346  })
347
348  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
349    const inner = await next(e)
350    if (e.props.hasSurvey) return inner
351    // Skip anything saved in an older format (plain strings) before a reload.
352    const [saved, picked, isComposing, isLoading] = await Promise.all([
353      read($, suggestions),
354      read($, selected),
355      read($, composing),
356      read($, loading),
357    ])
358    const list = saved.filter(x => x && typeof x === 'object' && x.label)
359    const { Box, Text, Button } = $.ui.resolve(e)
360
361    if (e.props.isWorking || list.length === 0) {
362      // Placeholder, so the spot never looks broken or empty.
363      const hint = e.props.isWorking
364        ? 'Next steps appear when Claude finishes'
365        : isLoading
366          ? 'Finding next steps…'
367          : 'No next steps yet'
368      return (
369        <Box flexDirection="row" alignItems="center" gap={1}>
370          <Text dimColor>{hint}</Text>
371          <Box flexGrow={1}>{inner}</Box>
372        </Box>
373      )
374    }
375
376    const item = (s: Suggestion, i: number) => (
377      <Button
378        key={`n${i}`}
379        label={`${picked.includes(i) ? '☑' : '☐'}  ${s.label}`}
380        dimColor={!picked.includes(i)}
381        onPress={() => toggle($, i)}
382      />
383    )
384    // Two columns, filled down then across: 1-3 left, 4-6 right.
385    const half = Math.ceil(list.length / 2)
386    return (
387      <Box flexDirection="row" alignItems="center" gap={1}>
388        {/* Shrinks and clips if ever too wide, so the right side is never pushed off. */}
389        <Box flexDirection="column" alignItems="flex-start" gap={1} flexShrink={1} minWidth={0} overflow="hidden">
390          <Box flexDirection="row" alignItems="flex-start" columnGap={2}>
391            <Box flexDirection="column" alignItems="flex-start" gap={1}>
392              {list.slice(0, half).map((s, i) => item(s, i))}
393            </Box>
394            <Box flexDirection="column" alignItems="flex-start" gap={1}>
395              {list.slice(half).map((s, i) => item(s, i + half))}
396            </Box>
397          </Box>
398          {/* Footer: what is ticked, a way to undo it, and the action, side by side. */}
399          <Box flexDirection="row" alignItems="center" gap={1}>
400            {isComposing ? (
401              <Text dimColor>Writing prompt…</Text>
402            ) : picked.length > 0 ? (
403              <Box flexDirection="row" alignItems="center" gap={1}>
404                <Button
405                  key="send"
406                  label={picked.length === 1 ? 'Send' : `Send ${picked.length}`}
407                  variant="primary"
408                  onPress={() => sendSelected($)}
409                />
410                <Button key="clear" label="Clear" dimColor onPress={() => update($, selected, () => [])} />
411                <Text dimColor>
412                  {picked.length} of {list.length} selected
413                </Text>
414              </Box>
415            ) : (
416              <Text dimColor>Tick one or more, then Send</Text>
417            )}
418          </Box>
419        </Box>
420        <Box flexGrow={1} flexShrink={0}>{inner}</Box>
421      </Box>
422    )
423  })
424}
425

1/** A suggested next step: the short button label and the full prompt it sends. */
2export type Suggestion = { label: string; prompt: string }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'next-steps': {
7      suggestions: Suggestion[]
8      lastPrompt: string
9      lastAnswer: string
10      loading: boolean
11      composing: boolean
12      selected: number[]
13      benchA: number
14      benchB: number
15      benchC: number
16    }
17  }
18}
19