feat: adding scope to budget limits table

.gitignore

@@ -99,6 +99,7 @@ dist/
 # IDE
 .idea/
 .vscode/
+.zed/
 
 # OS
 .DS_Store

framework/configstore/clientconfig.go

@@ -1068,10 +1068,18 @@ func GenerateTeamHash(t tables.TableTeam) (string, error) {
 // This is used to detect changes to model configs between config.json and database.
 // Skips: CreatedAt, UpdatedAt, and relationship objects (dynamic fields)
 func GenerateModelConfigHash(m tables.TableModelConfig) (string, error) {
+	// Normalize an empty scope to "global" so a config.json entry that omits scope
+	// hashes identically to the defaulted DB row.
+	scope := m.Scope
+	if scope == "" {
+		scope = tables.ModelConfigScopeGlobal
+	}
 	hash := sha256.New()
 	writeHashField(hash, "id", m.ID)
 	writeHashField(hash, "model_name", m.ModelName)
 	writeHashField(hash, "provider", derefStr(m.Provider))
+	writeHashField(hash, "scope", scope)
+	writeHashField(hash, "scope_id", derefStr(m.ScopeID))
 	writeHashField(hash, "budget_id", derefStr(m.BudgetID))
 	writeHashField(hash, "rate_limit_id", derefStr(m.RateLimitID))
 	return hex.EncodeToString(hash.Sum(nil)), nil

framework/configstore/migrations.go

@@ -828,6 +828,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
 	if err := migrationAddAdditionalAttributesToPricing(ctx, db); err != nil {
 		return err
 	}
+	if err := migrationAddModelConfigScopeColumns(ctx, db); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -3832,6 +3835,64 @@ func migrationAddProviderGovernanceColumns(ctx context.Context, db *gorm.DB) err
 	return nil
 }
 
+// migrationAddModelConfigScopeColumns adds the scope and scope_id columns to
+// governance_model_configs and swaps the unique index from (model_name, provider)
+// to (scope, scope_id, model_name, provider). Existing rows are backfilled to the
+// "global" scope, preserving pre-scope behavior. The new index is created before
+// the old one is dropped so uniqueness is never unenforced during the migration.
+func migrationAddModelConfigScopeColumns(ctx context.Context, db *gorm.DB) error {
+	m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
+		ID: "add_model_config_scope_columns",
+		Migrate: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
+			migrator := tx.Migrator()
+			modelConfig := &tables.TableModelConfig{}
+
+			// Add scope column (NOT NULL DEFAULT 'global' backfills existing rows).
+			if !migrator.HasColumn(modelConfig, "scope") {
+				if err := migrator.AddColumn(modelConfig, "scope"); err != nil {
+					return fmt.Errorf("failed to add scope column: %w", err)
+				}
+			}
+			// Add scope_id column (nullable).
+			if !migrator.HasColumn(modelConfig, "scope_id") {
+				if err := migrator.AddColumn(modelConfig, "scope_id"); err != nil {
+					return fmt.Errorf("failed to add scope_id column: %w", err)
+				}
+			}
+			// Belt-and-suspenders backfill in case the column default did not populate
+			// existing rows on this dialect.
+			if err := tx.Exec("UPDATE governance_model_configs SET scope = ? WHERE scope IS NULL OR scope = ''", tables.ModelConfigScopeGlobal).Error; err != nil {
+				return fmt.Errorf("failed to backfill scope: %w", err)
+			}
+
+			// Create the new composite unique index BEFORE dropping the old one. The
+			// composite index is strictly more selective, so already-unique rows stay
+			// unique under it; this ordering avoids any window where uniqueness is
+			// unenforced. CreateIndex reads the struct tags so it is dialect-safe.
+			if !migrator.HasIndex(modelConfig, "idx_model_scope_provider") {
+				if err := migrator.CreateIndex(modelConfig, "idx_model_scope_provider"); err != nil {
+					return fmt.Errorf("failed to create idx_model_scope_provider: %w", err)
+				}
+			}
+			// Drop the now-superseded (model_name, provider) unique index.
+			if migrator.HasIndex(modelConfig, "idx_model_provider") {
+				if err := migrator.DropIndex(modelConfig, "idx_model_provider"); err != nil {
+					return fmt.Errorf("failed to drop idx_model_provider: %w", err)
+				}
+			}
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			return fmt.Errorf("add_model_config_scope_columns is non-rollbackable: scope-aware rows and the previous uniqueness invariant cannot be restored safely")
+		},
+	}})
+	if err := m.Migrate(); err != nil {
+		return fmt.Errorf("error while running add model config scope columns 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

@@ -2381,3 +2381,64 @@ func assertNoCorruptedFKReferences(t *testing.T, db *gorm.DB) {
 }
 
 func strPtr(s string) *string { return &s }
+
+// TestMigrationAddModelConfigScopeColumns verifies the existing-install transition:
+// adding scope/scope_id columns, backfilling existing rows to "global", and swapping the
+// (model_name, provider) unique index for the composite (scope, scope_id, model_name, provider) one.
+func TestMigrationAddModelConfigScopeColumns(t *testing.T) {
+	db := setupTestDB(t)
+	ctx := context.Background()
+
+	// Create the OLD governance_model_configs schema: no scope/scope_id columns,
+	// with a unique index on (model_name, provider).
+	require.NoError(t, db.Exec(`
+		CREATE TABLE governance_model_configs (
+			id varchar(255) PRIMARY KEY,
+			model_name varchar(255) NOT NULL,
+			provider varchar(50),
+			budget_id varchar(255),
+			rate_limit_id varchar(255),
+			config_hash varchar(255),
+			created_at datetime NOT NULL,
+			updated_at datetime NOT NULL
+		)
+	`).Error)
+	require.NoError(t, db.Exec(`CREATE UNIQUE INDEX idx_model_provider ON governance_model_configs (model_name, provider)`).Error)
+
+	now := time.Now()
+	require.NoError(t, db.Exec(`
+		INSERT INTO governance_model_configs (id, model_name, created_at, updated_at)
+		VALUES (?, ?, ?, ?)
+	`, "mc-existing", "gpt-4", now, now).Error)
+
+	mc := &tables.TableModelConfig{}
+
+	// Pre-migration state.
+	assert.False(t, db.Migrator().HasColumn(mc, "scope"), "scope column should not exist yet")
+	assert.False(t, db.Migrator().HasColumn(mc, "scope_id"), "scope_id column should not exist yet")
+	assert.True(t, db.Migrator().HasIndex(mc, "idx_model_provider"), "old index should exist before migration")
+
+	require.NoError(t, migrationAddModelConfigScopeColumns(ctx, db))
+
+	// Post-migration state.
+	assert.True(t, db.Migrator().HasColumn(mc, "scope"), "scope column should exist after migration")
+	assert.True(t, db.Migrator().HasColumn(mc, "scope_id"), "scope_id column should exist after migration")
+	assert.True(t, db.Migrator().HasIndex(mc, "idx_model_scope_provider"), "new composite index should exist after migration")
+	assert.False(t, db.Migrator().HasIndex(mc, "idx_model_provider"), "old index should be dropped after migration")
+
+	// Existing row should be backfilled to the global scope.
+	var scope string
+	require.NoError(t, db.Table("governance_model_configs").Select("scope").Where("id = ?", "mc-existing").Scan(&scope).Error)
+	assert.Equal(t, tables.ModelConfigScopeGlobal, scope, "existing row should be backfilled to the global scope")
+
+	// Global-scope rows must have NULL scope_id for the composite unique index to work correctly.
+	var scopeID *string
+	require.NoError(t, db.Table("governance_model_configs").Select("scope_id").Where("id = ?", "mc-existing").Scan(&scopeID).Error)
+	assert.Nil(t, scopeID, "global scope rows must have NULL scope_id")
+
+	// Idempotency: running again must be a no-op (no error, state unchanged).
+	require.NoError(t, migrationAddModelConfigScopeColumns(ctx, db))
+	assert.True(t, db.Migrator().HasColumn(mc, "scope"))
+	assert.True(t, db.Migrator().HasIndex(mc, "idx_model_scope_provider"))
+	assert.False(t, db.Migrator().HasIndex(mc, "idx_model_provider"))
+}

framework/configstore/rdb.go

@@ -2956,6 +2956,32 @@ func (s *RDBConfigStore) DeleteVirtualKey(ctx context.Context, id string, tx ...
 		if err := txDB.WithContext(ctx).Where("virtual_key_id = ?", id).Delete(&tables.TableBudget{}).Error; err != nil {
 			return err
 		}
+		// Delete model configs scoped to this virtual key, along with their owned
+		// budgets/rate-limits. scope_id has no FK constraint, so this cleanup must be
+		// explicit; otherwise per-VK model limits would orphan and leak budget/rate-limit rows.
+		// Model configs are deleted first (matching DeleteModelConfig order) before their
+		// owned budget/rate-limit rows.
+		var scopedModelConfigs []tables.TableModelConfig
+		if err := txDB.WithContext(ctx).
+			Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id).
+			Find(&scopedModelConfigs).Error; err != nil {
+			return err
+		}
+		budgetIDs := make([]string, 0, len(scopedModelConfigs))
+		rateLimitIDs := make([]string, 0, len(scopedModelConfigs))
+		for _, mc := range scopedModelConfigs {
+			if mc.BudgetID != nil {
+				budgetIDs = append(budgetIDs, *mc.BudgetID)
+			}
+			if mc.RateLimitID != nil {
+				rateLimitIDs = append(rateLimitIDs, *mc.RateLimitID)
+			}
+		}
+		if err := txDB.WithContext(ctx).
+			Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id).
+			Delete(&tables.TableModelConfig{}).Error; err != nil {
+			return err
+		}
 		rateLimitID := virtualKey.RateLimitID
 		// Delete the virtual key
 		if err := txDB.WithContext(ctx).Delete(&tables.TableVirtualKey{}, "id = ?", id).Error; err != nil {
@@ -4230,7 +4256,7 @@ func (s *RDBConfigStore) GetModelConfigsPaginated(ctx context.Context, params Mo
 	if err := baseQuery.
 		Preload("Budget").
 		Preload("RateLimit").
-		Order("created_at ASC, id ASC").
+		Order("created_at DESC, id DESC").
 		Offset(offset).
 		Limit(limit).
 		Find(&modelConfigs).Error; err != nil {
@@ -4239,10 +4265,16 @@ func (s *RDBConfigStore) GetModelConfigsPaginated(ctx context.Context, params Mo
 	return modelConfigs, totalCount, nil
 }
 
-// GetModelConfig retrieves a specific model config from the database by model name and optional provider.
-func (s *RDBConfigStore) GetModelConfig(ctx context.Context, modelName string, provider *string) (*tables.TableModelConfig, error) {
+// GetModelConfig retrieves a specific model config from the database by its identity:
+// scope, optional scope ID, model name, and optional provider.
+func (s *RDBConfigStore) GetModelConfig(ctx context.Context, scope string, scopeID *string, modelName string, provider *string) (*tables.TableModelConfig, error) {
 	var modelConfig tables.TableModelConfig
-	query := s.DB().WithContext(ctx).Where("model_name = ?", modelName)
+	query := s.DB().WithContext(ctx).Where("model_name = ?", modelName).Where("scope = ?", scope)
+	if scopeID != nil {
+		query = query.Where("scope_id = ?", *scopeID)
+	} else {
+		query = query.Where("scope_id IS NULL")
+	}
 	if provider != nil {
 		query = query.Where("provider = ?", *provider)
 	} else {

framework/configstore/rdb_test.go

@@ -32,6 +32,7 @@ func setupRDBTestStore(t *testing.T) *RDBConfigStore {
 		&tables.TableVirtualKey{},
 		&tables.TableVirtualKeyProviderConfig{},
 		&tables.TableVirtualKeyProviderConfigKey{},
+		&tables.TableModelConfig{},
 		&tables.TableCustomer{},
 		&tables.TableTeam{},
 		&tables.TableClientConfig{},
@@ -641,6 +642,56 @@ func TestDeleteVirtualKey(t *testing.T) {
 	assert.Error(t, err, "Should not find deleted virtual key")
 }
 
+func TestDeleteVirtualKey_CleansUpScopedModelConfigs(t *testing.T) {
+	store := setupRDBTestStore(t)
+	ctx := context.Background()
+
+	vk := &tables.TableVirtualKey{
+		ID:       "vk-scoped",
+		Name:     "Scoped VK",
+		Value:    "vk-scoped-value",
+		IsActive: schemas.Ptr(true),
+	}
+	require.NoError(t, store.CreateVirtualKey(ctx, vk))
+
+	budget := &tables.TableBudget{ID: "b-scoped", MaxLimit: 100, ResetDuration: "1h"}
+	require.NoError(t, store.CreateBudget(ctx, budget))
+	rateLimit := &tables.TableRateLimit{
+		ID:                 "rl-scoped",
+		TokenMaxLimit:      schemas.Ptr(int64(1000)),
+		TokenResetDuration: schemas.Ptr("1h"),
+	}
+	require.NoError(t, store.CreateRateLimit(ctx, rateLimit))
+
+	mc := &tables.TableModelConfig{
+		ID:          "mc-scoped",
+		ModelName:   "gpt-4",
+		Scope:       tables.ModelConfigScopeVirtualKey,
+		ScopeID:     schemas.Ptr(vk.ID),
+		BudgetID:    &budget.ID,
+		RateLimitID: &rateLimit.ID,
+	}
+	require.NoError(t, store.CreateModelConfig(ctx, mc))
+
+	// Sanity: the scoped config exists before deletion.
+	_, err := store.GetModelConfigByID(ctx, "mc-scoped")
+	require.NoError(t, err)
+
+	// Deleting the VK must cascade-clean its scoped model config and owned budget/rate-limit.
+	require.NoError(t, store.DeleteVirtualKey(ctx, vk.ID))
+
+	_, err = store.GetModelConfigByID(ctx, "mc-scoped")
+	assert.Error(t, err, "scoped model config should be deleted with the VK")
+
+	var budgetCount int64
+	require.NoError(t, store.DB().Model(&tables.TableBudget{}).Where("id = ?", "b-scoped").Count(&budgetCount).Error)
+	assert.Equal(t, int64(0), budgetCount, "owned budget should be deleted")
+
+	var rlCount int64
+	require.NoError(t, store.DB().Model(&tables.TableRateLimit{}).Where("id = ?", "rl-scoped").Count(&rlCount).Error)
+	assert.Equal(t, int64(0), rlCount, "owned rate limit should be deleted")
+}
+
 // =============================================================================
 // Virtual Key Provider Config Tests
 // =============================================================================

framework/configstore/store.go

@@ -257,7 +257,7 @@ type ConfigStore interface {
 	// Model config CRUD
 	GetModelConfigs(ctx context.Context) ([]tables.TableModelConfig, error)
 	GetModelConfigsPaginated(ctx context.Context, params ModelConfigsQueryParams) ([]tables.TableModelConfig, int64, error)
-	GetModelConfig(ctx context.Context, modelName string, provider *string) (*tables.TableModelConfig, 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)
 	CreateModelConfig(ctx context.Context, modelConfig *tables.TableModelConfig, tx ...*gorm.DB) error
 	UpdateModelConfig(ctx context.Context, modelConfig *tables.TableModelConfig, tx ...*gorm.DB) error

framework/configstore/tables/modelconfig.go

@@ -8,14 +8,40 @@ import (
 	"gorm.io/gorm"
 )
 
+// Model config scope values. Scope determines where a model config applies.
+const (
+	ModelConfigScopeGlobal     = "global"
+	ModelConfigScopeVirtualKey = "virtual_key"
+)
+
+// validModelConfigScopes is the set of accepted scope values.
+var validModelConfigScopes = map[string]bool{
+	ModelConfigScopeGlobal:     true,
+	ModelConfigScopeVirtualKey: true,
+}
+
+// IsValidModelConfigScope reports whether scope is a recognized model config scope.
+func IsValidModelConfigScope(scope string) bool {
+	return validModelConfigScopes[scope]
+}
+
 // TableModelConfig represents a model configuration with rate limiting and budgeting
 type TableModelConfig struct {
-	ID          string  `gorm:"primaryKey;type:varchar(255)" json:"id"`
-	ModelName   string  `gorm:"type:varchar(255);not null;uniqueIndex:idx_model_provider" json:"model_name"`
-	Provider    *string `gorm:"type:varchar(50);uniqueIndex:idx_model_provider" json:"provider,omitempty"` // Optional provider, nullable
+	ID        string  `gorm:"primaryKey;type:varchar(255)" json:"id"`
+	ModelName string  `gorm:"type:varchar(255);not null;uniqueIndex:idx_model_scope_provider,priority:3" json:"model_name"`
+	Provider  *string `gorm:"type:varchar(50);uniqueIndex:idx_model_scope_provider,priority:4" json:"provider,omitempty"` // Optional provider, nullable
+	// Scope determines where this config applies: "global" (default) or "virtual_key".
+	Scope string `gorm:"type:varchar(50);not null;default:'global';uniqueIndex:idx_model_scope_provider,priority:1" json:"scope"`
+	// ScopeID is the target of a non-global scope (e.g. the virtual key ID). NULL for global.
+	ScopeID     *string `gorm:"type:varchar(255);uniqueIndex:idx_model_scope_provider,priority:2" json:"scope_id,omitempty"`
 	BudgetID    *string `gorm:"type:varchar(255);index:idx_model_config_budget" json:"budget_id,omitempty"`
 	RateLimitID *string `gorm:"type:varchar(255);index:idx_model_config_rate_limit" json:"rate_limit_id,omitempty"`
 
+	// ScopeName is a non-persisted, API-only field carrying the human-readable name of
+	// the scope target (e.g. the virtual key's name) so the UI can render a label
+	// instead of an opaque scope_id. Populated by the HTTP layer on read.
+	ScopeName string `gorm:"-" json:"scope_name,omitempty"`
+
 	// Relationships
 	Budget    *TableBudget    `gorm:"foreignKey:BudgetID;onDelete:CASCADE" json:"budget,omitempty"`
 	RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"`
@@ -35,6 +61,21 @@ func (TableModelConfig) TableName() string {
 
 // BeforeSave hook for ModelConfig to validate required fields
 func (mc *TableModelConfig) BeforeSave(tx *gorm.DB) error {
+	// Default and validate scope. Global is the implicit default (preserves
+	// pre-scope behavior for configs created without an explicit scope).
+	if strings.TrimSpace(mc.Scope) == "" {
+		mc.Scope = ModelConfigScopeGlobal
+	}
+	if !IsValidModelConfigScope(mc.Scope) {
+		return fmt.Errorf("invalid scope %q for model config", mc.Scope)
+	}
+	// Enforce scope_id rules: global must not have one; non-global requires it.
+	if mc.Scope == ModelConfigScopeGlobal {
+		mc.ScopeID = nil
+	} else if mc.ScopeID == nil || strings.TrimSpace(*mc.ScopeID) == "" {
+		return fmt.Errorf("scope_id is required when scope is %q", mc.Scope)
+	}
+
 	// Validate that ModelName is not empty
 	if strings.TrimSpace(mc.ModelName) == "" {
 		return fmt.Errorf("model_name cannot be empty")