From 2119010ef458bbe8706d77c313771bd04a78694b Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Sun, 9 Aug 2026 15:53:08 -0700 Subject: [PATCH 1/4] governance: extend reset_budget_usage to teams, customers, model limits and provider governance --- .../features/governance/budget-and-limits.mdx | 21 +++ docs/openapi/openapi.json | 18 ++- .../schemas/management/governance.yaml | 29 +++- .../bifrost-http/handlers/governance.go | 126 ++++++++++++++++-- .../handlers/pricing_override_test.go | 2 +- transports/bifrost-http/server/server.go | 10 +- .../governance/views/customerSheet.tsx | 34 +++++ .../workspace/governance/views/teamSheet.tsx | 39 ++++++ .../model-limits/views/modelLimitSheet.tsx | 34 +++++ .../fragments/governanceFormFragment.tsx | 35 +++++ ui/components/ui/budgetUsageResetDialog.tsx | 61 +++++++++ ui/hooks/useBudgetUsageResetPrompt.ts | 49 +++++++ ui/lib/types/governance.ts | 8 ++ 13 files changed, 447 insertions(+), 19 deletions(-) create mode 100644 ui/components/ui/budgetUsageResetDialog.tsx create mode 100644 ui/hooks/useBudgetUsageResetPrompt.ts diff --git a/docs/features/governance/budget-and-limits.mdx b/docs/features/governance/budget-and-limits.mdx index ca242699134..6d036c856a2 100644 --- a/docs/features/governance/budget-and-limits.mdx +++ b/docs/features/governance/budget-and-limits.mdx @@ -219,6 +219,27 @@ Quarter boundaries repeat every three months, so `quarter_start_month` only chan Changing `quarter_start_month` on a live calendar-aligned budget takes effect on the next reset tick rather than instantly. The new definition is stored immediately, which moves where the current window starts; if that boundary has moved forward, the budget reads as due and resets shortly after, zeroing usage. That is the honest outcome - under the new calendar the current quarter genuinely began on a later date. +### Resetting budget usage + +Changing a budget's amount or reset frequency does not clear the spend already recorded against it. To +clear it, send `reset_budget_usage: true` on the update request. The UI asks before saving whenever a +budget's configuration changed, offering **Preserve Usage** or **Reset Usage**. + +```json +{ + "budgets": [{ "max_limit": 5000, "reset_duration": "1M" }], + "reset_budget_usage": true +} +``` + +Supported on virtual keys, teams, customers, model limits, and provider governance. + + +The reset clears **usage only**. The reset window keeps its existing start and end, because `last_reset` +only ever advances and never moves as a side effect of a configuration write. That rule is what keeps +every node in a cluster agreeing on which window is currently open. + + ### Budget overrides A budget can carry a temporary **override** that adds spending capacity on top of its configured limit without touching the base limit, current usage, or reset schedule. While an override is active, enforcement uses: diff --git a/docs/openapi/openapi.json b/docs/openapi/openapi.json index d81ac482668..acb6f27b55e 100644 --- a/docs/openapi/openapi.json +++ b/docs/openapi/openapi.json @@ -82486,7 +82486,7 @@ }, "reset_budget_usage": { "type": "boolean", - "description": "When true, resets usage for virtual-key and provider-config budget records reconciled by this update." + "description": "When true, zeroes current usage on the budgets reconciled by this update. The reset window is left alone: last_reset only ever advances, and never as a side effect of a configuration write, so the current window keeps its existing start and end.\n" }, "expires_at": { "type": "string", @@ -82620,6 +82620,10 @@ "items": { "$ref": "#/components/schemas/CreateBudgetRequest" } + }, + "reset_budget_usage": { + "type": "boolean", + "description": "When true, zeroes current usage on the budgets reconciled by this update. The reset window is left alone: last_reset only ever advances, and never as a side effect of a configuration write, so the current window keeps its existing start and end.\n" } } }, @@ -82723,6 +82727,10 @@ }, "budget": { "$ref": "#/components/schemas/UpdateBudgetRequest" + }, + "reset_budget_usage": { + "type": "boolean", + "description": "When true, zeroes current usage on the budgets reconciled by this update. The reset window is left alone: last_reset only ever advances, and never as a side effect of a configuration write, so the current window keeps its existing start and end.\n" } } }, @@ -84010,6 +84018,10 @@ "rate_limit": { "$ref": "#/components/schemas/UpdateRateLimitRequest", "description": "Rate limit configuration" + }, + "reset_budget_usage": { + "type": "boolean", + "description": "When true, zeroes current usage on the budgets reconciled by this update. The reset window is left alone: last_reset only ever advances, and never as a side effect of a configuration write, so the current window keeps its existing start and end.\n" } } }, @@ -84101,6 +84113,10 @@ "type": "boolean", "nullable": true, "description": "When true, all budgets reset at clean calendar boundaries. Omit to leave unchanged." + }, + "reset_budget_usage": { + "type": "boolean", + "description": "When true, zeroes current usage on the budgets reconciled by this update. The reset window is left alone: last_reset only ever advances, and never as a side effect of a configuration write, so the current window keeps its existing start and end.\n" } } }, diff --git a/docs/openapi/schemas/management/governance.yaml b/docs/openapi/schemas/management/governance.yaml index 004946f1439..922833f2e18 100644 --- a/docs/openapi/schemas/management/governance.yaml +++ b/docs/openapi/schemas/management/governance.yaml @@ -563,7 +563,10 @@ UpdateVirtualKeyRequest: type: boolean reset_budget_usage: type: boolean - description: When true, resets usage for virtual-key and provider-config budget records reconciled by this update. + description: > + When true, zeroes current usage on the budgets reconciled by this update. The reset window is + left alone: last_reset only ever advances, and never as a side effect of a configuration write, + so the current window keeps its existing start and end. expires_at: type: string description: RFC3339 timestamp to set a new expiry (must be in the future), empty string to clear an existing expiry, omitted to leave it unchanged. @@ -787,6 +790,12 @@ UpdateTeamRequest: description: Replacement set of budgets for this team; reset durations must be unique items: $ref: '#/CreateBudgetRequest' + reset_budget_usage: + type: boolean + description: > + When true, zeroes current usage on the budgets reconciled by this update. The reset window is + left alone: last_reset only ever advances, and never as a side effect of a configuration write, + so the current window keeps its existing start and end. ListTeamsResponse: type: object @@ -864,6 +873,12 @@ UpdateCustomerRequest: type: string budget: $ref: '#/UpdateBudgetRequest' + reset_budget_usage: + type: boolean + description: > + When true, zeroes current usage on the budgets reconciled by this update. The reset window is + left alone: last_reset only ever advances, and never as a side effect of a configuration write, + so the current window keeps its existing start and end. ListCustomersResponse: type: object @@ -1518,6 +1533,12 @@ UpdateModelConfigRequest: rate_limit: $ref: '#/UpdateRateLimitRequest' description: Rate limit configuration + reset_budget_usage: + type: boolean + description: > + When true, zeroes current usage on the budgets reconciled by this update. The reset window is + left alone: last_reset only ever advances, and never as a side effect of a configuration write, + so the current window keeps its existing start and end. # Provider Governance @@ -1593,6 +1614,12 @@ UpdateProviderGovernanceRequest: type: boolean nullable: true description: When true, all budgets reset at clean calendar boundaries. Omit to leave unchanged. + reset_budget_usage: + type: boolean + description: > + When true, zeroes current usage on the budgets reconciled by this update. The reset window is + left alone: last_reset only ever advances, and never as a side effect of a configuration write, + so the current window keeps its existing start and end. # Pricing Overrides diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 5c5c07ca8b1..3e0c8396ee5 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -49,8 +49,8 @@ type GovernanceManager interface { ReloadVirtualKey(ctx context.Context, id string) (*configstoreTables.TableVirtualKey, error) // ResetBudgetUsageInMemory clears usage for the given budgets in the store that // enforces spend, leaving each reset boundary untouched. Enterprise also - // propagates the reset to cluster peers. - ResetBudgetUsageInMemory(ctx context.Context, vkID string, budgetIDs []string) error + // propagates the reset to cluster peers, which is what owner addresses. + ResetBudgetUsageInMemory(ctx context.Context, owner BudgetUsageResetOwner, budgetIDs []string) error RemoveVirtualKey(ctx context.Context, id string) error ReloadTeam(ctx context.Context, id string) (*configstoreTables.TableTeam, error) RemoveTeam(ctx context.Context, id string) error @@ -66,6 +66,27 @@ type GovernanceManager interface { DeletePricingOverride(ctx context.Context, id string) error } +// BudgetUsageResetOwner identifies the entity whose budgets had their usage reset. +// +// Enterprise needs this to address the cluster broadcast: a peer has to reload the +// right entity and then apply the reset, because a plain reload deliberately +// preserves that peer's cached usage. Kind deliberately uses the same strings as +// the enterprise cluster entity types ("virtual_key", "team", "customer", +// "model_config", "provider") so the mapping is a direct conversion rather than a +// lookup table that can drift. +type BudgetUsageResetOwner struct { + Kind string + ID string +} + +// Budget owner kinds, matching the enterprise cluster entity type values. +const ( + BudgetOwnerVirtualKey = "virtual_key" + BudgetOwnerTeam = "team" + BudgetOwnerCustomer = "customer" + BudgetOwnerModelConfig = "model_config" +) + type complexityAnalyzerConfigReloader interface { // HTTP server bridge signature: BifrostHTTPServer implements this and adapts // to the governance plugin's in-memory ReloadComplexityAnalyzerConfig(config). @@ -642,7 +663,7 @@ func (h *GovernanceHandler) reconcileModelConfigBudgets(ctx context.Context, tx // reconcileCustomerBudgets upserts the desired set of budgets owned by a customer // (via TableBudget.CustomerID), preserving usage on matched rows and deleting removed ones. // It mutates customer.Budgets to the reconciled set. Mirrors reconcileModelConfigBudgets. -func (h *GovernanceHandler) reconcileCustomerBudgets(ctx context.Context, tx *gorm.DB, customer *configstoreTables.TableCustomer, requests []CreateBudgetRequest) error { +func (h *GovernanceHandler) reconcileCustomerBudgets(ctx context.Context, tx *gorm.DB, customer *configstoreTables.TableCustomer, requests []CreateBudgetRequest, usageReset *budgetUsageReset) error { seenDurations := make(map[string]bool, len(requests)) for _, b := range requests { if b.MaxLimit < 0 { @@ -675,6 +696,12 @@ func (h *GovernanceHandler) reconcileCustomerBudgets(ctx context.Context, tx *go if err := h.configStore.UpdateBudget(ctx, &existing, tx); err != nil { return err } + if err := usageReset.apply(ctx, h.configStore, existing.ID, tx); err != nil { + return err + } + if usageReset != nil && usageReset.requested { + existing.CurrentUsage = 0 + } reconciled = append(reconciled, existing) matchedIDs[existing.ID] = true } else { @@ -1071,6 +1098,9 @@ type UpdateTeamRequest struct { Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: replaces all team budgets RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` CalendarAligned *bool `json:"calendar_aligned,omitempty"` // Team-wide setting; nil means "leave unchanged" + // ResetBudgetUsage zeroes current usage on the reconciled budgets when true. + // The reset boundary is left alone; only accumulated spend is cleared. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // CreateCustomerRequest represents the request body for creating a customer @@ -1089,6 +1119,8 @@ type UpdateCustomerRequest struct { Budget *UpdateBudgetRequest `json:"budget,omitempty"` // Deprecated: use budgets RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` CalendarAligned *bool `json:"calendar_aligned,omitempty"` + // ResetBudgetUsage zeroes current usage on the reconciled budgets when true. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // CreateModelConfigRequest represents the request body for creating a model config @@ -1109,6 +1141,8 @@ type UpdateModelConfigRequest struct { Provider *string `json:"provider,omitempty"` // Optional provider, nil means no change Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Full desired set of budgets (reconciled against existing) RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` + // ResetBudgetUsage zeroes current usage on the reconciled budgets when true. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // UpdateProviderGovernanceRequest represents the request body for updating provider governance @@ -1117,6 +1151,8 @@ type UpdateProviderGovernanceRequest struct { Budgets *[]CreateBudgetRequest `json:"budgets,omitempty"` // nil=no change, []=remove all RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` CalendarAligned *bool `json:"calendar_aligned,omitempty"` + // ResetBudgetUsage zeroes current usage on the reconciled budgets when true. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // RegisterRoutes registers all governance-related routes for the new hierarchical system @@ -2189,7 +2225,7 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { // and deliberately carries the cached CurrentUsage forward, which would undo // the reset. Enterprise additionally propagates this to cluster peers. if len(usageReset.budgetIDs) > 0 { - if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, vk.ID, usageReset.budgetIDs); err != nil { + 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 @@ -2541,6 +2577,10 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Invalid JSON") return } + // The operator's explicit "reset usage" choice. Carried into budget + // reconciliation and collected there, so the in-memory store can be cleared + // once the transaction commits. + usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage} // Fetching team from database team, err := h.configStore.GetTeam(ctx, teamID) if err != nil { @@ -2607,15 +2647,23 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { if existing, found := existingByDuration[b.ResetDuration]; found { existing.MaxLimit = b.MaxLimit applyResetConfigToExistingBudget(&existing, b) - // LastReset / CurrentUsage are preserved on update; if calendar - // alignment was just enabled in this request, the post-reconciliation - // snap block below resets them. + // LastReset is preserved on update: the reset boundary only ever + // moves forward, and never as a side effect of a config write. + // CurrentUsage is preserved too unless the operator explicitly + // asked for a reset, which UpdateBudget cannot carry and so goes + // through the store method that owns the column. if err := validateBudget(&existing); err != nil { return err } if err := h.configStore.UpdateBudget(ctx, &existing, tx); err != nil { return err } + if err := usageReset.apply(ctx, h.configStore, existing.ID, tx); err != nil { + return err + } + if usageReset.requested { + existing.CurrentUsage = 0 + } reconciledBudgets = append(reconciledBudgets, existing) matchedIDs[existing.ID] = true } else { @@ -2720,6 +2768,16 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { logger.Error("failed to reload team: %v", err) preloadedTeam = team } + // Clear usage in the store that enforces spend. This runs after the reload, + // which deliberately carries each budget's cached usage forward and would + // otherwise undo the reset. Enterprise propagates it to cluster peers. + if len(usageReset.budgetIDs) > 0 { + if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, BudgetUsageResetOwner{Kind: BudgetOwnerTeam, ID: team.ID}, usageReset.budgetIDs); err != nil { + logger.Error("failed to reset in-memory budget usage for Team: %v", err) + SendError(ctx, 500, "Team updated but budget usage reset did not take effect") + return + } + } SendJSON(ctx, map[string]interface{}{ "message": "Team updated successfully", "team": preloadedTeam, @@ -2858,7 +2916,7 @@ func (h *GovernanceHandler) createCustomer(ctx *fasthttp.RequestCtx) { return err } if len(budgetRequests) > 0 { - if err := h.reconcileCustomerBudgets(ctx, tx, &customer, budgetRequests); err != nil { + if err := h.reconcileCustomerBudgets(ctx, tx, &customer, budgetRequests, nil); err != nil { return err } } @@ -2930,6 +2988,10 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Invalid JSON") return } + // The operator's explicit "reset usage" choice. Carried into budget + // reconciliation and collected there, so the in-memory store can be cleared + // once the transaction commits. + usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage} if req.Budgets != nil && req.Budget != nil { SendError(ctx, 400, "only one of 'budget' or 'budgets' may be set") return @@ -2971,7 +3033,7 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { } } if effectiveBudgets != nil { - if err := h.reconcileCustomerBudgets(ctx, tx, customer, *effectiveBudgets); err != nil { + if err := h.reconcileCustomerBudgets(ctx, tx, customer, *effectiveBudgets, usageReset); err != nil { return err } } @@ -3051,6 +3113,16 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { logger.Error("failed to reload customer: %v", err) preloadedCustomer = customer } + // Clear usage in the store that enforces spend. This runs after the reload, + // which deliberately carries each budget's cached usage forward and would + // otherwise undo the reset. Enterprise propagates it to cluster peers. + if len(usageReset.budgetIDs) > 0 { + if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, BudgetUsageResetOwner{Kind: BudgetOwnerCustomer, ID: customer.ID}, usageReset.budgetIDs); err != nil { + logger.Error("failed to reset in-memory budget usage for Customer: %v", err) + SendError(ctx, 500, "Customer updated but budget usage reset did not take effect") + return + } + } SendJSON(ctx, map[string]interface{}{ "message": "Customer updated successfully", @@ -3559,6 +3631,10 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Invalid JSON") return } + // The operator's explicit "reset usage" choice. Carried into budget + // reconciliation and collected there, so the in-memory store can be cleared + // once the transaction commits. + usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage} mc, err := h.configStore.GetModelConfigByID(ctx, mcID) if err != nil { if errors.Is(err, configstore.ErrNotFound) { @@ -3584,7 +3660,7 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) { // slice removes all budgets; omitting the field leaves them unchanged. Budgets // are owned via ModelConfigID, so no model-config FK juggling is needed. if req.Budgets != nil { - if err := h.reconcileModelConfigBudgets(ctx, tx, mc, req.Budgets, nil); err != nil { + if err := h.reconcileModelConfigBudgets(ctx, tx, mc, req.Budgets, usageReset); err != nil { return err } } @@ -3662,6 +3738,16 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) { logger.Error("failed to reload model config in memory: %v", err) updatedMC = mc } + // Clear usage in the store that enforces spend. This runs after the reload, + // which deliberately carries each budget's cached usage forward and would + // otherwise undo the reset. Enterprise propagates it to cluster peers. + if len(usageReset.budgetIDs) > 0 { + if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, BudgetUsageResetOwner{Kind: BudgetOwnerModelConfig, ID: mc.ID}, usageReset.budgetIDs); err != nil { + logger.Error("failed to reset in-memory budget usage for Model config: %v", err) + SendError(ctx, 500, "Model config updated but budget usage reset did not take effect") + return + } + } h.resolveModelConfigScopeName(ctx, updatedMC, map[string]string{}) SendJSON(ctx, map[string]interface{}{ "message": "Model config updated successfully", @@ -3783,6 +3869,10 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Invalid JSON") return } + // The operator's explicit "reset usage" choice. Carried into budget + // reconciliation and collected there, so the in-memory store can be cleared + // once the transaction commits. + usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage} if req.Budget != nil && req.Budgets != nil { SendError(ctx, 400, "only one of 'budget' or 'budgets' may be set") return @@ -3932,7 +4022,7 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) { // Budget reconciliation (mc row exists at this point for create cases). if !deleted && effectiveBudgets != nil { - if err := h.reconcileModelConfigBudgets(ctx, tx, &mc, *effectiveBudgets, nil); err != nil { + if err := h.reconcileModelConfigBudgets(ctx, tx, &mc, *effectiveBudgets, usageReset); err != nil { return err } } @@ -3971,6 +4061,20 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) { resp = r } } + // Clear usage in the store that enforces spend. This runs after the reload, + // which deliberately carries each budget's cached usage forward and would + // otherwise undo the reset. Skipped when the governance was deleted outright, + // since there is nothing left to reset. + if !deleted && len(usageReset.budgetIDs) > 0 { + // Provider governance stores its budgets on a VK-agnostic model config, which + // is also what the reload above refreshes, so the reset is addressed to that + // model config rather than to the provider itself. + if err := h.governanceManager.ResetBudgetUsageInMemory(ctx, BudgetUsageResetOwner{Kind: BudgetOwnerModelConfig, ID: mc.ID}, usageReset.budgetIDs); err != nil { + logger.Error("failed to reset in-memory budget usage for provider governance: %v", err) + SendError(ctx, 500, "Provider governance updated but budget usage reset did not take effect") + return + } + } SendJSON(ctx, map[string]interface{}{ "message": "Provider governance updated successfully", "provider": resp, diff --git a/transports/bifrost-http/handlers/pricing_override_test.go b/transports/bifrost-http/handlers/pricing_override_test.go index 17649143282..4cb571e6e31 100644 --- a/transports/bifrost-http/handlers/pricing_override_test.go +++ b/transports/bifrost-http/handlers/pricing_override_test.go @@ -23,7 +23,7 @@ type pricingOverrideTestGovernanceManager struct{} func (pricingOverrideTestGovernanceManager) GetGovernanceData(ctx context.Context) *governance.GovernanceData { return nil } -func (pricingOverrideTestGovernanceManager) ResetBudgetUsageInMemory(context.Context, string, []string) error { +func (pricingOverrideTestGovernanceManager) ResetBudgetUsageInMemory(context.Context, BudgetUsageResetOwner, []string) error { return nil } func (pricingOverrideTestGovernanceManager) ReloadVirtualKey(context.Context, string) (*configstoreTables.TableVirtualKey, error) { diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 66c1ddda3bf..72750d39a08 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -101,10 +101,10 @@ type ServerCallbacks interface { // Virtual key related callbacks ReloadVirtualKey(ctx context.Context, id string) (*tables.TableVirtualKey, error) // ResetBudgetUsageInMemory clears usage for the given budgets in the governance - // store, leaving each reset boundary untouched. vkID identifies the owning - // virtual key so enterprise can address the cluster broadcast that propagates - // the reset to peers. - ResetBudgetUsageInMemory(ctx context.Context, vkID string, budgetIDs []string) error + // store, leaving each reset boundary untouched. owner identifies the entity that + // owns them so enterprise can address the cluster broadcast that propagates the + // reset to peers. + ResetBudgetUsageInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string) error RemoveVirtualKey(ctx context.Context, id string) error // Provider related callbacks GetModelsForProvider(provider schemas.ModelProvider) []string @@ -482,7 +482,7 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t // // Missing budgets are not an error: a budget can legitimately have been deleted // in the same request that asked for the reset. -func (s *BifrostHTTPServer) ResetBudgetUsageInMemory(ctx context.Context, vkID string, budgetIDs []string) error { +func (s *BifrostHTTPServer) ResetBudgetUsageInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string) error { if len(budgetIDs) == 0 { return nil } diff --git a/ui/app/workspace/governance/views/customerSheet.tsx b/ui/app/workspace/governance/views/customerSheet.tsx index e780560aabd..7c47e74d127 100644 --- a/ui/app/workspace/governance/views/customerSheet.tsx +++ b/ui/app/workspace/governance/views/customerSheet.tsx @@ -12,6 +12,8 @@ import { CopyableId } from "@/components/copyableId"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import BudgetUsageResetDialog from "@/components/ui/budgetUsageResetDialog"; +import { useBudgetUsageResetPrompt } from "@/hooks/useBudgetUsageResetPrompt"; import MultiBudgetLines, { BudgetLineEntry } from "@/components/ui/multibudgets"; import NumberAndSelect from "@/components/ui/numberAndSelect"; import { Sheet, SheetContent, SheetDescription, SheetFooter, SheetHeader, SheetTitle } from "@/components/ui/sheet"; @@ -70,6 +72,8 @@ export default function CustomerSheet({ open, onOpenChange, customer, onSuccess const [nameError, setNameError] = useState(null); const [showCalendarAlignWarning, setShowCalendarAlignWarning] = useState(false); + // Defers the save until the operator says whether to clear accumulated spend. + const resetPrompt = useBudgetUsageResetPrompt(); const hasCreateAccess = useRbac(RbacResource.Customers, RbacOperation.Create); const hasUpdateAccess = useRbac(RbacResource.Customers, RbacOperation.Update); @@ -188,6 +192,19 @@ export default function CustomerSheet({ open, onOpenChange, customer, onSuccess setFormData((prev) => ({ ...prev, [field]: value })); }; + // A budget config change on an existing customer is when clearing accumulated + // spend becomes a meaningful choice; creating one has no usage to reset. + const budgetsChanged = () => { + if (!isEditing || !customer) return false; + const signature = (rows: { max_limit?: number | null; reset_duration?: string }[]) => + [...rows] + .map((r) => `${r.max_limit ?? ""}:${r.reset_duration ?? ""}`) + .sort() + .join("|"); + const next = formData.budgets.filter((b) => b.max_limit !== undefined && b.max_limit !== null); + return signature(next) !== signature(customer.budgets ?? []); + }; + const handleSubmit = async (e: React.FormEvent) => { e.preventDefault(); @@ -196,6 +213,14 @@ export default function CustomerSheet({ open, onOpenChange, customer, onSuccess return; } + if (budgetsChanged()) { + resetPrompt.ask(true); + return; + } + await saveCustomer(false); + }; + + const saveCustomer = async (resetBudgetUsage: boolean) => { const budgetRequests: CreateBudgetRequest[] = formData.budgets .filter((b) => b.max_limit !== undefined && b.max_limit !== null) .map((b) => ({ id: b.id, max_limit: b.max_limit!, reset_duration: b.reset_duration })); @@ -206,6 +231,8 @@ export default function CustomerSheet({ open, onOpenChange, customer, onSuccess name: formData.name, calendar_aligned: formData.calendarAligned, budgets: budgetRequests, + // Only sent when the operator explicitly chose to clear spend. + reset_budget_usage: resetBudgetUsage || undefined, }; const hadRateLimit = !!customer.rate_limit; @@ -377,6 +404,13 @@ export default function CustomerSheet({ open, onOpenChange, customer, onSuccess + resetPrompt.resolve(() => saveCustomer(resetUsage))} + /> diff --git a/ui/app/workspace/governance/views/teamSheet.tsx b/ui/app/workspace/governance/views/teamSheet.tsx index 4d14e137938..1d3e4c8385c 100644 --- a/ui/app/workspace/governance/views/teamSheet.tsx +++ b/ui/app/workspace/governance/views/teamSheet.tsx @@ -18,6 +18,8 @@ import NumberAndSelect from "@/components/ui/numberAndSelect"; import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import BudgetUsageResetDialog from "@/components/ui/budgetUsageResetDialog"; +import { useBudgetUsageResetPrompt } from "@/hooks/useBudgetUsageResetPrompt"; import QuarterStartSelect from "@/components/ui/quarterStartSelect"; import { budgetResetDurationOptions, resetDurationOptions, supportsCalendarAlignment } from "@/lib/constants/governance"; import { getErrorMessage, useCreateTeamMutation, useUpdateTeamMutation } from "@/lib/store"; @@ -115,6 +117,10 @@ export default function TeamSheet({ team, onSave, onCancel }: TeamSheetProps) { // Team-wide calendar-align toggle: confirmation only fires on the off→on // transition for an existing team (mirrors the VK sheet behavior). const [showCalendarAlignWarning, setShowCalendarAlignWarning] = useState(false); + // Defers the save until the operator says whether to clear accumulated spend. + // The payload is a marker rather than the form data: this sheet keeps its own + // formData state, which the save reads directly. + const resetPrompt = useBudgetUsageResetPrompt(); const updateBudgetRow = (idx: number, patch: Partial) => { setFormData((prev) => { @@ -230,6 +236,22 @@ export default function TeamSheet({ team, onSave, onCancel }: TeamSheetProps) { setFormData((prev) => ({ ...prev, [field]: value })); }; + // Whether this save changes budget configuration on an existing team, which is + // when clearing accumulated spend is a meaningful choice. Creating a team has + // no usage to reset, and a save that leaves budgets alone should not ask. + const budgetsChanged = () => { + if (!isEditing || !team) return false; + const signature = (rows: { max_limit?: number | null; reset_duration?: string }[]) => + [...rows] + .map((r) => `${r.max_limit ?? ""}:${r.reset_duration ?? ""}`) + .sort() + .join("|"); + const next = formData.budgets + .filter((r) => r.maxLimit !== undefined && r.maxLimit !== null) + .map((r) => ({ max_limit: r.maxLimit, reset_duration: r.resetDuration })); + return signature(next) !== signature(team.budgets ?? []); + }; + const handleSubmit = async (e: React.FormEvent) => { e.preventDefault(); @@ -238,6 +260,14 @@ export default function TeamSheet({ team, onSave, onCancel }: TeamSheetProps) { return; } + if (budgetsChanged()) { + resetPrompt.ask(true); + return; + } + await saveTeam(false); + }; + + const saveTeam = async (resetBudgetUsage: boolean) => { // Serialize budget rows whose max_limit was filled in — rows left blank // are silently dropped (the backend treats the slice as authoritative). const submittableBudgets = formData.budgets @@ -259,6 +289,8 @@ export default function TeamSheet({ team, onSave, onCancel }: TeamSheetProps) { budgets: submittableBudgets, // Team-wide setting that governs both team budgets and the team rate limit. calendar_aligned: formData.calendarAligned, + // Only sent when the operator explicitly chose to clear spend. + reset_budget_usage: resetBudgetUsage || undefined, }; // Detect rate limit changes using had/has pattern @@ -528,6 +560,13 @@ export default function TeamSheet({ team, onSave, onCancel }: TeamSheetProps) { + resetPrompt.resolve(() => saveTeam(resetUsage))} + /> {/* Current Usage Section (only shown when editing with existing limits) */} {isEditing && ((team?.budgets && team.budgets.length > 0) || team?.rate_limit) && ( diff --git a/ui/app/workspace/model-limits/views/modelLimitSheet.tsx b/ui/app/workspace/model-limits/views/modelLimitSheet.tsx index 7dffeaf24bf..3409f33ceee 100644 --- a/ui/app/workspace/model-limits/views/modelLimitSheet.tsx +++ b/ui/app/workspace/model-limits/views/modelLimitSheet.tsx @@ -3,6 +3,8 @@ import { Form, FormControl, FormField, FormItem, FormLabel, FormMessage } from " import { Label } from "@/components/ui/label"; import { ModelMultiselect } from "@/components/ui/modelMultiselect"; import NumberAndSelect from "@/components/ui/numberAndSelect"; +import BudgetUsageResetDialog from "@/components/ui/budgetUsageResetDialog"; +import { useBudgetUsageResetPrompt } from "@/hooks/useBudgetUsageResetPrompt"; import MultiBudgetLines from "@/components/ui/multibudgets"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { DottedSeparator } from "@/components/ui/separator"; @@ -81,6 +83,8 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model const { data: providersData } = useGetProvidersQuery(); const [createModelConfig, { isLoading: isCreating }] = useCreateModelConfigMutation(); + // Defers the save until the operator says whether to clear accumulated spend. + const resetPrompt = useBudgetUsageResetPrompt(); const [updateModelConfig, { isLoading: isUpdating }] = useUpdateModelConfigMutation(); const [getModels] = useLazyGetModelsQuery(); const isLoading = isCreating || isUpdating; @@ -164,12 +168,33 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model } }, [modelConfig, form]); + // A budget config change on an existing limit is when clearing accumulated + // spend becomes a meaningful choice; creating one has no usage to reset. + const budgetsChanged = (data: FormData) => { + if (!isEditing || !modelConfig) return false; + const signature = (rows: { max_limit?: number | null; reset_duration?: string }[]) => + [...rows] + .map((r) => `${r.max_limit ?? ""}:${r.reset_duration ?? ""}`) + .sort() + .join("|"); + const next = (data.budgets ?? []).filter((b) => b.max_limit !== undefined && b.max_limit !== null); + return signature(next) !== signature(modelConfig.budgets ?? []); + }; + const onSubmit = async (data: FormData) => { if (!canSubmit) { toast.error("You don't have permission to perform this action"); return; } + if (budgetsChanged(data)) { + resetPrompt.ask(data); + return; + } + await saveModelLimit(data, false); + }; + + const saveModelLimit = async (data: FormData, resetBudgetUsage: boolean) => { if (!hasAnyLimit) { form.setError("root", { message: "At least one budget or rate limit is required" }); return; @@ -222,6 +247,8 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model provider: provider, budgets: budgetsPayload, rate_limit: rateLimitPayload, + // Only sent when the operator explicitly chose to clear spend. + reset_budget_usage: resetBudgetUsage || undefined, }, }).unwrap(); toast.success("Limit updated successfully"); @@ -546,6 +573,13 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model + resetPrompt.resolve((data) => saveModelLimit(data, resetUsage))} + /> ); diff --git a/ui/app/workspace/providers/fragments/governanceFormFragment.tsx b/ui/app/workspace/providers/fragments/governanceFormFragment.tsx index c26095ac8ab..f5c8846b19f 100644 --- a/ui/app/workspace/providers/fragments/governanceFormFragment.tsx +++ b/ui/app/workspace/providers/fragments/governanceFormFragment.tsx @@ -1,6 +1,8 @@ import { Button } from "@/components/ui/button"; import { Form } from "@/components/ui/form"; import { Label } from "@/components/ui/label"; +import BudgetUsageResetDialog from "@/components/ui/budgetUsageResetDialog"; +import { useBudgetUsageResetPrompt } from "@/hooks/useBudgetUsageResetPrompt"; import MultiBudgetLines, { BudgetLineEntry } from "@/components/ui/multibudgets"; import NumberAndSelect from "@/components/ui/numberAndSelect"; import { DottedSeparator } from "@/components/ui/separator"; @@ -69,6 +71,8 @@ function governanceToFormValues(provGov: ProviderGovernance | undefined): FormDa export function GovernanceFormFragment({ provider }: GovernanceFormFragmentProps) { const hasUpdateProviderAccess = useRbac(RbacResource.ModelProvider, RbacOperation.Update); + // Defers the save until the operator says whether to clear accumulated spend. + const resetPrompt = useBudgetUsageResetPrompt(); const hasViewAccess = useRbac(RbacResource.Governance, RbacOperation.View); const { data: providerGovernanceData } = useGetProviderGovernanceQuery(undefined, { @@ -111,7 +115,29 @@ export function GovernanceFormFragment({ provider }: GovernanceFormFragmentProps } }, [showCalendarAlignment, watchedCalendarAligned, form]); + // A budget config change on existing provider governance is when clearing + // accumulated spend becomes a meaningful choice. + const budgetsChanged = (data: FormData) => { + const signature = (rows: { max_limit?: number | null; reset_duration?: string }[]) => + [...rows] + .map((r) => `${r.max_limit ?? ""}:${r.reset_duration ?? ""}`) + .sort() + .join("|"); + const existing = providerGovernance?.budgets ?? []; + if (existing.length === 0) return false; + const next = data.budgets.filter((b) => b.max_limit !== undefined && b.max_limit > 0); + return signature(next) !== signature(existing); + }; + const onSubmit = async (data: FormData) => { + if (budgetsChanged(data)) { + resetPrompt.ask(data); + return; + } + await saveGovernance(data, false); + }; + + const saveGovernance = async (data: FormData, resetBudgetUsage: boolean) => { try { const validBudgets = data.budgets.filter((b) => b.max_limit !== undefined && b.max_limit > 0); const hasAlignableBudget = validBudgets.some((b) => supportsCalendarAlignment(b.reset_duration)); @@ -155,6 +181,8 @@ export function GovernanceFormFragment({ provider }: GovernanceFormFragmentProps budgets: budgetsPayload, ...(budgetsPayload !== undefined ? { calendar_aligned: hasAlignableBudget && data.calendarAligned } : {}), rate_limit: rateLimitPayload, + // Only sent when the operator explicitly chose to clear spend. + reset_budget_usage: resetBudgetUsage || undefined, }, }).unwrap(); @@ -287,6 +315,13 @@ export function GovernanceFormFragment({ provider }: GovernanceFormFragmentProps + resetPrompt.resolve((data) => saveGovernance(data, resetUsage))} + /> ); } \ No newline at end of file diff --git a/ui/components/ui/budgetUsageResetDialog.tsx b/ui/components/ui/budgetUsageResetDialog.tsx new file mode 100644 index 00000000000..ffb2b5bbb38 --- /dev/null +++ b/ui/components/ui/budgetUsageResetDialog.tsx @@ -0,0 +1,61 @@ +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alertDialog"; + +interface BudgetUsageResetDialogProps { + "data-testid"?: string; + open: boolean; + onOpenChange: (open: boolean) => void; + /** Called with the operator's choice: true to zero usage, false to preserve it. */ + onChoice: (resetUsage: boolean) => void; + /** Name of the thing being edited, e.g. "team", used in the prompt copy. */ + ownerLabel: string; +} + +/** + * Asks whether to clear accumulated spend when a budget's configuration changes. + * + * Shared by every owner's edit sheet rather than reimplemented per sheet: the + * wording is a promise about what the backend does, and four copies would drift + * from each other and from the API. + * + * Note the reset clears usage only. The reset boundary is never moved - it is + * guarded to advance forward only so that cluster nodes agree on which window is + * current - so the window keeps its existing start and end either way. + */ +export default function BudgetUsageResetDialog({ + "data-testid": testId, + open, + onOpenChange, + onChoice, + ownerLabel, +}: BudgetUsageResetDialogProps) { + return ( + + + + Reset budget usage? + + You changed a budget amount or reset frequency on this {ownerLabel}. Reset current usage to 0, or preserve the existing + counters. The reset window keeps its current start and end either way. + + + + onChoice(false)} data-testid={testId ? `${testId}-preserve-btn` : undefined}> + Preserve Usage + + onChoice(true)} data-testid={testId ? `${testId}-confirm-btn` : undefined}> + Reset Usage + + + + + ); +} \ No newline at end of file diff --git a/ui/hooks/useBudgetUsageResetPrompt.ts b/ui/hooks/useBudgetUsageResetPrompt.ts new file mode 100644 index 00000000000..170b0c9080a --- /dev/null +++ b/ui/hooks/useBudgetUsageResetPrompt.ts @@ -0,0 +1,49 @@ +import { useCallback, useState } from "react"; + +/** + * Defers a save until the operator has said whether to clear accumulated budget + * spend. + * + * The save payload is parked rather than recomputed after the choice, so the + * form cannot change underneath the dialog and produce a submission the operator + * never saw. Generic over that payload so each owner's sheet keeps its own shape. + * + * Usage: + * const resetPrompt = useBudgetUsageResetPrompt(); + * ... + * if (budgetsChanged) return resetPrompt.ask(data); + * await save(data, false); + * ... + * resetPrompt.resolve((payload) => save(payload, resetUsage))} + * /> + */ +export function useBudgetUsageResetPrompt() { + const [isOpen, setOpen] = useState(false); + const [pending, setPending] = useState(null); + + // Park the payload and open the dialog. Returns nothing so a caller can + // `return resetPrompt.ask(data)` straight out of a submit handler. + const ask = useCallback((payload: T) => { + setPending(payload); + setOpen(true); + }, []); + + // Hand the parked payload to the caller's save and clear the prompt. Clearing + // before awaiting keeps a slow save from leaving the dialog open and + // re-submittable. + const resolve = useCallback( + (save: (payload: T) => void | Promise) => { + if (pending === null) return; + const payload = pending; + setPending(null); + setOpen(false); + void save(payload); + }, + [pending], + ); + + return { isOpen, setOpen, ask, resolve }; +} \ No newline at end of file diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index bc25ed4b4f2..31f44bdd4b6 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -241,6 +241,8 @@ export interface UpdateTeamRequest { budgets?: CreateBudgetRequest[]; // Replaces all team budgets; empty array clears rate_limit?: UpdateRateLimitRequest; calendar_aligned?: boolean; + /** Zero current usage on the reconciled budgets. The reset window is unchanged. */ + reset_budget_usage?: boolean; } export interface CreateCustomerRequest { @@ -257,6 +259,8 @@ export interface UpdateCustomerRequest { budget?: UpdateBudgetRequest; // deprecated: use budgets rate_limit?: UpdateRateLimitRequest; calendar_aligned?: boolean; + /** Zero current usage on the reconciled budgets. The reset window is unchanged. */ + reset_budget_usage?: boolean; } export interface CreateBudgetRequest { @@ -422,6 +426,8 @@ export interface UpdateModelConfigRequest { provider?: string; // Optional provider - if empty/null, applies to all providers budgets?: CreateBudgetRequest[]; // Full desired set; reconciled against existing rate_limit?: UpdateRateLimitRequest; + /** Zero current usage on the reconciled budgets. The reset window is unchanged. */ + reset_budget_usage?: boolean; } export interface GetModelConfigsParams { @@ -606,6 +612,8 @@ export interface UpdateProviderGovernanceRequest { budgets?: CreateBudgetRequest[]; // [] = remove all rate_limit?: UpdateRateLimitRequest; calendar_aligned?: boolean; + /** Zero current usage on the reconciled budgets. The reset window is unchanged. */ + reset_budget_usage?: boolean; } export interface GetProviderGovernanceResponse { From febbb6c41cc6c9090dbd34a1bf8186b7506cbf97 Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Sun, 9 Aug 2026 16:27:01 -0700 Subject: [PATCH 2/4] governance: e2e coverage for reset_budget_usage, with an overridable gateway URL --- tests/governance/test_utils.go | 20 ++++- tests/governance/vkbudget_test.go | 133 ++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 2 deletions(-) diff --git a/tests/governance/test_utils.go b/tests/governance/test_utils.go index cb37197be73..2a77a9b00aa 100644 --- a/tests/governance/test_utils.go +++ b/tests/governance/test_utils.go @@ -7,6 +7,7 @@ import ( "io" "math/rand" "net/http" + "os" "strings" "testing" "time" @@ -88,9 +89,22 @@ type APIResponse struct { } // MakeRequest makes an HTTP request to the Bifrost API +// baseURL is the gateway the governance suite exercises. +// +// Overridable so a run can target a gateway other than the default dev one, for +// example an instance started with provider credentials in scope while a +// developer's own gateway keeps port 8080. Defaults to the historical value, so +// nothing changes for an unconfigured run. +func baseURL() string { + if override := os.Getenv("BIFROST_TEST_BASE_URL"); override != "" { + return strings.TrimSuffix(override, "/") + } + return "http://localhost:8080" +} + func MakeRequest(t *testing.T, req APIRequest) *APIResponse { client := &http.Client{} - url := fmt.Sprintf("http://localhost:8080%s", req.Path) + url := fmt.Sprintf("%s%s", baseURL(), req.Path) var body io.Reader if req.Body != nil { @@ -144,7 +158,7 @@ func MakeRequest(t *testing.T, req APIRequest) *APIResponse { // Use this when you need to test specific header formats (e.g., Authorization, x-api-key) func MakeRequestWithCustomHeaders(t *testing.T, req APIRequest, customHeaders map[string]string) *APIResponse { client := &http.Client{} - url := fmt.Sprintf("http://localhost:8080%s", req.Path) + url := fmt.Sprintf("%s%s", baseURL(), req.Path) var body io.Reader if req.Body != nil { @@ -289,6 +303,8 @@ type UpdateVirtualKeyRequest struct { IsActive *bool `json:"is_active,omitempty"` ProviderConfigs []ProviderConfigRequest `json:"provider_configs,omitempty"` CalendarAligned *bool `json:"calendar_aligned,omitempty"` + // ResetBudgetUsage zeroes accumulated spend on the reconciled budgets. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // UpdateTeamRequest represents a request to update a team diff --git a/tests/governance/vkbudget_test.go b/tests/governance/vkbudget_test.go index 4e5f42f8293..054f312d889 100644 --- a/tests/governance/vkbudget_test.go +++ b/tests/governance/vkbudget_test.go @@ -404,3 +404,136 @@ func TestVKQuarterlyBudgetRejectsInvalidResetConfig(t *testing.T) { }) } } + +// vkBudgetUsage reads the first budget's current usage off a virtual key's live +// in-memory governance state, which is what actually enforces spend. +func vkBudgetUsage(t *testing.T, vkID string) float64 { + t.Helper() + budget := budgetFromVK(t, vkID) + usage, ok := budget["current_usage"].(float64) + if !ok { + t.Fatalf("budget for VK %s has no current_usage: %v", vkID, budget) + } + return usage +} + +// spendOnVirtualKey makes one small real completion so the budget has usage to +// clear, and returns the usage the gateway recorded. +// +// One short request is enough: the assertions below care that usage is non-zero +// and then zero, not about the amount, so there is no reason to spend more. +func spendOnVirtualKey(t *testing.T, vkID, vkValue string) float64 { + t.Helper() + resp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/v1/chat/completions", + Body: ChatCompletionRequest{ + Model: "openai/gpt-4o", + Messages: []ChatMessage{{Role: "user", Content: "Reply with the single word: ok"}}, + }, + VKHeader: &vkValue, + }) + if resp.StatusCode >= 400 { + t.Fatalf("failed to spend against VK %s: status %d, body %v", vkID, resp.StatusCode, resp.Body) + } + + // Usage is recorded asynchronously relative to the response, so poll briefly. + deadline := time.Now().Add(20 * time.Second) + for { + if usage := vkBudgetUsage(t, vkID); usage > 0 { + return usage + } + if time.Now().After(deadline) { + t.Fatalf("VK %s recorded no budget usage after a successful completion", vkID) + } + time.Sleep(time.Second) + } +} + +// TestVKResetBudgetUsageClearsSpend is the full round trip for the operator's +// Reset choice: HTTP request, through the database, into the in-memory store that +// enforcement reads. +// +// This is the only layer that can catch the bug it guards. reset_budget_usage was +// accepted by the API, documented, and sent by the UI for a long time while no +// handler read it, and the value would have been discarded by UpdateBudget even if +// one had. Both failures are invisible to a unit test with a mocked store. +// +// Not parallel: it issues writes and polls, which contends with the rest of the +// suite on the default SQLite store. +func TestVKResetBudgetUsageClearsSpend(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + createResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-reset-usage-" + generateRandomID(), + ProviderConfigs: defaultProviderConfigs(), + Budgets: []BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}, + }, + }) + if createResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createResp.StatusCode, createResp.Body) + } + vkID := ExtractIDFromResponse(t, createResp) + testData.AddVirtualKey(vkID) + vkValue := createResp.Body["virtual_key"].(map[string]interface{})["value"].(string) + + spent := spendOnVirtualKey(t, vkID, vkValue) + t.Logf("recorded $%.6f of usage before the reset", spent) + + before := budgetFromVK(t, vkID) + lastResetBefore, err := time.Parse(time.RFC3339, before["last_reset"].(string)) + if err != nil { + t.Fatalf("could not parse last_reset %q: %v", before["last_reset"], err) + } + + // A budget edit WITHOUT the flag must leave the spend alone. Asserting this + // first matters: without it, a handler that always reset would still pass the + // positive case below and silently wipe usage on every unrelated edit. + preserveResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/virtual-keys/" + vkID, + Body: UpdateVirtualKeyRequest{ + Budgets: []BudgetRequest{{MaxLimit: 60, ResetDuration: "1M"}}, + }, + }) + if preserveResp.StatusCode != 200 { + t.Fatalf("Failed to update VK: status %d, body %v", preserveResp.StatusCode, preserveResp.Body) + } + if usage := vkBudgetUsage(t, vkID); usage != spent { + t.Errorf("an edit without reset_budget_usage changed usage from %v to %v; spend must be preserved", spent, usage) + } + + // Now with the flag. + resetUsage := true + resetResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/virtual-keys/" + vkID, + Body: UpdateVirtualKeyRequest{ + Budgets: []BudgetRequest{{MaxLimit: 70, ResetDuration: "1M"}}, + ResetBudgetUsage: &resetUsage, + }, + }) + if resetResp.StatusCode != 200 { + t.Fatalf("Failed to reset budget usage: status %d, body %v", resetResp.StatusCode, resetResp.Body) + } + + if usage := vkBudgetUsage(t, vkID); usage != 0 { + t.Errorf("reset_budget_usage did not clear spend: usage is still %v", usage) + } + + // The reset clears usage only. Moving the boundary is ruled out by the + // forward-only invariant that keeps cluster nodes agreeing on the open window. + after := budgetFromVK(t, vkID) + lastResetAfter, err := time.Parse(time.RFC3339, after["last_reset"].(string)) + if err != nil { + t.Fatalf("could not parse last_reset %q: %v", after["last_reset"], err) + } + if !lastResetAfter.Equal(lastResetBefore) { + t.Errorf("resetting usage moved last_reset from %s to %s; the window must be left alone", + lastResetBefore.Format(time.RFC3339), lastResetAfter.Format(time.RFC3339)) + } +} From 9e2ceb79fd1f419148365e7e9efd252f2e961709 Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Sun, 9 Aug 2026 16:33:18 -0700 Subject: [PATCH 3/4] governance: e2e coverage for budget usage resets and calendar alignment across owners --- tests/governance/customerbudget_test.go | 160 ++++++++++++++++++++++++ tests/governance/teambudget_test.go | 116 +++++++++++++++++ tests/governance/test_utils.go | 4 + tests/governance/vkbudget_test.go | 94 ++++++++++++++ 4 files changed, 374 insertions(+) diff --git a/tests/governance/customerbudget_test.go b/tests/governance/customerbudget_test.go index d20f0c4a05d..ef462d6c744 100644 --- a/tests/governance/customerbudget_test.go +++ b/tests/governance/customerbudget_test.go @@ -412,3 +412,163 @@ func customerBudgetLastReset(t *testing.T, customerID string) time.Time { } return parsed.UTC() } + +// TestCustomerResetBudgetUsageClearsSpend covers the customer owner's reset path, +// which goes through reconcileCustomerBudgets rather than the team inline loop or +// the model-config reconciler. +func TestCustomerResetBudgetUsageClearsSpend(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + createCustomerResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/customers", + Body: CreateCustomerRequest{ + Name: "test-customer-reset-usage-" + generateRandomID(), + Budgets: []BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}, + }, + }) + if createCustomerResp.StatusCode != 200 { + t.Fatalf("Failed to create customer: status %d, body %v", createCustomerResp.StatusCode, createCustomerResp.Body) + } + customerID := ExtractIDFromResponse(t, createCustomerResp) + testData.AddCustomer(customerID) + + createVKResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-customer-reset-" + generateRandomID(), + CustomerID: &customerID, + ProviderConfigs: defaultProviderConfigs(), + }, + }) + if createVKResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createVKResp.StatusCode, createVKResp.Body) + } + vkID := ExtractIDFromResponse(t, createVKResp) + testData.AddVirtualKey(vkID) + vkValue := createVKResp.Body["virtual_key"].(map[string]interface{})["value"].(string) + + spendVia(t, vkValue) + + deadline := time.Now().Add(30 * time.Second) + var spent float64 + for { + spent = customerBudgetUsage(t, customerID) + if spent > 0 { + break + } + if time.Now().After(deadline) { + t.Fatalf("customer %s recorded no budget usage after a successful completion", customerID) + } + time.Sleep(time.Second) + } + t.Logf("customer recorded $%.8f of usage before the reset", spent) + + resetUsage := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/customers/" + customerID, + Body: UpdateCustomerRequest{ + Budgets: []BudgetRequest{{MaxLimit: 60, ResetDuration: "1M"}}, + ResetBudgetUsage: &resetUsage, + }, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to reset customer budget usage: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + if usage := customerBudgetUsage(t, customerID); usage != 0 { + t.Errorf("reset_budget_usage did not clear customer spend: usage is still %v", usage) + } +} + +// customerBudgetUsage reads the first budget's current usage off a customer. +func customerBudgetUsage(t *testing.T, customerID string) float64 { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/customers/" + customerID}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read customer %s: status %d", customerID, resp.StatusCode) + } + customer, ok := resp.Body["customer"].(map[string]interface{}) + if !ok { + customer = resp.Body + } + budgets, ok := customer["budgets"].([]interface{}) + if !ok || len(budgets) == 0 { + t.Fatalf("customer %s has no budgets: %v", customerID, customer) + } + usage, _ := budgets[0].(map[string]interface{})["current_usage"].(float64) + return usage +} + +// TestCustomerCalendarAlignmentPreservesUsage pins the interaction between the two +// features: enabling calendar alignment must not clear spend as a side effect. +// +// The handler used to attempt exactly that, and it never worked because the write +// went through UpdateBudget, which carries usage forward. Now that a working reset +// path exists, an implementation could plausibly wire alignment into it and start +// silently wiping usage on a toggle. Clearing spend is only ever the operator's +// explicit choice via reset_budget_usage. +func TestCustomerCalendarAlignmentPreservesUsage(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + createCustomerResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/customers", + Body: CreateCustomerRequest{ + Name: "test-customer-align-usage-" + generateRandomID(), + Budgets: []BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}, + }, + }) + if createCustomerResp.StatusCode != 200 { + t.Fatalf("Failed to create customer: status %d, body %v", createCustomerResp.StatusCode, createCustomerResp.Body) + } + customerID := ExtractIDFromResponse(t, createCustomerResp) + testData.AddCustomer(customerID) + + createVKResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-customer-align-" + generateRandomID(), + CustomerID: &customerID, + ProviderConfigs: defaultProviderConfigs(), + }, + }) + if createVKResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createVKResp.StatusCode, createVKResp.Body) + } + vkID := ExtractIDFromResponse(t, createVKResp) + testData.AddVirtualKey(vkID) + spendVia(t, createVKResp.Body["virtual_key"].(map[string]interface{})["value"].(string)) + + deadline := time.Now().Add(30 * time.Second) + var spent float64 + for { + spent = customerBudgetUsage(t, customerID) + if spent > 0 { + break + } + if time.Now().After(deadline) { + t.Fatalf("customer %s recorded no usage after a successful completion", customerID) + } + time.Sleep(time.Second) + } + + aligned := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/customers/" + customerID, + Body: UpdateCustomerRequest{CalendarAligned: &aligned}, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to enable calendar alignment: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + if usage := customerBudgetUsage(t, customerID); usage != spent { + t.Errorf("enabling calendar alignment changed usage from %v to %v; only reset_budget_usage may clear spend", spent, usage) + } +} diff --git a/tests/governance/teambudget_test.go b/tests/governance/teambudget_test.go index 001dfaeab51..648ebc6e9f9 100644 --- a/tests/governance/teambudget_test.go +++ b/tests/governance/teambudget_test.go @@ -330,3 +330,119 @@ func teamRateLimitTokenLastReset(t *testing.T, teamID string) time.Time { } return parsed.UTC() } + +// TestTeamResetBudgetUsageClearsSpend covers the team owner's reset path. +// +// Each owner reconciles budgets through different code - teams use an inline +// loop, customers and model configs use their own reconcilers - so they are +// covered separately rather than by analogy from the virtual key case. +// +// Spend is accrued through a virtual key attached to the team, which is how team +// budgets are charged in the first place. +func TestTeamResetBudgetUsageClearsSpend(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + createTeamResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/teams", + Body: CreateTeamRequest{ + Name: "test-team-reset-usage-" + generateRandomID(), + Budgets: []BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}, + }, + }) + if createTeamResp.StatusCode != 200 { + t.Fatalf("Failed to create team: status %d, body %v", createTeamResp.StatusCode, createTeamResp.Body) + } + teamID := ExtractIDFromResponse(t, createTeamResp) + testData.AddTeam(teamID) + + createVKResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-team-reset-" + generateRandomID(), + TeamID: &teamID, + ProviderConfigs: defaultProviderConfigs(), + }, + }) + if createVKResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createVKResp.StatusCode, createVKResp.Body) + } + vkID := ExtractIDFromResponse(t, createVKResp) + testData.AddVirtualKey(vkID) + vkValue := createVKResp.Body["virtual_key"].(map[string]interface{})["value"].(string) + + spendVia(t, vkValue) + spent := waitForTeamBudgetUsage(t, teamID) + t.Logf("team recorded $%.8f of usage before the reset", spent) + + resetUsage := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/teams/" + teamID, + Body: UpdateTeamRequest{ + Budgets: &[]BudgetRequest{{MaxLimit: 60, ResetDuration: "1M"}}, + ResetBudgetUsage: &resetUsage, + }, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to reset team budget usage: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + if usage := teamBudgetUsage(t, teamID); usage != 0 { + t.Errorf("reset_budget_usage did not clear team spend: usage is still %v", usage) + } +} + +// spendVia makes one small completion against a virtual key so its owners accrue usage. +func spendVia(t *testing.T, vkValue string) { + t.Helper() + resp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/v1/chat/completions", + Body: ChatCompletionRequest{ + Model: "openai/gpt-4o", + Messages: []ChatMessage{{Role: "user", Content: "Reply with the single word: ok"}}, + }, + VKHeader: &vkValue, + }) + if resp.StatusCode >= 400 { + t.Fatalf("failed to spend: status %d, body %v", resp.StatusCode, resp.Body) + } +} + +// teamBudgetUsage reads the first budget's current usage off a team. +func teamBudgetUsage(t *testing.T, teamID string) float64 { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/teams/" + teamID}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read team %s: status %d", teamID, resp.StatusCode) + } + team, ok := resp.Body["team"].(map[string]interface{}) + if !ok { + team = resp.Body + } + budgets, ok := team["budgets"].([]interface{}) + if !ok || len(budgets) == 0 { + t.Fatalf("team %s has no budgets: %v", teamID, team) + } + usage, _ := budgets[0].(map[string]interface{})["current_usage"].(float64) + return usage +} + +// waitForTeamBudgetUsage polls until spend is recorded, which happens +// asynchronously relative to the completion response. +func waitForTeamBudgetUsage(t *testing.T, teamID string) float64 { + t.Helper() + deadline := time.Now().Add(30 * time.Second) + for { + if usage := teamBudgetUsage(t, teamID); usage > 0 { + return usage + } + if time.Now().After(deadline) { + t.Fatalf("team %s recorded no budget usage after a successful completion", teamID) + } + time.Sleep(time.Second) + } +} diff --git a/tests/governance/test_utils.go b/tests/governance/test_utils.go index 2a77a9b00aa..9232de13a59 100644 --- a/tests/governance/test_utils.go +++ b/tests/governance/test_utils.go @@ -319,6 +319,8 @@ type UpdateTeamRequest struct { // rate limit. Pointer so a test can distinguish "leave unchanged" from an // explicit false. CalendarAligned *bool `json:"calendar_aligned,omitempty"` + // ResetBudgetUsage zeroes accumulated spend on the reconciled budgets. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // UpdateCustomerRequest represents a request to update a customer @@ -326,6 +328,8 @@ type UpdateCustomerRequest struct { Name *string `json:"name,omitempty"` Budgets []BudgetRequest `json:"budgets,omitempty"` CalendarAligned *bool `json:"calendar_aligned,omitempty"` + // ResetBudgetUsage zeroes accumulated spend on the reconciled budgets. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } // ChatCompletionRequest represents an OpenAI-compatible chat completion request diff --git a/tests/governance/vkbudget_test.go b/tests/governance/vkbudget_test.go index 054f312d889..dba9c8f3fd7 100644 --- a/tests/governance/vkbudget_test.go +++ b/tests/governance/vkbudget_test.go @@ -537,3 +537,97 @@ func TestVKResetBudgetUsageClearsSpend(t *testing.T) { lastResetBefore.Format(time.RFC3339), lastResetAfter.Format(time.RFC3339)) } } + +// TestVKResetBudgetUsageRestoresCapacity answers the question a usage counter +// alone cannot: does clearing spend actually let requests through again. +// +// Zeroing current_usage is only meaningful if enforcement follows it. The +// enforcement path reads the in-memory governance store, not the database, so a +// reset that reached the database but not memory would show usage at zero on +// every read while requests kept getting rejected. This drives the budget to +// exhaustion, confirms rejection, resets, and confirms the next request succeeds. +// +// Not parallel: it writes and polls, contending with the rest of the suite on +// the default SQLite store. +func TestVKResetBudgetUsageRestoresCapacity(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + // Small enough that a single completion exhausts it, so the test costs one + // extra request rather than a loop. + createResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-reset-capacity-" + generateRandomID(), + ProviderConfigs: defaultProviderConfigs(), + Budgets: []BudgetRequest{{MaxLimit: 0.00001, ResetDuration: "1M"}}, + }, + }) + if createResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createResp.StatusCode, createResp.Body) + } + vkID := ExtractIDFromResponse(t, createResp) + testData.AddVirtualKey(vkID) + vkValue := createResp.Body["virtual_key"].(map[string]interface{})["value"].(string) + + completion := func() *APIResponse { + return MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/v1/chat/completions", + Body: ChatCompletionRequest{ + Model: "openai/gpt-4o", + Messages: []ChatMessage{{Role: "user", Content: "Reply with the single word: ok"}}, + }, + VKHeader: &vkValue, + }) + } + + // The first request is admitted because usage starts at zero, and its cost + // pushes usage past the cap. + if resp := completion(); resp.StatusCode >= 400 { + t.Fatalf("first request should be admitted on an unused budget: status %d, body %v", resp.StatusCode, resp.Body) + } + + // Usage is recorded asynchronously, so wait for enforcement to start rejecting. + deadline := time.Now().Add(30 * time.Second) + var blocked *APIResponse + for { + resp := completion() + if resp.StatusCode >= 400 { + blocked = resp + break + } + if time.Now().After(deadline) { + t.Fatalf("budget of $0.00001 never blocked a request; usage is %v", vkBudgetUsage(t, vkID)) + } + time.Sleep(time.Second) + } + if !CheckErrorMessage(t, blocked, "budget") { + t.Fatalf("request was rejected for a reason other than budget: %v", blocked.Body) + } + t.Logf("budget exhausted as expected at $%.8f of usage", vkBudgetUsage(t, vkID)) + + // Reset usage and raise the cap, so the next request is admitted only if the + // reset actually reached the store that enforces spend. + resetUsage := true + resetResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/virtual-keys/" + vkID, + Body: UpdateVirtualKeyRequest{ + Budgets: []BudgetRequest{{MaxLimit: 5, ResetDuration: "1M"}}, + ResetBudgetUsage: &resetUsage, + }, + }) + if resetResp.StatusCode != 200 { + t.Fatalf("Failed to reset budget usage: status %d, body %v", resetResp.StatusCode, resetResp.Body) + } + if usage := vkBudgetUsage(t, vkID); usage != 0 { + t.Fatalf("usage was not cleared: %v", usage) + } + + if resp := completion(); resp.StatusCode >= 400 { + t.Errorf("request still rejected after the usage reset, so enforcement did not see it: status %d, body %v", + resp.StatusCode, resp.Body) + } +} From bbdc667cdb7b4a947d4a68ce08238683c0f81218 Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Sun, 9 Aug 2026 16:41:46 -0700 Subject: [PATCH 4/4] governance: e2e budget coverage for provider governance and model limits --- .../features/governance/budget-and-limits.mdx | 4 +- tests/governance/customerbudget_test.go | 12 +- tests/governance/providerbudget_test.go | 279 ++++++++++++++++++ tests/governance/teambudget_test.go | 12 +- tests/governance/test_utils.go | 27 ++ .../bifrost-http/handlers/governance.go | 25 +- .../workspace/governance/views/teamSheet.tsx | 20 +- .../model-limits/views/modelLimitSheet.tsx | 21 +- ui/lib/utils/governance.test.ts | 25 ++ 9 files changed, 402 insertions(+), 23 deletions(-) diff --git a/docs/features/governance/budget-and-limits.mdx b/docs/features/governance/budget-and-limits.mdx index 6d036c856a2..b037e41205c 100644 --- a/docs/features/governance/budget-and-limits.mdx +++ b/docs/features/governance/budget-and-limits.mdx @@ -237,7 +237,9 @@ Supported on virtual keys, teams, customers, model limits, and provider governan The reset clears **usage only**. The reset window keeps its existing start and end, because `last_reset` only ever advances and never moves as a side effect of a configuration write. That rule is what keeps -every node in a cluster agreeing on which window is currently open. +every node in a cluster agreeing on which window is currently open. The normal scheduled reset is +unaffected and still advances `last_reset` when the window closes, so a manually cleared budget resets +again at its usual boundary rather than starting a fresh window from the moment you cleared it. ### Budget overrides diff --git a/tests/governance/customerbudget_test.go b/tests/governance/customerbudget_test.go index ef462d6c744..e21d7532bed 100644 --- a/tests/governance/customerbudget_test.go +++ b/tests/governance/customerbudget_test.go @@ -499,7 +499,17 @@ func customerBudgetUsage(t *testing.T, customerID string) float64 { if !ok || len(budgets) == 0 { t.Fatalf("customer %s has no budgets: %v", customerID, customer) } - usage, _ := budgets[0].(map[string]interface{})["current_usage"].(float64) + budget, ok := budgets[0].(map[string]interface{}) + if !ok { + t.Fatalf("customer %s first budget is not an object: %v", customerID, budgets[0]) + } + // Checked rather than defaulted to 0: every reset assertion is "usage is now + // zero", so a helper that silently returns 0 for a missing field would make + // those assertions pass without the reset having done anything. + usage, ok := budget["current_usage"].(float64) + if !ok { + t.Fatalf("customer %s budget has no current_usage: %v", customerID, budget) + } return usage } diff --git a/tests/governance/providerbudget_test.go b/tests/governance/providerbudget_test.go index b081b752b8b..13186d7d30b 100644 --- a/tests/governance/providerbudget_test.go +++ b/tests/governance/providerbudget_test.go @@ -3,6 +3,7 @@ package governance import ( "strconv" "testing" + "time" ) // TestProviderBudgetExceeded tests provider-specific budgets within a VK by making requests until budget is consumed @@ -238,3 +239,281 @@ func TestProviderBudgetExceeded(t *testing.T) { requestNum-1, consumedBudget, providerBudget) }) } + +// TestProviderGovernanceResetBudgetUsageClearsSpend covers provider-level +// governance, which is a distinct owner from a virtual key's per-provider budgets +// above: it applies to every request to that provider regardless of which key +// made it, and it is stored on a VK-agnostic model config. +// +// That storage detail is why it needs its own case. The reset is addressed to the +// model config that actually owns the budgets, not to the provider, and an +// implementation that addressed the provider would leave peers reloading an +// entity that owns nothing. +func TestProviderGovernanceResetBudgetUsageClearsSpend(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + // Provider governance is global, so leave it clean regardless of outcome. + defer func() { + MakeRequest(t, APIRequest{Method: "DELETE", Path: "/api/governance/providers/openai"}) + }() + + setResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/providers/openai", + Body: UpdateProviderGovernanceRequest{Budgets: &[]BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}}, + }) + if setResp.StatusCode != 200 { + t.Fatalf("Failed to set provider governance: status %d, body %v", setResp.StatusCode, setResp.Body) + } + + createVKResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-provider-gov-reset-" + generateRandomID(), + ProviderConfigs: defaultProviderConfigs(), + }, + }) + if createVKResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createVKResp.StatusCode, createVKResp.Body) + } + testData.AddVirtualKey(ExtractIDFromResponse(t, createVKResp)) + spendVia(t, createVKResp.Body["virtual_key"].(map[string]interface{})["value"].(string)) + + deadline := time.Now().Add(30 * time.Second) + var spent float64 + for { + spent = providerGovernanceUsage(t, "openai") + if spent > 0 { + break + } + if time.Now().After(deadline) { + t.Fatalf("provider governance recorded no usage after a successful completion") + } + time.Sleep(time.Second) + } + t.Logf("provider governance recorded $%.8f of usage before the reset", spent) + + resetUsage := true + resetResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/providers/openai", + Body: UpdateProviderGovernanceRequest{ + Budgets: &[]BudgetRequest{{MaxLimit: 60, ResetDuration: "1M"}}, + ResetBudgetUsage: &resetUsage, + }, + }) + if resetResp.StatusCode != 200 { + t.Fatalf("Failed to reset provider governance usage: status %d, body %v", resetResp.StatusCode, resetResp.Body) + } + + if usage := providerGovernanceUsage(t, "openai"); usage != 0 { + t.Errorf("reset_budget_usage did not clear provider governance spend: usage is still %v", usage) + } +} + +// TestProviderGovernanceCalendarAlignmentAppliesFromNextPeriod is the third snap +// site, alongside teams and customers. All three shared the same discarded write, +// and all three persist through different store methods, so each is pinned. +func TestProviderGovernanceCalendarAlignmentAppliesFromNextPeriod(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + defer func() { + MakeRequest(t, APIRequest{Method: "DELETE", Path: "/api/governance/providers/openai"}) + }() + + setResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/providers/openai", + Body: UpdateProviderGovernanceRequest{Budgets: &[]BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}}, + }) + if setResp.StatusCode != 200 { + t.Fatalf("Failed to set provider governance: status %d, body %v", setResp.StatusCode, setResp.Body) + } + + before := providerGovernanceLastReset(t, "openai") + + aligned := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/providers/openai", + Body: UpdateProviderGovernanceRequest{CalendarAligned: &aligned}, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to enable calendar alignment: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + after := providerGovernanceLastReset(t, "openai") + if !after.Equal(before) { + t.Errorf("enabling calendar alignment moved last_reset from %s to %s; the current window must be left alone", + before.Format(time.RFC3339), after.Format(time.RFC3339)) + } +} + +// providerGovernanceBudget returns the first budget recorded against a provider. +func providerGovernanceBudget(t *testing.T, provider string) map[string]interface{} { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/providers"}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read provider governance: status %d, body %v", resp.StatusCode, resp.Body) + } + providers, ok := resp.Body["providers"].([]interface{}) + if !ok { + t.Fatalf("unexpected provider governance response: %v", resp.Body) + } + for _, raw := range providers { + entry, ok := raw.(map[string]interface{}) + if !ok || entry["provider"] != provider { + continue + } + budgets, ok := entry["budgets"].([]interface{}) + if !ok || len(budgets) == 0 { + t.Fatalf("provider %s has no budgets: %v", provider, entry) + } + budget, _ := budgets[0].(map[string]interface{}) + return budget + } + t.Fatalf("provider %s not found in governance response", provider) + return nil +} + +func providerGovernanceUsage(t *testing.T, provider string) float64 { + t.Helper() + budget := providerGovernanceBudget(t, provider) + // Checked rather than defaulted to 0: every reset assertion is "usage is now + // zero", so a helper that silently returns 0 for a missing field would make + // those assertions pass without the reset having done anything. + usage, ok := budget["current_usage"].(float64) + if !ok { + t.Fatalf("provider %s governance budget has no current_usage: %v", provider, budget) + } + return usage +} + +func providerGovernanceLastReset(t *testing.T, provider string) time.Time { + t.Helper() + raw, _ := providerGovernanceBudget(t, provider)["last_reset"].(string) + parsed, err := time.Parse(time.RFC3339, raw) + if err != nil { + t.Fatalf("could not parse last_reset %q: %v", raw, err) + } + return parsed.UTC() +} + +// TestModelConfigResetBudgetUsageClearsSpend covers model-scoped limits, the last +// budget owner. Its budgets are reconciled by reconcileModelConfigBudgets, the +// same function the virtual key path uses, but reached through a different +// handler, so the wiring is verified independently. +func TestModelConfigResetBudgetUsageClearsSpend(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + provider := "openai" + createResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/model-configs", + Body: CreateModelConfigRequest{ + ModelName: "gpt-4o", + Provider: &provider, + Scope: "global", + Budgets: []BudgetRequest{{MaxLimit: 50, ResetDuration: "1M"}}, + }, + }) + if createResp.StatusCode != 200 { + t.Fatalf("Failed to create model config: status %d, body %v", createResp.StatusCode, createResp.Body) + } + mcID := modelConfigIDFromResponse(t, createResp) + defer func() { + MakeRequest(t, APIRequest{Method: "DELETE", Path: "/api/governance/model-configs/" + mcID}) + }() + + createVKResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/virtual-keys", + Body: CreateVirtualKeyRequest{ + Name: "test-vk-model-config-reset-" + generateRandomID(), + ProviderConfigs: defaultProviderConfigs(), + }, + }) + if createVKResp.StatusCode != 200 { + t.Fatalf("Failed to create VK: status %d, body %v", createVKResp.StatusCode, createVKResp.Body) + } + testData.AddVirtualKey(ExtractIDFromResponse(t, createVKResp)) + spendVia(t, createVKResp.Body["virtual_key"].(map[string]interface{})["value"].(string)) + + deadline := time.Now().Add(30 * time.Second) + var spent float64 + for { + spent = modelConfigUsage(t, mcID) + if spent > 0 { + break + } + if time.Now().After(deadline) { + t.Fatalf("model config %s recorded no usage after a successful completion on its model", mcID) + } + time.Sleep(time.Second) + } + t.Logf("model config recorded $%.8f of usage before the reset", spent) + + resetUsage := true + resetResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/model-configs/" + mcID, + Body: UpdateModelConfigRequest{ + Budgets: []BudgetRequest{{MaxLimit: 60, ResetDuration: "1M"}}, + ResetBudgetUsage: &resetUsage, + }, + }) + if resetResp.StatusCode != 200 { + t.Fatalf("Failed to reset model config usage: status %d, body %v", resetResp.StatusCode, resetResp.Body) + } + + if usage := modelConfigUsage(t, mcID); usage != 0 { + t.Errorf("reset_budget_usage did not clear model config spend: usage is still %v", usage) + } +} + +// modelConfigIDFromResponse pulls the id out of a model-config create response, +// which nests it under a key the shared extractor does not know about. +func modelConfigIDFromResponse(t *testing.T, resp *APIResponse) string { + t.Helper() + if mc, ok := resp.Body["model_config"].(map[string]interface{}); ok { + if id, ok := mc["id"].(string); ok { + return id + } + } + if id, ok := resp.Body["id"].(string); ok { + return id + } + t.Fatalf("could not find model config id in response: %v", resp.Body) + return "" +} + +func modelConfigUsage(t *testing.T, mcID string) float64 { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/model-configs/" + mcID}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read model config %s: status %d, body %v", mcID, resp.StatusCode, resp.Body) + } + mc, ok := resp.Body["model_config"].(map[string]interface{}) + if !ok { + mc = resp.Body + } + budgets, ok := mc["budgets"].([]interface{}) + if !ok || len(budgets) == 0 { + t.Fatalf("model config %s has no budgets: %v", mcID, mc) + } + budget, ok := budgets[0].(map[string]interface{}) + if !ok { + t.Fatalf("model config %s first budget is not an object: %v", mcID, budgets[0]) + } + // Checked rather than defaulted to 0: every reset assertion is "usage is now + // zero", so a helper that silently returns 0 for a missing field would make + // those assertions pass without the reset having done anything. + usage, ok := budget["current_usage"].(float64) + if !ok { + t.Fatalf("model config %s budget has no current_usage: %v", mcID, budget) + } + return usage +} diff --git a/tests/governance/teambudget_test.go b/tests/governance/teambudget_test.go index 648ebc6e9f9..cad9fdddb56 100644 --- a/tests/governance/teambudget_test.go +++ b/tests/governance/teambudget_test.go @@ -427,7 +427,17 @@ func teamBudgetUsage(t *testing.T, teamID string) float64 { if !ok || len(budgets) == 0 { t.Fatalf("team %s has no budgets: %v", teamID, team) } - usage, _ := budgets[0].(map[string]interface{})["current_usage"].(float64) + budget, ok := budgets[0].(map[string]interface{}) + if !ok { + t.Fatalf("team %s first budget is not an object: %v", teamID, budgets[0]) + } + // Checked rather than defaulted to 0: every reset assertion is "usage is now + // zero", so a helper that silently returns 0 for a missing field would make + // those assertions pass without the reset having done anything. + usage, ok := budget["current_usage"].(float64) + if !ok { + t.Fatalf("team %s budget has no current_usage: %v", teamID, budget) + } return usage } diff --git a/tests/governance/test_utils.go b/tests/governance/test_utils.go index 9232de13a59..92c27352b9d 100644 --- a/tests/governance/test_utils.go +++ b/tests/governance/test_utils.go @@ -332,6 +332,33 @@ type UpdateCustomerRequest struct { ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` } +// CreateModelConfigRequest represents a request to create a model-scoped limit. +type CreateModelConfigRequest struct { + ModelName string `json:"model_name"` + Provider *string `json:"provider,omitempty"` + Scope string `json:"scope,omitempty"` + ScopeID *string `json:"scope_id,omitempty"` + Budgets []BudgetRequest `json:"budgets,omitempty"` + RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` +} + +// UpdateModelConfigRequest represents a request to update a model-scoped limit. +type UpdateModelConfigRequest struct { + ModelName *string `json:"model_name,omitempty"` + Provider *string `json:"provider,omitempty"` + Budgets []BudgetRequest `json:"budgets,omitempty"` + // ResetBudgetUsage zeroes accumulated spend on the reconciled budgets. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` +} + +// UpdateProviderGovernanceRequest represents a request to update provider-level governance. +type UpdateProviderGovernanceRequest struct { + Budgets *[]BudgetRequest `json:"budgets,omitempty"` + CalendarAligned *bool `json:"calendar_aligned,omitempty"` + // ResetBudgetUsage zeroes accumulated spend on the reconciled budgets. + ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` +} + // ChatCompletionRequest represents an OpenAI-compatible chat completion request type ChatCompletionRequest struct { Model string `json:"model"` diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 3e0c8396ee5..62278f0c94a 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -2762,11 +2762,17 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, "Failed to update team") return } - // Reloading team from in-memory store + // Reloading team from in-memory store. + // + // A failure here cannot be swallowed on this path: the usage reset below is + // ordered after the reload and depends on it having happened, so continuing + // would clear usage against stale in-memory config while the database already + // holds the new one, and still answer 200 with the pre-update team. preloadedTeam, err := h.governanceManager.ReloadTeam(ctx, team.ID) if err != nil { logger.Error("failed to reload team: %v", err) - preloadedTeam = team + SendError(ctx, 500, "Team updated in database but failed to reload in-memory state") + return } // Clear usage in the store that enforces spend. This runs after the reload, // which deliberately carries each budget's cached usage forward and would @@ -3108,10 +3114,14 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { return } + // Not swallowed on this path: the usage reset below is ordered after the reload + // and depends on it, so continuing would clear usage against stale in-memory + // config and still answer 200 with the pre-update customer. preloadedCustomer, err := h.governanceManager.ReloadCustomer(ctx, customer.ID) if err != nil { logger.Error("failed to reload customer: %v", err) - preloadedCustomer = customer + SendError(ctx, 500, "Customer updated in database but failed to reload in-memory state") + return } // Clear usage in the store that enforces spend. This runs after the reload, // which deliberately carries each budget's cached usage forward and would @@ -3732,11 +3742,16 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, fmt.Sprintf("Failed to update model config: %v", err)) return } - // Reload model config in memory (also reloads from DB to get full relationships) + // Reload model config in memory (also reloads from DB to get full relationships). + // + // Not swallowed on this path: the usage reset below is ordered after the reload + // and depends on it, so continuing would clear usage against stale in-memory + // config and still answer 200 with the pre-update model config. updatedMC, err := h.governanceManager.ReloadModelConfig(ctx, mc.ID) if err != nil { logger.Error("failed to reload model config in memory: %v", err) - updatedMC = mc + SendError(ctx, 500, "Model config updated in database but failed to reload in-memory state") + return } // Clear usage in the store that enforces spend. This runs after the reload, // which deliberately carries each budget's cached usage forward and would diff --git a/ui/app/workspace/governance/views/teamSheet.tsx b/ui/app/workspace/governance/views/teamSheet.tsx index 1d3e4c8385c..0a34cf289ac 100644 --- a/ui/app/workspace/governance/views/teamSheet.tsx +++ b/ui/app/workspace/governance/views/teamSheet.tsx @@ -24,7 +24,7 @@ import QuarterStartSelect from "@/components/ui/quarterStartSelect"; import { budgetResetDurationOptions, resetDurationOptions, supportsCalendarAlignment } from "@/lib/constants/governance"; import { getErrorMessage, useCreateTeamMutation, useUpdateTeamMutation } from "@/lib/store"; import { CreateTeamRequest, Team, UpdateTeamRequest } from "@/lib/types/governance"; -import { formatCurrency } from "@/lib/utils/governance"; +import { budgetSignature, formatCurrency } from "@/lib/utils/governance"; import { Validator } from "@/lib/utils/validation"; import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { formatDistanceToNow } from "date-fns"; @@ -239,17 +239,21 @@ export default function TeamSheet({ team, onSave, onCancel }: TeamSheetProps) { // Whether this save changes budget configuration on an existing team, which is // when clearing accumulated spend is a meaningful choice. Creating a team has // no usage to reset, and a save that leaves budgets alone should not ask. + // Compared without ids on purpose: form rows carry none while persisted rows do, + // so including them would report a change on every save. The shared signature + // folds in the fiscal quarter, which a limit-and-duration comparison misses - + // moving Q1 from April to July reschedules the reset without touching either. const budgetsChanged = () => { if (!isEditing || !team) return false; - const signature = (rows: { max_limit?: number | null; reset_duration?: string }[]) => - [...rows] - .map((r) => `${r.max_limit ?? ""}:${r.reset_duration ?? ""}`) - .sort() - .join("|"); const next = formData.budgets .filter((r) => r.maxLimit !== undefined && r.maxLimit !== null) - .map((r) => ({ max_limit: r.maxLimit, reset_duration: r.resetDuration })); - return signature(next) !== signature(team.budgets ?? []); + .map((r) => ({ max_limit: r.maxLimit, reset_duration: r.resetDuration, reset_config: r.resetConfig })); + const current = (team.budgets ?? []).map((b) => ({ + max_limit: b.max_limit ?? undefined, + reset_duration: b.reset_duration, + reset_config: b.reset_config, + })); + return budgetSignature(next) !== budgetSignature(current); }; const handleSubmit = async (e: React.FormEvent) => { diff --git a/ui/app/workspace/model-limits/views/modelLimitSheet.tsx b/ui/app/workspace/model-limits/views/modelLimitSheet.tsx index 3409f33ceee..d659254f624 100644 --- a/ui/app/workspace/model-limits/views/modelLimitSheet.tsx +++ b/ui/app/workspace/model-limits/views/modelLimitSheet.tsx @@ -10,6 +10,7 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { DottedSeparator } from "@/components/ui/separator"; import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; import { resetDurationOptions } from "@/lib/constants/governance"; +import { budgetSignature } from "@/lib/utils/governance"; import { RenderProviderIcon } from "@/lib/constants/icons"; import { ProviderLabels, ProviderName } from "@/lib/constants/logs"; import { getModelLimitScope, getModelLimitScopes } from "@/lib/registries/modelLimitScopes"; @@ -170,15 +171,21 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model // A budget config change on an existing limit is when clearing accumulated // spend becomes a meaningful choice; creating one has no usage to reset. + // Compared without ids on purpose: form rows carry none while persisted rows do, + // so including them would report a change on every save. The shared signature + // folds in the fiscal quarter, which a limit-and-duration comparison misses - + // moving Q1 from April to July reschedules the reset without touching either. const budgetsChanged = (data: FormData) => { if (!isEditing || !modelConfig) return false; - const signature = (rows: { max_limit?: number | null; reset_duration?: string }[]) => - [...rows] - .map((r) => `${r.max_limit ?? ""}:${r.reset_duration ?? ""}`) - .sort() - .join("|"); - const next = (data.budgets ?? []).filter((b) => b.max_limit !== undefined && b.max_limit !== null); - return signature(next) !== signature(modelConfig.budgets ?? []); + const next = (data.budgets ?? []) + .filter((b) => b.max_limit !== undefined && b.max_limit !== null) + .map((b) => ({ max_limit: b.max_limit, reset_duration: b.reset_duration, reset_config: b.reset_config })); + const current = (modelConfig.budgets ?? []).map((b) => ({ + max_limit: b.max_limit ?? undefined, + reset_duration: b.reset_duration, + reset_config: b.reset_config, + })); + return budgetSignature(next) !== budgetSignature(current); }; const onSubmit = async (data: FormData) => { diff --git a/ui/lib/utils/governance.test.ts b/ui/lib/utils/governance.test.ts index 5593d9c7e0b..1bc7dc32bda 100644 --- a/ui/lib/utils/governance.test.ts +++ b/ui/lib/utils/governance.test.ts @@ -156,4 +156,29 @@ describe("budgetSignature", () => { it("skips budgets with no limit set", () => { expect(budgetSignature([{ id: "b-1", reset_duration: "1Q" }])).toBe(""); }); +}); +// The owner sheets compare budgets without ids: form rows have no id while +// persisted rows do, so feeding ids in would make every save look like a change. +describe("budgetSignature without ids", () => { + const row = (quarterStartMonth?: number) => [ + { + max_limit: 500, + reset_duration: "1Q", + ...(quarterStartMonth === undefined ? {} : { reset_config: { quarter_start_month: quarterStartMonth } }), + }, + ]; + + it("detects a fiscal quarter change on an otherwise identical budget", () => { + expect(budgetSignature(row(4))).not.toBe(budgetSignature(row(7))); + }); + + it("reports no change when only ids differ", () => { + const persisted = [{ id: "b-1", max_limit: 500, reset_duration: "1Q", reset_config: { quarter_start_month: 4 } }]; + const stripped = persisted.map((b) => ({ + max_limit: b.max_limit, + reset_duration: b.reset_duration, + reset_config: b.reset_config, + })); + expect(budgetSignature(stripped)).toBe(budgetSignature(row(4))); + }); }); \ No newline at end of file