| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "fmt" |
| 6 | "os" |
| 7 | "os/exec" |
| 8 | "path/filepath" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| 12 | // runSegments drives a task's legs and folds their accounting into one result. |
| 13 | // Each leg writes its own metrics file: sharing one path would let the last |
| 14 | // leg's numbers stand in for the whole run, and the tokens the earlier legs |
| 15 | // spent would simply vanish. Only the final leg's trajectory digest is kept; |
| 16 | // Segments records how many legs it does not cover. |
| 17 | func runSegments(ctx context.Context, cfg suiteConfig, t task, work, trajDir string, env []string, r *result) error { |
| 18 | segs := planSegments(t, cfg.segments, cfg.steers) |
| 19 | r.Segments = len(segs) |
| 20 | var runErr error |
| 21 | for _, seg := range segs { |
| 22 | metricsPath := filepath.Join(work, fmt.Sprintf(".run-metrics-%d.json", seg.index)) |
| 23 | trajPath := segmentTrajectoryPath(trajDir, t.ID, seg, len(segs)) |
| 24 | args := buildSegmentArgs(cfg, seg, metricsPath, trajPath) |
| 25 | |
| 26 | cmd := exec.CommandContext(ctx, cfg.bin, args...) |
| 27 | cmd.Dir = work |
| 28 | if len(env) > 0 { |
| 29 | cmd.Env = append(os.Environ(), env...) |
| 30 | } |
| 31 | cmd.Stdout = os.Stderr |
| 32 | cmd.Stderr = os.Stderr |
| 33 | cmd.WaitDelay = 10 * time.Second |
| 34 | runErr = cmd.Run() |
| 35 | |
| 36 | if m, err := readMetrics(metricsPath); err == nil { |
| 37 | foldSegmentMetrics(r, m, seg.index) |
| 38 | } |
| 39 | // A leg that died takes the run with it: resuming a session the child |
| 40 | // never finished writing would measure the harness's crash recovery, |
| 41 | // which is a different experiment. |
| 42 | if runErr != nil || ctx.Err() != nil { |
| 43 | break |
| 44 | } |
| 45 | } |
| 46 | return runErr |
| 47 | } |
| 48 | |
| 49 | // segmentTrajectoryPath keeps one file per leg so a resumed leg cannot truncate |
| 50 | // the record of the one before it. |
| 51 | func segmentTrajectoryPath(dir, id string, seg segment, total int) string { |
| 52 | if dir == "" { |
| 53 | return "" |
| 54 | } |
| 55 | if total < 2 { |
| 56 | return filepath.Join(dir, id+".trajectory.jsonl") |
| 57 | } |
| 58 | return filepath.Join(dir, fmt.Sprintf("%s.seg%d.trajectory.jsonl", id, seg.index)) |
| 59 | } |
| 60 | |
| 61 | // lastSegmentTrajectory names the file whose digest represents the run. Only |
| 62 | // the final leg's is read; Segments in the JSON says how many were not. |
| 63 | func lastSegmentTrajectory(dir, id string, segments int) string { |
| 64 | if dir == "" { |
| 65 | return "" |
| 66 | } |
| 67 | if segments < 2 { |
| 68 | return filepath.Join(dir, id+".trajectory.jsonl") |
| 69 | } |
| 70 | return filepath.Join(dir, fmt.Sprintf("%s.seg%d.trajectory.jsonl", id, segments)) |
| 71 | } |
| 72 | |
| 73 | // foldSegmentMetrics adds one leg's spend to the run. The first leg's metrics |
| 74 | // establish the non-additive fields (currency, outcome); later legs contribute |
| 75 | // their totals, and the last leg's outcome wins because it is the one that |
| 76 | // ended the run. |
| 77 | func foldSegmentMetrics(r *result, m runMetrics, index int) { |
| 78 | if index == 1 { |
| 79 | r.runMetrics = m |
| 80 | return |
| 81 | } |
| 82 | r.PromptTokens += m.PromptTokens |
| 83 | r.CompletionTokens += m.CompletionTokens |
| 84 | r.CacheHitTokens += m.CacheHitTokens |
| 85 | r.CacheMissTokens += m.CacheMissTokens |
| 86 | r.Cost += m.Cost |
| 87 | r.Steps += m.Steps |
| 88 | r.Compactions += m.Compactions |
| 89 | r.ToolCalls += m.ToolCalls |
| 90 | r.ToolFailures += m.ToolFailures |
| 91 | r.SubagentToolCalls += m.SubagentToolCalls |
| 92 | r.Retries += m.Retries |
| 93 | r.Complete = m.Complete |
| 94 | if m.Outcome != "" { |
| 95 | r.Outcome = m.Outcome |
| 96 | } |
| 97 | for name, n := range m.ToolCallsByName { |
| 98 | if r.ToolCallsByName == nil { |
| 99 | r.ToolCallsByName = map[string]int{} |
| 100 | } |
| 101 | r.ToolCallsByName[name] += n |
| 102 | } |
| 103 | for reason, n := range m.PrefixChangeReasonCounts { |
| 104 | if r.PrefixChangeReasonCounts == nil { |
| 105 | r.PrefixChangeReasonCounts = map[string]int{} |
| 106 | } |
| 107 | r.PrefixChangeReasonCounts[reason] += n |
| 108 | } |
| 109 | accumulateSources(r.UsageBySource, m.UsageBySource) |
| 110 | } |
| 111 |