| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "strings" |
| 7 | "testing" |
| 8 | "time" |
| 9 | |
| 10 | "reasonix/internal/event" |
| 11 | "reasonix/internal/provider" |
| 12 | "reasonix/internal/tool" |
| 13 | ) |
| 14 | |
| 15 | // The gate must stop a runaway on the axis it actually runs away on. This |
| 16 | // provider never repeats itself and never fails, so every adaptive guard stays |
| 17 | // quiet — the shape that burned four hours — and only spend can catch it. |
| 18 | func TestTaskBudgetGateLandsARunawayOnCost(t *testing.T) { |
| 19 | sink := newBudgetSink() |
| 20 | reg := tool.NewRegistry() |
| 21 | reg.Add(readProbe{}) |
| 22 | pricing := &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2, Currency: "CNY"} |
| 23 | // Each round bills 900 hits + 100 misses + 100 output = 3.18e-4. |
| 24 | // A 1e-3 budget lands on the fourth round; 500 rounds are available. |
| 25 | a := New(&spendingProvider{max: 500}, reg, NewSession("sys"), |
| 26 | Options{Pricing: pricing, TaskBudget: TaskBudget{Cost: 1e-3}}, sink) |
| 27 | |
| 28 | err := a.Run(context.Background(), "read everything") |
| 29 | if err == nil { |
| 30 | t.Fatal("Run returned nil; want a resumable task-budget pause") |
| 31 | } |
| 32 | info, ok := InspectRunPause(err) |
| 33 | if !ok || info.Kind != "task_budget" || info.Key != "cost" { |
| 34 | t.Fatalf("pause = %+v (%v), want a host-owned task_budget pause on cost", info, err) |
| 35 | } |
| 36 | if !info.HostOwned { |
| 37 | t.Fatal("a host-imposed budget must report HostOwned") |
| 38 | } |
| 39 | last := sink.samples[len(sink.samples)-1] |
| 40 | if last.Task.Rounds > 20 { |
| 41 | t.Fatalf("ran %d rounds before landing; the gate should fire on spend, not drift", last.Task.Rounds) |
| 42 | } |
| 43 | if last.Task.Cost < 1e-3 { |
| 44 | t.Fatalf("task cost %v landed below its own budget", last.Task.Cost) |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | // Landing is one tool-free summary, not a truncated turn: the work stays in |
| 49 | // the session and the next message continues it. |
| 50 | func TestTaskBudgetGateKeepsTheWorkAndAsksForASummary(t *testing.T) { |
| 51 | sink := newBudgetSink() |
| 52 | reg := tool.NewRegistry() |
| 53 | reg.Add(readProbe{}) |
| 54 | a := New(&spendingProvider{max: 500}, reg, NewSession("sys"), |
| 55 | Options{ |
| 56 | Pricing: &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2}, |
| 57 | TaskBudget: TaskBudget{Cost: 1e-3}, |
| 58 | }, sink) |
| 59 | |
| 60 | _ = a.Run(context.Background(), "read everything") |
| 61 | |
| 62 | var sawNudge, sawToolResult bool |
| 63 | for _, m := range a.sess.conversation.Messages { |
| 64 | if m.Role == provider.RoleUser && strings.Contains(m.Content, "reached its cost budget") { |
| 65 | sawNudge = true |
| 66 | } |
| 67 | if m.Role == provider.RoleTool { |
| 68 | sawToolResult = true |
| 69 | } |
| 70 | } |
| 71 | if !sawNudge { |
| 72 | t.Fatal("no finalization request in the session; the model was cut off instead of asked to land") |
| 73 | } |
| 74 | if !sawToolResult { |
| 75 | t.Fatal("completed tool work was dropped from the session") |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | // An unpriced model must not read as free-and-therefore-fine, nor as instantly |
| 80 | // over budget. Cost simply does not gate it; wall clock still can. |
| 81 | func TestTaskBudgetGateIgnoresCostWhenUnpriced(t *testing.T) { |
| 82 | var b runBudget |
| 83 | b.observe(&provider.Usage{PromptTokens: 10_000_000, CompletionTokens: 1_000_000, RequestCount: 1}, nil) |
| 84 | if axis, _ := b.exceeded(TaskBudget{Cost: 1e-9}); axis != "" { |
| 85 | t.Fatalf("unpriced turn crossed the %q axis; cost cannot judge what it cannot price", axis) |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | func TestTaskBudgetGateFiresOnWallClock(t *testing.T) { |
| 90 | b := runBudget{started: time.Now().Add(-2 * time.Hour)} |
| 91 | b.observe(&provider.Usage{PromptTokens: 1, RequestCount: 1}, nil) |
| 92 | axis, detail := b.exceeded(TaskBudget{Wall: time.Hour}) |
| 93 | if axis != "time" { |
| 94 | t.Fatalf("axis = %q, want the wall-clock crossing", axis) |
| 95 | } |
| 96 | if !strings.Contains(detail, "past the 1h0m0s budget") { |
| 97 | t.Fatalf("detail = %q, want it to name the budget it crossed", detail) |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | // Nothing is bounded out of the box. Stopping a task is the user's call: only |
| 102 | // they know which model they are paying for, and whether a long task is a |
| 103 | // runaway or the job they asked for. |
| 104 | func TestTaskBudgetShipsNoLimits(t *testing.T) { |
| 105 | if got := normalizeTaskBudget(TaskBudget{}); got.Cost != 0 || got.Wall != 0 { |
| 106 | t.Fatalf("default budget = %+v, want both axes off", got) |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | func TestTaskBudgetAxesSetIndependently(t *testing.T) { |
| 111 | got := normalizeTaskBudget(TaskBudget{Cost: 1.5, Wall: -1}) |
| 112 | if got.Cost != 1.5 || got.Wall != 0 { |
| 113 | t.Fatalf("budget = %+v, want the explicit cost kept and wall clock off", got) |
| 114 | } |
| 115 | if got := normalizeTaskBudget(TaskBudget{Wall: time.Minute}); got.Wall != time.Minute || got.Cost != 0 { |
| 116 | t.Fatalf("budget = %+v, want the explicit wall kept and cost off", got) |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | // However long an unconfigured task runs and however much it spends, nothing |
| 121 | // lands it. This is the promise the defaults make. |
| 122 | func TestUnconfiguredBudgetNeverCrosses(t *testing.T) { |
| 123 | b := runBudget{started: time.Now().Add(-8 * time.Hour)} |
| 124 | b.observe(&provider.Usage{PromptTokens: 50_000_000, CompletionTokens: 5_000_000, RequestCount: 1}, |
| 125 | &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2}) |
| 126 | if axis, detail := b.exceeded(normalizeTaskBudget(TaskBudget{})); axis != "" { |
| 127 | t.Fatalf("unconfigured budget crossed %q (%s); nothing should stop by default", axis, detail) |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | // An explicit max_steps still bounds a run by rounds when the user asks for |
| 132 | // that specifically, with every spend axis disabled. |
| 133 | func TestExplicitMaxStepsStillLandsWhenNoBudgetApplies(t *testing.T) { |
| 134 | sink := event.FuncSink(func(event.Event) {}) |
| 135 | reg := tool.NewRegistry() |
| 136 | reg.Add(readProbe{}) |
| 137 | a := New(&spendingProvider{max: 500}, reg, NewSession("sys"), |
| 138 | Options{MaxSteps: 3, TaskBudget: TaskBudget{Cost: 0, Wall: -1}}, sink) |
| 139 | |
| 140 | err := a.Run(context.Background(), "read everything") |
| 141 | var pause *maxStepsPause |
| 142 | if !errors.As(err, &pause) { |
| 143 | t.Fatalf("err = %v, want an explicit max_steps to still stop an unbudgeted runaway", err) |
| 144 | } |
| 145 | } |
| 146 |