diff --git a/internal/audit/audit.go b/internal/audit/audit.go index bf11ee1..396bd77 100644 --- a/internal/audit/audit.go +++ b/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. diff --git a/internal/budget/budget.go b/internal/budget/budget.go new file mode 100644 index 0000000..91948c2 --- /dev/null +++ b/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{} +} diff --git a/internal/budget/budget_test.go b/internal/budget/budget_test.go new file mode 100644 index 0000000..2fa8025 --- /dev/null +++ b/internal/budget/budget_test.go @@ -0,0 +1,222 @@ +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_RollingYearBoundary(t *testing.T) { + now := time.Date(2026, 1, 10, 0, 0, 0, 0, time.UTC) + start := budget.CurrentPeriodStart(now, "rolling", 15) + assert.Equal(t, time.Date(2025, 12, 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) +} diff --git a/internal/budget/store.go b/internal/budget/store.go new file mode 100644 index 0000000..75f60da --- /dev/null +++ b/internal/budget/store.go @@ -0,0 +1,164 @@ +package budget + +import ( + "context" + "database/sql" + "time" +) + +// ProviderUsage holds aggregated token usage for a team+provider in a period. +type ProviderUsage struct { + Provider string + PromptTokens int64 + CompletionTokens int64 + TotalTokens int64 +} + +// TeamBudget holds a team's budget configuration for a provider. +type TeamBudget struct { + ID string `json:"id"` + TeamID string `json:"team_id"` + Provider string `json:"provider"` + MonthlyTokenLimit int64 `json:"monthly_token_limit"` + Enforcement string `json:"enforcement"` +} + +// Store handles budget-related database operations. +type Store struct { + db *sql.DB +} + +// NewStore creates a new budget store. +func NewStore(db *sql.DB) *Store { + return &Store{db: db} +} + +// IncrementUsage atomically adds tokens to the counter for a (team, provider, model, period). +func (s *Store) IncrementUsage(ctx context.Context, teamID, provider, model string, periodStart time.Time, promptTokens, completionTokens int64) error { + _, err := s.db.ExecContext(ctx, ` + INSERT INTO ai_token_usage (team_id, provider, model, period_start, prompt_tokens, completion_tokens, total_tokens, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, now()) + ON CONFLICT (team_id, provider, model, period_start) + DO UPDATE SET + prompt_tokens = ai_token_usage.prompt_tokens + EXCLUDED.prompt_tokens, + completion_tokens = ai_token_usage.completion_tokens + EXCLUDED.completion_tokens, + total_tokens = ai_token_usage.total_tokens + EXCLUDED.total_tokens, + updated_at = now() + `, teamID, provider, model, periodStart, promptTokens, completionTokens, promptTokens+completionTokens) + return err +} + +// GetUsage returns aggregated token usage for a team+provider in a period. +func (s *Store) GetUsage(ctx context.Context, teamID, provider string, periodStart time.Time) (*ProviderUsage, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(prompt_tokens), 0), + COALESCE(SUM(completion_tokens), 0), + COALESCE(SUM(total_tokens), 0) + FROM ai_token_usage + WHERE team_id = $1 AND provider = $2 AND period_start = $3 + `, teamID, provider, periodStart) + + var u ProviderUsage + u.Provider = provider + if err := row.Scan(&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens); err != nil { + return nil, err + } + return &u, nil +} + +// GetTotalUsage returns aggregated token usage across all providers for a team in a period. +func (s *Store) GetTotalUsage(ctx context.Context, teamID string, periodStart time.Time) (*ProviderUsage, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(prompt_tokens), 0), + COALESCE(SUM(completion_tokens), 0), + COALESCE(SUM(total_tokens), 0) + FROM ai_token_usage + WHERE team_id = $1 AND period_start = $2 + `, teamID, periodStart) + + var u ProviderUsage + u.Provider = "*" + if err := row.Scan(&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens); err != nil { + return nil, err + } + return &u, nil +} + +// GetGlobalTotalUsage returns aggregated token usage across ALL teams in a period. +func (s *Store) GetGlobalTotalUsage(ctx context.Context, periodStart time.Time) (*ProviderUsage, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(prompt_tokens), 0), + COALESCE(SUM(completion_tokens), 0), + COALESCE(SUM(total_tokens), 0) + FROM ai_token_usage + WHERE period_start = $1 + `, periodStart) + + var u ProviderUsage + u.Provider = "*" + if err := row.Scan(&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens); err != nil { + return nil, err + } + return &u, nil +} + +// GetTeamBudget returns the budget for a team+provider. Returns nil if no explicit budget is set. +func (s *Store) GetTeamBudget(ctx context.Context, teamID, provider string) (*TeamBudget, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT id, team_id, provider, monthly_token_limit, enforcement + FROM team_budgets + WHERE team_id = $1 AND provider = $2 + `, teamID, provider) + + var b TeamBudget + if err := row.Scan(&b.ID, &b.TeamID, &b.Provider, &b.MonthlyTokenLimit, &b.Enforcement); err != nil { + if err == sql.ErrNoRows { + return nil, nil + } + return nil, err + } + return &b, nil +} + +// SetTeamBudget upserts a team's budget for a provider. +func (s *Store) SetTeamBudget(ctx context.Context, teamID, provider string, monthlyLimit int64, enforcement string) error { + _, err := s.db.ExecContext(ctx, ` + INSERT INTO team_budgets (team_id, provider, monthly_token_limit, enforcement, updated_at) + VALUES ($1, $2, $3, $4, now()) + ON CONFLICT (team_id, provider) + DO UPDATE SET monthly_token_limit = EXCLUDED.monthly_token_limit, + enforcement = EXCLUDED.enforcement, + updated_at = now() + `, teamID, provider, monthlyLimit, enforcement) + return err +} + +// DeleteTeamBudget removes a team's budget for a provider. +func (s *Store) DeleteTeamBudget(ctx context.Context, teamID, provider string) error { + _, err := s.db.ExecContext(ctx, ` + DELETE FROM team_budgets WHERE team_id = $1 AND provider = $2 + `, teamID, provider) + return err +} + +// ListTeamBudgets returns all budgets for a team. +func (s *Store) ListTeamBudgets(ctx context.Context, teamID string) ([]TeamBudget, error) { + rows, err := s.db.QueryContext(ctx, ` + SELECT id, team_id, provider, monthly_token_limit, enforcement + FROM team_budgets WHERE team_id = $1 ORDER BY provider + `, teamID) + if err != nil { + return nil, err + } + defer rows.Close() + + var budgets []TeamBudget + for rows.Next() { + var b TeamBudget + if err := rows.Scan(&b.ID, &b.TeamID, &b.Provider, &b.MonthlyTokenLimit, &b.Enforcement); err != nil { + return nil, err + } + budgets = append(budgets, b) + } + return budgets, rows.Err() +} diff --git a/internal/budget/store_test.go b/internal/budget/store_test.go new file mode 100644 index 0000000..99bd164 --- /dev/null +++ b/internal/budget/store_test.go @@ -0,0 +1,147 @@ +package budget_test + +import ( + "context" + "database/sql" + "testing" + "time" + + "github.com/dvflw/mantle/internal/budget" + "github.com/dvflw/mantle/internal/config" + "github.com/dvflw/mantle/internal/db" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/modules/postgres" + "github.com/testcontainers/testcontainers-go/wait" +) + +func setupTestDB(t *testing.T) *sql.DB { + t.Helper() + ctx := context.Background() + pgContainer, err := postgres.Run(ctx, + "postgres:16-alpine", + postgres.WithDatabase("mantle_test"), + postgres.WithUsername("mantle"), + postgres.WithPassword("mantle"), + testcontainers.WithWaitStrategy( + wait.ForLog("database system is ready to accept connections"). + WithOccurrence(2). + WithStartupTimeout(30*time.Second), + ), + ) + if err != nil { + t.Skipf("Could not start Postgres container: %v", err) + } + t.Cleanup(func() { + if err := pgContainer.Terminate(ctx); err != nil { + t.Logf("Failed to terminate container: %v", err) + } + }) + connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable") + if err != nil { + t.Fatalf("Failed to get connection string: %v", err) + } + database, err := db.Open(config.DatabaseConfig{URL: connStr}) + if err != nil { + t.Fatalf("Failed to open database: %v", err) + } + t.Cleanup(func() { database.Close() }) + if err := db.Migrate(ctx, database); err != nil { + t.Fatalf("Failed to run migrations: %v", err) + } + return database +} + +func TestStore_IncrementUsage(t *testing.T) { + database := setupTestDB(t) + store := budget.NewStore(database) + ctx := context.Background() + teamID := "00000000-0000-0000-0000-000000000001" + period := time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC) + + err := store.IncrementUsage(ctx, teamID, "openai", "gpt-4o", period, 100, 50) + require.NoError(t, err) + + err = store.IncrementUsage(ctx, teamID, "openai", "gpt-4o", period, 200, 100) + require.NoError(t, err) + + usage, err := store.GetUsage(ctx, teamID, "openai", period) + require.NoError(t, err) + assert.Equal(t, int64(300), usage.PromptTokens) + assert.Equal(t, int64(150), usage.CompletionTokens) + assert.Equal(t, int64(450), usage.TotalTokens) +} + +func TestStore_GetTeamBudget_NotFound(t *testing.T) { + database := setupTestDB(t) + store := budget.NewStore(database) + ctx := context.Background() + + b, err := store.GetTeamBudget(ctx, "00000000-0000-0000-0000-000000000001", "openai") + require.NoError(t, err) + assert.Nil(t, b) +} + +func TestStore_SetTeamBudget_Upsert(t *testing.T) { + database := setupTestDB(t) + store := budget.NewStore(database) + ctx := context.Background() + teamID := "00000000-0000-0000-0000-000000000001" + + err := store.SetTeamBudget(ctx, teamID, "openai", 1000000, "hard") + require.NoError(t, err) + + b, err := store.GetTeamBudget(ctx, teamID, "openai") + require.NoError(t, err) + require.NotNil(t, b) + assert.Equal(t, int64(1000000), b.MonthlyTokenLimit) + assert.Equal(t, "hard", b.Enforcement) + + // Upsert with new values + err = store.SetTeamBudget(ctx, teamID, "openai", 2000000, "warn") + require.NoError(t, err) + + b, err = store.GetTeamBudget(ctx, teamID, "openai") + require.NoError(t, err) + assert.Equal(t, int64(2000000), b.MonthlyTokenLimit) + assert.Equal(t, "warn", b.Enforcement) +} + +func TestStore_GetTotalUsage(t *testing.T) { + database := setupTestDB(t) + store := budget.NewStore(database) + ctx := context.Background() + teamID := "00000000-0000-0000-0000-000000000001" + period := time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC) + + err := store.IncrementUsage(ctx, teamID, "openai", "gpt-4o", period, 100, 50) + require.NoError(t, err) + err = store.IncrementUsage(ctx, teamID, "bedrock", "claude-3", period, 200, 100) + require.NoError(t, err) + + usage, err := store.GetTotalUsage(ctx, teamID, period) + require.NoError(t, err) + assert.Equal(t, int64(300), usage.PromptTokens) + assert.Equal(t, int64(150), usage.CompletionTokens) + assert.Equal(t, int64(450), usage.TotalTokens) + assert.Equal(t, "*", usage.Provider) +} + +func TestStore_ListTeamBudgets(t *testing.T) { + database := setupTestDB(t) + store := budget.NewStore(database) + ctx := context.Background() + teamID := "00000000-0000-0000-0000-000000000001" + + err := store.SetTeamBudget(ctx, teamID, "openai", 1000000, "hard") + require.NoError(t, err) + err = store.SetTeamBudget(ctx, teamID, "bedrock", 500000, "warn") + require.NoError(t, err) + + budgets, err := store.ListTeamBudgets(ctx, teamID) + require.NoError(t, err) + require.Len(t, budgets, 2) + assert.Equal(t, "bedrock", budgets[0].Provider) // sorted + assert.Equal(t, "openai", budgets[1].Provider) +} diff --git a/internal/cli/budget.go b/internal/cli/budget.go new file mode 100644 index 0000000..6766a84 --- /dev/null +++ b/internal/cli/budget.go @@ -0,0 +1,158 @@ +package cli + +import ( + "fmt" + "text/tabwriter" + "time" + + "github.com/dvflw/mantle/internal/auth" + "github.com/dvflw/mantle/internal/budget" + "github.com/dvflw/mantle/internal/config" + "github.com/dvflw/mantle/internal/db" + "github.com/spf13/cobra" +) + +func newBudgetCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "budget", + Short: "Manage AI token budgets", + } + cmd.AddCommand( + newBudgetSetCommand(), + newBudgetGetCommand(), + newBudgetUsageCommand(), + newBudgetDeleteCommand(), + ) + return cmd +} + +func newBudgetStore(cmd *cobra.Command) (*budget.Store, func(), error) { + cfg := config.FromContext(cmd.Context()) + database, err := db.Open(cfg.Database) + if err != nil { + return nil, nil, fmt.Errorf("opening database: %w", err) + } + return budget.NewStore(database), func() { database.Close() }, nil +} + +func newBudgetSetCommand() *cobra.Command { + var enforcement string + cmd := &cobra.Command{ + Use: "set ", + Short: "Set a monthly token budget for a provider", + Long: "Set a monthly token budget. Provider can be 'openai', 'bedrock', or '*' for all providers.", + Args: cobra.ExactArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + provider := args[0] + var limit int64 + if _, err := fmt.Sscanf(args[1], "%d", &limit); err != nil { + return fmt.Errorf("invalid token limit: %s", args[1]) + } + + teamID := auth.TeamIDFromContext(cmd.Context()) + if err := store.SetTeamBudget(cmd.Context(), teamID, provider, limit, enforcement); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "Budget set: %s → %d tokens/month (enforcement: %s)\n", provider, limit, enforcement) + return nil + }, + } + cmd.Flags().StringVar(&enforcement, "enforcement", "hard", "Enforcement mode: 'hard' (block) or 'warn' (log only)") + return cmd +} + +func newBudgetGetCommand() *cobra.Command { + return &cobra.Command{ + Use: "get", + Short: "List all team budgets", + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + teamID := auth.TeamIDFromContext(cmd.Context()) + budgets, err := store.ListTeamBudgets(cmd.Context(), teamID) + if err != nil { + return err + } + if len(budgets) == 0 { + fmt.Fprintln(cmd.OutOrStdout(), "No budgets configured (global defaults apply)") + return nil + } + w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0) + fmt.Fprintln(w, "PROVIDER\tLIMIT\tENFORCEMENT") + for _, b := range budgets { + fmt.Fprintf(w, "%s\t%d tokens/month\t%s\n", b.Provider, b.MonthlyTokenLimit, b.Enforcement) + } + return w.Flush() + }, + } +} + +func newBudgetUsageCommand() *cobra.Command { + var provider string + cmd := &cobra.Command{ + Use: "usage", + Short: "Show current period token usage", + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + cfg := config.FromContext(cmd.Context()) + teamID := auth.TeamIDFromContext(cmd.Context()) + period := budget.CurrentPeriodStart(time.Now(), cfg.Engine.Budget.ResetMode, cfg.Engine.Budget.ResetDay) + + var usage *budget.ProviderUsage + if provider != "" { + usage, err = store.GetUsage(cmd.Context(), teamID, provider, period) + } else { + usage, err = store.GetTotalUsage(cmd.Context(), teamID, period) + } + if err != nil { + return err + } + + fmt.Fprintf(cmd.OutOrStdout(), "Period: %s\n", period.Format("2006-01-02")) + fmt.Fprintf(cmd.OutOrStdout(), "Provider: %s\n", usage.Provider) + fmt.Fprintf(cmd.OutOrStdout(), "Prompt: %d tokens\n", usage.PromptTokens) + fmt.Fprintf(cmd.OutOrStdout(), "Completion: %d tokens\n", usage.CompletionTokens) + fmt.Fprintf(cmd.OutOrStdout(), "Total: %d tokens\n", usage.TotalTokens) + return nil + }, + } + cmd.Flags().StringVar(&provider, "provider", "", "Filter by provider (e.g., 'openai', 'bedrock')") + return cmd +} + +func newBudgetDeleteCommand() *cobra.Command { + return &cobra.Command{ + Use: "delete ", + Short: "Remove a team budget for a provider", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + teamID := auth.TeamIDFromContext(cmd.Context()) + if err := store.DeleteTeamBudget(cmd.Context(), teamID, args[0]); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "Budget removed for provider: %s\n", args[0]) + return nil + }, + } +} diff --git a/internal/cli/root.go b/internal/cli/root.go index 2b3cfe1..125ecd5 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -77,6 +77,7 @@ Full documentation: https://mantle.dvflw.co/docs`, newPluginsCommand(), newLibraryCommand(), newCleanupCommand(), + newBudgetCommand(), ) // Info. diff --git a/internal/cli/serve.go b/internal/cli/serve.go index 02c8d98..9c45479 100644 --- a/internal/cli/serve.go +++ b/internal/cli/serve.go @@ -1,12 +1,15 @@ package cli import ( + "context" "fmt" "os/signal" "syscall" + "time" "github.com/dvflw/mantle/internal/audit" "github.com/dvflw/mantle/internal/auth" + "github.com/dvflw/mantle/internal/budget" "github.com/dvflw/mantle/internal/config" "github.com/dvflw/mantle/internal/connector" "github.com/dvflw/mantle/internal/db" @@ -78,11 +81,39 @@ func newServeCommand() *cobra.Command { } } + // Wire budget enforcement into the engine. + budgetStore := budget.NewStore(database) + eng.BudgetStore = budgetStore + eng.BudgetChecker = &budget.Checker{ + GlobalMonthlyTokenLimit: cfg.Engine.Budget.GlobalMonthlyTokenLimit, + DefaultTeamMonthlyTokenLimit: cfg.Engine.Budget.DefaultTeamMonthlyTokenLimit, + ResetMode: cfg.Engine.Budget.ResetMode, + ResetDay: cfg.Engine.Budget.ResetDay, + GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) { + u, err := budgetStore.GetGlobalTotalUsage(ctx, period) + if err != nil { + return 0, err + } + return u.TotalTokens, nil + }, + GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) { + u, err := budgetStore.GetUsage(ctx, teamID, provider, period) + if err != nil { + return 0, err + } + return u.TotalTokens, nil + }, + GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) { + return budgetStore.GetTeamBudget(ctx, teamID, provider) + }, + } + srv := server.New(database, eng, cfg.API.Address) srv.Auditor = auditor srv.TLSCertFile = cfg.API.TLS.CertFile srv.TLSKeyFile = cfg.API.TLS.KeyFile srv.AuthStore = &auth.Store{DB: database} + srv.BudgetStore = budgetStore // Wire up OIDC validator if configured. if cfg.Auth.OIDC.IssuerURL != "" { diff --git a/internal/config/config.go b/internal/config/config.go index b0fbffa..2825021 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -96,6 +96,14 @@ type LogConfig struct { Level string `mapstructure:"level"` } +// BudgetConfig holds AI cost control settings. +type BudgetConfig struct { + ResetMode string `mapstructure:"reset_mode"` // "calendar" or "rolling" + ResetDay int `mapstructure:"reset_day"` // 1-28, used when reset_mode is "rolling" + GlobalMonthlyTokenLimit int64 `mapstructure:"global_monthly_token_limit"` // 0 = unlimited, hard block + DefaultTeamMonthlyTokenLimit int64 `mapstructure:"default_team_monthly_token_limit"` // 0 = unlimited, applies to teams without explicit budget +} + // EngineConfig holds distributed engine settings. type EngineConfig struct { NodeID string `mapstructure:"node_id"` @@ -112,6 +120,7 @@ type EngineConfig struct { AllowedBaseURLs []string `mapstructure:"allowed_base_urls"` AllowedModels []string `mapstructure:"allowed_models"` // empty = all allowed MaxToolRoundsLimit int `mapstructure:"max_tool_rounds_limit"` // 0 = no limit + Budget BudgetConfig `mapstructure:"budget"` } type contextKey struct{} @@ -152,6 +161,12 @@ func Load(cmd *cobra.Command) (*Config, error) { v.SetDefault("engine.default_max_tool_rounds", 10) v.SetDefault("engine.default_max_tool_calls_per_round", 10) + // Budget defaults + v.SetDefault("engine.budget.reset_mode", "calendar") + v.SetDefault("engine.budget.reset_day", 1) + v.SetDefault("engine.budget.global_monthly_token_limit", 0) + v.SetDefault("engine.budget.default_team_monthly_token_limit", 0) + // Config file configPath, _ := cmd.Flags().GetString("config") if configPath != "" { @@ -219,6 +234,12 @@ func Load(cmd *cobra.Command) (*Config, error) { _ = v.BindEnv("engine.allowed_models", "MANTLE_ENGINE_ALLOWED_MODELS") _ = v.BindEnv("engine.max_tool_rounds_limit", "MANTLE_ENGINE_MAX_TOOL_ROUNDS_LIMIT") + // Budget env var bindings + _ = v.BindEnv("engine.budget.reset_mode", "MANTLE_ENGINE_BUDGET_RESET_MODE") + _ = v.BindEnv("engine.budget.reset_day", "MANTLE_ENGINE_BUDGET_RESET_DAY") + _ = v.BindEnv("engine.budget.global_monthly_token_limit", "MANTLE_ENGINE_BUDGET_GLOBAL_MONTHLY_TOKEN_LIMIT") + _ = v.BindEnv("engine.budget.default_team_monthly_token_limit", "MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT") + // CLI flag binding if f := cmd.Flags().Lookup("database-url"); f != nil { _ = v.BindPFlag("database.url", f) @@ -235,6 +256,15 @@ func Load(cmd *cobra.Command) (*Config, error) { return nil, err } + // Validate budget reset_day range. + if cfg.Engine.Budget.ResetDay < 1 || cfg.Engine.Budget.ResetDay > 28 { + if cfg.Engine.Budget.ResetMode == "rolling" { + return nil, fmt.Errorf("engine.budget.reset_day must be between 1 and 28, got %d", cfg.Engine.Budget.ResetDay) + } + // For calendar mode, reset_day is ignored, so just clamp it silently. + cfg.Engine.Budget.ResetDay = 1 + } + // Warn if database URL uses sslmode=prefer on a non-loopback host. if dbURL := cfg.Database.URL; dbURL != "" { if parsed, err := url.Parse(dbURL); err == nil { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 616289b..ae1572a 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -194,3 +194,56 @@ log: t.Errorf("Log.Level = %q, want error (flag override)", cfg.Log.Level) } } + +func TestLoad_BudgetDefaults(t *testing.T) { + cmd := newTestCommand() + + cfg, err := Load(cmd) + require.NoError(t, err) + + assert.Equal(t, "calendar", cfg.Engine.Budget.ResetMode) + assert.Equal(t, 1, cfg.Engine.Budget.ResetDay) + assert.Equal(t, int64(0), cfg.Engine.Budget.GlobalMonthlyTokenLimit) + assert.Equal(t, int64(0), cfg.Engine.Budget.DefaultTeamMonthlyTokenLimit) +} + +func TestLoad_BudgetResetDay_RollingInvalid(t *testing.T) { + // ResetDay 0 with rolling mode should error + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_MODE", "rolling") + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_DAY", "0") + + cmd := newTestCommand() + _, err := Load(cmd) + require.Error(t, err) + assert.Contains(t, err.Error(), "reset_day must be between 1 and 28") + + // ResetDay 29 with rolling mode should error + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_DAY", "29") + cmd = newTestCommand() + _, err = Load(cmd) + require.Error(t, err) + assert.Contains(t, err.Error(), "reset_day must be between 1 and 28") +} + +func TestLoad_BudgetResetDay_CalendarClamps(t *testing.T) { + // Invalid ResetDay with calendar mode should silently clamp to 1 (lower bound) + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_MODE", "calendar") + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_DAY", "0") + + cmd := newTestCommand() + cfg, err := Load(cmd) + require.NoError(t, err) + assert.Equal(t, 1, cfg.Engine.Budget.ResetDay) +} + +func TestLoad_BudgetResetDay_CalendarClampsUpperBound(t *testing.T) { + // ResetDay >28 with calendar mode should silently clamp to 1 + // (calendar mode ignores reset_day entirely, so any invalid value is safe to clamp) + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_MODE", "calendar") + t.Setenv("MANTLE_ENGINE_BUDGET_RESET_DAY", "100") + + cmd := newTestCommand() + cfg, err := Load(cmd) + require.NoError(t, err) + assert.Equal(t, 1, cfg.Engine.Budget.ResetDay) +} diff --git a/internal/db/migrations/012_ai_cost_controls.sql b/internal/db/migrations/012_ai_cost_controls.sql new file mode 100644 index 0000000..98a2b60 --- /dev/null +++ b/internal/db/migrations/012_ai_cost_controls.sql @@ -0,0 +1,36 @@ +-- +goose Up + +-- Tracks cumulative AI token usage per team/provider/model/period. +-- One row per (team, provider, model, period_start) — atomically incremented. +CREATE TABLE ai_token_usage ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + team_id UUID NOT NULL REFERENCES teams(id) ON DELETE RESTRICT, + provider TEXT NOT NULL, -- "openai", "bedrock" + model TEXT NOT NULL, -- "gpt-4o", "claude-3-sonnet", etc. + period_start DATE NOT NULL, -- bucket start date + prompt_tokens BIGINT NOT NULL DEFAULT 0, + completion_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE(team_id, provider, model, period_start) +); + +CREATE INDEX idx_ai_token_usage_team_period ON ai_token_usage(team_id, period_start); +CREATE INDEX idx_ai_token_usage_team_provider_period ON ai_token_usage(team_id, provider, period_start); + +-- Per-team, per-provider budget configuration. +-- If no row exists for a team+provider, the global default applies. +CREATE TABLE team_budgets ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + team_id UUID NOT NULL REFERENCES teams(id) ON DELETE RESTRICT, + provider TEXT NOT NULL, -- "openai", "bedrock", or "*" for all providers + monthly_token_limit BIGINT NOT NULL, -- 0 = unlimited + enforcement TEXT NOT NULL DEFAULT 'hard' CHECK (enforcement IN ('hard', 'warn')), -- "hard" (block) or "warn" (log + continue) + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE(team_id, provider) +); + +-- +goose Down +DROP TABLE IF EXISTS team_budgets; +DROP TABLE IF EXISTS ai_token_usage; diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 38bcdc9..55cf7a1 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -5,10 +5,13 @@ import ( "database/sql" "encoding/json" "fmt" + "log" + "strings" "time" "github.com/dvflw/mantle/internal/audit" "github.com/dvflw/mantle/internal/auth" + "github.com/dvflw/mantle/internal/budget" mantleCEL "github.com/dvflw/mantle/internal/cel" "github.com/dvflw/mantle/internal/connector" "github.com/dvflw/mantle/internal/metrics" @@ -23,7 +26,16 @@ type Engine struct { Auditor audit.Emitter CEL *mantleCEL.Evaluator Resolver *secret.Resolver - MaxToolRoundsLimit int // admin ceiling for max_rounds; 0 = no limit + MaxToolRoundsLimit int // admin ceiling for max_rounds; 0 = no limit + BudgetChecker *budget.Checker // nil = budget enforcement disabled + BudgetStore *budget.Store // nil = token usage recording disabled +} + +// StepContext carries workflow-level metadata needed by step execution. +type StepContext struct { + WorkflowTokenBudget int64 + CompletedSteps map[string]map[string]any + TeamID string } // New creates an Engine with sensible defaults. @@ -137,6 +149,13 @@ func (e *Engine) resumeExecution(ctx context.Context, execID string, workflowNam return nil, fmt.Errorf("checkpoint: %w", err) } + // Build StepContext for budget enforcement. + sc := StepContext{ + WorkflowTokenBudget: wf.TokenBudget, + CompletedSteps: completedSteps, + TeamID: auth.TeamIDFromContext(ctx), + } + // Execute steps sequentially. for _, step := range wf.Steps { // Skip already-completed steps (checkpoint recovery). @@ -149,12 +168,13 @@ func (e *Engine) resumeExecution(ctx context.Context, execID string, workflowNam } stepStart := time.Now() - stepResult := e.executeStep(ctx, execID, workflowName, step, celCtx) + stepResult := e.executeStep(ctx, execID, workflowName, step, celCtx, sc) stepResult.Duration = time.Since(stepStart) result.Steps[step.Name] = stepResult if stepResult.Status == "completed" { celCtx.Steps[step.Name] = map[string]any{"output": stepResult.Output} + sc.CompletedSteps[step.Name] = stepResult.Output } else if stepResult.Status == "skipped" { celCtx.Steps[step.Name] = map[string]any{"output": map[string]any{}} } else if stepResult.Status == "failed" { @@ -194,7 +214,7 @@ func checkOutputSize(output map[string]any, maxBytes int) error { return nil } -func (e *Engine) executeStep(ctx context.Context, execID string, workflowName string, step workflow.Step, celCtx *mantleCEL.Context) StepResult { +func (e *Engine) executeStep(ctx context.Context, execID string, workflowName string, step workflow.Step, celCtx *mantleCEL.Context, sc StepContext) StepResult { // Evaluate conditional. if step.If != "" { shouldRun, err := e.CEL.EvalBool(step.If, celCtx) @@ -232,7 +252,7 @@ func (e *Engine) executeStep(ctx context.Context, execID string, workflowName st }) // Delegate to the shared step logic (connector lookup, CEL resolution, retry). - output, lastErr := e.executeStepLogic(ctx, execID, step, celCtx, workflowName) + output, lastErr := e.executeStepLogic(ctx, execID, step, celCtx, workflowName, sc) if lastErr != nil { errMsg := lastErr.Error() @@ -283,7 +303,78 @@ func (e *Engine) executeStep(ctx context.Context, execID string, workflowName st // executeStepLogic contains the core connector invocation logic: CEL param resolution, // credential resolution, connector lookup, and retry/timeout handling. It is used by // both the sequential executeStep path and the distributed MakeStepExecutor bridge. -func (e *Engine) executeStepLogic(ctx context.Context, execID string, step workflow.Step, celCtx *mantleCEL.Context, workflowName string) (map[string]any, error) { +func (e *Engine) executeStepLogic(ctx context.Context, execID string, step workflow.Step, celCtx *mantleCEL.Context, workflowName string, sc StepContext) (map[string]any, error) { + // Budget enforcement — check before dispatching AI steps. + if strings.HasPrefix(step.Action, "ai/") { + provider := "openai" + if p, ok := step.Params["provider"].(string); ok && p != "" { + provider = p + } + + // 1. Workflow execution budget + if sc.WorkflowTokenBudget > 0 { + var usedTokens int64 + for _, stepOutput := range sc.CompletedSteps { + usedTokens += extractTotalTokens(stepOutput) + } + result := budget.CheckExecutionBudget(sc.WorkflowTokenBudget, usedTokens) + if result.Blocked { + e.Auditor.Emit(ctx, audit.Event{ + Timestamp: time.Now(), + Actor: "engine", + Action: audit.ActionBudgetExceeded, + Resource: audit.Resource{Type: "workflow_execution", ID: execID}, + Metadata: map[string]string{ + "blocked_by": result.BlockedBy, + "message": result.Message, + }, + TeamID: sc.TeamID, + }) + return nil, fmt.Errorf("%s", result.Message) + } + } + + // 2. Team+provider and global budget + if e.BudgetChecker != nil && sc.TeamID != "" { + result := e.BudgetChecker.Check(ctx, budget.CheckInput{ + TeamID: sc.TeamID, + Provider: provider, + }) + if result.Blocked { + e.Auditor.Emit(ctx, audit.Event{ + Timestamp: time.Now(), + Actor: "engine", + Action: audit.ActionBudgetExceeded, + Resource: audit.Resource{Type: "workflow_execution", ID: execID}, + Metadata: map[string]string{ + "blocked_by": result.BlockedBy, + "message": result.Message, + "provider": provider, + }, + TeamID: sc.TeamID, + }) + metrics.BudgetCheckTotal.WithLabelValues(sc.TeamID, provider, "blocked").Inc() + return nil, fmt.Errorf("%s", result.Message) + } + if result.Warning { + e.Auditor.Emit(ctx, audit.Event{ + Timestamp: time.Now(), + Actor: "engine", + Action: audit.ActionBudgetWarning, + Resource: audit.Resource{Type: "workflow_execution", ID: execID}, + Metadata: map[string]string{ + "message": result.Message, + "provider": provider, + }, + TeamID: sc.TeamID, + }) + metrics.BudgetCheckTotal.WithLabelValues(sc.TeamID, provider, "warning").Inc() + } else { + metrics.BudgetCheckTotal.WithLabelValues(sc.TeamID, provider, "pass").Inc() + } + } + } + // Resolve CEL expressions in params. resolvedParams, err := e.CEL.ResolveParams(step.Params, celCtx) if err != nil { @@ -360,6 +451,28 @@ func (e *Engine) executeStepLogic(ctx context.Context, execID string, step workf } } + // Record token usage for budget tracking. + if lastErr == nil && strings.HasPrefix(step.Action, "ai/") && e.BudgetStore != nil && sc.TeamID != "" && output != nil { + promptTok, completionTok := extractTokenCounts(output) + if promptTok > 0 || completionTok > 0 { + provider := "openai" + if p, ok := step.Params["provider"].(string); ok && p != "" { + provider = p + } + model, _ := output["model"].(string) + resetMode, resetDay := "calendar", 1 + if e.BudgetChecker != nil { + resetMode = e.BudgetChecker.ResetMode + resetDay = e.BudgetChecker.ResetDay + } + period := budget.CurrentPeriodStart(time.Now(), resetMode, resetDay) + if err := e.BudgetStore.IncrementUsage(ctx, sc.TeamID, provider, model, period, promptTok, completionTok); err != nil { + // Log but don't fail the step — token recording is best-effort. + log.Printf("budget: failed to record token usage: %v", err) + } + } + } + return output, lastErr } @@ -374,7 +487,17 @@ func (e *Engine) MakeStepExecutor(wf *workflow.Workflow, execID string, celCtx * if step == nil { return nil, fmt.Errorf("step %q not found in workflow", stepName) } - output, err := e.executeStepLogic(ctx, execID, *step, celCtx, wf.Name) + sc := StepContext{ + WorkflowTokenBudget: wf.TokenBudget, + TeamID: auth.TeamIDFromContext(ctx), + CompletedSteps: make(map[string]map[string]any), + } + for name, stepData := range celCtx.Steps { + if out, ok := stepData["output"].(map[string]any); ok { + sc.CompletedSteps[name] = out + } + } + output, err := e.executeStepLogic(ctx, execID, *step, celCtx, wf.Name, sc) if err != nil { return output, err } @@ -445,7 +568,12 @@ func (e *Engine) MakeGlobalStepExecutor() StepExecutor { celCtx.Steps[name] = map[string]any{"output": output} } - output, err := e.executeStepLogic(ctx, execID, *step, celCtx, workflowName) + sc := StepContext{ + WorkflowTokenBudget: wf.TokenBudget, + CompletedSteps: completedSteps, + TeamID: teamID, + } + output, err := e.executeStepLogic(ctx, execID, *step, celCtx, workflowName, sc) if err != nil { return output, err } @@ -599,3 +727,43 @@ func (e *Engine) loadCompletedSteps(ctx context.Context, execID string) (map[str } return result, rows.Err() } + +func extractTokenCounts(output map[string]any) (prompt, completion int64) { + usage, ok := output["usage"].(map[string]any) + if !ok { + return 0, 0 + } + if v, ok := usage["prompt_tokens"]; ok { + prompt = toInt64(v) + } + if v, ok := usage["completion_tokens"]; ok { + completion = toInt64(v) + } + return prompt, completion +} + +func extractTotalTokens(output map[string]any) int64 { + usage, ok := output["usage"].(map[string]any) + if !ok { + return 0 + } + return toInt64(usage["total_tokens"]) +} + +func toInt64(v any) int64 { + switch n := v.(type) { + case int: + return int64(n) + case int64: + return n + case float64: + return int64(n) + case json.Number: + if i, err := n.Int64(); err == nil { + return i + } + return 0 + default: + return 0 + } +} diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index fc57710..c087e4e 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -123,6 +123,14 @@ func RecordLeaseRenewal() { LeaseRenewalsTotal.Inc() } func RecordLeaseExpiration() { LeaseExpirationsTotal.Inc() } func RecordReaperReclaimed(n int) { ReaperReclaimedTotal.Add(float64(n)) } +// Budget metrics. +var ( + BudgetCheckTotal = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "mantle_budget_check_total", + Help: "Total budget checks performed before AI step dispatch", + }, []string{"team_id", "provider", "result"}) +) + // Tool-use helper functions. func RecordToolCall(step, tool, status string) { diff --git a/internal/server/api.go b/internal/server/api.go index 6baa23e..4cd166b 100644 --- a/internal/server/api.go +++ b/internal/server/api.go @@ -8,7 +8,10 @@ import ( "strings" "time" + "github.com/dvflw/mantle/internal/audit" "github.com/dvflw/mantle/internal/auth" + "github.com/dvflw/mantle/internal/budget" + "github.com/dvflw/mantle/internal/config" "github.com/dvflw/mantle/internal/workflow" ) @@ -364,3 +367,115 @@ func (s *Server) handleGetWorkflowVersion(w http.ResponseWriter, r *http.Request "definition": def, }) } + +func (s *Server) handleListBudgets(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + budgets, err := s.BudgetStore.ListTeamBudgets(r.Context(), teamID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + writeJSON(w, http.StatusOK, budgets) +} + +func (s *Server) handleSetBudget(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + provider := r.PathValue("provider") + + var body struct { + MonthlyTokenLimit int64 `json:"monthly_token_limit"` + Enforcement string `json:"enforcement"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, "invalid request body", http.StatusBadRequest) + return + } + if body.Enforcement == "" { + body.Enforcement = "hard" + } + if body.Enforcement != "hard" && body.Enforcement != "warn" { + http.Error(w, "enforcement must be 'hard' or 'warn'", http.StatusBadRequest) + return + } + + if err := s.BudgetStore.SetTeamBudget(r.Context(), teamID, provider, body.MonthlyTokenLimit, body.Enforcement); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + actor := "api" + if u := auth.UserFromContext(r.Context()); u != nil { + actor = u.ID + } + + s.Auditor.Emit(r.Context(), audit.Event{ + Timestamp: time.Now(), + Actor: actor, + Action: audit.ActionBudgetUpdated, + Resource: audit.Resource{Type: "team_budget", ID: teamID}, + Metadata: map[string]string{ + "provider": provider, + "limit": fmt.Sprintf("%d", body.MonthlyTokenLimit), + "enforcement": body.Enforcement, + }, + TeamID: teamID, + }) + + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) +} + +func (s *Server) handleDeleteBudget(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + provider := r.PathValue("provider") + + if err := s.BudgetStore.DeleteTeamBudget(r.Context(), teamID, provider); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + actor := "api" + if u := auth.UserFromContext(r.Context()); u != nil { + actor = u.ID + } + s.Auditor.Emit(r.Context(), audit.Event{ + Timestamp: time.Now(), + Actor: actor, + Action: audit.ActionBudgetUpdated, + Resource: audit.Resource{Type: "team_budget", ID: teamID}, + Metadata: map[string]string{ + "provider": provider, + "action": "deleted", + }, + TeamID: teamID, + }) + + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) +} + +func (s *Server) handleGetUsage(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + cfg := config.FromContext(r.Context()) + + period := budget.CurrentPeriodStart(time.Now(), cfg.Engine.Budget.ResetMode, cfg.Engine.Budget.ResetDay) + + provider := r.URL.Query().Get("provider") + var usage *budget.ProviderUsage + var err error + if provider != "" { + usage, err = s.BudgetStore.GetUsage(r.Context(), teamID, provider, period) + } else { + usage, err = s.BudgetStore.GetTotalUsage(r.Context(), teamID, period) + } + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + writeJSON(w, http.StatusOK, map[string]any{ + "period_start": period.Format("2006-01-02"), + "provider": usage.Provider, + "prompt_tokens": usage.PromptTokens, + "completion_tokens": usage.CompletionTokens, + "total_tokens": usage.TotalTokens, + }) +} diff --git a/internal/server/server.go b/internal/server/server.go index 99f5f28..84ce032 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -14,6 +14,7 @@ import ( "github.com/dvflw/mantle/internal/api/health" "github.com/dvflw/mantle/internal/audit" "github.com/dvflw/mantle/internal/auth" + "github.com/dvflw/mantle/internal/budget" "github.com/dvflw/mantle/internal/config" "github.com/dvflw/mantle/internal/engine" "github.com/dvflw/mantle/internal/metrics" @@ -28,6 +29,7 @@ type Server struct { AuthStore *auth.Store OIDCValidator *auth.OIDCValidator Auditor audit.Emitter + BudgetStore *budget.Store Address string TLSCertFile string TLSKeyFile string @@ -136,6 +138,12 @@ func (s *Server) Start(ctx context.Context) error { mux.HandleFunc("GET /api/v1/workflows/{name}/versions", s.handleListWorkflowVersions) mux.HandleFunc("GET /api/v1/workflows/{name}/versions/{version}", s.handleGetWorkflowVersion) + // Budget endpoints. + mux.HandleFunc("GET /api/v1/budgets", s.handleListBudgets) + mux.HandleFunc("PUT /api/v1/budgets/{provider}", s.handleSetBudget) + mux.HandleFunc("DELETE /api/v1/budgets/{provider}", s.handleDeleteBudget) + mux.HandleFunc("GET /api/v1/budgets/usage", s.handleGetUsage) + // Start distributed engine components (worker + reaper). if cfg := config.FromContext(ctx); cfg != nil { claimer := &engine.Claimer{ diff --git a/internal/workflow/validate.go b/internal/workflow/validate.go index 2264c72..39d66c0 100644 --- a/internal/workflow/validate.go +++ b/internal/workflow/validate.go @@ -43,6 +43,15 @@ func Validate(result *ParseResult) []ValidationError { }) } + // Validate token_budget is non-negative. + if w.TokenBudget < 0 { + line, col := findFieldPosition(root, "token_budget") + errs = append(errs, ValidationError{ + Line: line, Column: col, Field: "token_budget", + Message: fmt.Sprintf("token_budget must be >= 0, got %d", w.TokenBudget), + }) + } + // Validate steps exist and are non-empty. if len(w.Steps) == 0 { line, col := findFieldPosition(root, "steps") diff --git a/internal/workflow/workflow.go b/internal/workflow/workflow.go index 335f46a..e9187a4 100644 --- a/internal/workflow/workflow.go +++ b/internal/workflow/workflow.go @@ -10,6 +10,7 @@ type Workflow struct { Triggers []Trigger `yaml:"triggers"` Steps []Step `yaml:"steps"` Timeout string `yaml:"timeout"` // Go duration string, e.g., "30m" + TokenBudget int64 `yaml:"token_budget"` // 0 = unlimited } // Trigger defines an automatic execution trigger for a workflow. diff --git a/internal/workflow/workflow_test.go b/internal/workflow/workflow_test.go index 0de422b..527997b 100644 --- a/internal/workflow/workflow_test.go +++ b/internal/workflow/workflow_test.go @@ -60,3 +60,40 @@ func TestWorkflowStructTags(t *testing.T) { t.Errorf("expected max_attempts 3, got %d", w.Steps[0].Retry.MaxAttempts) } } + +func TestParse_TokenBudget(t *testing.T) { + input := []byte(` +name: test-workflow +token_budget: 500000 +steps: + - name: step1 + action: http/request + params: + url: "https://example.com" +`) + result, err := ParseBytes(input) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.Workflow.TokenBudget != int64(500000) { + t.Errorf("token_budget: got %d, want 500000", result.Workflow.TokenBudget) + } +} + +func TestParse_TokenBudget_Zero(t *testing.T) { + input := []byte(` +name: test-workflow +steps: + - name: step1 + action: http/request + params: + url: "https://example.com" +`) + result, err := ParseBytes(input) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.Workflow.TokenBudget != int64(0) { + t.Errorf("token_budget: got %d, want 0", result.Workflow.TokenBudget) + } +} diff --git a/plans/ai-cost-controls.md b/plans/ai-cost-controls.md new file mode 100644 index 0000000..98572a5 --- /dev/null +++ b/plans/ai-cost-controls.md @@ -0,0 +1,1582 @@ +# AI Cost Controls Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Implement per-provider, per-team, and per-workflow token budget enforcement with configurable reset windows, hard/soft blocking, and audit trail integration. + +**Architecture:** A new `ai_token_usage` counter table tracks cumulative tokens per (team, provider, model, period). Budget enforcement runs as a pre-dispatch check in `executeStep()` before any `ai/*` action. Three budget levels compose as AND gates: workflow-level (per-execution, YAML field), team+provider (monthly, DB-stored), and global (monthly, config file). Global budgets hard-block; team/workflow budgets are configurable (block or warn). A CLI command and API endpoint manage team budgets. Configurable reset windows support calendar month or rolling periods with a custom start day. + +**Tech Stack:** Go, Postgres, Cobra CLI, Prometheus metrics, audit events + +**GitHub Issue:** #8 + +--- + +## Design Decisions (from /grill-me session) + +1. **Token-based budgets** — not dollar-based. Docs link to provider pricing pages. +2. **Three levels, AND-gated**: workflow (per-execution), team+provider (monthly), global (monthly). All must pass. +3. **Workflow budget** is a top-level YAML field (`token_budget`). Existing step-level `max_token_budget` is unchanged. +4. **Enforcement before next AI step dispatch**, not mid-step. Tokens already spent are never wasted. +5. **Global budget = hard block.** Team/workflow budgets are configurable: hard block or warn-only (log + audit event + metrics). +6. **Counter table** `ai_token_usage(team_id, provider, model, period_start)` — atomically incremented after each AI step. Budget enforcement reads from this. +7. **Configurable reset window** — calendar month or rolling with a start day (1-28). Global config in `mantle.yaml`. +8. **Three config surfaces**: `mantle.yaml` for global defaults, CLI command for team budgets, API endpoint for programmatic access. +9. **Enforcement check lives in `executeStep()`** in the engine, before dispatching `ai/*` actions. One DB read per AI step (negligible vs. LLM call cost). + +--- + +## File Structure + +| File | Responsibility | +|------|---------------| +| **Create:** `internal/budget/budget.go` | `BudgetChecker` interface, `TokenBudgetChecker` implementation, budget period calculation, enforcement logic | +| **Create:** `internal/budget/budget_test.go` | Unit tests for period calculation, enforcement logic, AND-gate composition | +| **Create:** `internal/budget/store.go` | DB operations: increment counters, query usage, CRUD for team budgets | +| **Create:** `internal/budget/store_test.go` | Integration tests against real Postgres (testcontainers) | +| **Create:** `internal/db/migrations/012_ai_cost_controls.sql` | `ai_token_usage` counter table, `team_budgets` table | +| **Create:** `internal/cli/budget.go` | `mantle budget` CLI subcommands (set, get, usage) | +| **Create:** `internal/cli/budget_test.go` | CLI command tests | +| **Modify:** `internal/config/config.go` | Add `BudgetConfig` to `EngineConfig` (global budget, reset mode, default team budget) | +| **Modify:** `internal/workflow/workflow.go` | Add `TokenBudget` field to `Workflow` struct | +| **Modify:** `internal/workflow/workflow_test.go` | Test parsing of `token_budget` field | +| **Modify:** `internal/engine/engine.go` | Inject `BudgetChecker`, call pre-dispatch check before AI steps, record usage after step completion | +| **Modify:** `internal/engine/engine_test.go` | Test budget enforcement in engine | +| **Modify:** `internal/server/api.go` | Add team budget API endpoints | +| **Modify:** `internal/server/server.go` | Register budget API routes | +| **Modify:** `internal/metrics/metrics.go` | Add budget enforcement metrics | +| **Modify:** `internal/audit/audit.go` | Add budget-related audit action constants | +| **Create:** `site/src/content/docs/ai-cost-controls.md` | User-facing docs: budget config, enforcement semantics, provider pricing links | + +**Important:** Tasks 1-3 are the data foundation (migration, store, config). Task 4 adds the core enforcement logic. Task 5 wires it into the engine. Task 6 adds the workflow-level budget. Task 7 adds API endpoints. Task 8 adds the CLI. Task 9 adds docs. Tasks must be executed sequentially. + +--- + +### Task 1: Database migration for token usage tracking and team budgets + +**Files:** +- Create: `internal/db/migrations/012_ai_cost_controls.sql` + +- [ ] **Step 1: Write the migration** + +```sql +-- +goose Up + +-- Tracks cumulative AI token usage per team/provider/model/period. +-- One row per (team, provider, model, period_start) — atomically incremented. +CREATE TABLE ai_token_usage ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + team_id UUID NOT NULL REFERENCES teams(id), + provider TEXT NOT NULL, -- "openai", "bedrock" + model TEXT NOT NULL, -- "gpt-4o", "claude-3-sonnet", etc. + period_start DATE NOT NULL, -- bucket start date + prompt_tokens BIGINT NOT NULL DEFAULT 0, + completion_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE(team_id, provider, model, period_start) +); + +CREATE INDEX idx_ai_token_usage_team_period ON ai_token_usage(team_id, period_start); +CREATE INDEX idx_ai_token_usage_team_provider_period ON ai_token_usage(team_id, provider, period_start); + +-- Per-team, per-provider budget configuration. +-- If no row exists for a team+provider, the global default applies. +CREATE TABLE team_budgets ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + team_id UUID NOT NULL REFERENCES teams(id), + provider TEXT NOT NULL, -- "openai", "bedrock", or "*" for all providers + monthly_token_limit BIGINT NOT NULL, -- 0 = unlimited + enforcement TEXT NOT NULL DEFAULT 'hard', -- "hard" (block) or "warn" (log + continue) + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE(team_id, provider) +); + +-- +goose Down +DROP TABLE IF EXISTS team_budgets; +DROP TABLE IF EXISTS ai_token_usage; +``` + +- [ ] **Step 2: Run migration to verify it applies cleanly** + +Run: `make migrate` +Expected: Migration 012 applies without errors. + +- [ ] **Step 3: Verify rollback works** + +Run: `GOOSE_COMMAND=down make migrate` (or equivalent — check Makefile for down target) +Expected: Tables dropped cleanly. + +- [ ] **Step 4: Re-apply and commit** + +Run: `make migrate` + +```bash +git add internal/db/migrations/012_ai_cost_controls.sql +git commit -m "feat(budget): add ai_token_usage and team_budgets tables (migration 012)" +``` + +--- + +### Task 2: Config — add budget settings to EngineConfig + +**Files:** +- Modify: `internal/config/config.go` +- Modify: `internal/config/config_test.go` (if exists, otherwise create) + +- [ ] **Step 1: Write the failing test** + +Add a test that loads config with budget fields and asserts they unmarshal correctly. + +```go +func TestLoad_BudgetDefaults(t *testing.T) { + cmd := &cobra.Command{} + cmd.Flags().String("config", "", "") + + cfg, err := Load(cmd) + require.NoError(t, err) + + // Budget defaults + assert.Equal(t, "calendar", cfg.Engine.Budget.ResetMode) + assert.Equal(t, 1, cfg.Engine.Budget.ResetDay) + assert.Equal(t, int64(0), cfg.Engine.Budget.GlobalMonthlyTokenLimit) + assert.Equal(t, int64(0), cfg.Engine.Budget.DefaultTeamMonthlyTokenLimit) +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./internal/config/ -run TestLoad_BudgetDefaults -v` +Expected: FAIL — `cfg.Engine.Budget` field does not exist. + +- [ ] **Step 3: Add BudgetConfig struct and wire into EngineConfig** + +In `internal/config/config.go`, add: + +```go +// BudgetConfig holds AI cost control settings. +type BudgetConfig struct { + ResetMode string `mapstructure:"reset_mode"` // "calendar" or "rolling" + ResetDay int `mapstructure:"reset_day"` // 1-28, used when reset_mode is "rolling" + GlobalMonthlyTokenLimit int64 `mapstructure:"global_monthly_token_limit"` // 0 = unlimited, hard block + DefaultTeamMonthlyTokenLimit int64 `mapstructure:"default_team_monthly_token_limit"` // 0 = unlimited, applies to teams without explicit budget +} +``` + +Add `Budget BudgetConfig \`mapstructure:"budget"\`` to `EngineConfig`. + +Add defaults in `Load()`: + +```go +v.SetDefault("engine.budget.reset_mode", "calendar") +v.SetDefault("engine.budget.reset_day", 1) +v.SetDefault("engine.budget.global_monthly_token_limit", 0) +v.SetDefault("engine.budget.default_team_monthly_token_limit", 0) +``` + +Add env var bindings: + +```go +_ = v.BindEnv("engine.budget.reset_mode", "MANTLE_ENGINE_BUDGET_RESET_MODE") +_ = v.BindEnv("engine.budget.reset_day", "MANTLE_ENGINE_BUDGET_RESET_DAY") +_ = v.BindEnv("engine.budget.global_monthly_token_limit", "MANTLE_ENGINE_BUDGET_GLOBAL_MONTHLY_TOKEN_LIMIT") +_ = v.BindEnv("engine.budget.default_team_monthly_token_limit", "MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT") +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `go test ./internal/config/ -run TestLoad_BudgetDefaults -v` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add internal/config/config.go internal/config/config_test.go +git commit -m "feat(budget): add BudgetConfig to engine configuration" +``` + +--- + +### Task 3: Budget store — DB operations for token usage and team budgets + +**Files:** +- Create: `internal/budget/store.go` +- Create: `internal/budget/store_test.go` + +- [ ] **Step 1: Write the integration test for IncrementUsage** + +Note: `testSetup(t)` should follow the testcontainers pattern from `internal/engine/test_helpers_test.go` — start a Postgres container, run migrations, and return a `*sql.DB`. Copy/adapt that helper. + +```go +func TestStore_IncrementUsage(t *testing.T) { + db := testSetup(t) // testcontainers Postgres with migrations applied + + store := budget.NewStore(db) + ctx := context.Background() + teamID := "00000000-0000-0000-0000-000000000001" + period := time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC) + + // First increment creates the row + err := store.IncrementUsage(ctx, teamID, "openai", "gpt-4o", period, 100, 50) + require.NoError(t, err) + + // Second increment adds to existing row + err = store.IncrementUsage(ctx, teamID, "openai", "gpt-4o", period, 200, 100) + require.NoError(t, err) + + // Query usage for the period + usage, err := store.GetUsage(ctx, teamID, "openai", period) + require.NoError(t, err) + assert.Equal(t, int64(300), usage.PromptTokens) // 100 + 200 + assert.Equal(t, int64(150), usage.CompletionTokens) // 50 + 100 + assert.Equal(t, int64(450), usage.TotalTokens) // (100+50) + (200+100) +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./internal/budget/ -run TestStore_IncrementUsage -v` +Expected: FAIL — package/types don't exist. + +- [ ] **Step 3: Implement the store** + +Create `internal/budget/store.go`: + +```go +package budget + +import ( + "context" + "database/sql" + "time" +) + +// ProviderUsage holds aggregated token usage for a team+provider in a period. +type ProviderUsage struct { + Provider string + PromptTokens int64 + CompletionTokens int64 + TotalTokens int64 +} + +// TeamBudget holds a team's budget configuration for a provider. +type TeamBudget struct { + ID string `json:"id"` + TeamID string `json:"team_id"` + Provider string `json:"provider"` // specific provider or "*" for all + MonthlyTokenLimit int64 `json:"monthly_token_limit"` + Enforcement string `json:"enforcement"` // "hard" or "warn" +} + +// Store handles budget-related database operations. +type Store struct { + db *sql.DB +} + +// NewStore creates a new budget store. +func NewStore(db *sql.DB) *Store { + return &Store{db: db} +} + +// IncrementUsage atomically adds tokens to the counter for a (team, provider, model, period). +// Uses INSERT ... ON CONFLICT UPDATE to upsert. +func (s *Store) IncrementUsage(ctx context.Context, teamID, provider, model string, periodStart time.Time, promptTokens, completionTokens int64) error { + _, err := s.db.ExecContext(ctx, ` + INSERT INTO ai_token_usage (team_id, provider, model, period_start, prompt_tokens, completion_tokens, total_tokens, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, now()) + ON CONFLICT (team_id, provider, model, period_start) + DO UPDATE SET + prompt_tokens = ai_token_usage.prompt_tokens + EXCLUDED.prompt_tokens, + completion_tokens = ai_token_usage.completion_tokens + EXCLUDED.completion_tokens, + total_tokens = ai_token_usage.total_tokens + EXCLUDED.total_tokens, + updated_at = now() + `, teamID, provider, model, periodStart, promptTokens, completionTokens, promptTokens+completionTokens) + return err +} + +// GetUsage returns aggregated token usage for a team+provider in a period. +// Sums across all models for the given provider. +func (s *Store) GetUsage(ctx context.Context, teamID, provider string, periodStart time.Time) (*ProviderUsage, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(prompt_tokens), 0), + COALESCE(SUM(completion_tokens), 0), + COALESCE(SUM(total_tokens), 0) + FROM ai_token_usage + WHERE team_id = $1 AND provider = $2 AND period_start = $3 + `, teamID, provider, periodStart) + + var u ProviderUsage + u.Provider = provider + if err := row.Scan(&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens); err != nil { + return nil, err + } + return &u, nil +} + +// GetTotalUsage returns aggregated token usage across all providers for a team in a period. +func (s *Store) GetTotalUsage(ctx context.Context, teamID string, periodStart time.Time) (*ProviderUsage, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(prompt_tokens), 0), + COALESCE(SUM(completion_tokens), 0), + COALESCE(SUM(total_tokens), 0) + FROM ai_token_usage + WHERE team_id = $1 AND period_start = $2 + `, teamID, periodStart) + + var u ProviderUsage + u.Provider = "*" + if err := row.Scan(&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens); err != nil { + return nil, err + } + return &u, nil +} + +// GetTeamBudget returns the budget for a team+provider. Returns nil if no explicit budget is set. +func (s *Store) GetTeamBudget(ctx context.Context, teamID, provider string) (*TeamBudget, error) { + row := s.db.QueryRowContext(ctx, ` + SELECT id, team_id, provider, monthly_token_limit, enforcement + FROM team_budgets + WHERE team_id = $1 AND provider = $2 + `, teamID, provider) + + var b TeamBudget + if err := row.Scan(&b.ID, &b.TeamID, &b.Provider, &b.MonthlyTokenLimit, &b.Enforcement); err != nil { + if err == sql.ErrNoRows { + return nil, nil + } + return nil, err + } + return &b, nil +} + +// SetTeamBudget upserts a team's budget for a provider. +func (s *Store) SetTeamBudget(ctx context.Context, teamID, provider string, monthlyLimit int64, enforcement string) error { + _, err := s.db.ExecContext(ctx, ` + INSERT INTO team_budgets (team_id, provider, monthly_token_limit, enforcement, updated_at) + VALUES ($1, $2, $3, $4, now()) + ON CONFLICT (team_id, provider) + DO UPDATE SET monthly_token_limit = EXCLUDED.monthly_token_limit, + enforcement = EXCLUDED.enforcement, + updated_at = now() + `, teamID, provider, monthlyLimit, enforcement) + return err +} + +// DeleteTeamBudget removes a team's budget for a provider. +func (s *Store) DeleteTeamBudget(ctx context.Context, teamID, provider string) error { + _, err := s.db.ExecContext(ctx, ` + DELETE FROM team_budgets WHERE team_id = $1 AND provider = $2 + `, teamID, provider) + return err +} + +// ListTeamBudgets returns all budgets for a team. +func (s *Store) ListTeamBudgets(ctx context.Context, teamID string) ([]TeamBudget, error) { + rows, err := s.db.QueryContext(ctx, ` + SELECT id, team_id, provider, monthly_token_limit, enforcement + FROM team_budgets WHERE team_id = $1 ORDER BY provider + `, teamID) + if err != nil { + return nil, err + } + defer rows.Close() + + var budgets []TeamBudget + for rows.Next() { + var b TeamBudget + if err := rows.Scan(&b.ID, &b.TeamID, &b.Provider, &b.MonthlyTokenLimit, &b.Enforcement); err != nil { + return nil, err + } + budgets = append(budgets, b) + } + return budgets, rows.Err() +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `go test ./internal/budget/ -run TestStore_IncrementUsage -v -timeout 120s` +Expected: PASS + +- [ ] **Step 5: Write additional store tests** + +Add tests for: +- `TestStore_GetTeamBudget_NotFound` — returns nil, nil when no budget set +- `TestStore_SetTeamBudget_Upsert` — set, update, verify +- `TestStore_GetTotalUsage` — multiple providers sum correctly +- `TestStore_ListTeamBudgets` — returns sorted list + +- [ ] **Step 6: Run all store tests** + +Run: `go test ./internal/budget/ -v -timeout 120s` +Expected: All PASS + +- [ ] **Step 7: Commit** + +```bash +git add internal/budget/store.go internal/budget/store_test.go +git commit -m "feat(budget): add token usage and team budget DB store" +``` + +--- + +### Task 4: Budget checker — core enforcement logic + +**Files:** +- Create: `internal/budget/budget.go` +- Create: `internal/budget/budget_test.go` + +- [ ] **Step 1: Write the failing test for period calculation** + +```go +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) { + // Before reset day + 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) + + // After reset day + 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) + + // On reset day + 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) { + // Feb has 28 days — day 28 exists + 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) +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./internal/budget/ -run TestCurrentPeriodStart -v` +Expected: FAIL — function does not exist. + +- [ ] **Step 3: Implement CurrentPeriodStart** + +```go +package budget + +import "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 { + 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) + } + // Go back to previous month's reset day + prev := now.AddDate(0, -1, 0) + return time.Date(prev.Year(), prev.Month(), resetDay, 0, 0, 0, 0, time.UTC) + } + // Default: calendar month + return time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `go test ./internal/budget/ -run TestCurrentPeriodStart -v` +Expected: PASS + +- [ ] **Step 5: Write the failing test for budget enforcement** + +```go +func TestChecker_Check_GlobalHardBlock(t *testing.T) { + checker := &budget.Checker{ + GlobalMonthlyTokenLimit: 1000, + ResetMode: "calendar", + ResetDay: 1, + GetTotalUsage: func(ctx context.Context, teamID string, period time.Time) (int64, error) { + return 1001, nil // over limit + }, + } + + 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, teamID string, 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) (*TeamBudget, error) { + return &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, teamID string, 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) (*TeamBudget, error) { + return &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) +} +``` + +- [ ] **Step 6: Run test to verify it fails** + +Run: `go test ./internal/budget/ -run TestChecker_Check -v` +Expected: FAIL — Checker type does not exist. + +- [ ] **Step 7: Implement the Checker** + +```go +// 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 // true if the step should be prevented from running + Warning bool // true if budget exceeded but enforcement is "warn" + BlockedBy string // "global", "team", or "" + Message string // human-readable explanation +} + +// 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 + + // DB access functions — injected by caller + GetTotalUsage func(ctx context.Context, teamID string, 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, input.TeamID, 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 && c.DefaultTeamMonthlyTokenLimit > 0 { + // No explicit budget — use global default with hard enforcement + 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{} +} +``` + +- [ ] **Step 8: Run tests to verify they pass** + +Run: `go test ./internal/budget/ -run TestChecker_Check -v` +Expected: All PASS + +- [ ] **Step 9: Add test for workflow execution budget check** + +```go +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, 49999) + assert.False(t, result.Blocked) + + // 0 = unlimited + result = budget.CheckExecutionBudget(0, 999999) + assert.False(t, result.Blocked) +} +``` + +- [ ] **Step 10: Implement CheckExecutionBudget** + +```go +// CheckExecutionBudget checks if a workflow execution's cumulative token usage +// has exceeded the workflow-level token_budget. Called before each AI step dispatch. +// 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{} +} +``` + +- [ ] **Step 11: Run all budget tests** + +Run: `go test ./internal/budget/ -v -timeout 120s` +Expected: All PASS + +- [ ] **Step 12: Commit** + +```bash +git add internal/budget/budget.go internal/budget/budget_test.go +git commit -m "feat(budget): add budget checker with period calculation and AND-gate enforcement" +``` + +--- + +### Task 5: Add audit actions and metrics for budget enforcement + +**Files:** +- Modify: `internal/audit/audit.go` +- Modify: `internal/metrics/metrics.go` + +- [ ] **Step 1: Add budget audit action constants** + +In `internal/audit/audit.go`, add to the Action constants: + +```go +ActionBudgetExceeded Action = "budget.exceeded" // hard block +ActionBudgetWarning Action = "budget.warning" // warn-only +ActionBudgetUpdated Action = "budget.updated" // team budget changed +``` + +- [ ] **Step 2: Add budget Prometheus metrics** + +In `internal/metrics/metrics.go`, add: + +```go +BudgetCheckTotal = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "mantle_budget_check_total", + Help: "Total budget checks performed before AI step dispatch", +}, []string{"team_id", "provider", "result"}) // result: "pass", "blocked", "warning" + +BudgetUsageGauge = promauto.NewGaugeVec(prometheus.GaugeOpts{ + Name: "mantle_budget_usage_tokens", + Help: "Current token usage within the budget period", +}, []string{"team_id", "provider"}) +``` + +- [ ] **Step 3: Verify build** + +Run: `go build ./...` +Expected: exit 0 + +- [ ] **Step 4: Commit** + +```bash +git add internal/audit/audit.go internal/metrics/metrics.go +git commit -m "feat(budget): add budget audit actions and Prometheus metrics" +``` + +--- + +### Task 6: Add token_budget field to Workflow struct + +**Files:** +- Modify: `internal/workflow/workflow.go` +- Modify: `internal/workflow/workflow_test.go` + +- [ ] **Step 1: Write the failing test** + +```go +func TestParse_TokenBudget(t *testing.T) { + yaml := ` +name: test-workflow +token_budget: 500000 +steps: + - name: step1 + action: http/request + params: + url: "https://example.com" +` + result, err := workflow.ParseBytes([]byte(yaml)) + require.NoError(t, err) + assert.Equal(t, int64(500000), result.Workflow.TokenBudget) +} + +func TestParse_TokenBudget_Zero(t *testing.T) { + yaml := ` +name: test-workflow +steps: + - name: step1 + action: http/request + params: + url: "https://example.com" +` + result, err := workflow.ParseBytes([]byte(yaml)) + require.NoError(t, err) + assert.Equal(t, int64(0), result.Workflow.TokenBudget) +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./internal/workflow/ -run TestParse_TokenBudget -v` +Expected: FAIL — `TokenBudget` field does not exist on Workflow. + +- [ ] **Step 3: Add TokenBudget to Workflow struct** + +In `internal/workflow/workflow.go`, add to the `Workflow` struct: + +```go +TokenBudget int64 `yaml:"token_budget"` // 0 = unlimited +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `go test ./internal/workflow/ -run TestParse_TokenBudget -v` +Expected: PASS + +- [ ] **Step 5: Run all workflow tests to check for regressions** + +Run: `go test ./internal/workflow/ -v` +Expected: All PASS + +- [ ] **Step 6: Commit** + +```bash +git add internal/workflow/workflow.go internal/workflow/workflow_test.go +git commit -m "feat(budget): add token_budget field to Workflow struct" +``` + +--- + +### Task 7: Wire budget enforcement into the engine + +**Files:** +- Modify: `internal/engine/engine.go` +- Modify: `internal/engine/engine_test.go` + +This is the critical integration task. The engine needs to: +1. Check team+provider and global budgets before each AI step +2. Check workflow execution budget before each AI step +3. Record token usage after each AI step completes +4. Emit audit events and metrics for budget violations + +**Important architecture note:** Budget enforcement goes in `executeStepLogic()` (not `executeStep()`), because `executeStepLogic` is the shared code path used by BOTH the sequential engine (`executeStep` → `executeStepLogic`) AND the distributed worker (`MakeStepExecutor`/`MakeGlobalStepExecutor` → `executeStepLogic`). Placing checks in `executeStep` alone would leave the distributed worker path unprotected. + +The workflow execution budget check requires access to completed steps and the workflow's `TokenBudget`. In the distributed path (`MakeGlobalStepExecutor`), these are loaded from DB at lines 424-446 of `engine.go`. We'll pass them to `executeStepLogic` via a new `StepContext` parameter. + +- [ ] **Step 1: Write the failing test for pre-dispatch budget blocking** + +In `internal/engine/engine_test.go`, write a test that sets up a workflow with `token_budget` and verifies the engine blocks an AI step when the budget is exhausted. The exact test structure depends on the existing test helpers — examine `engine_test.go` for patterns. + +The test should: +- Create a workflow with `token_budget: 100` +- Mock the budget checker to return `Blocked: true` +- Execute and assert the step result is `"failed"` with a budget error message + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./internal/engine/ -run TestEngine_BudgetBlock -v -timeout 120s` +Expected: FAIL + +- [ ] **Step 3: Add BudgetChecker and BudgetStore to Engine struct** + +In `internal/engine/engine.go`, add fields to the `Engine` struct: + +```go +type Engine struct { + // ... existing fields ... + BudgetChecker *budget.Checker // nil = budget enforcement disabled + BudgetStore *budget.Store // nil = token usage recording disabled +} +``` + +- [ ] **Step 4: Add a StepContext struct to pass workflow-level data to executeStepLogic** + +The current `executeStepLogic` signature is: +```go +func (e *Engine) executeStepLogic(ctx context.Context, execID string, step workflow.Step, celCtx *mantleCEL.Context, workflowName string) (map[string]any, error) +``` + +Add a `StepContext` struct and update the signature: + +```go +// StepContext carries workflow-level metadata needed by step execution. +// Fields are optional — zero values disable the corresponding feature. +type StepContext struct { + WorkflowTokenBudget int64 // from workflow.TokenBudget; 0 = unlimited + CompletedSteps map[string]map[string]any // outputs of completed steps (for summing tokens) + TeamID string // from auth.TeamIDFromContext +} +``` + +Update `executeStepLogic` to accept `StepContext`: +```go +func (e *Engine) executeStepLogic(ctx context.Context, execID string, step workflow.Step, celCtx *mantleCEL.Context, workflowName string, sc StepContext) (map[string]any, error) +``` + +Update all callers of `executeStepLogic`: + +1. **`executeStep`** (line 235): Build `StepContext` from the `resumeExecution` caller. `resumeExecution` has access to `completedSteps` and the `Workflow` struct — pass `wf.TokenBudget`, `completedSteps`, and `auth.TeamIDFromContext(ctx)` through `executeStep`'s signature. + +2. **`MakeStepExecutor`** (line 377): The caller has `wf *workflow.Workflow` and `celCtx`. Build `StepContext{WorkflowTokenBudget: wf.TokenBudget, TeamID: auth.TeamIDFromContext(ctx)}`. For `CompletedSteps`, extract from `celCtx.Steps` (each entry has `{"output": ...}`). + +3. **`MakeGlobalStepExecutor`** (line 448): Already loads `completedSteps` at line 424 and `wf` at line 413. Build `StepContext{WorkflowTokenBudget: wf.TokenBudget, CompletedSteps: completedSteps, TeamID: teamID}` (teamID is at line 401). + +- [ ] **Step 5: Add pre-dispatch budget check at the top of executeStepLogic** + +In `executeStepLogic()`, before the existing CEL resolution and connector dispatch, add: + +```go +// Budget enforcement — check before dispatching AI steps. +if strings.HasPrefix(step.Action, "ai/") { + provider := "openai" // default + if p, ok := step.Params["provider"].(string); ok && p != "" { + provider = p + } + + // 1. Workflow execution budget: sum tokens from completed steps + if sc.WorkflowTokenBudget > 0 { + var usedTokens int64 + for _, stepOutput := range sc.CompletedSteps { + usedTokens += extractTotalTokens(stepOutput) + } + result := budget.CheckExecutionBudget(sc.WorkflowTokenBudget, usedTokens) + if result.Blocked { + return nil, fmt.Errorf("%s", result.Message) + } + } + + // 2. Team+provider and global budget check + if e.BudgetChecker != nil && sc.TeamID != "" { + result := e.BudgetChecker.Check(ctx, budget.CheckInput{ + TeamID: sc.TeamID, + Provider: provider, + }) + if result.Blocked { + e.Auditor.Emit(ctx, audit.Event{ + Timestamp: time.Now(), + Actor: "engine", + Action: audit.ActionBudgetExceeded, + Resource: audit.Resource{Type: "workflow_execution", ID: execID}, + Metadata: map[string]string{ + "blocked_by": result.BlockedBy, + "message": result.Message, + "provider": provider, + }, + TeamID: sc.TeamID, + }) + metrics.BudgetCheckTotal.WithLabelValues(sc.TeamID, provider, "blocked").Inc() + return nil, fmt.Errorf("%s", result.Message) + } + if result.Warning { + e.Auditor.Emit(ctx, audit.Event{ + Timestamp: time.Now(), + Actor: "engine", + Action: audit.ActionBudgetWarning, + Resource: audit.Resource{Type: "workflow_execution", ID: execID}, + Metadata: map[string]string{ + "message": result.Message, + "provider": provider, + }, + TeamID: sc.TeamID, + }) + metrics.BudgetCheckTotal.WithLabelValues(sc.TeamID, provider, "warning").Inc() + } else { + metrics.BudgetCheckTotal.WithLabelValues(sc.TeamID, provider, "pass").Inc() + } + } +} +``` + +- [ ] **Step 6: Add post-completion token recording** + +After a successful connector execution in `executeStepLogic` (after `output, lastErr = conn.Execute(...)` succeeds), add token recording: + +```go +// Record token usage for budget tracking after successful AI step +if strings.HasPrefix(step.Action, "ai/") && e.BudgetStore != nil && e.BudgetChecker != nil && sc.TeamID != "" && output != nil { + promptTokens, completionTokens := extractTokenCounts(output) + if promptTokens > 0 || completionTokens > 0 { + provider := "openai" + if p, ok := step.Params["provider"].(string); ok && p != "" { + provider = p + } + model, _ := output["model"].(string) + period := budget.CurrentPeriodStart(time.Now(), e.BudgetChecker.ResetMode, e.BudgetChecker.ResetDay) + _ = e.BudgetStore.IncrementUsage(ctx, sc.TeamID, provider, model, period, promptTokens, completionTokens) + } +} +``` + +- [ ] **Step 7: Add helper functions to extract token counts from step output** + +```go +// extractTokenCounts pulls prompt and completion token counts from an AI step's output map. +func extractTokenCounts(output map[string]any) (prompt, completion int64) { + usage, ok := output["usage"].(map[string]any) + if !ok { + return 0, 0 + } + if v, ok := usage["prompt_tokens"]; ok { + prompt = toInt64(v) + } + if v, ok := usage["completion_tokens"]; ok { + completion = toInt64(v) + } + return prompt, completion +} + +// extractTotalTokens pulls the total token count from a step's output map. +func extractTotalTokens(output map[string]any) int64 { + usage, ok := output["usage"].(map[string]any) + if !ok { + return 0 + } + return toInt64(usage["total_tokens"]) +} + +func toInt64(v any) int64 { + switch n := v.(type) { + case int: + return int64(n) + case int64: + return n + case float64: + return int64(n) + default: + return 0 + } +} +``` + +- [ ] **Step 8: Run tests to verify they pass** + +Run: `go test ./internal/engine/ -v -timeout 120s` +Expected: All PASS (including new budget test) + +- [ ] **Step 9: Run full test suite to check for regressions** + +Run: `go test ./... -timeout 300s` +Expected: All PASS + +- [ ] **Step 10: Commit** + +```bash +git add internal/engine/engine.go internal/engine/engine_test.go +git commit -m "feat(budget): wire budget enforcement into engine step dispatch" +``` + +--- + +### Task 8: API endpoints for team budget management + +**Files:** +- Modify: `internal/server/api.go` +- Modify: `internal/server/server.go` + +- [ ] **Step 1: Register new routes** + +In `internal/server/server.go`, add route registrations: + +```go +mux.HandleFunc("GET /api/v1/budgets", s.handleListBudgets) +mux.HandleFunc("PUT /api/v1/budgets/{provider}", s.handleSetBudget) +mux.HandleFunc("DELETE /api/v1/budgets/{provider}", s.handleDeleteBudget) +mux.HandleFunc("GET /api/v1/budgets/usage", s.handleGetUsage) +``` + +- [ ] **Step 2: Implement handlers** + +In `internal/server/api.go`, add: + +```go +func (s *Server) handleListBudgets(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + budgets, err := s.BudgetStore.ListTeamBudgets(r.Context(), teamID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + writeJSON(w, http.StatusOK, budgets) +} + +func (s *Server) handleSetBudget(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + provider := r.PathValue("provider") + + var body struct { + MonthlyTokenLimit int64 `json:"monthly_token_limit"` + Enforcement string `json:"enforcement"` // "hard" or "warn" + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, "invalid request body", http.StatusBadRequest) + return + } + if body.Enforcement == "" { + body.Enforcement = "hard" + } + if body.Enforcement != "hard" && body.Enforcement != "warn" { + http.Error(w, "enforcement must be 'hard' or 'warn'", http.StatusBadRequest) + return + } + + if err := s.BudgetStore.SetTeamBudget(r.Context(), teamID, provider, body.MonthlyTokenLimit, body.Enforcement); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + // Build actor string from authenticated user (auth.UserFromContext returns *User or nil) + actor := "api" + if u := auth.UserFromContext(r.Context()); u != nil { + actor = u.ID + } + + s.Auditor.Emit(r.Context(), audit.Event{ + Timestamp: time.Now(), + Actor: actor, + Action: audit.ActionBudgetUpdated, + Resource: audit.Resource{Type: "team_budget", ID: teamID}, + Metadata: map[string]string{ + "provider": provider, + "limit": fmt.Sprintf("%d", body.MonthlyTokenLimit), + "enforcement": body.Enforcement, + }, + TeamID: teamID, + }) + + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) +} + +func (s *Server) handleDeleteBudget(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + provider := r.PathValue("provider") + + if err := s.BudgetStore.DeleteTeamBudget(r.Context(), teamID, provider); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) +} + +func (s *Server) handleGetUsage(w http.ResponseWriter, r *http.Request) { + teamID := auth.TeamIDFromContext(r.Context()) + cfg := config.FromContext(r.Context()) + + period := budget.CurrentPeriodStart(time.Now(), cfg.Engine.Budget.ResetMode, cfg.Engine.Budget.ResetDay) + + provider := r.URL.Query().Get("provider") + var usage *budget.ProviderUsage + var err error + if provider != "" { + usage, err = s.BudgetStore.GetUsage(r.Context(), teamID, provider, period) + } else { + usage, err = s.BudgetStore.GetTotalUsage(r.Context(), teamID, period) + } + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + writeJSON(w, http.StatusOK, map[string]any{ + "period_start": period.Format("2006-01-02"), + "provider": usage.Provider, + "prompt_tokens": usage.PromptTokens, + "completion_tokens": usage.CompletionTokens, + "total_tokens": usage.TotalTokens, + }) +} +``` + +- [ ] **Step 3: Add BudgetStore to Server struct** + +The `Server` struct is defined in `internal/server/server.go` at line 25. Add `BudgetStore *budget.Store` as a new field. The `Server` struct uses direct field assignment (not a constructor function), so add initialization wherever the server is instantiated — look for `&Server{` or `server.Server{` in `internal/cli/serve.go` or similar and add `BudgetStore: budget.NewStore(db)` there. + +- [ ] **Step 4: Verify build** + +Run: `go build ./...` +Expected: exit 0 + +- [ ] **Step 5: Commit** + +```bash +git add internal/server/api.go internal/server/server.go +git commit -m "feat(budget): add team budget API endpoints (list, set, delete, usage)" +``` + +--- + +### Task 9: CLI commands for budget management + +**Files:** +- Create: `internal/cli/budget.go` +- Modify: `internal/cli/root.go` + +- [ ] **Step 1: Implement the budget subcommand** + +**Important:** The existing CLI commands (e.g., `secrets`, `apply`) use direct DB access, NOT HTTP API calls. There is no `apiRequest` helper. Follow the same pattern as `internal/cli/secrets.go`: open a DB connection via `db.Open(cfg.Database)`, create a store, and operate directly. + +Create `internal/cli/budget.go`: + +```go +package cli + +import ( + "fmt" + "text/tabwriter" + "time" + + "github.com/dvflw/mantle/internal/budget" + "github.com/dvflw/mantle/internal/config" + "github.com/dvflw/mantle/internal/db" + "github.com/dvflw/mantle/internal/auth" + "github.com/spf13/cobra" +) + +func newBudgetCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "budget", + Short: "Manage AI token budgets", + } + cmd.AddCommand( + newBudgetSetCommand(), + newBudgetGetCommand(), + newBudgetUsageCommand(), + newBudgetDeleteCommand(), + ) + return cmd +} + +// newBudgetStore opens a DB connection and returns a budget store (mirrors newSecretStore pattern). +func newBudgetStore(cmd *cobra.Command) (*budget.Store, func(), error) { + cfg := config.FromContext(cmd.Context()) + database, err := db.Open(cfg.Database) + if err != nil { + return nil, nil, fmt.Errorf("opening database: %w", err) + } + return budget.NewStore(database), func() { database.Close() }, nil +} + +func newBudgetSetCommand() *cobra.Command { + var enforcement string + cmd := &cobra.Command{ + Use: "set ", + Short: "Set a monthly token budget for a provider", + Long: "Set a monthly token budget. Provider can be 'openai', 'bedrock', or '*' for all providers.", + Args: cobra.ExactArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + provider := args[0] + var limit int64 + if _, err := fmt.Sscanf(args[1], "%d", &limit); err != nil { + return fmt.Errorf("invalid token limit: %s", args[1]) + } + + teamID := auth.TeamIDFromContext(cmd.Context()) + if err := store.SetTeamBudget(cmd.Context(), teamID, provider, limit, enforcement); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "Budget set: %s → %d tokens/month (enforcement: %s)\n", provider, limit, enforcement) + return nil + }, + } + cmd.Flags().StringVar(&enforcement, "enforcement", "hard", "Enforcement mode: 'hard' (block) or 'warn' (log only)") + return cmd +} + +func newBudgetGetCommand() *cobra.Command { + return &cobra.Command{ + Use: "get", + Short: "List all team budgets", + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + teamID := auth.TeamIDFromContext(cmd.Context()) + budgets, err := store.ListTeamBudgets(cmd.Context(), teamID) + if err != nil { + return err + } + if len(budgets) == 0 { + fmt.Fprintln(cmd.OutOrStdout(), "No budgets configured (global defaults apply)") + return nil + } + w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0) + fmt.Fprintln(w, "PROVIDER\tLIMIT\tENFORCEMENT") + for _, b := range budgets { + fmt.Fprintf(w, "%s\t%d tokens/month\t%s\n", b.Provider, b.MonthlyTokenLimit, b.Enforcement) + } + return w.Flush() + }, + } +} + +func newBudgetUsageCommand() *cobra.Command { + var provider string + cmd := &cobra.Command{ + Use: "usage", + Short: "Show current period token usage", + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + cfg := config.FromContext(cmd.Context()) + teamID := auth.TeamIDFromContext(cmd.Context()) + period := budget.CurrentPeriodStart(time.Now(), cfg.Engine.Budget.ResetMode, cfg.Engine.Budget.ResetDay) + + var usage *budget.ProviderUsage + if provider != "" { + usage, err = store.GetUsage(cmd.Context(), teamID, provider, period) + } else { + usage, err = store.GetTotalUsage(cmd.Context(), teamID, period) + } + if err != nil { + return err + } + + fmt.Fprintf(cmd.OutOrStdout(), "Period: %s\n", period.Format("2006-01-02")) + fmt.Fprintf(cmd.OutOrStdout(), "Provider: %s\n", usage.Provider) + fmt.Fprintf(cmd.OutOrStdout(), "Prompt: %d tokens\n", usage.PromptTokens) + fmt.Fprintf(cmd.OutOrStdout(), "Completion: %d tokens\n", usage.CompletionTokens) + fmt.Fprintf(cmd.OutOrStdout(), "Total: %d tokens\n", usage.TotalTokens) + return nil + }, + } + cmd.Flags().StringVar(&provider, "provider", "", "Filter by provider (e.g., 'openai', 'bedrock')") + return cmd +} + +func newBudgetDeleteCommand() *cobra.Command { + return &cobra.Command{ + Use: "delete ", + Short: "Remove a team budget for a provider", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + store, cleanup, err := newBudgetStore(cmd) + if err != nil { + return err + } + defer cleanup() + + teamID := auth.TeamIDFromContext(cmd.Context()) + if err := store.DeleteTeamBudget(cmd.Context(), teamID, args[0]); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "Budget removed for provider: %s\n", args[0]) + return nil + }, + } +} +``` + +- [ ] **Step 2: Register in root command** + +In `internal/cli/root.go`, add the budget command to the admin group: + +```go +addToGroup(cmd, "admin", + // ... existing admin commands ... + newBudgetCommand(), +) +``` + +- [ ] **Step 3: Verify build** + +Run: `go build ./...` +Expected: exit 0 + +- [ ] **Step 4: Commit** + +```bash +git add internal/cli/budget.go internal/cli/root.go +git commit -m "feat(budget): add mantle budget CLI commands (set, get, usage, delete)" +``` + +--- + +### Task 10: Documentation + +**Files:** +- Create: `site/src/content/docs/ai-cost-controls.md` + +- [ ] **Step 1: Write the docs page** + +Create `site/src/content/docs/ai-cost-controls.md` with the following content: + +```mdx +--- +title: AI Cost Controls +description: Configure token budgets to control AI usage at the workflow, team, and global level. +--- + +# AI Cost Controls + +Mantle provides token-based budgets at three levels to help you control AI usage costs. Budgets are denominated in **tokens** (not dollars) to stay provider-agnostic and avoid stale pricing data. + +## Budget Levels + +All three budget levels compose as **AND gates** — every applicable level must pass before an AI step is dispatched. If any level is exceeded, the step is blocked (or warned, depending on configuration). + +| Level | Scope | Configured In | Enforcement | Reset | +|-------|-------|---------------|-------------|-------| +| **Global** | All teams, all providers | `mantle.yaml` | Hard block only | Monthly | +| **Team + Provider** | One team, one provider | API / CLI | Configurable (hard or warn) | Monthly | +| **Workflow** | Single execution | Workflow YAML | Configurable (hard or warn) | Per execution | + +### Enforcement Behavior + +- **Hard block (default):** The AI step fails with a budget error. The execution stops at that step. Tokens consumed by prior steps in the execution are preserved — Mantle never cancels a step mid-execution. +- **Warn only:** The AI step proceeds, but a warning is logged and an audit event is emitted. The `mantle_budget_check_total` Prometheus metric is incremented with `result="warning"`. + +## Configuration + +### Global Budget (mantle.yaml) + +```yaml +engine: + budget: + global_monthly_token_limit: 10000000 # 0 = unlimited + default_team_monthly_token_limit: 1000000 # applied to teams without explicit budgets + reset_mode: calendar # "calendar" or "rolling" + reset_day: 1 # 1-28, only used when reset_mode is "rolling" +``` + +**Environment variables:** + +| Variable | Description | +|----------|-------------| +| `MANTLE_ENGINE_BUDGET_GLOBAL_MONTHLY_TOKEN_LIMIT` | Global token cap (hard block) | +| `MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT` | Default per-team cap | +| `MANTLE_ENGINE_BUDGET_RESET_MODE` | `calendar` or `rolling` | +| `MANTLE_ENGINE_BUDGET_RESET_DAY` | Start day for rolling periods (1-28) | + +### Reset Windows + +- **Calendar month:** Budget resets on the 1st of each month (UTC). +- **Rolling period:** Budget resets on your configured `reset_day` each month. For example, `reset_day: 15` means the current period runs from the 15th of the current month to the 14th of the next. The maximum `reset_day` is 28 to avoid February edge cases. + +### Team + Provider Budget (API / CLI) + +Set per-provider budgets for your team: + +```bash +# Set a hard budget of 1M tokens/month for OpenAI +mantle budget set openai 1000000 + +# Set a warn-only budget for Bedrock +mantle budget set bedrock 500000 --enforcement warn + +# Set a budget for all providers +mantle budget set '*' 2000000 + +# View current budgets +mantle budget get + +# View current usage +mantle budget usage +mantle budget usage --provider openai + +# Remove a budget +mantle budget delete openai +``` + +**API:** + +```bash +# Set budget +curl -X PUT /api/v1/budgets/openai \ + -d '{"monthly_token_limit": 1000000, "enforcement": "hard"}' + +# List budgets +curl /api/v1/budgets + +# Get usage +curl /api/v1/budgets/usage?provider=openai + +# Delete budget +curl -X DELETE /api/v1/budgets/openai +``` + +### Workflow Budget (YAML) + +Add a `token_budget` field at the workflow level to cap total tokens consumed in a single execution: + +```yaml +name: my-analysis-workflow +token_budget: 500000 # max tokens across all AI steps in one execution + +steps: + - name: summarize + action: ai/completion + params: + model: gpt-4o + prompt: "Summarize this document: {{ inputs.document }}" + max_token_budget: 100000 # per-step limit (existing feature) + + - name: analyze + action: ai/completion + params: + model: gpt-4o + prompt: "Analyze the summary: {{ steps.summarize.output.text }}" +``` + +The workflow `token_budget` is checked before each AI step. If the cumulative tokens from prior steps exceed the budget, the next AI step is blocked. Non-AI steps (HTTP, etc.) are unaffected. + +The existing per-step `max_token_budget` param continues to work independently — it caps tokens within a single step's tool-use loop. + +## Monitoring + +### Prometheus Metrics + +| Metric | Labels | Description | +|--------|--------|-------------| +| `mantle_budget_check_total` | `team_id`, `provider`, `result` | Budget checks (result: pass/blocked/warning) | +| `mantle_budget_usage_tokens` | `team_id`, `provider` | Current token usage in the budget period | +| `mantle_ai_tokens_total` | `workflow`, `step`, `model`, `provider`, `token_type` | Raw token consumption (existing) | + +### Audit Events + +| Action | When | +|--------|------| +| `budget.exceeded` | AI step blocked by budget | +| `budget.warning` | AI step proceeded with warn-only budget exceeded | +| `budget.updated` | Team budget created, updated, or deleted | + +## Provider Pricing Reference + +Token budgets are provider-agnostic. To estimate costs, refer to your provider's pricing page: + +- **OpenAI:** [https://openai.com/api/pricing](https://openai.com/api/pricing) +- **AWS Bedrock:** [https://aws.amazon.com/bedrock/pricing](https://aws.amazon.com/bedrock/pricing) +- **Anthropic:** [https://www.anthropic.com/pricing](https://www.anthropic.com/pricing) +- **Google AI:** [https://ai.google.dev/pricing](https://ai.google.dev/pricing) + +> **Tip:** Most providers charge differently for prompt (input) vs. completion (output) tokens. Mantle tracks both separately in the `ai_token_usage` table and Prometheus metrics, so you can compute costs externally using your provider's rate card. +``` + +- [ ] **Step 2: Verify the site builds (if applicable)** + +Run: `cd site && npm run build` (or whatever the site build command is) +Expected: Build succeeds with the new docs page. + +- [ ] **Step 3: Commit** + +```bash +git add site/src/content/docs/ai-cost-controls.md +git commit -m "docs: add AI cost controls documentation with provider pricing links" +``` + +--- + +## Verification Checklist + +After all tasks are complete, verify end-to-end: + +- [ ] `make migrate` applies migration 012 cleanly +- [ ] `go build ./...` succeeds +- [ ] `go test ./... -timeout 300s` all pass +- [ ] `go vet ./...` clean +- [ ] Config loads budget defaults correctly +- [ ] Workflow YAML with `token_budget` parses correctly +- [ ] Budget store increments and queries work against real Postgres +- [ ] Engine blocks AI step when global budget exceeded +- [ ] Engine blocks AI step when team budget exceeded (hard mode) +- [ ] Engine warns but continues when team budget exceeded (warn mode) +- [ ] Engine blocks AI step when workflow execution budget exceeded +- [ ] API endpoints return correct data +- [ ] CLI commands work against running server +- [ ] Audit events emitted for budget violations +- [ ] Prometheus metrics recorded for budget checks +- [ ] Docs page builds and renders correctly diff --git a/site/src/content/docs/ai-cost-controls.md b/site/src/content/docs/ai-cost-controls.md new file mode 100644 index 0000000..f371a76 --- /dev/null +++ b/site/src/content/docs/ai-cost-controls.md @@ -0,0 +1,148 @@ +--- +title: AI Cost Controls +description: Configure token budgets to control AI usage at the workflow, team, and global level. +--- + +# AI Cost Controls + +Mantle provides token-based budgets at three levels to help you control AI usage costs. Budgets are denominated in **tokens** (not dollars) to stay provider-agnostic and avoid stale pricing data. + +## Budget Levels + +All three budget levels compose as **AND gates** — every applicable level must pass before an AI step is dispatched. If any level is exceeded, the step is blocked (or warned, depending on configuration). + +| Level | Scope | Configured In | Enforcement | Reset | +|-------|-------|---------------|-------------|-------| +| **Global** | All teams, all providers | `mantle.yaml` | Hard block only | Monthly | +| **Team + Provider** | One team, one provider | API / CLI | Configurable (hard or warn) | Monthly | +| **Workflow** | Single execution | Workflow YAML | Configurable (hard or warn) | Per execution | + +### Enforcement Behavior + +- **Hard block (default):** The AI step fails with a budget error. The execution stops at that step. Tokens consumed by prior steps in the execution are preserved — Mantle never cancels a step mid-execution. +- **Warn only:** The AI step proceeds, but a warning is logged and an audit event is emitted. The `mantle_budget_check_total` Prometheus metric is incremented with `result="warning"`. + +## Configuration + +### Global Budget (mantle.yaml) + +```yaml +engine: + budget: + global_monthly_token_limit: 10000000 # 0 = unlimited + default_team_monthly_token_limit: 1000000 # applied to teams without explicit budgets + reset_mode: calendar # "calendar" or "rolling" + reset_day: 1 # 1-28, only used when reset_mode is "rolling" +``` + +**Environment variables:** + +| Variable | Description | +|----------|-------------| +| `MANTLE_ENGINE_BUDGET_GLOBAL_MONTHLY_TOKEN_LIMIT` | Global token cap (hard block) | +| `MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT` | Default per-team cap | +| `MANTLE_ENGINE_BUDGET_RESET_MODE` | `calendar` or `rolling` | +| `MANTLE_ENGINE_BUDGET_RESET_DAY` | Start day for rolling periods (1-28) | + +### Reset Windows + +- **Calendar month:** Budget resets on the 1st of each month (UTC). +- **Rolling period:** Budget resets on your configured `reset_day` each month. For example, `reset_day: 15` means the current period runs from the 15th of the current month to the 14th of the next. The maximum `reset_day` is 28 to avoid February edge cases. + +### Team + Provider Budget (API / CLI) + +Set per-provider budgets for your team: + +```bash +# Set a hard budget of 1M tokens/month for OpenAI +mantle budget set openai 1000000 + +# Set a warn-only budget for Bedrock +mantle budget set bedrock 500000 --enforcement warn + +# Set a budget for all providers +mantle budget set '*' 2000000 + +# View current budgets +mantle budget get + +# View current usage +mantle budget usage +mantle budget usage --provider openai + +# Remove a budget +mantle budget delete openai +``` + +**API:** + +```bash +# Set budget +curl -X PUT /api/v1/budgets/openai \ + -d '{"monthly_token_limit": 1000000, "enforcement": "hard"}' + +# List budgets +curl /api/v1/budgets + +# Get usage +curl /api/v1/budgets/usage?provider=openai + +# Delete budget +curl -X DELETE /api/v1/budgets/openai +``` + +### Workflow Budget (YAML) + +Add a `token_budget` field at the workflow level to cap total tokens consumed in a single execution: + +```yaml +name: my-analysis-workflow +token_budget: 500000 # max tokens across all AI steps in one execution + +steps: + - name: summarize + action: ai/completion + params: + model: gpt-4o + prompt: "Summarize this document: {{ inputs.document }}" + max_token_budget: 100000 # per-step limit (existing feature) + + - name: analyze + action: ai/completion + params: + model: gpt-4o + prompt: "Analyze the summary: {{ steps.summarize.output.text }}" +``` + +The workflow `token_budget` is checked before each AI step. If the cumulative tokens from prior steps exceed the budget, the next AI step is blocked. Non-AI steps (HTTP, etc.) are unaffected. + +The existing per-step `max_token_budget` param continues to work independently — it caps tokens within a single step's tool-use loop. + +## Monitoring + +### Prometheus Metrics + +| Metric | Labels | Description | +|--------|--------|-------------| +| `mantle_budget_check_total` | `team_id`, `provider`, `result` | Budget checks (result: pass/blocked/warning) | +| `mantle_budget_usage_tokens` | `team_id`, `provider` | Current token usage in the budget period | +| `mantle_ai_tokens_total` | `workflow`, `step`, `model`, `provider`, `token_type` | Raw token consumption (existing) | + +### Audit Events + +| Action | When | +|--------|------| +| `budget.exceeded` | AI step blocked by budget | +| `budget.warning` | AI step proceeded with warn-only budget exceeded | +| `budget.updated` | Team budget created, updated, or deleted | + +## Provider Pricing Reference + +Token budgets are provider-agnostic. To estimate costs, refer to your provider's pricing page: + +- **OpenAI:** [https://openai.com/api/pricing](https://openai.com/api/pricing) +- **AWS Bedrock:** [https://aws.amazon.com/bedrock/pricing](https://aws.amazon.com/bedrock/pricing) +- **Anthropic:** [https://www.anthropic.com/pricing](https://www.anthropic.com/pricing) +- **Google AI:** [https://ai.google.dev/pricing](https://ai.google.dev/pricing) + +> **Tip:** Most providers charge differently for prompt (input) vs. completion (output) tokens. Mantle tracks both separately in the `ai_token_usage` table and Prometheus metrics, so you can compute costs externally using your provider's rate card.