diff options
| author | Peter Stone <thepeterstone@gmail.com> | 2026-06-03 23:01:39 +0000 |
|---|---|---|
| committer | Peter Stone <thepeterstone@gmail.com> | 2026-06-03 23:01:39 +0000 |
| commit | 68cf1e24dd2ca2612b72babc4b35280cd3512f92 (patch) | |
| tree | 14b73f21570a2f42ae729ca7e9676de6628d6819 /internal/budget | |
| parent | b28cfc6ff288d083f6c8e9c055b69bfcadbceccc (diff) | |
| parent | 7388337d3be5cc7f65ef547e30b2e39884dd165b (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')
| -rw-r--r-- | internal/budget/budget.go | 122 | ||||
| -rw-r--r-- | internal/budget/budget_test.go | 112 |
2 files changed, 234 insertions, 0 deletions
diff --git a/internal/budget/budget.go b/internal/budget/budget.go new file mode 100644 index 0000000..60ce9d2 --- /dev/null +++ b/internal/budget/budget.go @@ -0,0 +1,122 @@ +// Package budget tracks agent spend against rolling per-provider windows so the +// dispatcher can refuse (or reroute) paid work that would breach a cap. Local +// runners are free and simply carry no configured limit. +package budget + +import ( + "sort" + "time" +) + +// DefaultWindow is the rolling spend window when none is configured. +const DefaultWindow = 5 * time.Hour + +// Limits configures per-provider spend caps within a rolling window. +type Limits struct { + // Window is the rolling lookback; DefaultWindow when zero. + Window time.Duration + // PerProvider maps a provider (claude/gemini/…) to its USD cap within the + // window. A missing or non-positive entry means "unlimited" (e.g. local). + PerProvider map[string]float64 +} + +// SpendSource supplies total spend per provider since a given time. +type SpendSource interface { + SpendByProviderSince(since time.Time) (map[string]float64, error) +} + +// Headroom describes remaining budget for one provider in the current window. +type Headroom struct { + Provider string `json:"provider"` + Limited bool `json:"limited"` + Limit float64 `json:"limit_usd"` + Spent float64 `json:"spent_usd"` + Remaining float64 `json:"remaining_usd"` + // Fraction is the share of the cap still available, 0..1 (0 when unlimited). + Fraction float64 `json:"fraction_remaining"` +} + +// Accountant answers spend questions against a SpendSource and Limits. +type Accountant struct { + src SpendSource + lim Limits + now func() time.Time +} + +// New builds an Accountant. A zero Window defaults to DefaultWindow. +func New(src SpendSource, lim Limits) *Accountant { + if lim.Window <= 0 { + lim.Window = DefaultWindow + } + return &Accountant{src: src, lim: lim, now: time.Now} +} + +func (a *Accountant) since() time.Time { return a.now().Add(-a.lim.Window) } + +func (a *Accountant) headroomFrom(provider string, spends map[string]float64) Headroom { + limit := a.lim.PerProvider[provider] + if limit <= 0 { + return Headroom{Provider: provider, Limited: false} + } + spent := spends[provider] + remaining := limit - spent + if remaining < 0 { + remaining = 0 + } + return Headroom{ + Provider: provider, + Limited: true, + Limit: limit, + Spent: spent, + Remaining: remaining, + Fraction: remaining / limit, + } +} + +// Headroom returns the remaining budget for one provider in the current window. +func (a *Accountant) Headroom(provider string) (Headroom, error) { + if a.lim.PerProvider[provider] <= 0 { + return Headroom{Provider: provider, Limited: false}, nil + } + spends, err := a.src.SpendByProviderSince(a.since()) + if err != nil { + return Headroom{}, err + } + return a.headroomFrom(provider, spends), nil +} + +// Allow reports whether starting a task on provider that may cost up to estCost +// is permitted without breaching the window cap. Unlimited providers always +// pass. estCost is typically the task's max_budget_usd (0 = unknown). +func (a *Accountant) Allow(provider string, estCost float64) (bool, error) { + h, err := a.Headroom(provider) + if err != nil { + return false, err + } + if !h.Limited { + return true, nil + } + return h.Spent+estCost <= h.Limit, nil +} + +// All returns headroom for every configured provider, sorted by name, in a +// single spend query. +func (a *Accountant) All() ([]Headroom, error) { + if len(a.lim.PerProvider) == 0 { + return []Headroom{}, nil + } + spends, err := a.src.SpendByProviderSince(a.since()) + if err != nil { + return nil, err + } + providers := make([]string, 0, len(a.lim.PerProvider)) + for p := range a.lim.PerProvider { + providers = append(providers, p) + } + sort.Strings(providers) + out := make([]Headroom, 0, len(providers)) + for _, p := range providers { + out = append(out, a.headroomFrom(p, spends)) + } + return out, nil +} 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 +} |
