| 1 | package cli |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "os" |
| 6 | "path/filepath" |
| 7 | "sync" |
| 8 | "testing" |
| 9 | "time" |
| 10 | |
| 11 | "reasonix/internal/event" |
| 12 | "reasonix/internal/provider" |
| 13 | ) |
| 14 | |
| 15 | func usageEvent(source string, prompt, completion int) event.Event { |
| 16 | return event.Event{ |
| 17 | Kind: event.Usage, |
| 18 | UsageSource: source, |
| 19 | Usage: &provider.Usage{PromptTokens: prompt, CompletionTokens: completion, CacheMissTokens: prompt}, |
| 20 | Pricing: &provider.Pricing{Input: 1, Output: 2, CacheHit: 0.1, Currency: "$"}, |
| 21 | } |
| 22 | } |
| 23 | |
| 24 | func usageEventWithCacheReason(reason string) event.Event { |
| 25 | e := usageEvent(event.UsageSourceSubagent, 10, 1) |
| 26 | e.CacheDiagnostics = &event.CacheDiagnostics{PrefixChangeReasons: []string{reason}} |
| 27 | return e |
| 28 | } |
| 29 | |
| 30 | func TestSnapshotDeepCopiesPrefixChangeReasons(t *testing.T) { |
| 31 | s := &metricsSink{inner: event.Discard} |
| 32 | s.Emit(usageEventWithCacheReason("compact_auto")) |
| 33 | |
| 34 | snapshot := s.Snapshot() |
| 35 | s.Emit(usageEventWithCacheReason("snip")) |
| 36 | |
| 37 | if snapshot.PrefixChangeReasonCounts["compact_auto"] != 1 { |
| 38 | t.Fatalf("snapshot compact_auto = %d, want 1", snapshot.PrefixChangeReasonCounts["compact_auto"]) |
| 39 | } |
| 40 | if _, changed := snapshot.PrefixChangeReasonCounts["snip"]; changed { |
| 41 | t.Fatalf("snapshot changed after return: %v", snapshot.PrefixChangeReasonCounts) |
| 42 | } |
| 43 | } |
| 44 | |
| 45 | // A killed agent writes no final record. Everything it did before the kill is |
| 46 | // only recoverable if snapshots landed on disk while it ran. |
| 47 | func TestSnapshotSurvivesWithoutAFinalWrite(t *testing.T) { |
| 48 | dir := t.TempDir() |
| 49 | final := filepath.Join(dir, "metrics.json") |
| 50 | now := time.Unix(0, 0) |
| 51 | s := &metricsSink{ |
| 52 | inner: event.Discard, |
| 53 | partialPath: partialMetricsPath(final), |
| 54 | snapshotEvery: time.Second, |
| 55 | clock: func() time.Time { return now }, |
| 56 | } |
| 57 | |
| 58 | s.Emit(usageEvent(event.UsageSourceExecutor, 1000, 10)) |
| 59 | now = now.Add(2 * time.Second) |
| 60 | s.Emit(usageEvent(event.UsageSourceExecutor, 500, 5)) |
| 61 | |
| 62 | raw, err := os.ReadFile(partialMetricsPath(final)) |
| 63 | if err != nil { |
| 64 | t.Fatalf("no snapshot on disk: %v", err) |
| 65 | } |
| 66 | var got RunMetrics |
| 67 | if err := json.Unmarshal(raw, &got); err != nil { |
| 68 | t.Fatalf("snapshot is not parseable JSON: %v", err) |
| 69 | } |
| 70 | if got.Complete { |
| 71 | t.Error("a snapshot must never claim to be complete") |
| 72 | } |
| 73 | if got.PromptTokens == 0 || got.Steps == 0 { |
| 74 | t.Errorf("snapshot lost the accounting it exists to preserve: %+v", got) |
| 75 | } |
| 76 | if _, err := os.Stat(final); !os.IsNotExist(err) { |
| 77 | t.Error("no final record should exist for a run that never finished") |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | // Snapshots are throttled: a run makes thousands of events and must not make |
| 82 | // thousands of disk writes. |
| 83 | func TestSnapshotsAreThrottled(t *testing.T) { |
| 84 | dir := t.TempDir() |
| 85 | now := time.Unix(0, 0) |
| 86 | s := &metricsSink{ |
| 87 | inner: event.Discard, |
| 88 | partialPath: filepath.Join(dir, "m.json.partial"), |
| 89 | snapshotEvery: time.Minute, |
| 90 | clock: func() time.Time { return now }, |
| 91 | } |
| 92 | for i := 0; i < 50; i++ { |
| 93 | s.Emit(usageEvent(event.UsageSourceExecutor, 10, 1)) |
| 94 | } |
| 95 | raw, err := os.ReadFile(s.partialPath) |
| 96 | if err != nil { |
| 97 | t.Fatalf("first snapshot should still be written: %v", err) |
| 98 | } |
| 99 | var got RunMetrics |
| 100 | if err := json.Unmarshal(raw, &got); err != nil { |
| 101 | t.Fatalf("unmarshal: %v", err) |
| 102 | } |
| 103 | if got.Steps != 1 { |
| 104 | t.Fatalf("snapshot steps = %d, want 1 — later events must not have rewritten within the window", got.Steps) |
| 105 | } |
| 106 | } |
| 107 | |
| 108 | // A completed run must leave exactly one readable record, or a reader could |
| 109 | // count the run twice. |
| 110 | func TestFinalWriteRetiresTheSnapshot(t *testing.T) { |
| 111 | dir := t.TempDir() |
| 112 | final := filepath.Join(dir, "metrics.json") |
| 113 | now := time.Unix(0, 0) |
| 114 | s := &metricsSink{ |
| 115 | inner: event.Discard, |
| 116 | partialPath: partialMetricsPath(final), |
| 117 | snapshotEvery: time.Millisecond, |
| 118 | clock: func() time.Time { now = now.Add(time.Second); return now }, |
| 119 | } |
| 120 | s.Emit(usageEvent(event.UsageSourceExecutor, 100, 10)) |
| 121 | if _, err := os.Stat(partialMetricsPath(final)); err != nil { |
| 122 | t.Fatalf("expected a snapshot before the final write: %v", err) |
| 123 | } |
| 124 | |
| 125 | if err := writeMetrics(final, s.Snapshot()); err != nil { |
| 126 | t.Fatalf("writeMetrics: %v", err) |
| 127 | } |
| 128 | |
| 129 | raw, err := os.ReadFile(final) |
| 130 | if err != nil { |
| 131 | t.Fatalf("final record missing: %v", err) |
| 132 | } |
| 133 | var got RunMetrics |
| 134 | if err := json.Unmarshal(raw, &got); err != nil { |
| 135 | t.Fatalf("unmarshal: %v", err) |
| 136 | } |
| 137 | if !got.Complete { |
| 138 | t.Error("the final record must be marked complete") |
| 139 | } |
| 140 | if _, err := os.Stat(partialMetricsPath(final)); !os.IsNotExist(err) { |
| 141 | t.Error("the snapshot must be retired so it cannot be double-counted") |
| 142 | } |
| 143 | } |
| 144 | |
| 145 | // Steps counts every billed call; the breakdown is what makes a total above |
| 146 | // max_steps explicable. An unrecognised origin must survive rather than vanish |
| 147 | // from a total that is meant to reconcile. |
| 148 | func TestUsageBySourceReconcilesWithTheTotal(t *testing.T) { |
| 149 | s := &metricsSink{inner: event.Discard} |
| 150 | s.Emit(usageEvent(event.UsageSourceExecutor, 100, 10)) |
| 151 | s.Emit(usageEvent(event.UsageSourceSubagent, 200, 20)) |
| 152 | s.Emit(usageEvent(event.UsageSourceCompaction, 300, 30)) |
| 153 | s.Emit(usageEvent("some-future-origin", 400, 40)) |
| 154 | s.Emit(usageEvent("", 500, 50)) // empty means executor, per the Usage contract |
| 155 | |
| 156 | m := s.Snapshot() |
| 157 | if len(m.UsageBySource) != 4 { |
| 158 | t.Fatalf("sources = %v, want executor/subagent/compaction/some-future-origin", m.UsageBySource) |
| 159 | } |
| 160 | if got := m.UsageBySource[event.UsageSourceExecutor].Calls; got != 2 { |
| 161 | t.Errorf("executor calls = %d, want 2 (an empty source is the executor)", got) |
| 162 | } |
| 163 | if _, ok := m.UsageBySource["some-future-origin"]; !ok { |
| 164 | t.Error("an unknown origin must be kept, not dropped") |
| 165 | } |
| 166 | |
| 167 | var calls, prompt int |
| 168 | for _, u := range m.UsageBySource { |
| 169 | calls += u.Calls |
| 170 | prompt += u.PromptTokens |
| 171 | } |
| 172 | if calls != m.Steps { |
| 173 | t.Errorf("source calls sum to %d but Steps is %d — the breakdown must reconcile", calls, m.Steps) |
| 174 | } |
| 175 | if prompt != m.PromptTokens { |
| 176 | t.Errorf("source prompt tokens sum to %d but total is %d", prompt, m.PromptTokens) |
| 177 | } |
| 178 | } |
| 179 | |
| 180 | // Background jobs emit while the run command assembles the final record. |
| 181 | // Run with -race. |
| 182 | func TestConcurrentEmitAndSnapshotAreRaceFree(t *testing.T) { |
| 183 | dir := t.TempDir() |
| 184 | s := &metricsSink{ |
| 185 | inner: event.Discard, |
| 186 | partialPath: filepath.Join(dir, "m.json.partial"), |
| 187 | snapshotEvery: time.Millisecond, |
| 188 | } |
| 189 | const emitters, each = 8, 50 |
| 190 | |
| 191 | var wg sync.WaitGroup |
| 192 | for i := 0; i < emitters; i++ { |
| 193 | wg.Add(1) |
| 194 | go func() { |
| 195 | defer wg.Done() |
| 196 | for j := 0; j < each; j++ { |
| 197 | s.Emit(usageEventWithCacheReason("compact_auto")) |
| 198 | s.Emit(event.Event{Kind: event.ToolResult, Tool: event.Tool{Name: "bash"}}) |
| 199 | } |
| 200 | }() |
| 201 | } |
| 202 | wg.Add(1) |
| 203 | go func() { |
| 204 | defer wg.Done() |
| 205 | for i := 0; i < 200; i++ { |
| 206 | if _, err := json.Marshal(s.Snapshot()); err != nil { |
| 207 | t.Errorf("marshal snapshot: %v", err) |
| 208 | return |
| 209 | } |
| 210 | } |
| 211 | }() |
| 212 | wg.Wait() |
| 213 | |
| 214 | m := s.Snapshot() |
| 215 | if m.Steps != emitters*each { |
| 216 | t.Errorf("steps = %d, want %d — concurrent emission lost counts", m.Steps, emitters*each) |
| 217 | } |
| 218 | if m.ToolCalls != emitters*each { |
| 219 | t.Errorf("tool calls = %d, want %d", m.ToolCalls, emitters*each) |
| 220 | } |
| 221 | } |
| 222 |