feat: bringing provider governance to budget & limits

framework/configstore/migrations.go

@@ -831,6 +831,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
 	if err := migrationAddModelConfigScopeColumns(ctx, db); err != nil {
 		return err
 	}
+	if err := migrationMigrateProviderGovernanceToModelConfigs(ctx, db); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -3893,6 +3896,102 @@ func migrationAddModelConfigScopeColumns(ctx context.Context, db *gorm.DB) error
 	return nil
 }
 
+// migrationMigrateProviderGovernanceToModelConfigs folds provider-level governance
+// (config_providers.budget_id / rate_limit_id) into governance_model_configs as
+// (scope='global', provider=<name>, model_name='*') "all models on this provider" rows,
+// reusing the same budget/rate-limit rows. It then NULLs the provider FKs so the old
+// provider-governance enforcement path goes inert (single source of truth = model_configs).
+func migrationMigrateProviderGovernanceToModelConfigs(ctx context.Context, db *gorm.DB) error {
+	m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
+		ID: "migrate_provider_governance_to_model_configs",
+		Migrate: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
+
+			// Guard: only run once the model-config table + scope columns exist.
+			if !tx.Migrator().HasTable(&tables.TableModelConfig{}) || !tx.Migrator().HasColumn(&tables.TableModelConfig{}, "scope") {
+				return nil
+			}
+
+			var providers []tables.TableProvider
+			if err := tx.Where("budget_id IS NOT NULL OR rate_limit_id IS NOT NULL").Find(&providers).Error; err != nil {
+				return fmt.Errorf("failed to load providers with governance: %w", err)
+			}
+
+			now := time.Now()
+			for i := range providers {
+				p := &providers[i]
+
+				// Idempotency: skip if a global all-models row already exists for this provider.
+				var existing int64
+				if err := tx.Model(&tables.TableModelConfig{}).
+					Where("scope = ? AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, p.Name).
+					Count(&existing).Error; err != nil {
+					return fmt.Errorf("failed to check existing wildcard config for provider %q: %w", p.Name, err)
+				}
+				if existing == 0 {
+					providerName := p.Name
+					mc := tables.TableModelConfig{
+						ID:          uuid.NewString(),
+						ModelName:   tables.ModelConfigAllModels,
+						Provider:    &providerName,
+						Scope:       tables.ModelConfigScopeGlobal,
+						BudgetID:    p.BudgetID,
+						RateLimitID: p.RateLimitID,
+						CreatedAt:   now,
+						UpdatedAt:   now,
+					}
+					if err := tx.Create(&mc).Error; err != nil {
+						return fmt.Errorf("failed to create wildcard model config for provider %q: %w", p.Name, err)
+					}
+				}
+
+				// Detach governance from the provider (FK rows are reused by the model config above).
+				if err := tx.Model(&tables.TableProvider{}).Where("name = ?", p.Name).
+					Updates(map[string]any{"budget_id": nil, "rate_limit_id": nil}).Error; err != nil {
+					return fmt.Errorf("failed to clear governance FKs for provider %q: %w", p.Name, err)
+				}
+			}
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
+
+			// Nothing to reverse if the model-config table/scope columns are gone.
+			if !tx.Migrator().HasTable(&tables.TableModelConfig{}) || !tx.Migrator().HasColumn(&tables.TableModelConfig{}, "scope") {
+				return nil
+			}
+
+			// Reverse provider-level wildcards:
+			// (scope='global', scope_id IS NULL, model_name='*', provider IS NOT NULL).
+			var wildcards []tables.TableModelConfig
+			if err := tx.Where(
+				"scope = ? AND scope_id IS NULL AND model_name = ? AND provider IS NOT NULL",
+				tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels,
+			).Find(&wildcards).Error; err != nil {
+				return fmt.Errorf("failed to load provider wildcard configs: %w", err)
+			}
+
+			for i := range wildcards {
+				mc := &wildcards[i]
+				// Re-attach the budget/rate-limit FK rows to the provider row.
+				if err := tx.Model(&tables.TableProvider{}).Where("name = ?", *mc.Provider).
+					Updates(map[string]any{"budget_id": mc.BudgetID, "rate_limit_id": mc.RateLimitID}).Error; err != nil {
+					return fmt.Errorf("failed to restore governance FKs for provider %q: %w", *mc.Provider, err)
+				}
+				// Drop the wildcard model config; its FK rows now live on the provider again.
+				if err := tx.Delete(&tables.TableModelConfig{}, "id = ?", mc.ID).Error; err != nil {
+					return fmt.Errorf("failed to delete wildcard config for provider %q: %w", *mc.Provider, err)
+				}
+			}
+			return nil
+		},
+	}})
+	if err := m.Migrate(); err != nil {
+		return fmt.Errorf("error while running migrate provider governance to model configs migration: %s", err.Error())
+	}
+	return nil
+}
+
 // migrationAddAllowedHeadersJSONColumn adds the allowed_headers_json column to the client config table
 func migrationAddAllowedHeadersJSONColumn(ctx context.Context, db *gorm.DB) error {
 	m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{

framework/configstore/migrations_test.go

@@ -2442,3 +2442,43 @@ func TestMigrationAddModelConfigScopeColumns(t *testing.T) {
 	assert.True(t, db.Migrator().HasIndex(mc, "idx_model_scope_provider"))
 	assert.False(t, db.Migrator().HasIndex(mc, "idx_model_provider"))
 }
+
+// TestMigrationMigrateProviderGovernanceToModelConfigs verifies provider-level governance is
+// folded into a (global, provider, '*') model_config reusing the same budget/rate-limit rows,
+// and the provider FKs are cleared. Idempotent on re-run.
+func TestMigrationMigrateProviderGovernanceToModelConfigs(t *testing.T) {
+	db := setupTestDB(t)
+	ctx := context.Background()
+	require.NoError(t, db.AutoMigrate(
+		&tables.TableProvider{}, &tables.TableModelConfig{}, &tables.TableBudget{}, &tables.TableRateLimit{},
+	))
+
+	now := time.Now()
+	require.NoError(t, db.Create(&tables.TableBudget{ID: "b1", MaxLimit: 100, ResetDuration: "1M", LastReset: now, CreatedAt: now, UpdatedAt: now}).Error)
+	require.NoError(t, db.Create(&tables.TableRateLimit{ID: "rl1", TokenMaxLimit: schemas.Ptr(int64(1000)), TokenResetDuration: schemas.Ptr("1h"), TokenLastReset: now, RequestLastReset: now, CreatedAt: now, UpdatedAt: now}).Error)
+	require.NoError(t, db.Create(&tables.TableProvider{Name: "openai", BudgetID: schemas.Ptr("b1"), RateLimitID: schemas.Ptr("rl1"), CreatedAt: now, UpdatedAt: now}).Error)
+
+	require.NoError(t, migrationMigrateProviderGovernanceToModelConfigs(ctx, db))
+
+	// A (global, openai, '*') model config now exists reusing the same budget/rate-limit IDs.
+	var mc tables.TableModelConfig
+	require.NoError(t, db.Where("scope = ? AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, "openai").First(&mc).Error)
+	require.NotNil(t, mc.BudgetID)
+	assert.Equal(t, "b1", *mc.BudgetID)
+	require.NotNil(t, mc.RateLimitID)
+	assert.Equal(t, "rl1", *mc.RateLimitID)
+
+	// Provider governance FKs are cleared (old path now inert).
+	var prov tables.TableProvider
+	require.NoError(t, db.Where("name = ?", "openai").First(&prov).Error)
+	assert.Nil(t, prov.BudgetID, "provider budget_id should be cleared")
+	assert.Nil(t, prov.RateLimitID, "provider rate_limit_id should be cleared")
+
+	// Idempotency: re-run creates no duplicate wildcard row.
+	require.NoError(t, migrationMigrateProviderGovernanceToModelConfigs(ctx, db))
+	var count int64
+	require.NoError(t, db.Model(&tables.TableModelConfig{}).
+		Where("scope = ? AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, "openai").
+		Count(&count).Error)
+	assert.Equal(t, int64(1), count, "re-run must not duplicate the wildcard config")
+}

framework/configstore/rdb.go

@@ -1155,6 +1155,33 @@ func (s *RDBConfigStore) DeleteProvider(ctx context.Context, provider schemas.Mo
 		}
 	}
 
+	// Clean up model configs scoped to this provider
+	// Delete by snapshotted IDs rather than a second WHERE provider=? pass to avoid a race
+	// where a concurrent CreateModelConfig lands between the snapshot and the delete, leaving
+	// its owned budget/rate-limit rows dangling.
+	var providerModelConfigs []tables.TableModelConfig
+	if err := txDB.WithContext(ctx).Where("provider = ?", string(provider)).Find(&providerModelConfigs).Error; err != nil {
+		return err
+	}
+	for _, mc := range providerModelConfigs {
+		if mc.BudgetID != nil {
+			if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", *mc.BudgetID).Error; err != nil {
+				return err
+			}
+		}
+		if mc.RateLimitID != nil {
+			if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *mc.RateLimitID).Error; err != nil {
+				return err
+			}
+		}
+	}
+	if len(providerModelConfigs) > 0 {
+		var mcIDs []string
+		if err := txDB.WithContext(ctx).Where("id IN ?", mcIDs).Delete(&tables.TableModelConfig{}).Error; err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -4223,6 +4250,18 @@ func (s *RDBConfigStore) GetModelConfigs(ctx context.Context) ([]tables.TableMod
 	return modelConfigs, nil
 }
 
+// GetProviderGovernanceModelConfigs retrieves the wildcard "all models on a provider" configs
+func (s *RDBConfigStore) GetProviderGovernanceModelConfigs(ctx context.Context) ([]tables.TableModelConfig, error) {
+	var modelConfigs []tables.TableModelConfig
+	if err := s.DB().WithContext(ctx).
+		Preload("Budget").Preload("RateLimit").
+		Where("scope = ? AND model_name = ? AND provider IS NOT NULL", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels).
+		Find(&modelConfigs).Error; err != nil {
+		return nil, err
+	}
+	return modelConfigs, nil
+}
+
 // GetModelConfigsPaginated retrieves model configs with pagination, filtering, and search support.
 func (s *RDBConfigStore) GetModelConfigsPaginated(ctx context.Context, params ModelConfigsQueryParams) ([]tables.TableModelConfig, int64, error) {
 	baseQuery := s.DB().WithContext(ctx).Model(&tables.TableModelConfig{})

framework/configstore/rdb_test.go

@@ -692,6 +692,47 @@ func TestDeleteVirtualKey_CleansUpScopedModelConfigs(t *testing.T) {
 	assert.Equal(t, int64(0), rlCount, "owned rate limit should be deleted")
 }
 
+func TestDeleteProvider_CleansUpProviderModelConfigs(t *testing.T) {
+	store := setupRDBTestStore(t)
+	ctx := context.Background()
+
+	now := time.Now()
+	require.NoError(t, store.DB().Create(&tables.TableProvider{Name: "openai", CreatedAt: now, UpdatedAt: now}).Error)
+	require.NoError(t, store.DB().Create(&tables.TableBudget{ID: "pb", MaxLimit: 100, ResetDuration: "1M", LastReset: now, CreatedAt: now, UpdatedAt: now}).Error)
+	require.NoError(t, store.DB().Create(&tables.TableRateLimit{ID: "prl", TokenMaxLimit: schemas.Ptr(int64(1000)), TokenResetDuration: schemas.Ptr("1h"), TokenLastReset: now, RequestLastReset: now, CreatedAt: now, UpdatedAt: now}).Error)
+
+	providerName := "openai"
+	mc := &tables.TableModelConfig{
+		ID:          "mc-wildcard",
+		ModelName:   tables.ModelConfigAllModels,
+		Provider:    &providerName,
+		Scope:       tables.ModelConfigScopeGlobal,
+		BudgetID:    schemas.Ptr("pb"),
+		RateLimitID: schemas.Ptr("prl"),
+		CreatedAt:   now,
+		UpdatedAt:   now,
+	}
+	require.NoError(t, store.CreateModelConfig(ctx, mc))
+
+	require.NoError(t, store.DeleteProvider(ctx, schemas.ModelProvider("openai")))
+
+	// The provider's wildcard model config and its owned budget/rate-limit are cleaned up.
+	_, err := store.GetModelConfigByID(ctx, "mc-wildcard")
+	assert.Error(t, err, "provider-scoped model config should be deleted with the provider")
+	for _, q := range []struct {
+		model any
+		id    string
+		label string
+	}{
+		{&tables.TableBudget{}, "pb", "budget"},
+		{&tables.TableRateLimit{}, "prl", "rate limit"},
+	} {
+		var count int64
+		require.NoError(t, store.DB().Model(q.model).Where("id = ?", q.id).Count(&count).Error)
+		assert.Equal(t, int64(0), count, "owned "+q.label+" should be deleted")
+	}
+}
+
 // =============================================================================
 // Virtual Key Provider Config Tests
 // =============================================================================

framework/configstore/store.go

@@ -256,6 +256,7 @@ type ConfigStore interface {
 
 	// Model config CRUD
 	GetModelConfigs(ctx context.Context) ([]tables.TableModelConfig, error)
+	GetProviderGovernanceModelConfigs(ctx context.Context) ([]tables.TableModelConfig, error)
 	GetModelConfigsPaginated(ctx context.Context, params ModelConfigsQueryParams) ([]tables.TableModelConfig, int64, error)
 	GetModelConfig(ctx context.Context, scope string, scopeID *string, modelName string, provider *string) (*tables.TableModelConfig, error)
 	GetModelConfigByID(ctx context.Context, id string) (*tables.TableModelConfig, error)

framework/configstore/tables/modelconfig.go

@@ -14,6 +14,11 @@ const (
 	ModelConfigScopeVirtualKey = "virtual_key"
 )
 
+// ModelConfigAllModels is the model_name sentinel meaning "all models". Combined with a
+// specific provider it expresses provider-level governance (all models on that provider);
+// with a nil provider it means all models on all providers.
+const ModelConfigAllModels = "*"
+
 // validModelConfigScopes is the set of accepted scope values.
 var validModelConfigScopes = map[string]bool{
 	ModelConfigScopeGlobal:     true,

plugins/governance/modelprovidergovernance_test.go

@@ -2101,6 +2101,75 @@ func TestStore_CheckModelBudget_NoCatalog_NoMatch(t *testing.T) {
 	assert.Error(t, err, "Direct match should still work without catalog")
 }
 
+// ============================================================================
+// Store Tests - All-models ("*") wildcard tier (provider-level governance)
+// ============================================================================
+
+// TestStore_CheckModelBudget_AllModelsOnProvider_Exceeded verifies that an all-models
+// wildcard config (provider=openai, model_name="*") — the migrated provider-level budget —
+// applies to ANY model on that provider.
+func TestStore_CheckModelBudget_AllModelsOnProvider_Exceeded(t *testing.T) {
+	logger := NewMockLogger()
+	budget := buildBudgetWithUsage("b1", 100.0, 100.0, "1h") // exceeded
+	providerStr := "openai"
+	mc := buildModelConfig("mc-provider", configstoreTables.ModelConfigAllModels, &providerStr, budget, nil)
+	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+		ModelConfigs: []configstoreTables.TableModelConfig{*mc},
+		Budgets:      []configstoreTables.TableBudget{*budget},
+	}, nil)
+	require.NoError(t, err)
+
+	// A request for an arbitrary OpenAI model must be caught by the "*:openai" config.
+	_, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: schemas.OpenAI}, nil)
+	assert.Error(t, err, "all-models budget for the provider should apply to any model on it")
+	assert.Contains(t, err.Error(), "budget exceeded")
+}
+
+// TestStore_CheckModelBudget_AllModelsOnProvider_OtherProviderPasses confirms the wildcard
+// is provider-scoped: it must NOT affect a different provider.
+func TestStore_CheckModelBudget_AllModelsOnProvider_OtherProviderPasses(t *testing.T) {
+	logger := NewMockLogger()
+	budget := buildBudgetWithUsage("b1", 100.0, 100.0, "1h") // exceeded
+	providerStr := "openai"
+	mc := buildModelConfig("mc-provider", configstoreTables.ModelConfigAllModels, &providerStr, budget, nil)
+	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+		ModelConfigs: []configstoreTables.TableModelConfig{*mc},
+		Budgets:      []configstoreTables.TableBudget{*budget},
+	}, nil)
+	require.NoError(t, err)
+
+	decision, err := store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "claude-opus-4-7", Provider: schemas.Anthropic}, nil)
+	assert.NoError(t, err, "an OpenAI all-models budget must not affect an Anthropic request")
+	assert.Equal(t, DecisionAllow, decision)
+}
+
+// TestStore_UpdateProviderModelUsage_BumpsAllModelsWildcard verifies usage recording reaches
+// the all-models wildcard config (record-then-check loop for provider-level governance).
+func TestStore_UpdateProviderModelUsage_BumpsAllModelsWildcard(t *testing.T) {
+	logger := NewMockLogger()
+	rateLimit := buildRateLimitWithUsage("rl1", 100, 0, 1000000, 0) // 100-token cap
+	providerStr := "openai"
+	mc := buildModelConfig("mc-provider", configstoreTables.ModelConfigAllModels, &providerStr, nil, rateLimit)
+	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+		ModelConfigs: []configstoreTables.TableModelConfig{*mc},
+		RateLimits:   []configstoreTables.TableRateLimit{*rateLimit},
+	}, nil)
+	require.NoError(t, err)
+
+	// Within limit initially.
+	decision, err := store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: schemas.OpenAI}, nil, nil)
+	require.NoError(t, err)
+	require.Equal(t, DecisionAllow, decision)
+
+	// Record usage for a (different) model on the provider — must bump the "*:openai" config.
+	require.NoError(t, store.UpdateProviderAndModelRateLimitUsageInMemory(context.Background(), "gpt-4o", schemas.OpenAI, 150, true, true))
+
+	// Now the all-models rate limit trips for any model on the provider.
+	decision, err = store.CheckModelRateLimit(context.Background(), &EvaluationRequest{Model: "gpt-4o-mini", Provider: schemas.OpenAI}, nil, nil)
+	assert.Error(t, err)
+	assert.Equal(t, DecisionTokenLimited, decision)
+}
+
 // ============================================================================
 // Store Tests - Per-VK-Scoped Model Budget / Rate Limit
 // ============================================================================