feat: AI cost controls — per-provider, per-team, per-workflow token budgets

internal/audit/audit.go

@@ -29,6 +29,11 @@ const (
 	ActionCredentialDeleted Action = "credential.deleted"
 	ActionCredentialRotated Action = "credential.rotated"
 	ActionAuthFailed        Action = "auth.failed"
+
+	// Budget operations.
+	ActionBudgetExceeded Action = "budget.exceeded"
+	ActionBudgetWarning  Action = "budget.warning"
+	ActionBudgetUpdated  Action = "budget.updated"
 )
 
 // Resource identifies the target of an audit event.

internal/budget/budget.go

@@ -0,0 +1,117 @@
+package budget
+
+import (
+	"context"
+	"fmt"
+	"time"
+)
+
+// CurrentPeriodStart returns the start date of the current budget period.
+// For "calendar" mode: first day of the current month.
+// For "rolling" mode: the most recent occurrence of resetDay (1-28).
+func CurrentPeriodStart(now time.Time, mode string, resetDay int) time.Time {
+	now = now.UTC()
+	if mode == "rolling" && resetDay >= 1 && resetDay <= 28 {
+		if now.Day() >= resetDay {
+			return time.Date(now.Year(), now.Month(), resetDay, 0, 0, 0, 0, time.UTC)
+		}
+		prev := now.AddDate(0, -1, 0)
+		return time.Date(prev.Year(), prev.Month(), resetDay, 0, 0, 0, 0, time.UTC)
+	}
+	return time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
+}
+
+// CheckInput describes the context of a budget check.
+type CheckInput struct {
+	TeamID   string
+	Provider string
+}
+
+// CheckResult describes the outcome of a budget check.
+type CheckResult struct {
+	Blocked   bool
+	Warning   bool
+	BlockedBy string // "global", "team", "workflow", or ""
+	Message   string
+}
+
+// Checker evaluates budget limits before AI step dispatch.
+// Uses function fields for DB access to allow easy testing with mocks.
+type Checker struct {
+	GlobalMonthlyTokenLimit      int64
+	DefaultTeamMonthlyTokenLimit int64
+	ResetMode                    string
+	ResetDay                     int
+
+	GetTotalUsage    func(ctx context.Context, period time.Time) (int64, error)
+	GetProviderUsage func(ctx context.Context, teamID, provider string, period time.Time) (int64, error)
+	GetTeamBudget    func(ctx context.Context, teamID, provider string) (*TeamBudget, error)
+}
+
+// Check evaluates all budget levels (global, team+provider) and returns the result.
+// Budget levels compose as AND gates — all must pass.
+func (c *Checker) Check(ctx context.Context, input CheckInput) CheckResult {
+	now := time.Now()
+	period := CurrentPeriodStart(now, c.ResetMode, c.ResetDay)
+
+	// 1. Global budget — hard block only
+	if c.GlobalMonthlyTokenLimit > 0 && c.GetTotalUsage != nil {
+		total, err := c.GetTotalUsage(ctx, period)
+		if err == nil && total >= c.GlobalMonthlyTokenLimit {
+			return CheckResult{
+				Blocked:   true,
+				BlockedBy: "global",
+				Message:   fmt.Sprintf("global monthly token limit exceeded (%d/%d)", total, c.GlobalMonthlyTokenLimit),
+			}
+		}
+	}
+
+	// 2. Team+provider budget — configurable enforcement
+	if c.GetTeamBudget != nil && c.GetProviderUsage != nil {
+		tb, err := c.GetTeamBudget(ctx, input.TeamID, input.Provider)
+		if err == nil && tb == nil {
+			tb, err = c.GetTeamBudget(ctx, input.TeamID, "*")
+		}
+		if err == nil && tb == nil && c.DefaultTeamMonthlyTokenLimit > 0 {
+			tb = &TeamBudget{
+				MonthlyTokenLimit: c.DefaultTeamMonthlyTokenLimit,
+				Enforcement:       "hard",
+			}
+		}
+		if err == nil && tb != nil && tb.MonthlyTokenLimit > 0 {
+			usage, err := c.GetProviderUsage(ctx, input.TeamID, input.Provider, period)
+			if err == nil && usage >= tb.MonthlyTokenLimit {
+				if tb.Enforcement == "warn" {
+					return CheckResult{
+						Warning:   true,
+						BlockedBy: "team",
+						Message:   fmt.Sprintf("team budget exceeded for provider %s (%d/%d tokens) — enforcement is warn-only", input.Provider, usage, tb.MonthlyTokenLimit),
+					}
+				}
+				return CheckResult{
+					Blocked:   true,
+					BlockedBy: "team",
+					Message:   fmt.Sprintf("team budget exceeded for provider %s (%d/%d tokens)", input.Provider, usage, tb.MonthlyTokenLimit),
+				}
+			}
+		}
+	}
+
+	return CheckResult{}
+}
+
+// CheckExecutionBudget checks if a workflow execution's cumulative token usage
+// has exceeded the workflow-level token_budget. budgetLimit of 0 means unlimited.
+func CheckExecutionBudget(budgetLimit int64, usedTokens int64) CheckResult {
+	if budgetLimit <= 0 {
+		return CheckResult{}
+	}
+	if usedTokens >= budgetLimit {
+		return CheckResult{
+			Blocked:   true,
+			BlockedBy: "workflow",
+			Message:   fmt.Sprintf("workflow token budget exceeded (%d/%d)", usedTokens, budgetLimit),
+		}
+	}
+	return CheckResult{}
+}

internal/budget/budget_test.go

@@ -0,0 +1,216 @@
+package budget_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/dvflw/mantle/internal/budget"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCurrentPeriodStart_Calendar(t *testing.T) {
+	now := time.Date(2026, 3, 23, 14, 30, 0, 0, time.UTC)
+	start := budget.CurrentPeriodStart(now, "calendar", 1)
+	assert.Equal(t, time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC), start)
+}
+
+func TestCurrentPeriodStart_Rolling(t *testing.T) {
+	now := time.Date(2026, 3, 10, 14, 30, 0, 0, time.UTC)
+	start := budget.CurrentPeriodStart(now, "rolling", 15)
+	assert.Equal(t, time.Date(2026, 2, 15, 0, 0, 0, 0, time.UTC), start)
+
+	now = time.Date(2026, 3, 20, 14, 30, 0, 0, time.UTC)
+	start = budget.CurrentPeriodStart(now, "rolling", 15)
+	assert.Equal(t, time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC), start)
+
+	now = time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC)
+	start = budget.CurrentPeriodStart(now, "rolling", 15)
+	assert.Equal(t, time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC), start)
+}
+
+func TestCurrentPeriodStart_RollingDay28_February(t *testing.T) {
+	now := time.Date(2026, 3, 10, 0, 0, 0, 0, time.UTC)
+	start := budget.CurrentPeriodStart(now, "rolling", 28)
+	assert.Equal(t, time.Date(2026, 2, 28, 0, 0, 0, 0, time.UTC), start)
+}
+
+func TestChecker_Check_GlobalHardBlock(t *testing.T) {
+	checker := &budget.Checker{
+		GlobalMonthlyTokenLimit: 1000,
+		ResetMode:               "calendar",
+		ResetDay:                1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 1001, nil
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	assert.True(t, result.Blocked)
+	assert.Equal(t, "global", result.BlockedBy)
+	assert.Contains(t, result.Message, "global monthly token limit")
+}
+
+func TestChecker_Check_TeamWarnOnly(t *testing.T) {
+	checker := &budget.Checker{
+		ResetMode: "calendar",
+		ResetDay:  1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 0, nil
+		},
+		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
+			return 5001, nil
+		},
+		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
+			return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "warn"}, nil
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	assert.False(t, result.Blocked)
+	assert.True(t, result.Warning)
+	assert.Contains(t, result.Message, "team budget exceeded")
+}
+
+func TestChecker_Check_AllPass(t *testing.T) {
+	checker := &budget.Checker{
+		GlobalMonthlyTokenLimit: 100000,
+		ResetMode:               "calendar",
+		ResetDay:                1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 500, nil
+		},
+		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
+			return 200, nil
+		},
+		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
+			return &budget.TeamBudget{MonthlyTokenLimit: 10000, Enforcement: "hard"}, nil
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	assert.False(t, result.Blocked)
+	assert.False(t, result.Warning)
+}
+
+func TestChecker_Check_FailOpenOnGetTotalUsageError(t *testing.T) {
+	checker := &budget.Checker{
+		GlobalMonthlyTokenLimit: 1000,
+		ResetMode:               "calendar",
+		ResetDay:                1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 0, fmt.Errorf("db connection lost")
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	// Intentional fail-open: DB errors should not block AI steps.
+	assert.False(t, result.Blocked)
+	assert.False(t, result.Warning)
+}
+
+func TestChecker_Check_FailOpenOnGetTeamBudgetError(t *testing.T) {
+	checker := &budget.Checker{
+		ResetMode: "calendar",
+		ResetDay:  1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 0, nil
+		},
+		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
+			return 0, nil
+		},
+		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
+			return nil, fmt.Errorf("db timeout")
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	assert.False(t, result.Blocked)
+	assert.False(t, result.Warning)
+}
+
+func TestChecker_Check_FailOpenOnGetProviderUsageError(t *testing.T) {
+	checker := &budget.Checker{
+		ResetMode: "calendar",
+		ResetDay:  1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 0, nil
+		},
+		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
+			return 0, fmt.Errorf("provider usage query failed")
+		},
+		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
+			return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "hard"}, nil
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	assert.False(t, result.Blocked)
+	assert.False(t, result.Warning)
+}
+
+func TestChecker_Check_TeamHardBlock(t *testing.T) {
+	checker := &budget.Checker{
+		ResetMode: "calendar",
+		ResetDay:  1,
+		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
+			return 0, nil
+		},
+		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
+			return 5001, nil
+		},
+		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
+			return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "hard"}, nil
+		},
+	}
+
+	result := checker.Check(context.Background(), budget.CheckInput{
+		TeamID:   "team-1",
+		Provider: "openai",
+	})
+
+	assert.True(t, result.Blocked)
+	assert.Equal(t, "team", result.BlockedBy)
+	assert.Contains(t, result.Message, "team budget exceeded")
+}
+
+func TestChecker_CheckExecutionBudget(t *testing.T) {
+	result := budget.CheckExecutionBudget(50000, 50001)
+	assert.True(t, result.Blocked)
+	assert.Equal(t, "workflow", result.BlockedBy)
+
+	result = budget.CheckExecutionBudget(50000, 50000)
+	assert.True(t, result.Blocked)
+	assert.Equal(t, "workflow", result.BlockedBy)
+
+	result = budget.CheckExecutionBudget(50000, 49999)
+	assert.False(t, result.Blocked)
+
+	result = budget.CheckExecutionBudget(0, 999999)
+	assert.False(t, result.Blocked)
+}