From 26a50f0898e76b2cc6ef7db54ecc861cc766770f Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Sun, 9 Aug 2026 15:16:01 -0700 Subject: [PATCH] adds reset budget honor flow --- plugins/governance/budgetcycle_test.go | 124 ++++++++++++++++ plugins/governance/store.go | 132 +++++++++++++++++- .../bifrost-http/handlers/governance.go | 77 ++++++++-- .../handlers/pricing_override_test.go | 3 + transports/bifrost-http/server/server.go | 33 +++++ 5 files changed, 357 insertions(+), 12 deletions(-) diff --git a/plugins/governance/budgetcycle_test.go b/plugins/governance/budgetcycle_test.go index ac907aa1ca1..812800171a0 100644 --- a/plugins/governance/budgetcycle_test.go +++ b/plugins/governance/budgetcycle_test.go @@ -717,3 +717,127 @@ func TestQuarterDefinitionSurvivesVirtualKeyReload(t *testing.T) { assertQuarterly(t, store.LoadBudget(ctx, "vk-quarterly-budget"), 250) }) } + +// TestResetBudgetUsageInMemory covers the primitive behind an operator-triggered +// usage reset. +// +// It must zero usage without touching LastReset. Moving the boundary is ruled +// out: every persistence path guards it forward-only, and the intended target +// would often be earlier than the current value. Clearing the LastDBUsages +// baseline alongside is not optional - the dump path writes the delta between +// in-memory usage and that baseline, so a stale baseline would immediately +// re-add the spend that was just cleared. +func TestResetBudgetUsageInMemory(t *testing.T) { + ctx := context.Background() + store := newStandaloneStore(t) + + anchor := cycleTestAnchor() + budget := buildBudgetWithUsage("operator-reset", 1000, 425, "1M") + budget.LastReset = anchor + store.budgets.Store(budget.ID, budget) + store.LastDBUsagesBudgetsMu.Lock() + store.LastDBUsagesBudgets[budget.ID] = 300 + store.LastDBUsagesBudgetsMu.Unlock() + + reset, ok := store.ResetBudgetUsageInMemory(ctx, budget.ID) + require.True(t, ok, "reset should apply to a budget that exists") + require.NotNil(t, reset) + + assert.Zero(t, reset.CurrentUsage, "usage must be cleared") + assert.True(t, reset.LastReset.Equal(anchor), "the reset boundary must not move") + + loaded := store.LoadBudget(ctx, budget.ID) + require.NotNil(t, loaded) + assert.Zero(t, loaded.CurrentUsage, "the stored budget must reflect the reset") + assert.True(t, loaded.LastReset.Equal(anchor)) + + store.LastDBUsagesBudgetsMu.RLock() + baseline := store.LastDBUsagesBudgets[budget.ID] + store.LastDBUsagesBudgetsMu.RUnlock() + assert.Zero(t, baseline, "a stale baseline would re-add the cleared spend on the next dump") +} + +// TestResetBudgetUsageInMemoryMissingBudget verifies an unknown ID is reported +// rather than silently succeeding, so a caller cannot believe it reset something +// that does not exist. +func TestResetBudgetUsageInMemoryMissingBudget(t *testing.T) { + store := newStandaloneStore(t) + reset, ok := store.ResetBudgetUsageInMemory(context.Background(), "does-not-exist") + assert.False(t, ok) + assert.Nil(t, reset) +} + +// TestDumpBudgetsCannotUndoOperatorReset covers the interleaving where an operator +// reset lands after a dump has snapshotted a budget but before that snapshot +// reaches the database. +// +// The two halves of a dump are separated by one database transaction per batch, so +// this window is real work, not an instant. A reset installs a fresh pointer via +// CAS, which the snapshot cannot see, and deliberately leaves LastReset alone, so +// the "<=" guard still matches. Without the reset generation the stale usage is +// written straight back over the operator's reset, and nothing reports an error: +// the number simply reappears. +func TestDumpBudgetsCannotUndoOperatorReset(t *testing.T) { + ctx := context.Background() + logger := NewMockLogger() + now := time.Now().UTC().Truncate(time.Second) + + newStore := func(t *testing.T, seeded *configstoreTables.TableBudget) (*LocalGovernanceStore, configstore.ConfigStore) { + t.Helper() + configStore, err := configstore.NewConfigStore(ctx, &configstore.Config{ + Enabled: true, + Type: configstore.ConfigStoreTypeSQLite, + Config: &configstore.SQLiteConfig{Path: t.TempDir() + "/resetrace.db"}, + }, logger) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, configStore.Close(ctx)) }) + require.NoError(t, configStore.CreateBudget(ctx, seeded)) + store, err := NewLocalGovernanceStore(ctx, logger, configStore, nil, nil) + require.NoError(t, err) + return store, configStore + } + + seed := buildBudgetWithUsage("dump-reset-race-budget", 1000, 500, "24h") + seed.CreatedAt = now + seed.UpdatedAt = now + seed.LastReset = now + + store, configStore := newStore(t, seed) + + live := store.LoadBudget(ctx, seed.ID) + require.NotNil(t, live) + require.Equal(t, 500.0, live.CurrentUsage, "precondition: the budget carries spend to be reset") + require.True(t, live.LastReset.Equal(now), "precondition: no sweep moved the boundary") + + // The dump reads the pre-reset usage, then stalls before writing. + rows, gens := store.snapshotBudgetRows(nil) + require.NotEmpty(t, rows) + + // The operator reset lands in the gap: memory is cleared here, and the handler's + // own transaction clears the database. + reset, ok := store.ResetBudgetUsageInMemory(ctx, seed.ID) + require.True(t, ok) + require.Zero(t, reset.CurrentUsage) + require.True(t, reset.LastReset.Equal(now), + "an operator reset must leave the boundary alone, which is why the <= guard cannot catch this") + require.NoError(t, configStore.UpdateBudgetUsage(ctx, seed.ID, 0)) + + // The stalled dump now writes. Its rows are stale and must be dropped. + require.NoError(t, store.writeBudgetRows(ctx, rows, gens)) + + persisted, err := configStore.GetBudget(ctx, seed.ID) + require.NoError(t, err) + assert.Zero(t, persisted.CurrentUsage, + "a dump that snapshotted before the reset wrote %.2f back and silently undid the operator's reset", persisted.CurrentUsage) + assert.True(t, persisted.LastReset.UTC().Equal(now), + "dropping a stale row must not disturb the boundary") + + // The drop is scoped to the reset, not permanent: the next cycle persists the + // post-reset value normally. + require.NoError(t, store.BumpBudgetUsage(ctx, seed.ID, 7.5)) + require.NoError(t, store.DumpBudgets(ctx, nil)) + persisted, err = configStore.GetBudget(ctx, seed.ID) + require.NoError(t, err) + assert.Equal(t, 7.5, persisted.CurrentUsage, + "the dump after a reset must resume persisting usage, or a reset would stop accounting for good") +} diff --git a/plugins/governance/store.go b/plugins/governance/store.go index 6ee3a3844cb..0e459a176c5 100644 --- a/plugins/governance/store.go +++ b/plugins/governance/store.go @@ -46,6 +46,19 @@ type LocalGovernanceStore struct { LastDBUsagesRequestsRateLimits map[string]int64 // Map for last DB usages for rate limits requests LastDBUsagesTokensRateLimits map[string]int64 // Map for last DB usages for rate limits tokens + // budgetResetGens counts how many times each budget's in-memory usage has been + // zeroed. A dump snapshots budgets by pointer and writes them several + // transactions later; a reset installs a fresh pointer and deliberately leaves + // LastReset alone, so neither the snapshot nor the last_reset guard can tell + // that the usage about to be persisted has since been cleared. LastReset cannot + // carry this signal: an operator reset leaves it equal, and equal has to keep + // meaning "persist this row" or steady-state usage would stop saving. + // + // Counts resets only, never usage writes, so a budget under traffic still dumps + // every cycle. Guarded by LastDBUsagesBudgetsMu, which every reset already holds + // to clear the matching baseline, so this adds no new lock. + budgetResetGens map[string]uint64 + // CEL caching layer for routing rules compiledRoutingPrograms sync.Map // string -> cel.Program (key: ruleID -> compiled CEL program) routingCELEnv *cel.Env // Singleton CEL environment reused for all compilations @@ -148,6 +161,9 @@ type GovernanceStore interface { // In-memory reset checks (return items that need DB sync) ResetExpiredRateLimitsInMemory(ctx context.Context, refreshReferences bool, rateLimitIDs ...string) []*configstoreTables.TableRateLimit ResetExpiredBudgetsInMemory(ctx context.Context, refreshReferences bool, budgetIDs ...string) []*configstoreTables.TableBudget + // ResetBudgetUsageInMemory clears one budget's usage on operator request, + // leaving its reset boundary untouched. + ResetBudgetUsageInMemory(ctx context.Context, budgetID string) (*configstoreTables.TableBudget, bool) // DB sync for expired items ResetExpiredRateLimits(ctx context.Context, resetRateLimits []*configstoreTables.TableRateLimit) error ResetExpiredBudgets(ctx context.Context, resetBudgets []*configstoreTables.TableBudget) error @@ -364,6 +380,8 @@ func (gs *LocalGovernanceStore) DeleteBudget(ctx context.Context, budgetID strin // Clean up LastDB baselines so the gossip delta doesn't carry stale entries. gs.LastDBUsagesBudgetsMu.Lock() delete(gs.LastDBUsagesBudgets, budgetID) + // A dump holding this budget in an earlier snapshot must not resurrect its row. + gs.markBudgetResetLocked(budgetID) gs.LastDBUsagesBudgetsMu.Unlock() } @@ -378,6 +396,60 @@ func (gs *LocalGovernanceStore) SetBudgetDBBaseline(budgetID string, currentUsag gs.LastDBUsagesBudgetsMu.Unlock() } +// markBudgetResetLocked records that a budget's in-memory usage was cleared, so a +// dump that snapshotted the pre-reset value can recognise its row as stale. +// +// The caller must already hold LastDBUsagesBudgetsMu for writing. Every reset site +// does, because clearing the gossip baseline and invalidating in-flight dumps are +// the same event and must not be observable half-done. +func (gs *LocalGovernanceStore) markBudgetResetLocked(budgetID string) { + if gs.budgetResetGens == nil { + gs.budgetResetGens = make(map[string]uint64) + } + gs.budgetResetGens[budgetID]++ +} + +// budgetResetGensSnapshot copies the current reset generations. A dump takes this +// alongside its budget snapshot and re-compares before writing. +func (gs *LocalGovernanceStore) budgetResetGensSnapshot() map[string]uint64 { + gs.LastDBUsagesBudgetsMu.RLock() + defer gs.LastDBUsagesBudgetsMu.RUnlock() + gens := make(map[string]uint64, len(gs.budgetResetGens)) + for budgetID, gen := range gs.budgetResetGens { + gens[budgetID] = gen + } + return gens +} + +// keepUnresetRows returns the rows of a dump batch whose budget has not been reset +// since snapshot was taken. A row that has been reset carries the usage the reset +// cleared, so writing it would put that number straight back. +// +// Dropping is safe and self-correcting: the post-reset usage is already in memory, +// so the next dump cycle persists it a few seconds later. Re-checked per batch +// rather than once per dump because batches are separated by database round trips, +// which is time enough for a reset to land between them. +func (gs *LocalGovernanceStore) keepUnresetRows(batch []budgetDumpRow, snapshot map[string]uint64) []budgetDumpRow { + gs.LastDBUsagesBudgetsMu.RLock() + defer gs.LastDBUsagesBudgetsMu.RUnlock() + stale := 0 + for _, row := range batch { + if gs.budgetResetGens[row.ID] != snapshot[row.ID] { + stale++ + } + } + if stale == 0 { + return batch + } + kept := make([]budgetDumpRow, 0, len(batch)-stale) + for _, row := range batch { + if gs.budgetResetGens[row.ID] == snapshot[row.ID] { + kept = append(kept, row) + } + } + return kept +} + // LoadRateLimit loads a rate limit by its ID from the local store. func (gs *LocalGovernanceStore) LoadRateLimit(ctx context.Context, rateLimitID string) *configstoreTables.TableRateLimit { if rateLimit, ok := gs.rateLimits.Load(rateLimitID); ok { @@ -503,6 +575,7 @@ func (gs *LocalGovernanceStore) BumpBudgetUsage(ctx context.Context, budgetID st clone.RefreshOverrideCyclesRemaining() gs.LastDBUsagesBudgetsMu.Lock() gs.LastDBUsagesBudgets[budgetID] = 0 + gs.markBudgetResetLocked(budgetID) gs.LastDBUsagesBudgetsMu.Unlock() } clone.CurrentUsage += cost @@ -2019,11 +2092,41 @@ func (gs *LocalGovernanceStore) resetExpiredBudgetFromSnapshot(ctx context.Conte oldUsage := budget.CurrentUsage gs.LastDBUsagesBudgetsMu.Lock() gs.LastDBUsagesBudgets[resetBudget.ID] = 0 + gs.markBudgetResetLocked(resetBudget.ID) gs.LastDBUsagesBudgetsMu.Unlock() gs.logger.Debug(fmt.Sprintf("Reset budget %s (was %.2f, reset to 0)", resetBudget.ID, oldUsage)) return resetBudget } +// ResetBudgetUsageInMemory clears one budget's usage on operator request, without +// moving its reset boundary. +// +// This is the deliberate counterpart to resetExpiredBudgetFromSnapshot above. That +// one is driven by a window closing, so it advances LastReset; this one is driven +// by a person choosing "reset usage", where the window has not closed and the +// boundary must stay exactly where it is. Every persistence path guards LastReset +// forward-only, so a boundary move here would be refused anyway. +// +// Clearing the LastDBUsages baseline is not optional. The dump path writes the +// difference between in-memory usage and that baseline, and in a cluster each +// node contributes CurrentUsage - LastDBUsage to the shared total. Leaving the +// old baseline behind would re-add the spend that was just cleared, or drive the +// contribution negative. +// +// Returns the reset snapshot and true, or (nil, false) when the budget is unknown. +func (gs *LocalGovernanceStore) ResetBudgetUsageInMemory(ctx context.Context, budgetID string) (*configstoreTables.TableBudget, bool) { + reset, ok := gs.RebaseBudget(ctx, budgetID, 0, nil) + if !ok { + return nil, false + } + gs.LastDBUsagesBudgetsMu.Lock() + gs.LastDBUsagesBudgets[budgetID] = 0 + gs.markBudgetResetLocked(budgetID) + gs.LastDBUsagesBudgetsMu.Unlock() + gs.logger.Debug(fmt.Sprintf("Reset budget %s usage on operator request (boundary unchanged at %s)", budgetID, reset.LastReset)) + return reset, true +} + // ResetExpiredBudgetsInMemory checks and resets budgets that have exceeded their reset duration. // With no budgetIDs it scans every budget; with IDs it only checks those budgets. // refreshReferences controls whether embedded owner references (VK, team, customer) are updated @@ -2574,15 +2677,30 @@ func (gs *LocalGovernanceStore) DumpRateLimits(ctx context.Context, tokenBaselin return nil } -// DumpBudgets dumps all budgets to the database +// DumpBudgets dumps all budgets to the database. +// +// Split into a snapshot and a write so the gap between them is expressible in a +// test: that gap is where an operator reset can land, and the reset generations +// taken here are what let the write recognise a row cleared underneath it. func (gs *LocalGovernanceStore) DumpBudgets(ctx context.Context, baselines map[string]float64) error { if gs.configStore == nil { return nil } + rows, gens := gs.snapshotBudgetRows(baselines) + return gs.writeBudgetRows(ctx, rows, gens) +} + +// snapshotBudgetRows captures every budget's current usage as dump rows, together +// with the reset generation each row was read at. +func (gs *LocalGovernanceStore) snapshotBudgetRows(baselines map[string]float64) ([]budgetDumpRow, map[string]uint64) { // This is to prevent nil pointer dereference if baselines == nil { baselines = map[string]float64{} } + // Read the generations first. Taking them before the budgets means a reset that + // races this snapshot is either fully visible in both, or shows up as a bumped + // generation later; it can never look unchanged while the usage read was stale. + gens := gs.budgetResetGensSnapshot() budgets := make(map[string]*configstoreTables.TableBudget) gs.budgets.Range(func(key, value interface{}) bool { // Type-safe conversion @@ -2616,13 +2734,23 @@ func (gs *LocalGovernanceStore) DumpBudgets(ctx context.Context, baselines map[s // Stable ID order keeps concurrent dumpers taking row locks in the same // sequence, which is what keeps them deadlock-free rather than merely lucky. sort.Slice(rows, func(i, j int) bool { return rows[i].ID < rows[j].ID }) + return rows, gens +} +// writeBudgetRows persists a snapshot taken by snapshotBudgetRows. +func (gs *LocalGovernanceStore) writeBudgetRows(ctx context.Context, rows []budgetDumpRow, gens map[string]uint64) error { + if gs.configStore == nil { + return nil + } // Written in batches so row locks are released between chunks rather than // held for the whole sweep, and so each chunk costs one round trip instead // of two per row. for start := 0; start < len(rows); start += dumpBatchSize { end := min(start+dumpBatchSize, len(rows)) - batch := rows[start:end] + batch := gs.keepUnresetRows(rows[start:end], gens) + if len(batch) == 0 { + continue + } if err := gs.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { return gs.writeBudgetBatch(ctx, tx, batch, "<=") }); err != nil { diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 3b4b745e90d..96158c613d4 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -47,6 +47,10 @@ func dbForUpdate(db *gorm.DB) *gorm.DB { type GovernanceManager interface { GetGovernanceData(ctx context.Context) *governance.GovernanceData 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 RemoveVirtualKey(ctx context.Context, id string) error ReloadTeam(ctx context.Context, id string) (*configstoreTables.TableTeam, error) RemoveTeam(ctx context.Context, id string) error @@ -533,11 +537,41 @@ func isRateLimitRemovalRequest(req *UpdateRateLimitRequest) bool { req.TokenResetDuration == nil && req.RequestResetDuration == nil } +// budgetUsageReset carries an operator's explicit "reset usage" choice through +// budget reconciliation and collects the budgets it was applied to. +// +// Two things make this necessary rather than a plain bool. UpdateBudget cannot +// clear usage - it copies CurrentUsage and LastReset back from the stored row on +// every config write, deliberately, so a config.json sync cannot replay stale +// values over live accounting - so the reset has to go through UpdateBudgetUsage, +// the store method that owns that column. And the in-memory governance store is +// what actually enforces spend, so the caller needs the list of budgets to clear +// there once the transaction has committed. +// +// A nil pointer means the caller is not resetting anything. +type budgetUsageReset struct { + requested bool + budgetIDs []string +} + +// apply zeroes one budget's persisted usage and records it for the in-memory +// pass. No-op when the operator did not ask for a reset. +func (r *budgetUsageReset) apply(ctx context.Context, store configstore.ConfigStore, budgetID string, tx *gorm.DB) error { + if r == nil || !r.requested { + return nil + } + if err := store.UpdateBudgetUsage(ctx, budgetID, 0, tx); err != nil { + return fmt.Errorf("failed to reset usage for budget %s: %w", budgetID, err) + } + r.budgetIDs = append(r.budgetIDs, budgetID) + return nil +} + // reconcileModelConfigBudgets upserts the desired set of budgets owned by a model config // (via TableBudget.ModelConfigID), preserving usage on matched rows and deleting removed // ones. It mutates mc.Budgets to the reconciled set. The model config row must already // exist (callers create it first). Mirrors the VK/team multi-budget reconciliation. -func (h *GovernanceHandler) reconcileModelConfigBudgets(ctx context.Context, tx *gorm.DB, mc *configstoreTables.TableModelConfig, requests []CreateBudgetRequest) error { +func (h *GovernanceHandler) reconcileModelConfigBudgets(ctx context.Context, tx *gorm.DB, mc *configstoreTables.TableModelConfig, requests []CreateBudgetRequest, usageReset *budgetUsageReset) error { seenDurations := make(map[string]bool, len(requests)) for _, b := range requests { if b.MaxLimit < 0 { @@ -570,6 +604,14 @@ func (h *GovernanceHandler) reconcileModelConfigBudgets(ctx context.Context, tx if err := h.configStore.UpdateBudget(ctx, &existing, tx); err != nil { return err } + // Usage is runtime-owned, so UpdateBudget above cannot clear it. Route + // the operator's explicit choice through the method that owns the column. + 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 { @@ -680,8 +722,8 @@ type vkModelConfigDesired struct { // handling: true treats perProvider as the full desired set (reconciling and removing absent // providers); false leaves all per-provider configs untouched (for a partial VK update that // omits provider_configs). -func (h *GovernanceHandler) syncVKGovernanceToModelConfigs(ctx context.Context, tx *gorm.DB, vk *configstoreTables.TableVirtualKey, top vkModelConfigDesired, perProvider []vkModelConfigDesired, reconcileProviders bool) error { - if err := h.reconcileVKModelConfig(ctx, tx, vk, top); err != nil { +func (h *GovernanceHandler) syncVKGovernanceToModelConfigs(ctx context.Context, tx *gorm.DB, vk *configstoreTables.TableVirtualKey, top vkModelConfigDesired, perProvider []vkModelConfigDesired, reconcileProviders bool, usageReset *budgetUsageReset) error { + if err := h.reconcileVKModelConfig(ctx, tx, vk, top, usageReset); err != nil { return err } if !reconcileProviders { @@ -693,7 +735,7 @@ func (h *GovernanceHandler) syncVKGovernanceToModelConfigs(ctx context.Context, continue } keep[*pg.provider] = true - if err := h.reconcileVKModelConfig(ctx, tx, vk, pg); err != nil { + if err := h.reconcileVKModelConfig(ctx, tx, vk, pg, usageReset); err != nil { return err } } @@ -717,7 +759,7 @@ func (h *GovernanceHandler) syncVKGovernanceToModelConfigs(ctx context.Context, } // reconcileVKModelConfig reconciles a single VK-scoped model config to the desired state. -func (h *GovernanceHandler) reconcileVKModelConfig(ctx context.Context, tx *gorm.DB, vk *configstoreTables.TableVirtualKey, d vkModelConfigDesired) error { +func (h *GovernanceHandler) reconcileVKModelConfig(ctx context.Context, tx *gorm.DB, vk *configstoreTables.TableVirtualKey, d vkModelConfigDesired, usageReset *budgetUsageReset) error { q := tx.Preload("Budgets").Where("scope = ? AND scope_id = ? AND model_name = ?", configstoreTables.ModelConfigScopeVirtualKey, vk.ID, configstoreTables.ModelConfigAllModels) if d.provider == nil { @@ -835,7 +877,7 @@ func (h *GovernanceHandler) reconcileVKModelConfig(ctx context.Context, tx *gorm } if d.budgetsProvided { - if err := h.reconcileModelConfigBudgets(ctx, tx, &mc, d.budgets); err != nil { + if err := h.reconcileModelConfigBudgets(ctx, tx, &mc, d.budgets, usageReset); err != nil { return err } } @@ -1526,7 +1568,7 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { budgets: req.Budgets, rateLimitProvided: req.RateLimit != nil, rateLimit: topRateLimit, - }, vkGovProviders, true); err != nil { + }, vkGovProviders, true, nil); err != nil { return err } if req.MCPConfigs != nil { @@ -1738,6 +1780,10 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "VirtualKey cannot be attached to both Team and Customer") return } + // The operator's explicit "reset usage" choice, surfaced by the UI's + // Preserve / Reset dialog. 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} // Parse expires_at when provided: a timestamp must be in the future, "" clears the expiry. var newExpiresAt *time.Time if req.ExpiresAt != nil && *req.ExpiresAt != "" { @@ -2010,7 +2056,7 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { top.rateLimit = rateLimitFromRequestFields(req.RateLimit.TokenMaxLimit, req.RateLimit.TokenResetDuration, req.RateLimit.RequestMaxLimit, req.RateLimit.RequestResetDuration) } } - if err := h.syncVKGovernanceToModelConfigs(ctx, tx, vk, top, vkGovProviders, req.ProviderConfigs != nil); err != nil { + if err := h.syncVKGovernanceToModelConfigs(ctx, tx, vk, top, vkGovProviders, req.ProviderConfigs != nil, usageReset); err != nil { return err } if req.MCPConfigs != nil { @@ -2138,6 +2184,17 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, "Virtual key updated in database but failed to reload in-memory state") return } + // Clear usage in the store that actually enforces spend. This must run after + // the reload, not before: ReloadVirtualKey rebuilds the virtual key in memory + // 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 { + logger.Error("failed to reset in-memory budget usage after update: %v", err) + SendError(ctx, 500, "Virtual key updated but budget usage reset did not take effect") + return + } + } // Per-user credential reconciliation when the VK's MCP allowlist // changed. Mirrors the AP-propagation path: enterprise orphans / @@ -3603,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); err != nil { + if err := h.reconcileModelConfigBudgets(ctx, tx, mc, req.Budgets, nil); err != nil { return err } } @@ -3953,7 +4010,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); err != nil { + if err := h.reconcileModelConfigBudgets(ctx, tx, &mc, *effectiveBudgets, nil); err != nil { return err } } diff --git a/transports/bifrost-http/handlers/pricing_override_test.go b/transports/bifrost-http/handlers/pricing_override_test.go index adc2c63d3b1..17649143282 100644 --- a/transports/bifrost-http/handlers/pricing_override_test.go +++ b/transports/bifrost-http/handlers/pricing_override_test.go @@ -23,6 +23,9 @@ type pricingOverrideTestGovernanceManager struct{} func (pricingOverrideTestGovernanceManager) GetGovernanceData(ctx context.Context) *governance.GovernanceData { return nil } +func (pricingOverrideTestGovernanceManager) ResetBudgetUsageInMemory(context.Context, string, []string) error { + return nil +} func (pricingOverrideTestGovernanceManager) ReloadVirtualKey(context.Context, string) (*configstoreTables.TableVirtualKey, error) { return nil, nil } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 70282ea77c5..66c1ddda3bf 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -100,6 +100,11 @@ type ServerCallbacks interface { RemoveCustomer(ctx context.Context, id string) error // 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 RemoveVirtualKey(ctx context.Context, id string) error // Provider related callbacks GetModelsForProvider(provider schemas.ModelProvider) []string @@ -466,6 +471,34 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t return virtualKey, nil } +// ResetBudgetUsageInMemory clears usage for the given budgets in the governance +// store that enforces spend, leaving each reset boundary untouched. +// +// The database write happens inside the update transaction through +// UpdateBudgetUsage, but the in-memory store is what enforcement actually reads, +// and ReloadVirtualKey deliberately carries the cached usage forward on a config +// reload. Without this step the reset would be visible in the database and +// nowhere else, and the next dump tick would write the cached value back over it. +// +// 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 { + if len(budgetIDs) == 0 { + return nil + } + governancePlugin, err := s.getGovernancePlugin() + if err != nil { + return err + } + store := governancePlugin.GetGovernanceStore() + for _, budgetID := range budgetIDs { + if _, ok := store.ResetBudgetUsageInMemory(ctx, budgetID); !ok { + logger.Debug("budget %s not present in the governance store; skipping usage reset", budgetID) + } + } + return nil +} + // RemoveVirtualKey removes a virtual key from the in-memory store func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) error { governancePlugin, err := s.getGovernancePlugin()