feat(bench): JIT arm — real baseline-vs-JIT break-even

cmd/aj/bench_cmd.go

@@ -48,6 +48,19 @@ var benchCmd = &cobra.Command{
 			return fmt.Errorf("--rollouts must be >= 1")
 		}
 
+		var agent bench.AgentRunner = bench.ClaudeRunner{ExtraArgs: modelArgs(benchModel)}
+		if benchDryRun {
+			agent = dryRunAgent{}
+		}
+		runner := bench.Runner{Agent: agent, Verifier: bench.CommandVerifier{}}
+
+		// --gen --compare regenerates a fresh fixture per arm (baseline mutates
+		// the tree, so the JIT arm must start clean) and installs the skill for
+		// the JIT arm. This is the only path that yields a real break-even.
+		if benchGen != "" && benchCompare {
+			return runGenCompare(runner, benchGen, benchN, benchWorkdir)
+		}
+
 		var tasks []bench.Task
 		var err error
 		if benchGen != "" {
@@ -63,12 +76,6 @@ var benchCmd = &cobra.Command{
 			return nil
 		}
 
-		var agent bench.AgentRunner = bench.ClaudeRunner{ExtraArgs: modelArgs(benchModel)}
-		if benchDryRun {
-			agent = dryRunAgent{}
-		}
-		runner := bench.Runner{Agent: agent, Verifier: bench.CommandVerifier{}}
-
 		if benchCompare {
 			return runCompare(runner, tasks)
 		}
@@ -154,6 +161,77 @@ func printComparison(c bench.Comparison, compileCost int) {
 	fmt.Println()
 }
 
+// runGenCompare sweeps the repeat counts, regenerating a FRESH fixture for each
+// arm (baseline mutates the tree in place, so the JIT arm must not inherit it)
+// and installing the workflow skill for the JIT arm. Reports break-even per count.
+func runGenCompare(runner bench.Runner, shape string, counts []int, workdir string) error {
+	fixture, ok := bench.FixtureByShape(shape)
+	if !ok {
+		return fmt.Errorf("unknown --gen shape %q (have: %v)", shape, bench.FixtureShapes())
+	}
+	skilled, canSkill := fixture.(bench.SkilledFixture)
+	if !canSkill {
+		return fmt.Errorf("fixture %q has no skill to install; --gen --compare needs one", shape)
+	}
+	if len(counts) == 0 {
+		counts = []int{3}
+	}
+	base := workdir
+	if base == "" {
+		dir, err := os.MkdirTemp("", "aj-bench-")
+		if err != nil {
+			return err
+		}
+		base = dir
+		fmt.Printf("[AJ] Generated fixtures under %s\n", base)
+	}
+
+	// The JIT arm installs the skill before its episode.
+	runner.Setup = func(task bench.Task, arm bench.Arm) error {
+		if arm == bench.ArmJIT {
+			return skilled.InstallSkill(task)
+		}
+		return nil
+	}
+
+	comparisons := make([]bench.Comparison, 0, len(counts))
+	for _, n := range counts {
+		// Fresh, separate trees per arm so neither sees the other's edits.
+		baseTask, err := genInto(fixture, base, fmt.Sprintf("baseline-%s-%d", shape, n), n)
+		if err != nil {
+			return err
+		}
+		jitTask, err := genInto(fixture, base, fmt.Sprintf("jit-%s-%d", shape, n), n)
+		if err != nil {
+			return err
+		}
+		baseline := runner.RunTask(context.Background(), baseTask, bench.ArmBaseline, benchRollouts)
+		jit := runner.RunTask(context.Background(), jitTask, bench.ArmJIT, benchRollouts)
+		// Compare wants matching task IDs; both were generated at the same count.
+		jit.Task = baseline.Task
+		c := bench.Compare(baseline, jit)
+		comparisons = append(comparisons, c)
+		if !benchJSON {
+			printComparison(c, benchCompileCost)
+		}
+	}
+	if benchJSON {
+		enc := json.NewEncoder(os.Stdout)
+		enc.SetIndent("", "  ")
+		return enc.Encode(comparisons)
+	}
+	return nil
+}
+
+// genInto generates a fixture into base/sub and returns its Task.
+func genInto(fixture bench.Fixture, base, sub string, n int) (bench.Task, error) {
+	dir := filepath.Join(base, sub)
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return bench.Task{}, err
+	}
+	return fixture.Generate(dir, n)
+}
+
 // generateTasks materializes a repetition fixture at each requested repeat count
 // under workdir (a temp dir when empty), returning one Task per count so a
 // benchmark can sweep the break-even curve.

internal/bench/fixture.go

@@ -40,6 +40,17 @@ type Fixture interface {
 	Generate(dir string, n int) (Task, error)
 }
 
+// SkilledFixture is a Fixture that can install the compiled skill for its
+// workflow, so the JIT arm runs with the skill available (as if AgentJIT had
+// compiled it from prior runs). The skill is installed into the task's own
+// working tree (.claude/skills/), so `claude` picks it up in that cwd without
+// touching any global config — keeping arms isolated.
+type SkilledFixture interface {
+	Fixture
+	// InstallSkill writes the workflow's skill under task.RepoDir/.claude/skills.
+	InstallSkill(task Task) error
+}
+
 // NullCheckFixture generates N Go files, each with a function that dereferences
 // a pointer parameter without a nil guard. The workflow under test is "add a
 // nil-check to every such function" — a same-shape task that repeats N times,
@@ -76,10 +87,10 @@ func Length%d(s *string) int {
 	}
 
 	// Dual-gate verifier: build must pass, and exactly n guard lines must exist.
-	// `grep -rc` over the package, summed, must equal n. We check via a small
-	// shell pipeline so the fixture stays declarative.
+	// Scope the grep to Go sources (--include='*.go') so an installed skill's
+	// SKILL.md example (which contains the guard text) can't inflate the count.
 	guardCheck := fmt.Sprintf(
-		`test "$(grep -rho 'if s == nil' . | wc -l | tr -d '[:space:]')" = "%d"`, n)
+		`test "$(grep -rho --include='*.go' 'if s == nil' . | wc -l | tr -d '[:space:]')" = "%d"`, n)
 	verifyScript := fmt.Sprintf("go build ./... && %s", guardCheck)
 
 	return Task{
@@ -93,3 +104,30 @@ func Length%d(s *string) int {
 		Verify: Verification{Command: []string{"sh", "-c", verifyScript}},
 	}, nil
 }
+
+// InstallSkill writes a SKILL.md into task.RepoDir/.claude/skills/nullcheck-guard
+// that gives the exact, deterministic transformation for this workflow, so the
+// JIT arm can apply it mechanically instead of reasoning it out each time.
+func (NullCheckFixture) InstallSkill(task Task) error {
+	skillDir := filepath.Join(task.RepoDir, ".claude", "skills", "nullcheck-guard")
+	if err := os.MkdirAll(skillDir, 0o755); err != nil {
+		return err
+	}
+	skill := "---\n" +
+		"name: nullcheck-guard\n" +
+		"description: Add a nil guard to Go functions that dereference a *string parameter s.\n" +
+		"---\n\n" +
+		"# nullcheck-guard\n\n" +
+		"For every function of the form `func LengthN(s *string) int { return len(*s) }`,\n" +
+		"insert a guard as the first statement so a nil pointer is handled:\n\n" +
+		"```go\n" +
+		"func LengthN(s *string) int {\n" +
+		"\tif s == nil {\n" +
+		"\t\treturn 0\n" +
+		"\t}\n" +
+		"\treturn len(*s)\n" +
+		"}\n" +
+		"```\n\n" +
+		"Apply this edit to every matching function in the package, then ensure it still builds.\n"
+	return os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(skill), 0o644)
+}

internal/bench/fixture_test.go

@@ -101,3 +101,37 @@ func TestFixtureRegistry(t *testing.T) {
 		t.Error("FixtureShapes() is empty")
 	}
 }
+
+// Installing the skill must NOT fool the verifier: the SKILL.md contains the
+// guard text as an example, but the grep is scoped to *.go, so on an unmodified
+// fixture (0 guards in code) the verifier must still fail.
+func TestInstallSkillDoesNotContaminateVerifier(t *testing.T) {
+	if _, err := exec.LookPath("go"); err != nil {
+		t.Skip("go toolchain not available")
+	}
+	dir := t.TempDir()
+	task, err := NullCheckFixture{}.Generate(dir, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := (NullCheckFixture{}).InstallSkill(task); err != nil {
+		t.Fatalf("InstallSkill: %v", err)
+	}
+	// Skill file exists (and contains the guard text as an example).
+	skillMD := filepath.Join(task.RepoDir, ".claude", "skills", "nullcheck-guard", "SKILL.md")
+	b, err := os.ReadFile(skillMD)
+	if err != nil {
+		t.Fatalf("skill not installed: %v", err)
+	}
+	if !strings.Contains(string(b), "if s == nil") {
+		t.Fatal("test premise broken: SKILL.md should contain the guard example")
+	}
+	// Verifier must still FAIL — code has 0 guards despite the skill example.
+	ok, err := CommandVerifier{}.Verify(context.Background(), task)
+	if err != nil {
+		t.Fatalf("Verify: %v", err)
+	}
+	if ok {
+		t.Error("verifier passed due to SKILL.md contamination; grep must exclude non-.go files")
+	}
+}

internal/bench/runner.go

@@ -63,13 +63,27 @@ func (tr TaskResult) MeanTokensToSuccess() (float64, bool) {
 type Runner struct {
 	Agent    AgentRunner
 	Verifier Verifier
+	// Setup, if set, runs once before a task's rollouts begin — used to prepare
+	// arm-specific state (e.g. install the JIT arm's skill). A Setup error fails
+	// every rollout of that task (recorded, not counted).
+	Setup func(task Task, arm Arm) error
 }
 
 // RunTask runs a task under one arm for n rollouts, gating each on the verifier.
 // A rollout whose episode or verifier errors is recorded with Err set and
 // Verified=false; it never counts toward tokens-to-success.
 func (rn Runner) RunTask(ctx context.Context, task Task, arm Arm, n int) TaskResult {
 	res := TaskResult{Task: task.ID, Arm: arm, Shape: task.Shape}
+	if rn.Setup != nil {
+		if err := rn.Setup(task, arm); err != nil {
+			for i := 0; i < n; i++ {
+				res.Rollouts = append(res.Rollouts, RolloutResult{
+					Task: task.ID, Arm: arm, Index: i, Err: "setup: " + err.Error(),
+				})
+			}
+			return res
+		}
+	}
 	for i := 0; i < n; i++ {
 		rollout := RolloutResult{Task: task.ID, Arm: arm, Index: i}