Suggest

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Buttons, LastReply, Message } from '../types'
5import {
6  clip,
7  formatTranscript,
8  invitesReply,
9  isAnswerTail,
10  isShellCommand,
11  mergeReplies,
12  normalizeSpace,
13  parseReplies,
14  parseSlashCommand,
15  replyKey,
16  systemPrompt,
17  toBareCommand,
18} from './text'
19
20// Messages kept for the model to read, prompts and answers together.
21const RECENT_MESSAGES = 8
22// Longest `/suggest <hint>` passed on to the model.
23const HINT_CHARS = 200
24const MODEL_TIMEOUT_MS = 20_000
25const EFFORTS = ['low', 'medium', 'high', 'xhigh', 'max'] as const
26// Prompts that don't come from the person aren't part of the conversation the model reads.
27const NON_PERSON_ORIGINS = new Set(['task-notification', 'scheduled-trigger', 'peer'])
28// Drawn dim after a `!` shell command that was put in the input without its `!`.
29const SHELL_TIP = '  (add ! at the beginning to run it in your shell, or press Enter to ask Claude)'
30
31// Kept in session state, not module variables, so a settings change (which reloads the mod) keeps them.
32const buttons = atom({ plugin: 'suggest', key: 'buttons' } as const, null as Buttons | null)
33const recent = atom({ plugin: 'suggest', key: 'recent' } as const, [] as Message[])
34const lastReply = atom({ plugin: 'suggest', key: 'lastReply' } as const, null as LastReply | null)
35
36type Settings = {
37  mode: 'auto' | 'smart' | 'manual'
38  isFillOnClick: boolean
39  maxReplies: number
40  model: string
41  effort: (typeof EFFORTS)[number]
42  contextChars: number
43  messageChars: number
44}
45
46// Bumped on every prompt and every suggestion run, so a model answer that arrives late never lands on newer buttons.
47type Generation = { current: number }
48
49// What the person asked for when buttons are shown.
50type Request = {
51  // Show only the free buttons, without calling the model.
52  isFree?: boolean
53  // Steers the model, from `/suggest <hint>`, and leaves Claude Code's own suggestion out.
54  hint?: string
55}
56
57const readSettings = (options: Parameters<Register>[1]): Settings => ({
58  mode: options.mode === 'auto' || options.mode === 'manual' ? options.mode : 'smart',
59  isFillOnClick: options.click === 'fill',
60  maxReplies: Math.max(1, Math.floor(Number(options.maxReplies) || 3)),
61  model: String(options.model || 'haiku'),
62  effort: EFFORTS.find(effort => effort === options.effort) ?? 'low',
63  contextChars: Math.max(500, Number(options.contextChars) || 6000),
64  messageChars: Math.max(200, Number(options.messageChars) || 2500),
65})
66
67async function rememberMessage($: EngineInterface, message: Message, messageChars: number) {
68  const clipped = { ...message, text: clip(message.text, messageChars) }
69  await update($, recent, messages => [...messages, clipped].slice(-RECENT_MESSAGES))
70}
71
72// How many replies to ask the model for; 0 means no model call.
73function modelReplyCount(settings: Settings, request: Request) {
74  if (request.isFree) return 0
75  // A single button is Claude Code's own suggestion, unless a hint replaces it.
76  if (settings.maxReplies === 1 && !request.hint) return 0
77
78  // Asks for every button, not just the ones left: Claude Code's suggestion may never come or may repeat
79  // one of the model's, and the merge puts it first and cuts the list to the limit anyway.
80  return settings.maxReplies
81}
82
83async function askModel($: EngineInterface, settings: Settings, prompt: string, count: number) {
84  const result = await $.model
85    .complete({
86      model: settings.model,
87      system: systemPrompt(count),
88      prompt,
89      maxTokens: settings.effort === 'low' ? 300 : 4000,
90      effort: settings.effort,
91      timeoutMs: MODEL_TIMEOUT_MS,
92    })
93    .catch(() => null)
94
95  return result?.isAnswered ? parseReplies(result.text).slice(0, count) : []
96}
97
98async function addModelReplies(
99  $: EngineInterface,
100  settings: Settings,
101  generation: Generation,
102  runId: number,
103  prompt: string,
104  count: number,
105) {
106  const replies = await askModel($, settings, prompt, count)
107  if (runId !== generation.current) return
108  await update($, buttons, shown =>
109    shown ? { ...shown, replies: mergeReplies(shown.replies, replies, settings.maxReplies), isLoading: false } : shown,
110  )
111
112  // Without a suggestion from Claude Code, the model's first reply becomes the dim one in the input.
113  // Taken with Tab, unlike a fill, a `!` shell command there switches on shell mode.
114  const engineSuggestion = (await read($, lastReply))?.engineSuggestion
115  if (!engineSuggestion && replies[0]) void $.prompt.suggest({ text: replies[0] })
116}
117
118// Shows buttons under the last reply: Claude Code's own suggestion right away, then the model's replies
119// once they come. Returns false when there is no reply to suggest for yet.
120async function showButtons($: EngineInterface, settings: Settings, generation: Generation, request: Request = {}) {
121  const reply = await read($, lastReply)
122  if (!reply) return false
123
124  generation.current += 1
125  const runId = generation.current
126  const free = reply.engineSuggestion && !request.hint ? [reply.engineSuggestion] : []
127  const count = modelReplyCount(settings, request)
128  await update($, buttons, () => ({ answer: reply.answer, replies: free, isLoading: count > 0 }))
129  if (count === 0) return true
130
131  const transcript = formatTranscript(await read($, recent), settings.contextChars, settings.messageChars)
132  const hint = request.hint ? `\n\nThe user asked for replies that are: ${request.hint.slice(0, HINT_CHARS)}` : ''
133  // Runs after this dispatch, so the hook returns at once and the model call can outlive it.
134  $.clock.after(0, () => addModelReplies($, settings, generation, runId, transcript + hint, count))
135
136  return true
137}
138
139// Puts a reply in the input to edit; Enter sends it. `isShellMode` says the person already typed `!`.
140async function fillInput($: EngineInterface, text: string, isShellMode = false) {
141  if (!isShellCommand(text)) {
142    await $.prompt.fill({ text })
143    return
144  }
145
146  const command = text.slice(1).trim()
147  // Inserted into the prompt shell mode already holds, the command keeps the mode, so Enter runs it. A replace
148  // would drop the mode and Enter would send the command to Claude. The hint line can lag behind the prompt, so
149  // the prompt is read too: an insert into a draft would land in the middle of it.
150  if (isShellMode && (await $.prompt.read()).text === '') {
151    await $.prompt.fill({ text: command, mode: 'insert' })
152    return
153  }
154
155  // A filled `!` can't switch on shell mode, but typing one does. So the command goes in without its `!`,
156  // with the caret at the start where the `!` goes, and a dim tip after it. The tip is not part of the draft
157  // and goes away at the next edit.
158  await $.prompt.fill({
159    text: command,
160    cursor: 0,
161    decorations: [
162      {
163        start: command.length,
164        end: command.length,
165        virtualText: SHELL_TIP,
166        color: 'subtle',
167        dimColor: true,
168        italic: true,
169      },
170    ],
171  })
172}
173
174// Running a command sends no prompt, so the input keeps showing Claude Code's dim suggestion. An edit clears it,
175// but only an empty input is touched: a draft the person typed stays.
176async function clearInputSuggestion($: EngineInterface) {
177  if ((await $.prompt.read()).text !== '') return
178  await $.prompt.fill({ text: ' ' })
179  await $.prompt.fill({ text: '' })
180}
181
182// Sends a reply as the person's own words, not as the plugin's message, or runs it when it is a slash command.
183async function sendReply($: EngineInterface, text: string, isShellMode = false) {
184  await update($, buttons, () => null)
185
186  // A plugin's prompt can't run in shell mode, so a `!` shell command goes to the input, never to Claude.
187  if (isShellCommand(text)) {
188    await fillInput($, text, isShellMode)
189    return
190  }
191
192  const command = parseSlashCommand(text)
193  if (!command) {
194    await $.prompt.submit({ text, asUser: true })
195    return
196  }
197
198  const ran = await $.command.run(command).then(
199    () => true,
200    () => false,
201  )
202  if (ran) {
203    await clearInputSuggestion($)
204  } else {
205    // An unknown command goes to the input to edit.
206    await $.prompt.fill({ text })
207  }
208}
209
210export const register: Register = (on, options) => {
211  const settings = readSettings(options)
212  const generation: Generation = { current: 0 }
213  // The prompt events never say whether the prompt is in shell mode, but the hint line under the empty prompt
214  // does: it reads "! for shell mode" while it is in it. Kept per surface, since each draws its own hint line.
215  const shellModeSurfaces = new Set<string>()
216
217  on('session.start', async ($, e, next) => {
218    // A reload drops a model call in flight, which would leave "more suggestions…" up forever.
219    await update($, buttons, shown => (shown?.isLoading ? { ...shown, isLoading: false } : shown))
220    // A status line stays pinned across a reload. This version pins none, so clear one an older version left.
221    $.ui.status(undefined)
222    await $.command.register({
223      name: 'suggest',
224      description: 'Show suggested replies now (optionally steered: /suggest shorter), or set the mode',
225      argumentHint: '[hint | auto | smart | manual]',
226    })
227
228    return next(e)
229  })
230
231  on('command.run', { command: 'suggest' }, async ($, e) => {
232    const arg = e.args.trim()
233    const mode = arg.toLowerCase()
234    if (mode === 'auto' || mode === 'smart' || mode === 'manual') {
235      const { deny } = await $.config.set({ key: 'suggest.mode', value: mode })
236
237      return { text: deny ? `Couldn't switch mode: ${deny}` : `Suggest: ${mode} mode` }
238    }
239
240    // Nothing goes to the session: the buttons show under the last reply.
241    const isShown = await showButtons($, settings, generation, { hint: arg || undefined })
242
243    return isShown ? {} : { text: 'No reply to suggest for yet.' }
244  })
245
246  on('prompt.submit', async ($, e, next) => {
247    // A sent prompt leaves shell mode, whether or not the hint line redraws before the next click.
248    shellModeSurfaces.clear()
249    // `/suggest` reads the last reply, so its own prompt must not clear it.
250    if (/^\/suggest\b/.test(e.text.trim())) return next(e)
251
252    generation.current += 1
253    // Errors are swallowed: the buttons must never hold up the prompt.
254    if (!NON_PERSON_ORIGINS.has(e.origin.kind) && e.text.trim()) {
255      await rememberMessage($, { role: 'user', text: e.text }, settings.messageChars).catch(() => undefined)
256    }
257    await update($, buttons, () => null).catch(() => undefined)
258    await update($, lastReply, () => null).catch(() => undefined)
259
260    return next(e)
261  })
262
263  on('turn.complete', async ($, e, next) => {
264    const result = await next(e)
265    if (e.agentId !== undefined || e.reason !== 'answer' || !e.answer.trim()) return result
266
267    await rememberMessage($, { role: 'assistant', text: e.answer }, settings.messageChars)
268    await update($, lastReply, () => ({ answer: e.answer }))
269    // Every mode shows the free suggestion. Auto also asks the model every time, smart only when the answer
270    // invites a reply, and manual never.
271    const isFree = settings.mode === 'manual' || (settings.mode === 'smart' && !invitesReply(e.answer))
272    await showButtons($, settings, generation, { isFree })
273
274    return result
275  })
276
277  // Claude Code's own dim suggestion in the input becomes the first button, even after the model's replies show.
278  // It is marked as Claude's, so it reads as such when it moves to the top.
279  on('prompt.suggest', { origin: { kind: 'suggestion' } }, async ($, e, next) => {
280    const text = toBareCommand(normalizeSpace(e.text))
281    if (text) {
282      await update($, lastReply, reply => (reply ? { ...reply, engineSuggestion: text } : reply))
283      await update($, buttons, shown =>
284        shown ? { ...shown, replies: mergeReplies([text], shown.replies, settings.maxReplies) } : shown,
285      )
286    }
287
288    return next(e)
289  })
290
291  on('ui.render', { component: 'PromptHint' }, ($, e, next) => {
292    if (e.props.hint.includes('shell mode') && !e.props.isDraft) shellModeSurfaces.add(e.surface)
293    else shellModeSurfaces.delete(e.surface)
294
295    return next(e)
296  })
297
298  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
299    const shown = await read($, buttons)
300    if (!shown || e.props.isSummary || !isAnswerTail(shown.answer, e.props.text)) return next(e)
301    if (!shown.isLoading && shown.replies.length === 0) return next(e)
302
303    const engineSuggestion = (await read($, lastReply))?.engineSuggestion
304    const isEngineSuggestion = (text: string) =>
305      engineSuggestion !== undefined && replyKey(engineSuggestion) === replyKey(text)
306    // Shell mode is read at the press, not the draw: the person may type `!` after the buttons show.
307    const isShellMode = () => shellModeSurfaces.has(e.surface)
308    const onPress = (text: string) => () =>
309      settings.isFillOnClick ? fillInput($, text, isShellMode()) : sendReply($, text, isShellMode())
310
311    const { Box, Button, Text } = $.ui.resolve(e)
312
313    return (
314      <Box flexDirection="column">
315        {await next(e)}
316        <Box flexDirection="column" marginTop={1} gap={1}>
317          {shown.replies.map((text, i) => (
318            <Box key={`row-${i}`} flexDirection="row" gap={1}>
319              {/* Claude Code's own suggestion gets a dim ✻; every other row a blank, so the buttons line up. */}
320              <Text dimColor>{isEngineSuggestion(text) ? '✻' : ' '}</Text>
321              <Button
322                key={`reply-${i}`}
323                label={text}
324                variant={i === 0 ? 'primary' : 'secondary'}
325                onPress={onPress(text)}
326              />
327              {settings.isFillOnClick ? null : (
328                <Button key={`edit-${i}`} label="✎" plain dimColor onPress={() => fillInput($, text, isShellMode())} />
329              )}
330            </Box>
331          ))}
332          {shown.isLoading ? (
333            <Text dimColor>{shown.replies.length ? 'more suggestions…' : 'suggesting replies…'}</Text>
334          ) : null}
335        </Box>
336      </Box>
337    )
338  })
339}
340

1import type { Message } from '../types'
2
3// Pure text helpers: the prompt the model reads, parsing its answer, and comparing replies.
4// Nothing here calls Claude Code, so each one can be tested on its own.
5
6// Longest reply the model is asked to write.
7const REPLY_CHARS = 80
8// How much of the end of the answer must match a rendered block for the buttons to go under it.
9const TAIL_MATCH_CHARS = 200
10// How much of the end of an answer is searched for a question or an offer.
11const INVITE_TAIL_CHARS = 600
12
13export const systemPrompt = (count: number) => `You predict what a user will type next to an AI coding assistant.
14Reply with ONLY a JSON array of 0 to ${count} strings, most likely first.
15Each string is a reply the user could send as-is: written as the user, in the user's language, under ${REPLY_CHARS} characters, no quotes or numbering.
16If the assistant asked a question or offered options, each reply answers it or picks one option. If there is an obvious next step, suggest approving it.
17If the assistant told the user to run a slash command, suggest the bare command itself, starting with "/" (like "/reload-plugins"), never "run /reload-plugins".
18If the next reply is not clear, return []. Fewer good replies beat more guesses.`
19
20// Cuts a long message down to its last `max` characters, because a question or a next step is usually at the end.
21export const clip = (text: string, max: number) => (text.length > max ? '…' + text.slice(-max) : text)
22
23// The recent conversation as the model reads it. Messages are taken newest first until `contextChars` runs out,
24// each clipped to `messageChars`, then shown oldest first.
25export const formatTranscript = (messages: readonly Message[], contextChars: number, messageChars: number) => {
26  const picked: string[] = []
27  let charsLeft = contextChars
28  for (const message of [...messages].reverse()) {
29    if (charsLeft <= 0) break
30    const text = clip(message.text.trim(), Math.min(messageChars, charsLeft))
31    picked.unshift(`<${message.role}>\n${text}\n</${message.role}>`)
32    charsLeft -= text.length
33  }
34
35  return picked.join('\n\n')
36}
37
38export const normalizeSpace = (text: string) => text.replace(/\s+/g, ' ').trim()
39
40// Whether a rendered message block is the last block of `answer`. A long answer renders as several blocks,
41// and the buttons belong under the last one.
42export const isAnswerTail = (answer: string, blockText: string) => {
43  const a = normalizeSpace(answer)
44  const block = normalizeSpace(blockText)
45  if (!a || !block) return false
46
47  return a === block || a.endsWith(block) || block.endsWith(a.slice(-TAIL_MATCH_CHARS))
48}
49
50// Turns an instruction to run a slash command into the command itself, so a click runs it:
51// "run /reload-plugins" and "Please try `/compact now`" become "/reload-plugins" and "/compact now".
52export const toBareCommand = (text: string) =>
53  /^(?:please\s+)?(?:run|type|use|try)\s+`?(\/[\w:-]+(?:\s[^`]*)?)`?$/i.exec(text)?.[1]?.trim() ?? text
54
55// A `!` shell command, as typed to run in shell mode. A bare `!` is not one.
56export const isShellCommand = (text: string) => /^!\s*\S/.test(text)
57
58// An empty reply, or a bare number like "2", says nothing on its own.
59export const isSendable = (text: string) => text !== '' && !/^\d+$/.test(text)
60
61// The model's answer as a list of replies. Anything that isn't a JSON array of strings gives none.
62export const parseReplies = (modelText: string): string[] => {
63  const json = modelText.match(/\[[\s\S]*\]/)?.[0]
64  if (!json) return []
65  try {
66    const list: unknown = JSON.parse(json)
67    if (!Array.isArray(list)) return []
68
69    return list
70      .filter((item): item is string => typeof item === 'string')
71      .map(normalizeSpace)
72      .map(toBareCommand)
73      .filter(isSendable)
74  } catch {
75    return []
76  }
77}
78
79// Splits "/name args" into the shape `$.command.run` takes, or gives null when the text is not a slash command.
80export const parseSlashCommand = (text: string) => {
81  const match = /^\/([\w:-]+)\s*([\s\S]*)$/.exec(text)
82  if (!match?.[1]) return null
83
84  return { command: match[1], args: (match[2] ?? '').trim() }
85}
86
87// Replies that differ only in case or punctuation get the same key, so they count as one reply.
88export const replyKey = (text: string) =>
89  text
90    .toLowerCase()
91    .replace(/[^\p{L}\p{N}/]+/gu, ' ')
92    .trim()
93
94// `first` then `rest`, without empty replies or repeats, cut to `max`.
95export const mergeReplies = (first: readonly string[], rest: readonly string[], max: number) => {
96  const seen = new Set<string>()
97
98  return [...first, ...rest]
99    .filter(isSendable)
100    .filter(reply => {
101      const key = replyKey(reply)
102      if (seen.has(key)) return false
103      seen.add(key)
104
105      return true
106    })
107    .slice(0, max)
108}
109
110// Whether the answer ends on a question or an offer, which is where suggested replies help most.
111export const invitesReply = (answer: string) =>
112  /\?|\b(should i|shall i|want me to|would you like|let me know|which (one|option)|or should)\b/i.test(
113    answer.slice(-INVITE_TAIL_CHARS),
114  )
115

1// The buttons under the last reply.
2export type Buttons = {
3  // The final text of the turn the buttons belong to, used to find the message block to draw them under.
4  answer: string
5  replies: string[]
6  // The model has been asked for more replies and hasn't answered yet.
7  isLoading: boolean
8}
9
10export type Message = { role: 'user' | 'assistant'; text: string }
11
12// The last main-loop reply, plus Claude Code's own suggestion for it once that arrives.
13export type LastReply = { answer: string; engineSuggestion?: string }
14
15declare module 'claude-code' {
16  interface PluginState {
17    suggest: { buttons: Buttons | null; recent: Message[]; lastReply: LastReply | null }
18  }
19}
20