feat: AI cost controls — per-provider, per-team, per-workflow token budgets#9
Merged
Conversation
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…achability Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Design decisions from /grill-me session: - Token-based budgets (not dollar-based) - Three AND-gated levels: workflow, team+provider, global - Configurable enforcement (hard block vs warn-only) - Calendar or rolling reset windows - Counter table for fast enforcement queries Ref #8 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Adds BudgetConfig struct with ResetMode, ResetDay, GlobalMonthlyTokenLimit, and DefaultTeamMonthlyTokenLimit fields, nested under EngineConfig. Includes defaults, env var bindings, and a test asserting correct default values. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
… enforcement Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add pre-execution budget checks and post-execution token recording to executeStepLogic, covering both sequential and distributed execution paths. AI steps are checked against workflow token_budget, team+provider limits, and global limits before dispatch. Token usage is recorded to the budget store after successful completion. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
BudgetStore recording should work independently of BudgetChecker — allows token tracking for observability even when enforcement is disabled. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…cording errors Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a token-based AI cost-control system: period computation, budget checks (workflow, team+provider, global), Postgres-backed usage and budget storage with migrations, CLI/API management, engine pre-dispatch checks and post-success usage recording, audit events, and Prometheus metrics. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Engine as Execution Engine
participant Checker as Budget Checker
participant Store as Budget Store
participant Auditor as Audit Logger
participant Metrics as Metrics
Client->>Engine: Execute workflow with AI step
activate Engine
Engine->>Engine: Build StepContext (workflow budget, team ID)
alt Workflow Budget Check (if budget > 0)
Engine->>Checker: CheckExecutionBudget(completedTokens, limit)
Checker-->>Engine: Result (blocked/pass)
alt Blocked
Engine->>Auditor: Emit `budget.exceeded`
Engine->>Metrics: Increment (result=blocked)
Engine-->>Client: Error (halt)
end
end
alt Team+Provider Check (if Checker configured)
Engine->>Checker: Check(team, provider)
Checker->>Store: GetTeamBudget(team, provider)
Store-->>Checker: TeamBudget
Checker->>Store: GetProviderUsage(team, provider, period)
Store-->>Checker: Usage
alt Hard Block
Checker-->>Engine: Blocked
Engine->>Auditor: Emit `budget.exceeded`
Engine->>Metrics: Increment (result=blocked)
Engine-->>Client: Error (halt)
else Warn Only
Checker-->>Engine: Warning
Engine->>Auditor: Emit `budget.warning`
Engine->>Metrics: Increment (result=warning)
Engine->>Engine: Continue
end
end
alt Global Check (if global limit set)
Checker->>Store: GetGlobalTotalUsage(period)
Store-->>Checker: GlobalUsage
alt Exceeded
Checker-->>Engine: Blocked
Engine->>Auditor: Emit `budget.exceeded`
Engine->>Metrics: Increment (result=blocked)
Engine-->>Client: Error (halt)
end
end
Engine->>Engine: Execute AI step
Engine->>Metrics: Increment (result=pass)
Engine->>Store: IncrementUsage(prompt, completion, period)
Store-->>Engine: Ack
Engine->>Auditor: Emit `budget.updated`
Engine-->>Client: Success
deactivate Engine
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 18
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/budget/budget.go`:
- Around line 45-46: The global monthly token limit is being enforced per-team
because Check calls GetTotalUsage/GetProviderUsage with input.TeamID and the
concrete store aggregates by team_id; change Check to call the "global" usage
callbacks without a team filter (e.g., pass empty string or a dedicated global
key) and adjust the store implementation used for GetTotalUsage/GetProviderUsage
to aggregate across all teams (remove/team_id filter) when evaluating
GlobalMonthlyTokenLimit; ensure the same change is applied for provider-scoped
usage (GetProviderUsage) so the GlobalMonthlyTokenLimit is enforced
platform-wide rather than per-TeamID.
- Around line 57-60: The current logic treats errors from usage lookups as a
pass (fail-open); change it so any error from c.GetTotalUsage (and other usage
lookup calls in this file) causes an immediate non-allowing CheckResult instead
of falling through. Specifically, after calling c.GetTotalUsage(ctx,
input.TeamID, period) check if err != nil and return a CheckResult that denies
dispatch (set Allow to false or populate the CheckResult error field) so
transient read/store failures do not create unlimited-spend windows; apply the
same pattern for the other usage lookup blocks around the GetTotalUsage /
GlobalMonthlyTokenLimit checks.
- Around line 12-20: CurrentPeriodStart reads Year/Month/Day from the incoming
now which may carry a non-UTC location; normalize now to UTC first to ensure
period math uses UTC semantics. Update CurrentPeriodStart to compute using
utcNow := now.UTC() (or similar) and then use
utcNow.Year()/utcNow.Month()/utcNow.Day() and utcNow.AddDate(...) when
constructing the period start times (still passing time.UTC to time.Date) so all
date comparisons and returned times are consistent in UTC.
- Around line 70-77: The code currently stops after an exact-provider lookup and
falls back to DefaultTeamMonthlyTokenLimit without checking a wildcard provider;
modify the logic in the block around GetTeamBudget so that if tb == nil after
calling c.GetTeamBudget(ctx, input.TeamID, input.Provider), you retry with
c.GetTeamBudget(ctx, input.TeamID, "*") and use that result (if non-nil) before
applying c.DefaultTeamMonthlyTokenLimit; keep existing behavior for tb != nil
and tb.MonthlyTokenLimit > 0 and ensure you preserve error handling for the
GetTeamBudget calls.
In `@internal/cli/serve.go`:
- Around line 92-98: The GetTotalUsage callback currently restricts aggregation
to a single team (it calls budgetStore.GetTotalUsage with teamID), which
enforces a per-team ceiling instead of a platform-wide ceiling; update the
implementation to compute a global total instead — either add a new store method
like GetGlobalTotalUsage(ctx, period) that aggregates tokens across all teams
and call that from the GetTotalUsage callback (or change the
budgetStore.GetTotalUsage SQL to remove the WHERE team_id = $1 filter when used
for the global check), or if per-team limits were intended, rename the
budget/variable from "global" to "per-team monthly limit" to avoid confusion.
Ensure references to GetTotalUsage and budgetStore.GetTotalUsage are updated
accordingly.
In `@internal/config/config.go`:
- Around line 95-101: Add explicit validation to ensure BudgetConfig.ResetDay is
within 1..28 (inclusive); when loading or initializing configuration (e.g., in
your config validation function or startup init that constructs BudgetConfig)
check if ResetMode == "rolling" (or always) and if ResetDay < 1 || ResetDay > 28
return a descriptive error (or fail startup) so callers like
budget.CurrentPeriodStart() never receive out-of-range values; reference the
BudgetConfig struct and the budget.CurrentPeriodStart() usage to locate where to
add this validation.
In `@internal/db/migrations/012_ai_cost_controls.sql`:
- Around line 21-32: The enforcement column in team_budgets currently accepts
any TEXT; add a CHECK constraint to enforce only 'hard' or 'warn' values to
ensure integrity—update the CREATE TABLE (or ALTER TABLE) for table team_budgets
to add a constraint (e.g., team_budgets_enforcement_check) that validates
enforcement IN ('hard','warn') and ensure migrations include the constraint name
and a reversible ALTER for rollbacks.
- Around line 5-16: The team foreign key in ai_token_usage (team_id REFERENCES
teams(id)) — and any other tables in this migration that reference teams(id) —
currently lacks ON DELETE behavior; decide whether deletions should cascade or
be restricted and update the FK definitions accordingly (e.g., change REFERENCES
teams(id) to REFERENCES teams(id) ON DELETE CASCADE or REFERENCES teams(id) ON
DELETE RESTRICT) so the DB enforces the desired behavior for team deletions;
apply the same change to any other table in this diff that references teams(id).
In `@internal/engine/engine.go`:
- Around line 307-312: The provider value is being read directly from
step.Params before template/CEL interpolation, causing wrong per-provider budget
checks and usage records; update the code so ResolveParams (or the routine that
performs CEL/template interpolation for step.Params) runs before any reads of
step.Params["provider"] used for the pre-dispatch budget check (the block
guarded by strings.HasPrefix(step.Action, "ai/") that sets provider := "openai")
and before any post-step usage writes—i.e., resolve params first, then extract
provider (falling back to "openai" if empty) and use that resolved provider for
budget enforcement and usage recording.
- Around line 454-474: The token-usage write via BudgetStore.IncrementUsage is
not being audited; after a successful call to e.BudgetStore.IncrementUsage(...)
inside the ai/* branch (the block that computes provider, model, period,
promptTok, completionTok), emit an audit event using the engine's AuditEmitter
(e.AuditEmitter) recording the state change (e.g., event type "ai_token_usage",
team id sc.TeamID, provider, model, period, promptTok, completionTok and any
actor/context from sc). Ensure the audit emission occurs on successful
persistence (and is best-effort like the logging) so all state-changing writes
go through the AuditEmitter interface.
In `@internal/metrics/metrics.go`:
- Around line 126-137: BudgetUsageGauge is registered but never updated; either
remove it or update it from the budget enforcement logic: when computing token
usage for a team/provider (e.g., in your budget check/enforcement function such
as the code path that increments BudgetCheckTotal), call
BudgetUsageGauge.WithLabelValues(teamID, provider).Set(currentTokenUsage) and
also ensure you update or reset it on period rollover or when usage changes; if
you choose not to emit usage at all, delete the BudgetUsageGauge declaration to
avoid an unused metric.
In `@internal/server/api.go`:
- Around line 373-376: The handler is leaking internal error text via http.Error
and returning non-JSON responses; replace these http.Error(...) calls in the
budget endpoints (calls around BudgetStore.ListTeamBudgets and the other listed
branches) with a consistent JSON error response (e.g. write a 500 response body
like {"error":"internal server error"} and set Content-Type: application/json)
and do not expose err.Error() to the client; instead log the actual err
internally (using the server's logger, e.g. s.Logger.Error or log.Printf) for
debugging. Apply the same change for the other failing branches mentioned
(around lines handling ListTeamBudgets and the other budget routes) so all
budget endpoints return generic JSON errors and internal errors are only logged.
In `@internal/server/server.go`:
- Around line 141-145: The GET route for a single provider is missing which
breaks the per-provider API contract; add a provider-scoped GET handler
registration like mux.HandleFunc("GET /api/v1/budgets/{provider}",
s.handleGetBudget) alongside the PUT/DELETE entries, and if the handler function
s.handleGetBudget does not yet exist, implement it to return the budget for the
given {provider} parameter consistent with
s.handleSetBudget/s.handleDeleteBudget behavior and error handling.
- Line 32: The budget route handlers will panic if Server.BudgetStore is nil;
fix by (1) conditionally registering the budget routes (handleListBudgets,
handleSetBudget, handleDeleteBudget, handleGetUsage) only when s.BudgetStore !=
nil during server setup (e.g., in Server.New or the router registration code),
and (2) add a defensive nil check at the top of each handler that returns an
appropriate HTTP error (e.g., 501 or 500) if s.BudgetStore is nil to prevent
panics; alternatively, enforce BudgetStore initialization by validating non-nil
in Server.Start and failing fast if it is missing.
In `@internal/workflow/workflow.go`:
- Line 13: The TokenBudget field currently allows negative values which breaks
enforcement; add a validation check after YAML unmarshalling (or implement
UnmarshalYAML/Validate for the workflow config struct) that verifies TokenBudget
>= 0 and returns an error if it's negative; update the parsing flow to call this
validation so any negative token_budget in the input YAML is rejected at
parse/validation time (reference: TokenBudget in the workflow struct).
In `@site/src/components/DocsSidebar.astro`:
- Around line 174-178: The GitHub external anchor in DocsSidebar.astro (the <a>
element whose visible text is "GitHub") must include rel="noopener" to prevent
the opened page from accessing window.opener; update that anchor to add
rel="noopener" (or rel="noopener noreferrer" if you also want to suppress
referrer) so the external link is secured.
In `@site/src/components/Nav.astro`:
- Around line 54-57: The external GitHub anchor (<a ...> element with class
"bottom-nav-item" in Nav.astro) is missing the security attribute; update that
<a> tag for the GitHub link to include rel="noopener" (e.g., add rel="noopener"
to the existing anchor attributes) so the external link uses the same security
pattern as DocsSidebar.
In `@site/src/styles/global.css`:
- Around line 397-399: Replace the nonstandard rule in the .prose code selector:
remove the deprecated "word-break: break-word" and use the standard
overflow-wrap property instead (e.g., set "overflow-wrap: break-word;" or
"overflow-wrap: anywhere;" depending on desired break behavior) and ensure
word-break is left to its default (or explicitly set to "normal") so code blocks
wrap using the standard CSS property.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: d0c1a96d-f05e-4e5c-8b4e-4ab884c2a7a1
📒 Files selected for processing (24)
internal/audit/audit.gointernal/budget/budget.gointernal/budget/budget_test.gointernal/budget/store.gointernal/budget/store_test.gointernal/cli/budget.gointernal/cli/root.gointernal/cli/serve.gointernal/config/config.gointernal/config/config_test.gointernal/db/migrations/012_ai_cost_controls.sqlinternal/engine/engine.gointernal/metrics/metrics.gointernal/server/api.gointernal/server/server.gointernal/workflow/workflow.gointernal/workflow/workflow_test.goplans/ai-cost-controls.mdsite/src/components/DocsSidebar.astrosite/src/components/Nav.astrosite/src/content/docs/ai-cost-controls.mdsite/src/layouts/Docs.astrosite/src/pages/index.astrosite/src/styles/global.css
Comment on lines
+57
to
+60
| if c.GlobalMonthlyTokenLimit > 0 && c.GetTotalUsage != nil { | ||
| total, err := c.GetTotalUsage(ctx, input.TeamID, period) | ||
| if err == nil && total >= c.GlobalMonthlyTokenLimit { | ||
| return CheckResult{ |
Contributor
There was a problem hiding this comment.
Do not fail open when budget reads error out.
If any of these lookups returns an error, the function falls through to CheckResult{} and dispatch proceeds. For hard budgets, that turns a transient store/read failure into an unlimited-spend window.
Also applies to: 70-79
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/budget/budget.go` around lines 57 - 60, The current logic treats
errors from usage lookups as a pass (fail-open); change it so any error from
c.GetTotalUsage (and other usage lookup calls in this file) causes an immediate
non-allowing CheckResult instead of falling through. Specifically, after calling
c.GetTotalUsage(ctx, input.TeamID, period) check if err != nil and return a
CheckResult that denies dispatch (set Allow to false or populate the CheckResult
error field) so transient read/store failures do not create unlimited-spend
windows; apply the same pattern for the other usage lookup blocks around the
GetTotalUsage / GlobalMonthlyTokenLimit checks.
Comment on lines
+141
to
+145
| // 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) |
Contributor
There was a problem hiding this comment.
Missing GET /api/v1/budgets/{provider} route breaks the per-provider API contract.
Line 142 registers only GET /api/v1/budgets, while PUT/DELETE are provider-scoped. This leaves no provider-scoped GET route.
Suggested fix
// Budget endpoints.
mux.HandleFunc("GET /api/v1/budgets", s.handleListBudgets)
+ mux.HandleFunc("GET /api/v1/budgets/{provider}", s.handleGetBudget)
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)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/server/server.go` around lines 141 - 145, The GET route for a single
provider is missing which breaks the per-provider API contract; add a
provider-scoped GET handler registration like mux.HandleFunc("GET
/api/v1/budgets/{provider}", s.handleGetBudget) alongside the PUT/DELETE
entries, and if the handler function s.handleGetBudget does not yet exist,
implement it to return the budget for the given {provider} parameter consistent
with s.handleSetBudget/s.handleDeleteBudget behavior and error handling.
| 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 |
Contributor
There was a problem hiding this comment.
Reject negative token_budget values at parse/validation time.
Line 13 introduces a budget cap but currently permits negative values, which can produce invalid enforcement semantics.
Suggested fix
diff --git a/internal/workflow/parse.go b/internal/workflow/parse.go
@@
func ParseBytes(data []byte) (*ParseResult, error) {
@@
var w Workflow
if err := root.Decode(&w); err != nil {
return nil, fmt.Errorf("decoding workflow: %w", err)
}
+ if w.TokenBudget < 0 {
+ return nil, fmt.Errorf("decoding workflow: token_budget must be >= 0")
+ }
return &ParseResult{
Workflow: &w,
Root: &root,
}, nil
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/workflow/workflow.go` at line 13, The TokenBudget field currently
allows negative values which breaks enforcement; add a validation check after
YAML unmarshalling (or implement UnmarshalYAML/Validate for the workflow config
struct) that verifies TokenBudget >= 0 and returns an error if it's negative;
update the parsing flow to call this validation so any negative token_budget in
the input YAML is rejected at parse/validation time (reference: TokenBudget in
the workflow struct).
…emove unused metric Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…lback, UTC normalize Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Keep both budget ResetDay validation and sslmode=prefer warning in config.go. Take main's version for all site component files. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/budget/budget_test.go`:
- Around line 83-106: Add tests that assert fail-open behavior when usage/budget
callbacks return errors: create new test(s) (e.g.,
TestChecker_Check_OnGetTotalUsageError and/or
TestChecker_Check_OnGetTeamBudgetError) that construct a budget.Checker with
GetTotalUsage and/or GetTeamBudget returning an error, call checker.Check with a
sample CheckInput, and assert that result.Blocked is false (and Warning behavior
as expected). Use the same setup style as TestChecker_Check_AllPass and
reference Checker.Check, GetTotalUsage, GetTeamBudget, and GetProviderUsage to
locate where to inject the error-returning stubs.
In `@internal/budget/budget.go`:
- Around line 57-67: GetTotalUsage/GetTeamBudget/GetProviderUsage errors are
currently ignored causing checks to fail-open; change each lookup in the
global/team/provider budget checks to fail-closed for hard budgets by returning
a blocking CheckResult when the corresponding call returns an error (i.e., if
err != nil return CheckResult{Blocked:true,
BlockedBy:"global"|"team"|"provider", Message: fmt.Sprintf("error checking
<global|team|provider> budget: %v", err)}), and apply the same pattern used for
the global check to the team budget check (GetTeamBudget) and provider usage
check (GetProviderUsage) so transient DB errors do not permit unlimited spend.
In `@internal/config/config_test.go`:
- Around line 198-208: Add unit tests to cover budget reset_day validation edge
cases alongside TestLoad_BudgetDefaults: create tests that call Load (or helper
newTestCommand) with configuration overrides for
Engine.Budget.ResetMode="rolling" and ResetDay values of 0 and 29 and assert
Load returns an error for those invalid rolling values; also create a test with
Engine.Budget.ResetMode="calendar" and an invalid ResetDay (e.g., 0 or >28) and
assert Load succeeds and cfg.Engine.Budget.ResetDay is clamped to 1. Reference
the existing TestLoad_BudgetDefaults test and the Load function /
cfg.Engine.Budget fields to locate where to add these new tests.
In `@internal/config/config.go`:
- Line 241: The BindEnv call that wires the config key
"engine.budget.default_team_monthly_token_limit" to an environment variable has
a typo in the env name; update the v.BindEnv invocation (the call that binds
"engine.budget.default_team_monthly_token_limit") to use the correct environment
variable MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT so it matches the
DefaultTeamMonthlyTokenLimit field and avoids confusion.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 8972843b-c3a8-42de-af1f-fd973709e336
📒 Files selected for processing (10)
internal/budget/budget.gointernal/budget/budget_test.gointernal/budget/store.gointernal/cli/root.gointernal/cli/serve.gointernal/config/config.gointernal/config/config_test.gointernal/db/migrations/012_ai_cost_controls.sqlinternal/metrics/metrics.gointernal/workflow/validate.go
Comment on lines
+57
to
+67
| // 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), | ||
| } | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
Budget checks fail-open on database errors.
When GetTotalUsage, GetTeamBudget, or GetProviderUsage return errors, the code silently continues and eventually returns CheckResult{} (allow). This means transient database failures grant unlimited spend.
For hard budgets, consider returning a blocking result when lookups fail to avoid an unlimited-spend window during outages.
🛡️ Proposed approach to fail-closed on errors
// 1. Global budget — hard block only
if c.GlobalMonthlyTokenLimit > 0 && c.GetTotalUsage != nil {
total, err := c.GetTotalUsage(ctx, period)
+ if err != nil {
+ return CheckResult{
+ Blocked: true,
+ BlockedBy: "global",
+ Message: fmt.Sprintf("budget check failed: %v", err),
+ }
+ }
- if err == nil && total >= c.GlobalMonthlyTokenLimit {
+ if total >= c.GlobalMonthlyTokenLimit {Apply similar pattern to the team budget check.
Also applies to: 81-97
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/budget/budget.go` around lines 57 - 67,
GetTotalUsage/GetTeamBudget/GetProviderUsage errors are currently ignored
causing checks to fail-open; change each lookup in the global/team/provider
budget checks to fail-closed for hard budgets by returning a blocking
CheckResult when the corresponding call returns an error (i.e., if err != nil
return CheckResult{Blocked:true, BlockedBy:"global"|"team"|"provider", Message:
fmt.Sprintf("error checking <global|team|provider> budget: %v", err)}), and
apply the same pattern used for the global check to the team budget check
(GetTeamBudget) and provider usage check (GetProviderUsage) so transient DB
errors do not permit unlimited spend.
| _ = 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") |
Contributor
There was a problem hiding this comment.
Typo in environment variable name.
The environment variable MANTLE_ENGINE_BUDGET_DEFAULT_TEAMLY_MONTHLY_TOKEN_LIMIT contains a typo ("TEAMLY" instead of "TEAM"). This inconsistency with the field name DefaultTeamMonthlyTokenLimit may cause confusion for operators.
🐛 Proposed fix
- _ = v.BindEnv("engine.budget.default_team_monthly_token_limit", "MANTLE_ENGINE_BUDGET_DEFAULT_TEAMLY_MONTHLY_TOKEN_LIMIT")
+ _ = v.BindEnv("engine.budget.default_team_monthly_token_limit", "MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| _ = v.BindEnv("engine.budget.default_team_monthly_token_limit", "MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT") | |
| _ = v.BindEnv("engine.budget.default_team_monthly_token_limit", "MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT") |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@internal/config/config.go` at line 241, The BindEnv call that wires the
config key "engine.budget.default_team_monthly_token_limit" to an environment
variable has a typo in the env name; update the v.BindEnv invocation (the call
that binds "engine.budget.default_team_monthly_token_limit") to use the correct
environment variable MANTLE_ENGINE_BUDGET_DEFAULT_TEAM_MONTHLY_TOKEN_LIMIT so it
matches the DefaultTeamMonthlyTokenLimit field and avoids confusion.
- Document intentional fail-open: DB errors in budget checks don't block AI steps - Test rolling mode rejects ResetDay outside 1-28 - Test calendar mode silently clamps invalid ResetDay to 1 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/budget/budget_test.go`:
- Around line 109-151: Add a new unit test named
TestChecker_Check_FailOpenOnGetProviderUsageError that constructs a
budget.Checker with ResetMode/ResetDay set, GetTotalUsage returning (0, nil),
GetProviderUsage returning an error (e.g., fmt.Errorf("provider usage query
failed")), and GetTeamBudget returning a valid *budget.TeamBudget; call
checker.Check with a sample CheckInput and assert that result.Blocked and
result.Warning are both false to verify fail-open behavior for GetProviderUsage;
reference the Checker.Check method and the GetProviderUsage, GetTotalUsage, and
GetTeamBudget callbacks when locating where to add the test.
- Around line 39-82: Add a new unit test named TestChecker_Check_TeamHardBlock
that constructs a budget.Checker with ResetMode "calendar" and ResetDay 1, stub
GetTotalUsage to return 0, stub GetProviderUsage to return a usage above the
team limit (e.g., 5001), and stub GetTeamBudget to return
&budget.TeamBudget{MonthlyTokenLimit:5000, Enforcement:"hard"}; call
checker.Check with a team and provider and assert that result.Blocked is true,
result.BlockedBy equals "team", and result.Message contains "team budget". Use
the same test style as TestChecker_Check_TeamWarnOnly and place it alongside the
other TestChecker tests.
In `@internal/config/config_test.go`:
- Around line 228-237: The test TestLoad_BudgetResetDay_CalendarClamps only
checks lower-bound clamping (ResetDay=0 -> 1); add a complementary test that
sets MANTLE_ENGINE_BUDGET_RESET_MODE="calendar" and
MANTLE_ENGINE_BUDGET_RESET_DAY to a value >28 (e.g., "29" or "100"), then call
newTestCommand() and Load(cmd) and assert no error and that
cfg.Engine.Budget.ResetDay equals 28 to verify the upper-bound clamping
implemented in Load/config.go.
- Around line 210-226: Split the two assertions in
TestLoad_BudgetResetDay_RollingInvalid into separate subtests to improve
clarity: use t.Run("reset day 0", func(t *testing.T) { ... }) and t.Run("reset
day 29", func(t *testing.T) { ... }), and inside each subtest set the
appropriate environment variables (via t.Setenv), call newTestCommand() and
Load(cmd), then assert error presence and the expected message; keep the same
checks against Load and newTestCommand but isolate each scenario in its own
t.Run block.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 727f0b9b-c655-4fb5-bd8e-b0db29590303
📒 Files selected for processing (2)
internal/budget/budget_test.gointernal/config/config_test.go
…und clamp tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/budget/budget_test.go`:
- Around line 202-211: Add a boundary test for the exact equality case in
TestChecker_CheckExecutionBudget: call budget.CheckExecutionBudget with
usedTokens equal to budgetLimit (e.g., 50000, 50000) and assert that
result.Blocked is true and result.BlockedBy == "workflow" to ensure the
function's >= behavior is preserved; modify the TestChecker_CheckExecutionBudget
function to include this equality assertion alongside the existing > and < cases
for CheckExecutionBudget.
In `@internal/config/config_test.go`:
- Around line 210-226: In TestLoad_BudgetResetDay_RollingInvalid replace the
soft assertions that check for an error (assert.Error(t, err)) with
require.Error(t, err) before calling err.Error() so the test stops immediately
on a nil error; update both occurrences in the test (the ones before
assert.Contains checks) and add the require import if not already present so
require.Error is available while keeping the subsequent assert.Contains(t,
err.Error(), "...") assertions unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 6fd72529-c1d9-4966-bc4e-08c92e54b3dd
📒 Files selected for processing (2)
internal/budget/budget_test.gointernal/config/config_test.go
Add exact-equality boundary test for CheckExecutionBudget and upgrade assert.Error to require.Error in TestLoad_BudgetResetDay_RollingInvalid to fail fast on nil error. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/budget/budget_test.go`:
- Around line 13-37: Add a year-boundary unit test for budget.CurrentPeriodStart
to verify rolling mode crosses years: create
TestCurrentPeriodStart_RollingYearBoundary that sets now := time.Date(2026, 1,
10, 0, 0, 0, 0, time.UTC), calls start := budget.CurrentPeriodStart(now,
"rolling", 15) and asserts start equals time.Date(2025, 12, 15, 0, 0, 0, 0,
time.UTC); this ensures the AddDate(0, -1, 0) path correctly handles year
rollover.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: beaea79e-fe3c-4182-89ec-4eba99a665d0
📒 Files selected for processing (2)
internal/budget/budget_test.gointernal/config/config_test.go
Verify AddDate(0,-1,0) correctly crosses year boundary (Jan 10 with resetDay 15 yields Dec 15 of prior year). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This was referenced Mar 24, 2026
This was referenced Apr 5, 2026
This was referenced Apr 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements AI cost controls with token-based budgets at three levels, composing as AND gates:
token_budget) — per-execution cap across all AI stepsWhat's included
ai_token_usagecounter table +team_budgetsconfig tableexecuteStepLogic(covers sequential + distributed worker paths), post-completion token recordingengine.budget.*settings with env var bindingsGET/PUT/DELETE /api/v1/budgets/{provider},GET /api/v1/budgets/usagemantle budget set/get/usage/deletemantle_budget_check_total), audit events (budget.exceeded,budget.warning,budget.updated)Design decisions
Test plan
go build ./...passesgo vet ./...cleantoken_budgetfieldmantle budget set openai 1000000→mantle budget getshows itCloses #8
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
CLI & REST
Database / Backend
Configuration
Metrics & Audit
Tests