governance: extend reset_budget_usage to teams, customers, model limits and provider governance - #6004
Merged
Merged
Conversation
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds an optional ChangesBudget usage reset
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant GovernanceForm
participant BudgetUsageResetDialog
participant GovernanceAPI
participant BudgetUsageStore
participant InMemoryBudgets
GovernanceForm->>BudgetUsageResetDialog: request reset choice
BudgetUsageResetDialog-->>GovernanceForm: preserve or reset usage
GovernanceForm->>GovernanceAPI: submit update with reset_budget_usage
GovernanceAPI->>BudgetUsageStore: clear affected budget usage
GovernanceAPI->>InMemoryBudgets: clear in-memory usage by owner
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
akshaydeo
marked this pull request as ready for review
August 9, 2026 23:01
This was referenced Aug 9, 2026
Merged
Contributor
Author
This was referenced Aug 9, 2026
coderabbitai
Bot
requested review from
Pratham-Mishra04,
impoiler and
roroghost17
August 9, 2026 23:02
![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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
transports/bifrost-http/server/server.go (1)
485-500: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftPropagate
ownerthrough the cluster reset path.
BifrostHTTPServer.ResetBudgetUsageInMemoryignoresownerand only resets the local governance store. Peers cannot identify the entity whose usage they must reset and can retain stale usage. Invoke the enterprise broadcast hook withownerandbudgetIDs, and test each owner kind and ID.🤖 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 `@transports/bifrost-http/server/server.go` around lines 485 - 500, Update BifrostHTTPServer.ResetBudgetUsageInMemory to invoke the enterprise cluster broadcast hook with owner and budgetIDs after validating the governance plugin, while preserving the local store reset. Ensure propagation supports each owner kind and ID, and add coverage verifying peers receive the correct owner and budget IDs.
🧹 Nitpick comments (1)
transports/bifrost-http/handlers/governance.go (1)
580-588: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd handler-level reset coverage.
Cover omitted and
truereset_budget_usagefor virtual keys, teams, customers, model configs, and provider governance. Assert persistedCurrentUsage == 0, unchangedLastReset, and in-memory usage reset after reload.🤖 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 `@transports/bifrost-http/handlers/governance.go` around lines 580 - 588, Add handler-level tests covering omitted and true reset_budget_usage for virtual keys, teams, customers, model configs, and provider governance. For each case, verify persisted CurrentUsage is zero, LastReset remains unchanged, and reloaded in-memory usage is reset, using the relevant governance handlers and the budgetUsageReset.apply flow.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 `@docs/features/governance/budget-and-limits.mdx`:
- Around line 219-221: Update the reset behavior documentation near the existing
window-preservation note to explicitly state that manual reset_budget_usage
clears current_usage only and leaves last_reset unchanged. Also state that
normal scheduled reset processing may advance last_reset later, keeping the
documentation aligned with the configuration contract.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 2771-2780: In transports/bifrost-http/handlers/governance.go, make
each reload failure terminate the handler with HTTP 500 before any usage reset
or success response: return immediately after ReloadTeam fails at lines
2771-2780, after ReloadCustomer fails at lines 3116-3125, and after
ReloadModelConfig fails at lines 3741-3750. Preserve the existing error logging
and response behavior at all three sites.
In `@ui/app/workspace/governance/views/teamSheet.tsx`:
- Around line 239-253: Update the budgetsChanged signature logic in
ui/app/workspace/governance/views/teamSheet.tsx lines 239-253 and the
corresponding change-detection logic in
ui/app/workspace/model-limits/views/modelLimitSheet.tsx lines 171-182 to include
a stable representation of each row’s resetConfig and persisted reset_config.
Ensure changes such as reset_config.quarter_start_month are detected even when
max_limit and reset_duration are unchanged.
---
Outside diff comments:
In `@transports/bifrost-http/server/server.go`:
- Around line 485-500: Update BifrostHTTPServer.ResetBudgetUsageInMemory to
invoke the enterprise cluster broadcast hook with owner and budgetIDs after
validating the governance plugin, while preserving the local store reset. Ensure
propagation supports each owner kind and ID, and add coverage verifying peers
receive the correct owner and budget IDs.
---
Nitpick comments:
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 580-588: Add handler-level tests covering omitted and true
reset_budget_usage for virtual keys, teams, customers, model configs, and
provider governance. For each case, verify persisted CurrentUsage is zero,
LastReset remains unchanged, and reloaded in-memory usage is reset, using the
relevant governance handlers and the budgetUsageReset.apply flow.
🪄 Autofix
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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 5bdac1c0-49b8-44a8-a18d-29f3af8d6f09
📒 Files selected for processing (13)
docs/features/governance/budget-and-limits.mdxdocs/openapi/openapi.jsondocs/openapi/schemas/management/governance.yamltransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/server/server.goui/app/workspace/governance/views/customerSheet.tsxui/app/workspace/governance/views/teamSheet.tsxui/app/workspace/model-limits/views/modelLimitSheet.tsxui/app/workspace/providers/fragments/governanceFormFragment.tsxui/components/ui/budgetUsageResetDialog.tsxui/hooks/useBudgetUsageResetPrompt.tsui/lib/types/governance.ts
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/governance/vkbudget_test.go (1)
567-580: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBound the number of paid completions on the failure path.
The loop issues a real chat completion every second for up to 30 seconds. If enforcement never rejects, the test bills about 30 completions before it fails. Add an attempt counter alongside the deadline, and stop at a small number of attempts.
♻️ Proposed change
// Usage is recorded asynchronously, so wait for enforcement to start rejecting. deadline := time.Now().Add(30 * time.Second) var blocked *APIResponse - for { + for attempt := 1; ; attempt++ { resp := completion() if resp.StatusCode >= 400 { blocked = resp break } - if time.Now().After(deadline) { + if attempt >= 10 || time.Now().After(deadline) { t.Fatalf("budget of $0.00001 never blocked a request; usage is %v", vkBudgetUsage(t, vkID)) } time.Sleep(time.Second) }🤖 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 `@tests/governance/vkbudget_test.go` around lines 567 - 580, Update the retry loop around completion() to track the number of attempts alongside the existing deadline, and stop retrying after a small bounded maximum. Preserve immediate exit when a response is rejected and retain the existing diagnostic failure message when the deadline or attempt limit is reached.tests/governance/providerbudget_test.go (1)
256-259: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winThe cleanup deletes provider governance that the test did not create.
The deferred
DELETE /api/governance/providers/openairemoves the global OpenAI governance unconditionally. If the target gateway already had provider governance configured foropenai, the test destroys it and does not restore it.baseURL()now makes the target gateway configurable, so the suite can run against an instance an operator also uses.Read the existing configuration before the test and restore it in the deferred function, or skip the test when governance already exists for the provider.
Also applies to: 322-324
🤖 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 `@tests/governance/providerbudget_test.go` around lines 256 - 259, Update the test around its provider-governance setup to read and preserve the existing OpenAI governance before making changes. In the deferred cleanup, restore the saved configuration instead of unconditionally deleting it; if no governance existed originally, delete the test-created entry. Apply the same preservation logic to the additional cleanup block noted in the comment, using the existing governance request helpers and configurable baseURL flow.tests/governance/teambudget_test.go (1)
427-441: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe same 30-second usage-polling loop is written five times.
waitForTeamBudgetUsageis specific to teams, so each new budget owner copied the loop body instead of reusing it. Generalize it to accept a reader function, then call it from every owner. This also keeps the timeout and the failure message in one place.
tests/governance/teambudget_test.go#L427-L441: changewaitForTeamBudgetUsageinto a shared helper such aswaitForBudgetUsage(t *testing.T, label string, read func() float64) float64, and keep a thin team wrapper if callers want one.tests/governance/customerbudget_test.go#L448-L459: replace the inline loop with a call that passes a closure overcustomerBudgetUsage(t, customerID).tests/governance/customerbudget_test.go#L541-L552: replace the inline loop with the same call.tests/governance/providerbudget_test.go#L284-L295: replace the inline loop with a call that passes a closure overproviderGovernanceUsage(t, "openai").tests/governance/providerbudget_test.go#L438-L449: replace the inline loop with a call that passes a closure overmodelConfigUsage(t, mcID).🤖 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 `@tests/governance/teambudget_test.go` around lines 427 - 441, Generalize waitForTeamBudgetUsage into a shared waitForBudgetUsage helper accepting a label and reader function, centralizing the 30-second polling timeout and failure message; retain a thin team wrapper if useful. In tests/governance/teambudget_test.go:427-441 update the helper. Replace each inline loop with the shared helper in tests/governance/customerbudget_test.go:448-459 and :541-552 using customerBudgetUsage closures, and in tests/governance/providerbudget_test.go:284-295 and :438-449 using providerGovernanceUsage and modelConfigUsage closures respectively.
🤖 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 `@tests/governance/teambudget_test.go`:
- Around line 408-425: Update teamBudgetUsage in
tests/governance/teambudget_test.go (lines 408-425), customerBudgetUsage in
tests/governance/customerbudget_test.go (lines 480-497), providerGovernanceUsage
in tests/governance/providerbudget_test.go (lines 381-385), and modelConfigUsage
in tests/governance/providerbudget_test.go (lines 486-502) to validate the
current_usage type assertion and call t.Fatalf when the field is absent, null,
or incorrectly shaped, matching the existing vkBudgetUsage pattern instead of
returning zero.
---
Nitpick comments:
In `@tests/governance/providerbudget_test.go`:
- Around line 256-259: Update the test around its provider-governance setup to
read and preserve the existing OpenAI governance before making changes. In the
deferred cleanup, restore the saved configuration instead of unconditionally
deleting it; if no governance existed originally, delete the test-created entry.
Apply the same preservation logic to the additional cleanup block noted in the
comment, using the existing governance request helpers and configurable baseURL
flow.
In `@tests/governance/teambudget_test.go`:
- Around line 427-441: Generalize waitForTeamBudgetUsage into a shared
waitForBudgetUsage helper accepting a label and reader function, centralizing
the 30-second polling timeout and failure message; retain a thin team wrapper if
useful. In tests/governance/teambudget_test.go:427-441 update the helper.
Replace each inline loop with the shared helper in
tests/governance/customerbudget_test.go:448-459 and :541-552 using
customerBudgetUsage closures, and in
tests/governance/providerbudget_test.go:284-295 and :438-449 using
providerGovernanceUsage and modelConfigUsage closures respectively.
In `@tests/governance/vkbudget_test.go`:
- Around line 567-580: Update the retry loop around completion() to track the
number of attempts alongside the existing deadline, and stop retrying after a
small bounded maximum. Preserve immediate exit when a response is rejected and
retain the existing diagnostic failure message when the deadline or attempt
limit is reached.
🪄 Autofix
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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 61faa0bb-8021-4f53-a2cd-f4350d45a32b
📒 Files selected for processing (5)
tests/governance/customerbudget_test.gotests/governance/providerbudget_test.gotests/governance/teambudget_test.gotests/governance/test_utils.gotests/governance/vkbudget_test.go
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 04:37
d986b7a to
f4e3169
Compare
akshaydeo
force-pushed
the
calendar-align-from-next-period
branch
from
August 10, 2026 04:37
ce56c83 to
9f91c19
Compare
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
transports/bifrost-http/handlers/governance.go (3)
2223-2247: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winThe early return on reset failure skips the OAuth and MCP credential reconciliation.
If
ResetBudgetUsageInMemoryfails, the handler returns before theReconcileOauthAfterVKChangeandReconcileMCPHeadersAfterVKChangecalls at Lines 2240-2247. Those calls are best-effort today (errors are only logged), so a reset failure now also drops credential reconciliation for an MCP allowlist change that already committed. Move the reset block after the reconciliation calls, or run the reconciliation before the reset.🔀 Proposed reordering
+ // Per-user credential reconciliation when the VK's MCP allowlist changed. + if req.MCPConfigs != nil && h.configStore != nil { + if err := h.configStore.ReconcileOauthAfterVKChange(ctx, vk.ID); err != nil { + logger.Error("reconcile OAuth credentials after VK %s update failed: %v", vk.ID, err) + } + if err := h.configStore.ReconcileMCPHeadersAfterVKChange(ctx, vk.ID); err != nil { + logger.Error("reconcile per-user-headers credentials after VK %s update failed: %v", vk.ID, err) + } + } if len(usageReset.budgetIDs) > 0 { if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, BudgetUsageResetOwner{Kind: BudgetOwnerVirtualKey, ID: vk.ID}, usageReset.budgetIDs); err != nil { logger.Error("failed to reset in-memory budget usage after update: %v", err) SendError(ctx, 500, "Virtual key updated but budget usage reset did not take effect") return } } - - // Per-user credential reconciliation ... (remove the original block here)🤖 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 `@transports/bifrost-http/handlers/governance.go` around lines 2223 - 2247, Move the ResetBudgetUsageInMemory block and its early return in the virtual-key update handler to after the ReconcileOauthAfterVKChange and ReconcileMCPHeadersAfterVKChange calls. Preserve the existing reset error response while ensuring both best-effort credential reconciliations run after a committed MCP allowlist change, even when the budget reset fails.
4069-4092: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winA reload failure here still runs the reset and still answers 200.
The other three handlers in this change now return 500 when the reload fails, and their comments state the reason: the reset is ordered after the reload and depends on it. This handler keeps the log-and-continue fallback at Lines 4070-4074, then runs
ResetBudgetUsageInMemoryat Line 4087 against in-memory state that was not refreshed. The response is 200 with the pre-update governance body.Track whether the reload succeeded, and skip the reset with a 500 when it did not.
🛡️ Proposed fix
resp := ProviderGovernanceResponse{Provider: providerName} + reloaded := false if deleted { if err := h.governanceManager.RemoveModelConfig(ctx, mc.ID); err != nil { logger.Error("failed to remove provider governance from memory: %v", err) } } else if len(mc.Budgets) > 0 || mc.RateLimitID != nil { - if reloaded, err := h.governanceManager.ReloadModelConfig(ctx, mc.ID); err != nil { + if reloadedMC, err := h.governanceManager.ReloadModelConfig(ctx, mc.ID); err != nil { logger.Error("failed to reload provider governance in memory: %v", err) if r, ok := modelConfigToProviderGovernance(&mc); ok { resp = r } - } else if r, ok := modelConfigToProviderGovernance(reloaded); ok { - resp = r + } else { + reloaded = true + if r, ok := modelConfigToProviderGovernance(reloadedMC); ok { + resp = r + } } } if !deleted && len(usageReset.budgetIDs) > 0 { + if !reloaded { + logger.Error("skipping budget usage reset for provider %s: in-memory reload failed", providerName) + SendError(ctx, 500, "Provider governance updated in database but failed to reload in-memory state") + return + } if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, BudgetUsageResetOwner{Kind: BudgetOwnerModelConfig, ID: mc.ID}, usageReset.budgetIDs); err != 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 `@transports/bifrost-http/handlers/governance.go` around lines 4069 - 4092, Update the provider governance reload branch around ReloadModelConfig to track whether reloading succeeded; on failure, log the error, send a 500 response, and return before ResetBudgetUsageInMemory runs. Only execute the reset and successful governance response path after a successful reload, removing the current fallback response based on stale mc state.
630-640: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAvoid inheriting usage when resetting budget usage.
When
usageReset.requestedis true, both reconcilers passfalsetoinheritUsageFromClosestShorterBudget. A new budget can therefore inheritCurrentUsageeven though the reset applies only to matched budgets. Pass the reset flag so new budgets start at zero.🤖 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 `@transports/bifrost-http/handlers/governance.go` around lines 630 - 640, Update the new-budget reconciliation branch around newBudgetFromRequest and inheritUsageFromClosestShorterBudget to pass usageReset.requested instead of false, ensuring newly created budgets start with zero usage when a reset is requested while preserving inheritance otherwise.
🧹 Nitpick comments (1)
transports/bifrost-http/handlers/governance.go (1)
69-89: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a named type for
Kindand align the doc comment with the constants.
Kindis a plainstring, so any caller can pass an arbitrary value and the compiler will accept it. A named type plus typed constants makes the contract explicit at every call site. The doc comment also lists"provider"as one of the mirrored entity types, but noBudgetOwnerProviderconstant exists; provider governance addressesBudgetOwnerModelConfigat Line 4087.♻️ Proposed refactor
-type BudgetUsageResetOwner struct { - Kind string - ID string -} +// BudgetOwnerKind mirrors the enterprise cluster entity type values. +type BudgetOwnerKind string + +type BudgetUsageResetOwner struct { + Kind BudgetOwnerKind + ID string +} // Budget owner kinds, matching the enterprise cluster entity type values. const ( - BudgetOwnerVirtualKey = "virtual_key" - BudgetOwnerTeam = "team" - BudgetOwnerCustomer = "customer" - BudgetOwnerModelConfig = "model_config" + BudgetOwnerVirtualKey BudgetOwnerKind = "virtual_key" + BudgetOwnerTeam BudgetOwnerKind = "team" + BudgetOwnerCustomer BudgetOwnerKind = "customer" + BudgetOwnerModelConfig BudgetOwnerKind = "model_config" )🤖 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 `@transports/bifrost-http/handlers/governance.go` around lines 69 - 89, Introduce a named budget-owner kind type and declare the existing owner constants with that type, then change BudgetUsageResetOwner.Kind to use it. Update the surrounding BudgetUsageResetOwner documentation to list only the supported constants, removing the unsupported provider entry while preserving the model-config mapping.
🤖 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.
Outside diff comments:
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 2223-2247: Move the ResetBudgetUsageInMemory block and its early
return in the virtual-key update handler to after the
ReconcileOauthAfterVKChange and ReconcileMCPHeadersAfterVKChange calls. Preserve
the existing reset error response while ensuring both best-effort credential
reconciliations run after a committed MCP allowlist change, even when the budget
reset fails.
- Around line 4069-4092: Update the provider governance reload branch around
ReloadModelConfig to track whether reloading succeeded; on failure, log the
error, send a 500 response, and return before ResetBudgetUsageInMemory runs.
Only execute the reset and successful governance response path after a
successful reload, removing the current fallback response based on stale mc
state.
- Around line 630-640: Update the new-budget reconciliation branch around
newBudgetFromRequest and inheritUsageFromClosestShorterBudget to pass
usageReset.requested instead of false, ensuring newly created budgets start with
zero usage when a reset is requested while preserving inheritance otherwise.
---
Nitpick comments:
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 69-89: Introduce a named budget-owner kind type and declare the
existing owner constants with that type, then change BudgetUsageResetOwner.Kind
to use it. Update the surrounding BudgetUsageResetOwner documentation to list
only the supported constants, removing the unsupported provider entry while
preserving the model-config mapping.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: d7659c60-6e4a-40fd-b71a-890b48e29962
📒 Files selected for processing (9)
docs/features/governance/budget-and-limits.mdxtests/governance/customerbudget_test.gotests/governance/providerbudget_test.gotests/governance/teambudget_test.gotests/governance/vkbudget_test.gotransports/bifrost-http/handlers/governance.goui/app/workspace/governance/views/teamSheet.tsxui/app/workspace/model-limits/views/modelLimitSheet.tsxui/lib/utils/governance.test.ts
🚧 Files skipped from review as they are similar to previous changes (7)
- docs/features/governance/budget-and-limits.mdx
- tests/governance/vkbudget_test.go
- tests/governance/teambudget_test.go
- ui/app/workspace/governance/views/teamSheet.tsx
- tests/governance/customerbudget_test.go
- tests/governance/providerbudget_test.go
- ui/app/workspace/model-limits/views/modelLimitSheet.tsx
akshaydeo
force-pushed
the
calendar-align-from-next-period
branch
from
August 10, 2026 05:02
9f91c19 to
9c813ec
Compare
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 05:02
f4e3169 to
e9e9222
Compare
akshaydeo
force-pushed
the
calendar-align-from-next-period
branch
from
August 10, 2026 06:48
9c813ec to
bb014fd
Compare
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 06:48
e9e9222 to
8c9f092
Compare
danpiths
previously approved these changes
Aug 10, 2026
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 06:57
8c9f092 to
772dd41
Compare
akshaydeo
force-pushed
the
calendar-align-from-next-period
branch
2 times, most recently
from
August 10, 2026 07:19
653e1a9 to
d9d2e2b
Compare
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 07:19
772dd41 to
7cea97e
Compare
17 tasks
This was referenced Aug 10, 2026
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 21:38
7cea97e to
d6f98e1
Compare
akshaydeo
force-pushed
the
calendar-align-from-next-period
branch
from
August 10, 2026 21:39
d9d2e2b to
37ee88e
Compare
Contributor
Author
Merge activity
|
akshaydeo
changed the base branch from
calendar-align-from-next-period
to
graphite-base/6004
August 10, 2026 21:58
akshaydeo
dismissed stale reviews from danpiths and coderabbitai[bot]
August 10, 2026 22:01
The base branch was changed.
…ts and provider governance
akshaydeo
force-pushed
the
reset-budget-usage-all-owners
branch
from
August 10, 2026 22:01
d6f98e1 to
bbdc667
Compare
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
atharvamhaske
pushed a commit
to atharvamhaske/bifrost
that referenced
this pull request
Aug 13, 2026
…ts and provider governance (maximhq#6004) ## Summary Adds an explicit `reset_budget_usage` flag to budget update endpoints so operators can zero accumulated spend on a budget without waiting for the next scheduled reset. Previously, updating a budget's amount or reset frequency had no way to clear existing usage. The UI now intercepts saves that change budget configuration on existing entities and asks whether to preserve or reset the counters before submitting. ## Changes - Added `reset_budget_usage: boolean` to `UpdateVirtualKeyRequest`, `UpdateTeamRequest`, `UpdateCustomerRequest`, `UpdateModelConfigRequest`, and `UpdateProviderGovernanceRequest` in both the Go handler structs and the OpenAPI/YAML schemas. - Wired the flag through each entity's budget reconciliation path (`reconcileCustomerBudgets`, `reconcileModelConfigBudgets`, and the inline team reconciliation loop). When set, the reconciler calls the store's usage-zero method for each matched budget ID and reflects the cleared value on the in-memory struct. - After the DB transaction commits and the entity is reloaded (which deliberately carries cached usage forward), `ResetBudgetUsageInMemory` is called with a typed `BudgetUsageResetOwner` struct instead of a bare virtual-key ID string. This lets Enterprise address the cluster broadcast to the correct entity type (`virtual_key`, `team`, `customer`, `model_config`). - Introduced `BudgetUsageResetOwner` and its `BudgetOwner*` constants so the owner kind values match the enterprise cluster entity type strings directly, avoiding a translation table. - Added a shared `BudgetUsageResetDialog` React component that asks **Preserve Usage** / **Reset Usage** and is reused across all four edit sheets rather than duplicated. - Added a `useBudgetUsageResetPrompt<T>` hook that parks the form payload before the dialog opens, preventing the form from changing underneath the dialog and producing a submission the operator never reviewed. - Wired the dialog and hook into the team, customer, model limit, and provider governance edit sheets. Each sheet detects whether budget amounts or durations changed on an existing entity and, if so, defers the save until the operator makes a choice. - Updated the `budget-and-limits.mdx` doc with a **Resetting budget usage** section covering the API field, a JSON example, supported owners, and a note clarifying that only usage is cleared while the reset window is left intact. - Updated the OpenAPI description for `reset_budget_usage` to clarify that `last_reset` only ever advances and is never moved as a side effect of a configuration write. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [ ] Plugins - [x] UI (React) - [x] Docs ## How to test ```sh # Core/Transports go test ./transports/bifrost-http/... # UI cd ui pnpm i pnpm build ``` 1. Create a virtual key, team, customer, model limit, or provider governance entry with a budget and accumulate some spend. 2. Edit the budget amount or reset duration and save. 3. The UI should present a dialog asking **Preserve Usage** or **Reset Usage**. 4. Choosing **Reset Usage** should zero `current_usage` on the budget immediately without changing `last_reset` or the window boundaries. 5. Choosing **Preserve Usage** should save the configuration change and leave `current_usage` untouched. 6. Sending `reset_budget_usage: true` directly via the API should produce the same outcome without the dialog. ## Screenshots/Recordings _Add before/after screenshots of the reset dialog if available._ ## Breaking changes - [ ] Yes - [x] No ## Related issues ## Security considerations The reset is gated behind the same RBAC permissions as the update operation for each entity type. No new privilege surface is introduced. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
This was referenced Aug 13, 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
2 participants
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
Adds an explicit
reset_budget_usageflag to budget update endpoints so operators can zero accumulated spend on a budget without waiting for the next scheduled reset. Previously, updating a budget's amount or reset frequency had no way to clear existing usage. The UI now intercepts saves that change budget configuration on existing entities and asks whether to preserve or reset the counters before submitting.Changes
reset_budget_usage: booleantoUpdateVirtualKeyRequest,UpdateTeamRequest,UpdateCustomerRequest,UpdateModelConfigRequest, andUpdateProviderGovernanceRequestin both the Go handler structs and the OpenAPI/YAML schemas.reconcileCustomerBudgets,reconcileModelConfigBudgets, and the inline team reconciliation loop). When set, the reconciler calls the store's usage-zero method for each matched budget ID and reflects the cleared value on the in-memory struct.ResetBudgetUsageInMemoryis called with a typedBudgetUsageResetOwnerstruct instead of a bare virtual-key ID string. This lets Enterprise address the cluster broadcast to the correct entity type (virtual_key,team,customer,model_config).BudgetUsageResetOwnerand itsBudgetOwner*constants so the owner kind values match the enterprise cluster entity type strings directly, avoiding a translation table.BudgetUsageResetDialogReact component that asks Preserve Usage / Reset Usage and is reused across all four edit sheets rather than duplicated.useBudgetUsageResetPrompt<T>hook that parks the form payload before the dialog opens, preventing the form from changing underneath the dialog and producing a submission the operator never reviewed.budget-and-limits.mdxdoc with a Resetting budget usage section covering the API field, a JSON example, supported owners, and a note clarifying that only usage is cleared while the reset window is left intact.reset_budget_usageto clarify thatlast_resetonly ever advances and is never moved as a side effect of a configuration write.Type of change
Affected areas
How to test
current_usageon the budget immediately without changinglast_resetor the window boundaries.current_usageuntouched.reset_budget_usage: truedirectly via the API should produce the same outcome without the dialog.Screenshots/Recordings
Add before/after screenshots of the reset dialog if available.
Breaking changes
Related issues
Security considerations
The reset is gated behind the same RBAC permissions as the update operation for each entity type. No new privilege surface is introduced.
Checklist
docs/contributing/README.mdand followed the guidelines