From 39f74dbbc7adca0e2058112408bb99694deefcaf Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 26 May 2026 20:27:36 +0000 Subject: feat(budget,storage,config): per-provider spend accounting substrate (Phase 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the budget accountant foundation: - storage: a per-execution `agent` column (additive migration, populated by the dispatcher) and SpendByProviderSince(since), summing cost_usd per provider in a window. Accurate attribution survives a task running on different providers across retries. - internal/budget: Accountant over a SpendSource + Limits, exposing Headroom (remaining/fraction per provider), Allow (would estCost breach the cap), and All (one query, sorted). Unconfigured/local providers are unlimited. - config: a [budget] section (window + provider_5h_usd map). No default cap — gating is opt-in by configuring limits. Fully unit-tested; dispatcher gating and the API/UI surface follow. https://claude.ai/code/session_01SESwn7kQ7oP62trWw6pc39 --- internal/budget/budget.go | 122 +++++++++++++++++++++++++++++++++++++++++ internal/budget/budget_test.go | 112 +++++++++++++++++++++++++++++++++++++ 2 files changed, 234 insertions(+) create mode 100644 internal/budget/budget.go create mode 100644 internal/budget/budget_test.go (limited to 'internal/budget') 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 +} -- cgit v1.2.3