diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index d03feaf087e..74b89ccdd21 100644 --- a/framework/configstore/rdb.go +++ b/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 { diff --git a/framework/configstore/rdb_test.go b/framework/configstore/rdb_test.go index 2e016d8bb11..88a643cb6ec 100644 --- a/framework/configstore/rdb_test.go +++ b/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() diff --git a/framework/configstore/store.go b/framework/configstore/store.go index 26ec79ab47b..f46c5365635 100644 --- a/framework/configstore/store.go +++ b/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 diff --git a/framework/configstore/tables/budget.go b/framework/configstore/tables/budget.go index 5b3a19e1edc..cadf9bc66b4 100644 --- a/framework/configstore/tables/budget.go +++ b/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) { diff --git a/framework/configstore/tables/budget_test.go b/framework/configstore/tables/budget_test.go index d6d296f06f3..96bbb7a4f08 100644 --- a/framework/configstore/tables/budget_test.go +++ b/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 { diff --git a/plugins/governance/store.go b/plugins/governance/store.go index 18f72e52c80..7ea740177ce 100644 --- a/plugins/governance/store.go +++ b/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 } diff --git a/plugins/governance/store_test.go b/plugins/governance/store_test.go index 620dfb4dbda..dd629c3961a 100644 --- a/plugins/governance/store_test.go +++ b/plugins/governance/store_test.go @@ -163,6 +163,23 @@ func TestGovernanceStore_CheckBudget_SingleBudget(t *testing.T) { } } +// TestGovernanceStoreCheckBudgetUsesOverride verifies enforcement reads the additive effective limit. +func TestGovernanceStoreCheckBudgetUsesOverride(t *testing.T) { + store, err := NewLocalGovernanceStore(context.Background(), NewMockLogger(), nil, &configstore.GovernanceConfig{}, nil) + require.NoError(t, err) + budget := buildBudgetWithUsage("override-budget", 100, 110, "1d") + require.NoError(t, budget.SetOverride(25, configstoreTables.BudgetOverrideModeCycles, 2)) + + decision, err := store.CheckBudget(context.Background(), EntityWiseBudgets{"VirtualKey": {budget}}, nil) + require.NoError(t, err) + assert.Equal(t, DecisionAllow, decision) + + decision, err = store.CheckBudget(context.Background(), EntityWiseBudgets{"VirtualKey": {budget}}, map[string]float64{budget.ID: 15}) + require.Error(t, err) + assert.Equal(t, DecisionBudgetExceeded, decision) + assert.Contains(t, err.Error(), "125.0000 dollars") +} + // TestGovernanceStore_CheckBudget_HierarchyValidation tests multi-level budget hierarchy func TestGovernanceStore_CheckBudget_HierarchyValidation(t *testing.T) { logger := NewMockLogger() @@ -785,6 +802,88 @@ func TestGovernanceStore_ResetExpiredBudgets(t *testing.T) { assert.Equal(t, 0.0, updatedVK.Budgets[0].CurrentUsage, "Budget usage should be reset") } +// TestGovernanceStoreResetBudgetAdvancesOverride verifies each existing reset consumes one finite cycle. +func TestGovernanceStoreResetBudgetAdvancesOverride(t *testing.T) { + store := newStandaloneStore(t) + finite := buildBudgetWithUsage("finite-override", 100, 75, "1d") + require.NoError(t, finite.SetOverride(25, configstoreTables.BudgetOverrideModeCycles, 2)) + store.budgets.Store(finite.ID, finite) + + firstReset := finite.LastReset.Add(24 * time.Hour) + reset, ok := store.ResetBudgetAt(context.Background(), finite.ID, firstReset) + require.True(t, ok) + assert.Zero(t, reset.CurrentUsage) + assert.Equal(t, 1, reset.OverrideCyclesRemaining) + assert.Equal(t, 125.0, reset.EffectiveMaxLimit()) + + secondReset := firstReset.Add(24 * time.Hour) + reset, ok = store.ResetBudgetAt(context.Background(), finite.ID, secondReset) + require.True(t, ok) + assert.False(t, reset.HasActiveOverride()) + assert.Equal(t, 100.0, reset.EffectiveMaxLimit()) + + permanent := buildBudgetWithUsage("forever-override", 100, 75, "1d") + require.NoError(t, permanent.SetOverride(50, configstoreTables.BudgetOverrideModeForever, 0)) + store.budgets.Store(permanent.ID, permanent) + reset, ok = store.ResetBudgetAt(context.Background(), permanent.ID, permanent.LastReset.Add(24*time.Hour)) + require.True(t, ok) + assert.True(t, reset.HasActiveOverride()) + assert.Equal(t, 150.0, reset.EffectiveMaxLimit()) +} + +// TestGovernanceStoreUpsertBudgetConfigRefreshesOverride verifies cache refreshes config without clobbering runtime counters. +func TestGovernanceStoreUpsertBudgetConfigRefreshesOverride(t *testing.T) { + store := newStandaloneStore(t) + lastReset := time.Now().Add(-30 * time.Minute) + live := buildBudgetWithUsage("cache-override", 100, 40, "1h") + live.LastReset = lastReset + store.budgets.Store(live.ID, live) + + refreshed := buildBudgetWithUsage(live.ID, 120, 0, "2h") + require.NoError(t, refreshed.SetOverride(30, configstoreTables.BudgetOverrideModeCycles, 3)) + store.UpsertBudgetConfig(context.Background(), live.ID, refreshed) + + got := store.LoadBudget(context.Background(), live.ID) + require.NotNil(t, got) + assert.Equal(t, 40.0, got.CurrentUsage) + assert.True(t, got.LastReset.Equal(lastReset)) + assert.Equal(t, 120.0, got.MaxLimit) + assert.Equal(t, "2h", got.ResetDuration) + assert.Equal(t, 30.0, got.OverrideAmount) + assert.Equal(t, configstoreTables.BudgetOverrideModeCycles, got.OverrideMode) + assert.Equal(t, 3, got.OverrideCyclesRemaining) +} + +// TestGovernanceStoreResetPersistsOverrideLifecycle verifies the existing reset write stores the decremented cycle state. +func TestGovernanceStoreResetPersistsOverrideLifecycle(t *testing.T) { + ctx := context.Background() + logger := NewMockLogger() + configStore, err := configstore.NewConfigStore(ctx, &configstore.Config{ + Enabled: true, + Type: configstore.ConfigStoreTypeSQLite, + Config: &configstore.SQLiteConfig{Path: t.TempDir() + "/governance.db"}, + }, logger) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, configStore.Close(ctx)) }) + + budget := buildBudgetWithUsage("persisted-override", 100, 75, "1d") + require.NoError(t, budget.SetOverride(25, configstoreTables.BudgetOverrideModeCycles, 2)) + require.NoError(t, configStore.CreateBudget(ctx, budget)) + + store, err := NewLocalGovernanceStore(ctx, logger, configStore, nil, nil) + require.NoError(t, err) + reset, ok := store.ResetBudgetAt(ctx, budget.ID, budget.LastReset.Add(24*time.Hour)) + require.True(t, ok) + require.NoError(t, store.ResetExpiredBudgets(ctx, []*configstoreTables.TableBudget{reset})) + + persisted, err := configStore.GetBudget(ctx, budget.ID) + require.NoError(t, err) + assert.Zero(t, persisted.CurrentUsage) + assert.Equal(t, 25.0, persisted.OverrideAmount) + assert.Equal(t, configstoreTables.BudgetOverrideModeCycles, persisted.OverrideMode) + assert.Equal(t, 1, persisted.OverrideCyclesRemaining) +} + // TestGovernanceStore_GetAllBudgets tests retrieving all budgets func TestGovernanceStore_GetAllBudgets(t *testing.T) { logger := NewMockLogger() diff --git a/plugins/governance/storeconcurrency_test.go b/plugins/governance/storeconcurrency_test.go index f3168954265..489481e4ede 100644 --- a/plugins/governance/storeconcurrency_test.go +++ b/plugins/governance/storeconcurrency_test.go @@ -97,6 +97,9 @@ func TestResetBudgetAt_ConcurrentResettersCollapse(t *testing.T) { old := buildBudget(budgetID, 1000, "1h") old.LastReset = time.Now().Add(-2 * time.Hour) old.CurrentUsage = 999 + old.OverrideAmount = 25 + old.OverrideMode = "cycles" + old.OverrideCyclesRemaining = 5 store.budgets.Store(budgetID, old) const goroutines = 128 @@ -120,4 +123,5 @@ func TestResetBudgetAt_ConcurrentResettersCollapse(t *testing.T) { require.NotNil(t, final) assert.Equal(t, 0.0, final.CurrentUsage) assert.True(t, final.LastReset.Equal(newLastReset)) + assert.Equal(t, 4, final.OverrideCyclesRemaining, "the single winning reset should consume exactly one override cycle") } diff --git a/tests/cmd/e2eseed/go.mod b/tests/cmd/e2eseed/go.mod index 0f55bc9b6f2..9fa91bb19f2 100644 --- a/tests/cmd/e2eseed/go.mod +++ b/tests/cmd/e2eseed/go.mod @@ -85,7 +85,7 @@ require ( github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-sqlite3 v1.14.32 // indirect - github.com/maximhq/bifrost/core v1.7.1 // indirect + github.com/maximhq/bifrost/core v1.7.3 // indirect github.com/maximhq/bifrost/framework v1.3.16 // indirect github.com/paulmach/orb v0.11.1 // indirect github.com/pierrec/lz4/v4 v4.1.22 // indirect diff --git a/tests/cmd/seed/go.mod b/tests/cmd/seed/go.mod index 45073278793..97224cc2dd4 100644 --- a/tests/cmd/seed/go.mod +++ b/tests/cmd/seed/go.mod @@ -8,7 +8,7 @@ replace ( ) require ( - github.com/maximhq/bifrost/core v1.7.1 + github.com/maximhq/bifrost/core v1.7.3 github.com/maximhq/bifrost/framework v1.3.16 gorm.io/driver/postgres v1.6.0 gorm.io/driver/sqlite v1.6.0 diff --git a/tests/cmd/seedvks/go.mod b/tests/cmd/seedvks/go.mod index 3ee938ec7e8..04956ede0b6 100644 --- a/tests/cmd/seedvks/go.mod +++ b/tests/cmd/seedvks/go.mod @@ -9,7 +9,7 @@ replace ( require ( github.com/google/uuid v1.6.0 - github.com/maximhq/bifrost/core v1.7.1 + github.com/maximhq/bifrost/core v1.7.3 github.com/maximhq/bifrost/framework v1.3.16 gorm.io/driver/postgres v1.6.0 gorm.io/gorm v1.31.1 diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 931511b492e..a3afb03cc9f 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -265,6 +265,19 @@ type UpdateBudgetRequest struct { ResetDuration *string `json:"reset_duration,omitempty"` } +// BudgetOverrideRequest replaces the active override on one budget. +type BudgetOverrideRequest struct { + Amount float64 `json:"amount"` + Mode configstoreTables.BudgetOverrideMode `json:"mode"` + Cycles int `json:"cycles,omitempty"` +} + +// BudgetOverrideResponse returns the persisted budget and its additive effective limit. +type BudgetOverrideResponse struct { + Budget *configstoreTables.TableBudget `json:"budget"` + EffectiveMaxLimit float64 `json:"effective_max_limit"` +} + // RoutingTarget represents a single weighted routing target within a rule. // All fields except Weight are optional; nil means "use the incoming request's value". // Weights across all targets in a rule must sum to 1 (e.g. 0.7 + 0.3 = 1.0). @@ -997,6 +1010,8 @@ func (h *GovernanceHandler) RegisterRoutes(r *router.Router, middlewares ...sche r.GET("/api/governance/virtual-keys/{vk_id}", lib.ChainMiddlewares(h.getVirtualKey, middlewares...)) r.PUT("/api/governance/virtual-keys/{vk_id}", lib.ChainMiddlewares(h.updateVirtualKey, middlewares...)) r.POST("/api/governance/virtual-keys/{vk_id}/rotate", lib.ChainMiddlewares(h.rotateVirtualKey, middlewares...)) + r.PUT("/api/governance/virtual-keys/{vk_id}/budgets/{budget_id}/override", lib.ChainMiddlewares(h.updateVirtualKeyBudgetOverride, middlewares...)) + r.DELETE("/api/governance/virtual-keys/{vk_id}/budgets/{budget_id}/override", lib.ChainMiddlewares(h.deleteVirtualKeyBudgetOverride, middlewares...)) r.DELETE("/api/governance/virtual-keys/{vk_id}", lib.ChainMiddlewares(h.deleteVirtualKey, middlewares...)) // Team CRUD operations @@ -1515,6 +1530,94 @@ func (h *GovernanceHandler) getVirtualKey(ctx *fasthttp.RequestCtx) { }) } +// loadVirtualKeyBudget resolves only budgets owned by the virtual key's scoped model configs. +func (h *GovernanceHandler) loadVirtualKeyBudget(ctx context.Context, vkID, budgetID string) (*configstoreTables.TableBudget, error) { + if _, err := h.configStore.GetVirtualKey(ctx, vkID); err != nil { + return nil, err + } + modelConfigs, err := h.configStore.GetModelConfigsByScopeAndScopeIDs( + ctx, + configstoreTables.ModelConfigScopeVirtualKey, + []string{vkID}, + ) + if err != nil { + return nil, err + } + for i := range modelConfigs { + for j := range modelConfigs[i].Budgets { + if modelConfigs[i].Budgets[j].ID == budgetID { + return &modelConfigs[i].Budgets[j], nil + } + } + } + return nil, configstore.ErrNotFound +} + +// updateVirtualKeyBudgetOverride handles PUT for a standalone virtual-key budget override. +func (h *GovernanceHandler) updateVirtualKeyBudgetOverride(ctx *fasthttp.RequestCtx) { + h.mutateVirtualKeyBudgetOverride(ctx, false) +} + +// deleteVirtualKeyBudgetOverride handles DELETE for a standalone virtual-key budget override. +func (h *GovernanceHandler) deleteVirtualKeyBudgetOverride(ctx *fasthttp.RequestCtx) { + h.mutateVirtualKeyBudgetOverride(ctx, true) +} + +// mutateVirtualKeyBudgetOverride replaces or clears an override without changing base budget configuration or usage. +func (h *GovernanceHandler) mutateVirtualKeyBudgetOverride(ctx *fasthttp.RequestCtx, clear bool) { + vkID := ctx.UserValue("vk_id").(string) + budgetID := ctx.UserValue("budget_id").(string) + budget, err := h.loadVirtualKeyBudget(ctx, vkID, budgetID) + if err != nil { + if errors.Is(err, configstore.ErrNotFound) { + SendError(ctx, fasthttp.StatusNotFound, "virtual key or budget not found") + return + } + SendError(ctx, fasthttp.StatusInternalServerError, "failed to retrieve virtual key budget") + return + } + + if clear { + budget.ClearOverride() + } else { + var req BudgetOverrideRequest + if err := json.Unmarshal(ctx.PostBody(), &req); err != nil { + SendError(ctx, fasthttp.StatusBadRequest, "invalid request body") + return + } + if err := budget.SetOverride(req.Amount, req.Mode, req.Cycles); err != nil { + SendError(ctx, fasthttp.StatusBadRequest, err.Error()) + return + } + } + + budget, err = h.configStore.UpdateBudgetOverride( + ctx, + budget.ID, + budget.OverrideAmount, + budget.OverrideMode, + budget.OverrideCyclesRemaining, + ) + if err != nil { + if errors.Is(err, configstore.ErrNotFound) { + SendError(ctx, fasthttp.StatusNotFound, "budget not found") + return + } + SendError(ctx, fasthttp.StatusInternalServerError, "failed to update budget override") + return + } + if _, err := h.governanceManager.ReloadVirtualKey(ctx, vkID); err != nil { + logger.Error("failed to reload virtual key after budget override: %v", err) + SendError(ctx, fasthttp.StatusInternalServerError, "budget override saved but virtual key reload failed") + return + } + + SendJSON(ctx, BudgetOverrideResponse{ + Budget: budget, + EffectiveMaxLimit: budget.EffectiveMaxLimit(), + }) +} + // updateVirtualKey handles PUT /api/governance/virtual-keys/{vk_id} - Update a virtual key func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { vkID := ctx.UserValue("vk_id").(string) diff --git a/transports/bifrost-http/handlers/governance_test.go b/transports/bifrost-http/handlers/governance_test.go index ceae46eb92b..a01913d0339 100644 --- a/transports/bifrost-http/handlers/governance_test.go +++ b/transports/bifrost-http/handlers/governance_test.go @@ -120,6 +120,19 @@ type mockRotateGovernanceManager struct { reloadErr error } +// budgetOverrideTestGovernanceManager records reloads after budget override mutations. +type budgetOverrideTestGovernanceManager struct { + GovernanceManager + store configstore.ConfigStore + reloadIDs []string +} + +// ReloadVirtualKey records the reload and returns the current persisted virtual key. +func (m *budgetOverrideTestGovernanceManager) ReloadVirtualKey(ctx context.Context, id string) (*configstoreTables.TableVirtualKey, error) { + m.reloadIDs = append(m.reloadIDs, id) + return m.store.GetVirtualKey(ctx, id) +} + func (m *mockRotateGovernanceManager) ReloadVirtualKey(ctx context.Context, id string) (*configstoreTables.TableVirtualKey, error) { m.reloadIDs = append(m.reloadIDs, id) if m.reloadErr != nil { @@ -128,6 +141,128 @@ func (m *mockRotateGovernanceManager) ReloadVirtualKey(ctx context.Context, id s return m.store.GetVirtualKey(ctx, id) } +// TestVirtualKeyBudgetOverrideLifecycle verifies finite, replacement, and clear mutations preserve base budget state. +func TestVirtualKeyBudgetOverrideLifecycle(t *testing.T) { + SetLogger(&mockLogger{}) + store := setupPricingOverrideHandlerStore(t) + manager := &budgetOverrideTestGovernanceManager{store: store} + handler := &GovernanceHandler{configStore: store, governanceManager: manager} + ctx := context.Background() + + active := true + vk := &configstoreTables.TableVirtualKey{ + ID: "vk-budget-override", + Name: "override-test", + Value: *schemas.NewSecretVar("sk-bf-override-test"), + IsActive: &active, + } + if err := store.CreateVirtualKey(ctx, vk); err != nil { + t.Fatalf("create virtual key: %v", err) + } + scopeID := vk.ID + modelConfig := &configstoreTables.TableModelConfig{ + ID: "mc-budget-override", + ModelName: configstoreTables.ModelConfigAllModels, + Scope: configstoreTables.ModelConfigScopeVirtualKey, + ScopeID: &scopeID, + Budgets: []configstoreTables.TableBudget{{ + ID: "budget-override", + MaxLimit: 100, + CurrentUsage: 40, + ResetDuration: "1d", + }}, + } + if err := store.CreateModelConfig(ctx, modelConfig); err != nil { + t.Fatalf("create model config: %v", err) + } + + putCtx := newTestRequestCtx(`{"amount":25,"mode":"cycles","cycles":4}`) + putCtx.SetUserValue("vk_id", vk.ID) + putCtx.SetUserValue("budget_id", "budget-override") + handler.updateVirtualKeyBudgetOverride(putCtx) + if putCtx.Response.StatusCode() != fasthttp.StatusOK { + t.Fatalf("finite override status=%d body=%s", putCtx.Response.StatusCode(), putCtx.Response.Body()) + } + + budget, err := store.GetBudget(ctx, "budget-override") + if err != nil { + t.Fatalf("get finite override budget: %v", err) + } + if budget.MaxLimit != 100 || budget.CurrentUsage != 40 || budget.OverrideAmount != 25 || budget.OverrideMode != configstoreTables.BudgetOverrideModeCycles || budget.OverrideCyclesRemaining != 4 { + t.Fatalf("unexpected finite override budget: %+v", budget) + } + + replaceCtx := newTestRequestCtx(`{"amount":50,"mode":"forever"}`) + replaceCtx.SetUserValue("vk_id", vk.ID) + replaceCtx.SetUserValue("budget_id", "budget-override") + handler.updateVirtualKeyBudgetOverride(replaceCtx) + if replaceCtx.Response.StatusCode() != fasthttp.StatusOK { + t.Fatalf("forever override status=%d body=%s", replaceCtx.Response.StatusCode(), replaceCtx.Response.Body()) + } + + deleteCtx := newTestRequestCtx("") + deleteCtx.SetUserValue("vk_id", vk.ID) + deleteCtx.SetUserValue("budget_id", "budget-override") + handler.deleteVirtualKeyBudgetOverride(deleteCtx) + if deleteCtx.Response.StatusCode() != fasthttp.StatusOK { + t.Fatalf("clear override status=%d body=%s", deleteCtx.Response.StatusCode(), deleteCtx.Response.Body()) + } + budget, err = store.GetBudget(ctx, "budget-override") + if err != nil { + t.Fatalf("get cleared override budget: %v", err) + } + if budget.OverrideAmount != 0 || budget.OverrideMode != "" || budget.OverrideCyclesRemaining != 0 { + t.Fatalf("override was not cleared: %+v", budget) + } + if len(manager.reloadIDs) != 3 { + t.Fatalf("reload calls=%d, want 3", len(manager.reloadIDs)) + } +} + +// TestVirtualKeyBudgetOverrideRejectsDirectMirrorBudget verifies AP-style direct VK budgets cannot be overridden through the OSS endpoint. +func TestVirtualKeyBudgetOverrideRejectsDirectMirrorBudget(t *testing.T) { + SetLogger(&mockLogger{}) + store := setupPricingOverrideHandlerStore(t) + manager := &budgetOverrideTestGovernanceManager{store: store} + handler := &GovernanceHandler{configStore: store, governanceManager: manager} + ctx := context.Background() + + active := true + vk := &configstoreTables.TableVirtualKey{ + ID: "vk-ap-managed", + Name: "ap-managed", + Value: *schemas.NewSecretVar("sk-bf-ap-managed"), + IsActive: &active, + } + if err := store.CreateVirtualKey(ctx, vk); err != nil { + t.Fatalf("create virtual key: %v", err) + } + directBudget := &configstoreTables.TableBudget{ + ID: "ap-mirror-budget", + MaxLimit: 100, + ResetDuration: "1d", + VirtualKeyID: &vk.ID, + } + if err := store.CreateBudget(ctx, directBudget); err != nil { + t.Fatalf("create direct budget: %v", err) + } + + putCtx := newTestRequestCtx(`{"amount":25,"mode":"forever"}`) + putCtx.SetUserValue("vk_id", vk.ID) + putCtx.SetUserValue("budget_id", directBudget.ID) + handler.updateVirtualKeyBudgetOverride(putCtx) + if putCtx.Response.StatusCode() != fasthttp.StatusNotFound { + t.Fatalf("status=%d, want 404; body=%s", putCtx.Response.StatusCode(), putCtx.Response.Body()) + } + stored, err := store.GetBudget(ctx, directBudget.ID) + if err != nil { + t.Fatalf("get direct budget: %v", err) + } + if stored.OverrideMode != "" { + t.Fatalf("direct mirror override changed unexpectedly: %+v", stored) + } +} + type mockComplexityGovernanceManager struct { GovernanceManager reloadedConfig *complexity.AnalyzerConfig diff --git a/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts b/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts index 414743dafb1..000ac55e6ad 100644 --- a/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts +++ b/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts @@ -5,6 +5,9 @@ export interface AccessProfileBudgetLine { reset_duration: string; current_usage: number; last_reset: string; + override_amount?: number; + override_mode?: "cycles" | "forever"; + override_cycles_remaining?: number; alert_thresholds?: number[]; } @@ -38,4 +41,4 @@ export interface UserAccessProfile { export interface GetUserAccessProfilesResponse { access_profiles: UserAccessProfile[]; -} \ No newline at end of file +} diff --git a/ui/app/workspace/virtual-keys/hooks/useVirtualKeyUsage.ts b/ui/app/workspace/virtual-keys/hooks/useVirtualKeyUsage.ts index 9ea7ebad9b8..d6bc6662258 100644 --- a/ui/app/workspace/virtual-keys/hooks/useVirtualKeyUsage.ts +++ b/ui/app/workspace/virtual-keys/hooks/useVirtualKeyUsage.ts @@ -1,4 +1,5 @@ import { Budget, RateLimit, VirtualKey } from "@/lib/types/governance"; +import { getEffectiveBudgetLimit } from "@/lib/utils/governance"; import { useGetUserAccessProfilesQuery } from "@enterprise/lib/store/apis/accessProfileApi"; import { useGetVirtualKeyUsersQuery } from "@enterprise/lib/store/apis/virtualKeyUsersApi"; import { UserAccessProfile } from "@enterprise/lib/types/accessProfile"; @@ -47,6 +48,9 @@ export function useVirtualKeyUsage(vk: VirtualKey | null | undefined): { reset_duration: line.reset_duration, current_usage: line.current_usage, last_reset: line.last_reset, + override_amount: line.override_amount, + override_mode: line.override_mode, + override_cycles_remaining: line.override_cycles_remaining, })) : vk?.budgets; @@ -71,7 +75,7 @@ export function useVirtualKeyUsage(vk: VirtualKey | null | undefined): { : vk?.rate_limit; const isExhausted = - (displayBudgets?.some((b) => b.current_usage >= b.max_limit) ?? false) || + (displayBudgets?.some((b) => b.current_usage >= getEffectiveBudgetLimit(b)) ?? false) || (displayRateLimit?.token_current_usage != null && displayRateLimit?.token_max_limit != null && displayRateLimit.token_current_usage >= displayRateLimit.token_max_limit) || diff --git a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx index 0a4166dd0e2..150eb47192d 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx @@ -1,3 +1,4 @@ +import { BudgetOverrideDialog } from "@/components/budgetOverrideDialog"; import { SheetNavigationButtons } from "@/components/sheetNavigationButtons"; import { Badge } from "@/components/ui/badge"; import { Label } from "@/components/ui/label"; @@ -9,10 +10,18 @@ import { useSheetNavigation } from "@/hooks/useSheetNavigation"; import { supportsCalendarAlignment } from "@/lib/constants/governance"; import { ProviderIconType, RenderProviderIcon } from "@/lib/constants/icons"; import { ProviderLabels, ProviderName } from "@/lib/constants/logs"; -import { VirtualKey } from "@/lib/types/governance"; +import { useRemoveVirtualKeyBudgetOverrideMutation, useSetVirtualKeyBudgetOverrideMutation } from "@/lib/store/apis/governanceApi"; +import { BudgetOverrideRequest, VirtualKey } from "@/lib/types/governance"; import { cn } from "@/lib/utils"; -import { calculateUsagePercentage, formatCurrency, parseResetPeriod } from "@/lib/utils/governance"; +import { + calculateUsagePercentage, + formatCurrency, + getEffectiveBudgetLimit, + hasActiveBudgetOverride, + parseResetPeriod, +} from "@/lib/utils/governance"; import ManagedVirtualKeyNotice from "@enterprise/components/access-profiles/managedVirtualKeyNotice"; +import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { formatDistanceToNow } from "date-fns"; import { Users } from "lucide-react"; import { useVirtualKeyUsage } from "../hooks/useVirtualKeyUsage"; @@ -63,6 +72,15 @@ export default function VirtualKeyDetailSheet({ }: VirtualKeyDetailSheetProps) { const { assignedUsers, isManagedByProfile, managingProfile, hasApRateLimit, displayBudgets, displayRateLimit } = useVirtualKeyUsage(virtualKey); + const canUpdateVirtualKeys = useRbac(RbacResource.VirtualKeys, RbacOperation.Update); + const [setBudgetOverride] = useSetVirtualKeyBudgetOverrideMutation(); + const [removeBudgetOverride] = useRemoveVirtualKeyBudgetOverrideMutation(); + const saveBudgetOverride = async (budgetId: string, data: BudgetOverrideRequest) => { + await setBudgetOverride({ vkId: virtualKey.id, budgetId, data }).unwrap(); + }; + const clearBudgetOverride = async (budgetId: string) => { + await removeBudgetOverride({ vkId: virtualKey.id, budgetId }).unwrap(); + }; const { prev: prevKeys, next: nextKeys } = useSheetNavigation({ enabled: true, @@ -85,7 +103,7 @@ export default function VirtualKeyDetailSheet({ const isExhausted = // Budget exhausted (AP-mirrored when managed, VK-own otherwise) - displayBudgets?.some((b) => b.current_usage >= b.max_limit) || + displayBudgets?.some((b) => b.current_usage >= getEffectiveBudgetLimit(b)) || // Rate limits exhausted (displayRateLimit?.token_current_usage && displayRateLimit?.token_max_limit && @@ -290,7 +308,23 @@ export default function VirtualKeyDetailSheet({

Provider Budgets

{config.budgets.map((b, bIdx) => (
- + {!isManagedByProfile && b.id ? ( +
+ saveBudgetOverride(b.id, data)} + onRemove={() => clearBudgetOverride(b.id)} + disabled={!canUpdateVirtualKeys} + calendarAligned={virtualKey.calendar_aligned} + /> +
+ ) : null} + + {hasActiveBudgetOverride(b) ? ( +

+ Base {formatCurrency(b.max_limit)} + {formatCurrency(b.override_amount ?? 0)} override +

+ ) : null}
Resets {parseResetPeriod(b.reset_duration)} @@ -446,7 +480,24 @@ export default function VirtualKeyDetailSheet({
{displayBudgets.map((b, bIdx) => (
- + {!isManagedByProfile && b.id ? ( +
+ saveBudgetOverride(b.id, data)} + onRemove={() => clearBudgetOverride(b.id)} + disabled={!canUpdateVirtualKeys} + calendarAligned={virtualKey.calendar_aligned} + /> +
+ ) : null} + + {hasActiveBudgetOverride(b) ? ( +

+ Base {formatCurrency(b.max_limit)} + {formatCurrency(b.override_amount ?? 0)} override + {b.override_mode === "cycles" ? ` · ${b.override_cycles_remaining} cycles remaining` : " · until removed"} +

+ ) : null}
Resets {parseResetPeriod(b.reset_duration)} diff --git a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx index 3bb3f3fadb8..ff02acd4cae 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx @@ -1,4 +1,5 @@ import { useVirtualKeyUsage } from "@/app/workspace/virtual-keys/hooks/useVirtualKeyUsage"; +import { BudgetOverrideDialog } from "@/components/budgetOverrideDialog"; import { Accordion, AccordionContent, AccordionItem, AccordionTrigger } from "@/components/ui/accordion"; import { Alert, AlertDescription } from "@/components/ui/alert"; import { @@ -43,10 +44,20 @@ import { useGetMCPClientsQuery, useGetProvidersQuery, useRotateVirtualKeyMutation, + useRemoveVirtualKeyBudgetOverrideMutation, + useSetVirtualKeyBudgetOverrideMutation, useUpdateVirtualKeyMutation, } from "@/lib/store"; import { KnownProvider } from "@/lib/types/config"; -import { CreateVirtualKeyRequest, Customer, Team, UpdateVirtualKeyRequest, VirtualKey } from "@/lib/types/governance"; +import { + BudgetOverrideRequest, + CreateVirtualKeyRequest, + Customer, + Team, + UpdateVirtualKeyRequest, + VirtualKey, +} from "@/lib/types/governance"; +import { formatCurrency, getEffectiveBudgetLimit, hasActiveBudgetOverride, parseResetPeriod } from "@/lib/utils/governance"; import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { zodResolver } from "@hookform/resolvers/zod"; import { useNavigate } from "@tanstack/react-router"; @@ -271,9 +282,25 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT const [createVirtualKey, { isLoading: isCreating }] = useCreateVirtualKeyMutation(); const [updateVirtualKey, { isLoading: isUpdating }] = useUpdateVirtualKeyMutation(); const [rotateVirtualKey, { isLoading: isRotating }] = useRotateVirtualKeyMutation(); + const [setBudgetOverride] = useSetVirtualKeyBudgetOverrideMutation(); + const [removeBudgetOverride] = useRemoveVirtualKeyBudgetOverrideMutation(); const { data: mcpClientsResponse, error: mcpClientsError } = useGetMCPClientsQuery(); const mcpClientsData = mcpClientsResponse?.clients || []; const isLoading = isCreating || isUpdating || isRotating; + const persistedOverrideBudgets = [ + ...(virtualKey?.budgets ?? []).map((budget) => ({ budget, label: "Virtual key" })), + ...(virtualKey?.provider_configs ?? []).flatMap((config) => + (config.budgets ?? []).map((budget) => ({ budget, label: ProviderLabels[config.provider as ProviderName] ?? config.provider })), + ), + ]; + const saveBudgetOverride = async (budgetId: string, data: BudgetOverrideRequest) => { + if (!virtualKey) throw new Error("Virtual key is required"); + await setBudgetOverride({ vkId: virtualKey.id, budgetId, data }).unwrap(); + }; + const clearBudgetOverride = async (budgetId: string) => { + if (!virtualKey) throw new Error("Virtual key is required"); + await removeBudgetOverride({ vkId: virtualKey.id, budgetId }).unwrap(); + }; const availableKeys = keysData || []; const availableProviders = providersData || []; @@ -1665,6 +1692,39 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT showReset={isEditing && !!(virtualKey?.budgets?.length || (watchedBudgets && watchedBudgets.length > 0))} /> + {isEditing && !isManagedByProfile && persistedOverrideBudgets.length > 0 ? ( +
+
+

Budget Overrides

+

+ Add temporary capacity without changing the configured base budgets above. +

+
+
+ {persistedOverrideBudgets.map(({ budget, label }) => ( +
+
+

+ {label} · resets every {parseResetPeriod(budget.reset_duration)} +

+

+ Base {formatCurrency(budget.max_limit)} + {hasActiveBudgetOverride(budget) ? ` · effective ${formatCurrency(getEffectiveBudgetLimit(budget))}` : ""} +

+
+ saveBudgetOverride(budget.id, data)} + onRemove={() => clearBudgetOverride(budget.id)} + disabled={!hasUpdateAccess} + calendarAligned={virtualKey.calendar_aligned} + /> +
+ ))} +
+
+ ) : null} + {/* Reassign team confirmation dialog */} resetDurationLabels[duration] type ExportScope = "current_page" | "all"; -function virtualKeysToCSV(vks: VirtualKey[], accessProfileNames: Record = {}): string { +function virtualKeysToCSV(vks: VirtualKey[]): string { const headers = ["Name", "Status", "Assigned To", "Budget Limit", "Budget Spent", "Budget Reset", "Description", "Created At"]; const rows = vks.map((vk) => { const isExhausted = - vk.budgets?.some((b) => b.current_usage >= b.max_limit) || + vk.budgets?.some((b) => b.current_usage >= getEffectiveBudgetLimit(b)) || (vk.rate_limit?.token_current_usage && vk.rate_limit?.token_max_limit && vk.rate_limit.token_current_usage >= vk.rate_limit.token_max_limit) || @@ -83,7 +83,7 @@ function virtualKeysToCSV(vks: VirtualKey[], accessProfileNames: Record= new Date(vk.expires_at).getTime(); const status = !vk.is_active ? "Inactive" : isExpired ? "Expired" : isExhausted ? "Exhausted" : "Active"; const assignedTo = vk.team ? `Team: ${vk.team.name}` : vk.customer ? `Customer: ${vk.customer.name}` : ""; - const budgetLimit = vk.budgets?.length ? vk.budgets.map((b) => formatCurrency(b.max_limit)).join("; ") : ""; + const budgetLimit = vk.budgets?.length ? vk.budgets.map((b) => formatCurrency(getEffectiveBudgetLimit(b))).join("; ") : ""; const budgetSpent = vk.budgets?.length ? vk.budgets.map((b) => formatCurrency(b.current_usage)).join("; ") : ""; const budgetReset = vk.budgets?.length ? vk.budgets.map((b) => formatResetDuration(b.reset_duration)).join("; ") : ""; return [vk.name, status, assignedTo, budgetLimit, budgetSpent, budgetReset, vk.description || "", vk.created_at]; @@ -971,4 +971,4 @@ export default function VirtualKeysTable({
); -} \ No newline at end of file +} diff --git a/ui/components/budgetDisplay.tsx b/ui/components/budgetDisplay.tsx index 7b2233fba8f..e32c9d794cc 100644 --- a/ui/components/budgetDisplay.tsx +++ b/ui/components/budgetDisplay.tsx @@ -3,7 +3,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip import { resetDurationLabels, supportsCalendarAlignment } from "@/lib/constants/governance"; import { Budget } from "@/lib/types/governance"; import { cn } from "@/lib/utils"; -import { formatCurrency } from "@/lib/utils/governance"; +import { formatCurrency, getEffectiveBudgetLimit, hasActiveBudgetOverride } from "@/lib/utils/governance"; interface BudgetDisplayProps { budgets: Budget[] | null | undefined; @@ -30,8 +30,10 @@ export function BudgetDisplay({ budgets, calendarAligned }: BudgetDisplayProps) return (
{budgets.map((b, idx) => { - const pct = b.max_limit > 0 ? Math.min((b.current_usage / b.max_limit) * 100, 100) : 0; - const isExhausted = b.max_limit > 0 && b.current_usage >= b.max_limit; + const effectiveMaxLimit = getEffectiveBudgetLimit(b); + const hasOverride = hasActiveBudgetOverride(b); + const pct = effectiveMaxLimit > 0 ? Math.min((b.current_usage / effectiveMaxLimit) * 100, 100) : 0; + const isExhausted = effectiveMaxLimit > 0 && b.current_usage >= effectiveMaxLimit; const barClass = isExhausted ? "[&>div]:bg-red-500/70" : pct > 80 ? "[&>div]:bg-amber-500/70" : "[&>div]:bg-emerald-500/70"; return ( @@ -39,7 +41,10 @@ export function BudgetDisplay({ budgets, calendarAligned }: BudgetDisplayProps)
- {formatCurrency(b.max_limit)} + + {formatCurrency(effectiveMaxLimit)} + {hasOverride ? override : null} + {formatResetDuration(b.reset_duration, calendarAligned)}
@@ -47,8 +52,13 @@ export function BudgetDisplay({ budgets, calendarAligned }: BudgetDisplayProps)

- {formatCurrency(b.current_usage)} / {formatCurrency(b.max_limit)} + {formatCurrency(b.current_usage)} / {formatCurrency(effectiveMaxLimit)}

+ {hasOverride ? ( +

+ Base {formatCurrency(b.max_limit)} + {formatCurrency(b.override_amount ?? 0)} override +

+ ) : null} {b.reset_duration ? (

Resets {formatResetDuration(b.reset_duration, calendarAligned)}

) : null} diff --git a/ui/components/budgetOverrideDialog.tsx b/ui/components/budgetOverrideDialog.tsx new file mode 100644 index 00000000000..a738baea553 --- /dev/null +++ b/ui/components/budgetOverrideDialog.tsx @@ -0,0 +1,203 @@ +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogDescription, DialogFooter, DialogHeader, DialogTitle, DialogTrigger } from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { getErrorMessage } from "@/lib/store/apis/baseApi"; +import { Budget, BudgetOverrideRequest } from "@/lib/types/governance"; +import { budgetOverrideFormSchema } from "@/lib/types/schemas"; +import { formatCurrency, getBudgetOverrideValidUntil, getEffectiveBudgetLimit, hasActiveBudgetOverride } from "@/lib/utils/governance"; +import { Pencil, Plus } from "lucide-react"; +import { FormEvent, useEffect, useState } from "react"; +import { toast } from "sonner"; + +interface BudgetOverrideDialogProps { + budget: Budget; + onSave: (request: BudgetOverrideRequest) => Promise; + onRemove: () => Promise; + disabled?: boolean; + calendarAligned?: boolean; +} + +/** Lets an operator add, replace, or remove the additive override on one persisted budget. */ +export function BudgetOverrideDialog({ budget, onSave, onRemove, disabled, calendarAligned }: BudgetOverrideDialogProps) { + const active = hasActiveBudgetOverride(budget); + const [open, setOpen] = useState(false); + const [amount, setAmount] = useState(""); + const [mode, setMode] = useState<"cycles" | "forever">("cycles"); + const [cycles, setCycles] = useState("1"); + const [isSaving, setIsSaving] = useState(false); + const [error, setError] = useState(null); + const validUntil = mode === "cycles" ? getBudgetOverrideValidUntil(budget, Number(cycles), calendarAligned) : null; + + useEffect(() => { + if (!open) return; + setAmount(active ? String(budget.override_amount) : ""); + setMode(active && budget.override_mode ? budget.override_mode : "cycles"); + setCycles(active && budget.override_mode === "cycles" ? String(budget.override_cycles_remaining ?? 1) : "1"); + setError(null); + }, [active, budget.override_amount, budget.override_cycles_remaining, budget.override_mode, open]); + + const handleSubmit = async (event: FormEvent) => { + event.preventDefault(); + const parsedAmount = Number(amount); + const parsedCycles = Number(cycles); + const parsed = budgetOverrideFormSchema.safeParse({ + amount: parsedAmount, + mode, + ...(mode === "cycles" ? { cycles: parsedCycles } : {}), + }); + if (!parsed.success) { + setError(parsed.error.issues[0]?.message ?? "Invalid input"); + return; + } + + setIsSaving(true); + setError(null); + try { + await onSave({ amount: parsedAmount, mode, ...(mode === "cycles" ? { cycles: parsedCycles } : {}) }); + toast.success(active ? "Budget override updated" : "Budget override added"); + setOpen(false); + } catch (mutationError) { + setError(getErrorMessage(mutationError)); + } finally { + setIsSaving(false); + } + }; + + const handleRemove = async () => { + setIsSaving(true); + setError(null); + try { + await onRemove(); + toast.success("Budget override removed"); + setOpen(false); + } catch (mutationError) { + setError(getErrorMessage(mutationError)); + } finally { + setIsSaving(false); + } + }; + + return ( + + + + + +
+ + {active ? "Edit budget override" : "Add budget override"} + + Temporarily add spending capacity without changing the base {formatCurrency(budget.max_limit)} budget. + + + +
+
+ +
+ + setAmount(event.target.value)} + placeholder="0.00" + className="rounded-sm pl-7" + disabled={isSaving} + data-testid="budget-override-amount" + /> +
+
+ +
+ + +
+ + {mode === "cycles" ? ( +
+ + setCycles(event.target.value)} + className="rounded-sm" + disabled={isSaving} + data-testid="budget-override-cycles" + /> +

The current reset cycle counts as the first cycle.

+ {validUntil ? ( +

+ Valid until {validUntil.toLocaleString()} +

+ ) : null} +
+ ) : null} + + {active ? ( +
+ Current effective limit: {formatCurrency(getEffectiveBudgetLimit(budget))} +
+ ) : null} + + {error ? ( +

+ {error} +

+ ) : null} +
+ + + {active ? ( + + ) : ( + + )} + + +
+ + + ); +} \ No newline at end of file diff --git a/ui/lib/store/apis/governanceApi.ts b/ui/lib/store/apis/governanceApi.ts index 8488c90360c..b9905c9c278 100644 --- a/ui/lib/store/apis/governanceApi.ts +++ b/ui/lib/store/apis/governanceApi.ts @@ -1,5 +1,7 @@ import { Budget, + BudgetOverrideRequest, + BudgetOverrideResponse, BulkRotateVirtualKeysRequest, BulkRotateVirtualKeysResponse, CreateCustomerRequest, @@ -130,6 +132,23 @@ export const governanceApi = baseApi.injectEndpoints({ invalidatesTags: ["VirtualKeys", "ModelConfigs"], }), + setVirtualKeyBudgetOverride: builder.mutation({ + query: ({ vkId, budgetId, data }) => ({ + url: `/governance/virtual-keys/${encodeURIComponent(vkId)}/budgets/${encodeURIComponent(budgetId)}/override`, + method: "PUT", + body: data, + }), + invalidatesTags: ["VirtualKeys", "Budgets", "ModelConfigs"], + }), + + removeVirtualKeyBudgetOverride: builder.mutation({ + query: ({ vkId, budgetId }) => ({ + url: `/governance/virtual-keys/${encodeURIComponent(vkId)}/budgets/${encodeURIComponent(budgetId)}/override`, + method: "DELETE", + }), + invalidatesTags: ["VirtualKeys", "Budgets", "ModelConfigs"], + }), + // Teams getTeams: builder.query({ query: (params) => ({ @@ -866,6 +885,8 @@ export const { useRotateVirtualKeyMutation, useBulkRotateVirtualKeysMutation, useDeleteVirtualKeyMutation, + useSetVirtualKeyBudgetOverrideMutation, + useRemoveVirtualKeyBudgetOverrideMutation, // Teams useGetTeamsQuery, diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index c3534f2664c..0afb5acba84 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -8,6 +8,22 @@ export interface Budget { reset_duration: string; // e.g., "30s", "5m", "1h", "1d", "1w", "1M" current_usage: number; // In dollars last_reset: string; // ISO timestamp + override_amount?: number; + override_mode?: BudgetOverrideMode; + override_cycles_remaining?: number; +} + +export type BudgetOverrideMode = "cycles" | "forever"; + +export interface BudgetOverrideRequest { + amount: number; + mode: BudgetOverrideMode; + cycles?: number; +} + +export interface BudgetOverrideResponse { + budget: Budget; + effective_max_limit: number; } export interface RateLimit { diff --git a/ui/lib/types/schemas.ts b/ui/lib/types/schemas.ts index c41872b4125..083bd432205 100644 --- a/ui/lib/types/schemas.ts +++ b/ui/lib/types/schemas.ts @@ -1262,6 +1262,18 @@ export const routingRuleSchema = z path: ["scope_id"], }); +// Budget override form schema (BudgetOverrideDialog) +export const budgetOverrideFormSchema = z + .object({ + amount: z.number("Additional budget must be greater than 0.").positive("Additional budget must be greater than 0."), + mode: z.enum(["cycles", "forever"]), + cycles: z.number().optional(), + }) + .refine((data) => data.mode !== "cycles" || (data.cycles !== undefined && Number.isSafeInteger(data.cycles) && data.cycles > 0), { + message: "Reset cycles must be a positive whole number.", + path: ["cycles"], + }); + // Export type inference helpers export type SecretVar = z.infer; export type MCPClientUpdateSchema = z.infer; @@ -1284,4 +1296,5 @@ export type GlobalProxyConfigSchema = z.infer; export type GlobalProxyFormSchema = z.infer; export type GlobalHeaderFilterConfigSchema = z.infer; export type GlobalHeaderFilterFormSchema = z.infer; -export type RoutingRuleSchema = z.infer; \ No newline at end of file +export type RoutingRuleSchema = z.infer; +export type BudgetOverrideFormSchema = z.infer; \ No newline at end of file diff --git a/ui/lib/utils/governance.test.ts b/ui/lib/utils/governance.test.ts new file mode 100644 index 00000000000..f8ae1c919fb --- /dev/null +++ b/ui/lib/utils/governance.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from "vitest"; +import { getBudgetOverrideValidUntil, getEffectiveBudgetLimit, hasActiveBudgetOverride, validateBudgetOverride } from "./governance"; + +describe("budget overrides", () => { + it("adds active finite and permanent overrides to the base limit", () => { + expect(getEffectiveBudgetLimit({ max_limit: 100, override_amount: 25, override_mode: "cycles", override_cycles_remaining: 2 })).toBe( + 125, + ); + expect(getEffectiveBudgetLimit({ max_limit: 100, override_amount: 50, override_mode: "forever" })).toBe(150); + }); + + it("ignores incomplete or expired override state", () => { + expect(hasActiveBudgetOverride({ max_limit: 100, override_amount: 25, override_mode: "cycles", override_cycles_remaining: 0 })).toBe( + false, + ); + expect(getEffectiveBudgetLimit({ max_limit: 100, override_amount: 25, override_mode: "cycles", override_cycles_remaining: 0 })).toBe( + 100, + ); + }); + + it("validates positive amounts and whole finite cycle counts", () => { + expect(validateBudgetOverride(0, "forever", 0)).toMatch(/greater than 0/); + expect(validateBudgetOverride(25, "cycles", 1.5)).toMatch(/whole number/); + expect(validateBudgetOverride(25, "cycles", 1)).toBeNull(); + expect(validateBudgetOverride(25, "forever", 0)).toBeNull(); + }); + + it("calculates the validity date from the current reset schedule", () => { + expect( + getBudgetOverrideValidUntil({ max_limit: 100, last_reset: "2026-07-01T00:00:00.000Z", reset_duration: "1d" }, 4)?.toISOString(), + ).toBe("2026-07-05T00:00:00.000Z"); + expect( + getBudgetOverrideValidUntil({ max_limit: 100, last_reset: "2026-01-01T00:00:00.000Z", reset_duration: "1M" }, 2, true)?.toISOString(), + ).toBe("2026-03-01T00:00:00.000Z"); + }); +}); \ No newline at end of file diff --git a/ui/lib/utils/governance.ts b/ui/lib/utils/governance.ts index 578aba4c282..997fdd2b5a1 100644 --- a/ui/lib/utils/governance.ts +++ b/ui/lib/utils/governance.ts @@ -24,12 +24,97 @@ export function parseResetPeriod(duration: string): string { return `${timeValue} ${unitName}`; } +import { budgetOverrideFormSchema } from "@/lib/types/schemas"; + import { formatCompactNumber } from "./numbers"; export function formatCurrency(dollars: number) { return `$${dollars.toFixed(2)}`; } +export interface BudgetOverrideFields { + max_limit: number; + override_amount?: number; + override_mode?: "cycles" | "forever"; + override_cycles_remaining?: number; +} + +/** Returns whether a budget has a complete, currently active override. */ +export function hasActiveBudgetOverride(budget: BudgetOverrideFields): boolean { + if (!budget.override_amount || budget.override_amount <= 0) return false; + return budget.override_mode === "forever" || (budget.override_mode === "cycles" && (budget.override_cycles_remaining ?? 0) > 0); +} + +/** Returns the base limit plus the active additive override. */ +export function getEffectiveBudgetLimit(budget: BudgetOverrideFields): number { + return budget.max_limit + (hasActiveBudgetOverride(budget) ? (budget.override_amount ?? 0) : 0); +} + +/** Validates the operator-entered override fields before sending them to the API. Delegates to budgetOverrideFormSchema. */ +export function validateBudgetOverride(amount: number, mode: "cycles" | "forever", cycles: number): string | null { + const result = budgetOverrideFormSchema.safeParse({ amount, mode, ...(mode === "cycles" ? { cycles } : {}) }); + return result.success ? null : (result.error.issues[0]?.message ?? "Invalid input"); +} + +/** + * Adds months in UTC, clamping to the target month's last day instead of letting + * JS Date overflow spill into the following month (e.g. Jan 31 + 1M = Feb 28/29). + */ +function addUTCMonthsClamped(date: Date, months: number): void { + const day = date.getUTCDate(); + date.setUTCDate(1); + date.setUTCMonth(date.getUTCMonth() + months); + const lastDay = new Date(Date.UTC(date.getUTCFullYear(), date.getUTCMonth() + 1, 0)).getUTCDate(); + date.setUTCDate(Math.min(day, lastDay)); +} + +/** Calculates when a cycle-based override will expire on the budget's current reset schedule. */ +export function getBudgetOverrideValidUntil( + budget: Pick & { last_reset: string; reset_duration: string }, + cycles: number, + calendarAligned = false, +): Date | null { + if (!Number.isSafeInteger(cycles) || cycles <= 0) return null; + const match = /^(\d+)([smhdwMyY])$/.exec(budget.reset_duration); + const validUntil = new Date(budget.last_reset); + if (!match || Number.isNaN(validUntil.getTime())) return null; + + const durationValue = Number(match[1]) * cycles; + switch (match[2]) { + case "s": + validUntil.setTime(validUntil.getTime() + durationValue * 1000); + break; + case "m": + validUntil.setTime(validUntil.getTime() + durationValue * 60 * 1000); + break; + case "h": + validUntil.setTime(validUntil.getTime() + durationValue * 60 * 60 * 1000); + break; + case "d": + validUntil.setUTCDate(validUntil.getUTCDate() + durationValue); + break; + case "w": + validUntil.setUTCDate(validUntil.getUTCDate() + durationValue * 7); + break; + case "M": + if (calendarAligned) { + addUTCMonthsClamped(validUntil, durationValue); + } else { + validUntil.setTime(validUntil.getTime() + durationValue * 30 * 24 * 60 * 60 * 1000); + } + break; + case "y": + case "Y": + if (calendarAligned) { + addUTCMonthsClamped(validUntil, durationValue * 12); + } else { + validUntil.setTime(validUntil.getTime() + durationValue * 365 * 24 * 60 * 60 * 1000); + } + break; + } + return Number.isNaN(validUntil.getTime()) ? null : validUntil; +} + const shortDurationLabels: Record = { "1m": "/min", "5m": "/5min",