| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "fmt" |
| 6 | "sync/atomic" |
| 7 | "testing" |
| 8 | |
| 9 | "reasonix/internal/billing" |
| 10 | "reasonix/internal/event" |
| 11 | "reasonix/internal/provider" |
| 12 | "reasonix/internal/tool" |
| 13 | ) |
| 14 | |
| 15 | func TestRunBudgetUsesTheCanonicalOccurrenceTimeQuote(t *testing.T) { |
| 16 | usage := &provider.Usage{CompletionTokens: 1_000_000, TotalTokens: 1_000_000, RequestCount: 1} |
| 17 | quote := func(amount, band string) *billing.CostQuote { |
| 18 | return &billing.CostQuote{Original: billing.Money{Amount: amount, Currency: "CNY"}, CostComplete: true, RateBand: band} |
| 19 | } |
| 20 | var peak, off runBudget |
| 21 | peak.observeQuote(usage, quote("27", billing.RateBandPeak)) |
| 22 | off.observeQuote(usage, quote("13.5", billing.RateBandOffPeak)) |
| 23 | if peak.cost != 27 || off.cost != 13.5 || peak.cost != 2*off.cost { |
| 24 | t.Fatalf("peak=%v off_peak=%v", peak.cost, off.cost) |
| 25 | } |
| 26 | } |
| 27 | |
| 28 | // budgetSink opts into the shadow axis; an ordinary sink would receive nothing. |
| 29 | type budgetSink struct { |
| 30 | event.FuncSink |
| 31 | samples []event.RunBudgetSample |
| 32 | } |
| 33 | |
| 34 | func newBudgetSink() *budgetSink { |
| 35 | s := &budgetSink{} |
| 36 | s.FuncSink = event.FuncSink(func(event.Event) {}) |
| 37 | return s |
| 38 | } |
| 39 | |
| 40 | func (s *budgetSink) RecordRunBudget(sample event.RunBudgetSample) { |
| 41 | s.samples = append(s.samples, sample) |
| 42 | } |
| 43 | |
| 44 | // spendingProvider bills a fixed usage per round and reads one file, so a turn |
| 45 | // costs a predictable amount without depending on a real backend. |
| 46 | type spendingProvider struct { |
| 47 | rounds atomic.Int32 |
| 48 | max int32 |
| 49 | } |
| 50 | |
| 51 | func (p *spendingProvider) Name() string { return "spending" } |
| 52 | |
| 53 | func (p *spendingProvider) Stream(context.Context, provider.Request) (<-chan provider.Chunk, error) { |
| 54 | round := p.rounds.Add(1) |
| 55 | ch := make(chan provider.Chunk, 4) |
| 56 | usage := &provider.Usage{ |
| 57 | PromptTokens: 1000, CompletionTokens: 100, TotalTokens: 1100, |
| 58 | CacheHitTokens: 900, CacheMissTokens: 100, RequestCount: 1, |
| 59 | } |
| 60 | if round > p.max { |
| 61 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "Done."} |
| 62 | ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: usage} |
| 63 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 64 | close(ch) |
| 65 | return ch, nil |
| 66 | } |
| 67 | ch <- provider.Chunk{Type: provider.ChunkToolCall, ToolCall: &provider.ToolCall{ |
| 68 | ID: fmt.Sprintf("call-%d", round), |
| 69 | Name: "read_file", |
| 70 | Arguments: fmt.Sprintf(`{"path":"pkg%d/file.go"}`, round), |
| 71 | }} |
| 72 | ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: usage} |
| 73 | ch <- provider.Chunk{Type: provider.ChunkDone} |
| 74 | close(ch) |
| 75 | return ch, nil |
| 76 | } |
| 77 | |
| 78 | // The axis must read what the turn actually spent, through the real Run loop: |
| 79 | // a component-level accumulator that never reaches a sink proves nothing. |
| 80 | func TestRunBudgetTracksRealTurnSpend(t *testing.T) { |
| 81 | sink := newBudgetSink() |
| 82 | reg := tool.NewRegistry() |
| 83 | reg.Add(readProbe{}) |
| 84 | pricing := &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2, Currency: "CNY"} |
| 85 | a := New(&spendingProvider{max: 3}, reg, NewSession("sys"), Options{Pricing: pricing}, sink) |
| 86 | |
| 87 | if err := a.Run(context.Background(), "read a few files"); err != nil { |
| 88 | t.Fatalf("Run: %v", err) |
| 89 | } |
| 90 | if len(sink.samples) != 4 { |
| 91 | t.Fatalf("samples = %d, want one per model round (3 tool rounds + 1 final)", len(sink.samples)) |
| 92 | } |
| 93 | |
| 94 | last := sink.samples[len(sink.samples)-1] |
| 95 | if last.Turn.Rounds != 4 || last.Turn.Requests != 4 { |
| 96 | t.Fatalf("last sample = %+v, want 4 rounds and 4 requests", last.Turn) |
| 97 | } |
| 98 | if last.Turn.PromptTokens != 4000 || last.Turn.OutputTokens != 400 { |
| 99 | t.Fatalf("tokens = prompt %d output %d, want 4000/400", last.Turn.PromptTokens, last.Turn.OutputTokens) |
| 100 | } |
| 101 | if !last.Turn.Priced || last.Currency != "¥" { |
| 102 | t.Fatalf("sample = %+v, want a priced reading in ¥", last) |
| 103 | } |
| 104 | // Cache hits are 50x cheaper than misses; a turn that bills 900 hits per |
| 105 | // round must not read as if all 1000 prompt tokens were misses. |
| 106 | wantCost := 4 * (900*0.02 + 100*1 + 100*2) / 1e6 |
| 107 | if diff := last.Turn.Cost - wantCost; diff > 1e-12 || diff < -1e-12 { |
| 108 | t.Fatalf("cost = %v, want %v (cache-hit priced)", last.Turn.Cost, wantCost) |
| 109 | } |
| 110 | if last.Turn.ElapsedMs < 0 { |
| 111 | t.Fatalf("elapsed = %d, want a wall-clock reading", last.Turn.ElapsedMs) |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | // The whole point of the task scope: "continue" starts a new Run, and a |
| 116 | // per-Run total resets there. The four-hour failure this axis exists for was |
| 117 | // never one Run. |
| 118 | func TestTaskBudgetRestartsForOrdinaryNewUserTurn(t *testing.T) { |
| 119 | sink := newBudgetSink() |
| 120 | reg := tool.NewRegistry() |
| 121 | reg.Add(readProbe{}) |
| 122 | pricing := &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2, Currency: "CNY"} |
| 123 | a := New(&spendingProvider{max: 2}, reg, NewSession("sys"), Options{Pricing: pricing}, sink) |
| 124 | |
| 125 | if err := a.Run(context.Background(), "start the work"); err != nil { |
| 126 | t.Fatalf("first Run: %v", err) |
| 127 | } |
| 128 | afterFirst := sink.samples[len(sink.samples)-1] |
| 129 | |
| 130 | if err := a.Run(context.Background(), "continue"); err != nil { |
| 131 | t.Fatalf("continuation Run: %v", err) |
| 132 | } |
| 133 | afterSecond := sink.samples[len(sink.samples)-1] |
| 134 | |
| 135 | if afterSecond.Turn.Rounds >= afterFirst.Turn.Rounds { |
| 136 | t.Fatalf("turn rounds = %d, want the per-Run scope to restart below the first Run's %d", |
| 137 | afterSecond.Turn.Rounds, afterFirst.Turn.Rounds) |
| 138 | } |
| 139 | if afterSecond.Task.Rounds != afterSecond.Turn.Rounds { |
| 140 | t.Fatalf("task rounds = %d, want current turn rounds %d", |
| 141 | afterSecond.Task.Rounds, afterSecond.Turn.Rounds) |
| 142 | } |
| 143 | if afterSecond.Task.Cost >= afterFirst.Task.Cost { |
| 144 | t.Fatalf("task cost = %v, want a fresh ordinary-turn task below the first Run's %v", |
| 145 | afterSecond.Task.Cost, afterFirst.Task.Cost) |
| 146 | } |
| 147 | if afterSecond.Task.ElapsedMs < afterSecond.Turn.ElapsedMs { |
| 148 | t.Fatal("task elapsed must cover the current ordinary turn") |
| 149 | } |
| 150 | } |
| 151 | |
| 152 | // A genuinely new task starts from zero, because a fresh evidence ledger is |
| 153 | // what "new task" means here. |
| 154 | func TestTaskBudgetResetsWithTheEvidenceLedger(t *testing.T) { |
| 155 | sink := newBudgetSink() |
| 156 | reg := tool.NewRegistry() |
| 157 | reg.Add(readProbe{}) |
| 158 | a := New(&spendingProvider{max: 1}, reg, NewSession("sys"), |
| 159 | Options{Pricing: &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2}}, sink) |
| 160 | |
| 161 | if err := a.Run(context.Background(), "first task"); err != nil { |
| 162 | t.Fatalf("first Run: %v", err) |
| 163 | } |
| 164 | first := sink.samples[len(sink.samples)-1].Task |
| 165 | if first.Rounds == 0 { |
| 166 | t.Fatal("first Run recorded nothing; the reset assertion would be vacuous") |
| 167 | } |
| 168 | |
| 169 | if err := a.Run(context.Background(), "an unrelated second task"); err != nil { |
| 170 | t.Fatalf("second Run: %v", err) |
| 171 | } |
| 172 | second := sink.samples[len(sink.samples)-1] |
| 173 | |
| 174 | if second.Task.Rounds != second.Turn.Rounds { |
| 175 | t.Fatalf("task rounds = %d, want a reset to this Run's own %d", |
| 176 | second.Task.Rounds, second.Turn.Rounds) |
| 177 | } |
| 178 | if second.Task.Cost >= first.Cost+second.Turn.Cost { |
| 179 | t.Fatalf("task cost = %v, want the first task's %v dropped", second.Task.Cost, first.Cost) |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | // Every round counts even when its usage never arrived, so the axis never |
| 184 | // reads cheaper than the turn was. |
| 185 | func TestRunBudgetCountsRoundsWithoutUsage(t *testing.T) { |
| 186 | var b runBudget |
| 187 | b.observe(nil, nil) |
| 188 | b.observe(&provider.Usage{PromptTokens: 10, CompletionTokens: 1, RequestCount: 1}, nil) |
| 189 | got := b.totals() |
| 190 | if got.Rounds != 2 || got.Requests != 1 || got.PromptTokens != 10 { |
| 191 | t.Fatalf("sample = %+v, want 2 rounds / 1 request / 10 prompt tokens", got) |
| 192 | } |
| 193 | if got.Priced { |
| 194 | t.Fatal("an unpriced turn must not report a priced reading") |
| 195 | } |
| 196 | } |
| 197 | |
| 198 | func TestRunBudgetIgnoresSinksThatDoNotOptIn(t *testing.T) { |
| 199 | plain := event.FuncSink(func(event.Event) {}) |
| 200 | a := &Agent{svc: agentServices{sink: plain}} |
| 201 | state := &turnRuntime{} |
| 202 | a.observeRunBudget(state, &provider.Usage{PromptTokens: 5, RequestCount: 1}) |
| 203 | if state.budget.rounds != 1 || state.budget.promptTokens != 5 { |
| 204 | t.Fatalf("budget = %+v, want the round still accumulated locally", state.budget) |
| 205 | } |
| 206 | } |
| 207 |