From ad9a9ea6ed1029bf981f397a8b5d43726f7c4bcb Mon Sep 17 00:00:00 2001 From: Anuj Parihar Date: Fri, 31 Jul 2026 01:31:49 +0530 Subject: [PATCH 1/2] feat: add virtual key per model budget creation on the virtual key sheet --- framework/configstore/tables/virtualkey.go | 13 + .../bifrost-http/handlers/governance.go | 307 ++++++++++++++---- .../bifrost-http/handlers/governance_test.go | 21 +- .../views/virtualKeyDetailsSheet.tsx | 85 +++++ .../virtual-keys/views/virtualKeySheet.tsx | 115 +++++-- ui/lib/types/governance.ts | 18 + ui/package-lock.json | 3 + 7 files changed, 472 insertions(+), 90 deletions(-) diff --git a/framework/configstore/tables/virtualkey.go b/framework/configstore/tables/virtualkey.go index a0d12dcdc19..530c585e185 100644 --- a/framework/configstore/tables/virtualkey.go +++ b/framework/configstore/tables/virtualkey.go @@ -37,6 +37,19 @@ type TableVirtualKeyProviderConfig struct { RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"` Budgets []TableBudget `gorm:"foreignKey:ProviderConfigID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` // Multiple budgets with different reset intervals Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty means all keys allowed for this provider + + // ModelBudgets carries per-model budgets/rate-limits under this provider for serialization + // only. They live in VK-scoped model configs (the source of truth), not this table; the + // handler hydrates this field when returning a VK so the sheet can render/edit them. + ModelBudgets []VKProviderModelBudget `gorm:"-" json:"model_budgets,omitempty"` +} + +// VKProviderModelBudget is one per-model budget/rate-limit group under a VK provider config, +// used purely for serialization (reverse-mapped from a VK-scoped model config). +type VKProviderModelBudget struct { + ModelName string `json:"model_name"` + Budgets []TableBudget `json:"budgets,omitempty"` + RateLimit *TableRateLimit `json:"rate_limit,omitempty"` } // TableName sets the table name for each model diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 8e259e3084d..a793e7bbfc2 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -213,6 +213,7 @@ type CreateVirtualKeyRequest struct { BlacklistedModels schemas.BlackList `json:"blacklisted_models,omitempty"` // ["*"] blocks all models; empty blocks none Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget for provider config RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` // Provider-level rate limit + ModelBudgets []vkModelBudgetRequest `json:"model_budgets,omitempty"` // Per-model budgets/rate-limits under this provider KeyIDs schemas.WhiteList `json:"key_ids,omitempty"` // List of DBKey UUIDs to associate with this provider config } `json:"provider_configs,omitempty"` // Empty means no providers allowed (deny-by-default) MCPConfigs []struct { @@ -228,19 +229,36 @@ type CreateVirtualKeyRequest struct { ExpiresAt *time.Time `json:"expires_at,omitempty"` // Optional expiry; nil means never expires } +// vkModelBudgetRequest is one per-model budget/rate-limit group under a provider config +// on VK create. model_name must be a concrete model (not the "*" wildcard, which is the +// provider-level tier). +type vkModelBudgetRequest struct { + ModelName string `json:"model_name" validate:"required"` + Budgets []CreateBudgetRequest `json:"budgets,omitempty"` + RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` +} + +// vkModelBudgetUpdateRequest mirrors vkModelBudgetRequest for VK update (removable rate limit). +type vkModelBudgetUpdateRequest struct { + ModelName string `json:"model_name" validate:"required"` + Budgets []CreateBudgetRequest `json:"budgets,omitempty"` + RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` +} + // UpdateVirtualKeyRequest represents the request body for updating a virtual key type UpdateVirtualKeyRequest struct { Name *string `json:"name,omitempty"` Description *string `json:"description,omitempty"` ProviderConfigs []struct { - ID *uint `json:"id,omitempty"` // null for new entries - Provider string `json:"provider" validate:"required"` - Weight *float64 `json:"weight,omitempty"` - AllowedModels schemas.WhiteList `json:"allowed_models,omitempty"` // ["*"] allows all models; empty denies all - BlacklistedModels schemas.BlackList `json:"blacklisted_models,omitempty"` // ["*"] blocks all models; empty blocks none - Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget for provider config - RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` // Provider-level rate limit - KeyIDs schemas.WhiteList `json:"key_ids,omitempty"` // List of DBKey UUIDs to associate with this provider config + ID *uint `json:"id,omitempty"` // null for new entries + Provider string `json:"provider" validate:"required"` + Weight *float64 `json:"weight,omitempty"` + AllowedModels schemas.WhiteList `json:"allowed_models,omitempty"` // ["*"] allows all models; empty denies all + BlacklistedModels schemas.BlackList `json:"blacklisted_models,omitempty"` // ["*"] blocks all models; empty blocks none + Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget for provider config + RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"` // Provider-level rate limit + ModelBudgets []vkModelBudgetUpdateRequest `json:"model_budgets,omitempty"` // Per-model budgets/rate-limits under this provider (full desired set when provider_configs is supplied) + KeyIDs schemas.WhiteList `json:"key_ids,omitempty"` // List of DBKey UUIDs to associate with this provider config } `json:"provider_configs,omitempty"` MCPConfigs []struct { ID *uint `json:"id,omitempty"` // null for new entries @@ -772,12 +790,27 @@ func (h *GovernanceHandler) reconcileCustomerBudgets(ctx context.Context, tx *go // "leave unchanged" (false, used by partial VK updates) from "set to the given value" (true). // The rateLimit carries only the limit/duration fields (no ID/usage). type vkModelConfigDesired struct { - provider *string + provider *string + // modelName is the model tier this config governs. Empty defaults to the "*" all-models + // tier (VK top-level / per-provider); a concrete model name targets a per-model budget. + modelName string budgetsProvided bool budgets []CreateBudgetRequest rateLimitProvided bool rateLimitRemove bool rateLimit *configstoreTables.TableRateLimit + // reconcileModelBudgets applies to a provider's "*" tier entry: when true the request's + // per-model budgets for this provider are authoritative and configs absent from them are + // pruned; when false (model_budgets omitted) existing per-model configs are left untouched. + reconcileModelBudgets bool +} + +// modelNameOrAll returns the desired model tier, defaulting to the "*" all-models tier. +func (d vkModelConfigDesired) modelNameOrAll() string { + if d.modelName == "" { + return configstoreTables.ModelConfigAllModels + } + return d.modelName } // syncVKGovernanceToModelConfigs folds a virtual key's governance (top-level + per-provider @@ -794,30 +827,48 @@ func (h *GovernanceHandler) syncVKGovernanceToModelConfigs(ctx context.Context, if !reconcileProviders { return nil } + // Keyed by provider + model tier so both the per-provider "*" config and each per-model + // config are retained; everything else under this VK's providers is reconciled away. keep := make(map[string]bool, len(perProvider)) + configuredProviders := make(map[string]bool) + pruneModelBudgets := make(map[string]bool) for _, pg := range perProvider { if pg.provider == nil { continue } - keep[*pg.provider] = true + keep[*pg.provider+"\x00"+pg.modelNameOrAll()] = true + if pg.modelName == "" { + configuredProviders[*pg.provider] = true + if pg.reconcileModelBudgets { + pruneModelBudgets[*pg.provider] = true + } + } if err := h.reconcileVKModelConfig(ctx, tx, vk, pg, usageReset); err != nil { return err } } - // Delete VK-scoped provider model configs whose provider is no longer configured. + // Delete VK-scoped provider model configs that are no longer desired: all tiers for a + // provider dropped from the VK, plus per-model configs pruned from a provider whose + // model_budgets set was supplied. Per-model configs of a still-configured provider whose + // model_budgets were omitted are left untouched (e.g. ones set via the model-config API). var existing []configstoreTables.TableModelConfig if err := tx.Preload("Budgets"). - Where("scope = ? AND scope_id = ? AND model_name = ? AND provider IS NOT NULL", - configstoreTables.ModelConfigScopeVirtualKey, vk.ID, configstoreTables.ModelConfigAllModels). + Where("scope = ? AND scope_id = ? AND provider IS NOT NULL", + configstoreTables.ModelConfigScopeVirtualKey, vk.ID). Find(&existing).Error; err != nil { return err } for i := range existing { mc := &existing[i] - if mc.Provider != nil && !keep[*mc.Provider] { - if err := h.deleteVKModelConfig(ctx, tx, mc); err != nil { - return err - } + if mc.Provider == nil || keep[*mc.Provider+"\x00"+mc.ModelName] { + continue + } + isModelTier := mc.ModelName != configstoreTables.ModelConfigAllModels + if isModelTier && configuredProviders[*mc.Provider] && !pruneModelBudgets[*mc.Provider] { + continue + } + if err := h.deleteVKModelConfig(ctx, tx, mc); err != nil { + return err } } return nil @@ -825,8 +876,9 @@ func (h *GovernanceHandler) syncVKGovernanceToModelConfigs(ctx context.Context, // reconcileVKModelConfig reconciles a single VK-scoped model config to the desired state. func (h *GovernanceHandler) reconcileVKModelConfig(ctx context.Context, tx *gorm.DB, vk *configstoreTables.TableVirtualKey, d vkModelConfigDesired, usageReset *budgetUsageReset) error { + modelName := d.modelNameOrAll() q := tx.Preload("Budgets").Where("scope = ? AND scope_id = ? AND model_name = ?", - configstoreTables.ModelConfigScopeVirtualKey, vk.ID, configstoreTables.ModelConfigAllModels) + configstoreTables.ModelConfigScopeVirtualKey, vk.ID, modelName) if d.provider == nil { q = q.Where("provider IS NULL") } else { @@ -842,7 +894,7 @@ func (h *GovernanceHandler) reconcileVKModelConfig(ctx context.Context, tx *gorm if isNew { mc = configstoreTables.TableModelConfig{ ID: uuid.NewString(), - ModelName: configstoreTables.ModelConfigAllModels, + ModelName: modelName, Provider: d.provider, Scope: configstoreTables.ModelConfigScopeVirtualKey, ScopeID: &vk.ID, @@ -986,6 +1038,96 @@ func rateLimitFromRequestFields(tokenMax *int64, tokenDur *string, reqMax *int64 } } +// maxVKModelBudgetsPerProvider bounds how many per-model budget groups a single provider +// config on a VK may declare. Mirrors the access-profile limit. +const maxVKModelBudgetsPerProvider = 100 + +// validateVKModelBudgetNames enforces the per-provider model-budget invariants: at most +// maxVKModelBudgetsPerProvider groups, and non-empty, unique (trimmed), non-wildcard model names. +func validateVKModelBudgetNames(provider string, names []string) error { + if len(names) > maxVKModelBudgetsPerProvider { + return &badRequestError{err: fmt.Errorf("model_budgets for provider %s exceeds maximum of %d", provider, maxVKModelBudgetsPerProvider)} + } + seen := make(map[string]struct{}, len(names)) + for _, n := range names { + name := strings.TrimSpace(n) + if name == "" { + return &badRequestError{err: fmt.Errorf("model_budgets for provider %s contains an entry with an empty model name", provider)} + } + if name == configstoreTables.ModelConfigAllModels { + return &badRequestError{err: fmt.Errorf("model_budgets for provider %s cannot target the %q wildcard tier", provider, configstoreTables.ModelConfigAllModels)} + } + if _, dup := seen[name]; dup { + return &badRequestError{err: fmt.Errorf("model_budgets for provider %s contains duplicate model %q", provider, name)} + } + seen[name] = struct{}{} + } + return nil +} + +// buildVKCreateModelBudgets validates and folds a provider config's per-model budgets (create form) +// into per-model desired model-config states under the given provider. +func buildVKCreateModelBudgets(provider *string, providerName string, mbs []vkModelBudgetRequest) ([]vkModelConfigDesired, error) { + names := make([]string, len(mbs)) + for i := range mbs { + names[i] = mbs[i].ModelName + } + if err := validateVKModelBudgetNames(providerName, names); err != nil { + return nil, err + } + out := make([]vkModelConfigDesired, 0, len(mbs)) + for _, mb := range mbs { + var rl *configstoreTables.TableRateLimit + if mb.RateLimit != nil { + rl = rateLimitFromRequestFields(mb.RateLimit.TokenMaxLimit, mb.RateLimit.TokenResetDuration, mb.RateLimit.RequestMaxLimit, mb.RateLimit.RequestResetDuration) + } + out = append(out, vkModelConfigDesired{ + provider: provider, + modelName: strings.TrimSpace(mb.ModelName), + budgetsProvided: true, + budgets: mb.Budgets, + rateLimitProvided: mb.RateLimit != nil, + rateLimit: rl, + }) + } + return out, nil +} + +// buildVKUpdateModelBudgets is buildVKCreateModelBudgets for the update form, where a rate limit +// may be explicitly removed. The supplied list is the full desired set for the provider; models +// absent from it are reconciled away by syncVKGovernanceToModelConfigs. +func buildVKUpdateModelBudgets(provider *string, providerName string, mbs []vkModelBudgetUpdateRequest) ([]vkModelConfigDesired, error) { + names := make([]string, len(mbs)) + for i := range mbs { + names[i] = mbs[i].ModelName + } + if err := validateVKModelBudgetNames(providerName, names); err != nil { + return nil, err + } + out := make([]vkModelConfigDesired, 0, len(mbs)) + for _, mb := range mbs { + rlRemove := false + var rl *configstoreTables.TableRateLimit + if mb.RateLimit != nil { + if isRateLimitRemovalRequest(mb.RateLimit) { + rlRemove = true + } else { + rl = rateLimitFromRequestFields(mb.RateLimit.TokenMaxLimit, mb.RateLimit.TokenResetDuration, mb.RateLimit.RequestMaxLimit, mb.RateLimit.RequestResetDuration) + } + } + out = append(out, vkModelConfigDesired{ + provider: provider, + modelName: strings.TrimSpace(mb.ModelName), + budgetsProvided: true, + budgets: mb.Budgets, + rateLimitProvided: mb.RateLimit != nil, + rateLimitRemove: rlRemove, + rateLimit: rl, + }) + } + return out, nil +} + // vkModelConfigIndexKey builds a lookup key for a VK-scoped model config by scope target + provider. func vkModelConfigIndexKey(scopeID string, provider *string) string { if provider == nil { @@ -996,9 +1138,10 @@ func vkModelConfigIndexKey(scopeID string, provider *string) string { // applyVKGovernanceFromModelConfigs repopulates a VK's (and each provider config's) budgets and // rate-limit from the VK-scoped model configs that own them โ€” for serialization only (so the VK -// sheet still renders the governance it edits). byKey is keyed by vkModelConfigIndexKey. +// sheet still renders the governance it edits). byKey holds the "*" (all-models) configs keyed by +// vkModelConfigIndexKey; perModelByKey holds each provider's per-model budgets keyed the same way. // The reverse of syncVKGovernanceToModelConfigs. -func applyVKGovernanceFromModelConfigs(vk *configstoreTables.TableVirtualKey, byKey map[string]*configstoreTables.TableModelConfig) { +func applyVKGovernanceFromModelConfigs(vk *configstoreTables.TableVirtualKey, byKey map[string]*configstoreTables.TableModelConfig, perModelByKey map[string][]configstoreTables.VKProviderModelBudget) { if mc := byKey[vkModelConfigIndexKey(vk.ID, nil)]; mc != nil { vk.Budgets = mc.Budgets vk.RateLimit = mc.RateLimit @@ -1011,33 +1154,50 @@ func applyVKGovernanceFromModelConfigs(vk *configstoreTables.TableVirtualKey, by pc.RateLimit = mc.RateLimit pc.RateLimitID = mc.RateLimitID } + pc.ModelBudgets = perModelByKey[vkModelConfigIndexKey(vk.ID, &pc.Provider)] + } +} + +// buildVKModelBudgetsIndex groups the specific-model VK-scoped model configs by +// vkModelConfigIndexKey(vkID, provider), sorted by model name for stable output. +func buildVKModelBudgetsIndex(mcs []*configstoreTables.TableModelConfig) map[string][]configstoreTables.VKProviderModelBudget { + byKey := make(map[string][]configstoreTables.VKProviderModelBudget) + for _, mc := range mcs { + if mc == nil || mc.Scope != configstoreTables.ModelConfigScopeVirtualKey || mc.ScopeID == nil { + continue + } + if mc.ModelName == configstoreTables.ModelConfigAllModels || mc.Provider == nil { + continue + } + key := vkModelConfigIndexKey(*mc.ScopeID, mc.Provider) + byKey[key] = append(byKey[key], configstoreTables.VKProviderModelBudget{ + ModelName: mc.ModelName, + Budgets: mc.Budgets, + RateLimit: mc.RateLimit, + }) + } + for key := range byKey { + sort.Slice(byKey[key], func(i, j int) bool { return byKey[key][i].ModelName < byKey[key][j].ModelName }) } + return byKey } -// hydrateVKGovernance reverse-maps a single VK's governance from its VK-scoped model configs. +// hydrateVKGovernance reverse-maps a single VK's governance (top-level, per-provider, and +// per-model budgets) from its VK-scoped model configs in one bulk load. func (h *GovernanceHandler) hydrateVKGovernance(ctx context.Context, vk *configstoreTables.TableVirtualKey) { if vk == nil { return } - byKey := make(map[string]*configstoreTables.TableModelConfig) - add := func(provider *string) { - mc, err := h.configStore.GetModelConfig(ctx, configstoreTables.ModelConfigScopeVirtualKey, &vk.ID, configstoreTables.ModelConfigAllModels, provider) - if err != nil { - if !errors.Is(err, configstore.ErrNotFound) { - logger.Error("failed to get model config for VK governance hydration: %v", err) - } - return - } - if mc != nil { - byKey[vkModelConfigIndexKey(vk.ID, provider)] = mc - } + mcs, err := h.configStore.GetModelConfigsByScopeAndScopeIDs(ctx, configstoreTables.ModelConfigScopeVirtualKey, []string{vk.ID}) + if err != nil { + logger.Error("failed to load model configs for VK governance hydration: %v", err) + return } - add(nil) - for i := range vk.ProviderConfigs { - prov := vk.ProviderConfigs[i].Provider - add(&prov) + ptrs := make([]*configstoreTables.TableModelConfig, len(mcs)) + for i := range mcs { + ptrs[i] = &mcs[i] } - applyVKGovernanceFromModelConfigs(vk, byKey) + applyVKGovernanceFromModelConfigs(vk, buildVKModelConfigIndex(ptrs), buildVKModelBudgetsIndex(ptrs)) } // buildVKModelConfigIndex builds a lookup map of VK-scoped model configs keyed by @@ -1063,15 +1223,14 @@ func (h *GovernanceHandler) hydrateVKListGovernance(ctx context.Context, vks []c logger.Error("failed to load model configs for VK governance hydration: %v", err) return } - byKey := make(map[string]*configstoreTables.TableModelConfig) + ptrs := make([]*configstoreTables.TableModelConfig, len(allMCs)) for i := range allMCs { - mc := &allMCs[i] - if mc.Scope == configstoreTables.ModelConfigScopeVirtualKey && mc.ModelName == configstoreTables.ModelConfigAllModels && mc.ScopeID != nil { - byKey[vkModelConfigIndexKey(*mc.ScopeID, mc.Provider)] = mc - } + ptrs[i] = &allMCs[i] } + byKey := buildVKModelConfigIndex(ptrs) + perModelByKey := buildVKModelBudgetsIndex(ptrs) for i := range vks { - applyVKGovernanceFromModelConfigs(&vks[i], byKey) + applyVKGovernanceFromModelConfigs(&vks[i], byKey, perModelByKey) } } @@ -1371,13 +1530,14 @@ func (h *GovernanceHandler) getVirtualKeys(ctx *fasthttp.RequestCtx) { return virtualKeys[i].CreatedAt.Before(virtualKeys[j].CreatedAt) }) byKey := buildVKModelConfigIndex(data.ModelConfigs) + perModelByKey := buildVKModelBudgetsIndex(data.ModelConfigs) hydratedVKs := make([]*configstoreTables.TableVirtualKey, len(virtualKeys)) for i, vk := range virtualKeys { clone := *vk pcs := make([]configstoreTables.TableVirtualKeyProviderConfig, len(vk.ProviderConfigs)) copy(pcs, vk.ProviderConfigs) clone.ProviderConfigs = pcs - applyVKGovernanceFromModelConfigs(&clone, byKey) + applyVKGovernanceFromModelConfigs(&clone, byKey, perModelByKey) h.applyExternalBudgets(ctx, &clone) hydratedVKs[i] = &clone } @@ -1624,12 +1784,18 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { pcRateLimit = rateLimitFromRequestFields(pc.RateLimit.TokenMaxLimit, pc.RateLimit.TokenResetDuration, pc.RateLimit.RequestMaxLimit, pc.RateLimit.RequestResetDuration) } vkGovProviders = append(vkGovProviders, vkModelConfigDesired{ - provider: &providerNameStr, - budgetsProvided: true, - budgets: pc.Budgets, - rateLimitProvided: pc.RateLimit != nil, - rateLimit: pcRateLimit, + provider: &providerNameStr, + budgetsProvided: true, + budgets: pc.Budgets, + rateLimitProvided: pc.RateLimit != nil, + rateLimit: pcRateLimit, + reconcileModelBudgets: true, }) + modelDesired, err := buildVKCreateModelBudgets(&providerNameStr, providerNameStr, pc.ModelBudgets) + if err != nil { + return err + } + vkGovProviders = append(vkGovProviders, modelDesired...) } } // Fold VK top-level + per-provider governance into VK-scoped model configs. @@ -1712,13 +1878,14 @@ func (h *GovernanceHandler) getVirtualKey(ctx *fasthttp.RequestCtx) { return } byKey := buildVKModelConfigIndex(data.ModelConfigs) + perModelByKey := buildVKModelBudgetsIndex(data.ModelConfigs) for _, vk := range data.VirtualKeys { if vk.ID == vkID { clone := *vk pcs := make([]configstoreTables.TableVirtualKeyProviderConfig, len(vk.ProviderConfigs)) copy(pcs, vk.ProviderConfigs) clone.ProviderConfigs = pcs - applyVKGovernanceFromModelConfigs(&clone, byKey) + applyVKGovernanceFromModelConfigs(&clone, byKey, perModelByKey) h.applyExternalBudgets(ctx, &clone) SendJSON(ctx, map[string]interface{}{ "virtual_key": &clone, @@ -2029,12 +2196,18 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { pcRL = rateLimitFromRequestFields(pc.RateLimit.TokenMaxLimit, pc.RateLimit.TokenResetDuration, pc.RateLimit.RequestMaxLimit, pc.RateLimit.RequestResetDuration) } vkGovProviders = append(vkGovProviders, vkModelConfigDesired{ - provider: &pName, - budgetsProvided: true, - budgets: pc.Budgets, - rateLimitProvided: pc.RateLimit != nil, - rateLimit: pcRL, + provider: &pName, + budgetsProvided: true, + budgets: pc.Budgets, + rateLimitProvided: pc.RateLimit != nil, + rateLimit: pcRL, + reconcileModelBudgets: pc.ModelBudgets != nil, }) + modelDesired, err := buildVKUpdateModelBudgets(&pName, pName, pc.ModelBudgets) + if err != nil { + return err + } + vkGovProviders = append(vkGovProviders, modelDesired...) } else { // Update existing provider config existing, ok := existingConfigsMap[*pc.ID] @@ -2088,13 +2261,19 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { } } vkGovProviders = append(vkGovProviders, vkModelConfigDesired{ - provider: &pName, - budgetsProvided: pc.Budgets != nil, - budgets: pc.Budgets, - rateLimitProvided: pc.RateLimit != nil, - rateLimitRemove: rlRemove, - rateLimit: pcRL, + provider: &pName, + budgetsProvided: pc.Budgets != nil, + budgets: pc.Budgets, + rateLimitProvided: pc.RateLimit != nil, + rateLimitRemove: rlRemove, + rateLimit: pcRL, + reconcileModelBudgets: pc.ModelBudgets != nil, }) + modelDesired, err := buildVKUpdateModelBudgets(&pName, pName, pc.ModelBudgets) + if err != nil { + return err + } + vkGovProviders = append(vkGovProviders, modelDesired...) if err := h.configStore.UpdateVirtualKeyProviderConfig(ctx, &existing, tx); err != nil { return err } @@ -5163,7 +5342,7 @@ func (h *GovernanceHandler) collectVKModelUsage(ctx context.Context, vk *configs for i := range mcs { ptrs[i] = &mcs[i] } - applyVKGovernanceFromModelConfigs(vk, buildVKModelConfigIndex(ptrs)) + applyVKGovernanceFromModelConfigs(vk, buildVKModelConfigIndex(ptrs), buildVKModelBudgetsIndex(ptrs)) models := make([]quotaModelUsage, 0) for i := range mcs { diff --git a/transports/bifrost-http/handlers/governance_test.go b/transports/bifrost-http/handlers/governance_test.go index b9d1b49a035..a673a74b8b2 100644 --- a/transports/bifrost-http/handlers/governance_test.go +++ b/transports/bifrost-http/handlers/governance_test.go @@ -115,6 +115,23 @@ func (m *mockRotateConfigStore) GetModelConfig(_ context.Context, scope string, return lookupVKModelConfig(m.modelConfigs, scope, scopeID, modelName, provider) } +// GetModelConfigsByScopeAndScopeIDs returns the stored configs matching the scope and scope IDs, +// mirroring the bulk load hydrateVKGovernance performs. +func (m *mockRotateConfigStore) GetModelConfigsByScopeAndScopeIDs(_ context.Context, scope string, scopeIDs []string) ([]configstoreTables.TableModelConfig, error) { + idset := make(map[string]bool, len(scopeIDs)) + for _, id := range scopeIDs { + idset[id] = true + } + var out []configstoreTables.TableModelConfig + for _, mc := range m.modelConfigs { + if mc == nil || mc.Scope != scope || mc.ScopeID == nil || !idset[*mc.ScopeID] { + continue + } + out = append(out, *mc) + } + return out, nil +} + type mockRotateGovernanceManager struct { GovernanceManager store *mockRotateConfigStore @@ -3341,7 +3358,7 @@ func TestApplyVKGovernanceFromModelConfigs_PreservesDirectlyAttachedBudget(t *te } // No VK-scoped model config exists for this VK. - applyVKGovernanceFromModelConfigs(vk, map[string]*configstoreTables.TableModelConfig{}) + applyVKGovernanceFromModelConfigs(vk, map[string]*configstoreTables.TableModelConfig{}, nil) if len(vk.Budgets) != 1 || vk.Budgets[0].ID != "bud-direct" { t.Fatalf("directly attached budget was wiped: got %+v", vk.Budgets) @@ -3376,7 +3393,7 @@ func TestApplyVKGovernanceFromModelConfigs_OverlaysModelConfigGovernance(t *test }, } - applyVKGovernanceFromModelConfigs(vk, byKey) + applyVKGovernanceFromModelConfigs(vk, byKey, nil) if len(vk.Budgets) != 1 || vk.Budgets[0].ID != "bud-mc" { t.Fatalf("expected model-config budget overlaid, got %+v", vk.Budgets) diff --git a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx index c7348b5ec75..a80585a3d4d 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx @@ -414,6 +414,91 @@ export default function VirtualKeyDetailSheet({ )} + + {/* Model Budgets โ€” per-model caps/rate-limits under this provider */} + {config.model_budgets && config.model_budgets.length > 0 && ( + <> + +
+

Model Budgets

+ {config.model_budgets.map((mb, mbIdx) => ( +
+ {mb.model_name} + + {/* Budgets */} + {mb.budgets && mb.budgets.length > 0 + ? mb.budgets.map((b, bIdx) => ( +
+ {!isManagedByProfile && b.id ? ( +
+ saveBudgetOverride(b.id, data)} + onRemove={() => clearBudgetOverride(b.id)} + disabled={!canUpdateVirtualKeys} + calendarAligned={virtualKey.calendar_aligned} + /> +
+ ) : null} + + {hasActiveBudgetOverride(b) ? ( +

+ Base {formatCurrency(b.max_limit)} + {formatCurrency(b.override_amount ?? 0)} override +

+ ) : null} +
+ + Resets {parseResetPeriod(b.reset_duration)} + {virtualKey.calendar_aligned && supportsCalendarAlignment(b.reset_duration) && " (calendar)"} + + {b.last_reset ? ( + Last reset {formatDistanceToNow(new Date(b.last_reset), { addSuffix: true })} + ) : null} +
+
+ )) + : null} + + {/* Token Limits */} + {mb.rate_limit?.token_max_limit != null ? ( +
+ TOKEN LIMITS + n.toLocaleString()} + /> +
+ Resets {parseResetPeriod(mb.rate_limit.token_reset_duration || "")} + {virtualKey.calendar_aligned && + supportsCalendarAlignment(mb.rate_limit.token_reset_duration || "") && + " (calendar)"} +
+
+ ) : null} + + {/* Request Limits */} + {mb.rate_limit?.request_max_limit != null ? ( +
+ REQUEST LIMITS + n.toLocaleString()} + /> +
+ Resets {parseResetPeriod(mb.rate_limit.request_reset_duration || "")} + {virtualKey.calendar_aligned && + supportsCalendarAlignment(mb.rate_limit.request_reset_duration || "") && + " (calendar)"} +
+
+ ) : null} +
+ ))} +
+ + )} ))} diff --git a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx index ca0e1039ff2..d4dd7392342 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx @@ -108,6 +108,31 @@ const providerConfigSchema = z.object({ request_reset_duration: z.string().optional(), }) .optional(), + // Per-model budgets/rate-limits under this provider + model_budgets: z + .array( + z.object({ + model_name: z.string().trim().min(1, "Model name is required"), + budgets: z + .array( + z.object({ + id: z.string().optional(), + max_limit: z.number().nonnegative().optional(), + reset_duration: z.string().optional(), + }), + ) + .optional(), + rate_limit: z + .object({ + token_max_limit: z.number().int().nonnegative().optional(), + token_reset_duration: z.string().optional(), + request_max_limit: z.number().int().nonnegative().optional(), + request_reset_duration: z.string().optional(), + }) + .optional(), + }), + ) + .optional(), }); const mcpConfigSchema = z.object({ @@ -346,6 +371,22 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC request_reset_duration: config.rate_limit.request_reset_duration, } : undefined, + model_budgets: config.model_budgets?.map((mb) => ({ + model_name: mb.model_name, + budgets: mb.budgets?.map((b) => ({ + id: b.id, + max_limit: b.max_limit, + reset_duration: b.reset_duration, + })), + rate_limit: mb.rate_limit + ? { + token_max_limit: mb.rate_limit.token_max_limit ?? undefined, + token_reset_duration: mb.rate_limit.token_reset_duration, + request_max_limit: mb.rate_limit.request_max_limit ?? undefined, + request_reset_duration: mb.rate_limit.request_reset_duration, + } + : undefined, + })), })) || [], mcpConfigs: virtualKey?.mcp_configs?.map((config) => ({ @@ -534,31 +575,50 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC form.setValue("requestResetDuration", "1h", { shouldDirty: true }); }; - const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => { - return configs.map((config) => ({ - ...config, - budgets: config.budgets?.filter((b): b is { id?: string; max_limit: number; reset_duration: string } => b.max_limit !== undefined), - weight: config.weight ?? null, - rate_limit: (() => { - const hasTokenMaxLimit = config.rate_limit?.token_max_limit !== undefined; - const hasRequestMaxLimit = config.rate_limit?.request_max_limit !== undefined; - if (hasTokenMaxLimit || hasRequestMaxLimit) { - return { - token_max_limit: config.rate_limit?.token_max_limit ?? null, - token_reset_duration: hasTokenMaxLimit ? config.rate_limit?.token_reset_duration || "1h" : null, - request_max_limit: config.rate_limit?.request_max_limit ?? null, - request_reset_duration: hasRequestMaxLimit ? config.rate_limit?.request_reset_duration || "1h" : null, - }; - } - - const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider)); - if (existingConfig?.rate_limit) { - return {}; - } + // Build a request rate-limit payload from the form's rate-limit fields. Returns the field + // values when a limit is set, {} to clear an existing rate limit (removal), or undefined. + const normalizeRateLimit = ( + rl: { token_max_limit?: number; token_reset_duration?: string; request_max_limit?: number; request_reset_duration?: string } | undefined, + hadExisting: boolean, + ) => { + const hasToken = rl?.token_max_limit !== undefined; + const hasRequest = rl?.request_max_limit !== undefined; + if (hasToken || hasRequest) { + return { + token_max_limit: rl?.token_max_limit ?? null, + token_reset_duration: hasToken ? rl?.token_reset_duration || "1h" : null, + request_max_limit: rl?.request_max_limit ?? null, + request_reset_duration: hasRequest ? rl?.request_reset_duration || "1h" : null, + }; + } + return hadExisting ? {} : undefined; + }; - return undefined; - })(), - })); + const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => { + return configs.map((config) => { + const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider)); + return { + ...config, + budgets: config.budgets?.filter((b): b is { id?: string; max_limit: number; reset_duration: string } => b.max_limit !== undefined), + weight: config.weight ?? null, + rate_limit: normalizeRateLimit(config.rate_limit, !!existingConfig?.rate_limit), + // Full desired per-model set: drop unfilled models, keep an empty array so the + // backend prunes any per-model budgets removed here. + model_budgets: (config.model_budgets || []) + .filter((mb) => mb.model_name && mb.model_name.trim() !== "") + .map((mb) => { + const existingMB = existingConfig?.model_budgets?.find((m) => m.model_name === mb.model_name.trim()); + return { + model_name: mb.model_name.trim(), + budgets: (mb.budgets || []).filter( + (b): b is { id?: string; max_limit: number; reset_duration: string } => b.max_limit !== undefined, + ), + rate_limit: normalizeRateLimit(mb.rate_limit, !!existingMB?.rate_limit), + }; + }) + .filter((mb) => mb.budgets.length > 0 || mb.rate_limit !== undefined), + }; + }); }; const parseResetDurationMs = (duration?: string) => { @@ -1137,6 +1197,7 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC providerLabel={providerLabel} iconProvider={iconProvider} providerKeys={providerKeys} + showModelBudgets onRemove={() => handleRemoveProvider(index)} value={{ providerName: config.provider, @@ -1146,6 +1207,11 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC keyIds: config.key_ids || [], budgets: config.budgets || [], rateLimit: config.rate_limit ?? null, + modelBudgets: (config.model_budgets || []).map((mb) => ({ + model_name: mb.model_name, + budgets: mb.budgets || [], + rate_limit: mb.rate_limit, + })), }} onChange={(next) => { const updated = [...providerConfigs]; @@ -1162,6 +1228,7 @@ export default function VirtualKeySheet({ virtualKey, defaultTeamId, onSave, onC reset_config: l.reset_config, })), rate_limit: next.rateLimit ?? undefined, + model_budgets: next.modelBudgets, }; form.setValue("providerConfigs", updated, { shouldDirty: true, diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 31f44bdd4b6..0c624510b20 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -120,6 +120,13 @@ export interface VirtualKey { config_hash?: string; // Present when config is synced from config.json } +// Per-model budgets/rate-limits under a provider config, surfaced on the VK for display/edit. +export interface VirtualKeyModelBudget { + model_name: string; + budgets?: Budget[]; + rate_limit?: RateLimit; +} + export interface VirtualKeyProviderConfig { id?: number; provider: string; @@ -129,6 +136,7 @@ export interface VirtualKeyProviderConfig { allow_all_keys: boolean; // True means all keys allowed; false with empty keys means no keys allowed budgets?: Budget[]; rate_limit?: RateLimit; + model_budgets?: VirtualKeyModelBudget[]; // Per-model budgets/rate-limits under this provider keys?: DBKey[]; // Associated database keys for this provider (only used when allow_all_keys is false) } @@ -165,6 +173,14 @@ export interface UsageStats { requests_last_reset: string; } +// One per-model budget/rate-limit group in a provider-config request. model_name must be a +// concrete model (not the "*" wildcard, which is the provider-level tier). +export interface VirtualKeyModelBudgetRequest { + model_name: string; + budgets?: CreateBudgetRequest[]; + rate_limit?: CreateRateLimitRequest; +} + // Request interfaces for provider config operations export interface VirtualKeyProviderConfigRequest { provider: string; @@ -173,6 +189,7 @@ export interface VirtualKeyProviderConfigRequest { blacklisted_models?: string[]; budgets?: CreateBudgetRequest[]; rate_limit?: CreateRateLimitRequest; + model_budgets?: VirtualKeyModelBudgetRequest[]; key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config } @@ -184,6 +201,7 @@ export interface VirtualKeyProviderConfigUpdateRequest { blacklisted_models?: string[]; budgets?: CreateBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; + model_budgets?: VirtualKeyModelBudgetRequest[]; // Full desired per-model set when provider_configs is supplied key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config } diff --git a/ui/package-lock.json b/ui/package-lock.json index 30947bcb22d..91b148f4a0c 100644 --- a/ui/package-lock.json +++ b/ui/package-lock.json @@ -90,6 +90,9 @@ "vite": "8.0.16", "vite-plugin-monaco-editor": "1.1.0", "vitest": "4.1.6" + }, + "engines": { + "node": ">=22.12.0" } }, "node_modules/@alloc/quick-lru": { From 1e1f23a04583042b2fb4b85f2cf8a2b3c115c21b Mon Sep 17 00:00:00 2001 From: Anuj Parihar Date: Fri, 31 Jul 2026 01:32:17 +0530 Subject: [PATCH 2/2] feat: provider config card to show model budgets --- ui/components/ui/providerConfigCard.tsx | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/ui/components/ui/providerConfigCard.tsx b/ui/components/ui/providerConfigCard.tsx index 26462033d7c..11f656cc788 100644 --- a/ui/components/ui/providerConfigCard.tsx +++ b/ui/components/ui/providerConfigCard.tsx @@ -176,6 +176,17 @@ export function ProviderConfigCard({ const modelBudgets = value.modelBudgets || []; const capLabel = budgetLinesLabel(value.budgets); + // Header summary: the provider cap and/or a model-budget count, falling back to + // "No budget" only when neither is set โ€” so a provider with only model budgets + // doesn't read as "No budget". + const modelBudgetCount = showModelBudgets ? modelBudgets.length : 0; + const headerSummary = + [ + budgetLinesLabel(value.budgets, ""), + modelBudgetCount > 0 ? `${modelBudgetCount} model budget${modelBudgetCount === 1 ? "" : "s"}` : "", + ] + .filter(Boolean) + .join(" ยท ") || "No budget"; const ws = globalProviderCap; // Key scope handed to ModelMultiselect so model suggestions match the keys @@ -216,8 +227,9 @@ export function ProviderConfigCard({ > {providerLabel} - - {capLabel} + + {headerSummary} +