Skip to content

Commit 957996c

Browse files
michaelmcneesclaude
michaelmcnees
and
claude
authored
feat: AI cost controls — per-provider, per-team, per-workflow token budgets (#9)
* fix: improve mobile navigation and layouts on docs site Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: replace mobile top nav with bottom navigation bar for better reachability Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * docs: add AI cost controls implementation plan Design decisions from /grill-me session: - Token-based budgets (not dollar-based) - Three AND-gated levels: workflow, team+provider, global - Configurable enforcement (hard block vs warn-only) - Calendar or rolling reset windows - Counter table for fast enforcement queries Ref #8 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat(budget): add ai_token_usage and team_budgets tables (migration 012) * feat(budget): add BudgetConfig to engine configuration Adds BudgetConfig struct with ResetMode, ResetDay, GlobalMonthlyTokenLimit, and DefaultTeamMonthlyTokenLimit fields, nested under EngineConfig. Includes defaults, env var bindings, and a test asserting correct default values. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(budget): add token usage and team budget DB store Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(budget): add budget checker with period calculation and AND-gate enforcement Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(budget): add budget audit actions and Prometheus metrics * feat(budget): add token_budget field to Workflow struct * feat(budget): wire budget enforcement into engine step dispatch Add pre-execution budget checks and post-execution token recording to executeStepLogic, covering both sequential and distributed execution paths. AI steps are checked against workflow token_budget, team+provider limits, and global limits before dispatch. Token usage is recorded to the budget store after successful completion. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(budget): decouple token recording from budget enforcement BudgetStore recording should work independently of BudgetChecker — allows token tracking for observability even when enforcement is disabled. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat(budget): add mantle budget CLI commands (set, get, usage, delete) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(budget): add team budget API endpoints (list, set, delete, usage) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs: add AI cost controls documentation with provider pricing links * fix(budget): wire BudgetChecker into engine, add audit events, log recording errors Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(budget): add DB constraints, validate ResetDay and TokenBudget, remove unused metric Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(budget): global limit is platform-wide, add wildcard provider fallback, UTC normalize Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * test(budget): add fail-open error behavior and ResetDay validation tests - Document intentional fail-open: DB errors in budget checks don't block AI steps - Test rolling mode rejects ResetDay outside 1-28 - Test calendar mode silently clamps invalid ResetDay to 1 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test(budget): add team hard block, provider usage error, and upper-bound clamp tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test(budget): address PR9 review feedback Add exact-equality boundary test for CheckExecutionBudget and upgrade assert.Error to require.Error in TestLoad_BudgetResetDay_RollingInvalid to fail fast on nil error. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test(budget): add rolling year-boundary test for CurrentPeriodStart Verify AddDate(0,-1,0) correctly crosses year boundary (Jan 10 with resetDay 15 yields Dec 15 of prior year). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 5089655 commit 957996c

20 files changed

Lines changed: 3047 additions & 7 deletions

File tree

‎internal/audit/audit.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,11 @@ const (
2929
ActionCredentialDeleted Action = "credential.deleted"
3030
ActionCredentialRotated Action = "credential.rotated"
3131
ActionAuthFailed Action = "auth.failed"
32+
33+
// Budget operations.
34+
ActionBudgetExceeded Action = "budget.exceeded"
35+
ActionBudgetWarning Action = "budget.warning"
36+
ActionBudgetUpdated Action = "budget.updated"
3237
)
3338

3439
// Resource identifies the target of an audit event.

‎internal/budget/budget.go‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
package budget
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"time"
7+
)
8+
9+
// CurrentPeriodStart returns the start date of the current budget period.
10+
// For "calendar" mode: first day of the current month.
11+
// For "rolling" mode: the most recent occurrence of resetDay (1-28).
12+
func CurrentPeriodStart(now time.Time, mode string, resetDay int) time.Time {
13+
now = now.UTC()
14+
if mode == "rolling" && resetDay >= 1 && resetDay <= 28 {
15+
if now.Day() >= resetDay {
16+
return time.Date(now.Year(), now.Month(), resetDay, 0, 0, 0, 0, time.UTC)
17+
}
18+
prev := now.AddDate(0, -1, 0)
19+
return time.Date(prev.Year(), prev.Month(), resetDay, 0, 0, 0, 0, time.UTC)
20+
}
21+
return time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
22+
}
23+
24+
// CheckInput describes the context of a budget check.
25+
type CheckInput struct {
26+
TeamID string
27+
Provider string
28+
}
29+
30+
// CheckResult describes the outcome of a budget check.
31+
type CheckResult struct {
32+
Blocked bool
33+
Warning bool
34+
BlockedBy string // "global", "team", "workflow", or ""
35+
Message string
36+
}
37+
38+
// Checker evaluates budget limits before AI step dispatch.
39+
// Uses function fields for DB access to allow easy testing with mocks.
40+
type Checker struct {
41+
GlobalMonthlyTokenLimit int64
42+
DefaultTeamMonthlyTokenLimit int64
43+
ResetMode string
44+
ResetDay int
45+
46+
GetTotalUsage func(ctx context.Context, period time.Time) (int64, error)
47+
GetProviderUsage func(ctx context.Context, teamID, provider string, period time.Time) (int64, error)
48+
GetTeamBudget func(ctx context.Context, teamID, provider string) (*TeamBudget, error)
49+
}
50+
51+
// Check evaluates all budget levels (global, team+provider) and returns the result.
52+
// Budget levels compose as AND gates — all must pass.
53+
func (c *Checker) Check(ctx context.Context, input CheckInput) CheckResult {
54+
now := time.Now()
55+
period := CurrentPeriodStart(now, c.ResetMode, c.ResetDay)
56+
57+
// 1. Global budget — hard block only
58+
if c.GlobalMonthlyTokenLimit > 0 && c.GetTotalUsage != nil {
59+
total, err := c.GetTotalUsage(ctx, period)
60+
if err == nil && total >= c.GlobalMonthlyTokenLimit {
61+
return CheckResult{
62+
Blocked: true,
63+
BlockedBy: "global",
64+
Message: fmt.Sprintf("global monthly token limit exceeded (%d/%d)", total, c.GlobalMonthlyTokenLimit),
65+
}
66+
}
67+
}
68+
69+
// 2. Team+provider budget — configurable enforcement
70+
if c.GetTeamBudget != nil && c.GetProviderUsage != nil {
71+
tb, err := c.GetTeamBudget(ctx, input.TeamID, input.Provider)
72+
if err == nil && tb == nil {
73+
tb, err = c.GetTeamBudget(ctx, input.TeamID, "*")
74+
}
75+
if err == nil && tb == nil && c.DefaultTeamMonthlyTokenLimit > 0 {
76+
tb = &TeamBudget{
77+
MonthlyTokenLimit: c.DefaultTeamMonthlyTokenLimit,
78+
Enforcement: "hard",
79+
}
80+
}
81+
if err == nil && tb != nil && tb.MonthlyTokenLimit > 0 {
82+
usage, err := c.GetProviderUsage(ctx, input.TeamID, input.Provider, period)
83+
if err == nil && usage >= tb.MonthlyTokenLimit {
84+
if tb.Enforcement == "warn" {
85+
return CheckResult{
86+
Warning: true,
87+
BlockedBy: "team",
88+
Message: fmt.Sprintf("team budget exceeded for provider %s (%d/%d tokens) — enforcement is warn-only", input.Provider, usage, tb.MonthlyTokenLimit),
89+
}
90+
}
91+
return CheckResult{
92+
Blocked: true,
93+
BlockedBy: "team",
94+
Message: fmt.Sprintf("team budget exceeded for provider %s (%d/%d tokens)", input.Provider, usage, tb.MonthlyTokenLimit),
95+
}
96+
}
97+
}
98+
}
99+
100+
return CheckResult{}
101+
}
102+
103+
// CheckExecutionBudget checks if a workflow execution's cumulative token usage
104+
// has exceeded the workflow-level token_budget. budgetLimit of 0 means unlimited.
105+
func CheckExecutionBudget(budgetLimit int64, usedTokens int64) CheckResult {
106+
if budgetLimit <= 0 {
107+
return CheckResult{}
108+
}
109+
if usedTokens >= budgetLimit {
110+
return CheckResult{
111+
Blocked: true,
112+
BlockedBy: "workflow",
113+
Message: fmt.Sprintf("workflow token budget exceeded (%d/%d)", usedTokens, budgetLimit),
114+
}
115+
}
116+
return CheckResult{}
117+
}

‎internal/budget/budget_test.go‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
package budget_test
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"testing"
7+
"time"
8+
9+
"github.com/dvflw/mantle/internal/budget"
10+
"github.com/stretchr/testify/assert"
11+
)
12+
13+
func TestCurrentPeriodStart_Calendar(t *testing.T) {
14+
now := time.Date(2026, 3, 23, 14, 30, 0, 0, time.UTC)
15+
start := budget.CurrentPeriodStart(now, "calendar", 1)
16+
assert.Equal(t, time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC), start)
17+
}
18+
19+
func TestCurrentPeriodStart_Rolling(t *testing.T) {
20+
now := time.Date(2026, 3, 10, 14, 30, 0, 0, time.UTC)
21+
start := budget.CurrentPeriodStart(now, "rolling", 15)
22+
assert.Equal(t, time.Date(2026, 2, 15, 0, 0, 0, 0, time.UTC), start)
23+
24+
now = time.Date(2026, 3, 20, 14, 30, 0, 0, time.UTC)
25+
start = budget.CurrentPeriodStart(now, "rolling", 15)
26+
assert.Equal(t, time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC), start)
27+
28+
now = time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC)
29+
start = budget.CurrentPeriodStart(now, "rolling", 15)
30+
assert.Equal(t, time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC), start)
31+
}
32+
33+
func TestCurrentPeriodStart_RollingYearBoundary(t *testing.T) {
34+
now := time.Date(2026, 1, 10, 0, 0, 0, 0, time.UTC)
35+
start := budget.CurrentPeriodStart(now, "rolling", 15)
36+
assert.Equal(t, time.Date(2025, 12, 15, 0, 0, 0, 0, time.UTC), start)
37+
}
38+
39+
func TestCurrentPeriodStart_RollingDay28_February(t *testing.T) {
40+
now := time.Date(2026, 3, 10, 0, 0, 0, 0, time.UTC)
41+
start := budget.CurrentPeriodStart(now, "rolling", 28)
42+
assert.Equal(t, time.Date(2026, 2, 28, 0, 0, 0, 0, time.UTC), start)
43+
}
44+
45+
func TestChecker_Check_GlobalHardBlock(t *testing.T) {
46+
checker := &budget.Checker{
47+
GlobalMonthlyTokenLimit: 1000,
48+
ResetMode: "calendar",
49+
ResetDay: 1,
50+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
51+
return 1001, nil
52+
},
53+
}
54+
55+
result := checker.Check(context.Background(), budget.CheckInput{
56+
TeamID: "team-1",
57+
Provider: "openai",
58+
})
59+
60+
assert.True(t, result.Blocked)
61+
assert.Equal(t, "global", result.BlockedBy)
62+
assert.Contains(t, result.Message, "global monthly token limit")
63+
}
64+
65+
func TestChecker_Check_TeamWarnOnly(t *testing.T) {
66+
checker := &budget.Checker{
67+
ResetMode: "calendar",
68+
ResetDay: 1,
69+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
70+
return 0, nil
71+
},
72+
GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
73+
return 5001, nil
74+
},
75+
GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
76+
return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "warn"}, nil
77+
},
78+
}
79+
80+
result := checker.Check(context.Background(), budget.CheckInput{
81+
TeamID: "team-1",
82+
Provider: "openai",
83+
})
84+
85+
assert.False(t, result.Blocked)
86+
assert.True(t, result.Warning)
87+
assert.Contains(t, result.Message, "team budget exceeded")
88+
}
89+
90+
func TestChecker_Check_AllPass(t *testing.T) {
91+
checker := &budget.Checker{
92+
GlobalMonthlyTokenLimit: 100000,
93+
ResetMode: "calendar",
94+
ResetDay: 1,
95+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
96+
return 500, nil
97+
},
98+
GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
99+
return 200, nil
100+
},
101+
GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
102+
return &budget.TeamBudget{MonthlyTokenLimit: 10000, Enforcement: "hard"}, nil
103+
},
104+
}
105+
106+
result := checker.Check(context.Background(), budget.CheckInput{
107+
TeamID: "team-1",
108+
Provider: "openai",
109+
})
110+
111+
assert.False(t, result.Blocked)
112+
assert.False(t, result.Warning)
113+
}
114+
115+
func TestChecker_Check_FailOpenOnGetTotalUsageError(t *testing.T) {
116+
checker := &budget.Checker{
117+
GlobalMonthlyTokenLimit: 1000,
118+
ResetMode: "calendar",
119+
ResetDay: 1,
120+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
121+
return 0, fmt.Errorf("db connection lost")
122+
},
123+
}
124+
125+
result := checker.Check(context.Background(), budget.CheckInput{
126+
TeamID: "team-1",
127+
Provider: "openai",
128+
})
129+
130+
// Intentional fail-open: DB errors should not block AI steps.
131+
assert.False(t, result.Blocked)
132+
assert.False(t, result.Warning)
133+
}
134+
135+
func TestChecker_Check_FailOpenOnGetTeamBudgetError(t *testing.T) {
136+
checker := &budget.Checker{
137+
ResetMode: "calendar",
138+
ResetDay: 1,
139+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
140+
return 0, nil
141+
},
142+
GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
143+
return 0, nil
144+
},
145+
GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
146+
return nil, fmt.Errorf("db timeout")
147+
},
148+
}
149+
150+
result := checker.Check(context.Background(), budget.CheckInput{
151+
TeamID: "team-1",
152+
Provider: "openai",
153+
})
154+
155+
assert.False(t, result.Blocked)
156+
assert.False(t, result.Warning)
157+
}
158+
159+
func TestChecker_Check_FailOpenOnGetProviderUsageError(t *testing.T) {
160+
checker := &budget.Checker{
161+
ResetMode: "calendar",
162+
ResetDay: 1,
163+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
164+
return 0, nil
165+
},
166+
GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
167+
return 0, fmt.Errorf("provider usage query failed")
168+
},
169+
GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
170+
return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "hard"}, nil
171+
},
172+
}
173+
174+
result := checker.Check(context.Background(), budget.CheckInput{
175+
TeamID: "team-1",
176+
Provider: "openai",
177+
})
178+
179+
assert.False(t, result.Blocked)
180+
assert.False(t, result.Warning)
181+
}
182+
183+
func TestChecker_Check_TeamHardBlock(t *testing.T) {
184+
checker := &budget.Checker{
185+
ResetMode: "calendar",
186+
ResetDay: 1,
187+
GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
188+
return 0, nil
189+
},
190+
GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
191+
return 5001, nil
192+
},
193+
GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
194+
return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "hard"}, nil
195+
},
196+
}
197+
198+
result := checker.Check(context.Background(), budget.CheckInput{
199+
TeamID: "team-1",
200+
Provider: "openai",
201+
})
202+
203+
assert.True(t, result.Blocked)
204+
assert.Equal(t, "team", result.BlockedBy)
205+
assert.Contains(t, result.Message, "team budget exceeded")
206+
}
207+
208+
func TestChecker_CheckExecutionBudget(t *testing.T) {
209+
result := budget.CheckExecutionBudget(50000, 50001)
210+
assert.True(t, result.Blocked)
211+
assert.Equal(t, "workflow", result.BlockedBy)
212+
213+
result = budget.CheckExecutionBudget(50000, 50000)
214+
assert.True(t, result.Blocked)
215+
assert.Equal(t, "workflow", result.BlockedBy)
216+
217+
result = budget.CheckExecutionBudget(50000, 49999)
218+
assert.False(t, result.Blocked)
219+
220+
result = budget.CheckExecutionBudget(0, 999999)
221+
assert.False(t, result.Blocked)
222+
}

0 commit comments

Comments
 (0)