feat(budgets): scope budgets to request labels - #590
Conversation
|
Warning Review limit reached
Next review available in: 12 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: ⛔ Files ignored due to path filters (3)
📒 Files selected for processing (13)
📝 WalkthroughWalkthroughBudgets now support ChangesScoped budgeting
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
Confidence Score: 3/5This PR should not merge until the MongoDB upgrade migration drops or replaces the legacy unique index before removing Existing MongoDB installations with multiple user paths sharing a budget period encounter a duplicate-key error during startup migration because all migrated documents lose the field used by the still-active legacy unique index. Files Needing Attention: internal/budget/store_mongodb.go
 T-Rex LogsWhat T-Rex did
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Existing MongoDB budgets] --> B[Startup migration]
B --> C[Add scope and subject]
C --> D[Unset legacy user_path]
D --> E[Drop legacy unique index]
E --> F[Create scoped unique index]
Reviews (1): Last reviewed commit: "feat(budgets): scope budgets to request ..." | Re-trigger Greptile |
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| {Key: "scope", Value: string(ScopeUserPath)}, | ||
| {Key: "subject", Value: "$user_path"}, | ||
| }}}, | ||
| bson.D{{Key: "$unset", Value: "user_path"}}, |
There was a problem hiding this comment.
When an existing MongoDB deployment has budgets for different user paths with the same period, this update unsets user_path while the legacy (user_path, period_seconds) unique index is still active. The migrated documents then share the same missing-path index key, causing UpdateMany to fail with a duplicate-key error and preventing the gateway from starting after upgrade.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
internal/budget/store_sql_test.go (1)
190-256: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
SumSpendis only exercised on SQLite.Both spend tests open a raw
sqliteDB instead of going throughsqlxtest.Run, so the PostgreSQL branch ofspendSubjectMatch— notablyjsonb_exists(labels, ?)and the non-unixepochtime comparison insumSpendChunk— has no coverage at all. That is the highest-risk new SQL in this PR. Routing these throughsqlxtest.Run(asTestSQLStoreMigratesPreScopeTabledoes) would exercise both engines wherever a Postgres instance is available.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/budget/store_sql_test.go` around lines 190 - 256, Update the spend tests, including TestSQLStoreSumSpendHonorsSubjectBoundaryAndCacheType and the other SumSpend test, to run through sqlxtest.Run like TestSQLStoreMigratesPreScopeTable instead of opening a raw SQLite database. Keep the existing setup and assertions within the harness so PostgreSQL executes spendSubjectMatch and sumSpendChunk when available, while retaining SQLite coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/admin/handler_budgets.go`:
- Around line 374-383: Update the subject/user_path validation around rawSubject
and rawUserPath so payloads containing both fields are rejected for every scope,
including budget.ScopeUserPath, instead of silently preferring subject. Preserve
the existing fallback that uses rawUserPath when subject is empty, and keep the
label-budget-specific error behavior where applicable.
In `@internal/budget/store_mongodb.go`:
- Around line 419-441: Update MongoDBStore.SumSpend to process windows in
spendChunkSize-sized batches, matching SQLStore.SumSpend, and execute a bounded
aggregation per chunk. Preserve the existing mongoSpendCondition,
spendTotalField, and spendHasUsageField accumulation behavior, then combine
chunk results into the same []Spend output and error semantics.
- Around line 59-76: Update migratePreScopeDocuments to drop the legacy
user_path_1_period_seconds_1 index before calling UpdateMany, keeping the drop
best-effort for fresh databases. Preserve the existing migration filter,
pipeline, and error wrapping, and remove the post-migration index drop.
In `@internal/budget/store_sql_migrate.go`:
- Around line 38-60: Update the migration error handling in the InTx callback
around the scoped-schema rebuild so a failed migration re-inspects the budgets
table schema and returns nil when another process has already completed the
migration. Preserve the existing error propagation for failures where the schema
is still unmigrated or cannot be inspected, allowing concurrent startup attempts
to converge.
In `@internal/budget/store_sql.go`:
- Around line 155-162: The stale-config delete construction around the budgets
loop must avoid exceeding SQLite’s parameter limit because each budget adds
three bind arguments. Chunk the budget conditions using the existing chunking
approach (such as slices.Chunk), execute the delete per chunk while preserving
the NOT-matching semantics, and retain the current behavior when budgets is
empty.
In `@internal/server/budget_support_test.go`:
- Around line 21-27: Add assertions covering Subjects.Labels in both affected
server tests: in internal/server/budget_support_test.go:21-27, extend a
request-labels case to verify countingBudgetChecker.subjects.Labels alongside
UserPath; in internal/server/usage_status_handler_test.go:31-41, assert
budgets.gotSubjects.Labels equals the labels attached to the request context.
Update the existing table-driven test cases without changing the fakes’
behavior.
---
Outside diff comments:
In `@internal/budget/store_sql_test.go`:
- Around line 190-256: Update the spend tests, including
TestSQLStoreSumSpendHonorsSubjectBoundaryAndCacheType and the other SumSpend
test, to run through sqlxtest.Run like TestSQLStoreMigratesPreScopeTable instead
of opening a raw SQLite database. Keep the existing setup and assertions within
the harness so PostgreSQL executes spendSubjectMatch and sumSpendChunk when
available, while retaining SQLite coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: ef845c64-da50-4195-b965-c361dd2401e9
⛔ Files ignored due to path filters (3)
internal/admin/dashboard/static/dist/assets/index-BSE-2hNO.jsis excluded by!**/dist/**internal/admin/dashboard/static/dist/assets/index-C5b8F5Yi.jsis excluded by!**/dist/**internal/admin/dashboard/static/dist/index.htmlis excluded by!**/dist/**
📒 Files selected for processing (35)
CLAUDE.mdcmd/gomodel/docs/docs.goconfig/budget.goconfig/config.example.yamlconfig/config_test.godocs/features/budgets.mdxdocs/features/labelling.mdxdocs/openapi.jsoninternal/admin/handler_budgets.gointernal/admin/handler_budgets_test.gointernal/budget/factory.gointernal/budget/service.gointernal/budget/service_test.gointernal/budget/store.gointernal/budget/store_mongodb.gointernal/budget/store_mongodb_test.gointernal/budget/store_sql.gointernal/budget/store_sql_migrate.gointernal/budget/store_sql_test.gointernal/budget/types.gointernal/budget/types_test.gointernal/server/budget_support.gointernal/server/budget_support_test.gointernal/server/mcp_service_test.gointernal/server/usage_status_handler.gointernal/server/usage_status_handler_test.gotests/e2e/budget_test.gotests/e2e/upgrade-compat.shtests/integration/dbassert/budget.goweb/dashboard/src/pages/budgets/BudgetEditor.svelteweb/dashboard/src/pages/budgets/BudgetList.svelteweb/dashboard/src/pages/budgets/BudgetsPage.svelteweb/dashboard/src/pages/budgets/budgets-helpers.jsweb/dashboard/src/pages/budgets/budgets.svelte.jsweb/dashboard/tests/budgets.test.js
| rawSubject = strings.TrimSpace(rawSubject) | ||
| if rawSubject == "" { | ||
| if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" { | ||
| return "", "", errors.New("subject is required for label budgets; user_path only names user-path budgets") | ||
| } | ||
| rawSubject = rawUserPath | ||
| } else if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" { | ||
| // Silently dropping the conflicting field would mask authoring mistakes. | ||
| return "", "", errors.New("user_path must not be set alongside subject for label budgets") | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Conflicting subject + user_path is only rejected for label budgets.
For the default user_path scope, a payload carrying both fields silently drops user_path and keys off subject — the exact authoring mistake the label branch deliberately rejects. Since this identity feeds delete/reset, the silent winner can address the wrong budget.
🛡️ Proposed fix
rawSubject = strings.TrimSpace(rawSubject)
if rawSubject == "" {
if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" {
return "", "", errors.New("subject is required for label budgets; user_path only names user-path budgets")
}
rawSubject = rawUserPath
- } else if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" {
+ } else if strings.TrimSpace(rawUserPath) != "" {
// Silently dropping the conflicting field would mask authoring mistakes.
- return "", "", errors.New("user_path must not be set alongside subject for label budgets")
+ return "", "", errors.New("set either subject or user_path, not both")
}📝 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.
| rawSubject = strings.TrimSpace(rawSubject) | |
| if rawSubject == "" { | |
| if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" { | |
| return "", "", errors.New("subject is required for label budgets; user_path only names user-path budgets") | |
| } | |
| rawSubject = rawUserPath | |
| } else if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" { | |
| // Silently dropping the conflicting field would mask authoring mistakes. | |
| return "", "", errors.New("user_path must not be set alongside subject for label budgets") | |
| } | |
| rawSubject = strings.TrimSpace(rawSubject) | |
| if rawSubject == "" { | |
| if scope != budget.ScopeUserPath && strings.TrimSpace(rawUserPath) != "" { | |
| return "", "", errors.New("subject is required for label budgets; user_path only names user-path budgets") | |
| } | |
| rawSubject = rawUserPath | |
| } else if strings.TrimSpace(rawUserPath) != "" { | |
| // Silently dropping the conflicting field would mask authoring mistakes. | |
| return "", "", errors.New("set either subject or user_path, not both") | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/admin/handler_budgets.go` around lines 374 - 383, Update the
subject/user_path validation around rawSubject and rawUserPath so payloads
containing both fields are rejected for every scope, including
budget.ScopeUserPath, instead of silently preferring subject. Preserve the
existing fallback that uses rawUserPath when subject is empty, and keep the
label-budget-specific error behavior where applicable.
| func (s *MongoDBStore) SumSpend(ctx context.Context, windows []SpendWindow) ([]Spend, error) { | ||
| if len(windows) == 0 { | ||
| return nil, nil | ||
| } | ||
| group := bson.D{{Key: "_id", Value: nil}} | ||
| for i, window := range windows { | ||
| condition, err := mongoSpendCondition(window) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| group = append(group, | ||
| bson.E{Key: spendTotalField(i), Value: bson.D{{Key: "$sum", Value: bson.D{{Key: "$cond", Value: bson.A{ | ||
| condition, | ||
| bson.D{{Key: "$ifNull", Value: bson.A{"$total_cost", 0}}}, | ||
| 0, | ||
| }}}}}}, | ||
| bson.E{Key: spendHasUsageField(i), Value: bson.D{{Key: "$sum", Value: bson.D{{Key: "$cond", Value: bson.A{ | ||
| bson.D{{Key: "$and", Value: bson.A{condition, bson.D{{Key: "$gt", Value: bson.A{"$total_cost", nil}}}}}}, | ||
| 1, | ||
| 0, | ||
| }}}}}}, | ||
| ) | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
No chunking parity with the SQL store.
SQLStore.SumSpend splits windows into spendChunkSize batches; here every window adds two accumulators to one $group stage, so a request matching many label budgets builds an unbounded pipeline. Consider chunking the windows the same way to bound pipeline size.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/budget/store_mongodb.go` around lines 419 - 441, Update
MongoDBStore.SumSpend to process windows in spendChunkSize-sized batches,
matching SQLStore.SumSpend, and execute a bounded aggregation per chunk.
Preserve the existing mongoSpendCondition, spendTotalField, and
spendHasUsageField accumulation behavior, then combine chunk results into the
same []Spend output and error semantics.
| // One transaction: a crash mid-rebuild must not leave the table renamed | ||
| // away, or the next startup would create a fresh empty budgets table and | ||
| // orphan every limit. If concurrent replicas race here, one commits and the | ||
| // others fail fast and see the migrated schema on restart. | ||
| return db.InTx(ctx, func(q sqlx.Querier) error { | ||
| statements := []string{ | ||
| `ALTER TABLE budgets RENAME TO budgets_pre_scope`, | ||
| db.Dialect().ExpandTypes(sqlBudgetsSchema), | ||
| `INSERT INTO budgets (scope, subject, period_seconds, amount, source, last_reset_at, created_at, updated_at) | ||
| SELECT 'user_path', user_path, period_seconds, amount, ` + sourceExpr + `, ` + lastResetExpr + `, created_at, updated_at | ||
| FROM budgets_pre_scope`, | ||
| `DROP TABLE budgets_pre_scope`, | ||
| // SQLite carried the old index along with the renamed table; | ||
| // PostgreSQL drops it with the table, so this is a no-op there. | ||
| `DROP INDEX IF EXISTS idx_budgets_user_path`, | ||
| } | ||
| for _, statement := range statements { | ||
| if _, err := q.Exec(ctx, statement); err != nil { | ||
| return fmt.Errorf("migrate budgets to scoped schema: %w", err) | ||
| } | ||
| } | ||
| return nil | ||
| }) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Concurrent startup loses the race hard instead of converging.
On PostgreSQL (READ COMMITTED) two replicas can both observe user_path and both attempt the rebuild; the loser blocks on the lock and then fails (relation "budgets" does not exist / duplicate table), which propagates out of NewSQLStore and aborts startup. During a rolling deploy that turns into a crash-looping replica until it is restarted. Cheap fix: on error, re-inspect the columns and return nil if another process already migrated.
♻️ Converge instead of failing
- return db.InTx(ctx, func(q sqlx.Querier) error {
+ err = db.InTx(ctx, func(q sqlx.Querier) error {
statements := []string{
`ALTER TABLE budgets RENAME TO budgets_pre_scope`,
db.Dialect().ExpandTypes(sqlBudgetsSchema),
@@
return nil
})
+ if err != nil {
+ // Another replica may have committed the same rebuild first.
+ if after, inspectErr := budgetColumns(ctx, db); inspectErr == nil && after["subject"] {
+ return nil
+ }
+ return err
+ }
+ return nil🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/budget/store_sql_migrate.go` around lines 38 - 60, Update the
migration error handling in the InTx callback around the scoped-schema rebuild
so a failed migration re-inspects the budgets table schema and returns nil when
another process has already completed the migration. Preserve the existing error
propagation for failures where the schema is still unmigrated or cannot be
inspected, allowing concurrent startup attempts to converge.
| if len(budgets) > 0 { | ||
| conditions := make([]string, 0, len(budgets)) | ||
| for _, budget := range budgets { | ||
| conditions = append(conditions, `(user_path = ? AND period_seconds = ?)`) | ||
| args = append(args, budget.UserPath, budget.PeriodSeconds) | ||
| conditions = append(conditions, `(scope = ? AND subject = ? AND period_seconds = ?)`) | ||
| args = append(args, budget.Scope, budget.Subject, budget.PeriodSeconds) | ||
| } | ||
| query += ` AND NOT (` + strings.Join(conditions, " OR ") + `)` | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Parameter count in the stale-config delete now grows 3× per budget.
Each declared budget contributes three bind parameters, so the single DELETE exceeds SQLite's 999-parameter ceiling at ~333 config budgets (previously ~500). SumSpend explicitly chunks for this reason; this statement doesn't. Consider chunking with slices.Chunk here as well, or deleting by source = ? minus a temp/IN (VALUES ...) set.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/budget/store_sql.go` around lines 155 - 162, The stale-config delete
construction around the budgets loop must avoid exceeding SQLite’s parameter
limit because each budget adds three bind arguments. Chunk the budget conditions
using the existing chunking approach (such as slices.Chunk), execute the delete
per chunk while preserving the NOT-matching semantics, and retain the current
behavior when budgets is empty.
| subjects budget.Subjects | ||
| } | ||
|
|
||
| func (c *countingBudgetChecker) Check(_ context.Context, userPath string, _ time.Time) error { | ||
| func (c *countingBudgetChecker) Check(_ context.Context, subjects budget.Subjects, _ time.Time) error { | ||
| c.calls++ | ||
| c.userPath = userPath | ||
| c.subjects = subjects | ||
| return nil |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Subjects.Labels is unasserted across the server-layer budget tests. Both fakes were widened from a string user path to budget.Subjects, but every assertion still reads only UserPath, so the new label-propagation contract has no test guarding it in this package.
internal/server/budget_support_test.go#L21-L27: assertchecker.subjects.Labelsin a case that sets request labels, not justUserPath.internal/server/usage_status_handler_test.go#L31-L41: assertbudgets.gotSubjects.Labelsmatches the labels attached to the request context.
As per path instructions for **/*_test.go: "Add or update table-driven tests for behavior changes".
📍 Affects 2 files
internal/server/budget_support_test.go#L21-L27(this comment)internal/server/usage_status_handler_test.go#L31-L41
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/server/budget_support_test.go` around lines 21 - 27, Add assertions
covering Subjects.Labels in both affected server tests: in
internal/server/budget_support_test.go:21-27, extend a request-labels case to
verify countingBudgetChecker.subjects.Labels alongside UserPath; in
internal/server/usage_status_handler_test.go:31-41, assert
budgets.gotSubjects.Labels equals the labels attached to the request context.
Update the existing table-driven test cases without changing the fakes’
behavior.
Source: Path instructions
Budgets could only limit a `user_path` subtree, so an operator running several applications through one path had no way to cap their spend separately without minting a key per application (issue #319). A budget now carries a scope and a subject, mirroring the shape rate limit rules already use: - `user_path` — a path and its descendants, unchanged - `label` — every request carrying that label, matched verbatim Labels come from the tagging headers and from the managed key that authenticated the request, so a request carrying several labels is charged against every matching label budget. Label budgets are declared under `budgets.labels:` in config.yaml or through the dashboard and admin API; they have no env-var form, because labels are matched verbatim and are not env-name safe (the same rationale model rate limit rules already use). Enforcement no longer runs one SUM query per matching budget. `Store.SumSpend` takes every matching window and answers them in one scan — conditional aggregation on SQL, one `$group` on MongoDB. Timed against a 50k-row usage table: parity at one matching budget (~30ms), 96ms -> 35ms at five, 443ms -> 70ms at twenty-five. The budgets table is rekeyed to (scope, subject, period_seconds), so both SQL backends rebuild it in one transaction and MongoDB rewrites its documents in place. The rebuild also absorbs the two columns older releases added with ALTER TABLE, which removes the separate migration list. Also folds three near-identical evaluation loops in the budget service into one match/evaluate pair, moves the SQL-only path helpers out of the shared store file, and corrects the admin API examples in the budgets docs, which showed URL path parameters those endpoints never had. Part of #319 — this covers the budgeting integration the issue asked for. Label filtering and grouping in the audit log remain open. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The scope migration unset `user_path` while the pre-scope unique index on
(user_path, period_seconds) was still in place. A missing field indexes as
null, so every migrated document sharing a period collapsed onto the key
{user_path: null, period_seconds: N} and the rewrite failed with a duplicate
key error — aborting startup for any MongoDB deployment holding two budgets
with the same period, which is the ordinary case rather than an edge one.
`internal/ratelimit` has carried the same ordering since rule scopes shipped,
so it is fixed alongside. Both now drop the legacy index first, and both gain
a mongotest case that seeds two pre-scope documents sharing a period and
asserts they survive the upgrade.
Also asserts at the server layer that request labels reach the budget checker
with the user path, which is what lets a label budget match at all.
Reported by Greptile and CodeRabbit on #590.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
7921c5f to
d0921be
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
internal/budget/store_sql_test.go (1)
190-314: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBoth
SumSpendtests bypasssqlxtest.Run, so the PostgreSQL predicate is never exercised.
spendSubjectMatchemits a completely different label predicate per dialect (json_eachvsjsonb_exists), yet these tests open raw SQLite only. The PostgreSQL branch of the new aggregation ships untested.As per coding guidelines: "Add or update tests for behavior changes, using table-driven tests where appropriate. Cover request translation, response normalization, error handling, default configuration, and provider-specific parameter mapping."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/budget/store_sql_test.go` around lines 190 - 314, Update TestSQLStoreSumSpendHonorsSubjectBoundaryAndCacheType and TestSQLStoreSumSpendChunksLargeBatches to run through sqlxtest.Run so each case executes against the configured dialects, including PostgreSQL. Preserve the existing assertions and test data while ensuring the PostgreSQL-specific spendSubjectMatch label predicate is exercised; apply the same harness change to both SumSpend tests.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cmd/gomodel/docs/docs.go`:
- Around line 10178-10186: Update the `/v1/usage` endpoint description in the
generated documentation to state that the response may include matching
label-scoped budgets, in addition to statuses for the effective `user_path`,
consistent with the exposed `scope` and `subject` fields.
In `@internal/budget/store_sql.go`:
- Around line 316-338: Update label-budget tests in
internal/budget/store_sql_test.go (lines 190-314) to run through sqlxtest.Run,
ensuring both SQLite and PostgreSQL paths exercise spendSubjectMatch, including
the jsonb_exists(labels, ?) branch. Adjust SQLite seed data to use valid JSON
values in usage.labels rather than raw strings such as "iOS"; no direct change
is required in internal/budget/store_sql.go (lines 316-338) unless needed to
support the cross-database test coverage.
In `@internal/budget/store.go`:
- Around line 30-43: Document the non-empty precondition for spendBounds,
explicitly stating that callers must provide at least one SpendWindow because
the helper indexes windows[0]. Keep the existing range-union behavior unchanged.
In `@internal/server/usage_status_handler_test.go`:
- Around line 30-41: Extend the usage status handler tests to cover the
label-scope response branch: send request labels, assert the
fakeBudgetStatusChecker’s gotSubjects contains the translated labels, and verify
ScopeLabel results serialize scope and subject while omitting user_path. Use the
existing test structure and table-driven style where appropriate, including the
normal response behavior alongside this new case.
In `@web/dashboard/src/pages/budgets/budgets.svelte.js`:
- Line 244: Update the reset and delete confirmation labels in budgets.svelte.js
at lines 244-244 and 286-286 to include budgetScopeLabel(item) alongside the
subject and period. Update the override message in budgets-helpers.js at line
309-309 to include the selected budget’s scope, ensuring every budget
confirmation identifies subject, period, and scope.
---
Outside diff comments:
In `@internal/budget/store_sql_test.go`:
- Around line 190-314: Update
TestSQLStoreSumSpendHonorsSubjectBoundaryAndCacheType and
TestSQLStoreSumSpendChunksLargeBatches to run through sqlxtest.Run so each case
executes against the configured dialects, including PostgreSQL. Preserve the
existing assertions and test data while ensuring the PostgreSQL-specific
spendSubjectMatch label predicate is exercised; apply the same harness change to
both SumSpend tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: f05696e4-06a7-4f87-ad37-f99dc6dbe2ce
⛔ Files ignored due to path filters (3)
internal/admin/dashboard/static/dist/assets/index-BSE-2hNO.jsis excluded by!**/dist/**internal/admin/dashboard/static/dist/assets/index-C5b8F5Yi.jsis excluded by!**/dist/**internal/admin/dashboard/static/dist/index.htmlis excluded by!**/dist/**
📒 Files selected for processing (37)
CLAUDE.mdcmd/gomodel/docs/docs.goconfig/budget.goconfig/config.example.yamlconfig/config_test.godocs/features/budgets.mdxdocs/features/labelling.mdxdocs/openapi.jsoninternal/admin/handler_budgets.gointernal/admin/handler_budgets_test.gointernal/budget/factory.gointernal/budget/service.gointernal/budget/service_test.gointernal/budget/store.gointernal/budget/store_mongodb.gointernal/budget/store_mongodb_test.gointernal/budget/store_sql.gointernal/budget/store_sql_migrate.gointernal/budget/store_sql_test.gointernal/budget/types.gointernal/budget/types_test.gointernal/ratelimit/store_mongodb.gointernal/ratelimit/store_mongodb_test.gointernal/server/budget_support.gointernal/server/budget_support_test.gointernal/server/mcp_service_test.gointernal/server/usage_status_handler.gointernal/server/usage_status_handler_test.gotests/e2e/budget_test.gotests/e2e/upgrade-compat.shtests/integration/dbassert/budget.goweb/dashboard/src/pages/budgets/BudgetEditor.svelteweb/dashboard/src/pages/budgets/BudgetList.svelteweb/dashboard/src/pages/budgets/BudgetsPage.svelteweb/dashboard/src/pages/budgets/budgets-helpers.jsweb/dashboard/src/pages/budgets/budgets.svelte.jsweb/dashboard/tests/budgets.test.js
| "scope": { | ||
| "type": "string" | ||
| }, | ||
| "spent": { | ||
| "type": "number" | ||
| }, | ||
| "subject": { | ||
| "type": "string" | ||
| }, |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Document label budgets in /v1/usage.
The endpoint description still says statuses are for the effective user_path only. Update it to explain that matching label-scoped budgets can also appear now that this response exposes scope and subject.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cmd/gomodel/docs/docs.go` around lines 10178 - 10186, Update the `/v1/usage`
endpoint description in the generated documentation to state that the response
may include matching label-scoped budgets, in addition to statuses for the
effective `user_path`, consistent with the exposed `scope` and `subject` fields.
| // spendBounds returns the union of every window's time range, which is the | ||
| // range a batched spend query has to scan. | ||
| func spendBounds(windows []SpendWindow) (time.Time, time.Time) { | ||
| start, end := windows[0].Start, windows[0].End | ||
| for _, window := range windows[1:] { | ||
| if window.Start.Before(start) { | ||
| start = window.Start | ||
| } | ||
| if window.End.After(end) { | ||
| end = window.End | ||
| } | ||
| } | ||
| return start, end | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
State the non-empty precondition on spendBounds.
Both callers guard len(windows) == 0, but the helper panics on an empty slice and its doc comment does not say so. A one-line note (or an explicit guard) keeps a future caller from tripping it.
♻️ Proposed doc tweak
// spendBounds returns the union of every window's time range, which is the
-// range a batched spend query has to scan.
+// range a batched spend query has to scan. Callers must pass at least one
+// window.
func spendBounds(windows []SpendWindow) (time.Time, time.Time) {📝 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.
| // spendBounds returns the union of every window's time range, which is the | |
| // range a batched spend query has to scan. | |
| func spendBounds(windows []SpendWindow) (time.Time, time.Time) { | |
| start, end := windows[0].Start, windows[0].End | |
| for _, window := range windows[1:] { | |
| if window.Start.Before(start) { | |
| start = window.Start | |
| } | |
| if window.End.After(end) { | |
| end = window.End | |
| } | |
| } | |
| return start, end | |
| } | |
| // spendBounds returns the union of every window's time range, which is the | |
| // range a batched spend query has to scan. Callers must pass at least one | |
| // window. | |
| func spendBounds(windows []SpendWindow) (time.Time, time.Time) { | |
| start, end := windows[0].Start, windows[0].End | |
| for _, window := range windows[1:] { | |
| if window.Start.Before(start) { | |
| start = window.Start | |
| } | |
| if window.End.After(end) { | |
| end = window.End | |
| } | |
| } | |
| return start, end | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/budget/store.go` around lines 30 - 43, Document the non-empty
precondition for spendBounds, explicitly stating that callers must provide at
least one SpendWindow because the helper indexes windows[0]. Keep the existing
range-union behavior unchanged.
| type fakeBudgetStatusChecker struct { | ||
| results []budget.CheckResult | ||
| err error | ||
| gotPath string | ||
| results []budget.CheckResult | ||
| err error | ||
| gotSubjects budget.Subjects | ||
| } | ||
|
|
||
| func (f *fakeBudgetStatusChecker) Check(context.Context, string, time.Time) error { return nil } | ||
| func (f *fakeBudgetStatusChecker) Check(context.Context, budget.Subjects, time.Time) error { return nil } | ||
|
|
||
| func (f *fakeBudgetStatusChecker) StatusesForPath(_ context.Context, userPath string, _ time.Time) ([]budget.CheckResult, error) { | ||
| f.gotPath = userPath | ||
| func (f *fakeBudgetStatusChecker) StatusesFor(_ context.Context, subjects budget.Subjects, _ time.Time) ([]budget.CheckResult, error) { | ||
| f.gotSubjects = subjects | ||
| return f.results, f.err | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add coverage for the label-scope response branch.
The fake now captures gotSubjects, but no test asserts that request labels reach StatusesFor, nor that a ScopeLabel result serializes with scope/subject and omits user_path (usage_status_handler.go lines 209-227). That branch is new behavior and currently untested.
As per coding guidelines: "Add or update tests for behavior changes, using table-driven tests where appropriate. Cover request translation, response normalization, error handling, default configuration, and provider-specific parameter mapping."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/server/usage_status_handler_test.go` around lines 30 - 41, Extend
the usage status handler tests to cover the label-scope response branch: send
request labels, assert the fakeBudgetStatusChecker’s gotSubjects contains the
translated labels, and verify ScopeLabel results serialize scope and subject
while omitting user_path. Use the existing test structure and table-driven style
where appropriate, including the normal response behavior alongside this new
case.
Source: Coding guidelines
| return; | ||
| } | ||
| const label = String(item.user_path || "") + " " + budgetPeriodLabel(item); | ||
| const label = budgetSubject(item) + " " + budgetPeriodLabel(item); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Include scope in every budget confirmation label.
Budget identity now includes scope, but these messages display only subject and period. A user-path and label budget can therefore appear identical.
web/dashboard/src/pages/budgets/budgets.svelte.js#L244-L244: includebudgetScopeLabel(item)in the reset confirmation.web/dashboard/src/pages/budgets/budgets.svelte.js#L286-L286: includebudgetScopeLabel(item)in the delete confirmation.web/dashboard/src/pages/budgets/budgets-helpers.js#L309-L309: include the selected budget’s scope in the override message.
📍 Affects 2 files
web/dashboard/src/pages/budgets/budgets.svelte.js#L244-L244(this comment)web/dashboard/src/pages/budgets/budgets.svelte.js#L286-L286web/dashboard/src/pages/budgets/budgets-helpers.js#L309-L309
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@web/dashboard/src/pages/budgets/budgets.svelte.js` at line 244, Update the
reset and delete confirmation labels in budgets.svelte.js at lines 244-244 and
286-286 to include budgetScopeLabel(item) alongside the subject and period.
Update the override message in budgets-helpers.js at line 309-309 to include the
selected budget’s scope, ensuring every budget confirmation identifies subject,
period, and scope.
The label branch of SumSpend is the one part written twice — json_each on SQLite, jsonb_exists on PostgreSQL — and only the SQLite half was tested, because the test opened a raw SQLite database to get a usage store. A mistake in the PostgreSQL predicate would have silently matched nothing, so label budgets would never enforce there and no test would notice. The case now runs on both backends. Verified it catches a broken predicate by corrupting the bound subject and watching the PostgreSQL subtest fail. `usage` has not moved onto sqlx, so its PostgreSQL store still takes a *pgxpool.Pool directly; sqlxtest grows an exported NewPostgresPool so a test spanning both can point them at one throwaway schema. Entry ids become fixed UUIDs derived from their readable names, since PostgreSQL types usage.id as a UUID. Reported by CodeRabbit on #590. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
# Conflicts: # internal/admin/dashboard/static/dist/index.html
The budget list carried the scope twice: a "(tag) label" pill sat next to the period, while the subject rendered in .budget-user-path regardless of what it actually was. The pill said nothing the subject could not say itself, and the class name was wrong for half the rows. The tag icon now prefixes the label name inside the subject, and the shared styling moves to .budget-scope-value with .budget-user-path and .budget-label as hue-only modifiers. A label takes the same per-label colour the usage log, breakdown chart, and API keys already use, so one label reads the same everywhere. The rate limit list and inspector pick up the shared base class. Also rounds the provider status toggle to 6px, matching .btn and .table-action-btn rather than the full pill it was; the switch track and thumb stay round. And the budget filter placeholder now mentions labels, which it can now match on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
Part of #319 — this covers the budgeting integration the issue asked for. Header-to-label ingestion already shipped; label filtering and grouping in the audit log remain open, as do conditional (
if header X equals Y) tagging rules. Deliberately not auto-closing the issue.Budgets could only limit a
user_pathsubtree. An operator running several applications through one path — the exact case in the issue — had no way to cap their spend separately without minting a key per application.What changed
A budget now carries a scope and a subject, mirroring the shape rate limit rules already use:
user_path/team/alphalabelMobile-App-iOSLabels are matched verbatim — no subtree semantics, no case folding — the same way the usage label filter and by-label breakdown already compare them. They come from the tagging headers and from the managed key that authenticated the request, so a request carrying several labels is charged against every matching label budget. A breach returns the usual 429
budget_exceeded, now naming the subject:Surfaces:
labelchip on label rowsscope+subjecton/admin/budgets, withuser_pathkept as the shorthand spelling for user-path budgetsbudgets.labels:inconfig.yamlGET /v1/usage— label budgets appear in the caller's own statusLabel budgets are deliberately YAML/admin-only: labels are matched verbatim and are not env-name safe. That is the same rationale model rate limit rules already use for skipping an env form.
Performance
Enforcement runs on every request and used to issue one
SUMquery per matching budget. Adding label budgets on top of that would have multiplied the query count, soSumUsageCostis replaced by a batchedStore.SumSpend([]SpendWindow): one scan with conditional aggregation on SQL, one$groupwith conditional sums on MongoDB.That makes the pre-existing user-path path faster too. Timed against a 50k-row usage table, full budget evaluation:
Parity at one budget, ~6x at twenty-five. A unit test asserts the store is hit exactly once regardless of how many budgets match.
Migration
budgetsis rekeyed to(scope, subject, period_seconds). Both SQL backends rebuild the table inside one transaction and MongoDB rewrites its documents in place — the same patternratelimitused for its own scope migration. The rebuild also absorbs the two columns older releases added withALTER TABLE, so the separate migration list is gone.Verified on a real SQLite database written in the old shape: rows preserved, old index dropped, and a label budget with the same subject spelling now coexists (the old primary key would have rejected it).
tests/e2e/upgrade-compat.shcovers it going forward, and there is a migration test over both SQL dialects.Cleanups along the way
budget.Servicefolded into onematch+evaluatepair;StatusesForPathbecomesStatusesForstore.go; per-store key validation deduplicated intonormalizeBudgetKeyVerification
Live against a running gateway with a real provider:
/v1/usagereports the label budget for the callerFull Go suite, e2e, lint, 344 dashboard tests, swagger/openapi regenerated, dashboard bundle rebuilt.
Note
Response cache hits still return before budget enforcement, as documented. Worth knowing when validating budgets by hand — it briefly looked like a bug during testing.
🤖 Generated with Claude Code
Summary by CodeRabbit
Review round 1
Fixed — legacy MongoDB index blocked the migration (Greptile P1 / CodeRabbit Critical, independently found by both).
migratePreScopeDocumentsunsetuser_pathwhile the pre-scope unique index on(user_path, period_seconds)was still in place. A missing field indexes as null, so every migrated document sharing a period collapsed onto{user_path: null, period_seconds: N}. Reproduced against MongoDB 8:Startup would have aborted for any Mongo deployment holding two budgets with the same period — the ordinary case. Fixed by dropping the legacy index first, with a
mongotestcase that seeds two colliding pre-scope documents and asserts they survive.internal/ratelimithas carried the identical bug since rule scopes shipped — that is where this migration was copied from. Reproduced there too and fixed alongside, with its own regression test.Fixed — label propagation was untested at the server layer (CodeRabbit). The fakes were widened to
budget.Subjectsbut nothing asserted labels actually reach the checker. Added a table-driven case covering labels with and without a bound user path.Not changed, with reasons
ratelimitmigration this one mirrors; diverging in one of the two would be worse than the transient. Worth changing in both or neither.subject+user_pathboth set is only rejected for label budgets (CodeRabbit Minor) — exact parity with the shippedrate_limitsadmin API. Same call: both or neither.DELETEparameter count (CodeRabbit Minor) — real, and I narrowed it: 3 binds per budget instead of 2 moves SQLite's old 999-parameter ceiling from ~500 to ~333 config-declared budgets. Far beyond any realistic declaration, and restructuring the NOT-IN delete to chunk safely is not worth it here.SumSpendchunking parity (CodeRabbit Trivial) — SQL chunks for a bind-parameter limit that MongoDB does not have. Window count is bounded by matching budgets; the 16MB command limit is orders of magnitude away.MCP_ENABLEDin the generated spec (CodeRabbit out-of-scope warning) — three lines of pre-existing drift thatmake swaggercorrected. CI enforces spec-vs-source drift, so reverting it would just re-break that check.Review round 2
Fixed — the PostgreSQL label predicate was untested (CodeRabbit Major).
SumSpend's label branch is the one part written twice —json_eachon SQLite,jsonb_existson PostgreSQL — and only the SQLite half had a test, because that test opens a raw SQLite database to get a usage store. A mistake in the PostgreSQL predicate would have matched nothing silently: label budgets would simply never enforce on PG, with every test still green.The case now runs on both backends. I confirmed it is not vacuous by corrupting the bound subject and watching only the PostgreSQL subtest fail.
usagehas not moved ontosqlx, so its PostgreSQL store still takes a*pgxpool.Pooldirectly.sqlxtestgrows an exportedNewPostgresPoolso a test spanning a migrated and an unmigrated store can point both at one throwaway schema.The second half of that finding — that the SQLite seed inserts raw strings into
usage.labels— is not accurate: entries go through the usage store, which serializesLabelsto a valid JSON array. The assertions on label sums would fail otherwise.Not changed, with reasons
/v1/usageswagger (Minor) — the response model is generated fromusageStatusBudget, which now carriesscopeandsubject; the regenerated spec in this PR already reflects them.spendBoundsnon-empty precondition (Trivial) — its only caller isSumSpend, which returns early on an empty slice.usage_status_handler_test.go(Trivial) — the branch is three lines mirroring the admin handler, which this PR does cover with a label case.