diff --git a/core/utils.go b/core/utils.go index 8799c99183..dfbc7a95b6 100644 --- a/core/utils.go +++ b/core/utils.go @@ -193,6 +193,7 @@ func newBifrostMessageChan(message *schemas.BifrostResponse) chan *schemas.Bifro func clearCtxForFallback(ctx *schemas.BifrostContext) { ctx.ClearValue(schemas.BifrostContextKeyAPIKeyID) ctx.ClearValue(schemas.BifrostContextKeyAPIKeyName) + ctx.ClearValue(schemas.BifrostContextKeyGovernanceIncludeOnlyKeys) } var supportedBaseProvidersSet = func() map[schemas.ModelProvider]struct{} { diff --git a/framework/changelog.md b/framework/changelog.md index 257ebecae7..b3ee80b636 100644 --- a/framework/changelog.md +++ b/framework/changelog.md @@ -1 +1,2 @@ +- feat: migrate VK provider config allowed keys to explicit allow-list semantics — add AllowAllKeys bool to TableVirtualKeyProviderConfig; backfill existing configs with allow_all_keys=true; empty keys now denies all, ["*"] allows all - feat: add MCPDisableAutoToolInject column to TableClientConfig diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index d21c250017..c3cb588232 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -320,6 +320,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddMCPDisableAutoToolInjectColumn(ctx, db); err != nil { return err } + if err := migrationAddAllowAllKeysToProviderConfig(ctx, db); err != nil { + return err + } return nil } @@ -4338,6 +4341,94 @@ func migrationAddBedrockAssumeRoleColumns(ctx context.Context, db *gorm.DB) erro return nil } +// migrationAddAllowAllKeysToProviderConfig adds the allow_all_keys column to the provider config table +// and backfills existing rows: any provider config with no keys in the join table previously meant +// "allow all keys" (old semantic), so they get allow_all_keys = true to preserve behaviour. +func migrationAddAllowAllKeysToProviderConfig(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_allow_all_keys_to_provider_config", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migratorInstance := tx.Migrator() + + // Add the column if it doesn't exist + if !migratorInstance.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "allow_all_keys") { + if err := migratorInstance.AddColumn(&tables.TableVirtualKeyProviderConfig{}, "allow_all_keys"); err != nil { + return fmt.Errorf("failed to add allow_all_keys column: %w", err) + } + } + + // Backfill: find all provider configs that have no keys in the join table. + // These previously meant "allow all keys", so set allow_all_keys = true. + var allConfigs []tables.TableVirtualKeyProviderConfig + if err := tx.Find(&allConfigs).Error; err != nil { + return fmt.Errorf("failed to query provider configs: %w", err) + } + + // Track which VK IDs were modified so we can recompute their config_hash. + // Without this, subsequent config-sync diff logic would see a stale hash + // and attempt to re-reconcile the VK (potentially undoing the backfill). + modifiedVKIDs := make(map[string]struct{}) + + for _, pc := range allConfigs { + var keyCount int64 + if err := tx.Table("governance_virtual_key_provider_config_keys"). + Where("table_virtual_key_provider_config_id = ?", pc.ID). + Count(&keyCount).Error; err != nil { + return fmt.Errorf("failed to count keys for provider config %d: %w", pc.ID, err) + } + + if keyCount == 0 { + if err := tx.Model(&tables.TableVirtualKeyProviderConfig{}). + Where("id = ?", pc.ID). + Update("allow_all_keys", true).Error; err != nil { + return fmt.Errorf("failed to backfill allow_all_keys for provider config %d: %w", pc.ID, err) + } + modifiedVKIDs[pc.VirtualKeyID] = struct{}{} + } + } + + // Recompute and persist config_hash for every VK that was modified. + for vkID := range modifiedVKIDs { + var vk tables.TableVirtualKey + if err := tx. + Preload("ProviderConfigs"). + Preload("ProviderConfigs.Keys"). + Preload("MCPConfigs"). + First(&vk, "id = ?", vkID).Error; err != nil { + return fmt.Errorf("failed to reload VK %s for hash recomputation: %w", vkID, err) + } + newHash, err := GenerateVirtualKeyHash(vk) + if err != nil { + return fmt.Errorf("failed to generate hash for VK %s: %w", vkID, err) + } + if err := tx.Model(&tables.TableVirtualKey{}). + Where("id = ?", vkID). + Update("config_hash", newHash).Error; err != nil { + return fmt.Errorf("failed to update config_hash for VK %s: %w", vkID, err) + } + log.Printf("[Migration] Recomputed config_hash for VK '%s'", vk.Name) + } + + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migratorInstance := tx.Migrator() + if migratorInstance.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "allow_all_keys") { + if err := migratorInstance.DropColumn(&tables.TableVirtualKeyProviderConfig{}, "allow_all_keys"); err != nil { + return fmt.Errorf("failed to drop allow_all_keys column: %w", err) + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running allow_all_keys migration: %s", err.Error()) + } + return nil +} + // migrationAddMCPDisableAutoToolInjectColumn adds the mcp_disable_auto_tool_inject column to the client config table. // When true, MCP tools are not automatically injected into requests; only explicit context filters apply. func migrationAddMCPDisableAutoToolInjectColumn(ctx context.Context, db *gorm.DB) error { diff --git a/framework/configstore/tables/virtualkey.go b/framework/configstore/tables/virtualkey.go index 4d2f039d35..b8fee6ba14 100644 --- a/framework/configstore/tables/virtualkey.go +++ b/framework/configstore/tables/virtualkey.go @@ -29,13 +29,14 @@ type TableVirtualKeyProviderConfig struct { Provider string `gorm:"type:varchar(50);not null" json:"provider"` Weight *float64 `json:"weight"` AllowedModels []string `gorm:"type:text;serializer:json" json:"allowed_models"` // Empty means all models allowed + AllowAllKeys bool `gorm:"default:false" json:"allow_all_keys"` // True means all keys allowed; false with empty Keys means no keys allowed (deny-by-default) BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"` RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"` // Relationships Budget *TableBudget `gorm:"foreignKey:BudgetID;onDelete:CASCADE" json:"budget,omitempty"` RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"` - Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty means all keys allowed for this provider + Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Used when AllowAllKeys is false; empty means no keys allowed } // TableName sets the table name for each model @@ -60,12 +61,17 @@ func (pc *TableVirtualKeyProviderConfig) UnmarshalJSON(data []byte) error { // Copy all standard fields *pc = TableVirtualKeyProviderConfig(temp.Alias) - // If allowed_keys is provided (config file format), convert to Keys + // If allowed_keys is provided (config file format), convert to Keys or set AllowAllKeys // This takes precedence if Keys is empty but allowed_keys has values if len(temp.AllowedKeys) > 0 && len(pc.Keys) == 0 { - pc.Keys = make([]TableKey, len(temp.AllowedKeys)) - for i, keyName := range temp.AllowedKeys { - pc.Keys[i] = TableKey{Name: keyName} + // Check for wildcard — ["*"] means allow all keys + if len(temp.AllowedKeys) == 1 && temp.AllowedKeys[0] == "*" { + pc.AllowAllKeys = true + } else { + pc.Keys = make([]TableKey, len(temp.AllowedKeys)) + for i, keyName := range temp.AllowedKeys { + pc.Keys[i] = TableKey{Name: keyName} + } } } diff --git a/plugins/governance/changelog.md b/plugins/governance/changelog.md index be5706641b..f4a9d7c43c 100644 --- a/plugins/governance/changelog.md +++ b/plugins/governance/changelog.md @@ -1 +1,2 @@ +- feat: migrate VK provider config allowed keys to deny-by-default — `key_ids: ["*"]` (API) / `allowed_keys: ["*"]` (config) maps to `AllowAllKeys=true` without DB lookups; empty list sets `AllowAllKeys=false` and blocks all keys; resolver populates `includeOnlyKeys` for specific keys or clears it to nil for allow-all to prevent stale filters - feat: enforce VK MCPConfigs as an execution-time allow-list — empty MCPConfigs denies all MCP tools, non-empty validates each tool in both PreMCPHook and evaluateGovernanceRequest; respects disable_auto_tool_inject toggle (transport config key: mcp_disable_auto_tool_inject) and skips auto-injection header when caller already set it diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 3b5f312706..11f6048686 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1012,10 +1012,10 @@ func (p *GovernancePlugin) evaluateGovernanceRequest(ctx *schemas.BifrostContext if result.Decision == DecisionAllow && evaluationRequest.VirtualKey != "" { if evaluationRequest.UserID != "" { // User auth present: only use VK for routing/filtering (skip rate limits and budgets) - result = p.resolver.EvaluateVirtualKeyFiltering(ctx, evaluationRequest.VirtualKey, evaluationRequest.Provider, evaluationRequest.Model, requestType) + result = p.resolver.EvaluateVirtualKeyRequest(ctx, evaluationRequest.VirtualKey, evaluationRequest.Provider, evaluationRequest.Model, requestType, true) } else { // No user auth: full VK governance (routing + limits) - result = p.resolver.EvaluateVirtualKeyRequest(ctx, evaluationRequest.VirtualKey, evaluationRequest.Provider, evaluationRequest.Model, requestType) + result = p.resolver.EvaluateVirtualKeyRequest(ctx, evaluationRequest.VirtualKey, evaluationRequest.Provider, evaluationRequest.Model, requestType, false) } } diff --git a/plugins/governance/resolver.go b/plugins/governance/resolver.go index 978e807073..f78a5da3c9 100644 --- a/plugins/governance/resolver.go +++ b/plugins/governance/resolver.go @@ -171,7 +171,8 @@ func (r *BudgetResolver) isModelRequired(requestType schemas.RequestType) bool { } // EvaluateVirtualKeyRequest evaluates virtual key-specific checks including validation, filtering, rate limits, and budgets -func (r *BudgetResolver) EvaluateVirtualKeyRequest(ctx *schemas.BifrostContext, virtualKeyValue string, provider schemas.ModelProvider, model string, requestType schemas.RequestType) *EvaluationResult { +// skipRateLimitsAndBudgets evaluates to true when we want to skip rate limits and budgets. This is used when user auth is present (user governance handles limits). +func (r *BudgetResolver) EvaluateVirtualKeyRequest(ctx *schemas.BifrostContext, virtualKeyValue string, provider schemas.ModelProvider, model string, requestType schemas.RequestType, skipRateLimitsAndBudgets bool) *EvaluationResult { // 1. Validate virtual key exists and is active vk, exists := r.store.GetVirtualKey(virtualKeyValue) if !exists { @@ -225,100 +226,36 @@ func (r *BudgetResolver) EvaluateVirtualKeyRequest(ctx *schemas.BifrostContext, } // 4. Check rate limits hierarchy (VK level) - if rateLimitResult := r.checkRateLimitHierarchy(ctx, vk, evaluationRequest); rateLimitResult != nil { - return rateLimitResult - } - - // 5. Check budget hierarchy (VK → Team → Customer) - if budgetResult := r.checkBudgetHierarchy(ctx, vk, evaluationRequest); budgetResult != nil { - return budgetResult - } - - // Find the provider config that matches the request's provider and get its allowed keys - for _, pc := range vk.ProviderConfigs { - if schemas.ModelProvider(pc.Provider) == provider && len(pc.Keys) > 0 { - includeOnlyKeys := make([]string, 0, len(pc.Keys)) - for _, dbKey := range pc.Keys { - includeOnlyKeys = append(includeOnlyKeys, dbKey.KeyID) - } - ctx.SetValue(schemas.BifrostContextKeyGovernanceIncludeOnlyKeys, includeOnlyKeys) - break + if !skipRateLimitsAndBudgets { + if rateLimitResult := r.checkRateLimitHierarchy(ctx, vk, evaluationRequest); rateLimitResult != nil { + return rateLimitResult } - } - - // All checks passed - return &EvaluationResult{ - Decision: DecisionAllow, - Reason: "Request allowed by governance policy", - VirtualKey: vk, - } -} -// EvaluateVirtualKeyFiltering evaluates virtual key checks for routing and model/provider filtering only, -// skipping rate limits and budgets. Used when user auth is present (user governance handles limits). -func (r *BudgetResolver) EvaluateVirtualKeyFiltering(ctx *schemas.BifrostContext, virtualKeyValue string, provider schemas.ModelProvider, model string, requestType schemas.RequestType) *EvaluationResult { - // 1. Validate virtual key exists and is active - vk, exists := r.store.GetVirtualKey(virtualKeyValue) - if !exists { - return &EvaluationResult{ - Decision: DecisionVirtualKeyNotFound, - Reason: "Virtual key not found", - } - } - // Set virtual key id and name in context - ctx.SetValue(schemas.BifrostContextKeyGovernanceVirtualKeyID, vk.ID) - ctx.SetValue(schemas.BifrostContextKeyGovernanceVirtualKeyName, vk.Name) - if vk.Team != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamID, vk.Team.ID) - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamName, vk.Team.Name) - if vk.Team.Customer != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, vk.Team.Customer.ID) - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, vk.Team.Customer.Name) - } - } - if vk.Customer != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, vk.Customer.ID) - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, vk.Customer.Name) - } - if !vk.IsActive { - return &EvaluationResult{ - Decision: DecisionVirtualKeyBlocked, - Reason: "Virtual key is inactive", - } - } - // 2. Check provider filtering - if requestType != schemas.MCPToolExecutionRequest && !r.isProviderAllowed(vk, provider) { - return &EvaluationResult{ - Decision: DecisionProviderBlocked, - Reason: fmt.Sprintf("Provider '%s' is not allowed for this virtual key", provider), - VirtualKey: vk, - } - } - // 3. Check model filtering - if r.isModelRequired(requestType) && !r.isModelAllowed(vk, provider, model) { - return &EvaluationResult{ - Decision: DecisionModelBlocked, - Reason: fmt.Sprintf("Model '%s' is not allowed for this virtual key", model), - VirtualKey: vk, + // 5. Check budget hierarchy (VK → Team → Customer) + if budgetResult := r.checkBudgetHierarchy(ctx, vk, evaluationRequest); budgetResult != nil { + return budgetResult } } - // Set include-only keys for provider config routing + // Find the provider config that matches the request's provider and apply key filtering for _, pc := range vk.ProviderConfigs { - if schemas.ModelProvider(pc.Provider) == provider && len(pc.Keys) > 0 { - includeOnlyKeys := make([]string, 0, len(pc.Keys)) - for _, dbKey := range pc.Keys { - includeOnlyKeys = append(includeOnlyKeys, dbKey.KeyID) + if schemas.ModelProvider(pc.Provider) == provider { + if !pc.AllowAllKeys { + // Restrict to specific keys (empty slice = no keys allowed) + includeOnlyKeys := make([]string, 0, len(pc.Keys)) + for _, dbKey := range pc.Keys { + includeOnlyKeys = append(includeOnlyKeys, dbKey.KeyID) + } + ctx.SetValue(schemas.BifrostContextKeyGovernanceIncludeOnlyKeys, includeOnlyKeys) } - ctx.SetValue(schemas.BifrostContextKeyGovernanceIncludeOnlyKeys, includeOnlyKeys) break } } - // Skip rate limits and budgets — user auth handles those + // All checks passed return &EvaluationResult{ Decision: DecisionAllow, - Reason: "Request allowed by governance policy (VK filtering only)", + Reason: "Request allowed by governance policy", VirtualKey: vk, } } diff --git a/plugins/governance/resolver_test.go b/plugins/governance/resolver_test.go index f83c8a53c1..9ba1609ed7 100644 --- a/plugins/governance/resolver_test.go +++ b/plugins/governance/resolver_test.go @@ -26,7 +26,7 @@ func TestBudgetResolver_EvaluateRequest_AllowedRequest(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionAllow, result) assertVirtualKeyFound(t, result) @@ -41,7 +41,7 @@ func TestBudgetResolver_EvaluateRequest_VirtualKeyNotFound(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-nonexistent", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-nonexistent", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionVirtualKeyNotFound, result) } @@ -59,7 +59,7 @@ func TestBudgetResolver_EvaluateRequest_VirtualKeyBlocked(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionVirtualKeyBlocked, result) } @@ -83,7 +83,7 @@ func TestBudgetResolver_EvaluateRequest_ProviderBlocked(t *testing.T) { ctx := &schemas.BifrostContext{} // Try to use OpenAI (not allowed) - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionProviderBlocked, result) assertVirtualKeyFound(t, result) @@ -115,7 +115,7 @@ func TestBudgetResolver_EvaluateRequest_ModelBlocked(t *testing.T) { ctx := &schemas.BifrostContext{} // Try to use gpt-4o-mini (not in allowed list) - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4o-mini", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4o-mini", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionModelBlocked, result) } @@ -137,7 +137,7 @@ func TestBudgetResolver_EvaluateRequest_RateLimitExceeded_TokenLimit(t *testing. resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionTokenLimited, result) assertRateLimitInfo(t, result) @@ -160,7 +160,7 @@ func TestBudgetResolver_EvaluateRequest_RateLimitExceeded_RequestLimit(t *testin resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionRequestLimited, result) } @@ -198,7 +198,7 @@ func TestBudgetResolver_EvaluateRequest_RateLimitExpired(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) // Should allow because rate limit was expired and has been reset assertDecision(t, DecisionAllow, result) @@ -220,7 +220,7 @@ func TestBudgetResolver_EvaluateRequest_BudgetExceeded(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionBudgetExceeded, result) } @@ -247,7 +247,7 @@ func TestBudgetResolver_EvaluateRequest_BudgetExpired(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) // Should allow because budget is expired (will be reset) assertDecision(t, DecisionAllow, result) @@ -282,7 +282,7 @@ func TestBudgetResolver_EvaluateRequest_MultiLevelBudgetHierarchy(t *testing.T) ctx := &schemas.BifrostContext{} // Test: All under limit should pass - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionAllow, result) // Test: VK budget exceeds should fail @@ -293,7 +293,7 @@ func TestBudgetResolver_EvaluateRequest_MultiLevelBudgetHierarchy(t *testing.T) vkBudgetToUpdate.CurrentUsage = 100.0 store.budgets.Store("vk-budget", vkBudgetToUpdate) } - result = resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result = resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionBudgetExceeded, result) } @@ -315,7 +315,7 @@ func TestBudgetResolver_EvaluateRequest_ProviderLevelRateLimit(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionTokenLimited, result) assertRateLimitInfo(t, result) @@ -338,7 +338,7 @@ func TestBudgetResolver_CheckRateLimits_BothExceeded(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assertDecision(t, DecisionRateLimited, result) assert.Contains(t, result.Reason, "rate limit") @@ -476,7 +476,7 @@ func TestBudgetResolver_ContextPopulation(t *testing.T) { resolver := NewBudgetResolver(store, nil, logger) ctx := &schemas.BifrostContext{} - result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest) + result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false) assert.Equal(t, DecisionAllow, result.Decision) diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 560d6ec993..576e0fdd32 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -496,7 +496,10 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { // Get keys for this provider config if specified var keys []configstoreTables.TableKey - if len(pc.KeyIDs) > 0 { + allowAllKeys := false + if len(pc.KeyIDs) == 1 && pc.KeyIDs[0] == "*" { + allowAllKeys = true + } else if len(pc.KeyIDs) > 0 { var err error keys, err = h.configStore.GetKeysByIDs(ctx, pc.KeyIDs) if err != nil { @@ -512,6 +515,7 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { Provider: pc.Provider, Weight: pc.Weight, AllowedModels: pc.AllowedModels, + AllowAllKeys: allowAllKeys, Keys: keys, } @@ -833,7 +837,10 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { } // Get keys for this provider config if specified var keys []configstoreTables.TableKey - if len(pc.KeyIDs) > 0 { + allowAllKeys := false + if len(pc.KeyIDs) == 1 && pc.KeyIDs[0] == "*" { + allowAllKeys = true + } else if len(pc.KeyIDs) > 0 { var err error keys, err = h.configStore.GetKeysByIDs(ctx, pc.KeyIDs) if err != nil { @@ -850,6 +857,7 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { Provider: pc.Provider, Weight: pc.Weight, AllowedModels: pc.AllowedModels, + AllowAllKeys: allowAllKeys, Keys: keys, } // Create budget for provider config if provided @@ -904,7 +912,10 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { // Get keys for this provider config if specified var keys []configstoreTables.TableKey - if len(pc.KeyIDs) > 0 { + allowAllKeys := false + if len(pc.KeyIDs) == 1 && pc.KeyIDs[0] == "*" { + allowAllKeys = true + } else if len(pc.KeyIDs) > 0 { var err error keys, err = h.configStore.GetKeysByIDs(ctx, pc.KeyIDs) if err != nil { @@ -914,6 +925,7 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { return fmt.Errorf("some keys not found for provider %s: expected %d, found %d", pc.Provider, len(pc.KeyIDs), len(keys)) } } + existing.AllowAllKeys = allowAllKeys existing.Keys = keys // Handle budget updates for provider config diff --git a/transports/changelog.md b/transports/changelog.md index 9a8697d3cb..4a02442a62 100644 --- a/transports/changelog.md +++ b/transports/changelog.md @@ -1,2 +1,3 @@ +- feat: VK provider config key_ids now supports ["*"] wildcard to allow all keys; empty key_ids denies all; handler resolves wildcard to AllowAllKeys flag without DB key lookups - feat: add option to disable automatic MCP tool injection per request - feat: virtual key MCP configs now act as an execution-time allow-list — tools not permitted by the VK are blocked at inference and MCP tool execution diff --git a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx index 10276938a8..86ce8e97cc 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx @@ -92,7 +92,7 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe
{!virtualKey.provider_configs || virtualKey.provider_configs.length === 0 ? ( - All providers allowed with default settings + No providers configured (deny-by-default) ) : (
{virtualKey.provider_configs.map((config, index) => ( @@ -130,7 +130,9 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe
Allowed Keys
- {config.keys && config.keys.length > 0 ? ( + {config.allow_all_keys ? ( + All keys allowed + ) : config.keys && config.keys.length > 0 ? (
{config.keys.map((key) => ( @@ -139,7 +141,7 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe ))}
) : ( - All keys allowed + No keys allowed )}
@@ -295,7 +297,7 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe
{!virtualKey.mcp_configs || virtualKey.mcp_configs.length === 0 ? ( - All MCP clients allowed with default settings + No MCP clients configured (deny-by-default) ) : (
diff --git a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx index 9ad01fb082..bb938b2ffa 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx @@ -177,7 +177,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, virtualKey?.provider_configs?.map((config) => ({ ...config, weight: config.weight ?? "", - key_ids: config.keys?.map((key) => key.key_id) || [], + key_ids: config.allow_all_keys ? ["*"] : (config.keys?.map((key) => key.key_id) || []), budget: config.budget ? { max_limit: String(config.budget.max_limit), @@ -722,32 +722,42 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, {(() => { const providerKeys = availableKeys.filter((key) => key.provider === config.provider); const configKeyIds = config.key_ids || []; - const selectedProviderKeys = providerKeys - .filter((key) => configKeyIds.includes(key.key_id)) - .map((key) => ({ + const hasWildcard = configKeyIds.includes("*"); + const allKeyOptions = [ + { + label: "Allow All Keys", + value: "*", + description: "Allow all current and future keys for this provider", + provider: "", + }, + ...providerKeys.map((key) => ({ label: key.name, value: key.key_id, description: key.models?.join(", ") || "", provider: key.provider, - })); - - if (providerKeys.length === 0) return null; + })), + ]; + const selectedProviderKeys = hasWildcard + ? [allKeyOptions[0]] + : providerKeys + .filter((key) => configKeyIds.includes(key.key_id)) + .map((key) => ({ + label: key.name, + value: key.key_id, + description: key.models?.join(", ") || "", + provider: key.provider, + })); return (
-

Keep empty to use all available keys for the provider

+

Select specific keys or allow all. Leave empty to block all keys for this provider.

({ - label: key.name, - value: key.key_id, - description: key.models?.join(", ") || "", - provider: key.provider, - }))} + defaultOptions={allKeyOptions} views={{ multiValue: (multiValueProps: MultiValueProps) => { return ( @@ -790,11 +800,25 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, }} value={selectedProviderKeys} onChange={(keys) => { - // Update key_ids for this provider config - const newKeyIds = keys.map((key) => key.value as string); - handleUpdateProviderConfig(index, "key_ids", newKeyIds); + const hadStar = hasWildcard; + const hasStar = keys.some((k) => k.value === "*"); + if (!hadStar && hasStar) { + // Just selected "Allow All Keys" — set to ["*"] only + handleUpdateProviderConfig(index, "key_ids", ["*"]); + } else if (hadStar && hasStar && keys.length > 1) { + // Had "*", still has "*", but user also selected a specific key — drop "*" + handleUpdateProviderConfig(index, "key_ids", keys.filter((k) => k.value !== "*").map((k) => k.value as string)); + } else { + handleUpdateProviderConfig(index, "key_ids", keys.map((k) => k.value as string)); + } }} - placeholder="Select keys..." + placeholder={ + hasWildcard + ? "All keys allowed" + : configKeyIds.length === 0 + ? "No keys selected" + : "Select keys..." + } className="hover:bg-accent w-full" menuClassName="z-[60] max-h-[300px] overflow-y-auto w-full cursor-pointer custom-scrollbar" /> diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index cb35cfef3f..4ec77a057b 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -89,9 +89,10 @@ export interface VirtualKeyProviderConfig { provider: string; weight: number | null; allowed_models: string[]; + allow_all_keys: boolean; // True means all keys allowed; false with empty keys means no keys allowed budget?: Budget; rate_limit?: RateLimit; - keys?: DBKey[]; // Associated database keys for this provider + keys?: DBKey[]; // Associated database keys for this provider (only used when allow_all_keys is false) } export interface VirtualKeyMCPConfig {