1package subagents
2
3import (
4 "strings"
5 "testing"
6
7 "github.com/bnema/zut/packages/provider"
8)
9
10func TestBudgetLimitDefaultsToSeventyFivePercent(t *testing.T) {
11 limit, ratio, err := BudgetLimit(500_000, 0, 0)
12 if err != nil {
13 t.Fatal(err)
14 }
15 if limit != 375_000 || ratio != 0.75 {
16 t.Fatalf("BudgetLimit = %d, %v; want 375000, 0.75", limit, ratio)
17 }
18
19 limit, ratio, err = BudgetLimit(500_000, 0.5, 0)
20 if err != nil || limit != 250_000 || ratio != 0.5 {
21 t.Fatalf("ratio override = %d, %v, %v", limit, ratio, err)
22 }
23
24 limit, ratio, err = BudgetLimit(500_000, 0.5, 42_000)
25 if err != nil || limit != 42_000 || ratio != 0 {
26 t.Fatalf("token override = %d, %v, %v", limit, ratio, err)
27 }
28}
29
30func TestEffectiveBudgetLimitDerivesLegacyDefault(t *testing.T) {
31 if got := EffectiveBudgetLimit(0, 500_000); got != 375_000 {
32 t.Fatalf("legacy budget limit = %d, want 375000", got)
33 }
34 if got := EffectiveBudgetLimit(42_000, 500_000); got != 42_000 {
35 t.Fatalf("persisted budget limit = %d, want 42000", got)
36 }
37}
38
39func TestWeightedBudgetUsageDiscountsCacheReads(t *testing.T) {
40 usage := provider.Usage{InputTokens: 100, CacheReadTokens: 100, CacheWriteTokens: 20, OutputTokens: 30, ReasoningTokens: 10, ReasoningTokensKnown: true}
41 if got := WeightedBudgetUsage(usage); got != 175 {
42 t.Fatalf("WeightedBudgetUsage = %d, want 175", got)
43 }
44}
45
46func TestBudgetThresholdInstructionsAreProgressiveAndOneShot(t *testing.T) {
47 budget := NewRolloutBudget(1_000, provider.Usage{})
48 cases := []struct {
49 used int
50 state BudgetState
51 text string
52 }{
53 {500, BudgetNotice, "prioritize concrete evidence"},
54 {700, BudgetFocused, "Stop broad exploration"},
55 {850, BudgetVerifying, "Do not identify new candidates"},
56 {900, BudgetFinalizing, "Finalization is required now"},
57 }
58 for _, tc := range cases {
59 budget.Observe(provider.Usage{InputTokens: tc.used})
60 if got := budget.Snapshot().State; got != tc.state {
61 t.Fatalf("usage %d state = %q, want %q", tc.used, got, tc.state)
62 }
63 instruction, exceeded := budget.TurnContext()
64 if exceeded || !strings.Contains(instruction, tc.text) {
65 t.Fatalf("usage %d instruction = %q exceeded=%t", tc.used, instruction, exceeded)
66 }
67 if repeated, exceeded := budget.TurnContext(); repeated != "" || exceeded {
68 t.Fatalf("usage %d repeated instruction = %q exceeded=%t", tc.used, repeated, exceeded)
69 }
70 }
71
72 budget.Observe(provider.Usage{InputTokens: 1_000})
73 if instruction, exceeded := budget.TurnContext(); instruction != "" || !exceeded {
74 t.Fatalf("exhausted instruction = %q exceeded=%t", instruction, exceeded)
75 }
76}