word-predict

1import { atom, read } from 'claude-code'
2import type { EngineInterface, PromptBox, Register } from 'claude-code'
3
4import { LAYOUTS, literal, readPress, take } from './keys'
5import type { Press } from './keys'
6import { SOURCES, buildLexicon, fetchInRanges, isCurrent } from './lexicon'
7import { addition, buildModel, empty, learn, predict } from './predict'
8import type { Model } from './predict'
9
10import type { History, Layout, Learned, Lexicon, Prediction, Settings } from '../types'
11
12const NONE: Prediction = { prefix: '', words: [] }
13const DEFAULTS: Settings = { isOn: true, layout: 'us' }
14
15const PREDICTION = { plugin: 'word-predict', key: 'prediction' } as const
16const SETTINGS = { plugin: 'word-predict', key: 'settings' } as const
17const prediction = atom(PREDICTION, NONE)
18const shown = atom(SETTINGS, DEFAULTS)
19
20const HELP =
21  '/predict on, off, layout us|danish|german|british (which Option+1/2/3 take a word), import (learn from past sessions), forget (clear what it learned)'
22
23// The store holds the whole history; past this, an import is refused rather than
24// crowding out the word list and what is learned (the store's limit is 4 MiB).
25const MAX_HISTORY = 1_500_000
26
27/** Publishes what the strip offers; a plain write, as nothing else writes it. */
28function show($: EngineInterface, next: Prediction) {
29  void $.state.set(PREDICTION, next)
30}
31
32/** Reads what `/predict import`'s script printed; rejects output too large to keep. */
33function parseHistory(run: { exitCode: number; stdout: string; stderr: string; isStdoutTruncated: boolean }): History {
34  if (run.exitCode !== 0) throw new Error(run.stderr.trim().split('\n').pop() || `python3 exited with code ${run.exitCode}`)
35  if (run.isStdoutTruncated || run.stdout.length > MAX_HISTORY) throw new Error('your history is too large to keep')
36
37  return JSON.parse(run.stdout) as History
38}
39
40const LAYOUT_NAMES: Record<Layout, string> = { us: 'US', danish: 'Danish', german: 'German', british: 'British' }
41
42const isLayout =(value: string): value is Layout => Object.hasOwn(LAYOUTS, value)
43
44/** Reads what any session may have changed: settings, learned words, the imported history. */
45async function readStored($: EngineInterface): Promise<{ settings: Settings; learned: Learned; history?: History }> {
46  const [settings, learned, history] = await Promise.all([$.store.get('settings'), $.store.get('learned'), $.store.get('history')])
47  return {
48    settings: { ...DEFAULTS, ...(settings as Partial<Settings> | undefined) },
49    learned: (learned as Learned | undefined) ?? empty(),
50    history: history as History | undefined,
51  }
52}
53
54export const register: Register = on => {
55  let learned: Learned = empty()
56  let lexicon: Lexicon | undefined
57  let history: History | undefined
58  let settings: Settings = DEFAULTS
59  let model: Model = buildModel()
60  // Whether the strip is drawn: Option keys take only words the person can see.
61  let isStripDrawn = false
62  // The last Option press that took a word, so the same key again types its character.
63  let lastPress: { press: Press; after: PromptBox } | undefined
64
65  const rebuild = () => {
66    model = buildModel(lexicon, history)
67  }
68
69  /** Takes in what the store holds now, rebuilding the model only when the history changed. */
70  const adopt = (stored: Awaited<ReturnType<typeof readStored>>) => {
71    settings = stored.settings
72    learned = stored.learned
73    if (stored.history?.prompts !== history?.prompts || stored.history?.next.length !== history?.next.length) {
74      history = stored.history
75      rebuild()
76    }
77  }
78
79  on('session.start', async ($, e, next) => {
80    adopt(await readStored($))
81    void $.state.set(SETTINGS, settings)
82    lexicon = (await $.store.get('lexicon')) as Lexicon | undefined
83    rebuild()
84    await $.command.register({ name: 'predict', description: `Word prediction: ${HELP}` })
85
86    // First run, or a word list from an older version: fetch it in the background.
87    if (!isCurrent(lexicon)) {
88      void (async () => {
89        $.ui.status('Downloading the word list for word prediction…')
90        try {
91          // The only network use: two ranged GETs of Peter Norvig's count files.
92          const [words, pairs] = await Promise.all([
93            fetchInRanges(SOURCES.words, range =>
94              $.http.fetch('https://www.norvig.com/ngrams/count_1w.txt', { headers: { Range: range } }),
95            ),
96            fetchInRanges(SOURCES.pairs, range =>
97              $.http.fetch('https://www.norvig.com/ngrams/count_2w.txt', { headers: { Range: range } }),
98            ),
99          ])
100          const fresh = buildLexicon(words, pairs)
101          await $.store.set('lexicon', fresh)
102          lexicon = fresh
103          rebuild()
104          $.ui.toast('Word prediction is ready. Run /predict layout to pick your keyboard, /predict import to learn from past prompts.')
105        } catch (error) {
106          $.ui.toast(`Word prediction could not set up its word list: ${String(error)}`)
107        } finally {
108          $.ui.status(undefined)
109        }
110      })()
111    }
112
113    return next(e)
114  })
115
116  on('prompt.edit', async ($, e, next) => {
117    const press = settings.isOn && isStripDrawn ? readPress(settings.layout, e) : undefined
118    if (press) {
119      // The same key right after a pick: the person wanted the character, not a word.
120      const isRepeat = lastPress?.press.char === press.char && e.text.trimEnd() === lastPress.after.text.trimEnd()
121      const box = isRepeat && lastPress ? literal(lastPress.press) : take(model, learned, press)
122      lastPress = box && !isRepeat ? { press, after: box } : undefined
123      if (box) {
124        show($, predict(model, learned, box.text, box.cursor))
125        return box
126      }
127    }
128    lastPress = undefined
129
130    const box = await next(e)
131    if (!settings.isOn) return box
132    // The desktop holds a typed space back until the next character, so the box
133    // does not have it yet: predict as if it had, so next words show at once.
134    const isSpaceHeld = e.key?.key === 'space' && e.inputText === '' && box.text === e.text
135    const text = isSpaceHeld ? box.text.slice(0, box.cursor) + ' ' + box.text.slice(box.cursor) : box.text
136    show($, predict(model, learned, text, box.cursor + (isSpaceHeld ? 1 : 0)))
137
138    return box
139  })
140
141  on('prompt.submit', async ($, e, next) => {
142    const result = await next(e)
143    // The box is empty once a prompt is sent.
144    lastPress = undefined
145    show($, NONE)
146    // Learn what entered, as the hooks beneath left it; nothing from a refused prompt.
147    if (result.drop !== undefined || e.command || e.isMeta) return result
148
149    // Merge into what the store holds now, so other sessions' learning and a
150    // `/predict forget` elsewhere are kept, not overwritten by this session's copy.
151    adopt(await readStored($))
152    if (settings.isOn) {
153      learned = learn(learned, result.text)
154      await $.store.set('learned', learned).catch(() => $.ui.toast('Word prediction could not save what it learned: the store is full.'))
155    }
156
157    return result
158  })
159
160  on('command.run', { command: 'predict' }, async ($, e) => {
161    const [arg = '', value = ''] = e.args.trim().toLowerCase().split(/\s+/)
162    adopt(await readStored($))
163
164    if (arg === 'on' || arg === 'off' || arg === 'layout') {
165      if (arg === 'layout' && !isLayout(value)) return { text: `Pick one of: ${Object.keys(LAYOUTS).join(', ')}.` }
166      settings = arg === 'layout' ? { ...settings, layout: value as Layout } : { ...settings, isOn: arg === 'on' }
167      await $.store.set('settings', settings)
168      void $.state.set(SETTINGS, settings)
169      if (!settings.isOn) show($, NONE)
170      return { text: arg === 'layout' ? `Option+1/2/3 now follow the ${value} layout.` : `Word prediction ${arg}.` }
171    }
172    if (arg === 'import') {
173      let imported: History
174      try {
175        imported = parseHistory(
176          await $.process.run(['python3', '-I', `${$.plugin.root}/scripts/import_history.py`], { timeoutMs: 300_000 }),
177        )
178        await $.store.set('history', imported)
179      } catch (error) {
180        return { text: `Could not import past sessions: ${String(error)}` }
181      }
182      history = imported
183      rebuild()
184      return { text: `Learned from ${imported.prompts} prompts in past sessions.` }
185    }
186    if (arg === 'forget') {
187      await Promise.all([$.store.delete('learned'), $.store.delete('history')])
188      learned = empty()
189      history = undefined
190      rebuild()
191      return { text: 'Forgot everything learned from your prompts.' }
192    }
193    const state = settings.isOn ? 'on' : 'off'
194    const words = model.vocab.length > 0 ? `${model.vocab.length} words` : 'no word list yet'
195    const own = history ? `${history.prompts} past prompts imported` : 'past prompts not imported'
196    const sent = `${Object.keys(learned.words).length} words learned since`
197
198    return { text: `Word prediction is ${state} (${settings.layout} layout): ${words}, ${own}, ${sent}.\n${HELP}` }
199  })
200
201  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
202    const { prefix, words } = await read($, prediction)
203    const { isOn, layout } = (await read($, shown)) as Settings
204    isStripDrawn = isOn && !e.props.hasSurvey && !e.props.isWorking && words.length > 0
205    if (!isStripDrawn) return next(e)
206    const { Box, Button, Text } = $.ui.resolve(e)
207
208    // A press takes the word against the box as it is now.
209    const press = async (word: string) => {
210      const box = await $.prompt.read()
211      const added = addition(box.text.slice(0, box.cursor), word, prefix)
212      const { isFilled } = await $.prompt.fill({ text: added, mode: 'insert' })
213      // A fill is no keystroke, so offer the next words here, for the draft as the person now sees it.
214      const text = box.text.slice(0, box.cursor) + added + box.text.slice(box.cursor)
215      if (isFilled) show($, predict(model, learned, text, box.cursor + added.length))
216    }
217
218    return (
219      <Box flexDirection="row" justifyContent="space-between">
220        <Box flexDirection="row" gap={3}>
221          <Text dimColor>{prefix ? 'complete' : 'next'}</Text>
222          {words.map((word, i) => (
223            <Box key={`s${i}`} flexDirection="row" gap={1}>
224              <Text dimColor>{`⌥${i + 1}`}</Text>
225              <Button key={`w${i}`} label={word} plain onPress={() => press(word)} />
226            </Box>
227          ))}
228        </Box>
229        {/* Which keyboard layout the ⌥ keys follow (/predict layout). */}
230        <Text dimColor>{LAYOUT_NAMES[layout]}</Text>
231      </Box>
232    )
233  })
234}
235

1import type { PromptBox, PromptEditInput } from 'claude-code'
2
3import { accept, predict } from './predict'
4import type { Model } from './predict'
5
6import type { Layout, Learned } from '../types'
7
8// What Option+1, Option+2 and Option+3 type on each Mac layout. The desktop app
9// keeps Tab, the arrows and Ctrl chords for itself, but sends an Option key as
10// the character it types, with `meta`. Only the person's own layout counts, so
11// another layout's Option characters are never taken.
12export const LAYOUTS: Record<Layout, readonly string[]> = {
13  us: ['¡', '™', '£'],
14  danish: ['¡', '“', '§'],
15  german: ['¡', '“', '¶'],
16  british: ['¡', '€', '#'],
17}
18
19/** An Option+1..3 press: which place it takes, its character, and the draft as the person sees it. */
20export type Press = { place: number; char: string; text: string; cursor: number }
21
22/** Reads an edit as an Option+1..3 press on `layout`, or undefined for any other edit. */
23export const readPress = (layout: Layout, e: PromptEditInput): Press | undefined => {
24  const char = e.key?.meta && !e.key.ctrl ? e.key.key : ''
25  const place = char ? LAYOUTS[layout].indexOf(char) : -1
26  // Only the key itself, after the space the desktop may hold back: not an
27  // undo or a paste that happens to end in the same character.
28  const typed = e.inputText.slice(0, -char.length)
29  if (place < 0 || !e.inputText.endsWith(char) || typed.trim() !== '') return undefined
30
31  // The desktop may send the key as replacing the space before it: then keep
32  // the draft as it was, else add the held space it carried.
33  const text = typed ? e.text.slice(0, e.start) + typed + e.text.slice(e.end) : e.text
34  const cursor = typed ? e.start + typed.length : e.cursor
35
36  return { place, char, text, cursor }
37}
38
39/** The box with the suggestion at the press's place taken, or undefined when there is none. */
40export const take = (model: Model, learned: Learned, press: Press): PromptBox | undefined => {
41  const { prefix, words } = predict(model, learned, press.text, press.cursor)
42  const word = words[press.place]
43
44  return word ? accept(press.text, press.cursor, word, prefix) : undefined
45}
46
47/** The press's own character typed into its draft, for when the person wanted the character. */
48export const literal = (press: Press): PromptBox => ({
49  text: press.text.slice(0, press.cursor) + press.char + press.text.slice(press.cursor),
50  cursor: press.cursor + press.char.length,
51})
52

1import type { Lexicon } from '../types'
2
3// Peter Norvig's English word and word-pair counts, derived from the Google Web
4// Trillion Word Corpus. Each person's copy is downloaded from his site on first
5// run and built here, so this repository ships none of the data. The hashes pin
6// the exact files the mod was built and tested against: anything else is refused.
7export const SOURCES = {
8  words: { bytes: 4956241, sha256: '51df159fd3de12b20e403c108f526e96dbd723d9cabdd5f17955cdc16059e690' },
9  pairs: { bytes: 5566017, sha256: '781c0596c3eea532d30bef9f3dba1d5137d652f00376260822c761a7584dfb8c' },
10} as const
11
12/** One ranged GET of a source file, as `register.tsx` writes it with the file's fixed address. */
13export type FetchRange = (range: string) => Promise<{ status: number; text: string }>
14
15// Bumped whenever buildLexicon's output changes, so stored copies are rebuilt.
16export const LEXICON_VERSION = 2
17
18// `$.http.fetch` cuts a body at 4 MiB, so each file is fetched in ranges below that.
19const CHUNK = 2 * 1024 * 1024
20
21const SIZE = 25000
22const NEXT_PER_WORD = 5
23const MAX_WORD = 24
24
25// Kept out of suggestions, as a phone keyboard does; still typeable. A word is
26// blocked when it starts with one of PREFIXES, or is one of ROOTS with an
27// inflection added.
28const PREFIXES = `porn xxx fuck shit cunt pussy pussies milf hentai slut whore dildo blowjob cumshot
29  upskirt shemale tranny faggot bdsm gangbang bukkake masturbat orgasm erotic horny
30  voyeur incest bestiality squirting creampie handjob threesome`.split(/\s+/)
31const ROOTS = `sex sexy nude naked anal cum cock dick tit boob escort viagra cialis penis
32  vagina fetish rape retard fag bondage bitch nigger nigga`.split(/\s+/)
33const INFLECTIONS = ['', 's', 'es', 'ed', 'ing', 'er', 'ers', 'y', 'ies']
34
35export const isBlocked = (word: string): boolean =>
36  PREFIXES.some(prefix => word.startsWith(prefix)) ||
37  ROOTS.some(root => word.startsWith(root) && INFLECTIONS.includes(word.slice(root.length)))
38
39// Common, but absent from the corpus, which drops apostrophes.
40const CONTRACTIONS = `don't can't won't isn't doesn't didn't wasn't aren't shouldn't couldn't wouldn't
41  haven't hasn't i'm i've i'll i'd it's that's there's what's let's you're we're they're`.split(/\s+/)
42
43// What the corpus keeps of a contraction once it splits it at the apostrophe.
44const FRAGMENTS = new Set(['re', 'll', 've', 'don', 'haven', 'isn', 'doesn', 'didn', 'wasn', 'aren', 'shouldn', 'couldn', 'wouldn', 'hasn'])
45
46/** Packs `word:next1,next2|...` from counts, each word's most frequent followers first. */
47const pack = (counts: Map<string, Map<string, number>>, keep: number): string => {
48  const rows: string[] = []
49  for (const [first, nexts] of counts) {
50    const top = [...nexts]
51      .sort((a, b) => b[1] - a[1])
52      .slice(0, keep)
53      .map(([word]) => word)
54    if (top.length > 0) rows.push(`${first}:${top.join(',')}`)
55  }
56  return rows.join('|')
57}
58
59/** Builds the vocabulary and next-word table from the text of the two count files. */
60export const buildLexicon = (wordCounts: string, pairCounts: string): Lexicon => {
61  // A contraction's followers are counted under its apostrophe-less spelling
62  // (`dont`), which leaves the word list for the contraction itself.
63  const spelled = new Map(CONTRACTIONS.map(c => [c.replace(/'/g, ''), c]))
64  const vocab: string[] = []
65  for (const line of wordCounts.split('\n')) {
66    const word = line.split('\t')[0] ?? ''
67    const isWord = /^[a-z]+$/.test(word) && word.length <= MAX_WORD && (word.length > 1 || word === 'a' || word === 'i')
68    if (isWord && !isBlocked(word) && !FRAGMENTS.has(word) && !spelled.has(word)) {
69      vocab.push(word)
70      if (vocab.length === SIZE) break
71    }
72  }
73  const words = [...vocab.slice(0, 300), ...CONTRACTIONS, ...vocab.slice(300)]
74  const known = new Set(words)
75
76  const pairs = new Map<string, Map<string, number>>()
77  for (const line of pairCounts.split('\n')) {
78    const [phrase, count] = line.split('\t')
79    const n = Number(count)
80    const [written, second, extra] = (phrase ?? '').toLowerCase().split(' ')
81    const first = spelled.get(written ?? '') ?? written
82    if (!first || !second || extra !== undefined || !Number.isFinite(n) || !known.has(first) || !known.has(second)) continue
83    const nexts = pairs.get(first) ?? new Map<string, number>()
84    nexts.set(second, (nexts.get(second) ?? 0) + n)
85    pairs.set(first, nexts)
86  }
87
88  return { version: LEXICON_VERSION, words: words.join(' '), next: pack(pairs, NEXT_PER_WORD) }
89}
90
91/** True for a stored lexicon this version built, not a stale or broken copy. */
92export const isCurrent = (lexicon: Lexicon | undefined): lexicon is Lexicon =>
93  lexicon?.version === LEXICON_VERSION && typeof lexicon.words === 'string' && typeof lexicon.next === 'string'
94
95const hex = (bytes: ArrayBuffer): string => [...new Uint8Array(bytes)].map(b => b.toString(16).padStart(2, '0')).join('')
96
97/**
98 * Fetches one source in ranges through `get` and checks its length and hash;
99 * rejects on any mismatch. A range can end inside a multi-byte character, which
100 * arrives broken as text, so each range keeps whole lines and the next starts after them.
101 */
102export async function fetchInRanges(source: { bytes: number; sha256: string }, get: FetchRange): Promise<string> {
103  const encoder = new TextEncoder()
104  const parts: string[] = []
105  let start = 0
106  while (start < source.bytes) {
107    const end = Math.min(start + CHUNK, source.bytes) - 1
108    const response = await get(`bytes=${start}-${end}`)
109    if (response.status !== 206) throw new Error(`the server answered ${response.status}`)
110    const kept = end === source.bytes - 1 ? response.text : response.text.slice(0, response.text.lastIndexOf('\n') + 1)
111    if (kept === '') throw new Error('the server sent no whole line')
112    parts.push(kept)
113    start += encoder.encode(kept).length
114  }
115  const text = parts.join('')
116  const bytes = encoder.encode(text)
117  const digest = hex(await crypto.subtle.digest('SHA-256', bytes))
118  if (bytes.length !== source.bytes || digest !== source.sha256) throw new Error('the downloaded file is not the expected one')
119
120  return text
121}
122

1import type { History, Learned, Lexicon, Prediction } from '../types'
2
3// Stands for "the start of a sentence" as a previous word.
4export const START = '<s>'
5const WORD = /^[a-z][a-z']*$/
6const STARTERS = ['i', 'can', 'please', 'the', 'what', 'how']
7const MAX_WORDS = 5000
8const MAX_PAIRS = 3000
9const MAX_NEXT = 12
10
11// Web words that follow almost anything on a web page and nothing in a prompt.
12const FILLER = new Set(
13  'click page site home online search contact posted comments reviews rss login http www com print reply cart shipping'.split(' '),
14)
15
16/** What predictions are made from: the word list and, when imported, past prompts. */
17export type Model = {
18  vocab: string[]
19  inVocab: Set<string>
20  next: Map<string, string[]>
21  mine: Map<string, string[]>
22  mine2: Map<string, string[]>
23}
24
25const table = (packed: string): Map<string, string[]> => {
26  const map = new Map<string, string[]>()
27  for (const entry of packed ? packed.split('|') : []) {
28    const at = entry.lastIndexOf(':')
29    if (at > 0) map.set(entry.slice(0, at), entry.slice(at + 1).split(','))
30  }
31  return map
32}
33
34export const buildModel = (lexicon?: Lexicon, history?: History): Model => {
35  const vocab = lexicon?.words ? lexicon.words.split(' ') : []
36  return {
37    vocab,
38    inVocab: new Set(vocab),
39    next: table(lexicon?.next ?? ''),
40    mine: table(history?.next ?? ''),
41    mine2: table(history?.next2 ?? ''),
42  }
43}
44
45export const empty = (): Learned => ({ words: {}, pairs: {} })
46
47// Learned counts are plain objects read back from JSON, so a word such as
48// `constructor` must be read as an own property, never through the prototype.
49const count = (counts: Record<string, number> | undefined, word: string): number => {
50  const value = counts && Object.hasOwn(counts, word) ? counts[word] : undefined
51  return typeof value === 'number' ? value : 0
52}
53
54const followersOf = (learned: Learned, word: string): Record<string, number> | undefined =>
55  Object.hasOwn(learned.pairs, word) ? learned.pairs[word] : undefined
56
57const byCount = (counts: Record<string, number> | undefined): string[] =>
58  Object.keys(counts ?? {})
59    .map(word => [word, count(counts, word)] as const)
60    .sort((a, b) => b[1] - a[1])
61    .map(([word]) => word)
62
63/** Writes `word` in the case the person used for `like`: `Pre` -> `Prediction`. */
64const cased = (word: string, like: string): string => {
65  if (word === 'i' || word.startsWith("i'")) return capital(word)
66  if (like.length > 1 && like === like.toUpperCase()) return word.toUpperCase()
67  if (like[0] && like[0] !== like[0].toLowerCase()) return capital(word)
68  return word
69}
70
71const capital = (word: string): string => word.charAt(0).toUpperCase() + word.slice(1)
72
73const unique = (words: string[], limit: number): string[] => [...new Set(words)].slice(0, limit)
74
75/** The bare lowercase word of a whitespace token, or undefined for code, paths and the like. */
76const bare = (token: string): string | undefined => {
77  const word = token.toLowerCase().replace(/^["'(\[]+|[.,!?;:"')\]]+$/g, '')
78  return WORD.test(word) && word.length <= 24 ? word : undefined
79}
80
81/** True where a new sentence starts after `before`: an empty draft, a full stop, a new line. */
82const isSentenceStart = (before: string): boolean => {
83  const trimmed = before.trimEnd()
84  return trimmed === '' || /[.!?]["')]*$/.test(trimmed) || /\n\s*$/.test(before)
85}
86
87/** True for a word worth offering: in the word list or past prompts, or sent at least twice. */
88const isKnown = (model: Model, learned: Learned, word: string): boolean =>
89  model.inVocab.has(word) || model.mine.has(word) || count(learned.words, word) >= 2
90
91/** Words likely after `previous` (and `earlier` before it), best first: the person's own, then the web's. */
92const candidates = (model: Model, learned: Learned, previous: string, earlier?: string): string[] =>
93  [
94    ...(earlier ? (model.mine2.get(`${earlier} ${previous}`) ?? []) : []),
95    ...byCount(followersOf(learned, previous)).filter(word => isKnown(model, learned, word)),
96    ...(model.mine.get(previous) ?? []),
97    ...(model.next.get(previous) ?? []).filter(word => !FILLER.has(word)),
98    ...(previous === START ? STARTERS : []),
99  ].filter(word => word !== previous)
100
101/** Ways to finish `prefix`, given the words before it (`START` at a sentence start). */
102export const complete = (
103  model: Model,
104  learned: Learned,
105  prefix: string,
106  previous: string | undefined,
107  limit: number,
108  earlier?: string,
109): string[] => {
110  const lower = prefix.toLowerCase()
111  const fits = (word: string) => word.startsWith(lower) && word !== lower
112
113  const found: string[] = []
114  if (previous) found.push(...candidates(model, learned, previous, earlier).filter(fits))
115  found.push(...byCount(learned.words).filter(word => fits(word) && count(learned.words, word) >= 2))
116  for (const word of model.vocab) {
117    if (found.length >= limit * 4) break
118    if (fits(word)) found.push(word)
119  }
120
121  return unique(found, limit).map(word => cased(word, prefix))
122}
123
124/** Likely next words after `previous`; sentence starters, capitalised, after `START`. */
125export const following = (model: Model, learned: Learned, previous: string, limit: number, earlier?: string): string[] => {
126  const found = unique(candidates(model, learned, previous, earlier), limit)
127  return previous === START ? found.map(capital) : found.map(word => cased(word, word))
128}
129
130/** What the strip offers for a draft with the caret at `cursor`. */
131export const predict = (model: Model, learned: Learned, text: string, cursor: number, limit = 3): Prediction => {
132  const before = text.slice(0, cursor)
133  const tokens = before.split(/\s+/)
134  const current = tokens[tokens.length - 1] ?? ''
135  const head = before.slice(0, before.length - current.length)
136  const previousToken = tokens[tokens.length - 2] ?? ''
137  // START at a sentence start; undefined after a token that isn't a word (a path, a number).
138  const previous = isSentenceStart(head) ? START : bare(previousToken)
139  const lead = head.slice(0, head.trimEnd().length - previousToken.length)
140  const earlier = previous && previous !== START ? (isSentenceStart(lead) ? START : bare(tokens[tokens.length - 3] ?? '')) : undefined
141
142  if (current === '') {
143    if (text.trim() === '' || !previous) return { prefix: '', words: [] }
144    return { prefix: '', words: following(model, learned, previous, limit, earlier) }
145  }
146  if (!/^[A-Za-z][A-Za-z']*$/.test(current)) return { prefix: current, words: [] }
147
148  return { prefix: current, words: complete(model, learned, current, previous, limit, earlier) }
149}
150
151/**
152 * What taking `word` adds after `before`: for a completion (`prefix` set), the
153 * rest of the word being typed; for a next word, the word with a space in
154 * front where `before` doesn't end in one. Then a space.
155 */
156export const addition = (before: string, word: string, prefix: string): string => {
157  if (prefix) {
158    const typed = before.match(/[A-Za-z']+$/)?.[0] ?? ''
159    const rest = typed && word.toLowerCase().startsWith(typed.toLowerCase()) ? word.slice(typed.length) : word
160    return rest + ' '
161  }
162  return (before === '' || /\s$/.test(before) ? '' : ' ') + word + ' '
163}
164
165/** The draft with `word` taken at the cursor. */
166export const accept = (text: string, cursor: number, word: string, prefix: string): { text: string; cursor: number } => {
167  const added = addition(text.slice(0, cursor), word, prefix)
168  return { text: text.slice(0, cursor) + added + text.slice(cursor), cursor: cursor + added.length }
169}
170
171/** Adds `n` to `word` in `counts`, moving it to the end so insertion order is recency. */
172const bump = (counts: Record<string, number>, word: string, n: number): void => {
173  const total = count(counts, word) + n
174  delete counts[word]
175  counts[word] = total
176}
177
178/** Counts the words and word pairs of a sent prompt into `learned`, leaving out code and paths. */
179export const learn = (learned: Learned, text: string): Learned => {
180  const words: Record<string, number> = Object.assign(Object.create(null), learned.words)
181  const pairs: Record<string, Record<string, number>> = Object.assign(Object.create(null), learned.pairs)
182  // Closed code fences, then an unclosed one to the end, then inline code.
183  const prose = text
184    .replace(/(```|~~~)[\s\S]*?\1/g, '\n')
185    .replace(/(```|~~~)[\s\S]*$/, '\n')
186    .replace(/`[^`\n]*`/g, ' ')
187
188  for (const sentence of prose.split(/(?<=[.!?]["')]*)\s+|\n+/)) {
189    let previous = START
190    for (const token of sentence.trim().split(/\s+/)) {
191      const word = bare(token)
192      if (!word) {
193        previous = ''
194        continue
195      }
196      bump(words, word, 1)
197      if (previous) {
198        const nexts: Record<string, number> = Object.assign(Object.create(null), followersOf({ words, pairs }, previous))
199        bump(nexts, word, 1)
200        delete pairs[previous]
201        pairs[previous] = nexts
202      }
203      previous = /[.!?,;:]["')]*$/.test(token) ? '' : word
204    }
205  }
206
207  return { words: keepTop(words, MAX_WORDS), pairs: prunePairs(pairs) }
208}
209
210/** Keeps the `limit` highest counts; on a tie the more recent entry stays, so new words get a chance. */
211const keepTop = (counts: Record<string, number>, limit: number): Record<string, number> => {
212  const entries = Object.keys(counts).map((word, at) => ({ word, n: count(counts, word), at }))
213  if (entries.length <= limit) return counts
214  const kept = entries.sort((a, b) => b.n - a.n || b.at - a.at).slice(0, limit).sort((a, b) => a.at - b.at)
215  const result: Record<string, number> = Object.create(null)
216  for (const { word, n } of kept) result[word] = n
217  return result
218}
219
220const prunePairs = (pairs: Record<string, Record<string, number>>): Record<string, Record<string, number>> => {
221  const entries = Object.keys(pairs).map((word, at) => {
222    const nexts = keepTop(pairs[word] ?? {}, MAX_NEXT)
223    const total = Object.keys(nexts).reduce((sum, next) => sum + count(nexts, next), 0)
224    return { word, nexts, total, at }
225  })
226  const kept = entries.length > MAX_PAIRS ? entries.sort((a, b) => b.total - a.total || b.at - a.at).slice(0, MAX_PAIRS) : entries
227  const result: Record<string, Record<string, number>> = Object.create(null)
228  for (const { word, nexts } of kept.sort((a, b) => a.at - b.at)) result[word] = nexts
229  return result
230}
231

1/** Counts kept across sessions from the prompts sent; `<s>` keys sentence starters. */
2export type Learned = {
3  words: Record<string, number>
4  pairs: Record<string, Record<string, number>>
5}
6
7/** The downloaded word list, packed: words by frequency, and `word:next1,next2|...`. */
8export type Lexicon = { version: number; words: string; next: string }
9
10/** What `/predict import` counted in past sessions' prompts, packed like `Lexicon.next`. */
11export type History = { prompts: number; next: string; next2: string }
12
13/** The Mac keyboard layout whose Option+1/2/3 characters take a suggestion. */
14export type Layout = 'us' | 'danish' | 'german' | 'british'
15
16/** What `/predict` sets, kept across sessions. */
17export type Settings = { isOn: boolean; layout: Layout }
18
19/** What the strip offers: completions of `prefix`, or next words where `prefix` is empty. */
20export type Prediction = { prefix: string; words: string[] }
21
22declare module 'claude-code' {
23  interface PluginState {
24    'word-predict': { prediction: Prediction; settings: Settings }
25  }
26}
27