adds reset budget honor flow

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")
+}

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 {