pill-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: 'pill-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: 'pill-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  // 點膠囊：數字或文字都把那句放進輸入框當草稿（不會直接送出）；點「略過」收起
63  on('ui.press', { plugin: 'pill-steps' }, async ($, e, next) => {
64    const m = /^step:(\d+)(?::text)?$/.exec(e.element)
65    if (e.element === 'step:dismiss' || e.element === 'step:dismiss:text') {
66      await setSteps($, () => null)
67      return { element: e.element }
68    }
69    if (!m) return next(e)
70    const s = await read($, steps)
71    const pick = s?.items[Number(m[1])]
72    if (!pick) return next(e)
73    await $.prompt.fill({ text: pick })
74    return { element: e.element }
75  })
76
77  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
78    const rest = await next(e) // what other mods and Claude Code draw here stays
79    hasSurvey = e.props.hasSurvey
80    const s = await read($, steps)
81    if (e.props.hasSurvey || e.props.isWorking || !s || s.items.length === 0) return rest
82
83    const { Box, Text, Button } = $.ui.resolve(e)
84    // 圓角膠囊（Nerd Font 的  ）：Claude 橘的數字＋深墨底米白字；滑鼠移上去整顆變亮
85    const ORANGE = '#D97757'
86    const ORANGE_HI = '#EE9677'
87    const INK_HI = '#454860'
88    const GREY = '#5B5F6E'
89    const GREY_HI = '#7A7F90'
90    // 一列一個：數字是小圓角膠囊，文字不上底色（疊起來才不會黏成一塊）；滑鼠移上去整列亮起
91    const row = (id: string, key: string, text: string, dim = false) => {
92      const head = dim ? GREY : ORANGE
93      const headHi = dim ? GREY_HI : ORANGE_HI
94      return (
95        <Box key={`row:${id}`} flexDirection="row">
96          <Text color={head} hover={{ color: headHi }}>{'\uE0B6'}</Text>
97          <Box backgroundColor={head} hover={{ backgroundColor: headHi }}>
98            <Button key={`step:${id}`} plain label={key} hover={{ bold: true }} onPress={() => undefined} />
99          </Box>
100          <Text color={head} hover={{ color: headHi }}>{'\uE0B4'}</Text>
101          <Box paddingX={1} hover={{ backgroundColor: INK_HI }}>
102            <Button key={`step:${id}:text`} plain dimColor={dim} label={text} hover={{ bold: true }} onPress={() => undefined} />
103          </Box>
104        </Box>
105      )
106    }
107
108    return (
109      <Box flexDirection="column">
110        <Box flexDirection="row" paddingX={1} columnGap={2}>
111          <Text color={ORANGE} bold>{'✻ 下一步'}</Text>
112          <Box flexDirection="column">
113            {s.items.map((t, i) => row(String(i), String(i + 1), t))}
114            {row('dismiss', '0', '略過', true)}
115          </Box>
116        </Box>
117        {rest}
118      </Box>
119    )
120  })
121}
122
123// Writes the list, then whether it shows, for where-am-i.
124async function setSteps($: EngineInterface, fn: (before: NextSteps | null) => NextSteps | null) {
125  const after = await update($, steps, fn)
126  const showing = after !== null && after.items.length > 0
127  if ((await read($, active)) !== showing) await update($, active, () => showing)
128}
129
130// The whole list on one row, or one suggestion per row when it would not fit.
131// 中文字在終端占 2 格；每顆膠囊另有兩端圓角、數字段與內距約 8 格，標籤與「略過」約 20 格
132const cells = (s: string) => [...s].reduce((n, ch) => n + (/[\u1100-\u115F\u2E80-\uA4CF\uAC00-\uD7A3\uF900-\uFAFF\uFE30-\uFE4F\uFF00-\uFF60\uFFE0-\uFFE6]/.test(ch) ? 2 : 1), 0)
133export function fitsOnOneLine(items: string[], columns: number) {
134  return items.reduce((n, s) => n + cells(s) + 11, 20) <= columns
135}
136
137async function suggest($: EngineInterface, turnId: string, answer: string, current: () => string) {
138  // Wait for background agents: the main agent picks up again once they finish.
139  const agents = await $.agent.list().catch(() => [])
140  if (agents.some(a => a.status === 'running')) return
141
142  const all = (await $.session.messages().catch(() => [])).filter(m => m.text.trim() !== '')
143  if (all.at(-1)?.role === 'assistant') all.pop() // the latest reply goes in on its own below
144  const messages = all.slice(-6)
145  const r = await $.model.complete({
146    model: MODEL,
147    maxTokens: 200,
148    timeoutMs: 20_000,
149    system:
150      'You predict what a person will type next to their coding assistant, after reading a transcript of their session. ' +
151      'Give 2 or 3 short, concrete follow-up prompts, each an instruction in the imperative, written in Traditional Chinese (Taiwan usage, 繁體中文), ' +
152      'at most 24 Chinese characters each; keep file names, commands and code identifiers as they are. ' +
153      'Specific to what just happened (name the file, test or feature). No em dashes, no numbering, no quotes. ' +
154      'One per line and nothing else. If nothing sensible follows, reply with the single word NONE.\n' +
155      'Example reply:\n再跑一次登入測試\n把同樣的檢查加到 signup.ts\n開一個 draft PR',
156    prompt: [
157      `<transcript>\n${messages.map(m => `[${m.role === 'user' ? 'person' : 'assistant'}] ${m.text.slice(0, 800)}`).join('\n')}\n</transcript>`,
158      `<latest_reply>\n${answer.slice(0, 2000)}\n</latest_reply>`,
159      'Write the follow-up prompts the person is likely to send next: the lines only, not a reply to the transcript.',
160    ].join('\n\n'),
161  })
162  if (!r.isAnswered) return
163  const items = parseSteps(r.text)
164  // Checked inside the write, so a submit or a new turn that lands meanwhile always wins.
165  await setSteps($, before => (current() !== turnId ? before : items.length > 0 ? { turnId, items } : null))
166}
167
168// The model's lines, cleaned: no fences, bullets, numbers or em dashes; at most three.
169// A preamble, a "nothing to add" or a line too long to draft whole is dropped.
170export function parseSteps(text: string): string[] {
171  const body = text.replace(/```[a-z]*\n?/gi, '').trim()
172  let lines = body.split('\n')
173  if (body.startsWith('[') || body.startsWith('{')) {
174    try {
175      const o = JSON.parse(body) as unknown
176      const arr = Array.isArray(o) ? o : Object.values(o as Record<string, unknown>).find(Array.isArray)
177      lines = Array.isArray(arr) ? arr.filter((x): x is string => typeof x === 'string') : []
178    } catch {
179      // not JSON after all: read it as lines
180    }
181  }
182  const out: string[] = []
183  for (const line of lines) {
184    const clean = line
185      .replace(/\*\*/g, '')
186      .replace(/^\s*(?:[-*•]|\d+[.):])\s*/, '')
187      .replace(/^["'`]+|["'`,]+$/g, '')
188      .replace(/\s*\u2014\s*/g, ', ')
189      .replace(/\.$/, '')
190      .trim()
191    if (clean === '' || clean.endsWith(':') || clean.length > MAX_KEPT) continue
192    if (/^(none|no |nothing|here are|here's|sure)/i.test(clean) || /^[[\]{}]$/.test(clean)) continue
193    if (!out.some(o => o.toLowerCase() === clean.toLowerCase())) out.push(clean)
194    if (out.length === MAX_ITEMS) break
195  }
196  return out
197}
198

1export type NextSteps = { turnId: string; items: string[] }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'pill-steps': { steps: NextSteps | null; active: boolean }
6  }
7}
8