summaryrefslogtreecommitdiff
path: root/internal/budget/budget_test.go
diff options
context:
space:
mode:
authorPeter Stone <thepeterstone@gmail.com>2026-06-03 23:01:39 +0000
committerPeter Stone <thepeterstone@gmail.com>2026-06-03 23:01:39 +0000
commit68cf1e24dd2ca2612b72babc4b35280cd3512f92 (patch)
tree14b73f21570a2f42ae729ca7e9676de6628d6819 /internal/budget/budget_test.go
parentb28cfc6ff288d083f6c8e9c055b69bfcadbceccc (diff)
parent7388337d3be5cc7f65ef547e30b2e39884dd165b (diff)
merge: integrate oss branch — budget gating, MCP back-channel, events, loopback-only bind
Key features from OSS branch: - Budget gating: rolling per-provider spend caps with BUDGET_EXCEEDED task state - Agent MCP back-channel: runners mint tokens; /mcp endpoint resolves them - Chatbot MCP server at /chatbot/mcp (requires api_token in config) - Event timeline: unified event stream replacing ad-hoc question/summary flows - Loopback-only default bind (127.0.0.1:8484); external_bind_allowed=true to expose - repository_url required on task creation (enforces traceability) - Auto-checker: spawns verification task after each execution Conflict resolutions: - serve.go: keep cfg.Runners.XEnabled() conditional registration from main + agentRegistry from oss - config.go: keep RunnersConfig from main + BudgetConfig/ExternalBindAllowed from oss - server.go: handleStaticFiles (base-path rewrite) kept; deduplicated duplicate from both sides - executor/claude.go: accepted oss deletion (ClaudeRunner replaced by ContainerRunner) - container_test.go: added noopChannel + AgentChannel parameter to Run call Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Diffstat (limited to 'internal/budget/budget_test.go')
-rw-r--r--internal/budget/budget_test.go112
1 files changed, 112 insertions, 0 deletions
diff --git a/internal/budget/budget_test.go b/internal/budget/budget_test.go
new file mode 100644
index 0000000..c257c91
--- /dev/null
+++ b/internal/budget/budget_test.go
@@ -0,0 +1,112 @@
+package budget
+
+import (
+ "testing"
+ "time"
+)
+
+type fakeSpend struct {
+ spends map[string]float64
+ since time.Time
+}
+
+func (f *fakeSpend) SpendByProviderSince(since time.Time) (map[string]float64, error) {
+ f.since = since
+ return f.spends, nil
+}
+
+func newAcct(spends map[string]float64, lim Limits) (*Accountant, *fakeSpend) {
+ src := &fakeSpend{spends: spends}
+ a := New(src, lim)
+ a.now = func() time.Time { return time.Date(2026, 5, 26, 12, 0, 0, 0, time.UTC) }
+ return a, src
+}
+
+func TestNew_DefaultsWindow(t *testing.T) {
+ a, _ := newAcct(nil, Limits{})
+ if a.lim.Window != DefaultWindow {
+ t.Errorf("want default window %v, got %v", DefaultWindow, a.lim.Window)
+ }
+}
+
+func TestHeadroom_UnlimitedWhenNoLimit(t *testing.T) {
+ a, _ := newAcct(map[string]float64{"local": 999}, Limits{PerProvider: map[string]float64{"claude": 10}})
+ h, err := a.Headroom("local")
+ if err != nil {
+ t.Fatal(err)
+ }
+ if h.Limited {
+ t.Errorf("local should be unlimited, got %+v", h)
+ }
+}
+
+func TestHeadroom_ComputesRemainingAndFraction(t *testing.T) {
+ a, src := newAcct(map[string]float64{"claude": 4}, Limits{Window: 5 * time.Hour, PerProvider: map[string]float64{"claude": 10}})
+ h, err := a.Headroom("claude")
+ if err != nil {
+ t.Fatal(err)
+ }
+ if !h.Limited || h.Limit != 10 || h.Spent != 4 || h.Remaining != 6 {
+ t.Errorf("unexpected headroom: %+v", h)
+ }
+ if h.Fraction != 0.6 {
+ t.Errorf("fraction: want 0.6, got %v", h.Fraction)
+ }
+ // Window lookback is now-5h.
+ want := time.Date(2026, 5, 26, 7, 0, 0, 0, time.UTC)
+ if !src.since.Equal(want) {
+ t.Errorf("since: want %v, got %v", want, src.since)
+ }
+}
+
+func TestHeadroom_ClampsNegativeRemaining(t *testing.T) {
+ a, _ := newAcct(map[string]float64{"claude": 15}, Limits{PerProvider: map[string]float64{"claude": 10}})
+ h, _ := a.Headroom("claude")
+ if h.Remaining != 0 {
+ t.Errorf("remaining should clamp to 0, got %v", h.Remaining)
+ }
+}
+
+func TestAllow(t *testing.T) {
+ a, _ := newAcct(map[string]float64{"claude": 8}, Limits{PerProvider: map[string]float64{"claude": 10}})
+
+ if ok, _ := a.Allow("claude", 1.5); !ok {
+ t.Error("8 + 1.5 <= 10 should be allowed")
+ }
+ if ok, _ := a.Allow("claude", 3); ok {
+ t.Error("8 + 3 > 10 should be denied")
+ }
+ if ok, _ := a.Allow("local", 1000); !ok {
+ t.Error("unlimited provider should always allow")
+ }
+}
+
+func TestAll_SortedAndSingleQuery(t *testing.T) {
+ a, src := newAcct(map[string]float64{"claude": 2, "gemini": 1},
+ Limits{PerProvider: map[string]float64{"gemini": 5, "claude": 10}})
+ calls := 0
+ orig := src.spends
+ a2 := New(&countingSpend{spends: orig, calls: &calls}, a.lim)
+ a2.now = a.now
+
+ hs, err := a2.All()
+ if err != nil {
+ t.Fatal(err)
+ }
+ if len(hs) != 2 || hs[0].Provider != "claude" || hs[1].Provider != "gemini" {
+ t.Errorf("want [claude gemini] sorted, got %+v", hs)
+ }
+ if calls != 1 {
+ t.Errorf("All should issue exactly one spend query, got %d", calls)
+ }
+}
+
+type countingSpend struct {
+ spends map[string]float64
+ calls *int
+}
+
+func (c *countingSpend) SpendByProviderSince(time.Time) (map[string]float64, error) {
+ *c.calls++
+ return c.spends, nil
+}