cc-spell-check

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PromptDecoration, Register } from 'claude-code'
3
4import type { SpellActive, SpellDictStatus } from '../types'
5import {
6  HUNSPELL_DIC,
7  parseAspellDump,
8  parseHunspellDic,
9  type DictData,
10  type Dictionary,
11} from './dictionary'
12import { suggestionsFor } from './suggest'
13import { flagged, tokenize, type Token } from './tokenizer'
14
15const active = atom({ plugin: 'cc-spell-check', key: 'active' } as const, null)
16const isEnabled = atom({ plugin: 'cc-spell-check', key: 'isEnabled' } as const, true)
17const dictStatus = atom({ plugin: 'cc-spell-check', key: 'dictStatus' } as const, 'loading')
18
19// Module scope: reset on hot reload; session.start re-fires then and rebuilds.
20let dict: Dictionary | null = null
21let contractions = new Map<string, string>()
22let personal = new Set<string>()
23let enabledNow = true
24const suggestCache = new Map<string, string[]>()
25let lastActiveKey = ''
26
27// Above this, flagging is skipped entirely so huge pastes never lag a keystroke.
28const MAX_CHARS = 10_000
29
30// A flagged token: a plain misspelling, or a contraction missing its
31// apostrophe ("dont" -> "don't"), which gets its own color and exact fix.
32type Flag = Token & { kind: 'spelling' | 'contraction'; fix?: string }
33
34const CONTRACTION_COLOR = '#ffa500'
35
36function recase(word: string, like: string): string {
37  return /^[A-Z]/.test(like) ? word.charAt(0).toUpperCase() + word.slice(1) : word
38}
39
40function classify(tokens: Token[]): Flag[] {
41  return tokens.map(t => {
42    const fix = contractions.get(t.word.toLowerCase())
43    return fix
44      ? { ...t, kind: 'contraction' as const, fix: recase(fix, t.word) }
45      : { ...t, kind: 'spelling' as const }
46  })
47}
48
49function pickNearest(bad: Flag[], cursor: number): Flag | null {
50  let nearest: Flag | null = null
51  for (const t of bad) {
52    if (t.start <= cursor && cursor <= t.end) return t
53    if (t.end <= cursor && (!nearest || t.end > nearest.end)) nearest = t
54  }
55  return nearest
56}
57
58function toDecorations(bad: Flag[]): PromptDecoration[] {
59  return bad.map(t => ({
60    start: t.start,
61    end: t.end,
62    color: t.kind === 'contraction' ? CONTRACTION_COLOR : 'red',
63    underline: true,
64  }))
65}
66
67function decorationsFor(text: string, cursor: number): PromptDecoration[] {
68  if (!dict) return []
69  const bad = classify(flagged(tokenize(text), dict, personal))
70  return toDecorations(bad.filter(t => cursor < t.start || cursor > t.end))
71}
72
73async function loadDictionary(
74  $: EngineInterface,
75  lang: string,
76): Promise<{ data: DictData | null; status: SpellDictStatus }> {
77  try {
78    const run = await $.process.run(['aspell', '-l', lang, 'dump', 'master'], {
79      timeoutMs: 15_000,
80    })
81    if (run.exitCode === 0) {
82      const data = parseAspellDump(run.stdout, run.isStdoutTruncated)
83      if (data.dict.size > 0) return { data, status: 'ready' }
84    }
85  } catch {
86    // aspell missing or failed to start: fall through
87  }
88  try {
89    const data = parseHunspellDic(await $.fs.read(HUNSPELL_DIC))
90    if (data.dict.size > 0) return { data, status: 'fallback' }
91  } catch {
92    // no hunspell wordlist either
93  }
94  return { data: null, status: 'missing' }
95}
96
97async function clearActive($: EngineInterface) {
98  lastActiveKey = ''
99  await update($, active, () => null)
100}
101
102async function applyFix($: EngineInterface, cur: NonNullable<SpellActive>, replacement: string) {
103  if (!dict) return
104  const box = await $.prompt.read()
105  let { start, end } = cur
106  if (box.text.slice(start, end) !== cur.word) {
107    // The draft moved since the word was flagged: find it again.
108    const again = flagged(tokenize(box.text), dict, personal).find(t => t.word === cur.word)
109    if (!again) return void clearActive($)
110    start = again.start
111    end = again.end
112  }
113  const text = box.text.slice(0, start) + replacement + box.text.slice(end)
114  await $.prompt.fill({ text, mode: 'replace', decorations: decorationsFor(text, text.length) })
115  await clearActive($)
116}
117
118async function addWord($: EngineInterface, word: string) {
119  const w = word.toLowerCase()
120  personal.add(w)
121  await $.store.set('personalWords', [...personal].sort())
122  await clearActive($)
123  $.ui.toast(`"${w}" added to dictionary`)
124}
125
126export const register: Register = (on, options) => {
127  const lang = () => String(options.lang ?? 'en_AU')
128
129  on('session.start', async ($, e, next) => {
130    await $.command.register({
131      name: 'spell',
132      description: 'Spell check: toggle, status, personal dictionary',
133      argumentHint: '[on|off|status|add <word>|remove <word>]',
134    })
135    enabledNow = options.enabled !== false
136    void update($, isEnabled, () => enabledNow)
137    // Dictionary load runs off the critical path; no flagging until it lands.
138    void (async () => {
139      personal = new Set(((await $.store.get('personalWords')) as string[] | undefined) ?? [])
140      const loaded = await loadDictionary($, lang())
141      dict = loaded.data?.dict ?? null
142      contractions = loaded.data?.contractions ?? new Map()
143      await update($, dictStatus, () => loaded.status)
144      if (loaded.status === 'missing') {
145        $.ui.toast('cc-spell-check: no dictionary found (sudo apt install aspell aspell-en)')
146      } else if (loaded.status === 'fallback') {
147        $.ui.toast('cc-spell-check: aspell missing, using hunspell en_US wordlist')
148      }
149    })()
150    return next(e)
151  })
152
153  on('prompt.edit', async ($, e, next) => {
154    const r = await next(e)
155    if (!enabledNow || !dict) return r
156    if (r.text.length > MAX_CHARS) {
157      if (lastActiveKey !== '') void clearActive($)
158      return r
159    }
160    const bad = classify(flagged(tokenize(r.text), dict, personal))
161    // A word still under the cursor is mid-typing: flag it only once the
162    // cursor has moved past it (space, punctuation, click elsewhere).
163    const settled = bad.filter(t => r.cursor < t.start || r.cursor > t.end)
164    const near = pickNearest(settled, r.cursor)
165    const key = near ? `${near.word}:${near.start}` : ''
166    if (key !== lastActiveKey) {
167      lastActiveKey = key
168      void update($, active, () => near)
169    }
170    if (settled.length === 0) return r
171    return { ...r, decorations: toDecorations(settled) }
172  })
173
174  on('prompt.submit', ($, e, next) => {
175    // Visual only: the submitted prompt is never rewritten.
176    void clearActive($)
177    return next(e)
178  })
179
180  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
181    if (e.props.hasSurvey) return next(e)
182    const cur = await read($, active)
183    if (!cur || !dict || !(await read($, isEnabled))) return next(e)
184
185    const sugg =
186      cur.kind === 'contraction' && cur.fix
187        ? [cur.fix]
188        : suggestionsFor(cur.word, dict, suggestCache)
189    const { Box, Button, Text } = $.ui.resolve(e)
190    return (
191      <Box>
192        <Text color={cur.kind === 'contraction' ? CONTRACTION_COLOR : 'red'}>✗ {cur.word} </Text>
193        {sugg.map((s, i) => (
194          <Button
195            key={`fix-${i}`}
196            hotkey={String(i + 1)}
197            label={s}
198            variant={i === 0 ? 'primary' : 'secondary'}
199            onPress={() => void applyFix($, cur, s)}
200          />
201        ))}
202        {sugg.length === 0 && <Text dimColor>no suggestions </Text>}
203        <Button key="add" hotkey="a" label="+dict" onPress={() => void addWord($, cur.word)} />
204        <Button key="dismiss" role="dismiss" label="x" onPress={() => void clearActive($)} />
205      </Box>
206    )
207  })
208
209  on('command.run', { command: 'spell' }, async ($, e) => {
210    const [verb, ...rest] = e.args.trim().split(/\s+/).filter(Boolean)
211    switch (verb) {
212      case 'on':
213      case 'off': {
214        enabledNow = verb === 'on'
215        await update($, isEnabled, () => enabledNow)
216        if (!enabledNow) await clearActive($)
217        return { text: `Spell check ${verb} (this session; use /config for the permanent default).` }
218      }
219      case 'add': {
220        if (rest.length === 0) return { text: 'Usage: /spell add <word> [...]' }
221        for (const w of rest) personal.add(w.toLowerCase())
222        await $.store.set('personalWords', [...personal].sort())
223        return { text: `Added to personal dictionary: ${rest.join(', ')}` }
224      }
225      case 'remove': {
226        if (rest.length === 0) return { text: 'Usage: /spell remove <word> [...]' }
227        for (const w of rest) personal.delete(w.toLowerCase())
228        await $.store.set('personalWords', [...personal].sort())
229        return { text: `Removed from personal dictionary: ${rest.join(', ')}` }
230      }
231      default: {
232        const status = await read($, dictStatus)
233        return {
234          text:
235            `Spell check: ${enabledNow ? 'on' : 'off'} | dict: ${status}` +
236            ` (${dict?.size ?? 0} words, lang ${lang()}) | personal: ${personal.size}`,
237        }
238      }
239    }
240  })
241}
242

1export type Dictionary = Set<string>
2
3/** Accept set plus a map of de-apostrophized contractions: "dont" -> "don't". */
4export type DictData = { dict: Dictionary; contractions: Map<string, string> }
5
6export const HUNSPELL_DIC = '/usr/share/hunspell/en_US.dic'
7
8function parseWords(lines: string[], stripFlags: boolean): DictData {
9  const dict: Dictionary = new Set()
10  const contractions = new Map<string, string>()
11  for (const line of lines) {
12    const word = stripFlags ? line.split('/')[0] : line
13    if (!word) continue
14    const w = word.toLowerCase()
15    dict.add(w)
16    if (w.includes("'")) {
17      const bare = w.replace(/'/g, '')
18      if (bare.length > 2 && !contractions.has(bare)) contractions.set(bare, w)
19    }
20  }
21  return { dict, contractions }
22}
23
24/** Wordlist from `aspell dump master` output, one word per line. */
25export function parseAspellDump(stdout: string, isTruncated: boolean): DictData {
26  const lines = stdout.split('\n')
27  if (isTruncated) lines.pop() // last line may be cut
28  return parseWords(lines, false)
29}
30
31/** Wordlist from a hunspell .dic file: count line first, `word/FLAGS` after. */
32export function parseHunspellDic(text: string): DictData {
33  const lines = text.split('\n')
34  lines.shift() // first line is the entry count
35  return parseWords(lines, true)
36}
37

1const LETTERS = "abcdefghijklmnopqrstuvwxyz'"
2const MAX_SUGGESTIONS = 3
3const MAX_EDITS2_HITS = 10
4const MAX_EDITS2_LENGTH = 8
5export const SUGGEST_CACHE_CAP = 500
6
7function edits1(word: string): string[] {
8  const out: string[] = []
9  for (let i = 0; i <= word.length; i++) {
10    const left = word.slice(0, i)
11    const right = word.slice(i)
12    if (right) out.push(left + right.slice(1)) // delete
13    if (right.length > 1) out.push(left + right[1] + right[0] + right.slice(2)) // transpose
14    for (const c of LETTERS) {
15      if (right) out.push(left + c + right.slice(1)) // replace
16      out.push(left + c + right) // insert
17    }
18  }
19  return out
20}
21
22function rank(word: string, candidates: Map<string, number>): string[] {
23  return [...candidates.entries()]
24    .sort(([a, da], [b, db]) => {
25      if (da !== db) return da - db
26      const fa = a[0] === word[0] ? 0 : 1
27      const fb = b[0] === word[0] ? 0 : 1
28      if (fa !== fb) return fa - fb
29      const la = Math.abs(a.length - word.length)
30      const lb = Math.abs(b.length - word.length)
31      if (la !== lb) return la - lb
32      return a < b ? -1 : 1
33    })
34    .map(([w]) => w)
35    .slice(0, MAX_SUGGESTIONS)
36}
37
38/**
39 * Up to 3 corrections for a flagged word: dictionary words at edit distance 1,
40 * falling back to distance 2 for short words. Cached per lowercased word.
41 */
42export function suggestionsFor(
43  word: string,
44  dict: Set<string>,
45  cache: Map<string, string[]>,
46): string[] {
47  const lower = word.toLowerCase()
48  const cached = cache.get(lower)
49  if (cached) return recase(cached, word)
50
51  const candidates = new Map<string, number>()
52  const e1 = edits1(lower)
53  for (const c of e1) {
54    if (dict.has(c)) candidates.set(c, 1)
55  }
56  if (candidates.size === 0 && lower.length <= MAX_EDITS2_LENGTH) {
57    outer: for (const c1 of e1) {
58      for (const c2 of edits1(c1)) {
59        if (dict.has(c2) && !candidates.has(c2)) {
60          candidates.set(c2, 2)
61          if (candidates.size >= MAX_EDITS2_HITS) break outer
62        }
63      }
64    }
65  }
66
67  const ranked = rank(lower, candidates)
68  if (cache.size >= SUGGEST_CACHE_CAP) cache.clear()
69  cache.set(lower, ranked)
70  return recase(ranked, word)
71}
72
73function recase(suggestions: string[], word: string): string[] {
74  if (!/^[A-Z]/.test(word)) return suggestions
75  return suggestions.map(s => s.charAt(0).toUpperCase() + s.slice(1))
76}
77

1export type Token = {
2  word: string
3  start: number
4  end: number
5  isCapitalized: boolean
6}
7
8export type Accepts = { has: (word: string) => boolean }
9
10// Fenced blocks (``` ... ``` — an unclosed trailing fence masks to end of text)
11// and inline `code` spans. Alternation order makes fences win over inline ticks.
12const CODE_RE = /```[\s\S]*?(?:```|$)|`[^`\n]*`/g
13
14// Chunks containing any of these are never spell-checked: URLs, paths,
15// snake_case, digits (versions, hex, ids).
16const CHUNK_SKIP_RE = /[/\\_\d]/
17const CHUNK_PREFIX_SKIP_RE = /^[@#/]/
18
19const CORE_RE = /^[A-Za-z]+(?:'[A-Za-z]+)*$/
20const ALLCAPS_RE = /^[A-Z]+$/
21const INTERNAL_CAPS_RE = /[A-Z]/
22const HEX_RE = /^[a-f]{6,}$/
23
24function codeRanges(text: string): Array<[number, number]> {
25  const ranges: Array<[number, number]> = []
26  for (const m of text.matchAll(CODE_RE)) {
27    ranges.push([m.index, m.index + m[0].length])
28  }
29  return ranges
30}
31
32function inRanges(ranges: Array<[number, number]>, start: number, end: number): boolean {
33  return ranges.some(([a, b]) => start < b && end > a)
34}
35
36/**
37 * Candidate words for spell checking, with UTF-16 offsets into `text`
38 * (the units PromptDecoration and PromptBox.cursor use). Emits only pure
39 * alphabetic lowercase or Capitalized words of 3+ characters outside code.
40 */
41export function tokenize(text: string): Token[] {
42  const masked = codeRanges(text)
43  const tokens: Token[] = []
44
45  for (const m of text.matchAll(/\S+/g)) {
46    const chunk = m[0]
47    const chunkStart = m.index
48    if (inRanges(masked, chunkStart, chunkStart + chunk.length)) continue
49    if (CHUNK_PREFIX_SKIP_RE.test(chunk) || CHUNK_SKIP_RE.test(chunk)) continue
50
51    // Strip edge punctuation (quotes, commas, brackets, stray apostrophes).
52    const lead = /^[^A-Za-z]*/.exec(chunk)![0].length
53    const core = chunk.slice(lead).replace(/[^A-Za-z]+$/, '')
54    if (core.length <= 2) continue
55    if (!CORE_RE.test(core)) continue
56    if (ALLCAPS_RE.test(core)) continue
57    if (INTERNAL_CAPS_RE.test(core.slice(1))) continue
58    if (HEX_RE.test(core)) continue
59
60    tokens.push({
61      word: core,
62      start: chunkStart + lead,
63      end: chunkStart + lead + core.length,
64      isCapitalized: /^[A-Z]/.test(core),
65    })
66  }
67  return tokens
68}
69
70export function accept(word: string, dict: Accepts, personal: Accepts): boolean {
71  const w = word.toLowerCase()
72  return dict.has(w) || personal.has(w)
73}
74
75export function flagged(tokens: Token[], dict: Accepts, personal: Accepts): Token[] {
76  return tokens.filter(t => !accept(t.word, dict, personal))
77}
78

1export type SpellActive = {
2  word: string
3  start: number
4  end: number
5  isCapitalized: boolean
6  kind: 'spelling' | 'contraction'
7  fix?: string
8} | null
9
10export type SpellDictStatus = 'loading' | 'ready' | 'fallback' | 'missing'
11
12declare module 'claude-code' {
13  interface PluginState {
14    'cc-spell-check': {
15      active: SpellActive
16      isEnabled: boolean
17      dictStatus: SpellDictStatus
18    }
19  }
20}
21