| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "os" |
| 5 | "path/filepath" |
| 6 | "runtime" |
| 7 | "strconv" |
| 8 | "strings" |
| 9 | "testing" |
| 10 | ) |
| 11 | |
| 12 | // requireShellStub skips where a #!/usr/bin/env bash stub cannot be executed. |
| 13 | // Every test that writes its own stub must call it: exec fails silently enough |
| 14 | // on Windows that the assertion which follows blames the code instead. |
| 15 | func requireShellStub(t *testing.T) { |
| 16 | t.Helper() |
| 17 | if runtime.GOOS == "windows" { |
| 18 | t.Skip("the stub agent is a shell script") |
| 19 | } |
| 20 | } |
| 21 | |
| 22 | // fakeAgent stands in for the reasonix binary: it records every invocation's |
| 23 | // argv and writes the metrics file it was told to, making segment accounting |
| 24 | // verifiable without a provider. The failure it guards — a later leg replacing |
| 25 | // an earlier leg's numbers — is invisible in any single run. |
| 26 | func fakeAgent(t *testing.T, promptTokens, completionTokens int) (bin, argvLog string) { |
| 27 | t.Helper() |
| 28 | requireShellStub(t) |
| 29 | dir := t.TempDir() |
| 30 | argvLog = filepath.Join(dir, "argv.log") |
| 31 | bin = filepath.Join(dir, "fake-agent") |
| 32 | script := `#!/usr/bin/env bash |
| 33 | printf '%s\n' "$*" >> ` + argvLog + ` |
| 34 | metrics="" |
| 35 | while [ $# -gt 0 ]; do |
| 36 | case "$1" in |
| 37 | --metrics) metrics="$2"; shift 2;; |
| 38 | *) shift;; |
| 39 | esac |
| 40 | done |
| 41 | [ -n "$metrics" ] && cat > "$metrics" <<JSON |
| 42 | {"prompt_tokens":` + strconv.Itoa(promptTokens) + `,"completion_tokens":` + strconv.Itoa(completionTokens) + `,"steps":2,"cost":0.5,"currency":"USD","complete":true,"tool_calls":3} |
| 43 | JSON |
| 44 | exit 0 |
| 45 | ` |
| 46 | if err := os.WriteFile(bin, []byte(script), 0o755); err != nil { |
| 47 | t.Fatal(err) |
| 48 | } |
| 49 | return bin, argvLog |
| 50 | } |
| 51 | |
| 52 | func TestSegmentedRunAddsEveryLegsSpend(t *testing.T) { |
| 53 | bin, argvLog := fakeAgent(t, 100, 20) |
| 54 | cfg := suiteConfig{bin: bin, segments: 3} |
| 55 | task := task{ID: "seg", Prompt: "fix it", MaxSteps: 9, TimeoutSec: 60} |
| 56 | var r result |
| 57 | if err := runSegments(t.Context(), cfg, task, t.TempDir(), "", nil, &r); err != nil { |
| 58 | t.Fatalf("runSegments: %v", err) |
| 59 | } |
| 60 | |
| 61 | if r.Segments != 3 { |
| 62 | t.Fatalf("segments = %d, want 3", r.Segments) |
| 63 | } |
| 64 | if r.PromptTokens != 300 || r.CompletionTokens != 60 { |
| 65 | t.Fatalf("tokens = %d/%d, want 300/60 — a later leg replaced an earlier one instead of adding", |
| 66 | r.PromptTokens, r.CompletionTokens) |
| 67 | } |
| 68 | if r.Steps != 6 || r.ToolCalls != 9 { |
| 69 | t.Fatalf("steps=%d tools=%d, want 6/9 summed across legs", r.Steps, r.ToolCalls) |
| 70 | } |
| 71 | if r.Cost != 1.5 { |
| 72 | t.Fatalf("cost = %v, want 1.5", r.Cost) |
| 73 | } |
| 74 | if r.Currency != "USD" || !r.Complete { |
| 75 | t.Fatalf("non-additive fields lost: %+v", r.runMetrics) |
| 76 | } |
| 77 | |
| 78 | argv, err := os.ReadFile(argvLog) |
| 79 | if err != nil { |
| 80 | t.Fatal(err) |
| 81 | } |
| 82 | lines := strings.Split(strings.TrimSpace(string(argv)), "\n") |
| 83 | if len(lines) != 3 { |
| 84 | t.Fatalf("invocations = %d, want 3", len(lines)) |
| 85 | } |
| 86 | if strings.Contains(lines[0], "--continue") { |
| 87 | t.Fatalf("leg 1 must start a fresh session: %s", lines[0]) |
| 88 | } |
| 89 | for i, line := range lines[1:] { |
| 90 | if !strings.Contains(line, "--continue") { |
| 91 | t.Fatalf("leg %d must resume: %s", i+2, line) |
| 92 | } |
| 93 | if strings.Contains(line, "fix it") { |
| 94 | t.Fatalf("leg %d restated the task: %s", i+2, line) |
| 95 | } |
| 96 | } |
| 97 | } |
| 98 | |
| 99 | func TestSegmentedRunKeepsPerLegMetricsApart(t *testing.T) { |
| 100 | bin, _ := fakeAgent(t, 10, 1) |
| 101 | work := t.TempDir() |
| 102 | var r result |
| 103 | if err := runSegments(t.Context(), suiteConfig{bin: bin, segments: 2}, task{ID: "x", Prompt: "p", MaxSteps: 4, TimeoutSec: 60}, work, "", nil, &r); err != nil { |
| 104 | t.Fatalf("runSegments: %v", err) |
| 105 | } |
| 106 | for _, name := range []string{".run-metrics-1.json", ".run-metrics-2.json"} { |
| 107 | if _, err := os.Stat(filepath.Join(work, name)); err != nil { |
| 108 | t.Fatalf("%s missing: one shared path would have let the last leg overwrite the rest", name) |
| 109 | } |
| 110 | } |
| 111 | } |
| 112 | |
| 113 | func TestSegmentedRunStopsAtTheFirstFailedLeg(t *testing.T) { |
| 114 | requireShellStub(t) |
| 115 | dir := t.TempDir() |
| 116 | bin := filepath.Join(dir, "failing-agent") |
| 117 | log := filepath.Join(dir, "calls.log") |
| 118 | script := "#!/usr/bin/env bash\necho x >> " + log + "\nexit 3\n" |
| 119 | if err := os.WriteFile(bin, []byte(script), 0o755); err != nil { |
| 120 | t.Fatal(err) |
| 121 | } |
| 122 | var r result |
| 123 | err := runSegments(t.Context(), suiteConfig{bin: bin, segments: 3}, task{ID: "x", Prompt: "p", MaxSteps: 6, TimeoutSec: 60}, t.TempDir(), "", nil, &r) |
| 124 | if err == nil { |
| 125 | t.Fatal("a failed leg must surface as the run's error") |
| 126 | } |
| 127 | calls, _ := os.ReadFile(log) |
| 128 | if got := len(strings.Fields(string(calls))); got != 1 { |
| 129 | t.Fatalf("invocations = %d, want 1 — resuming a session the child never finished is a different experiment", got) |
| 130 | } |
| 131 | } |
| 132 | |
| 133 | func TestUnsegmentedRunKeepsTheOriginalTrajectoryPath(t *testing.T) { |
| 134 | if got := lastSegmentTrajectory("t", "task-a", 1); got != filepath.Join("t", "task-a.trajectory.jsonl") { |
| 135 | t.Fatalf("path = %q, want the unsegmented name", got) |
| 136 | } |
| 137 | if got := lastSegmentTrajectory("t", "task-a", 3); got != filepath.Join("t", "task-a.seg3.trajectory.jsonl") { |
| 138 | t.Fatalf("path = %q, want the last leg's file", got) |
| 139 | } |
| 140 | if got := lastSegmentTrajectory("", "task-a", 3); got != "" { |
| 141 | t.Fatalf("path = %q, want empty when trajectories are off", got) |
| 142 | } |
| 143 | } |
| 144 |