fix: evict stale VK-scoped model configs from in-memory store on ReloadVirtualKey

framework/configstore/rdb.go

@@ -3992,8 +3992,12 @@ func (s *RDBConfigStore) UpdateBudgetUsage(ctx context.Context, id string, curre
 
 // UpdateRateLimitUsage updates only the usage fields of a rate limit.
 // Uses SkipHooks to avoid triggering BeforeSave validation since we're only updating usage.
-func (s *RDBConfigStore) UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64) error {
-	result := s.DB().WithContext(ctx).
+func (s *RDBConfigStore) UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64, tx ...*gorm.DB) error {
+	db := s.DB()
+	if len(tx) > 0 && tx[0] != nil {
+		db = tx[0]
+	}
+	result := db.WithContext(ctx).
 		Session(&gorm.Session{SkipHooks: true}).
 		Model(&tables.TableRateLimit{}).
 		Where("id = ?", id).

framework/configstore/store.go

@@ -244,7 +244,7 @@ type ConfigStore interface {
 	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
-	UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64) error
+	UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64, tx ...*gorm.DB) error
 
 	// Routing Rules CRUD
 	GetRoutingRules(ctx context.Context) ([]tables.TableRoutingRule, error)

plugins/governance/store.go

@@ -176,6 +176,7 @@ type GovernanceStore interface {
 	// Model config in-memory operations
 	UpdateModelConfigInMemory(ctx context.Context, mc *configstoreTables.TableModelConfig) *configstoreTables.TableModelConfig
 	DeleteModelConfigInMemory(ctx context.Context, mcID string)
+	ScopedModelConfigIDs(scope, scopeID string) []string
 	// Provider in-memory operations
 	UpdateProviderInMemory(ctx context.Context, provider *configstoreTables.TableProvider) *configstoreTables.TableProvider
 	DeleteProviderInMemory(ctx context.Context, providerName string)
@@ -450,6 +451,35 @@ func (gs *LocalGovernanceStore) BumpRateLimitUsage(ctx context.Context, rateLimi
 	}
 }
 
+// BumpRateLimitUsageBy atomically adds arbitrary token and request deltas to the
+// rate limit identified by rateLimitID. Unlike BumpRateLimitUsage (which adds a
+// token count and a single request), this adds caller-supplied counts on both
+// dimensions — used to fold accumulated usage carried from another rate limit.
+// Same CAS-retry contract: no increment is dropped under concurrent callers.
+// No window-reset side effect, since carried deltas are not request traffic.
+// No-op when the rate limit is absent or both deltas are zero.
+func (gs *LocalGovernanceStore) BumpRateLimitUsageBy(ctx context.Context, rateLimitID string, tokenDelta, requestDelta int64) error {
+	if tokenDelta == 0 && requestDelta == 0 {
+		return nil
+	}
+	for {
+		raw, exists := gs.rateLimits.Load(rateLimitID)
+		if !exists || raw == nil {
+			return nil
+		}
+		old, ok := raw.(*configstoreTables.TableRateLimit)
+		if !ok || old == nil {
+			return nil
+		}
+		clone := *old
+		clone.TokenCurrentUsage += tokenDelta
+		clone.RequestCurrentUsage += requestDelta
+		if gs.rateLimits.CompareAndSwap(rateLimitID, raw, &clone) {
+			return nil
+		}
+	}
+}
+
 // ResetBudgetAt atomically zeros the budget's CurrentUsage and advances its
 // LastReset to newLastReset, provided the currently-stored budget has an
 // older LastReset. Returns the reset budget and true when the CAS succeeds;
@@ -3219,6 +3249,28 @@ func (gs *LocalGovernanceStore) DeleteModelConfigInMemory(ctx context.Context, m
 	})
 }
 
+// ScopedModelConfigIDs returns the IDs of all in-memory model configs for the
+// given (scope, scopeID). Callers use this to diff against the DB result and
+// evict stale entries via DeleteModelConfigInMemory.
+func (gs *LocalGovernanceStore) ScopedModelConfigIDs(scope, scopeID string) []string {
+	var ids []string
+	gs.modelConfigs.Range(func(key, value interface{}) bool {
+		mc, ok := value.(*configstoreTables.TableModelConfig)
+		if !ok || mc == nil {
+			return true
+		}
+		mcScopeID := ""
+		if mc.ScopeID != nil {
+			mcScopeID = *mc.ScopeID
+		}
+		if mc.Scope == scope && mcScopeID == scopeID {
+			ids = append(ids, mc.ID)
+		}
+		return true
+	})
+	return ids
+}
+
 // UpdateProviderInMemory adds or updates a provider in the in-memory store (lock-free)
 // Preserves existing usage values when updating budgets and rate limits
 // Returns the updated provider with potentially modified usage values

transports/bifrost-http/lib/config_test.go

@@ -1156,7 +1156,7 @@ func (m *MockConfigStore) UpdateBudgetUsage(ctx context.Context, id string, curr
 	return nil
 }
 
-func (m *MockConfigStore) UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64) error {
+func (m *MockConfigStore) UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64, tx ...*gorm.DB) error {
 	return nil
 }
 

transports/bifrost-http/server/server.go

@@ -380,6 +380,17 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t
 	if err != nil {
 		return nil, err
 	}
+	// Fetch VK-scoped model configs up front, alongside the VK load, so that a DB
+	// failure here aborts before we mutate any in-memory state. Reloading these
+	// reflects governance changes made via the VK sheet (syncVKGovernanceToModelConfigs)
+	// in memory immediately — both on the node that handled the update and on peers
+	// that receive this reload via the cluster gossip broadcast.
+	mcs, err := s.Config.ConfigStore.GetModelConfigsByScopeAndScopeIDs(
+		ctx, tables.ModelConfigScopeVirtualKey, []string{id},
+	)
+	if err != nil {
+		return virtualKey, fmt.Errorf("failed to reload VK-scoped model configs for VK %s: %w", id, err)
+	}
 	if governanceData := governancePlugin.GetGovernanceStore().GetGovernanceData(ctx); governanceData != nil {
 		for _, existingVK := range governanceData.VirtualKeys {
 			if existingVK != nil && existingVK.ID == virtualKey.ID && existingVK.Value != "" && existingVK.Value != virtualKey.Value {
@@ -388,17 +399,22 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t
 			}
 		}
 	}
-	governancePlugin.GetGovernanceStore().UpdateVirtualKeyInMemory(ctx, virtualKey, nil, nil, nil)
-	// Also reload any model configs scoped to this VK so governance changes made
-	// via the VK sheet (syncVKGovernanceToModelConfigs) are reflected in memory
-	// immediately — both on the node that handled the update and on peers that
-	// receive this reload via the cluster gossip broadcast.
-	if mcs, err := s.Config.ConfigStore.GetModelConfigsByScopeAndScopeIDs(
-		ctx, tables.ModelConfigScopeVirtualKey, []string{id},
-	); err == nil {
-		for i := range mcs {
-			governancePlugin.GetGovernanceStore().UpdateModelConfigInMemory(ctx, &mcs[i])
-		}
+	store := governancePlugin.GetGovernanceStore()
+	store.UpdateVirtualKeyInMemory(ctx, virtualKey, nil, nil, nil)
+	// Snapshot in-memory VK-scoped config IDs before the upserts so we can evict
+	// the ones that no longer exist in the DB (e.g. a standalone VK adopted into
+	// an access profile has its VK-scoped governance model configs deleted).
+	// Without this their stale budgets keep enforcing.
+	staleIDs := make(map[string]bool)
+	for _, mcID := range store.ScopedModelConfigIDs(tables.ModelConfigScopeVirtualKey, id) {
+		staleIDs[mcID] = true
+	}
+	for i := range mcs {
+		delete(staleIDs, mcs[i].ID)
+		store.UpdateModelConfigInMemory(ctx, &mcs[i])
+	}
+	for mcID := range staleIDs {
+		store.DeleteModelConfigInMemory(ctx, mcID)
 	}
 	s.MCPServerHandler.SyncVKMCPServer(virtualKey)
 	return virtualKey, nil