proofread

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PromptOrigin, Register } from 'claude-code'
3
4import type { Corrections } from '../types'
5
6const corrections = atom({ plugin: 'proofread', key: 'corrections' } as const, {})
7const MAX_REMEMBERED = 50
8const ACCENT = 'suggestion'
9const FIX = 'success'
10
11const MODEL = 'haiku'
12const MAX_LENGTH = 4000
13const MIN_WORDS = 3
14
15const SYSTEM = `You are a proofreader. The user message is a prompt someone typed to a coding assistant.
16Fix only spelling, typos and clear grammar mistakes. Do not rephrase, reword, shorten, expand, change tone, or answer it.
17Leave untouched: code, code blocks, inline code, file paths, URLs, identifiers, class/function/variable names, CLI commands, slang, technical jargon, product names, and markdown.
18Reply with the corrected prompt text only, nothing else. If nothing needs fixing, reply with the text exactly as given.`
19
20const isHuman = (origin: PromptOrigin | undefined): boolean =>
21  origin === undefined ||
22  origin.kind === 'composer' ||
23  origin.kind === 'bridge' ||
24  (origin.kind === 'plugin' && origin.asUser === true)
25
26const shouldCheck = (text: string): boolean => {
27  const trimmed = text.trim()
28
29  return (
30    trimmed.length <= MAX_LENGTH &&
31    !trimmed.startsWith('/') &&
32    !trimmed.startsWith('!') &&
33    trimmed.split(/\s+/).length >= MIN_WORDS
34  )
35}
36
37const words = (text: string): string[] => text.split(/\s+/).filter(word => word !== '')
38
39const changedPositions = (before: string, after: string): number[] => {
40  const a = words(before)
41  const b = words(after)
42  const lengths = Array.from({ length: a.length + 1 }, () => new Array<number>(b.length + 1).fill(0))
43
44  for (let i = a.length - 1; i >= 0; i--) {
45    for (let j = b.length - 1; j >= 0; j--) {
46      const row = lengths[i]!
47      row[j] = a[i] === b[j] ? lengths[i + 1]![j + 1]! + 1 : Math.max(lengths[i + 1]![j]!, row[j + 1]!)
48    }
49  }
50
51  const changed: number[] = []
52  let i = 0
53  let j = 0
54  while (j < b.length) {
55    if (i < a.length && a[i] === b[j]) {
56      i++
57      j++
58    } else if (i < a.length && lengths[i + 1]![j]! >= lengths[i]![j + 1]!) {
59      i++
60    } else {
61      changed.push(j)
62      j++
63    }
64  }
65
66  return changed
67}
68
69const OFF_KEY = 'isOff'
70const OFF_STATUS = 'proofread: off'
71
72const isOff = async ($: EngineInterface): Promise<boolean> => (await $.store.get(OFF_KEY)) === true
73
74const setOff = async ($: EngineInterface, off: boolean): Promise<void> => {
75  await $.store.set(OFF_KEY, off)
76  $.ui.status(off ? OFF_STATUS : undefined)
77}
78
79export const register: Register = on => {
80  on('session.start', async ($, e, next) => {
81    await $.command
82      .register({
83        name: 'proofread',
84        description: 'Turn prompt proofreading on or off',
85        argumentHint: '[on|off|status]',
86        immediate: true,
87      })
88      .catch(() => $.ui.toast('proofread: could not register /proofread'))
89    $.ui.status((await isOff($)) ? OFF_STATUS : undefined)
90
91    return next(e)
92  })
93
94  on('command.run', { command: 'proofread' }, async ($, e) => {
95    const arg = e.args.trim().toLowerCase()
96    const wasOff = await isOff($)
97
98    if (arg === 'status') {
99      return { text: `Proofreading is ${wasOff ? 'off' : 'on'}.` }
100    }
101
102    const off = arg === 'off' ? true : arg === 'on' ? false : !wasOff
103    await setOff($, off)
104
105    return { text: `Proofreading is now ${off ? 'off' : 'on'}.` }
106  })
107
108  on('prompt.submit', async ($, e, next) => {
109    if (!isHuman(e.origin) || e.turnId !== undefined || !shouldCheck(e.text) || (await isOff($))) {
110      return next(e)
111    }
112
113    $.ui.status('proofread: checking…')
114    const reply = await $.model.complete({
115      model: MODEL,
116      system: SYSTEM,
117      prompt: e.text,
118      maxTokens: Math.ceil(e.text.length / 2) + 256,
119      timeoutMs: 8000,
120    })
121    $.ui.status(undefined)
122
123    if (!reply.isAnswered) {
124      return next(e)
125    }
126
127    const corrected = reply.text.trim()
128    const ratio = corrected.length / e.text.trim().length
129
130    if (corrected === e.text.trim() || ratio < 0.8 || ratio > 1.2) {
131      return next(e)
132    }
133
134    const positions = changedPositions(e.text, corrected)
135    const correctedWords = words(corrected)
136    const fixes = positions.map(position => correctedWords[position] ?? '')
137    await update($, corrections, (all: Corrections) => {
138      const kept = Object.entries(all).slice(-(MAX_REMEMBERED - 1))
139
140      return Object.fromEntries([...kept, [corrected, positions]])
141    })
142    $.ui.toast(
143      fixes.length > 0
144        ? `proofread fixed: ${fixes.slice(0, 6).join(', ')}${fixes.length > 6 ? '…' : ''}`
145        : 'proofread fixed punctuation',
146    )
147
148    return next({ ...e, text: corrected })
149  }).catch(async ($, e, next) => {
150    $.ui.status((await isOff($)) ? OFF_STATUS : undefined)
151
152    return next(e)
153  })
154
155  on('ui.render', { component: 'UserMessage', props: { origin: { kind: 'composer' } } }, async ($, e, next) => {
156    const text = e.props.text.trim()
157    const { Box, Text } = $.ui.resolve(e)
158    const fixes = (await read($, corrections))[text]
159    const fixed = new Set(fixes ?? [])
160    let position = -1
161
162    return (
163      <Box flexDirection="column" borderStyle="round" borderColor={ACCENT} paddingX={1}>
164        <Text>
165          <Text color={ACCENT} bold>
166            {'❯ '}
167          </Text>
168          {text.split(/(\s+)/).map(part =>
169            part.trim() !== '' && fixed.has(++position) ? (
170              <Text color={FIX} bold underline>
171                {part}
172              </Text>
173            ) : (
174              part
175            ),
176          )}
177        </Text>
178        {fixes !== undefined ? (
179          <Text dimColor>
180            ✓ proofread fixed {fixes.length > 0 ? fixes.length : 'punctuation'}
181            {fixes.length === 1 ? ' word' : fixes.length > 1 ? ' words' : ''}
182          </Text>
183        ) : null}
184      </Box>
185    )
186  })
187}
188

1export type Corrections = Record<string, readonly number[]>
2
3declare module 'claude-code' {
4  interface PluginState {
5    proofread: { corrections: Corrections }
6  }
7}
8