next-steps

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionUsage } from 'claude-code'
3
4import type { Baseline, Limit, Usage } from '../types'
5
6const suggestions = atom({ plugin: 'next-steps', key: 'suggestions' } as const, [])
7const isLoading = atom({ plugin: 'next-steps', key: 'isLoading' } as const, false)
8const usage = atom({ plugin: 'next-steps', key: 'usage' } as const, null)
9const baseline = atom({ plugin: 'next-steps', key: 'baseline' } as const, null)
10const tick = atom({ plugin: 'next-steps', key: 'tick' } as const, 0)
11
12// Asked over the session's own transcript ($.model.fork), so the prefix comes
13// from the prompt cache: each turn costs a cache read plus a few output tokens.
14const ASK = [
15  'Without using any tools and without doing any of the work, suggest the 3 most useful prompts I could send you next in this conversation.',
16  'Write each one as I would type it: imperative, specific to this work, at most 70 characters, in the language I have been writing in.',
17  'Answer with exactly 3 lines, one prompt per line, and nothing else: no numbering, no quotes, no preamble.',
18].join(' ')
19
20const BAR_CELLS = 10
21
22// Bumped by every prompt: a fork that answers after a newer prompt went in
23// is stale and dropped.
24let generation = 0
25
26export const register: Register = on => {
27  const measure = (u: Pick<SessionUsage, 'context' | 'rateLimits' | 'cost'>): Usage => ({
28    contextTokens: u.context.tokens,
29    contextWindow: u.context.window,
30    contextPercent: u.context.percent,
31    limits: u.rateLimits.map(({ kind, percentUsed, resetsAt }) => ({ kind, percentUsed, resetsAt })),
32    usd: u.cost?.usd,
33  })
34
35  on('session.start', async ($, e, next) => {
36    const started = await next(e)
37    const now = await $.session.usage()
38    await update($, usage, () => measure(now))
39    // Redraws the "resets in" countdowns once a minute.
40    $.clock.every(60_000, () => void update($, tick, n => n + 1))
41
42    return started
43  })
44
45  on('session.measure', async ($, e, next) => {
46    await update($, usage, () => measure(e))
47
48    return next(e)
49  })
50
51  on('prompt.submit', async ($, e, next) => {
52    generation += 1
53    await update($, suggestions, () => [])
54    await update($, isLoading, () => false)
55    // A prompt typed over a running turn folds into it: the baseline stays
56    // the one its turn started from.
57    if (e.turnId === undefined) {
58      const now = await read($, usage)
59      const base: Baseline = {
60        contextTokens: now?.contextTokens,
61        fiveHour: limit(now, 'five_hour')?.percentUsed,
62        usd: now?.usd,
63      }
64      await update($, baseline, () => base)
65    }
66
67    return next(e)
68  })
69
70  on('turn.complete', async ($, e, next) => {
71    const done = await next(e)
72    const isMainAnswer = e.agentId === undefined && e.reason === 'answer' && e.answer !== ''
73    if (isMainAnswer) {
74      void suggest($, generation)
75    }
76
77    return done
78  })
79
80  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
81    if (e.props.hasSurvey) {
82      return next(e)
83    }
84    const list = await read($, suggestions)
85    const isThinking = await read($, isLoading)
86    const now = await read($, usage)
87    const base = await read($, baseline)
88    await read($, tick)
89    const showsList = !e.props.isWorking && (list.length > 0 || isThinking)
90    if (!showsList && now === null) {
91      return next(e)
92    }
93    const at = await $.clock.now()
94    const { Box, Button, Text } = $.ui.resolve(e)
95
96    const pick = async (text: string) => {
97      const box = await $.prompt.read()
98      const isEmpty = box.text.trim() === ''
99      await $.prompt.fill(isEmpty ? { text } : { text: `\n${text}`, mode: 'append' })
100      await update($, suggestions, () => [])
101    }
102    const dismiss = () => update($, suggestions, () => [])
103
104    const sep = () => <Text dimColor> · </Text>
105    const bar = (one: Limit | undefined, label: string) => {
106      if (one === undefined) {
107        return null
108      }
109      const filled = Math.min(BAR_CELLS, Math.round((one.percentUsed / 100) * BAR_CELLS))
110      const color = levelColor(one.percentUsed)
111      const left = one.resetsAt === undefined ? '' : untilText(Date.parse(one.resetsAt) - at)
112
113      return (
114        <Text>
115          {sep()}
116          <Text dimColor>{label} </Text>
117          <Text color={color}>{'█'.repeat(filled)}</Text>
118          <Text dimColor>{'░'.repeat(BAR_CELLS - filled)}</Text>
119          <Text color={color}>
120            {' '}
121            {formatPercent(one.percentUsed)}
122            {left === '' ? '' : ` (resets ${left})`}
123          </Text>
124        </Text>
125      )
126    }
127
128    const usageLine = () => {
129      if (now === null) {
130        return null
131      }
132      const ctxPercent = now.contextPercent ?? 0
133      const added =
134        now.contextTokens !== undefined && base?.contextTokens !== undefined
135          ? now.contextTokens - base.contextTokens
136          : undefined
137      const fiveHour = limit(now, 'five_hour')?.percentUsed
138      const spent = now.usd !== undefined && base?.usd !== undefined ? now.usd - base.usd : undefined
139      const fiveHourAdded =
140        fiveHour !== undefined && base?.fiveHour !== undefined ? fiveHour - base.fiveHour : undefined
141      const last = [
142        added !== undefined && added >= 0 ? `+${kilo(added)}` : undefined,
143        fiveHourAdded !== undefined ? `+${formatPercent(Math.max(0, fiveHourAdded))}` : undefined,
144        spent !== undefined ? `+$${spent.toFixed(2)}` : undefined,
145      ].filter(part => part !== undefined)
146      // While the turn runs the figures have not moved yet.
147      const hasLast = !e.props.isWorking && base !== null && last.length > 0
148
149      return (
150        <Text wrap="wrap">
151          {now.usd !== undefined && <Text color="success">${now.usd.toFixed(2)}</Text>}
152          {now.usd !== undefined && sep()}
153          <Text dimColor>ctx </Text>
154          <Text>
155            {now.contextTokens === undefined ? '–' : kilo(now.contextTokens)}/{kilo(now.contextWindow)}
156          </Text>
157          <Text color={levelColor(ctxPercent)}> ({formatPercent(ctxPercent)})</Text>
158          {bar(limit(now, 'seven_day'), '7d')}
159          {bar(limit(now, 'five_hour'), '5h')}
160          {hasLast && sep()}
161          {hasLast && <Text color="success">{last.join(' ')}</Text>}
162          {hasLast && <Text dimColor> (last prompt)</Text>}
163        </Text>
164      )
165    }
166
167    return (
168      <Box flexDirection="column">
169        {showsList && <Text dimColor>next:</Text>}
170        {showsList && isThinking && list.length === 0 && (
171          <Box paddingLeft={2}>
172            <Text dimColor>thinking of next steps…</Text>
173          </Box>
174        )}
175        {showsList &&
176          list.map((text, i) => (
177            <Box key={`row${i + 1}`} paddingLeft={2}>
178              <Button key={`pick${i + 1}`} plain hotkey={String(i + 1)} label={text} onPress={() => void pick(text)} />
179            </Box>
180          ))}
181        {showsList && list.length > 0 && (
182          <Box paddingLeft={2}>
183            <Button key="dismiss" plain dimColor hotkey="0" label="dismiss" role="dismiss" onPress={() => void dismiss()} />
184          </Box>
185        )}
186        {usageLine()}
187      </Box>
188    )
189  })
190}
191
192async function suggest($: EngineInterface, mine: number) {
193  await update($, isLoading, () => true)
194  try {
195    const reply = await $.model.fork({ prompt: ASK })
196    if (generation === mine) {
197      await update($, suggestions, () => (reply.isAnswered ? parse(reply.text) : []))
198    }
199  } finally {
200    if (generation === mine) {
201      await update($, isLoading, () => false)
202    }
203  }
204}
205
206const limit = (u: Usage | null, kind: string): Limit | undefined => u?.limits.find(one => one.kind === kind)
207
208// One prompt per line, numbering, bullets and quotes stripped, at most three.
209const parse = (text: string): string[] =>
210  text
211    .split('\n')
212    .map(line =>
213      line
214        .trim()
215        .replace(/^(?:[-*•]|\d+[.):]|\[\d+\])\s*/, '')
216        .replace(/^["'`“”]+|["'`“”]+$/g, '')
217        .trim(),
218    )
219    .filter(line => line !== '')
220    .slice(0, 3)
221
222const kilo = (tokens: number) => `${(tokens / 1000).toFixed(1)}k`
223
224const formatPercent = (n: number) => `${Math.round(n)}%`
225
226const levelColor = (percent: number) => (percent >= 90 ? 'error' : percent >= 60 ? 'warning' : 'success')
227
228const untilText = (ms: number) => {
229  if (!Number.isFinite(ms) || ms <= 0) {
230    return 'now'
231  }
232  const minutes = Math.round(ms / 60_000)
233  const days = Math.floor(minutes / 1440)
234  const hours = Math.floor((minutes % 1440) / 60)
235  const rest = minutes % 60
236  if (days > 0) {
237    return `${days}d${hours}h`
238  }
239
240  return hours > 0 ? `${hours}h${rest}m` : `${rest}m`
241}
242

1// The figures of `$.session.usage()` the band draws (the status line's).
2export type Usage = {
3  contextTokens?: number
4  contextWindow: number
5  contextPercent?: number
6  limits: Limit[]
7  usd?: number
8}
9
10// One rate-limit window: `five_hour`, `seven_day`, or a gateway's.
11export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
12
13// The usage when the person's last prompt went in: the "last prompt"
14// figures are the current usage minus this.
15export type Baseline = { contextTokens?: number; fiveHour?: number; usd?: number }
16
17declare module 'claude-code' {
18  interface PluginState {
19    'next-steps': {
20      suggestions: string[]
21      isLoading: boolean
22      usage: Usage | null
23      baseline: Baseline | null
24      tick: number
25    }
26  }
27}
28