budget overrides ui

framework/configstore/rdb.go

@@ -4473,6 +4473,45 @@ func (s *RDBConfigStore) UpdateBudget(ctx context.Context, budget *tables.TableB
 	return nil
 }
 
+// UpdateBudgetOverride atomically updates only override columns so concurrent usage changes are preserved.
+func (s *RDBConfigStore) UpdateBudgetOverride(ctx context.Context, id string, amount float64, mode tables.BudgetOverrideMode, cyclesRemaining int, tx ...*gorm.DB) (*tables.TableBudget, error) {
+	if len(tx) == 0 {
+		var updated *tables.TableBudget
+		err := s.DB().WithContext(ctx).Transaction(func(transaction *gorm.DB) error {
+			var err error
+			updated, err = s.UpdateBudgetOverride(ctx, id, amount, mode, cyclesRemaining, transaction)
+			return err
+		})
+		return updated, err
+	}
+
+	txDB := tx[0].WithContext(ctx)
+	var budget tables.TableBudget
+	if err := dbForUpdate(txDB).First(&budget, "id = ?", id).Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, ErrNotFound
+		}
+		return nil, err
+	}
+	if err := budget.SetOverride(amount, mode, cyclesRemaining); err != nil {
+		return nil, err
+	}
+	if err := txDB.Session(&gorm.Session{SkipHooks: true}).Model(&tables.TableBudget{}).
+		Where("id = ?", id).
+		Updates(map[string]any{
+			"override_amount":           budget.OverrideAmount,
+			"override_mode":             budget.OverrideMode,
+			"override_cycles_remaining": budget.OverrideCyclesRemaining,
+			"updated_at":                time.Now(),
+		}).Error; err != nil {
+		return nil, s.parseGormError(err)
+	}
+	if err := txDB.First(&budget, "id = ?", id).Error; err != nil {
+		return nil, err
+	}
+	return &budget, nil
+}
+
 // DeleteBudget deletes a budget from the database.
 func (s *RDBConfigStore) DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error {
 	if len(tx) == 0 {

framework/configstore/rdb_test.go

@@ -780,6 +780,36 @@ func TestCreateBudgetWithOverride(t *testing.T) {
 	}
 }
 
+// TestUpdateBudgetOverridePreservesBudgetState verifies the partial update cannot clobber usage or base configuration.
+func TestUpdateBudgetOverridePreservesBudgetState(t *testing.T) {
+	store := setupRDBTestStore(t)
+	ctx := context.Background()
+	budget := &tables.TableBudget{
+		ID:            "budget-override-partial-update",
+		MaxLimit:      100,
+		ResetDuration: "1d",
+		CurrentUsage:  40,
+	}
+	require.NoError(t, store.CreateBudget(ctx, budget))
+
+	updated, err := store.UpdateBudgetOverride(ctx, budget.ID, 25, tables.BudgetOverrideModeCycles, 4)
+	require.NoError(t, err)
+	assert.Equal(t, 100.0, updated.MaxLimit)
+	assert.Equal(t, "1d", updated.ResetDuration)
+	assert.Equal(t, 40.0, updated.CurrentUsage)
+	assert.Equal(t, 25.0, updated.OverrideAmount)
+	assert.Equal(t, tables.BudgetOverrideModeCycles, updated.OverrideMode)
+	assert.Equal(t, 4, updated.OverrideCyclesRemaining)
+
+	cleared, err := store.UpdateBudgetOverride(ctx, budget.ID, 0, "", 0)
+	require.NoError(t, err)
+	assert.Equal(t, 100.0, cleared.MaxLimit)
+	assert.Equal(t, 40.0, cleared.CurrentUsage)
+	assert.Zero(t, cleared.OverrideAmount)
+	assert.Empty(t, cleared.OverrideMode)
+	assert.Zero(t, cleared.OverrideCyclesRemaining)
+}
+
 func TestGetBudgets(t *testing.T) {
 	store := setupRDBTestStore(t)
 	ctx := context.Background()

framework/configstore/store.go

@@ -353,6 +353,8 @@ type ConfigStore interface {
 	GetBudget(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableBudget, error)
 	CreateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error
 	UpdateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error
+	// UpdateBudgetOverride updates only the override state and returns the refreshed budget.
+	UpdateBudgetOverride(ctx context.Context, id string, amount float64, mode tables.BudgetOverrideMode, cyclesRemaining int, tx ...*gorm.DB) (*tables.TableBudget, error)
 	UpdateBudgets(ctx context.Context, budgets []*tables.TableBudget, tx ...*gorm.DB) error
 	DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error
 	UpdateBudgetUsage(ctx context.Context, id string, currentUsage float64, tx ...*gorm.DB) error

framework/configstore/tables/budget.go

@@ -64,6 +64,68 @@ type TableBudget struct {
 // TableName sets the table name for each model
 func (TableBudget) TableName() string { return "governance_budgets" }
 
+// HasActiveOverride reports whether the budget currently has a valid configured override.
+func (b *TableBudget) HasActiveOverride() bool {
+	if b == nil || b.OverrideAmount <= 0 {
+		return false
+	}
+	return b.OverrideMode == BudgetOverrideModeForever ||
+		(b.OverrideMode == BudgetOverrideModeCycles && b.OverrideCyclesRemaining > 0)
+}
+
+// EffectiveMaxLimit returns the base limit plus any active override amount.
+func (b *TableBudget) EffectiveMaxLimit() float64 {
+	if b == nil {
+		return 0
+	}
+	if !b.HasActiveOverride() {
+		return b.MaxLimit
+	}
+	return b.MaxLimit + b.OverrideAmount
+}
+
+// SetOverride replaces the budget's current override after validating the complete state.
+func (b *TableBudget) SetOverride(amount float64, mode BudgetOverrideMode, cyclesRemaining int) error {
+	if b == nil {
+		return fmt.Errorf("budget is required")
+	}
+	previousAmount := b.OverrideAmount
+	previousMode := b.OverrideMode
+	previousCyclesRemaining := b.OverrideCyclesRemaining
+	b.OverrideAmount = amount
+	b.OverrideMode = mode
+	b.OverrideCyclesRemaining = cyclesRemaining
+	if err := b.validateOverride(); err != nil {
+		b.OverrideAmount = previousAmount
+		b.OverrideMode = previousMode
+		b.OverrideCyclesRemaining = previousCyclesRemaining
+		return err
+	}
+	return nil
+}
+
+// ClearOverride removes any finite or permanent override from the budget.
+func (b *TableBudget) ClearOverride() {
+	if b == nil {
+		return
+	}
+	b.OverrideAmount = 0
+	b.OverrideMode = ""
+	b.OverrideCyclesRemaining = 0
+}
+
+// ConsumeOverrideCycle advances a finite override by one completed budget reset cycle.
+func (b *TableBudget) ConsumeOverrideCycle() {
+	if b == nil || b.OverrideMode != BudgetOverrideModeCycles || !b.HasActiveOverride() {
+		return
+	}
+	if b.OverrideCyclesRemaining == 1 {
+		b.ClearOverride()
+		return
+	}
+	b.OverrideCyclesRemaining--
+}
+
 // validateOverride checks that the persisted override fields form one unambiguous state.
 func (b *TableBudget) validateOverride() error {
 	if math.IsNaN(b.OverrideAmount) || math.IsInf(b.OverrideAmount, 0) {

framework/configstore/tables/budget_test.go

@@ -4,9 +4,62 @@ import (
 	"math"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+// TestTableBudgetOverrideLifecycle verifies replacement, effective-limit, and clear behavior.
+func TestTableBudgetOverrideLifecycle(t *testing.T) {
+	budget := &TableBudget{MaxLimit: 100}
+	require.NoError(t, budget.SetOverride(25, BudgetOverrideModeCycles, 4))
+	assert.True(t, budget.HasActiveOverride())
+	assert.Equal(t, 125.0, budget.EffectiveMaxLimit())
+	assert.Equal(t, 4, budget.OverrideCyclesRemaining)
+
+	require.NoError(t, budget.SetOverride(50, BudgetOverrideModeForever, 0))
+	assert.True(t, budget.HasActiveOverride())
+	assert.Equal(t, 150.0, budget.EffectiveMaxLimit())
+	assert.Equal(t, BudgetOverrideModeForever, budget.OverrideMode)
+
+	budget.ClearOverride()
+	assert.False(t, budget.HasActiveOverride())
+	assert.Equal(t, 100.0, budget.EffectiveMaxLimit())
+	assert.Zero(t, budget.OverrideAmount)
+	assert.Empty(t, budget.OverrideMode)
+	assert.Zero(t, budget.OverrideCyclesRemaining)
+}
+
+// TestTableBudgetSetOverrideRestoresPreviousState verifies invalid replacements are non-mutating.
+func TestTableBudgetSetOverrideRestoresPreviousState(t *testing.T) {
+	budget := &TableBudget{MaxLimit: 100}
+	require.NoError(t, budget.SetOverride(25, BudgetOverrideModeCycles, 2))
+
+	require.Error(t, budget.SetOverride(50, BudgetOverrideModeCycles, 0))
+	assert.Equal(t, 25.0, budget.OverrideAmount)
+	assert.Equal(t, BudgetOverrideModeCycles, budget.OverrideMode)
+	assert.Equal(t, 2, budget.OverrideCyclesRemaining)
+}
+
+// TestTableBudgetConsumeOverrideCycle verifies finite overrides expire while permanent overrides survive resets.
+func TestTableBudgetConsumeOverrideCycle(t *testing.T) {
+	finite := &TableBudget{MaxLimit: 100}
+	require.NoError(t, finite.SetOverride(25, BudgetOverrideModeCycles, 2))
+
+	finite.ConsumeOverrideCycle()
+	assert.Equal(t, 1, finite.OverrideCyclesRemaining)
+	assert.Equal(t, 125.0, finite.EffectiveMaxLimit())
+
+	finite.ConsumeOverrideCycle()
+	assert.False(t, finite.HasActiveOverride())
+	assert.Equal(t, 100.0, finite.EffectiveMaxLimit())
+
+	permanent := &TableBudget{MaxLimit: 100}
+	require.NoError(t, permanent.SetOverride(50, BudgetOverrideModeForever, 0))
+	permanent.ConsumeOverrideCycle()
+	assert.True(t, permanent.HasActiveOverride())
+	assert.Equal(t, 150.0, permanent.EffectiveMaxLimit())
+}
+
 // TestTableBudgetValidateOverride verifies that persisted override fields cannot form an ambiguous state.
 func TestTableBudgetValidateOverride(t *testing.T) {
 	tests := []struct {

plugins/governance/store.go

@@ -448,6 +448,7 @@ func (gs *LocalGovernanceStore) BumpBudgetUsage(ctx context.Context, budgetID st
 			gs.logger.Error("budget %s reset target not converging after %d resets; applying inline reset to avoid request-path spin", budgetID, resetAttempts)
 			clone.CurrentUsage = 0
 			clone.LastReset = *target
+			clone.ConsumeOverrideCycle()
 			gs.LastDBUsagesBudgetsMu.Lock()
 			gs.LastDBUsagesBudgets[budgetID] = 0
 			gs.LastDBUsagesBudgetsMu.Unlock()
@@ -574,6 +575,7 @@ func (gs *LocalGovernanceStore) ResetBudgetAt(ctx context.Context, budgetID stri
 		clone := *old
 		clone.CurrentUsage = 0
 		clone.LastReset = newLastReset
+		clone.ConsumeOverrideCycle()
 		if gs.budgets.CompareAndSwap(budgetID, raw, &clone) {
 			return &clone, true
 		}
@@ -1069,13 +1071,14 @@ func (gs *LocalGovernanceStore) CheckBudget(ctx context.Context, entityWiseBudge
 			if !exists {
 				baseline = 0
 			}
+			effectiveMaxLimit := budget.EffectiveMaxLimit()
 			gs.logger.Debug("LocalStore CheckBudget: Checking %s budget %s: local=%.4f, remote=%.4f, total=%.4f, limit=%.4f",
-				entity, budget.ID, budget.CurrentUsage, baseline, budget.CurrentUsage+baseline, budget.MaxLimit)
+				entity, budget.ID, budget.CurrentUsage, baseline, budget.CurrentUsage+baseline, effectiveMaxLimit)
 			// Check if current usage (local + remote baseline) exceeds budget limit
-			if budget.CurrentUsage+baseline >= budget.MaxLimit {
+			if budget.CurrentUsage+baseline >= effectiveMaxLimit {
 				gs.logger.Debug("LocalStore CheckBudget: Budget %s EXCEEDED", budget.ID)
 				return DecisionBudgetExceeded, fmt.Errorf("%s budget exceeded: %.4f >= %.4f dollars",
-					entity, budget.CurrentUsage+baseline, budget.MaxLimit)
+					entity, budget.CurrentUsage+baseline, effectiveMaxLimit)
 			}
 		}
 	}
@@ -2037,6 +2040,28 @@ func (gs *LocalGovernanceStore) ResetExpiredBudgets(ctx context.Context, resetBu
 	if len(resetBudgets) > 0 && gs.configStore != nil {
 		if err := gs.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error {
 			for _, budget := range resetBudgets {
+				// Persist the finite-override lifecycle first, guarded on last_reset:
+				// a snapshot that lost the persistence race to a newer reset must not
+				// restore older lifecycle state (extra override cycles would survive a
+				// restart otherwise). This must run BEFORE the usage write below, which
+				// advances last_reset and would close the guard within this transaction.
+				overrideResult := tx.WithContext(ctx).
+					Session(&gorm.Session{SkipHooks: true}).
+					Model(&configstoreTables.TableBudget{}).
+					Where("id = ? AND last_reset < ?", budget.ID, budget.LastReset).
+					Updates(map[string]interface{}{
+						"override_amount":           budget.OverrideAmount,
+						"override_mode":             budget.OverrideMode,
+						"override_cycles_remaining": budget.OverrideCyclesRemaining,
+					})
+
+				if overrideResult.Error != nil {
+					return fmt.Errorf("failed to reset budget override lifecycle %s: %w", budget.ID, overrideResult.Error)
+				}
+				if overrideResult.RowsAffected == 0 {
+					gs.logger.Debug("skipping stale override reset persistence for budget %s: database already holds newer reset state", budget.ID)
+				}
+
 				// Direct UPDATE only resets current_usage and last_reset
 				// This prevents overwriting max_limit or reset_duration that may have been changed by other nodes/requests
 				result := tx.WithContext(ctx).
@@ -2228,6 +2253,27 @@ func (gs *LocalGovernanceStore) DumpBudgets(ctx context.Context, baselines map[s
 					newUsage += baseline
 				}
 
+				// The override trio mutated by ConsumeOverrideCycle is flushed first
+				// under a monotonic last_reset guard: it only fires when this node
+				// performed a reset the database has not seen, which is exactly when
+				// its override snapshot is authoritative. An unguarded write would let
+				// a node with a stale in-memory override clobber a newer admin update
+				// every dump cycle. This must run BEFORE the usage write below, which
+				// advances last_reset and would close the guard within this transaction.
+				overrideResult := tx.WithContext(ctx).
+					Session(&gorm.Session{SkipHooks: true}).
+					Model(&configstoreTables.TableBudget{}).
+					Where("id = ? AND last_reset < ?", inMemoryBudget.ID, inMemoryBudget.LastReset).
+					Updates(map[string]interface{}{
+						"override_amount":           inMemoryBudget.OverrideAmount,
+						"override_mode":             inMemoryBudget.OverrideMode,
+						"override_cycles_remaining": inMemoryBudget.OverrideCyclesRemaining,
+					})
+
+				if overrideResult.Error != nil {
+					return fmt.Errorf("failed to update budget override lifecycle %s: %w", inMemoryBudget.ID, overrideResult.Error)
+				}
+
 				// Direct UPDATE avoids read-then-write lock escalation that causes deadlocks
 				// Use Session with SkipHooks to avoid triggering BeforeSave hook validation
 				result := tx.WithContext(ctx).
@@ -4126,8 +4172,8 @@ func (gs *LocalGovernanceStore) GetBudgetAndRateLimitStatus(ctx context.Context,
 		for bi := range modelConfig.Budgets {
 			if budgetValue, ok := gs.budgets.Load(modelConfig.Budgets[bi].ID); ok && budgetValue != nil {
 				if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
-					if budget.MaxLimit > 0 {
-						budgetPercent := float64(budget.CurrentUsage+budgetBaselines[budget.ID]) / budget.MaxLimit * 100
+					if effectiveMaxLimit := budget.EffectiveMaxLimit(); effectiveMaxLimit > 0 {
+						budgetPercent := float64(budget.CurrentUsage+budgetBaselines[budget.ID]) / effectiveMaxLimit * 100
 						if budgetPercent > result.BudgetPercentUsed {
 							result.BudgetPercentUsed = budgetPercent
 						}
@@ -4185,8 +4231,8 @@ func (gs *LocalGovernanceStore) GetBudgetAndRateLimitStatus(ctx context.Context,
 			if providerTable.BudgetID != nil {
 				if budgetValue, ok := gs.budgets.Load(*providerTable.BudgetID); ok && budgetValue != nil {
 					if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
-						if budget.MaxLimit > 0 {
-							budgetPercent := float64(budget.CurrentUsage+budgetBaselines[budget.ID]) / budget.MaxLimit * 100
+						if effectiveMaxLimit := budget.EffectiveMaxLimit(); effectiveMaxLimit > 0 {
+							budgetPercent := float64(budget.CurrentUsage+budgetBaselines[budget.ID]) / effectiveMaxLimit * 100
 							if budgetPercent > result.BudgetPercentUsed {
 								result.BudgetPercentUsed = budgetPercent
 							}
@@ -4229,8 +4275,8 @@ func (gs *LocalGovernanceStore) GetBudgetAndRateLimitStatus(ctx context.Context,
 					for _, b := range pc.Budgets {
 						if budgetValue, ok := gs.budgets.Load(b.ID); ok && budgetValue != nil {
 							if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
-								if budget.MaxLimit > 0 {
-									budgetPercent := float64(budget.CurrentUsage+budgetBaselines[budget.ID]) / budget.MaxLimit * 100
+								if effectiveMaxLimit := budget.EffectiveMaxLimit(); effectiveMaxLimit > 0 {
+									budgetPercent := float64(budget.CurrentUsage+budgetBaselines[budget.ID]) / effectiveMaxLimit * 100
 									if budgetPercent > result.BudgetPercentUsed {
 										result.BudgetPercentUsed = budgetPercent
 									}