cc-code-format-mod

1// cc-code-format-mod
2//
3// Rewrites every fenced code block in an LLM response so it conforms to the
4// language's standard style. Hooks `session.append { door: 'response' }`,
5// walks the text blocks, finds ```lang\n<body>\n``` fences, dispatches each
6// body to a language-appropriate formatter, and rewrites the row before the
7// transcript stores it.
8//
9// The formatting logic lives in `./format.ts` so unit tests can exercise it
10// directly without firing session.append through the test kit.
11
12import { atom, update } from 'claude-code'
13import type { Register } from 'claude-code'
14import type { Stats } from '../types'
15import { formatText, type FmtDeps } from './format'
16
17// ---------------------------------------------------------------------------
18// Persistent state (survives hot reload; keyed by plugin + key)
19// ---------------------------------------------------------------------------
20
21const statsAtom = atom(
22  { plugin: 'cc-code-format-mod', key: 'stats' } as const,
23  { formatted: 0, skipped: 0, failed: 0 } as Stats,
24)
25
26// Mirrors of `userConfig.enabled` (re-read on each (re)load) and a session
27// toggle for verbose toasts. Module-scoped: a reload re-reads userConfig, and
28// verbose is purely a /format-code runtime switch.
29let configuredEnabled = true
30let configuredVerbose = false
31
32// ---------------------------------------------------------------------------
33// Register
34// ---------------------------------------------------------------------------
35
36export const register: Register = (on, options) => {
37  configuredEnabled = options.enabled !== false
38
39  on('session.start', async ($, e, next) => {
40    await $.command.register({
41      name: 'format-code',
42      description: 'Toggle / inspect the LLM-output code formatter',
43    })
44    return next(e)
45  })
46
47  on('command.run', { command: 'format-code' }, async ($, e) => {
48    const raw = (e.args ?? '').trim()
49    const tokens = raw.toLowerCase().split(/\s+/).filter(Boolean)
50    const verb = tokens[0] ?? ''
51    const sub = tokens[1] ?? ''
52
53    if (verb === 'help') {
54      return {
55        text: [
56          '/format-code                  toggle on/off',
57          '/format-code on               enable',
58          '/format-code off              disable',
59          '/format-code status           show on/off and verbose state',
60          '/format-code verbose          toggle verbose toasts',
61          '/format-code verbose on|off   set verbose mode',
62        ].join('\n'),
63      }
64    }
65
66    if (verb === 'verbose') {
67      if (sub === 'on') {
68        configuredVerbose = true
69        return { text: 'verbose on (toasts after each rewrite).' }
70      }
71      if (sub === 'off') {
72        configuredVerbose = false
73        return { text: 'verbose off.' }
74      }
75      configuredVerbose = !configuredVerbose
76      return {
77        text: `verbose ${configuredVerbose ? 'on' : 'off'}.`,
78      }
79    }
80
81    if (verb === 'on') {
82      configuredEnabled = true
83      return { text: 'enabled.' }
84    }
85
86    if (verb === 'off') {
87      configuredEnabled = false
88      return { text: 'disabled.' }
89    }
90
91    if (verb === 'status') {
92      return {
93        text:
94          `${configuredEnabled ? 'enabled' : 'disabled'}. ` +
95          `verbose: ${configuredVerbose ? 'on' : 'off'}.`,
96      }
97    }
98
99    if (verb === '' || verb === 'toggle') {
100      configuredEnabled = !configuredEnabled
101      return {
102        text: `${configuredEnabled ? 'enabled' : 'disabled'}.`,
103      }
104    }
105
106    return {
107      text: `unknown argument "${e.args}". Try /format-code help.`,
108    }
109  })
110
111  on('session.append', { door: 'response' }, async ($, e, next) => {
112    if (!configuredEnabled) return next(e)
113
114    const content = e.message.content as Array<{
115      type: string
116      text?: unknown
117      [field: string]: unknown
118    }>
119
120    let totalFormatted = 0
121    let totalFailed = 0
122    let anyChanged = false
123
124    // $ cannot cross an import boundary (the engine's capability check), so
125    // detach the methods we use and pass them as plain function values.
126    const deps: FmtDeps = {
127      runProcess: (argv, init) => $.process.run(argv, init),
128      writeFile: (path, text) => $.fs.write(path, text),
129      readFile: (path) => $.fs.read(path),
130    }
131
132    // Sequential so per-block file-based formatters don't race on /tmp.
133    const newContent: typeof content = []
134    for (const block of content) {
135      if (block.type !== 'text' || typeof block.text !== 'string') {
136        newContent.push(block)
137        continue
138      }
139      const result = await formatText(block.text, deps)
140      totalFormatted += result.formatted
141      totalFailed += result.failed
142      if (result.text === block.text) {
143        newContent.push(block)
144      } else {
145        newContent.push({ ...block, text: result.text })
146        anyChanged = true
147      }
148    }
149
150    if (totalFormatted > 0 || totalFailed > 0) {
151      await update($, statsAtom, prev => ({
152        formatted: prev.formatted + totalFormatted,
153        skipped: prev.skipped,
154        failed: prev.failed + totalFailed,
155      }))
156
157      if (configuredVerbose) {
158        if (totalFormatted > 0) {
159          $.ui.toast(
160            `formatted ${totalFormatted} code block${totalFormatted === 1 ? '' : 's'}`,
161          )
162        }
163        if (totalFailed > 0) {
164          $.ui.toast(
165            `skipped ${totalFailed} block${totalFailed === 1 ? '' : 's'} (formatter error)`,
166            { timeoutMs: 6000 },
167          )
168        }
169      }
170    }
171
172    if (!anyChanged) return next(e)
173    return next({ ...e, message: { ...e.message, content: newContent } })
174  })
175}

1// Pure formatting helpers. The session.append hook in register.tsx extracts
2// the methods it needs from $ and passes them as plain function values so the
3// engine's capability check (which forbids passing $ across an import) stays
4// happy. unit tests pass mock functions in directly.
5
6export interface FmtDeps {
7  runProcess: (
8    argv: readonly string[],
9    init?: { stdin?: string; timeoutMs?: number; cwd?: string },
10  ) => Promise<{
11    exitCode: number
12    stdout: string
13    stderr: string
14    isStdoutTruncated: boolean
15    isStderrTruncated: boolean
16  }>
17  writeFile: (path: string, text: string) => Promise<void>
18  readFile: (path: string) => Promise<string>
19}
20
21// ---------------------------------------------------------------------------
22// Formatter tables
23// ---------------------------------------------------------------------------
24
25export const PRETTIER_PARSERS: Record<string, string> = {
26  javascript: 'babel',
27  js: 'babel',
28  jsx: 'babel',
29  mjs: 'babel',
30  cjs: 'babel',
31  flow: 'flow',
32  typescript: 'typescript',
33  ts: 'typescript',
34  tsx: 'tsx',
35  json: 'json',
36  json5: 'json5',
37  jsonc: 'json',
38  css: 'css',
39  scss: 'scss',
40  less: 'less',
41  html: 'html',
42  htm: 'html',
43  vue: 'vue',
44  svelte: 'svelte',
45  astro: 'astro',
46  angular: 'angular',
47  yaml: 'yaml',
48  yml: 'yaml',
49  markdown: 'markdown',
50  md: 'markdown',
51  mdx: 'mdx',
52  graphql: 'graphql',
53  gql: 'graphql',
54  toml: 'toml',
55}
56
57export const TOOL_LANGS: Record<string, { tool: string; argv: readonly string[] }> = {
58  // `--code` is for inline mode and conflicts with stdin `-`. With `-`, black
59  // reads from stdin and prints the formatted result to stdout.
60  python: { tool: 'python3', argv: ['-m', 'black', '-', '--quiet'] },
61  py: { tool: 'python3', argv: ['-m', 'black', '-', '--quiet'] },
62  go: { tool: 'gofmt', argv: [] },
63  golang: { tool: 'gofmt', argv: [] },
64  bash: { tool: 'shfmt', argv: ['-i', '2', '-'] },
65  sh: { tool: 'shfmt', argv: ['-i', '2', '-'] },
66  shell: { tool: 'shfmt', argv: ['-i', '2', '-'] },
67  zsh: { tool: 'shfmt', argv: ['-i', '2', '-'] },
68}
69
70export const FILE_LANGS: Record<string, { tool: string; ext: string }> = {
71  rust: { tool: 'rustfmt', ext: 'rs' },
72  rs: { tool: 'rustfmt', ext: 'rs' },
73}
74
75export const SKIP_LANGS = new Set<string>([
76  '',
77  'diff',
78  'plaintext',
79  'text',
80  'console',
81  'log',
82  'shell-session',
83  'shellsession',
84  'ansi',
85  'output',
86  'txt',
87])
88
89// ---------------------------------------------------------------------------
90// Fence regex: ```<lang> [attrs]\n<body>```
91//   group 1 = lang tag, group 2 = body
92// ---------------------------------------------------------------------------
93
94export const FENCE_RE = /```([A-Za-z0-9_+\-]*)[^\n]*\n([\s\S]*?)```/g
95
96// ---------------------------------------------------------------------------
97// Timeouts
98// ---------------------------------------------------------------------------
99
100export const PRETTIER_TIMEOUT_MS = 15_000
101export const NATIVE_TIMEOUT_MS = 10_000
102export const TEMP_DIR = `/tmp/cc-code-format-mod`
103
104// ---------------------------------------------------------------------------
105// Pure helpers (testable without firing session.append)
106// ---------------------------------------------------------------------------
107
108export interface TextResult {
109  text: string
110  formatted: number
111  failed: number
112}
113
114interface Fence {
115  lang: string
116  body: string
117  start: number
118  end: number
119  raw: string
120}
121
122export async function formatText(text: string, deps: FmtDeps): Promise<TextResult> {
123  const fences: Fence[] = []
124  FENCE_RE.lastIndex = 0
125  let m: RegExpExecArray | null
126  while ((m = FENCE_RE.exec(text)) !== null) {
127    fences.push({
128      lang: (m[1] ?? '').toLowerCase(),
129      body: m[2],
130      start: m.index,
131      end: m.index + m[0].length,
132      raw: m[0],
133    })
134  }
135  if (fences.length === 0) return { text, formatted: 0, failed: 0 }
136
137  let out = ''
138  let cursor = 0
139  let formatted = 0
140  let failed = 0
141  for (const f of fences) {
142    out += text.slice(cursor, f.start)
143    try {
144      const replacement = await formatBlock(f.lang, f.body, deps)
145      if (replacement !== null) {
146        out += '```' + f.lang + '\n' + replacement + '```'
147        formatted += 1
148      } else {
149        out += f.raw
150      }
151    } catch {
152      out += f.raw
153      failed += 1
154    }
155    cursor = f.end
156  }
157  out += text.slice(cursor)
158  return { text: out, formatted, failed }
159}
160
161export async function formatBlock(
162  lang: string,
163  code: string,
164  deps: FmtDeps,
165): Promise<string | null> {
166  if (SKIP_LANGS.has(lang)) return null
167  if (!code.trim()) return null
168
169  const parser = PRETTIER_PARSERS[lang]
170  if (parser !== undefined) {
171    return runStdin(
172      'npx',
173      [
174        '--yes',
175        'prettier@3',
176        '--stdin-filepath',
177        `file.${lang || 'txt'}`,
178        '--parser',
179        parser,
180      ],
181      code,
182      deps,
183      PRETTIER_TIMEOUT_MS,
184    )
185  }
186
187  const tool = TOOL_LANGS[lang]
188  if (tool !== undefined) {
189    return runStdin(tool.tool, tool.argv, code, deps, NATIVE_TIMEOUT_MS)
190  }
191
192  const fileTool = FILE_LANGS[lang]
193  if (fileTool !== undefined) {
194    return formatWithFile(fileTool.tool, fileTool.ext, code, deps)
195  }
196
197  return null
198}
199
200export async function runStdin(
201  cmd: string,
202  argv: readonly string[],
203  stdin: string,
204  deps: FmtDeps,
205  timeoutMs: number,
206): Promise<string | null> {
207  try {
208    const result = await deps.runProcess([cmd, ...argv], { stdin, timeoutMs })
209    if (result.exitCode === 0 && result.stdout.length > 0) {
210      return result.stdout.replace(/\s+$/, '') + '\n'
211    }
212  } catch {
213    /* formatter unavailable, network down, timeout, etc. */
214  }
215  return null
216}
217
218export async function formatWithFile(
219  tool: string,
220  ext: string,
221  code: string,
222  deps: FmtDeps,
223): Promise<string | null> {
224  const fileName = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}.${ext}`
225  const path = `${TEMP_DIR}/${fileName}`
226  try {
227    await deps.writeFile(path, code)
228    const result = await deps.runProcess([tool, path], { timeoutMs: NATIVE_TIMEOUT_MS })
229    if (result.exitCode === 0) {
230      const out = await deps.readFile(path)
231      return out.replace(/\s+$/, '') + '\n'
232    }
233  } catch {
234    /* ignore */
235  }
236  return null
237}

1export type Stats = { formatted: number; skipped: number; failed: number }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'cc-code-format-mod': {
6      stats: { value: Stats }
7    }
8  }
9}