ablation

1import type { EngineInterface, InstructionFile, Register } from 'claude-code'
2
3// Harness A/B without editing CLAUDE.md: each active experiment in
4// ~/.claude/harness-ablation.json assigns this session to `treatment` or
5// `control` by a hash of (session id, experiment id), so the arm is stable
6// across reloads, compactions and subagents. Treatment drops the instruction
7// files whose path ends with one of `drop`. Every assignment is recorded in
8// ~/.claude/harness-ablation/<session id>.json for agentlogs to join on.
9// No config, or no active experiment: the context passes through untouched.
10
11export type Experiment = {
12  id: string
13  drop: string[]
14  treatmentShare: number
15  active: boolean
16  hypothesis?: string
17}
18
19export type Assignment = { experiment: string; arm: 'treatment' | 'control'; dropped: string[]; droppedChars: number }
20
21type $ = EngineInterface
22
23export async function bucket(sessionId: string, experimentId: string): Promise<number> {
24  const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(`${sessionId}\0${experimentId}`))
25  return new DataView(digest).getUint32(0) / 2 ** 32
26}
27
28async function home($: $): Promise<string> {
29  return (await $.env.get('HOME')) ?? '/Users/alien'
30}
31
32async function experiments($: $): Promise<Experiment[]> {
33  const path = `${await home($)}/.claude/harness-ablation.json`
34  if (!(await $.fs.exists(path))) return []
35  const config = JSON.parse(await $.fs.read(path)) as { experiments?: Experiment[] }
36  return (config.experiments ?? []).filter(x => x.active && x.drop.length > 0)
37}
38
39export async function assign(
40  sessionId: string,
41  list: readonly Experiment[],
42  files: readonly InstructionFile[],
43): Promise<{ kept: InstructionFile[]; assignments: Assignment[] }> {
44  let kept = [...files]
45  const assignments: Assignment[] = []
46  for (const x of list) {
47    const arm = (await bucket(sessionId, x.id)) < x.treatmentShare ? 'treatment' : 'control'
48    const hit = arm === 'treatment' ? kept.filter(f => x.drop.some(d => f.path.endsWith(d))) : []
49    kept = kept.filter(f => !hit.includes(f))
50    assignments.push({
51      experiment: x.id,
52      arm,
53      dropped: hit.map(f => f.path),
54      droppedChars: hit.reduce((n, f) => n + f.content.length, 0),
55    })
56  }
57  return { kept, assignments }
58}
59
60export const register: Register = on => {
61  on('session.start', async ($, e, next) => {
62    await $.command.register({ name: 'ablation', description: "Show this session's harness-ablation arms" })
63    return next(e)
64  })
65
66  on('prompt.context', async ($, e, next) => {
67    const below = await next(e)
68    const list = await experiments($)
69    const files = below.instructionFiles ?? e.instructionFiles
70    if (list.length === 0 || files === undefined) return below
71    const sessionId = await $.session.id()
72    const { kept, assignments } = await assign(sessionId, list, files)
73    const root = `${await home($)}/.claude/harness-ablation`
74    await $.fs.write(
75      `${root}/${sessionId}.json`,
76      JSON.stringify({ session_id: sessionId, cwd: await $.session.cwd(), at: new Date(await $.clock.now()).toISOString(), assignments }, null, 2),
77    )
78    return kept.length === files.length ? below : { ...below, instructionFiles: kept }
79  })
80
81  on('command.run', { command: 'ablation' }, async $ => {
82    const path = `${await home($)}/.claude/harness-ablation/${await $.session.id()}.json`
83    if (!(await $.fs.exists(path))) return { text: 'No ablation experiment applies to this session.' }
84    const rec = JSON.parse(await $.fs.read(path)) as { assignments: Assignment[] }
85    return {
86      text: rec.assignments
87        .map(a => `${a.experiment}: ${a.arm}${a.dropped.length ? ` (dropped ${a.dropped.length} files, ${a.droppedChars} chars)` : ''}`)
88        .join('\n'),
89    }
90  })
91}
92