diff --git a/cmd/aj/bench_cmd.go b/cmd/aj/bench_cmd.go index 1c70027..68a2002 100644 --- a/cmd/aj/bench_cmd.go +++ b/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. diff --git a/internal/bench/fixture.go b/internal/bench/fixture.go index 5f96b65..65b303c 100644 --- a/internal/bench/fixture.go +++ b/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) +} diff --git a/internal/bench/fixture_test.go b/internal/bench/fixture_test.go index b6b8ecb..d8f3fde 100644 --- a/internal/bench/fixture_test.go +++ b/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") + } +} diff --git a/internal/bench/runner.go b/internal/bench/runner.go index b2425ca..8a07e73 100644 --- a/internal/bench/runner.go +++ b/internal/bench/runner.go @@ -63,6 +63,10 @@ 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. @@ -70,6 +74,16 @@ type Runner struct { // 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}