From d7b225142a7dce0465855b17e632b68116a01ae3 Mon Sep 17 00:00:00 2001 From: roroghost17 Date: Tue, 26 May 2026 18:23:05 +0530 Subject: [PATCH] feat: adding scope to budget limits table --- .gitignore | 1 + framework/configstore/clientconfig.go | 8 + framework/configstore/migrations.go | 61 ++++ framework/configstore/migrations_test.go | 61 ++++ framework/configstore/rdb.go | 40 ++- framework/configstore/rdb_test.go | 51 ++++ framework/configstore/store.go | 2 +- framework/configstore/tables/modelconfig.go | 47 ++- .../modelprovidergovernance_test.go | 183 ++++++++++++ plugins/governance/resolver.go | 21 ++ plugins/governance/store.go | 281 ++++++++++++++++-- plugins/governance/test_utils.go | 8 + plugins/governance/tracker.go | 13 + .../bifrost-http/handlers/governance.go | 107 ++++++- transports/bifrost-http/lib/config_test.go | 2 +- .../model-limits/views/modelLimitSheet.tsx | 107 ++++++- .../model-limits/views/modelLimitsTable.tsx | 7 +- .../views/routingRuleInfoSheet.tsx | 33 +- .../routing-rules/views/routingRulesTable.tsx | 12 +- ui/lib/constants/governance.ts | 7 + ui/lib/types/governance.ts | 5 + ui/lib/utils/labels.ts | 14 + ui/lib/utils/routingRules.ts | 15 - 23 files changed, 1007 insertions(+), 79 deletions(-) create mode 100644 ui/lib/utils/labels.ts diff --git a/.gitignore b/.gitignore index 3702ed6ec6..4decb80183 100644 --- a/.gitignore +++ b/.gitignore @@ -99,6 +99,7 @@ dist/ # IDE .idea/ .vscode/ +.zed/ # OS .DS_Store diff --git a/framework/configstore/clientconfig.go b/framework/configstore/clientconfig.go index f6b933d08f..e6a64994cd 100644 --- a/framework/configstore/clientconfig.go +++ b/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 diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 7392a56b94..9ecfac92cc 100644 --- a/framework/configstore/migrations.go +++ b/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{{ diff --git a/framework/configstore/migrations_test.go b/framework/configstore/migrations_test.go index 4675db3ea2..706652461e 100644 --- a/framework/configstore/migrations_test.go +++ b/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")) +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 51483396e0..2b5878a32e 100644 --- a/framework/configstore/rdb.go +++ b/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 { diff --git a/framework/configstore/rdb_test.go b/framework/configstore/rdb_test.go index a73ad60227..35033e9844 100644 --- a/framework/configstore/rdb_test.go +++ b/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 // ============================================================================= diff --git a/framework/configstore/store.go b/framework/configstore/store.go index 789ddf8796..2b3f38cc85 100644 --- a/framework/configstore/store.go +++ b/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 diff --git a/framework/configstore/tables/modelconfig.go b/framework/configstore/tables/modelconfig.go index 5e6b5ba6dc..5e27874cee 100644 --- a/framework/configstore/tables/modelconfig.go +++ b/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") diff --git a/plugins/governance/modelprovidergovernance_test.go b/plugins/governance/modelprovidergovernance_test.go index f0fb5b1d2e..fa206f82e8 100644 --- a/plugins/governance/modelprovidergovernance_test.go +++ b/plugins/governance/modelprovidergovernance_test.go @@ -2100,3 +2100,186 @@ func TestStore_CheckModelBudget_NoCatalog_NoMatch(t *testing.T) { _, err = store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4o", Provider: schemas.OpenAI}, nil) assert.Error(t, err, "Direct match should still work without catalog") } + +// ============================================================================ +// Store Tests - Per-VK-Scoped Model Budget / Rate Limit +// ============================================================================ + +func TestStore_CheckVirtualKeyScopedModelBudget_NilVK(t *testing.T) { + logger := NewMockLogger() + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) + require.NoError(t, err) + + decision, err := store.CheckVirtualKeyScopedModelBudget(context.Background(), nil, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.NoError(t, err) + assert.Equal(t, DecisionAllow, decision) +} + +func TestStore_CheckVirtualKeyScopedModelBudget_NoConfig(t *testing.T) { + logger := NewMockLogger() + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil) + require.NoError(t, err) + + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + decision, err := store.CheckVirtualKeyScopedModelBudget(context.Background(), vk, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.NoError(t, err) + assert.Equal(t, DecisionAllow, decision) +} + +func TestStore_CheckVirtualKeyScopedModelBudget_WithinLimit(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + budget := buildBudget("b1", 100.0, "1h") + mc := buildVKScopedModelConfig("mc1", "gpt-4", nil, vk.ID, budget, nil) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*mc}, + Budgets: []configstoreTables.TableBudget{*budget}, + }, nil) + require.NoError(t, err) + + _, err = store.CheckVirtualKeyScopedModelBudget(context.Background(), vk, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.NoError(t, err, "Should allow when per-VK model budget is within limit") +} + +func TestStore_CheckVirtualKeyScopedModelBudget_Exceeded(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + budget := buildBudgetWithUsage("b1", 100.0, 100.0, "1h") // At limit + mc := buildVKScopedModelConfig("mc1", "gpt-4", nil, vk.ID, budget, nil) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*mc}, + Budgets: []configstoreTables.TableBudget{*budget}, + }, nil) + require.NoError(t, err) + + _, err = store.CheckVirtualKeyScopedModelBudget(context.Background(), vk, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.Error(t, err, "Should reject when per-VK model budget is exceeded") + assert.Contains(t, err.Error(), "budget exceeded") +} + +func TestStore_CheckVirtualKeyScopedModelBudget_OnlyAppliesToMatchingVK(t *testing.T) { + logger := NewMockLogger() + ownerVK := buildVirtualKey("vk1", "vk1-value", "vk1", true) + otherVK := buildVirtualKey("vk2", "vk2-value", "vk2", true) + budget := buildBudgetWithUsage("b1", 100.0, 100.0, "1h") // exceeded + mc := buildVKScopedModelConfig("mc1", "gpt-4", nil, ownerVK.ID, 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 made with a DIFFERENT virtual key must not be affected by vk1's scoped config. + decision, err := store.CheckVirtualKeyScopedModelBudget(context.Background(), otherVK, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.NoError(t, err) + assert.Equal(t, DecisionAllow, decision) +} + +func TestStore_CheckVirtualKeyScopedModelBudget_IgnoresGlobalConfig(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + // A GLOBAL (scope defaults to global) model config that is exceeded. The per-VK scoped + // check must ignore it — global is enforced separately by EvaluateModelAndProviderRequest, + // so the scoped path must not double-count it. + budget := buildBudgetWithUsage("b1", 100.0, 100.0, "1h") + globalMC := buildModelConfig("mc-global", "gpt-4", nil, budget, nil) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*globalMC}, + Budgets: []configstoreTables.TableBudget{*budget}, + }, nil) + require.NoError(t, err) + + decision, err := store.CheckVirtualKeyScopedModelBudget(context.Background(), vk, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.NoError(t, err, "Scoped check must not pick up the global config") + assert.Equal(t, DecisionAllow, decision) + + // Sanity: the global model check DOES still catch the exceeded global budget. + _, gErr := store.CheckModelBudget(context.Background(), &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil) + assert.Error(t, gErr, "Global model check should catch the exceeded global budget") +} + +func TestStore_CheckVirtualKeyScopedModelRateLimit_TokenLimitExceeded(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // tokens at max + mc := buildVKScopedModelConfig("mc1", "gpt-4", nil, vk.ID, nil, rateLimit) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*mc}, + RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, + }, nil) + require.NoError(t, err) + + decision, err := store.CheckVirtualKeyScopedModelRateLimit(context.Background(), vk, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil, nil) + assert.Error(t, err, "Should reject when per-VK model token limit is exceeded") + assert.Equal(t, DecisionTokenLimited, decision) +} + +func TestStore_CheckVirtualKeyScopedModelRateLimit_WithinLimit(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + rateLimit := buildRateLimitWithUsage("rl1", 10000, 100, 1000, 10) // well within limits + mc := buildVKScopedModelConfig("mc1", "gpt-4", nil, vk.ID, nil, rateLimit) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*mc}, + RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, + }, nil) + require.NoError(t, err) + + decision, err := store.CheckVirtualKeyScopedModelRateLimit(context.Background(), vk, &EvaluationRequest{Model: "gpt-4", Provider: schemas.OpenAI}, nil, nil) + assert.NoError(t, err) + assert.Equal(t, DecisionAllow, decision) +} + +// TestStore_VirtualKeyScopedModel_RecordThenCheck_TokenLimitTrips reproduces the reported bug: +// recording usage against a per-VK scoped model config must increment the scoped counter so a +// subsequent check trips. (The original bug only wired the check, not the usage recording.) +func TestStore_VirtualKeyScopedModel_RecordThenCheck_TokenLimitTrips(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + rateLimit := buildRateLimitWithUsage("rl1", 100, 0, 1000000, 0) // 100 token cap, request cap effectively unlimited + mc := buildVKScopedModelConfig("mc1", "claude-opus-4-7", nil, vk.ID, nil, rateLimit) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*mc}, + RateLimits: []configstoreTables.TableRateLimit{*rateLimit}, + }, nil) + require.NoError(t, err) + + req := &EvaluationRequest{Model: "claude-opus-4-7", Provider: schemas.Anthropic} + + // Initially within limit. + decision, err := store.CheckVirtualKeyScopedModelRateLimit(context.Background(), vk, req, nil, nil) + require.NoError(t, err) + require.Equal(t, DecisionAllow, decision) + + // Record usage above the limit (what the tracker does post-response). Provider differs from + // the config's (which is all-providers), exercising the model-only scoped lookup. + require.NoError(t, store.UpdateVirtualKeyScopedModelRateLimitUsageInMemory(context.Background(), vk, "claude-opus-4-7", schemas.Anthropic, 150, true, true)) + + // Now the scoped check must trip. + decision, err = store.CheckVirtualKeyScopedModelRateLimit(context.Background(), vk, req, nil, nil) + assert.Error(t, err) + assert.Equal(t, DecisionTokenLimited, decision) +} + +func TestStore_VirtualKeyScopedModel_RecordThenCheck_BudgetTrips(t *testing.T) { + logger := NewMockLogger() + vk := buildVirtualKey("vk1", "vk1-value", "vk1", true) + budget := buildBudget("b1", 10.0, "1h") // $10 cap, 0 usage + mc := buildVKScopedModelConfig("mc1", "claude-opus-4-7", nil, vk.ID, budget, nil) + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + ModelConfigs: []configstoreTables.TableModelConfig{*mc}, + Budgets: []configstoreTables.TableBudget{*budget}, + }, nil) + require.NoError(t, err) + + req := &EvaluationRequest{Model: "claude-opus-4-7", Provider: schemas.Anthropic} + + decision, err := store.CheckVirtualKeyScopedModelBudget(context.Background(), vk, req, nil) + require.NoError(t, err) + require.Equal(t, DecisionAllow, decision) + + require.NoError(t, store.UpdateVirtualKeyScopedModelBudgetUsageInMemory(context.Background(), vk, "claude-opus-4-7", schemas.Anthropic, 15.0)) + + _, err = store.CheckVirtualKeyScopedModelBudget(context.Background(), vk, req, nil) + assert.Error(t, err, "scoped budget should trip once usage exceeds the cap") +} diff --git a/plugins/governance/resolver.go b/plugins/governance/resolver.go index 3bd169c6ee..0f76dbab8a 100644 --- a/plugins/governance/resolver.go +++ b/plugins/governance/resolver.go @@ -286,6 +286,27 @@ func (r *BudgetResolver) EvaluateVirtualKeyRequest(ctx *schemas.BifrostContext, if budgetResult := r.checkBudgetHierarchy(ctx, vk, evaluationRequest); budgetResult != nil { return budgetResult } + + // 6. Check per-VK-scoped model config rate limits and budgets. These aggregate with + // the global model checks already enforced in EvaluateModelAndProviderRequest — the + // request must satisfy both (most-restrictive wins). Gated on a model being present, + // mirroring the global model checks. + if model != "" { + if decision, err := r.store.CheckVirtualKeyScopedModelRateLimit(ctx, vk, evaluationRequest, nil, nil); err != nil || isRateLimitViolation(decision) { + return &EvaluationResult{ + Decision: decision, + Reason: fmt.Sprintf("Model-level rate limit check failed (virtual key scope): %s", reasonFromErr(err, decision)), + VirtualKey: vk, + } + } + if decision, err := r.store.CheckVirtualKeyScopedModelBudget(ctx, vk, evaluationRequest, nil); err != nil || isBudgetViolation(decision) { + return &EvaluationResult{ + Decision: decision, + Reason: fmt.Sprintf("Model-level budget exceeded (virtual key scope): %s", reasonFromErr(err, decision)), + VirtualKey: vk, + } + } + } } // Find the provider config that matches the request's provider and apply key filtering diff --git a/plugins/governance/store.go b/plugins/governance/store.go index 5da25ca63f..980a4301f8 100644 --- a/plugins/governance/store.go +++ b/plugins/governance/store.go @@ -118,12 +118,18 @@ type GovernanceStore interface { // Model-level governance checks CheckModelBudget(ctx context.Context, request *EvaluationRequest, baselines map[string]float64) (Decision, error) CheckModelRateLimit(ctx context.Context, request *EvaluationRequest, tokensBaselines map[string]int64, requestsBaselines map[string]int64) (Decision, error) + // Per-VK-scoped model-level governance checks (aggregate with the global model checks above) + CheckVirtualKeyScopedModelBudget(ctx context.Context, vk *configstoreTables.TableVirtualKey, request *EvaluationRequest, baselines map[string]float64) (Decision, error) + CheckVirtualKeyScopedModelRateLimit(ctx context.Context, vk *configstoreTables.TableVirtualKey, request *EvaluationRequest, tokensBaselines map[string]int64, requestsBaselines map[string]int64) (Decision, error) // VK-level governance checks CheckVirtualKeyBudget(ctx context.Context, vk *configstoreTables.TableVirtualKey, request *EvaluationRequest, baselines map[string]float64) (Decision, error) CheckVirtualKeyRateLimit(ctx context.Context, vk *configstoreTables.TableVirtualKey, request *EvaluationRequest, tokensBaselines map[string]int64, requestsBaselines map[string]int64) (Decision, error) // In-memory usage updates (for VK-level) UpdateVirtualKeyBudgetUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, provider schemas.ModelProvider, cost float64) error UpdateVirtualKeyRateLimitUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, provider schemas.ModelProvider, tokensUsed int64, shouldUpdateTokens bool, shouldUpdateRequests bool) error + // In-memory usage updates for per-VK-scoped model configs (mirror the global model updates) + UpdateVirtualKeyScopedModelBudgetUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, model string, provider schemas.ModelProvider, cost float64) error + UpdateVirtualKeyScopedModelRateLimitUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, model string, provider schemas.ModelProvider, tokensUsed int64, shouldUpdateTokens bool, shouldUpdateRequests bool) error // In-memory reset checks (return items that need DB sync) ResetExpiredRateLimitsInMemory(ctx context.Context) []*configstoreTables.TableRateLimit ResetExpiredBudgetsInMemory(ctx context.Context) []*configstoreTables.TableBudget @@ -1003,27 +1009,58 @@ func (gs *LocalGovernanceStore) CheckProviderRateLimit(ctx context.Context, requ return gs.CheckRateLimit(ctx, EntityWiseRateLimits{providerKey: []*configstoreTables.TableRateLimit{rateLimit}}, tokensBaselines, requestsBaselines) } -// findModelOnlyConfig looks up a model-only config (no provider) with cross-provider model name normalization. -// Returns the matching config and the display name for error messages. -func (gs *LocalGovernanceStore) findModelOnlyConfig(ctx context.Context, model string) (*configstoreTables.TableModelConfig, string) { - // If modelMatcher is available, try normalized base model name first (cross-provider matching) +// modelConfigStoreKey builds the in-memory cache key for a model config. +func modelConfigStoreKey(scope, scopeID, modelKey string, provider *string) string { + base := modelKey + if provider != nil { + base = fmt.Sprintf("%s:%s", modelKey, *provider) + } + if scope == "" || scope == configstoreTables.ModelConfigScopeGlobal { + return base + } + return fmt.Sprintf("%s:%s:%s", scope, scopeID, base) +} + +// modelConfigScope is one level of the model-config scope chain (name + target ID). +type modelConfigScope struct { + name string + id string +} + +// nonGlobalModelConfigScopeChain returns the non-global scopes that apply to a request +// made with the given virtual key, most specific first. The global scope is intentionally +// excluded because it is enforced separately (and unconditionally) by EvaluateModelAndProviderRequest. +func nonGlobalModelConfigScopeChain(vk *configstoreTables.TableVirtualKey) []modelConfigScope { + if vk == nil { + return nil + } + return []modelConfigScope{{name: configstoreTables.ModelConfigScopeVirtualKey, id: vk.ID}} +} + +// findScopedModelOnlyConfig looks up a model-only config (no provider) within a specific +// scope, preserving cross-provider model-name normalization. scope=="global" reproduces the +// historical global lookup exactly. Returns the matching config and the display name. +func (gs *LocalGovernanceStore) findScopedModelOnlyConfig(ctx context.Context, scope, scopeID, model string) (*configstoreTables.TableModelConfig, string) { + tryKey := func(modelKey string) (*configstoreTables.TableModelConfig, string) { + key := modelConfigStoreKey(scope, scopeID, modelKey, nil) + if value, exists := gs.modelConfigs.Load(key); exists && value != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil { + return mc, modelKey + } + } + return nil, "" + } + // If modelCatalog is available, try normalized base model name first (cross-provider matching) if gs.modelCatalog != nil { baseName := gs.modelCatalog.GetBaseModelName(model) if baseName != model { - if value, exists := gs.modelConfigs.Load(baseName); exists && value != nil { - if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil { - return mc, baseName - } + if mc, name := tryKey(baseName); mc != nil { + return mc, name } } } // Always try direct lookup by original model name as fallback - if value, exists := gs.modelConfigs.Load(model); exists && value != nil { - if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil { - return mc, model - } - } - return nil, "" + return tryKey(model) } // CheckModelBudget performs budget checking for model-level configs (lock-free for high performance) @@ -1224,6 +1261,96 @@ func (gs *LocalGovernanceStore) CheckModelRateLimit(ctx context.Context, request return gs.CheckRateLimit(ctx, entityWiseRateLimits, tokensBaselines, requestsBaselines) } +// CheckVirtualKeyScopedModelBudget enforces budgets from model configs scoped to the +// request's virtual key. These are checked in addition to the global model budgets. A request +// must satisfy both the global model config and any matching per-VK model config. +func (gs *LocalGovernanceStore) CheckVirtualKeyScopedModelBudget(ctx context.Context, vk *configstoreTables.TableVirtualKey, request *EvaluationRequest, baselines map[string]float64) (Decision, error) { + if vk == nil { + return DecisionAllow, nil + } + if baselines == nil { + baselines = map[string]float64{} + } + var model string + var providerStr *string + if request != nil { + model = request.Model + if request.Provider != "" { + p := string(request.Provider) + providerStr = &p + } + } + entityWiseBudgets := EntityWiseBudgets{} + for _, scope := range nonGlobalModelConfigScopeChain(vk) { + // Scoped model+provider config first (more specific) - if provider is provided + if providerStr != nil { + key := modelConfigStoreKey(scope.name, scope.id, model, providerStr) + if value, exists := gs.modelConfigs.Load(key); exists && value != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.Budget != nil { + if budget := gs.LoadBudget(ctx, *mc.BudgetID); budget != nil { + ewKey := fmt.Sprintf("Model:%s:Provider:%s:%s:%s", mc.ModelName, *providerStr, scope.name, scope.id) + entityWiseBudgets[ewKey] = []*configstoreTables.TableBudget{budget} + } + } + } + } + // Always check scoped model-only config (cross-provider normalization preserved) + if mc, _ := gs.findScopedModelOnlyConfig(ctx, scope.name, scope.id, model); mc != nil && mc.Budget != nil { + if budget := gs.LoadBudget(ctx, *mc.BudgetID); budget != nil { + ewKey := fmt.Sprintf("Model:%s:%s:%s", mc.ModelName, scope.name, scope.id) + entityWiseBudgets[ewKey] = []*configstoreTables.TableBudget{budget} + } + } + } + return gs.CheckBudget(ctx, entityWiseBudgets, baselines) +} + +// CheckVirtualKeyScopedModelRateLimit enforces rate limits from model configs scoped to the +// request's virtual key, in addition to the global model rate limits. +func (gs *LocalGovernanceStore) CheckVirtualKeyScopedModelRateLimit(ctx context.Context, vk *configstoreTables.TableVirtualKey, request *EvaluationRequest, tokensBaselines map[string]int64, requestsBaselines map[string]int64) (Decision, error) { + if vk == nil { + return DecisionAllow, nil + } + if tokensBaselines == nil { + tokensBaselines = map[string]int64{} + } + if requestsBaselines == nil { + requestsBaselines = map[string]int64{} + } + var model string + var providerStr *string + if request != nil { + model = request.Model + if request.Provider != "" { + p := string(request.Provider) + providerStr = &p + } + } + entityWiseRateLimits := make(EntityWiseRateLimits) + for _, scope := range nonGlobalModelConfigScopeChain(vk) { + // Scoped model+provider config first - if provider is provided + if providerStr != nil { + key := modelConfigStoreKey(scope.name, scope.id, model, providerStr) + if value, exists := gs.modelConfigs.Load(key); exists && value != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.RateLimitID != nil { + if rateLimit := gs.LoadRateLimit(ctx, *mc.RateLimitID); rateLimit != nil { + ewKey := fmt.Sprintf("Model:%s:Provider:%s:%s:%s", mc.ModelName, *providerStr, scope.name, scope.id) + entityWiseRateLimits[ewKey] = []*configstoreTables.TableRateLimit{rateLimit} + } + } + } + } + // Always check scoped model-only config (cross-provider normalization preserved) + if mc, configKey := gs.findScopedModelOnlyConfig(ctx, scope.name, scope.id, model); mc != nil && mc.RateLimitID != nil { + if rateLimit := gs.LoadRateLimit(ctx, *mc.RateLimitID); rateLimit != nil { + ewKey := fmt.Sprintf("Model:%s:%s:%s", configKey, scope.name, scope.id) + entityWiseRateLimits[ewKey] = []*configstoreTables.TableRateLimit{rateLimit} + } + } + } + return gs.CheckRateLimit(ctx, entityWiseRateLimits, tokensBaselines, requestsBaselines) +} + // CheckUserRateLimit checks if user's rate limit allows the request (enterprise-only) // Community build: silent no-op so user-governance absence never silently denies requests. func (gs *LocalGovernanceStore) CheckUserRateLimit(ctx context.Context, userID string, request *EvaluationRequest, tokensBaselines map[string]int64, requestsBaselines map[string]int64) (Decision, error) { @@ -1346,6 +1473,73 @@ func (gs *LocalGovernanceStore) UpdateProviderAndModelRateLimitUsageInMemory(ctx return nil } +// UpdateVirtualKeyScopedModelBudgetUsageInMemory bumps budget usage for model configs scoped +// to the request's virtual key. This is the post-response counterpart to +// CheckVirtualKeyScopedModelBudget — without it, scoped budgets never increase and never trip. +func (gs *LocalGovernanceStore) UpdateVirtualKeyScopedModelBudgetUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, model string, provider schemas.ModelProvider, cost float64) error { + if vk == nil || model == "" { + return nil + } + var providerStr *string + if provider != "" { + p := string(provider) + providerStr = &p + } + for _, scope := range nonGlobalModelConfigScopeChain(vk) { + // Scoped model+provider config first (more specific) - if provider is set + if providerStr != nil { + key := modelConfigStoreKey(scope.name, scope.id, model, providerStr) + if value, exists := gs.modelConfigs.Load(key); exists && value != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.BudgetID != nil { + if err := gs.BumpBudgetUsage(ctx, *mc.BudgetID, cost); err != nil { + return err + } + } + } + } + // Always bump scoped model-only config (cross-provider normalization preserved) + if mc, _ := gs.findScopedModelOnlyConfig(ctx, scope.name, scope.id, model); mc != nil && mc.BudgetID != nil { + if err := gs.BumpBudgetUsage(ctx, *mc.BudgetID, cost); err != nil { + return err + } + } + } + return nil +} + +// UpdateVirtualKeyScopedModelRateLimitUsageInMemory bumps rate limit counters for model configs +// scoped to the request's virtual key. Post-response counterpart to CheckVirtualKeyScopedModelRateLimit. +func (gs *LocalGovernanceStore) UpdateVirtualKeyScopedModelRateLimitUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, model string, provider schemas.ModelProvider, tokensUsed int64, shouldUpdateTokens bool, shouldUpdateRequests bool) error { + if vk == nil || model == "" { + return nil + } + var providerStr *string + if provider != "" { + p := string(provider) + providerStr = &p + } + for _, scope := range nonGlobalModelConfigScopeChain(vk) { + // Scoped model+provider config first (more specific) - if provider is set + if providerStr != nil { + key := modelConfigStoreKey(scope.name, scope.id, model, providerStr) + if value, exists := gs.modelConfigs.Load(key); exists && value != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.RateLimitID != nil { + if err := gs.BumpRateLimitUsage(ctx, *mc.RateLimitID, tokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { + return err + } + } + } + } + // Always bump scoped model-only config (cross-provider normalization preserved) + if mc, _ := gs.findScopedModelOnlyConfig(ctx, scope.name, scope.id, model); mc != nil && mc.RateLimitID != nil { + if err := gs.BumpRateLimitUsage(ctx, *mc.RateLimitID, tokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { + return err + } + } + } + return nil +} + // UpdateVirtualKeyRateLimitUsageInMemory updates rate limit counters for VK-level rate limits. func (gs *LocalGovernanceStore) UpdateVirtualKeyRateLimitUsageInMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, provider schemas.ModelProvider, tokensUsed int64, shouldUpdateTokens bool, shouldUpdateRequests bool) error { if vk == nil { @@ -1944,22 +2138,29 @@ func (gs *LocalGovernanceStore) rebuildInMemoryStructures(ctx context.Context, c gs.virtualKeys.Store(vk.Value, vk) } - // Build model configs map - // Key format: "modelName" for global configs, "modelName:provider" for provider-specific configs + // Build model configs map. + // Key format (global scope): "modelName" for all-provider configs, "modelName:provider" + // for provider-specific configs. Non-global scopes (e.g. virtual_key) prefix the key with + // "::" via modelConfigStoreKey so they never collide with global configs. // Model names are normalized using GetBaseModelName to prevent duplicate config leakage - // (e.g., "openai/gpt-4o" and "gpt-4o" both store under key "gpt-4o") + // (e.g., "openai/gpt-4o" and "gpt-4o" both store under base "gpt-4o"). for i := range modelConfigs { mc := &modelConfigs[i] + scopeID := "" + if mc.ScopeID != nil { + scopeID = *mc.ScopeID + } if mc.Provider != nil { - // Store under provider-specific key - key := fmt.Sprintf("%s:%s", mc.ModelName, *mc.Provider) + // Provider-specific: store under (scope-prefixed) "modelName:provider" key + key := modelConfigStoreKey(mc.Scope, scopeID, mc.ModelName, mc.Provider) gs.modelConfigs.Store(key, mc) } else { - // Global config (applies to all providers) - store under normalized model name - key := mc.ModelName + // All-provider config - store under normalized (scope-prefixed) model name + modelKey := mc.ModelName if gs.modelCatalog != nil { - key = gs.modelCatalog.GetBaseModelName(mc.ModelName) + modelKey = gs.modelCatalog.GetBaseModelName(mc.ModelName) } + key := modelConfigStoreKey(mc.Scope, scopeID, modelKey, nil) gs.modelConfigs.Store(key, mc) } } @@ -2593,6 +2794,26 @@ func (gs *LocalGovernanceStore) DeleteVirtualKeyInMemory(ctx context.Context, vk } return true // continue iteration }) + + // Evict any model configs scoped to this virtual key (and their budgets/rate-limits). + // Mirrors the DB-side cleanup in DeleteVirtualKey and keeps the in-memory store + // consistent even when the VK entry was already removed. + gs.modelConfigs.Range(func(key, value any) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true + } + if mc.Scope == configstoreTables.ModelConfigScopeVirtualKey && mc.ScopeID != nil && *mc.ScopeID == vkID { + if mc.BudgetID != nil { + gs.DeleteBudget(ctx, *mc.BudgetID) + } + if mc.RateLimitID != nil { + gs.DeleteRateLimit(ctx, *mc.RateLimitID) + } + gs.modelConfigs.Delete(key) + } + return true + }) } // CreateTeamInMemory adds a new team to the in-memory store (lock-free) @@ -2913,16 +3134,22 @@ func (gs *LocalGovernanceStore) UpdateModelConfigInMemory(ctx context.Context, m gs.rateLimits.Store(clone.RateLimit.ID, clone.RateLimit) } - // Determine the key based on whether provider is specified - // Key format: "modelName" for global configs, "modelName:provider" for provider-specific configs + // Determine the (scope-aware) key. Global scope keeps the historical key format; + // non-global scopes are namespaced by modelConfigStoreKey. Scope/scope_id are part of + // a config's identity and do not change on update, so this matches the stored key. + scopeID := "" + if clone.ScopeID != nil { + scopeID = *clone.ScopeID + } if clone.Provider != nil { - key := fmt.Sprintf("%s:%s", clone.ModelName, *clone.Provider) + key := modelConfigStoreKey(clone.Scope, scopeID, clone.ModelName, clone.Provider) gs.modelConfigs.Store(key, &clone) } else { - key := clone.ModelName + modelKey := clone.ModelName if gs.modelCatalog != nil { - key = gs.modelCatalog.GetBaseModelName(clone.ModelName) + modelKey = gs.modelCatalog.GetBaseModelName(clone.ModelName) } + key := modelConfigStoreKey(clone.Scope, scopeID, modelKey, nil) gs.modelConfigs.Store(key, &clone) } diff --git a/plugins/governance/test_utils.go b/plugins/governance/test_utils.go index e19b4197d5..f7313c0d08 100644 --- a/plugins/governance/test_utils.go +++ b/plugins/governance/test_utils.go @@ -263,6 +263,14 @@ func buildModelConfig(id, modelName string, provider *string, budget *configstor return mc } +// buildVKScopedModelConfig builds a model config scoped to a specific virtual key. +func buildVKScopedModelConfig(id, modelName string, provider *string, vkID string, budget *configstoreTables.TableBudget, rateLimit *configstoreTables.TableRateLimit) *configstoreTables.TableModelConfig { + mc := buildModelConfig(id, modelName, provider, budget, rateLimit) + mc.Scope = configstoreTables.ModelConfigScopeVirtualKey + mc.ScopeID = &vkID + return mc +} + func buildProviderWithGovernance(name string, budget *configstoreTables.TableBudget, rateLimit *configstoreTables.TableRateLimit) *configstoreTables.TableProvider { provider := &configstoreTables.TableProvider{ Name: name, diff --git a/plugins/governance/tracker.go b/plugins/governance/tracker.go index ecc1b04513..d445416067 100644 --- a/plugins/governance/tracker.go +++ b/plugins/governance/tracker.go @@ -124,6 +124,19 @@ func (t *UsageTracker) UpdateUsage(ctx context.Context, update *UsageUpdate) { return } + // Update per-VK-scoped model config usage (counterpart to the global model updates above). + // Without this, per-VK model limits never increment and so never trip. + if update.Model != "" { + if err := t.store.UpdateVirtualKeyScopedModelRateLimitUsageInMemory(ctx, vk, update.Model, update.Provider, update.TokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { + t.logger.Error("failed to update scoped model rate limit usage for VK %s: %v", vk.ID, err) + } + if shouldUpdateBudget && update.Cost > 0 { + if err := t.store.UpdateVirtualKeyScopedModelBudgetUsageInMemory(ctx, vk, update.Model, update.Provider, update.Cost); err != nil { + t.logger.Error("failed to update scoped model budget usage for VK %s: %v", vk.ID, err) + } + } + } + // Update rate limit usage (VK-level, provider-config-level, team-level, customer-level) if applicable // Include TeamID and CustomerID checks since rate limits can be configured at those levels if vk.RateLimit != nil || len(vk.ProviderConfigs) > 0 || vk.TeamID != nil || vk.CustomerID != nil { diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 70a97a66c3..ac43cd1241 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -392,11 +392,15 @@ type UpdateCustomerRequest struct { type CreateModelConfigRequest struct { ModelName string `json:"model_name" validate:"required"` Provider *string `json:"provider,omitempty"` // Optional provider, nil means all providers + Scope string `json:"scope,omitempty"` // Defaults to "global" if not provided + ScopeID *string `json:"scope_id,omitempty"` // Required for non-global scopes (e.g. the virtual key ID) Budget *CreateBudgetRequest `json:"budget,omitempty"` RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` } -// UpdateModelConfigRequest represents the request body for updating a model config +// UpdateModelConfigRequest represents the request body for updating a model config. +// Scope and scope_id are part of a config's identity and are intentionally not +// editable here (mirroring model_name/provider) — change them by recreating the config. type UpdateModelConfigRequest struct { ModelName *string `json:"model_name,omitempty"` Provider *string `json:"provider,omitempty"` // Optional provider, nil means no change @@ -2489,6 +2493,33 @@ func validateBudget(budget *configstoreTables.TableBudget) error { // getModelConfigs handles GET /api/governance/model-configs - Get all model configs func (h *GovernanceHandler) getModelConfigs(ctx *fasthttp.RequestCtx) { + fromMemory := string(ctx.QueryArgs().Peek("from_memory")) == "true" + if fromMemory { + data := h.governanceManager.GetGovernanceData(ctx) + if data == nil { + SendError(ctx, 500, "Governance data is not available") + return + } + // Copy into a value slice before enriching: never mutate ScopeName on the + // pointers returned here, so we don't risk racing the in-memory store regardless + // of whether the manager hands back shared pointers or clones. + modelConfigs := make([]configstoreTables.TableModelConfig, 0, len(data.ModelConfigs)) + for _, mc := range data.ModelConfigs { + if mc != nil { + modelConfigs = append(modelConfigs, *mc) + } + } + h.enrichModelConfigScopeNames(ctx, modelConfigs) + SendJSON(ctx, map[string]any{ + "model_configs": modelConfigs, + "count": len(modelConfigs), + "total_count": len(modelConfigs), + "limit": len(modelConfigs), + "offset": 0, + }) + return + } + // Check for pagination parameters limitStr := string(ctx.QueryArgs().Peek("limit")) offsetStr := string(ctx.QueryArgs().Peek("offset")) @@ -2531,6 +2562,7 @@ func (h *GovernanceHandler) getModelConfigs(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, "Failed to retrieve model configs") return } + h.enrichModelConfigScopeNames(ctx, modelConfigs) SendJSON(ctx, map[string]any{ "model_configs": modelConfigs, "count": len(modelConfigs), @@ -2548,6 +2580,7 @@ func (h *GovernanceHandler) getModelConfigs(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, "Failed to retrieve model configs") return } + h.enrichModelConfigScopeNames(ctx, modelConfigs) SendJSON(ctx, map[string]any{ "model_configs": modelConfigs, "count": len(modelConfigs), @@ -2569,11 +2602,38 @@ func (h *GovernanceHandler) getModelConfig(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, "Failed to retrieve model config") return } + h.resolveModelConfigScopeName(ctx, mc, map[string]string{}) SendJSON(ctx, map[string]interface{}{ "model_config": mc, }) } +// resolveModelConfigScopeName populates the transient ScopeName for a single non-global +// model config (currently resolves a virtual_key scope_id to the VK's name). The cache +// lets callers dedupe lookups across many configs. Resolution failures are non-fatal. +func (h *GovernanceHandler) resolveModelConfigScopeName(ctx context.Context, mc *configstoreTables.TableModelConfig, cache map[string]string) { + if mc == nil || mc.Scope != configstoreTables.ModelConfigScopeVirtualKey || mc.ScopeID == nil { + return + } + id := *mc.ScopeID + name, ok := cache[id] + if !ok { + if vk, err := h.configStore.GetVirtualKey(ctx, id); err == nil && vk != nil { + name = vk.Name + } + cache[id] = name + } + mc.ScopeName = name +} + +// enrichModelConfigScopeNames populates ScopeName for each non-global config in the slice. +func (h *GovernanceHandler) enrichModelConfigScopeNames(ctx context.Context, configs []configstoreTables.TableModelConfig) { + cache := map[string]string{} + for i := range configs { + h.resolveModelConfigScopeName(ctx, &configs[i], cache) + } +} + // createModelConfig handles POST /api/governance/model-configs - Create a new model config func (h *GovernanceHandler) createModelConfig(ctx *fasthttp.RequestCtx) { var req CreateModelConfigRequest @@ -2586,18 +2646,51 @@ func (h *GovernanceHandler) createModelConfig(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Model name is required") return } - // Check if model config with same (model_name, provider) already exists - existing, err := h.configStore.GetModelConfig(ctx, req.ModelName, req.Provider) + // Default and validate scope. Global is the implicit default (preserves + // pre-scope behavior). Non-global scopes require a scope_id naming the target. + if req.Scope == "" { + req.Scope = configstoreTables.ModelConfigScopeGlobal + } + if !configstoreTables.IsValidModelConfigScope(req.Scope) { + SendError(ctx, 400, fmt.Sprintf("Invalid scope %q", req.Scope)) + return + } + if req.Scope == configstoreTables.ModelConfigScopeGlobal { + req.ScopeID = nil // normalize: global configs must not carry a scope_id + } else { + if req.ScopeID == nil || *req.ScopeID == "" { + SendError(ctx, 400, "scope_id is required when scope is not global") + return + } + // For the virtual_key scope, the scope_id must reference an existing VK. + if req.Scope == configstoreTables.ModelConfigScopeVirtualKey { + if _, vkErr := h.configStore.GetVirtualKey(ctx, *req.ScopeID); vkErr != nil { + if errors.Is(vkErr, configstore.ErrNotFound) { + SendError(ctx, 400, fmt.Sprintf("Virtual key '%s' not found", *req.ScopeID)) + } else { + logger.Error("failed to verify virtual key for model config scope: %v", vkErr) + SendError(ctx, 500, "Failed to verify virtual key") + } + return + } + } + } + // Check if a model config with the same identity (scope, scope_id, model_name, provider) already exists + existing, err := h.configStore.GetModelConfig(ctx, req.Scope, req.ScopeID, req.ModelName, req.Provider) if err != nil && err != configstore.ErrNotFound { logger.Error("failed to check existing model config: %v", err) SendError(ctx, 500, fmt.Sprintf("Failed to check existing model config: %v", err)) return } if existing != nil { + scopeDesc := "global" + if req.Scope != configstoreTables.ModelConfigScopeGlobal { + scopeDesc = fmt.Sprintf("%s '%s'", req.Scope, *req.ScopeID) + } if req.Provider != nil { - SendError(ctx, 409, fmt.Sprintf("Model config for model '%s' with provider '%s' already exists", req.ModelName, *req.Provider)) + SendError(ctx, 409, fmt.Sprintf("Model config for model '%s' with provider '%s' (%s) already exists", req.ModelName, *req.Provider, scopeDesc)) } else { - SendError(ctx, 409, fmt.Sprintf("Model config for model '%s' (global) already exists", req.ModelName)) + SendError(ctx, 409, fmt.Sprintf("Model config for model '%s' (%s) already exists", req.ModelName, scopeDesc)) } return } @@ -2618,6 +2711,8 @@ func (h *GovernanceHandler) createModelConfig(ctx *fasthttp.RequestCtx) { ID: uuid.NewString(), ModelName: req.ModelName, Provider: req.Provider, + Scope: req.Scope, + ScopeID: req.ScopeID, CreatedAt: time.Now(), UpdatedAt: time.Now(), } @@ -2674,6 +2769,7 @@ func (h *GovernanceHandler) createModelConfig(ctx *fasthttp.RequestCtx) { logger.Error("failed to reload model config in memory: %v", err) preloadedMC = &mc } + h.resolveModelConfigScopeName(ctx, preloadedMC, map[string]string{}) SendJSON(ctx, map[string]interface{}{ "message": "Model config created successfully", "model_config": preloadedMC, @@ -2846,6 +2942,7 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) { logger.Error("failed to reload model config in memory: %v", err) updatedMC = mc } + h.resolveModelConfigScopeName(ctx, updatedMC, map[string]string{}) SendJSON(ctx, map[string]interface{}{ "message": "Model config updated successfully", "model_config": updatedMC, diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index 0ba5d48830..cba6765b33 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -1102,7 +1102,7 @@ func (m *MockConfigStore) GetModelConfigsPaginated(ctx context.Context, params c return nil, 0, nil } -func (m *MockConfigStore) GetModelConfig(ctx context.Context, modelName string, provider *string) (*tables.TableModelConfig, error) { +func (m *MockConfigStore) GetModelConfig(ctx context.Context, scope string, scopeID *string, modelName string, provider *string) (*tables.TableModelConfig, error) { return nil, nil } diff --git a/ui/app/workspace/model-limits/views/modelLimitSheet.tsx b/ui/app/workspace/model-limits/views/modelLimitSheet.tsx index 518ce35de1..a460918865 100644 --- a/ui/app/workspace/model-limits/views/modelLimitSheet.tsx +++ b/ui/app/workspace/model-limits/views/modelLimitSheet.tsx @@ -6,13 +6,15 @@ import NumberAndSelect from "@/components/ui/numberAndSelect"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { DottedSeparator } from "@/components/ui/separator"; import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; -import { resetDurationOptions } from "@/lib/constants/governance"; +import { ComboboxSelect } from "@/components/ui/combobox"; +import { MODEL_LIMIT_SCOPES, resetDurationOptions } from "@/lib/constants/governance"; import { RenderProviderIcon } from "@/lib/constants/icons"; import { ProviderLabels, ProviderName } from "@/lib/constants/logs"; import { getErrorMessage, useCreateModelConfigMutation, useGetProvidersQuery, + useGetVirtualKeysQuery, useLazyGetModelsQuery, useUpdateModelConfigMutation, } from "@/lib/store"; @@ -31,16 +33,23 @@ interface ModelLimitSheetProps { onCancel: () => void; } -const formSchema = z.object({ - modelName: z.string().min(1, "Model name is required"), - provider: z.string().optional(), - budgetMaxLimit: z.number().nonnegative().optional(), - budgetResetDuration: z.string().optional(), - tokenMaxLimit: z.number().int().nonnegative().optional(), - tokenResetDuration: z.string().optional(), - requestMaxLimit: z.number().int().nonnegative().optional(), - requestResetDuration: z.string().optional(), -}); +const formSchema = z + .object({ + modelName: z.string().min(1, "Model name is required"), + provider: z.string().optional(), + scope: z.string().optional(), + scopeId: z.string().optional(), + budgetMaxLimit: z.number().nonnegative().optional(), + budgetResetDuration: z.string().optional(), + tokenMaxLimit: z.number().int().nonnegative().optional(), + tokenResetDuration: z.string().optional(), + requestMaxLimit: z.number().int().nonnegative().optional(), + requestResetDuration: z.string().optional(), + }) + .refine((data) => data.scope !== "virtual_key" || !!data.scopeId, { + message: "Virtual key is required for the Virtual Key scope", + path: ["scopeId"], + }); type FormData = z.infer; @@ -60,6 +69,7 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model }; const { data: providersData } = useGetProvidersQuery(); + const { data: vksData = { virtual_keys: [] } } = useGetVirtualKeysQuery(); const [createModelConfig, { isLoading: isCreating }] = useCreateModelConfigMutation(); const [updateModelConfig, { isLoading: isUpdating }] = useUpdateModelConfigMutation(); const [getModels] = useLazyGetModelsQuery(); @@ -94,6 +104,8 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model defaultValues: { modelName: modelConfig?.model_name || "", provider: modelConfig?.provider || "", + scope: modelConfig?.scope || "global", + scopeId: modelConfig?.scope_id || "", budgetMaxLimit: modelConfig?.budget?.max_limit ?? undefined, budgetResetDuration: modelConfig?.budget?.reset_duration || "1M", tokenMaxLimit: modelConfig?.rate_limit?.token_max_limit ?? undefined, @@ -121,6 +133,8 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model form.reset({ modelName: modelConfig.model_name || "", provider: modelConfig.provider || "", + scope: modelConfig.scope || "global", + scopeId: modelConfig.scope_id || "", budgetMaxLimit: modelConfig.budget?.max_limit ?? undefined, budgetResetDuration: modelConfig.budget?.reset_duration || "1M", tokenMaxLimit: modelConfig.rate_limit?.token_max_limit ?? undefined, @@ -197,6 +211,8 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model await createModelConfig({ model_name: data.modelName, provider, + scope: data.scope || "global", + scope_id: data.scope === "virtual_key" ? data.scopeId : undefined, budget: data.budgetMaxLimit !== undefined && data.budgetMaxLimit !== null ? { @@ -315,6 +331,75 @@ export default function ModelLimitSheet({ modelConfig, onSave, onCancel }: Model )} /> + {/* Scope */} + ( + + Scope + + + + )} + /> + + {/* Virtual Key picker (only for the Virtual Key scope) */} + {form.watch("scope") === "virtual_key" && ( + ( + + Virtual Key + +
+ vk.id !== modelConfig?.scope_id) + .map((vk) => ({ label: vk.name, value: vk.id })), + ]} + value={field.value || null} + onValueChange={(value) => field.onChange(value ?? "")} + placeholder="Select a virtual key..." + disabled={isEditing} + noPortal + /> +
+ + + + )} + /> + )} + {/* Budget Configuration */} diff --git a/ui/app/workspace/model-limits/views/modelLimitsTable.tsx b/ui/app/workspace/model-limits/views/modelLimitsTable.tsx index e8a1d0d5ef..e06a985666 100644 --- a/ui/app/workspace/model-limits/views/modelLimitsTable.tsx +++ b/ui/app/workspace/model-limits/views/modelLimitsTable.tsx @@ -28,6 +28,7 @@ import { useMemo, useState } from "react"; import { toast } from "sonner"; import ModelLimitSheet from "./modelLimitSheet"; import { ModelLimitsEmptyState } from "./modelLimitsEmptyState"; +import { getScopeLabel } from "@/lib/utils/labels"; // Helper to format reset duration for display const formatResetDuration = (duration: string) => { @@ -246,6 +247,7 @@ export default function ModelLimitsTable({ Model Provider + Scope Budget Rate Limit @@ -254,7 +256,7 @@ export default function ModelLimitsTable({ {modelConfigs.length === 0 ? ( - + No matching model limits found. @@ -311,6 +313,9 @@ export default function ModelLimitsTable({ All Providers )} + + {getScopeLabel(config.scope ?? "global")} + {config.budget ? ( diff --git a/ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsx b/ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsx index 5e6a6d91e3..816bf41237 100644 --- a/ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsx +++ b/ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsx @@ -4,14 +4,14 @@ import { Button } from "@/components/ui/button"; import { DottedSeparator } from "@/components/ui/separator"; import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { useSheetNavigation } from "@/hooks/useSheetNavigation"; import { baseRoutingFields } from "@/lib/config/celFieldsRouting"; import { getOperatorLabel } from "@/lib/config/celOperatorsRouting"; import { ProviderIconType, RenderProviderIcon } from "@/lib/constants/icons"; import { getProviderLabel } from "@/lib/constants/logs"; import { useGetCustomersQuery, useGetTeamsQuery, useGetVirtualKeysQuery } from "@/lib/store/apis/governanceApi"; import { RoutingRule } from "@/lib/types/routingRules"; -import { getScopeLabel } from "@/lib/utils/routingRules"; -import { useSheetNavigation } from "@/hooks/useSheetNavigation"; +import { getScopeLabel } from "@/lib/utils/labels"; import { formatDistanceToNow } from "date-fns"; import { Check, Copy, GitMerge, Key } from "lucide-react"; import { useMemo, useState } from "react"; @@ -41,9 +41,15 @@ function formatRuleValue(value: any): string { } function useScopeName(scope: string, scopeId?: string): string | undefined { - const { data: teamsData } = useGetTeamsQuery(undefined, { skip: scope !== "team" || !scopeId }); - const { data: customersData } = useGetCustomersQuery(undefined, { skip: scope !== "customer" || !scopeId }); - const { data: vksData } = useGetVirtualKeysQuery(undefined, { skip: scope !== "virtual_key" || !scopeId }); + const { data: teamsData } = useGetTeamsQuery(undefined, { + skip: scope !== "team" || !scopeId, + }); + const { data: customersData } = useGetCustomersQuery(undefined, { + skip: scope !== "customer" || !scopeId, + }); + const { data: vksData } = useGetVirtualKeysQuery(undefined, { + skip: scope !== "virtual_key" || !scopeId, + }); return useMemo(() => { if (!scopeId) return undefined; @@ -105,7 +111,10 @@ function ConditionRow({ rule }: { rule: RuleType }) { const bareKeyValue = !keyMatch && (isHeader || isParam) && value ? value.includes(":") - ? { key: value.slice(0, value.indexOf(":")), val: value.slice(value.indexOf(":") + 1) } + ? { + key: value.slice(0, value.indexOf(":")), + val: value.slice(value.indexOf(":") + 1), + } : { key: value, val: "" } : null; const keyName = keyMatch?.[1] ?? bareKeyValue?.key; @@ -364,11 +373,19 @@ export function RoutingRuleInfoSheet({ rule, open, onOpenChange, onNavigate, has

Created

- {formatDistanceToNow(new Date(rule.created_at), { addSuffix: true })} + + {formatDistanceToNow(new Date(rule.created_at), { + addSuffix: true, + })} +

Last Updated

- {formatDistanceToNow(new Date(rule.updated_at), { addSuffix: true })} + + {formatDistanceToNow(new Date(rule.updated_at), { + addSuffix: true, + })} +
diff --git a/ui/app/workspace/routing-rules/views/routingRulesTable.tsx b/ui/app/workspace/routing-rules/views/routingRulesTable.tsx index 051ed548ce..495b77550f 100644 --- a/ui/app/workspace/routing-rules/views/routingRulesTable.tsx +++ b/ui/app/workspace/routing-rules/views/routingRulesTable.tsx @@ -24,7 +24,8 @@ import { getProviderLabel } from "@/lib/constants/logs"; import { getErrorMessage } from "@/lib/store"; import { useDeleteRoutingRuleMutation, useUpdateRoutingRuleMutation } from "@/lib/store/apis/routingRulesApi"; import { RoutingRule, RoutingTarget } from "@/lib/types/routingRules"; -import { getPriorityBadgeClass, getScopeLabel, truncateCELExpression } from "@/lib/utils/routingRules"; +import { getScopeLabel } from "@/lib/utils/labels"; +import { getPriorityBadgeClass, truncateCELExpression } from "@/lib/utils/routingRules"; import { ChevronLeft, ChevronRight, Edit, MoreHorizontal, Search, Trash2 } from "lucide-react"; import { useState } from "react"; import { toast } from "sonner"; @@ -240,13 +241,18 @@ export function RoutingRulesTable({ size="md" disabled={!canUpdate} onAsyncCheckedChange={async (checked) => { - await updateRoutingRule({ id: rule.id, data: { enabled: checked } }) + await updateRoutingRule({ + id: rule.id, + data: { enabled: checked }, + }) .unwrap() .then(() => { toast.success(`Rule ${checked ? "enabled" : "disabled"} successfully`); }) .catch((err) => { - toast.error("Failed to update rule", { description: getErrorMessage(err) }); + toast.error("Failed to update rule", { + description: getErrorMessage(err), + }); }); }} /> diff --git a/ui/lib/constants/governance.ts b/ui/lib/constants/governance.ts index 1a494e4914..a674101fc6 100644 --- a/ui/lib/constants/governance.ts +++ b/ui/lib/constants/governance.ts @@ -23,6 +23,13 @@ export const budgetDurationOptions = [ // Must stay in sync with IsCalendarAlignableDuration in framework/configstore/tables/utils.go. export const supportsCalendarAlignment = (duration: string): boolean => duration.length > 0 && /[dwMY]$/.test(duration); +// Scopes available for model limits. "global" applies to all traffic; "virtual_key" +// applies only when a specific virtual key is used. Extensible (team/customer/user) later. +export const MODEL_LIMIT_SCOPES = [ + { label: "Global", value: "global" }, + { label: "Virtual Key", value: "virtual_key" }, +]; + // Map of duration values to short labels for display export const resetDurationLabels: Record = { "1m": "Every Minute", diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 68ff8e1051..1d6022cc9c 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -347,6 +347,9 @@ export interface ModelConfig { id: string; model_name: string; provider?: string; // Optional provider - if empty/null, applies to all providers + scope?: string; // "global" (default) or "virtual_key" + scope_id?: string; // Target of a non-global scope (e.g. the virtual key ID) + scope_name?: string; // Resolved, human-readable name of the scope target (read-only) budget_id?: string; rate_limit_id?: string; // Populated relationships @@ -360,6 +363,8 @@ export interface ModelConfig { export interface CreateModelConfigRequest { model_name: string; provider?: string; // Optional provider - if empty/null, applies to all providers + scope?: string; // Defaults to "global" if omitted + scope_id?: string; // Required for non-global scopes (e.g. the virtual key ID) budget?: CreateBudgetRequest; rate_limit?: CreateRateLimitRequest; } diff --git a/ui/lib/utils/labels.ts b/ui/lib/utils/labels.ts new file mode 100644 index 0000000000..7f1e2a9c5b --- /dev/null +++ b/ui/lib/utils/labels.ts @@ -0,0 +1,14 @@ +/** + * Gets a friendly display name for a scope + * @param scope - The scope value (global|team|customer|virtual_key) + * @returns Friendly display name + */ +export function getScopeLabel(scope: string): string { + const labels: Record = { + global: "Global", + team: "Team", + customer: "Customer", + virtual_key: "Virtual Key", + }; + return labels[scope] || scope; +} \ No newline at end of file diff --git a/ui/lib/utils/routingRules.ts b/ui/lib/utils/routingRules.ts index 56c5f91297..3e78ca42bc 100644 --- a/ui/lib/utils/routingRules.ts +++ b/ui/lib/utils/routingRules.ts @@ -82,21 +82,6 @@ export function stringToFallbacks(str: string): string[] { .filter((s) => s.length > 0); } -/** - * Gets a friendly display name for a scope - * @param scope - The scope value (global|team|customer|virtual_key) - * @returns Friendly display name - */ -export function getScopeLabel(scope: string): string { - const labels: Record = { - global: "Global", - team: "Team", - customer: "Customer", - virtual_key: "Virtual Key", - }; - return labels[scope] || scope; -} - /** * Truncates CEL expression for table display * @param expression - The CEL expression