test-judge

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { Baseline, TestEntry } from '../types'
5import { changedTests, extractTests, isTestFile, parseVerdict } from './parse'
6
7const PANE = 'test-judge'
8const TITLE = 'Tests'
9// No engine key reaches a plugin directly: the band's Button takes this
10// action, which ~/.claude/keybindings.json binds to ctrl+x t.
11const ACTION = 'app:toggleDiffPreSession'
12const MODEL = 'claude-haiku-5-5'
13
14const tests = atom({ plugin: 'test-judge', key: 'tests' } as const, [])
15const baseline = atom({ plugin: 'test-judge', key: 'baseline' } as const, null)
16const cost = atom({ plugin: 'test-judge', key: 'cost' } as const, 0)
17
18// Haiku 5.5's list prices in US dollars per million tokens, by prompt length:
19// up to 100K tokens, and beyond. Cache reads are 0.1x input, 5-minute cache
20// writes 1.25x.
21const LONG_PROMPT = 100_000
22const RATES = { short: { input: 0.1, output: 0.5 }, long: { input: 0.5, output: 2.5 } } as const
23
24function priceOf(usage: ModelUsage) {
25  const prompt = usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
26  const { input, output } = prompt > LONG_PROMPT ? RATES.long : RATES.short
27  const tokens =
28    usage.input_tokens * input +
29    usage.cache_read_input_tokens * input * 0.1 +
30    usage.cache_creation_input_tokens * input * 1.25 +
31    usage.output_tokens * output
32
33  return tokens / 1_000_000
34}
35
36const RUBRIC = `You review unit tests written during test-driven development.
37A BAD test only exercises hardcoded values: it compares literals with literals, checks a constant or a value the test itself just assigned, re-states the implementation's constants, or would pass with any implementation or none.
38A GOOD test drives the code under test and asserts on behaviour a wrong implementation would fail. Literal inputs and expected values are fine when they go through the code under test.
39Answer on one line, nothing else: "GOOD: <reason, under 12 words>" or "BAD: <reason, under 12 words>".`
40
41const COLORS = { good: 'success', bad: 'error', pending: 'warning', unknown: 'subtle' } as const
42const MARKS = { good: '●', bad: '●', pending: '○', unknown: '?' } as const
43
44type $ = EngineInterface
45
46async function readOr(file: string, $: $) {
47  try {
48    return await $.fs.read(file)
49  } catch {
50    return ''
51  }
52}
53
54// A git command's output, or undefined when git failed or could not start.
55async function git($: $, args: string[], cwd?: string) {
56  try {
57    const ran = await $.process.run(['git', ...args], cwd === undefined ? {} : { cwd })
58
59    return ran.exitCode === 0 ? ran.stdout : undefined
60  } catch {
61    return undefined
62  }
63}
64
65const nulPaths = (stdout = '') => stdout.split('\0').filter(path => path !== '')
66
67// The test files that differ from `head` (every file before the first
68// commit), untracked ones included, by path from the root.
69async function changedTestFiles($: $, root: string, head?: string) {
70  const tracked = head === undefined ? '' : await git($, ['diff', '--name-only', '-z', '--no-renames', head], root)
71  const untracked = await git($, ['ls-files', '-z', '--others', '--exclude-standard'], root)
72
73  return [...new Set([...nulPaths(tracked), ...nulPaths(untracked)])].filter(isTestFile)
74}
75
76async function takeBaseline($: $): Promise<Baseline | null> {
77  const root = (await git($, ['rev-parse', '--show-toplevel']))?.trim()
78  if (!root) {
79    return null
80  }
81
82  const head = (await git($, ['rev-parse', '--verify', '--quiet', 'HEAD'], root))?.trim() || undefined
83  const paths = await changedTestFiles($, root, head)
84  const texts = await Promise.all(paths.map(path => readOr(`${root}/${path}`, $)))
85
86  return { root, head, files: Object.fromEntries(paths.map((path, at) => [path, texts[at] ?? ''])) }
87}
88
89async function textBefore($: $, base: Baseline, path: string) {
90  const kept = base.files[path]
91  if (kept !== undefined) {
92    return kept
93  }
94
95  return base.head === undefined ? '' : ((await git($, ['show', `${base.head}:${path}`], base.root)) ?? '')
96}
97
98// Records every test the turn added or rewrote since `base`, and drops the
99// ones it deleted; answers whether any is left to judge.
100async function recordTurn($: $, base: Baseline) {
101  const paths = [...new Set([...(await changedTestFiles($, base.root, base.head)), ...Object.keys(base.files)])]
102  let isJudgeable = false
103
104  for (const path of paths) {
105    const file = `${base.root}/${path}`
106    const before = extractTests(file, await textBefore($, base, path))
107    const after = extractTests(file, await readOr(file, $))
108    const { changed, removed } = changedTests(before, after)
109
110    await update($, tests, list => {
111      const gone = new Set(removed.map(name => `${file}::${name}`))
112      const fresh: TestEntry[] = changed.map(one => ({
113        id: `${file}::${one.name}`,
114        file: path,
115        name: one.name,
116        body: one.body,
117        verdict: 'pending',
118        reason: '',
119      }))
120      const replaced = new Set(fresh.map(one => one.id))
121      const kept = list.filter(one => !gone.has(one.id) && !replaced.has(one.id))
122
123      return [...kept, ...fresh]
124    })
125    isJudgeable ||= changed.length > 0
126  }
127
128  return isJudgeable
129}
130
131async function togglePane($: $) {
132  const isOpen = (await $.ui.panes()).some(pane => pane.id === PANE)
133
134  if (isOpen) {
135    await $.ui.close({ id: PANE })
136  } else {
137    await $.ui.open({ id: PANE, title: TITLE })
138  }
139}
140
141let isJudging = false
142
143// One test at a time, oldest first, until none is pending; a reload starts
144// over from what `$.state` still holds.
145async function judgePending($: $) {
146  if (isJudging) {
147    return
148  }
149  isJudging = true
150  try {
151    for (;;) {
152      const next = (await read($, tests)).find(one => one.verdict === 'pending')
153      if (!next) {
154        return
155      }
156
157      const reply = await $.model.complete({
158        model: MODEL,
159        system: RUBRIC,
160        prompt: `File: ${next.file}\n\n${next.body.slice(0, 6000)}`,
161        maxTokens: 60,
162      })
163      await update($, cost, spent => spent + priceOf(reply.usage))
164      const judged = reply.isAnswered
165        ? parseVerdict(reply.text)
166        : { verdict: 'unknown' as const, reason: `not judged: ${reply.reason}` }
167
168      await update($, tests, list =>
169        list.map(one => (one.id === next.id && one.body === next.body ? { ...one, ...judged } : one)),
170      )
171    }
172  } finally {
173    isJudging = false
174  }
175}
176
177export const register: Register = on => {
178  on('session.start', async ($, e, next) => {
179    await $.command.register({ name: 'tests', description: 'Toggle the pane of tests written this session' })
180    $.clock.after(0, () => judgePending($))
181
182    return next(e)
183  })
184
185  on('command.run', { command: 'tests' }, async $ => {
186    await togglePane($)
187
188    return { text: 'Toggled the tests pane.' }
189  })
190
191  // A main turn's edits are found at its end against the repository as it
192  // began, so tests written by any tool count, Bash included; a subagent's
193  // turns fall inside the main turn that spawned them.
194  on('turn.start', async ($, e, next) => {
195    if (e.agentId === undefined) {
196      const found = await takeBaseline($)
197      await update($, baseline, () => found)
198    }
199
200    return next(e)
201  }).catch(($, e, next) => next(e))
202
203  on('turn.complete', async ($, e, next) => {
204    const base = e.agentId === undefined ? await read($, baseline) : null
205    if (base) {
206      await update($, baseline, () => null)
207      if (await recordTurn($, base)) {
208        $.clock.after(0, () => judgePending($))
209      }
210    }
211
212    return next(e)
213  }).catch(($, e, next) => next(e))
214
215  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
216    const { Box, Text } = $.ui.resolve(e)
217    const list = await read($, tests)
218    const spent = <Text dimColor>Haiku cost: ${(await read($, cost)).toFixed(4)}</Text>
219
220    if (list.length === 0) {
221      return (
222        <Box flexDirection="column">
223          {spent}
224          <Text dimColor>No tests written this session yet.</Text>
225        </Box>
226      )
227    }
228
229    const files = [...new Set(list.map(one => one.file))]
230
231    return (
232      <Box flexDirection="column">
233        <Box marginBottom={1}>{spent}</Box>
234        {files.map(file => (
235          <Box key={file} flexDirection="column" marginBottom={1}>
236            <Text color="inactive">{file}</Text>
237            {list
238              .filter(one => one.file === file)
239              .map(one => (
240                <Box key={one.id} flexDirection="column">
241                  <Text color={COLORS[one.verdict]}>
242                    {MARKS[one.verdict]} {one.name}
243                  </Text>
244                  {one.verdict !== 'good' && one.reason !== '' && <Text dimColor>  {one.reason}</Text>}
245                </Box>
246              ))}
247          </Box>
248        ))}
249      </Box>
250    )
251  })
252
253  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
254    if (e.props.hasSurvey) {
255      return next(e)
256    }
257
258    const { Button, Text } = $.ui.resolve(e)
259    const list = await read($, tests)
260
261    // The count itself is the button, so ctrl+x t still reaches the pane.
262    return (
263      <Button key="toggle" plain action={ACTION} onPress={() => togglePane($)}>
264        <Text dimColor>Tests: {list.length}</Text>
265      </Button>
266    )
267  })
268}
269

1import type { Verdict } from '../types'
2
3export type FoundTest = { name: string; body: string }
4
5const TEST_FILES = [
6  /_test\.go$/,
7  /\.(test|spec)\.[cm]?[jt]sx?$/,
8  /(^|\/)test_[^/]*\.py$/,
9  /_test\.py$/,
10]
11
12export const isTestFile = (path: string) => TEST_FILES.some(rule => rule.test(path))
13
14export function extractTests(path: string, text: string): FoundTest[] {
15  if (path.endsWith('.go')) {
16    return extractBraced(text, /^func\s+(Test\w*)\s*\(/gm, '{', '}')
17  }
18  if (path.endsWith('.py')) {
19    return extractPython(text)
20  }
21
22  return extractBraced(
23    text,
24    /\b(?:test|it)(?:\.only|\.skip|\.todo)?\s*\(\s*(['"`])((?:\\.|(?!\1)[^\\])*)\1/g,
25    '(',
26    ')',
27  )
28}
29
30// Each match's body runs from the match to the bracket closing the first
31// `open` at or after it (for JS, the call's own parenthesis).
32function extractBraced(text: string, rule: RegExp, open: string, close: string): FoundTest[] {
33  const found: FoundTest[] = []
34
35  for (const match of text.matchAll(rule)) {
36    const start = match.index ?? 0
37    const name = match[2] ?? match[1] ?? ''
38    const from = open === '(' ? text.indexOf('(', start) : text.indexOf(open, start + match[0].length)
39    const end = from < 0 ? -1 : closing(text, from, open, close)
40    found.push({ name, body: text.slice(start, end < 0 ? text.length : end + 1) })
41  }
42
43  return found
44}
45
46// The index of the bracket closing the one at `from`, skipping strings and
47// line comments; -1 when the text ends first.
48function closing(text: string, from: number, open: string, close: string) {
49  let depth = 0
50
51  for (let i = from; i < text.length; i += 1) {
52    const char = text.charAt(i)
53
54    if (char === '"' || char === "'" || char === '`') {
55      i = stringEnd(text, i)
56    } else if (char === '/' && text[i + 1] === '/') {
57      i = text.indexOf('\n', i)
58      if (i < 0) {
59        return -1
60      }
61    } else if (char === open) {
62      depth += 1
63    } else if (char === close) {
64      depth -= 1
65      if (depth === 0) {
66        return i
67      }
68    }
69  }
70
71  return -1
72}
73
74function stringEnd(text: string, from: number) {
75  const quote = text[from]
76
77  for (let i = from + 1; i < text.length; i += 1) {
78    if (text.charAt(i) === '\\' && quote !== '`') {
79      i += 1
80    } else if (text.charAt(i) === quote) {
81      return i
82    } else if (text.charAt(i) === '\n' && quote !== '`') {
83      return i
84    }
85  }
86
87  return text.length
88}
89
90function extractPython(text: string): FoundTest[] {
91  const lines = text.split('\n')
92  const found: FoundTest[] = []
93
94  lines.forEach((line, at) => {
95    const match = /^([ \t]*)(?:async\s+)?def\s+(test\w*)\s*\(/.exec(line)
96    if (!match) {
97      return
98    }
99
100    const indent = (match[1] ?? '').length
101    let last = at
102    for (let i = at + 1; i < lines.length; i += 1) {
103      const next = lines[i] ?? ''
104      if (next.trim() === '') {
105        continue
106      }
107      if (next.length - next.trimStart().length <= indent) {
108        break
109      }
110      last = i
111    }
112    found.push({ name: match[2] ?? '', body: lines.slice(at, last + 1).join('\n') })
113  })
114
115  return found
116}
117
118export function changedTests(before: FoundTest[], after: FoundTest[]) {
119  const was = new Map(before.map(one => [one.name, one.body]))
120  const names = new Set(after.map(one => one.name))
121
122  return {
123    changed: after.filter(one => was.get(one.name) !== one.body),
124    removed: before.filter(one => !names.has(one.name)).map(one => one.name),
125  }
126}
127
128export function parseVerdict(reply: string): { verdict: Verdict; reason: string } {
129  const match = /^\s*(good|bad)\b\s*[:\-–]?\s*(.*)/i.exec(reply.split('\n')[0] ?? '')
130
131  if (!match) {
132    return { verdict: 'unknown', reason: reply.trim().slice(0, 80) }
133  }
134
135  return {
136    verdict: match[1]?.toLowerCase() === 'good' ? 'good' : 'bad',
137    reason: (match[2] ?? '').trim(),
138  }
139}
140

1export type Verdict = 'pending' | 'good' | 'bad' | 'unknown'
2
3export type TestEntry = {
4  id: string
5  file: string
6  name: string
7  body: string
8  verdict: Verdict
9  reason: string
10}
11
12// The repository as a main turn found it: its root, its HEAD commit (absent
13// before the first commit) and the test files it already had changed, by
14// path from the root, with their text.
15export type Baseline = {
16  root: string
17  head?: string
18  files: Record<string, string>
19}
20
21declare module 'claude-code' {
22  interface PluginState {
23    // `cost`: US dollars spent judging with Haiku this session.
24    'test-judge': { tests: TestEntry[]; baseline: Baseline | null; cost: number }
25  }
26}
27