-
-
-3. Note the **Issuer URI** for your authorization server (e.g., `https://your-domain.okta.com/oauth2/default`)
-
-
-
-
-5. Click **Save**
-
-
-
-
-2. Under **OpenID Connect ID Token**, configure:
- - **Groups claim type**: Expression
- - **Groups claim expression**: `Arrays.flatten(Groups.startsWith("OKTA", "bifrost", 100))`
-
-
@@ -277,6 +211,22 @@ Role claims are available only when you configure custom claims on your authoriz
---
+## Step 7: Create API token for bulk user and team sync
+
+To create an API token, navigate to **Security** → **API** → **Tokens**.
+
+
+
+
+
+1. Click on "Create token"
+
+
+ 
+
+
+2. Copy token to be used in the next step.
+
## Step 8: Configure Bifrost
Now configure Bifrost to use Okta as the identity provider.
@@ -297,9 +247,9 @@ Now configure Bifrost to use Okta as the identity provider.
4. Toggle **Enabled** to activate the provider
5. Click **Save Configuration**
-### Group-to-Role Mappings (Optional)
+### Group-to-Role Mappings
-If you configured groups in Okta (Step 6), you can map Okta group names directly to Bifrost roles. This is an alternative to using custom role claims (Steps 4-5) and works with all Okta plans.
+If you configured groups in Okta (Step 5), you can map Okta group names directly to Bifrost roles. This is an alternative to using custom role claims (Steps 3-4) and works with all Okta plans.
1. In the User Provisioning configuration, scroll down to **Group-to-Role Mappings**
2. Click **Add Mapping**
diff --git a/docs/media/user-provisioning/okta-api-token-created.png b/docs/media/user-provisioning/okta-api-token-created.png
new file mode 100644
index 0000000000..e442519f8f
Binary files /dev/null and b/docs/media/user-provisioning/okta-api-token-created.png differ
diff --git a/docs/media/user-provisioning/okta-create-token-form.png b/docs/media/user-provisioning/okta-create-token-form.png
new file mode 100644
index 0000000000..2888d28da7
Binary files /dev/null and b/docs/media/user-provisioning/okta-create-token-form.png differ
diff --git a/docs/media/user-provisioning/okta-tokens-screen.png b/docs/media/user-provisioning/okta-tokens-screen.png
new file mode 100644
index 0000000000..6530a8a6d7
Binary files /dev/null and b/docs/media/user-provisioning/okta-tokens-screen.png differ
diff --git a/docs/media/user-provisioning/zitadel-add-role.png b/docs/media/user-provisioning/zitadel-add-role.png
deleted file mode 100644
index f00212f6d2..0000000000
Binary files a/docs/media/user-provisioning/zitadel-add-role.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-add-user-select-key.png b/docs/media/user-provisioning/zitadel-add-user-select-key.png
deleted file mode 100644
index ba8d8e52a9..0000000000
Binary files a/docs/media/user-provisioning/zitadel-add-user-select-key.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-create-app-auth-method.png b/docs/media/user-provisioning/zitadel-create-app-auth-method.png
deleted file mode 100644
index c27a6e5772..0000000000
Binary files a/docs/media/user-provisioning/zitadel-create-app-auth-method.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-create-app-namne.png b/docs/media/user-provisioning/zitadel-create-app-namne.png
deleted file mode 100644
index 7e220ce193..0000000000
Binary files a/docs/media/user-provisioning/zitadel-create-app-namne.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-create-app-uri.png b/docs/media/user-provisioning/zitadel-create-app-uri.png
deleted file mode 100644
index 8796e77ec5..0000000000
Binary files a/docs/media/user-provisioning/zitadel-create-app-uri.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-create-app.png b/docs/media/user-provisioning/zitadel-create-app.png
new file mode 100644
index 0000000000..0400316932
Binary files /dev/null and b/docs/media/user-provisioning/zitadel-create-app.png differ
diff --git a/docs/media/user-provisioning/zitadel-role-assignemnt.png b/docs/media/user-provisioning/zitadel-role-assignemnt.png
deleted file mode 100644
index 5f233eb436..0000000000
Binary files a/docs/media/user-provisioning/zitadel-role-assignemnt.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-select-project.png b/docs/media/user-provisioning/zitadel-select-project.png
deleted file mode 100644
index 824c48dd83..0000000000
Binary files a/docs/media/user-provisioning/zitadel-select-project.png and /dev/null differ
diff --git a/docs/media/user-provisioning/zitadel-token-config.png b/docs/media/user-provisioning/zitadel-token-config.png
deleted file mode 100644
index 1354a62830..0000000000
Binary files a/docs/media/user-provisioning/zitadel-token-config.png and /dev/null differ
diff --git a/docs/openapi/openapi.yaml b/docs/openapi/openapi.yaml
index cc125b63a7..8a816a3cc5 100644
--- a/docs/openapi/openapi.yaml
+++ b/docs/openapi/openapi.yaml
@@ -548,6 +548,12 @@ paths:
/api/pricing/force-sync:
$ref: './paths/management/config.yaml#/force-sync-pricing'
+ # Users
+ /api/users:
+ $ref: './paths/management/users.yaml#/users'
+ /api/users/{id}:
+ $ref: './paths/management/users.yaml#/users-by-id'
+
# Session
/api/session/login:
$ref: './paths/management/session.yaml#/login'
diff --git a/docs/openapi/paths/management/users.yaml b/docs/openapi/paths/management/users.yaml
new file mode 100644
index 0000000000..5c6df756e6
--- /dev/null
+++ b/docs/openapi/paths/management/users.yaml
@@ -0,0 +1,106 @@
+users:
+ get:
+ operationId: listUsers
+ summary: List users
+ description: Returns a paginated list of users with optional search.
+ tags:
+ - Users
+ parameters:
+ - name: page
+ in: query
+ description: Page number (1-based)
+ schema:
+ type: integer
+ minimum: 1
+ default: 1
+ - name: limit
+ in: query
+ description: Number of users per page (max 100)
+ schema:
+ type: integer
+ minimum: 1
+ maximum: 100
+ default: 20
+ - name: search
+ in: query
+ description: Search by name or email
+ schema:
+ type: string
+ responses:
+ '200':
+ description: Successful response
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/users.yaml#/ListUsersResponse'
+ '500':
+ $ref: '../../openapi.yaml#/components/responses/InternalError'
+
+ post:
+ operationId: createUser
+ summary: Create user
+ description: Manually creates a new user in the organization.
+ tags:
+ - Users
+ requestBody:
+ required: true
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/users.yaml#/CreateUserRequest'
+ responses:
+ '200':
+ description: User created successfully
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/users.yaml#/UserResponse'
+ '400':
+ $ref: '../../openapi.yaml#/components/responses/BadRequest'
+ '409':
+ description: User with this email already exists
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/common.yaml#/ErrorResponse'
+ '500':
+ $ref: '../../openapi.yaml#/components/responses/InternalError'
+
+users-by-id:
+ delete:
+ operationId: deleteUser
+ summary: Delete user
+ description: >
+ Permanently removes a user from the organization. This cascades to delete
+ the user's governance settings (budget/rate limits), team memberships,
+ access profiles, and OIDC sessions. Cannot delete yourself.
+ tags:
+ - Users
+ parameters:
+ - name: id
+ in: path
+ required: true
+ description: User ID
+ schema:
+ type: string
+ responses:
+ '200':
+ description: User deleted successfully
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/common.yaml#/MessageResponse'
+ '400':
+ description: Bad request (e.g. cannot delete yourself)
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/common.yaml#/ErrorResponse'
+ '404':
+ description: User not found
+ content:
+ application/json:
+ schema:
+ $ref: '../../schemas/management/common.yaml#/ErrorResponse'
+ '500':
+ $ref: '../../openapi.yaml#/components/responses/InternalError'
diff --git a/docs/openapi/schemas/management/users.yaml b/docs/openapi/schemas/management/users.yaml
new file mode 100644
index 0000000000..e57c550ad5
--- /dev/null
+++ b/docs/openapi/schemas/management/users.yaml
@@ -0,0 +1,82 @@
+UserObject:
+ type: object
+ properties:
+ id:
+ type: string
+ description: Unique user identifier
+ name:
+ type: string
+ description: User's display name
+ email:
+ type: string
+ format: email
+ description: User's email address
+ role_id:
+ type: integer
+ nullable: true
+ description: ID of the assigned RBAC role
+ role:
+ type: object
+ nullable: true
+ description: RBAC role details
+ properties:
+ id:
+ type: integer
+ name:
+ type: string
+ description:
+ type: string
+ is_system_role:
+ type: boolean
+ created_at:
+ type: string
+ format: date-time
+ updated_at:
+ type: string
+ format: date-time
+
+CreateUserRequest:
+ type: object
+ required:
+ - name
+ - email
+ properties:
+ name:
+ type: string
+ description: User's display name
+ email:
+ type: string
+ format: email
+ description: User's email address (must be unique)
+ role_id:
+ type: integer
+ description: Optional RBAC role ID to assign
+
+UserResponse:
+ type: object
+ properties:
+ user:
+ $ref: '#/UserObject'
+
+ListUsersResponse:
+ type: object
+ properties:
+ users:
+ type: array
+ items:
+ $ref: '#/UserObject'
+ total:
+ type: integer
+ description: Total number of users matching the query
+ page:
+ type: integer
+ description: Current page number
+ limit:
+ type: integer
+ description: Number of users per page
+ total_pages:
+ type: integer
+ description: Total number of pages
+ has_more:
+ type: boolean
+ description: Whether more pages are available
diff --git a/framework/configstore/clientconfig.go b/framework/configstore/clientconfig.go
index 7e2b3fe6d1..6fecc296a5 100644
--- a/framework/configstore/clientconfig.go
+++ b/framework/configstore/clientconfig.go
@@ -646,7 +646,6 @@ type VirtualKeyHashInput struct {
IsActive bool
TeamID *string
CustomerID *string
- BudgetID *string
RateLimitID *string
// ProviderConfigs and MCPConfigs are hashed separately as they contain nested data
ProviderConfigs []VirtualKeyProviderConfigHashInput
@@ -658,7 +657,6 @@ type VirtualKeyProviderConfigHashInput struct {
Provider string
Weight *float64
AllowedModels []string
- BudgetID *string
RateLimitID *string
KeyIDs []string // Only key IDs, not full key objects
}
@@ -694,10 +692,6 @@ func GenerateVirtualKeyHash(vk tables.TableVirtualKey) (string, error) {
if vk.CustomerID != nil {
hash.Write([]byte("customerID:" + *vk.CustomerID))
}
- // Hash BudgetID
- if vk.BudgetID != nil {
- hash.Write([]byte("budgetID:" + *vk.BudgetID))
- }
// Hash RateLimitID
if vk.RateLimitID != nil {
hash.Write([]byte("rateLimitID:" + *vk.RateLimitID))
@@ -711,16 +705,6 @@ func GenerateVirtualKeyHash(vk tables.TableVirtualKey) (string, error) {
if sortedProviderConfigs[i].Provider != sortedProviderConfigs[j].Provider {
return sortedProviderConfigs[i].Provider < sortedProviderConfigs[j].Provider
}
- bi, bj := "", ""
- if sortedProviderConfigs[i].BudgetID != nil {
- bi = *sortedProviderConfigs[i].BudgetID
- }
- if sortedProviderConfigs[j].BudgetID != nil {
- bj = *sortedProviderConfigs[j].BudgetID
- }
- if bi != bj {
- return bi < bj
- }
ri, rj := "", ""
if sortedProviderConfigs[i].RateLimitID != nil {
ri = *sortedProviderConfigs[i].RateLimitID
@@ -758,7 +742,6 @@ func GenerateVirtualKeyHash(vk tables.TableVirtualKey) (string, error) {
Provider: pc.Provider,
Weight: pc.Weight,
AllowedModels: sortedAllowedModels,
- BudgetID: pc.BudgetID,
RateLimitID: pc.RateLimitID,
KeyIDs: keyIDs,
}
diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go
index d9592fdd3e..995cad189d 100644
--- a/framework/configstore/migrations.go
+++ b/framework/configstore/migrations.go
@@ -364,6 +364,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
if err := migrationAddOllamaSGLConfigColumns(ctx, db); err != nil {
return err
}
+ if err := migrationAddMultiBudgetTables(ctx, db); err != nil {
+ return err
+ }
return nil
}
@@ -939,8 +942,8 @@ func migrationCleanupMCPClientToolsConfig(ctx context.Context, db *gorm.DB) erro
// Step 2: Update empty ToolsToExecuteJSON arrays to wildcard ["*"]
// Convert "[]" (empty array) to "[\"*\"]" (wildcard array) for backward compatibility
updateSQL := `
- UPDATE config_mcp_clients
- SET tools_to_execute_json = '["*"]'
+ UPDATE config_mcp_clients
+ SET tools_to_execute_json = '["*"]'
WHERE tools_to_execute_json = '[]' OR tools_to_execute_json = '' OR tools_to_execute_json IS NULL
`
if err := tx.Exec(updateSQL).Error; err != nil {
@@ -955,8 +958,8 @@ func migrationCleanupMCPClientToolsConfig(ctx context.Context, db *gorm.DB) erro
tx = tx.WithContext(ctx)
revertSQL := `
- UPDATE config_mcp_clients
- SET tools_to_execute_json = '[]'
+ UPDATE config_mcp_clients
+ SET tools_to_execute_json = '[]'
WHERE tools_to_execute_json = '["*"]'
`
if err := tx.Exec(revertSQL).Error; err != nil {
@@ -1011,12 +1014,13 @@ func migrationAddProviderConfigBudgetRateLimit(ctx context.Context, db *gorm.DB)
tx = tx.WithContext(ctx)
migrator := tx.Migrator()
- // Add BudgetID column if it doesn't exist
+ // Add budget_id and rate_limit_id columns if they don't exist
+ // Note: budget_id is added via raw SQL because the field was later removed from the struct
+ // (migrated to governance_budgets.provider_config_id in add_multi_budget_tables)
if migrator.HasTable(&tables.TableVirtualKeyProviderConfig{}) {
- if !migrator.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") {
- if err := migrator.AddColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id"); err != nil {
- return fmt.Errorf("failed to add budget_id column: %w", err)
- }
+ if err := tx.Exec("ALTER TABLE governance_virtual_key_provider_configs ADD COLUMN IF NOT EXISTS budget_id VARCHAR(255)").Error; err != nil {
+ // Ignore error for databases that don't support IF NOT EXISTS (e.g., SQLite)
+ // The column may already exist from a previous run
}
// Add RateLimitID column if it doesn't exist
@@ -1027,10 +1031,8 @@ func migrationAddProviderConfigBudgetRateLimit(ctx context.Context, db *gorm.DB)
}
// Create foreign key indexes for better performance
- if !migrator.HasIndex(&tables.TableVirtualKeyProviderConfig{}, "idx_provider_config_budget") {
- if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_config_budget ON governance_virtual_key_provider_configs (budget_id)").Error; err != nil {
- return fmt.Errorf("failed to create budget_id index: %w", err)
- }
+ if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_config_budget ON governance_virtual_key_provider_configs (budget_id)").Error; err != nil {
+ // Ignore - index may already exist or column may not exist yet
}
if !migrator.HasIndex(&tables.TableVirtualKeyProviderConfig{}, "idx_provider_config_rate_limit") {
@@ -1039,12 +1041,7 @@ func migrationAddProviderConfigBudgetRateLimit(ctx context.Context, db *gorm.DB)
}
}
- // Create FK constraints (dialect‑agnostic)
- if !migrator.HasConstraint(&tables.TableVirtualKeyProviderConfig{}, "Budget") {
- if err := migrator.CreateConstraint(&tables.TableVirtualKeyProviderConfig{}, "Budget"); err != nil {
- return fmt.Errorf("failed to create Budget FK constraint: %w", err)
- }
- }
+ // Create FK constraint for RateLimit (Budget FK is no longer needed - budgets use direct FK on budget table)
if !migrator.HasConstraint(&tables.TableVirtualKeyProviderConfig{}, "RateLimit") {
if err := migrator.CreateConstraint(&tables.TableVirtualKeyProviderConfig{}, "RateLimit"); err != nil {
return fmt.Errorf("failed to create RateLimit FK constraint: %w", err)
@@ -1058,32 +1055,19 @@ func migrationAddProviderConfigBudgetRateLimit(ctx context.Context, db *gorm.DB)
tx = tx.WithContext(ctx)
migrator := tx.Migrator()
- // Drop indexes first
- if err := tx.Exec("DROP INDEX IF EXISTS idx_provider_config_budget").Error; err != nil {
- return fmt.Errorf("failed to drop budget_id index: %w", err)
- }
- if err := tx.Exec("DROP INDEX IF EXISTS idx_provider_config_rate_limit").Error; err != nil {
- return fmt.Errorf("failed to drop rate_limit_id index: %w", err)
- }
+ // Drop indexes
+ _ = tx.Exec("DROP INDEX IF EXISTS idx_provider_config_budget")
+ _ = tx.Exec("DROP INDEX IF EXISTS idx_provider_config_rate_limit")
// Drop FK constraints
- if migrator.HasConstraint(&tables.TableVirtualKeyProviderConfig{}, "Budget") {
- if err := migrator.DropConstraint(&tables.TableVirtualKeyProviderConfig{}, "Budget"); err != nil {
- return fmt.Errorf("failed to drop Budget FK constraint: %w", err)
- }
- }
if migrator.HasConstraint(&tables.TableVirtualKeyProviderConfig{}, "RateLimit") {
if err := migrator.DropConstraint(&tables.TableVirtualKeyProviderConfig{}, "RateLimit"); err != nil {
return fmt.Errorf("failed to drop RateLimit FK constraint: %w", err)
}
}
- // Drop columns
- if migrator.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") {
- if err := migrator.DropColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id"); err != nil {
- return fmt.Errorf("failed to drop budget_id column: %w", err)
- }
- }
+ // Drop columns via raw SQL (budget_id no longer on struct)
+ _ = tx.Exec("ALTER TABLE governance_virtual_key_provider_configs DROP COLUMN IF EXISTS budget_id")
if migrator.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "rate_limit_id") {
if err := migrator.DropColumn(&tables.TableVirtualKeyProviderConfig{}, "rate_limit_id"); err != nil {
return fmt.Errorf("failed to drop rate_limit_id column: %w", err)
@@ -5338,6 +5322,7 @@ func migrationMakeBasePricingColumnsNullable(ctx context.Context, db *gorm.DB) e
return nil
}
+// migrationAddAllowOnAllVirtualKeysColumn adds the allow_on_all_virtual_keys column to the mcp_client table
func migrationAddAllowOnAllVirtualKeysColumn(ctx context.Context, db *gorm.DB) error {
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
ID: "add_allow_on_all_virtual_keys_column",
@@ -5561,30 +5546,14 @@ func migrationAddReplicateKeyConfigColumn(ctx context.Context, db *gorm.DB) erro
return nil
}
-// migrationAddBudgetCalendarAlignedColumn adds the calendar_aligned column to the governance_budgets table.
+// migrationAddBudgetCalendarAlignedColumn was originally for adding calendar_aligned to governance_budgets.
+// Calendar alignment is now a VK-level field (governance_virtual_keys.calendar_aligned) added in migrationAddMultiBudgetTables.
+// This migration is kept as a no-op so the migrator doesn't try to re-run it.
func migrationAddBudgetCalendarAlignedColumn(ctx context.Context, db *gorm.DB) error {
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
- ID: "add_budget_calendar_aligned_column",
- Migrate: func(tx *gorm.DB) error {
- tx = tx.WithContext(ctx)
- mg := tx.Migrator()
- if !mg.HasColumn(&tables.TableBudget{}, "calendar_aligned") {
- if err := mg.AddColumn(&tables.TableBudget{}, "calendar_aligned"); err != nil {
- return fmt.Errorf("failed to add calendar_aligned column: %w", err)
- }
- }
- return nil
- },
- Rollback: func(tx *gorm.DB) error {
- tx = tx.WithContext(ctx)
- mg := tx.Migrator()
- if mg.HasColumn(&tables.TableBudget{}, "calendar_aligned") {
- if err := mg.DropColumn(&tables.TableBudget{}, "calendar_aligned"); err != nil {
- return fmt.Errorf("failed to drop calendar_aligned column: %w", err)
- }
- }
- return nil
- },
+ ID: "add_budget_calendar_aligned_column",
+ Migrate: func(tx *gorm.DB) error { return nil },
+ Rollback: func(tx *gorm.DB) error { return nil },
}})
if err := m.Migrate(); err != nil {
return fmt.Errorf("error running add_budget_calendar_aligned_column migration: %s", err.Error())
@@ -5815,3 +5784,128 @@ func migrationAddOllamaSGLConfigColumns(ctx context.Context, db *gorm.DB) error
}
return nil
}
+
+// migrationAddMultiBudgetTables creates junction tables for multi-budget support and backfills existing data.
+func migrationAddMultiBudgetTables(ctx context.Context, db *gorm.DB) error {
+ m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
+ ID: "add_multi_budget_tables",
+ Migrate: func(tx *gorm.DB) error {
+ tx = tx.WithContext(ctx)
+ mg := tx.Migrator()
+
+ // Add calendar_aligned to governance_virtual_keys (VK-level setting)
+ if !mg.HasColumn(&tables.TableVirtualKey{}, "calendar_aligned") {
+ if err := mg.AddColumn(&tables.TableVirtualKey{}, "CalendarAligned"); err != nil {
+ return fmt.Errorf("failed to add calendar_aligned column to governance_virtual_keys: %w", err)
+ }
+ }
+
+ // Add FK columns on governance_budgets for multi-budget ownership
+ if !mg.HasColumn(&tables.TableBudget{}, "virtual_key_id") {
+ if err := mg.AddColumn(&tables.TableBudget{}, "VirtualKeyID"); err != nil {
+ return fmt.Errorf("failed to add virtual_key_id column to governance_budgets: %w", err)
+ }
+ }
+ if !mg.HasColumn(&tables.TableBudget{}, "provider_config_id") {
+ if err := mg.AddColumn(&tables.TableBudget{}, "ProviderConfigID"); err != nil {
+ return fmt.Errorf("failed to add provider_config_id column to governance_budgets: %w", err)
+ }
+ }
+
+ // Create indexes on the new FK columns (AddColumn doesn't create indexes from struct tags)
+ if !mg.HasIndex(&tables.TableBudget{}, "idx_governance_budgets_virtual_key_id") {
+ if err := mg.CreateIndex(&tables.TableBudget{}, "VirtualKeyID"); err != nil {
+ return fmt.Errorf("failed to create index on governance_budgets.virtual_key_id: %w", err)
+ }
+ }
+ if !mg.HasIndex(&tables.TableBudget{}, "idx_governance_budgets_provider_config_id") {
+ if err := mg.CreateIndex(&tables.TableBudget{}, "ProviderConfigID"); err != nil {
+ return fmt.Errorf("failed to create index on governance_budgets.provider_config_id: %w", err)
+ }
+ }
+
+ // Create FK constraints with CASCADE delete (defined on parent structs)
+ if !mg.HasConstraint(&tables.TableVirtualKey{}, "Budgets") {
+ if err := mg.CreateConstraint(&tables.TableVirtualKey{}, "Budgets"); err != nil {
+ return fmt.Errorf("failed to create FK constraint for VirtualKey -> Budgets: %w", err)
+ }
+ }
+ if !mg.HasConstraint(&tables.TableVirtualKeyProviderConfig{}, "Budgets") {
+ if err := mg.CreateConstraint(&tables.TableVirtualKeyProviderConfig{}, "Budgets"); err != nil {
+ return fmt.Errorf("failed to create FK constraint for ProviderConfig -> Budgets: %w", err)
+ }
+ }
+
+ // Backfill: set virtual_key_id from legacy VK budget_id (if column still exists)
+ if mg.HasColumn(&tables.TableVirtualKey{}, "budget_id") {
+ if err := tx.Exec(`
+ UPDATE governance_budgets SET virtual_key_id = (
+ SELECT id FROM governance_virtual_keys
+ WHERE governance_virtual_keys.budget_id = governance_budgets.id
+ ) WHERE virtual_key_id IS NULL AND EXISTS (
+ SELECT 1 FROM governance_virtual_keys
+ WHERE governance_virtual_keys.budget_id = governance_budgets.id
+ )
+ `).Error; err != nil {
+ return fmt.Errorf("failed to backfill VK budget virtual_key_id: %w", err)
+ }
+ }
+
+ // Backfill: set provider_config_id from legacy PC budget_id (if column still exists)
+ if mg.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") {
+ if err := tx.Exec(`
+ UPDATE governance_budgets SET provider_config_id = (
+ SELECT id FROM governance_virtual_key_provider_configs
+ WHERE governance_virtual_key_provider_configs.budget_id = governance_budgets.id
+ ) WHERE provider_config_id IS NULL AND EXISTS (
+ SELECT 1 FROM governance_virtual_key_provider_configs
+ WHERE governance_virtual_key_provider_configs.budget_id = governance_budgets.id
+ )
+ `).Error; err != nil {
+ return fmt.Errorf("failed to backfill PC budget provider_config_id: %w", err)
+ }
+ }
+
+ // Backfill: copy calendar_aligned from legacy budget column to VK-level field
+ // (governance_budgets.calendar_aligned was added by add_budget_calendar_aligned_column on main)
+ if mg.HasColumn(&tables.TableBudget{}, "calendar_aligned") {
+ if err := tx.Exec(`
+ UPDATE governance_virtual_keys SET calendar_aligned = true
+ WHERE id IN (
+ SELECT DISTINCT virtual_key_id FROM governance_budgets
+ WHERE calendar_aligned = true AND virtual_key_id IS NOT NULL
+ ) AND calendar_aligned = false
+ `).Error; err != nil {
+ return fmt.Errorf("failed to backfill calendar_aligned from budgets to virtual keys: %w", err)
+ }
+ // Drop the legacy calendar_aligned column from governance_budgets
+ _ = tx.Exec("ALTER TABLE governance_budgets DROP COLUMN IF EXISTS calendar_aligned")
+ }
+
+ // Drop legacy budget_id columns from VK and ProviderConfig (raw SQL to avoid GORM FK lookup issues)
+ _ = tx.Exec("ALTER TABLE governance_virtual_keys DROP COLUMN IF EXISTS budget_id")
+ _ = tx.Exec("ALTER TABLE governance_virtual_key_provider_configs DROP COLUMN IF EXISTS budget_id")
+
+ return nil
+ },
+ Rollback: func(tx *gorm.DB) error {
+ tx = tx.WithContext(ctx)
+ mg := tx.Migrator()
+ if mg.HasColumn(&tables.TableBudget{}, "virtual_key_id") {
+ if err := mg.DropColumn(&tables.TableBudget{}, "virtual_key_id"); err != nil {
+ return err
+ }
+ }
+ if mg.HasColumn(&tables.TableBudget{}, "provider_config_id") {
+ if err := mg.DropColumn(&tables.TableBudget{}, "provider_config_id"); err != nil {
+ return err
+ }
+ }
+ return nil
+ },
+ }})
+ if err := m.Migrate(); err != nil {
+ return fmt.Errorf("error running add_multi_budget_tables migration: %s", err.Error())
+ }
+ return nil
+}
diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go
index ade838add8..f3b5d45b68 100644
--- a/framework/configstore/rdb.go
+++ b/framework/configstore/rdb.go
@@ -35,6 +35,7 @@ func getWeight(w *float64) float64 {
return *w
}
+// schemaKeyFromTableKey converts a database key to a schema key.
func schemaKeyFromTableKey(dbKey tables.TableKey) schemas.Key {
return schemas.Key{
ID: dbKey.KeyID,
@@ -59,6 +60,7 @@ func schemaKeyFromTableKey(dbKey tables.TableKey) schemas.Key {
}
}
+// tableKeyFromSchemaKey converts a schema key to a database key.
func tableKeyFromSchemaKey(provider tables.TableProvider, key schemas.Key) (tables.TableKey, error) {
dbKey := tables.TableKey{
Provider: provider.Name,
@@ -175,13 +177,10 @@ func (s *RDBConfigStore) parseGormError(err error) error {
if err == nil {
return nil
}
-
if errors.Is(err, gorm.ErrRecordNotFound) {
return ErrNotFound
}
-
errMsg := err.Error()
-
// Check for unique constraint violations
// SQLite format: "UNIQUE constraint failed: table_name.column_name"
// PostgreSQL format: "ERROR: duplicate key value violates unique constraint"
@@ -914,7 +913,6 @@ func (s *RDBConfigStore) CreateProviderKey(ctx context.Context, provider schemas
} else {
txDB = s.db
}
-
var dbProvider tables.TableProvider
if err := txDB.WithContext(ctx).Where("name = ?", string(provider)).First(&dbProvider).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -922,16 +920,13 @@ func (s *RDBConfigStore) CreateProviderKey(ctx context.Context, provider schemas
}
return err
}
-
dbKey, err := tableKeyFromSchemaKey(dbProvider, key)
if err != nil {
return err
}
-
if err := txDB.WithContext(ctx).Create(&dbKey).Error; err != nil {
return s.parseGormError(err)
}
-
return nil
}
@@ -1787,35 +1782,31 @@ func (s *RDBConfigStore) UpdatePlugin(ctx context.Context, plugin *tables.TableP
txDB = s.db.Begin()
localTx = true
}
-
// Mark plugin as custom if path is not empty
if plugin.Path != nil && strings.TrimSpace(*plugin.Path) != "" {
plugin.IsCustom = true
} else {
plugin.IsCustom = false
}
-
if err := txDB.WithContext(ctx).Delete(&tables.TablePlugin{}, "name = ?", plugin.Name).Error; err != nil {
if localTx {
txDB.Rollback()
}
return err
}
-
if err := txDB.WithContext(ctx).Create(plugin).Error; err != nil {
if localTx {
txDB.Rollback()
}
return s.parseGormError(err)
}
-
if localTx {
return txDB.Commit().Error
}
-
return nil
}
+// DeletePlugin deletes a plugin from the database.
func (s *RDBConfigStore) DeletePlugin(ctx context.Context, name string, tx ...*gorm.DB) error {
var txDB *gorm.DB
if len(tx) > 0 {
@@ -1828,6 +1819,7 @@ func (s *RDBConfigStore) DeletePlugin(ctx context.Context, name string, tx ...*g
// GOVERNANCE METHODS
+// GetRedactedVirtualKeys retrieves redacted virtual keys from the database.
func (s *RDBConfigStore) GetRedactedVirtualKeys(ctx context.Context, ids []string) ([]tables.TableVirtualKey, error) {
var virtualKeys []tables.TableVirtualKey
@@ -1845,6 +1837,7 @@ func (s *RDBConfigStore) GetRedactedVirtualKeys(ctx context.Context, ids []strin
return virtualKeys, nil
}
+// preloadCustomerRelations preloads the customer relations for a virtual key.
func preloadCustomerRelations(db *gorm.DB, prefix string) *gorm.DB {
relation := func(name string) string {
if prefix == "" {
@@ -1852,7 +1845,6 @@ func preloadCustomerRelations(db *gorm.DB, prefix string) *gorm.DB {
}
return prefix + name
}
-
return db.
Preload(relation("Teams")).
Preload(relation("Budget")).
@@ -1860,16 +1852,16 @@ func preloadCustomerRelations(db *gorm.DB, prefix string) *gorm.DB {
Preload(relation("VirtualKeys"))
}
+// preloadVirtualKeyBaseRelations preloads the base relationships for a virtual key.
func preloadVirtualKeyBaseRelations(db *gorm.DB) *gorm.DB {
- db = db.Preload("Team").Preload("Team.Customer")
-
- db = db.Preload("Customer")
-
return db.
- Preload("Budget").
+ Preload("Team").
+ Preload("Team.Customer").
+ Preload("Customer").
+ Preload("Budgets").
Preload("RateLimit").
Preload("ProviderConfigs").
- Preload("ProviderConfigs.Budget").
+ Preload("ProviderConfigs.Budgets").
Preload("ProviderConfigs.RateLimit").
Preload("ProviderConfigs.Keys", func(db *gorm.DB) *gorm.DB {
return db.Select("id, name, key_id, models_json, provider")
@@ -1878,6 +1870,7 @@ func preloadVirtualKeyBaseRelations(db *gorm.DB) *gorm.DB {
Preload("MCPConfigs.MCPClient")
}
+// preloadVirtualKeyDetailRelations preloads the detail relationships for a virtual key.
func preloadVirtualKeyDetailRelations(db *gorm.DB) *gorm.DB {
return preloadCustomerRelations(preloadVirtualKeyBaseRelations(db), "Customer.")
}
@@ -1964,7 +1957,6 @@ func (s *RDBConfigStore) GetVirtualKeyByValue(ctx context.Context, value string)
valueHash := encrypt.HashSHA256(value)
var virtualKey tables.TableVirtualKey
query := preloadVirtualKeyBaseRelations(s.db.WithContext(ctx))
-
// Use hash-based lookup if hash column is populated, fall back to plaintext for backward compat
if err := query.Where("value_hash = ?", valueHash).First(&virtualKey).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -1982,6 +1974,7 @@ func (s *RDBConfigStore) GetVirtualKeyByValue(ctx context.Context, value string)
return &virtualKey, nil
}
+// CreateVirtualKey creates a new virtual key in the database.
func (s *RDBConfigStore) CreateVirtualKey(ctx context.Context, virtualKey *tables.TableVirtualKey, tx ...*gorm.DB) error {
var txDB *gorm.DB
if len(tx) > 0 {
@@ -1995,6 +1988,7 @@ func (s *RDBConfigStore) CreateVirtualKey(ctx context.Context, virtualKey *table
return nil
}
+// UpdateVirtualKey updates an existing virtual key in the database.
func (s *RDBConfigStore) UpdateVirtualKey(ctx context.Context, virtualKey *tables.TableVirtualKey, tx ...*gorm.DB) error {
var txDB *gorm.DB
if len(tx) > 0 {
@@ -2095,33 +2089,26 @@ func (s *RDBConfigStore) DeleteVirtualKey(ctx context.Context, id string) error
return err
}
- // Collect budget and rate limit IDs from provider configs before deletion
- var providerConfigBudgetIDs []string
+ // Delete provider config resources before deleting the configs themselves
var providerConfigRateLimitIDs []string
for _, pc := range virtualKey.ProviderConfigs {
// Delete the keys join table entries
if err := tx.WithContext(ctx).Exec("DELETE FROM governance_virtual_key_provider_config_keys WHERE table_virtual_key_provider_config_id = ?", pc.ID).Error; err != nil {
return err
}
- // Collect budget and rate limit IDs for deletion after provider config
- if pc.BudgetID != nil {
- providerConfigBudgetIDs = append(providerConfigBudgetIDs, *pc.BudgetID)
+ // Delete budgets owned by this provider config
+ if err := tx.WithContext(ctx).Where("provider_config_id = ?", pc.ID).Delete(&tables.TableBudget{}).Error; err != nil {
+ return err
}
if pc.RateLimitID != nil {
providerConfigRateLimitIDs = append(providerConfigRateLimitIDs, *pc.RateLimitID)
}
}
- // Delete all provider configs associated with the virtual key first
+ // Delete all provider configs associated with the virtual key
if err := tx.WithContext(ctx).Delete(&tables.TableVirtualKeyProviderConfig{}, "virtual_key_id = ?", id).Error; err != nil {
return err
}
- // Now delete the collected budgets and rate limits
- for _, budgetID := range providerConfigBudgetIDs {
- if err := tx.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", budgetID).Error; err != nil {
- return err
- }
- }
for _, rateLimitID := range providerConfigRateLimitIDs {
if err := tx.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", rateLimitID).Error; err != nil {
return err
@@ -2131,8 +2118,10 @@ func (s *RDBConfigStore) DeleteVirtualKey(ctx context.Context, id string) error
if err := tx.WithContext(ctx).Delete(&tables.TableVirtualKeyMCPConfig{}, "virtual_key_id = ?", id).Error; err != nil {
return err
}
- // Delete the budget associated with the virtual key
- budgetID := virtualKey.BudgetID
+ // Delete budgets owned by this virtual key
+ if err := tx.WithContext(ctx).Where("virtual_key_id = ?", id).Delete(&tables.TableBudget{}).Error; err != nil {
+ return err
+ }
rateLimitID := virtualKey.RateLimitID
// Delete the virtual key
if err := tx.WithContext(ctx).Delete(&tables.TableVirtualKey{}, "id = ?", id).Error; err != nil {
@@ -2141,11 +2130,6 @@ func (s *RDBConfigStore) DeleteVirtualKey(ctx context.Context, id string) error
}
return err
}
- if budgetID != nil {
- if err := tx.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", *budgetID).Error; err != nil {
- return err
- }
- }
// Delete the rate limit associated with the virtual key
if rateLimitID != nil {
if err := tx.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil {
@@ -2343,18 +2327,15 @@ func (s *RDBConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id
}
return err
}
- // Store the budget and rate limit IDs before deleting
- budgetID := providerConfig.BudgetID
+ // Store the rate limit ID before deleting
rateLimitID := providerConfig.RateLimitID
- // Delete the provider config first
- if err := txDB.WithContext(ctx).Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id).Error; err != nil {
+ // Delete budgets owned by this provider config
+ if err := txDB.WithContext(ctx).Where("provider_config_id = ?", id).Delete(&tables.TableBudget{}).Error; err != nil {
return err
}
- // Delete the budget if it exists
- if budgetID != nil {
- if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", *budgetID).Error; err != nil {
- return err
- }
+ // Delete the provider config
+ if err := txDB.WithContext(ctx).Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id).Error; err != nil {
+ return err
}
// Delete the rate limit if it exists
if rateLimitID != nil {
@@ -3178,6 +3159,7 @@ func (s *RDBConfigStore) GetModelConfigs(ctx context.Context) ([]tables.TableMod
return modelConfigs, nil
}
+// GetModelConfigsPaginated retrieves model configs with pagination, filtering, and search support.
func (s *RDBConfigStore) GetModelConfigsPaginated(ctx context.Context, params ModelConfigsQueryParams) ([]tables.TableModelConfig, int64, error) {
baseQuery := s.db.WithContext(ctx).Model(&tables.TableModelConfig{})
diff --git a/framework/configstore/rdb_test.go b/framework/configstore/rdb_test.go
index 796539a3a9..406e7a5cfd 100644
--- a/framework/configstore/rdb_test.go
+++ b/framework/configstore/rdb_test.go
@@ -520,24 +520,28 @@ func TestCreateVirtualKey_WithBudgetAndRateLimit(t *testing.T) {
require.NoError(t, err)
// Create virtual key with references
- budgetID := "budget-for-vk"
rateLimitID := "rate-limit-for-vk"
+ vkID := "vk-with-refs"
vk := &tables.TableVirtualKey{
- ID: "vk-with-refs",
+ ID: vkID,
Name: "VK With References",
Value: "vk-refs-value",
IsActive: true,
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
}
err = store.CreateVirtualKey(ctx, vk)
require.NoError(t, err)
+ // Link the existing budget to the VK via FK
+ budget.VirtualKeyID = &vkID
+ err = store.UpdateBudget(ctx, budget)
+ require.NoError(t, err)
+
result, err := store.GetVirtualKey(ctx, "vk-with-refs")
require.NoError(t, err)
- assert.NotNil(t, result.BudgetID)
- assert.Equal(t, "budget-for-vk", *result.BudgetID)
+ assert.Len(t, result.Budgets, 1)
+ assert.Equal(t, "budget-for-vk", result.Budgets[0].ID)
assert.NotNil(t, result.RateLimitID)
assert.Equal(t, "rate-limit-for-vk", *result.RateLimitID)
}
@@ -999,19 +1003,23 @@ func TestFullVirtualKeyFlow(t *testing.T) {
require.NoError(t, err)
// Step 4: Create virtual key
- budgetID := "integration-budget"
rateLimitID := "integration-rate-limit"
+ integrationVKID := "integration-vk"
vk := &tables.TableVirtualKey{
- ID: "integration-vk",
+ ID: integrationVKID,
Name: "Integration Virtual Key",
Value: "vk-integration-xyz",
IsActive: true,
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
}
err = store.CreateVirtualKey(ctx, vk)
require.NoError(t, err)
+ // Link the existing budget to the VK via FK
+ budget.VirtualKeyID = &integrationVKID
+ err = store.UpdateBudget(ctx, budget)
+ require.NoError(t, err)
+
// Step 5: Create provider config with key reference
weight := 1.0
pc := &tables.TableVirtualKeyProviderConfig{
@@ -1029,7 +1037,7 @@ func TestFullVirtualKeyFlow(t *testing.T) {
result, err := store.GetVirtualKey(ctx, "integration-vk")
require.NoError(t, err)
assert.Equal(t, "Integration Virtual Key", result.Name)
- assert.NotNil(t, result.BudgetID)
+ assert.Len(t, result.Budgets, 1)
assert.NotNil(t, result.RateLimitID)
configs, err := store.GetVirtualKeyProviderConfigs(ctx, "integration-vk")
diff --git a/framework/configstore/tables/budget.go b/framework/configstore/tables/budget.go
index e35c530f3d..2d7d397d26 100644
--- a/framework/configstore/tables/budget.go
+++ b/framework/configstore/tables/budget.go
@@ -15,9 +15,9 @@ type TableBudget struct {
LastReset time.Time `gorm:"index" json:"last_reset"` // Last time budget was reset
CurrentUsage float64 `gorm:"default:0" json:"current_usage"` // Current usage in dollars
- // CalendarAligned snaps LastReset to the start of the current calendar period (day, week, month, year)
- // instead of the exact creation/update time, so budgets reset at clean calendar boundaries.
- CalendarAligned bool `gorm:"default:false" json:"calendar_aligned"`
+ // Owner FKs: a budget belongs to at most one VK or one ProviderConfig
+ VirtualKeyID *string `gorm:"type:varchar(255);index" json:"virtual_key_id,omitempty"`
+ ProviderConfigID *uint `gorm:"index" json:"provider_config_id,omitempty"`
// Config hash is used to detect the changes synced from config.json file
// Every time we sync the config.json file, we will update the config hash
@@ -32,6 +32,11 @@ func (TableBudget) TableName() string { return "governance_budgets" }
// BeforeSave hook for Budget to validate reset duration format and max limit
func (b *TableBudget) BeforeSave(tx *gorm.DB) error {
+ // A budget belongs to at most one owner type
+ if b.VirtualKeyID != nil && b.ProviderConfigID != nil {
+ return fmt.Errorf("budget cannot belong to both a virtual key and a provider config")
+ }
+
// Validate that ResetDuration is in correct format (e.g., "30s", "5m", "1h", "1d", "1w", "1M", "1Y")
if d, err := ParseDuration(b.ResetDuration); err != nil {
return fmt.Errorf("invalid reset duration format: %s", b.ResetDuration)
diff --git a/framework/configstore/tables/virtualkey.go b/framework/configstore/tables/virtualkey.go
index d7b1c92a6a..8c70d2e4bf 100644
--- a/framework/configstore/tables/virtualkey.go
+++ b/framework/configstore/tables/virtualkey.go
@@ -30,13 +30,12 @@ type TableVirtualKeyProviderConfig struct {
Weight *float64 `json:"weight"`
AllowedModels schemas.WhiteList `gorm:"type:text;serializer:json" json:"allowed_models"` // ["*"] allows all models; empty denies all (deny-by-default)
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"`
+ 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"` // Used when AllowAllKeys is false; empty means no keys allowed
+ 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
}
// TableName sets the table name for each model
@@ -211,16 +210,16 @@ type TableVirtualKey struct {
MCPConfigs []TableVirtualKeyMCPConfig `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"mcp_configs"`
// Foreign key relationships (mutually exclusive: either TeamID or CustomerID, not both)
- TeamID *string `gorm:"type:varchar(255);index" json:"team_id,omitempty"`
- CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"`
- BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"`
- RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`
+ TeamID *string `gorm:"type:varchar(255);index" json:"team_id,omitempty"`
+ CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"`
+ RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`
+ CalendarAligned bool `gorm:"default:false" json:"calendar_aligned"` // When true, all budgets under this VK reset at clean calendar boundaries
// Relationships
Team *TableTeam `gorm:"foreignKey:TeamID" json:"team,omitempty"`
Customer *TableCustomer `gorm:"foreignKey:CustomerID" json:"customer,omitempty"`
- Budget *TableBudget `gorm:"foreignKey:BudgetID;onDelete:CASCADE" json:"budget,omitempty"`
RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"`
+ Budgets []TableBudget `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` // Multiple budgets with different reset intervals
// Config hash is used to detect the changes synced from config.json file
// Every time we sync the config.json file, we will update the config hash
diff --git a/plugins/governance/model_provider_governance_test.go b/plugins/governance/modelprovidergovernance_test.go
similarity index 99%
rename from plugins/governance/model_provider_governance_test.go
rename to plugins/governance/modelprovidergovernance_test.go
index e363c15435..a7889f14bd 100644
--- a/plugins/governance/model_provider_governance_test.go
+++ b/plugins/governance/modelprovidergovernance_test.go
@@ -1479,9 +1479,10 @@ func TestPreLLMHook_ModelProviderPass_VirtualKeyChecksPass(t *testing.T) {
logger := NewMockLogger()
// Model/provider checks pass (no limits)
// Virtual key checks also pass
- vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK", []configstoreTables.TableVirtualKeyProviderConfig{
+ vk := buildVirtualKey("vk1", "sk-bf-test", "Test VK", true)
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
buildProviderConfig("openai", []string{"*"}),
- })
+ }
store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
}, nil)
diff --git a/plugins/governance/resolver.go b/plugins/governance/resolver.go
index 65607ed86e..bf184908d3 100644
--- a/plugins/governance/resolver.go
+++ b/plugins/governance/resolver.go
@@ -353,7 +353,7 @@ func (r *BudgetResolver) isProviderBudgetViolated(ctx context.Context, vk *confi
}
// 2. Check VK-level provider config budget
- if config.Budget == nil {
+ if len(config.Budgets) == 0 {
return false
}
if err := r.store.CheckBudget(ctx, vk, request, nil); err != nil {
diff --git a/plugins/governance/resolver_test.go b/plugins/governance/resolver_test.go
index 4d06c29a82..fceb69a737 100644
--- a/plugins/governance/resolver_test.go
+++ b/plugins/governance/resolver_test.go
@@ -16,9 +16,10 @@ import (
// TestBudgetResolver_EvaluateRequest_AllowedRequest tests happy path
func TestBudgetResolver_EvaluateRequest_AllowedRequest(t *testing.T) {
logger := NewMockLogger()
- vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK", []configstoreTables.TableVirtualKeyProviderConfig{
+ vk := buildVirtualKey("vk1", "sk-bf-test", "Test VK", true)
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
buildProviderConfig("openai", []string{"*"}),
- })
+ }
store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
@@ -102,7 +103,6 @@ func TestBudgetResolver_EvaluateRequest_ModelBlocked(t *testing.T) {
AllowedModels: []string{"gpt-4", "gpt-4-turbo"}, // Only these models
Weight: bifrost.Ptr(1.0),
RateLimit: nil,
- Budget: nil,
Keys: []configstoreTables.TableKey{},
},
}
@@ -469,9 +469,10 @@ func TestBudgetResolver_IsModelAllowed(t *testing.T) {
// TestBudgetResolver_ContextPopulation tests context values are set correctly
func TestBudgetResolver_ContextPopulation(t *testing.T) {
logger := NewMockLogger()
- vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK", []configstoreTables.TableVirtualKeyProviderConfig{
+ vk := buildVirtualKey("vk1", "sk-bf-test", "Test VK", true)
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
buildProviderConfig("openai", []string{"*"}),
- })
+ }
customer := buildCustomer("cust1", "Customer 1", nil)
team := buildTeam("team1", "Team 1", nil)
team.CustomerID = &customer.ID
diff --git a/plugins/governance/store.go b/plugins/governance/store.go
index 6158cc1302..49086dc9e4 100644
--- a/plugins/governance/store.go
+++ b/plugins/governance/store.go
@@ -195,12 +195,19 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData {
if vk == nil {
return
}
- if vk.BudgetID != nil {
- if liveBudget, exists := gs.budgets.Load(*vk.BudgetID); exists && liveBudget != nil {
- if b, ok := liveBudget.(*configstoreTables.TableBudget); ok {
- vk.Budget = b
+ // Cross-reference live budget/rate limit from standalone maps
+ // (usage updates clone into budgets/rateLimits maps, so embedded pointers go stale)
+ // Hydrate multi-budgets from live sync.Map
+ if len(vk.Budgets) > 0 {
+ liveBudgets := make([]configstoreTables.TableBudget, 0, len(vk.Budgets))
+ for _, b := range vk.Budgets {
+ if lb, exists := gs.budgets.Load(b.ID); exists && lb != nil {
+ if budget, ok := lb.(*configstoreTables.TableBudget); ok {
+ liveBudgets = append(liveBudgets, *budget)
+ }
}
}
+ vk.Budgets = liveBudgets
}
if vk.RateLimitID != nil {
if liveRL, exists := gs.rateLimits.Load(*vk.RateLimitID); exists && liveRL != nil {
@@ -213,12 +220,17 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData {
configs := make([]configstoreTables.TableVirtualKeyProviderConfig, len(vk.ProviderConfigs))
copy(configs, vk.ProviderConfigs)
for i := range configs {
- if configs[i].BudgetID != nil {
- if liveBudget, exists := gs.budgets.Load(*configs[i].BudgetID); exists && liveBudget != nil {
- if b, ok := liveBudget.(*configstoreTables.TableBudget); ok {
- configs[i].Budget = b
+ // Hydrate provider config multi-budgets
+ if len(configs[i].Budgets) > 0 {
+ liveBudgets := make([]configstoreTables.TableBudget, 0, len(configs[i].Budgets))
+ for _, b := range configs[i].Budgets {
+ if lb, exists := gs.budgets.Load(b.ID); exists && lb != nil {
+ if budget, ok := lb.(*configstoreTables.TableBudget); ok {
+ liveBudgets = append(liveBudgets, *budget)
+ }
}
}
+ configs[i].Budgets = liveBudgets
}
if configs[i].RateLimitID != nil {
if liveRL, exists := gs.rateLimits.Load(*configs[i].RateLimitID); exists && liveRL != nil {
@@ -1527,10 +1539,34 @@ func (gs *LocalGovernanceStore) ResetExpiredBudgetsInMemory(ctx context.Context)
var shouldReset bool
var newLastReset time.Time
- if budget.CalendarAligned {
+ // Check if the owning VK has calendar alignment enabled
+ // virtualKeys map is keyed by VK value (not ID), so we scan to find by VirtualKeyID
+ calendarAligned := false
+ if budget.VirtualKeyID != nil {
+ gs.virtualKeys.Range(func(_, v interface{}) bool {
+ if vk, ok := v.(*configstoreTables.TableVirtualKey); ok && vk != nil && vk.ID == *budget.VirtualKeyID {
+ calendarAligned = vk.CalendarAligned
+ return false // stop
+ }
+ return true
+ })
+ } else if budget.ProviderConfigID != nil {
+ // Provider config budgets: look up the VK that owns this provider config
+ gs.virtualKeys.Range(func(_, v interface{}) bool {
+ if vk, ok := v.(*configstoreTables.TableVirtualKey); ok && vk != nil {
+ for _, pc := range vk.ProviderConfigs {
+ if pc.ID == *budget.ProviderConfigID {
+ calendarAligned = vk.CalendarAligned
+ return false // stop
+ }
+ }
+ }
+ return true
+ })
+ }
+
+ if calendarAligned {
// Calendar-aligned: reset when we've entered a genuinely new calendar period.
- // This avoids the double-reset bug with rolling durations in months with
- // more days than ParseDuration approximates (e.g. 31-day months with "1M" = 30 days).
currentPeriodStart := configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, now)
if currentPeriodStart.After(budget.LastReset) {
shouldReset = true
@@ -2102,33 +2138,17 @@ func (gs *LocalGovernanceStore) loadFromConfigMemory(ctx context.Context, config
}
}
- for i := range budgets {
- if vk.BudgetID != nil && budgets[i].ID == *vk.BudgetID {
- vk.Budget = &budgets[i]
- }
- }
-
for i := range rateLimits {
if vk.RateLimitID != nil && rateLimits[i].ID == *vk.RateLimitID {
vk.RateLimit = &rateLimits[i]
}
}
- // Populate provider config relationships with budgets and rate limits
+ // Populate provider config relationships with rate limits
if vk.ProviderConfigs != nil {
for j := range vk.ProviderConfigs {
pc := &vk.ProviderConfigs[j]
- // Populate budget
- if pc.BudgetID != nil {
- for k := range budgets {
- if budgets[k].ID == *pc.BudgetID {
- pc.Budget = &budgets[k]
- break
- }
- }
- }
-
// Populate rate limit
if pc.RateLimitID != nil {
for k := range rateLimits {
@@ -2373,22 +2393,36 @@ func (gs *LocalGovernanceStore) collectBudgetsFromHierarchy(vk *configstoreTable
var budgetNames []string
// Collect all budgets in hierarchy order using lock-free sync.Map access (Provider Configs → VK → Team → Customer)
+ seen := make(map[string]bool)
for _, pc := range vk.ProviderConfigs {
- if pc.BudgetID != nil && pc.Provider == string(requestedProvider) {
- if budgetValue, exists := gs.budgets.Load(*pc.BudgetID); exists && budgetValue != nil {
+ if pc.Provider != string(requestedProvider) {
+ continue
+ }
+ // Multi-budgets
+ for _, b := range pc.Budgets {
+ if seen[b.ID] {
+ continue
+ }
+ if budgetValue, exists := gs.budgets.Load(b.ID); exists && budgetValue != nil {
if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
budgets = append(budgets, budget)
budgetNames = append(budgetNames, pc.Provider)
+ seen[budget.ID] = true
}
}
}
}
- if vk.BudgetID != nil {
- if budgetValue, exists := gs.budgets.Load(*vk.BudgetID); exists && budgetValue != nil {
+ // VK-level multi-budgets
+ for _, b := range vk.Budgets {
+ if seen[b.ID] {
+ continue
+ }
+ if budgetValue, exists := gs.budgets.Load(b.ID); exists && budgetValue != nil {
if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
budgets = append(budgets, budget)
budgetNames = append(budgetNames, "VK")
+ seen[budget.ID] = true
}
}
}
@@ -2477,9 +2511,9 @@ func (gs *LocalGovernanceStore) CreateVirtualKeyInMemory(vk *configstoreTables.T
return // Nothing to create
}
- // Create associated budget if exists
- if vk.Budget != nil {
- gs.budgets.Store(vk.Budget.ID, vk.Budget)
+ // Store budgets
+ for i := range vk.Budgets {
+ gs.budgets.Store(vk.Budgets[i].ID, &vk.Budgets[i])
}
// Create associated rate limit if exists
@@ -2490,8 +2524,8 @@ func (gs *LocalGovernanceStore) CreateVirtualKeyInMemory(vk *configstoreTables.T
// Create provider config budgets and rate limits if they exist
if vk.ProviderConfigs != nil {
for _, pc := range vk.ProviderConfigs {
- if pc.Budget != nil {
- gs.budgets.Store(pc.Budget.ID, pc.Budget)
+ for i := range pc.Budgets {
+ gs.budgets.Store(pc.Budgets[i].ID, &pc.Budgets[i])
}
if pc.RateLimit != nil {
gs.rateLimits.Store(pc.RateLimit.ID, pc.RateLimit)
@@ -2518,23 +2552,37 @@ func (gs *LocalGovernanceStore) UpdateVirtualKeyInMemory(vk *configstoreTables.T
// Create clone to avoid modifying the original
clone := *vk
- // Update Budget for VK in memory store
- if clone.Budget != nil {
- // Preserve existing usage from memory when updating budget config
- // The usage tracker maintains current usage in memory, and we only want to update
- // the configuration fields (max_limit, reset_duration) from the database
- if existingBudgetValue, exists := gs.budgets.Load(clone.Budget.ID); exists && existingBudgetValue != nil {
+
+ // Collect all incoming budget IDs across VK + provider configs to avoid
+ // deleting a budget that was moved between VK-level and PC-level in one update.
+ allNewBudgetIDs := make(map[string]bool)
+ for i := range clone.Budgets {
+ allNewBudgetIDs[clone.Budgets[i].ID] = true
+ }
+ for i := range clone.ProviderConfigs {
+ for j := range clone.ProviderConfigs[i].Budgets {
+ allNewBudgetIDs[clone.ProviderConfigs[i].Budgets[j].ID] = true
+ }
+ }
+
+ // Update multi-budgets for VK
+ for i := range clone.Budgets {
+ // Preserve existing usage from memory
+ if existingBudgetValue, exists := gs.budgets.Load(clone.Budgets[i].ID); exists && existingBudgetValue != nil {
if existingBudget, ok := existingBudgetValue.(*configstoreTables.TableBudget); ok && existingBudget != nil {
- // Preserve current usage and last reset time from existing in-memory budget
- clone.Budget.CurrentUsage = existingBudget.CurrentUsage
- clone.Budget.LastReset = existingBudget.LastReset
+ clone.Budgets[i].CurrentUsage = existingBudget.CurrentUsage
+ clone.Budgets[i].LastReset = existingBudget.LastReset
}
}
- gs.budgets.Store(clone.Budget.ID, clone.Budget)
- } else if existingVK.Budget != nil {
- // Budget was removed from the virtual key, delete it from memory
- gs.budgets.Delete(existingVK.Budget.ID)
+ gs.budgets.Store(clone.Budgets[i].ID, &clone.Budgets[i])
}
+ // Delete removed multi-budgets
+ for _, oldBudget := range existingVK.Budgets {
+ if !allNewBudgetIDs[oldBudget.ID] {
+ gs.budgets.Delete(oldBudget.ID)
+ }
+ }
+
if clone.RateLimit != nil {
// Preserve existing usage from memory when updating rate limit config
// The usage tracker maintains current usage in memory, and we only want to update
@@ -2584,22 +2632,23 @@ func (gs *LocalGovernanceStore) UpdateVirtualKeyInMemory(vk *configstoreTables.T
clone.ProviderConfigs[i].RateLimit = nil
}
}
- // Update Budget for provider config in memory store
- if pc.Budget != nil {
- // Preserve existing usage from memory when updating provider config budget
- if existingBudgetValue, exists := gs.budgets.Load(pc.Budget.ID); exists && existingBudgetValue != nil {
+ // Update multi-budgets for provider config
+ for j := range clone.ProviderConfigs[i].Budgets {
+ b := &clone.ProviderConfigs[i].Budgets[j]
+ if existingBudgetValue, exists := gs.budgets.Load(b.ID); exists && existingBudgetValue != nil {
if existingBudget, ok := existingBudgetValue.(*configstoreTables.TableBudget); ok && existingBudget != nil {
- // Preserve current usage and last reset time from existing in-memory budget
- clone.ProviderConfigs[i].Budget.CurrentUsage = existingBudget.CurrentUsage
- clone.ProviderConfigs[i].Budget.LastReset = existingBudget.LastReset
+ b.CurrentUsage = existingBudget.CurrentUsage
+ b.LastReset = existingBudget.LastReset
}
}
- gs.budgets.Store(clone.ProviderConfigs[i].Budget.ID, clone.ProviderConfigs[i].Budget)
- } else {
- // Budget was removed from provider config, delete it from memory if it existed
- if existingPC, exists := existingProviderConfigs[pc.ID]; exists && existingPC.Budget != nil {
- gs.budgets.Delete(existingPC.Budget.ID)
- clone.ProviderConfigs[i].Budget = nil
+ gs.budgets.Store(b.ID, b)
+ }
+ // Delete removed multi-budgets for this provider config
+ if existingPC, exists := existingProviderConfigs[pc.ID]; exists {
+ for _, oldBudget := range existingPC.Budgets {
+ if !allNewBudgetIDs[oldBudget.ID] {
+ gs.budgets.Delete(oldBudget.ID)
+ }
}
}
}
@@ -2625,9 +2674,9 @@ func (gs *LocalGovernanceStore) DeleteVirtualKeyInMemory(vkID string) {
}
if vk.ID == vkID {
- // Delete associated budget if exists
- if vk.BudgetID != nil {
- gs.budgets.Delete(*vk.BudgetID)
+ // Delete budgets
+ for _, b := range vk.Budgets {
+ gs.budgets.Delete(b.ID)
}
// Delete associated rate limit if exists
@@ -2638,8 +2687,8 @@ func (gs *LocalGovernanceStore) DeleteVirtualKeyInMemory(vkID string) {
// Delete provider config budgets and rate limits
if vk.ProviderConfigs != nil {
for _, pc := range vk.ProviderConfigs {
- if pc.BudgetID != nil {
- gs.budgets.Delete(*pc.BudgetID)
+ for _, b := range pc.Budgets {
+ gs.budgets.Delete(b.ID)
}
if pc.RateLimitID != nil {
gs.rateLimits.Delete(*pc.RateLimitID)
@@ -3165,18 +3214,22 @@ func (gs *LocalGovernanceStore) updateBudgetReferences(resetBudget *configstoreT
needsUpdate := false
clone := *vk
- // Check VK-level budget
- if vk.BudgetID != nil && *vk.BudgetID == budgetID {
- clone.Budget = resetBudget
- needsUpdate = true
+ // Check VK-level budgets
+ for i, b := range clone.Budgets {
+ if b.ID == budgetID {
+ clone.Budgets[i] = *resetBudget
+ needsUpdate = true
+ }
}
// Check provider config budgets
if vk.ProviderConfigs != nil {
- for i, pc := range clone.ProviderConfigs {
- if pc.BudgetID != nil && *pc.BudgetID == budgetID {
- clone.ProviderConfigs[i].Budget = resetBudget
- needsUpdate = true
+ for i := range clone.ProviderConfigs {
+ for j, b := range clone.ProviderConfigs[i].Budgets {
+ if b.ID == budgetID {
+ clone.ProviderConfigs[i].Budgets[j] = *resetBudget
+ needsUpdate = true
+ }
}
}
}
@@ -3638,9 +3691,9 @@ func (gs *LocalGovernanceStore) GetBudgetAndRateLimitStatus(ctx context.Context,
}
}
}
- // Get budget status
- if pc.BudgetID != nil {
- if budgetValue, ok := gs.budgets.Load(*pc.BudgetID); ok && budgetValue != nil {
+ // Get budget status from multi-budgets
+ for _, b := range pc.Budgets {
+ if budgetValue, ok := gs.budgets.Load(b.ID); ok && budgetValue != nil {
if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
baseline, exists := budgetBaselines[budget.ID]
if !exists {
diff --git a/plugins/governance/store_test.go b/plugins/governance/store_test.go
index 0640052d56..21a54e1bd5 100644
--- a/plugins/governance/store_test.go
+++ b/plugins/governance/store_test.go
@@ -198,16 +198,397 @@ func TestGovernanceStore_CheckBudget_HierarchyValidation(t *testing.T) {
// Test: If VK budget exceeds limit, should fail
// Update the budget directly in the budgets map (since UpdateVirtualKeyInMemory preserves usage)
- if vk.BudgetID != nil {
- if budgetValue, exists := store.budgets.Load(*vk.BudgetID); exists && budgetValue != nil {
+ if len(vk.Budgets) > 0 {
+ budgetID := vk.Budgets[0].ID
+ if budgetValue, exists := store.budgets.Load(budgetID); exists && budgetValue != nil {
if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil {
budget.CurrentUsage = 100.0
- store.budgets.Store(*vk.BudgetID, budget)
+ store.budgets.Store(budgetID, budget)
}
}
}
err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
- assert.Error(t, err, "Should fail when VK budget exceeds limit")
+ require.Error(t, err, "Should fail when VK budget exceeds limit")
+}
+
+// TestGovernanceStore_MultiBudget_AllUnderLimit tests that requests pass when all budgets are under their limits
+func TestGovernanceStore_MultiBudget_AllUnderLimit(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Create VK with hourly ($10) and daily ($100) budgets
+ hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 5.0, "1h")
+ dailyBudget := buildBudgetWithUsage("daily", 100.0, 40.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ // Add provider config so the resolver allows the provider
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ assert.NoError(t, err, "Should pass when all budgets are under limit")
+}
+
+// TestGovernanceStore_MultiBudget_SmallBudgetExceeded tests that request is blocked when the smaller budget exceeds its limit
+func TestGovernanceStore_MultiBudget_SmallBudgetExceeded(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Hourly at limit, daily still has room
+ hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 10.0, "1h")
+ dailyBudget := buildBudgetWithUsage("daily", 100.0, 40.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ require.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine")
+ assert.Contains(t, err.Error(), "budget exceeded")
+}
+
+// TestGovernanceStore_MultiBudget_LargeBudgetExceeded tests that request is blocked when only the larger budget exceeds
+func TestGovernanceStore_MultiBudget_LargeBudgetExceeded(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Hourly has room, but daily is at limit
+ hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 3.0, "1h")
+ dailyBudget := buildBudgetWithUsage("daily", 100.0, 100.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ require.Error(t, err, "Should fail when daily budget is exceeded even though hourly is fine")
+ assert.Contains(t, err.Error(), "budget exceeded")
+}
+
+// TestGovernanceStore_MultiBudget_UsageUpdatesAllBudgets tests that usage updates are applied to every budget in the hierarchy
+func TestGovernanceStore_MultiBudget_UsageUpdatesAllBudgets(t *testing.T) {
+ logger := NewMockLogger()
+
+ hourlyBudget := buildBudget("hourly", 10.0, "1h")
+ dailyBudget := buildBudget("daily", 100.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+
+ // Simulate a $3.50 request
+ err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 3.50)
+ require.NoError(t, err)
+
+ // Both budgets should reflect the cost
+ hourlyVal, exists := store.budgets.Load("hourly")
+ require.True(t, exists)
+ assert.InDelta(t, 3.50, hourlyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Hourly budget should reflect usage")
+
+ dailyVal, exists := store.budgets.Load("daily")
+ require.True(t, exists)
+ assert.InDelta(t, 3.50, dailyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Daily budget should reflect usage")
+
+ // Second request: $7.00 — should push hourly over limit
+ err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 7.00)
+ require.NoError(t, err)
+
+ hourlyVal, _ = store.budgets.Load("hourly")
+ assert.InDelta(t, 10.50, hourlyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Hourly budget should accumulate")
+
+ dailyVal, _ = store.budgets.Load("daily")
+ assert.InDelta(t, 10.50, dailyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01, "Daily budget should accumulate")
+
+ // Now CheckBudget should fail (hourly exceeded)
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ require.Error(t, err, "Should fail after usage exceeds hourly budget")
+ assert.Contains(t, err.Error(), "budget exceeded")
+}
+
+// TestGovernanceStore_MultiBudget_ProviderConfigBudgets tests that provider-config-level multi-budgets are enforced
+func TestGovernanceStore_MultiBudget_ProviderConfigBudgets(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Provider-level budgets: hourly $5 (exceeded), daily $50 (ok)
+ pcHourly := buildBudgetWithUsage("pc-hourly", 5.0, 5.0, "1h")
+ pcDaily := buildBudgetWithUsage("pc-daily", 50.0, 10.0, "1d")
+
+ pc := buildProviderConfigWithBudgets("openai", []string{"*"},
+ []configstoreTables.TableBudget{*pcHourly, *pcDaily})
+
+ vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableVirtualKeyProviderConfig{pc})
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*pcHourly, *pcDaily},
+ }, nil)
+ require.NoError(t, err)
+
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ require.Error(t, err, "Should fail when provider config hourly budget is exceeded")
+ assert.Contains(t, err.Error(), "budget exceeded")
+}
+
+// TestGovernanceStore_MultiBudget_VKAndProviderConfigCombined tests budgets at both VK and provider config levels
+func TestGovernanceStore_MultiBudget_VKAndProviderConfigCombined(t *testing.T) {
+ logger := NewMockLogger()
+
+ // VK-level budgets: all under limit
+ vkMonthly := buildBudgetWithUsage("vk-monthly", 1000.0, 200.0, "1M")
+
+ // Provider-config-level budgets: hourly at limit
+ pcHourly := buildBudgetWithUsage("pc-hourly", 5.0, 5.0, "1h")
+
+ pc := buildProviderConfigWithBudgets("openai", []string{"*"},
+ []configstoreTables.TableBudget{*pcHourly})
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*vkMonthly})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{pc}
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*vkMonthly, *pcHourly},
+ }, nil)
+ require.NoError(t, err)
+
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+
+ // Provider config budget exceeded → should block even though VK budget is fine
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ require.Error(t, err, "Should fail: provider config budget exceeded even though VK budget is fine")
+ assert.Contains(t, err.Error(), "budget exceeded")
+}
+
+// TestGovernanceStore_MultiBudget_ResolverBlocksOnBudgetExceeded tests that the full resolver flow blocks when any budget is exceeded
+func TestGovernanceStore_MultiBudget_ResolverBlocksOnBudgetExceeded(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Two VK-level budgets: hourly at limit, daily has room
+ hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 10.0, "1h")
+ dailyBudget := buildBudgetWithUsage("daily", 100.0, 30.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ resolver := NewBudgetResolver(store, nil, logger)
+ ctx := &schemas.BifrostContext{}
+
+ result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
+ assertDecision(t, DecisionBudgetExceeded, result)
+ assert.Contains(t, result.Reason, "budget exceeded")
+}
+
+// TestGovernanceStore_MultiBudget_ResolverAllowsUnderLimit tests that the full resolver flow allows requests when all budgets are under limit
+func TestGovernanceStore_MultiBudget_ResolverAllowsUnderLimit(t *testing.T) {
+ logger := NewMockLogger()
+
+ hourlyBudget := buildBudgetWithUsage("hourly", 10.0, 5.0, "1h")
+ dailyBudget := buildBudgetWithUsage("daily", 100.0, 30.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ resolver := NewBudgetResolver(store, nil, logger)
+ ctx := &schemas.BifrostContext{}
+
+ result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
+ assertDecision(t, DecisionAllow, result)
+}
+
+// TestGovernanceStore_MultiBudget_UsageDrivesBlockAfterRequests tests the full lifecycle:
+// start under limit → accumulate usage → eventually hit a budget → get blocked
+func TestGovernanceStore_MultiBudget_UsageDrivesBlockAfterRequests(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Tight hourly ($2), generous daily ($100)
+ hourlyBudget := buildBudget("hourly", 2.0, "1h")
+ dailyBudget := buildBudget("daily", 100.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*hourlyBudget, *dailyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ resolver := NewBudgetResolver(store, nil, logger)
+
+ // Request 1: $0.80 — both budgets fine
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 0.80)
+ require.NoError(t, err)
+
+ ctx := &schemas.BifrostContext{}
+ result := resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
+ assertDecision(t, DecisionAllow, result)
+
+ // Request 2: $0.80 — still fine ($1.60 total)
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 0.80)
+ require.NoError(t, err)
+
+ ctx = &schemas.BifrostContext{}
+ result = resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
+ assertDecision(t, DecisionAllow, result)
+
+ // Request 3: $0.80 — pushes hourly to $2.40 > $2.00 limit → blocked
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ err = store.UpdateVirtualKeyBudgetUsageInMemory(context.Background(), vk, schemas.OpenAI, 0.80)
+ require.NoError(t, err)
+
+ ctx = &schemas.BifrostContext{}
+ result = resolver.EvaluateVirtualKeyRequest(ctx, "sk-bf-test", schemas.OpenAI, "gpt-4", schemas.ChatCompletionRequest, false)
+ assertDecision(t, DecisionBudgetExceeded, result)
+ assert.Contains(t, result.Reason, "budget exceeded")
+
+ // Verify daily budget is still under limit
+ dailyVal, exists := store.budgets.Load("daily")
+ require.True(t, exists)
+ assert.InDelta(t, 2.40, dailyVal.(*configstoreTables.TableBudget).CurrentUsage, 0.01,
+ "Daily budget should be at $2.40, well under $100 limit")
+}
+
+// TestGovernanceStore_MultiBudget_CalendarAligned tests that calendar-aligned budgets are stored and retrievable
+func TestGovernanceStore_MultiBudget_CalendarAligned(t *testing.T) {
+ logger := NewMockLogger()
+
+ // Calendar alignment is a VK-level setting — budgets don't have it
+ dailyBudget := &configstoreTables.TableBudget{
+ ID: "daily-cal",
+ MaxLimit: 50.0,
+ CurrentUsage: 10.0,
+ ResetDuration: "1d",
+ LastReset: time.Now(),
+ }
+ monthlyBudget := &configstoreTables.TableBudget{
+ ID: "monthly-cal",
+ MaxLimit: 1000.0,
+ CurrentUsage: 200.0,
+ ResetDuration: "1M",
+ LastReset: time.Now(),
+ }
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*dailyBudget, *monthlyBudget})
+ vk.CalendarAligned = true // VK-level setting applies to all budgets
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
+ VirtualKeys: []configstoreTables.TableVirtualKey{*vk},
+ Budgets: []configstoreTables.TableBudget{*dailyBudget, *monthlyBudget},
+ }, nil)
+ require.NoError(t, err)
+
+ // Verify VK-level calendar_aligned is set
+ vk, _ = store.GetVirtualKey("sk-bf-test")
+ assert.True(t, vk.CalendarAligned, "VK should have calendar_aligned=true")
+
+ // Both under limit — should pass
+ err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil)
+ assert.NoError(t, err)
+}
+
+// TestGovernanceStore_MultiBudget_InMemoryCreateAndDelete tests CreateVirtualKeyInMemory and DeleteVirtualKeyInMemory
+// properly store and clean up multi-budget entries
+func TestGovernanceStore_MultiBudget_InMemoryCreateAndDelete(t *testing.T) {
+ logger := NewMockLogger()
+
+ store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{}, nil)
+ require.NoError(t, err)
+
+ b1 := buildBudget("b1", 10.0, "1h")
+ b2 := buildBudget("b2", 100.0, "1d")
+
+ vk := buildVirtualKeyWithMultiBudgets("vk1", "sk-bf-test", "Test VK",
+ []configstoreTables.TableBudget{*b1, *b2})
+ vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
+ buildProviderConfig("openai", []string{"*"}),
+ }
+
+ // Create
+ store.CreateVirtualKeyInMemory(vk)
+
+ _, exists := store.budgets.Load("b1")
+ assert.True(t, exists, "Budget b1 should be in memory after create")
+ _, exists = store.budgets.Load("b2")
+ assert.True(t, exists, "Budget b2 should be in memory after create")
+
+ retrieved, found := store.GetVirtualKey("sk-bf-test")
+ require.True(t, found)
+ assert.Len(t, retrieved.Budgets, 2, "VK should have 2 budgets")
+
+ // Delete
+ store.DeleteVirtualKeyInMemory("vk1")
+
+ _, exists = store.budgets.Load("b1")
+ assert.False(t, exists, "Budget b1 should be removed after delete")
+ _, exists = store.budgets.Load("b2")
+ assert.False(t, exists, "Budget b2 should be removed after delete")
+
+ _, found = store.GetVirtualKey("sk-bf-test")
+ assert.False(t, found, "VK should not be found after delete")
}
// TestGovernanceStore_UpdateRateLimitUsage_TokensAndRequests tests atomic rate limit usage updates
@@ -319,9 +700,9 @@ func TestGovernanceStore_ResetExpiredBudgets(t *testing.T) {
// Retrieve the updated VK to check budget changes
updatedVK, _ := store.GetVirtualKey("sk-bf-test")
require.NotNil(t, updatedVK)
- require.NotNil(t, updatedVK.Budget)
+ require.True(t, len(updatedVK.Budgets) > 0, "VK should have budgets")
- assert.Equal(t, 0.0, updatedVK.Budget.CurrentUsage, "Budget usage should be reset")
+ assert.Equal(t, 0.0, updatedVK.Budgets[0].CurrentUsage, "Budget usage should be reset")
}
// TestGovernanceStore_GetAllBudgets tests retrieving all budgets
diff --git a/plugins/governance/test_utils.go b/plugins/governance/test_utils.go
index ccd1e5968d..b7ad6dbad3 100644
--- a/plugins/governance/test_utils.go
+++ b/plugins/governance/test_utils.go
@@ -15,7 +15,6 @@ import (
configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables"
"github.com/maximhq/bifrost/framework/modelcatalog"
"github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
)
// MockLogger implements schemas.Logger for testing
@@ -89,9 +88,10 @@ func buildVirtualKey(id, value, name string, isActive bool) *configstoreTables.T
func buildVirtualKeyWithBudget(id, value, name string, budget *configstoreTables.TableBudget) *configstoreTables.TableVirtualKey {
vk := buildVirtualKey(id, value, name, true)
- vk.Budget = budget
- budgetID := budget.ID
- vk.BudgetID = &budgetID
+ vkID := id
+ budget.VirtualKeyID = &vkID
+ vk.Budgets = []configstoreTables.TableBudget{*budget}
+ // Add a default provider config so the resolver doesn't block at provider check
vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
buildProviderConfig("openai", []string{"*"}),
}
@@ -103,6 +103,7 @@ func buildVirtualKeyWithRateLimit(id, value, name string, rateLimit *configstore
vk.RateLimit = rateLimit
rateLimitID := rateLimit.ID
vk.RateLimitID = &rateLimitID
+ // Add a default provider config so the resolver doesn't block at provider check
vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{
buildProviderConfig("openai", []string{"*"}),
}
@@ -195,11 +196,26 @@ func buildProviderConfig(provider string, allowedModels []string) configstoreTab
AllowedModels: allowedModels,
Weight: bifrost.Ptr(1.0),
RateLimit: nil,
- Budget: nil,
Keys: []configstoreTables.TableKey{},
}
}
+func buildProviderConfigWithBudgets(provider string, allowedModels []string, budgets []configstoreTables.TableBudget) configstoreTables.TableVirtualKeyProviderConfig {
+ pc := buildProviderConfig(provider, allowedModels)
+ pc.Budgets = budgets
+ return pc
+}
+
+func buildVirtualKeyWithMultiBudgets(id, value, name string, budgets []configstoreTables.TableBudget) *configstoreTables.TableVirtualKey {
+ vk := buildVirtualKey(id, value, name, true)
+ for i := range budgets {
+ vkID := id
+ budgets[i].VirtualKeyID = &vkID
+ }
+ vk.Budgets = budgets
+ return vk
+}
+
func buildProviderConfigWithRateLimit(provider string, allowedModels []string, rateLimit *configstoreTables.TableRateLimit) configstoreTables.TableVirtualKeyProviderConfig {
pc := buildProviderConfig(provider, allowedModels)
pc.RateLimit = rateLimit
@@ -227,15 +243,6 @@ func assertRateLimitInfo(t *testing.T, result *EvaluationResult) {
assert.NotNil(t, result.RateLimitInfo, "RateLimitInfo should be present in result")
}
-func requireNoError(t *testing.T, err error, msg string) {
- t.Helper()
- require.NoError(t, err, msg)
-}
-
-func requireError(t *testing.T, err error, msg string) {
- t.Helper()
- require.Error(t, err, msg)
-}
func buildModelConfig(id, modelName string, provider *string, budget *configstoreTables.TableBudget, rateLimit *configstoreTables.TableRateLimit) *configstoreTables.TableModelConfig {
mc := &configstoreTables.TableModelConfig{
diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go
index b21008983b..cd5b2250de 100644
--- a/transports/bifrost-http/handlers/governance.go
+++ b/transports/bifrost-http/handlers/governance.go
@@ -73,7 +73,7 @@ type CreateVirtualKeyRequest struct {
Provider string `json:"provider" validate:"required"`
Weight *float64 `json:"weight,omitempty"`
AllowedModels schemas.WhiteList `json:"allowed_models,omitempty"` // ["*"] allows all models; empty denies all
- Budget *CreateBudgetRequest `json:"budget,omitempty"` // Provider-level budget
+ Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget for provider config
RateLimit *CreateRateLimitRequest `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
} `json:"provider_configs,omitempty"` // Empty means no providers allowed (deny-by-default)
@@ -83,9 +83,10 @@ type CreateVirtualKeyRequest struct {
} `json:"mcp_configs,omitempty"` // Empty means no MCP clients allowed (deny-by-default)
TeamID *string `json:"team_id,omitempty"` // Mutually exclusive with CustomerID
CustomerID *string `json:"customer_id,omitempty"` // Mutually exclusive with TeamID
- Budget *CreateBudgetRequest `json:"budget,omitempty"`
- RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"`
- IsActive *bool `json:"is_active,omitempty"`
+ Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: each must have a unique reset_duration
+ RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"`
+ IsActive *bool `json:"is_active,omitempty"`
+ CalendarAligned bool `json:"calendar_aligned,omitempty"` // When true, all budgets reset at clean calendar boundaries
}
// UpdateVirtualKeyRequest represents the request body for updating a virtual key
@@ -97,7 +98,7 @@ type UpdateVirtualKeyRequest struct {
Provider string `json:"provider" validate:"required"`
Weight *float64 `json:"weight,omitempty"`
AllowedModels schemas.WhiteList `json:"allowed_models,omitempty"` // ["*"] allows all models; empty denies all
- Budget *UpdateBudgetRequest `json:"budget,omitempty"` // Provider-level budget
+ 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
} `json:"provider_configs,omitempty"`
@@ -108,9 +109,10 @@ type UpdateVirtualKeyRequest struct {
} `json:"mcp_configs,omitempty"`
TeamID *string `json:"team_id,omitempty"`
CustomerID *string `json:"customer_id,omitempty"`
- Budget *UpdateBudgetRequest `json:"budget,omitempty"`
- RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"`
- IsActive *bool `json:"is_active,omitempty"`
+ Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: replaces all VK-level budgets
+ RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"`
+ IsActive *bool `json:"is_active,omitempty"`
+ CalendarAligned *bool `json:"calendar_aligned,omitempty"` // When true, all budgets reset at clean calendar boundaries
}
// CreateBudgetRequest represents the request body for creating a budget
@@ -207,8 +209,8 @@ func collectProviderConfigDeleteIDs(
budgetIDs []string,
rateLimitIDs []string,
) ([]string, []string) {
- if config.BudgetID != nil {
- budgetIDs = append(budgetIDs, *config.BudgetID)
+ for _, b := range config.Budgets {
+ budgetIDs = append(budgetIDs, b.ID)
}
if config.RateLimitID != nil {
rateLimitIDs = append(rateLimitIDs, *config.RateLimitID)
@@ -439,16 +441,23 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
SendError(ctx, 400, "VirtualKey cannot be attached to both Team and Customer")
return
}
- // Validate budget if provided
- if req.Budget != nil {
- if req.Budget.MaxLimit < 0 {
- SendError(ctx, 400, fmt.Sprintf("Budget max_limit cannot be negative: %.2f", req.Budget.MaxLimit))
- return
- }
- // Validate reset duration format
- if _, err := configstoreTables.ParseDuration(req.Budget.ResetDuration); err != nil {
- SendError(ctx, 400, fmt.Sprintf("Invalid reset duration format: %s", req.Budget.ResetDuration))
- return
+ // Validate budgets if provided
+ if len(req.Budgets) > 0 {
+ seenDurations := make(map[string]bool)
+ for _, b := range req.Budgets {
+ if b.MaxLimit < 0 {
+ SendError(ctx, 400, fmt.Sprintf("Budget max_limit cannot be negative: %.2f", b.MaxLimit))
+ return
+ }
+ if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil {
+ SendError(ctx, 400, fmt.Sprintf("Invalid reset duration format: %s", b.ResetDuration))
+ return
+ }
+ if seenDurations[b.ResetDuration] {
+ SendError(ctx, 400, fmt.Sprintf("Duplicate reset_duration in budgets: %s", b.ResetDuration))
+ return
+ }
+ seenDurations[b.ResetDuration] = true
}
}
// Set defaults
@@ -459,30 +468,14 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
var vk configstoreTables.TableVirtualKey
if err := h.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error {
vk = configstoreTables.TableVirtualKey{
- ID: uuid.NewString(),
- Name: req.Name,
- Value: governance.GenerateVirtualKey(),
- Description: req.Description,
- TeamID: req.TeamID,
- CustomerID: req.CustomerID,
- IsActive: isActive,
- }
- if req.Budget != nil {
- budget := configstoreTables.TableBudget{
- ID: uuid.NewString(),
- MaxLimit: req.Budget.MaxLimit,
- ResetDuration: req.Budget.ResetDuration,
- CalendarAligned: req.Budget.CalendarAligned,
- LastReset: budgetLastReset(req.Budget.CalendarAligned, req.Budget.ResetDuration),
- CurrentUsage: 0,
- }
- if err := validateBudget(&budget); err != nil {
- return err
- }
- if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
- return err
- }
- vk.BudgetID = &budget.ID
+ ID: uuid.NewString(),
+ Name: req.Name,
+ Value: governance.GenerateVirtualKey(),
+ Description: req.Description,
+ TeamID: req.TeamID,
+ CustomerID: req.CustomerID,
+ IsActive: isActive,
+ CalendarAligned: req.CalendarAligned,
}
if req.RateLimit != nil {
rateLimit := configstoreTables.TableRateLimit{
@@ -505,19 +498,27 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
if err := h.configStore.CreateVirtualKey(ctx, &vk, tx); err != nil {
return err
}
+ // Create multi-budgets for VK
+ if len(req.Budgets) > 0 {
+ for _, b := range req.Budgets {
+ budget := configstoreTables.TableBudget{
+ ID: uuid.NewString(),
+ MaxLimit: b.MaxLimit,
+ ResetDuration: b.ResetDuration,
+ LastReset: budgetLastReset(vk.CalendarAligned, b.ResetDuration),
+ CurrentUsage: 0,
+ VirtualKeyID: &vk.ID,
+ }
+ if err := validateBudget(&budget); err != nil {
+ return err
+ }
+ if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
+ return err
+ }
+ }
+ }
if req.ProviderConfigs != nil {
for _, pc := range req.ProviderConfigs {
- // Validate budget if provided
- if pc.Budget != nil {
- if pc.Budget.MaxLimit < 0 {
- return fmt.Errorf("provider config budget max_limit cannot be negative: %.2f", pc.Budget.MaxLimit)
- }
- // Validate reset duration format
- if _, err := configstoreTables.ParseDuration(pc.Budget.ResetDuration); err != nil {
- return fmt.Errorf("invalid provider config budget reset duration format: %s", pc.Budget.ResetDuration)
- }
- }
-
if err := pc.AllowedModels.Validate(); err != nil {
return &badRequestError{err: fmt.Errorf("invalid allowed_models for provider %s: %w", pc.Provider, err)}
}
@@ -550,24 +551,6 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
Keys: keys,
}
- // Create budget for provider config if provided
- if pc.Budget != nil {
- budget := configstoreTables.TableBudget{
- ID: uuid.NewString(),
- MaxLimit: pc.Budget.MaxLimit,
- ResetDuration: pc.Budget.ResetDuration,
- CalendarAligned: pc.Budget.CalendarAligned,
- LastReset: budgetLastReset(pc.Budget.CalendarAligned, pc.Budget.ResetDuration),
- CurrentUsage: 0,
- }
- if err := validateBudget(&budget); err != nil {
- return err
- }
- if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
- return err
- }
- providerConfig.BudgetID = &budget.ID
- }
// Create rate limit for provider config if provided
if pc.RateLimit != nil {
rateLimit := configstoreTables.TableRateLimit{
@@ -591,6 +574,30 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
if err := h.configStore.CreateVirtualKeyProviderConfig(ctx, providerConfig, tx); err != nil {
return err
}
+ // Create multi-budgets for provider config
+ if len(pc.Budgets) > 0 {
+ seenDurations := make(map[string]bool)
+ for _, b := range pc.Budgets {
+ if seenDurations[b.ResetDuration] {
+ return &badRequestError{err: fmt.Errorf("duplicate reset_duration in provider config budgets: %s", b.ResetDuration)}
+ }
+ seenDurations[b.ResetDuration] = true
+ budget := configstoreTables.TableBudget{
+ ID: uuid.NewString(),
+ MaxLimit: b.MaxLimit,
+ ResetDuration: b.ResetDuration,
+ LastReset: budgetLastReset(vk.CalendarAligned, b.ResetDuration),
+ CurrentUsage: 0,
+ ProviderConfigID: &providerConfig.ID,
+ }
+ if err := validateBudget(&budget); err != nil {
+ return err
+ }
+ if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
+ return err
+ }
+ }
+ }
}
}
if req.MCPConfigs != nil {
@@ -702,8 +709,9 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
return
}
if err := h.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error {
- var budgetIDToDelete, rateLimitIDToDelete string
- var providerBudgetIDsToDelete, providerRateLimitIDsToDelete []string
+ var rateLimitIDToDelete string
+ var providerBudgetIDsToDelete []string
+ var providerRateLimitIDsToDelete []string
// Update fields if provided
if req.Name != nil {
@@ -728,72 +736,77 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
if req.IsActive != nil {
vk.IsActive = *req.IsActive
}
- // Handle budget updates
- if req.Budget != nil {
- if isBudgetRemovalRequest(req.Budget) {
- if vk.BudgetID != nil {
- budgetIDToDelete = *vk.BudgetID
- vk.BudgetID = nil
- vk.Budget = nil
- }
- } else if vk.BudgetID != nil {
- // Update existing budget
- budget := configstoreTables.TableBudget{}
- if err := tx.First(&budget, "id = ?", *vk.BudgetID).Error; err != nil {
- return err
- }
-
- if req.Budget.MaxLimit != nil {
- budget.MaxLimit = *req.Budget.MaxLimit
- }
- if req.Budget.ResetDuration != nil {
- budget.ResetDuration = *req.Budget.ResetDuration
- }
- if req.Budget.CalendarAligned != nil {
- wasCalendarAligned := budget.CalendarAligned
- budget.CalendarAligned = *req.Budget.CalendarAligned
- if *req.Budget.CalendarAligned && !wasCalendarAligned {
- budget.LastReset = configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, time.Now())
- budget.CurrentUsage = 0
+ if req.CalendarAligned != nil {
+ vk.CalendarAligned = *req.CalendarAligned
+ }
+ // Handle multi-budget updates
+ if req.Budgets != nil {
+ // Validate multi-budgets
+ seenDurations := make(map[string]bool)
+ for _, b := range req.Budgets {
+ if b.MaxLimit < 0 {
+ return &badRequestError{err: fmt.Errorf("budget max_limit cannot be negative: %.2f", b.MaxLimit)}
}
- }
- if err := validateBudget(&budget); err != nil {
- return err
- }
- if err := h.configStore.UpdateBudget(ctx, &budget, tx); err != nil {
- return err
- }
- vk.Budget = &budget
- } else {
- // Create new budget
- if req.Budget.MaxLimit == nil || req.Budget.ResetDuration == nil {
- return fmt.Errorf("both max_limit and reset_duration are required when creating a new budget")
+ if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil {
+ return &badRequestError{err: fmt.Errorf("invalid reset duration format: %s", b.ResetDuration)}
}
- if *req.Budget.MaxLimit < 0 {
- return fmt.Errorf("budget max_limit cannot be negative: %.2f", *req.Budget.MaxLimit)
+ if seenDurations[b.ResetDuration] {
+ return &badRequestError{err: fmt.Errorf("duplicate reset_duration in budgets: %s", b.ResetDuration)}
}
- if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil {
- return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration)
- }
- calAligned := req.Budget.CalendarAligned != nil && *req.Budget.CalendarAligned
- budget := configstoreTables.TableBudget{
- ID: uuid.NewString(),
- MaxLimit: *req.Budget.MaxLimit,
- ResetDuration: *req.Budget.ResetDuration,
- CalendarAligned: calAligned,
- LastReset: budgetLastReset(calAligned, *req.Budget.ResetDuration),
- CurrentUsage: 0,
+ seenDurations[b.ResetDuration] = true
}
- if err := validateBudget(&budget); err != nil {
- return err
+
+ // Build map of existing budgets by reset_duration for matching
+ existingByDuration := make(map[string]configstoreTables.TableBudget)
+ for _, existing := range vk.Budgets {
+ existingByDuration[existing.ResetDuration] = existing
}
- if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
- return err
+
+ // Reconcile: preserve existing budgets where possible, create new ones where needed
+ var reconciledBudgets []configstoreTables.TableBudget
+ matchedIDs := make(map[string]bool)
+ for _, b := range req.Budgets {
+ if existing, found := existingByDuration[b.ResetDuration]; found {
+ // Budget with same duration exists — update max_limit, preserve usage
+ existing.MaxLimit = b.MaxLimit
+ if err := validateBudget(&existing); err != nil {
+ return err
+ }
+ if err := h.configStore.UpdateBudget(ctx, &existing, tx); err != nil {
+ return err
+ }
+ reconciledBudgets = append(reconciledBudgets, existing)
+ matchedIDs[existing.ID] = true
+ } else {
+ // New budget duration — create fresh
+ budget := configstoreTables.TableBudget{
+ ID: uuid.NewString(),
+ MaxLimit: b.MaxLimit,
+ ResetDuration: b.ResetDuration,
+ LastReset: budgetLastReset(vk.CalendarAligned, b.ResetDuration),
+ CurrentUsage: 0,
+ VirtualKeyID: &vk.ID,
+ }
+ if err := validateBudget(&budget); err != nil {
+ return err
+ }
+ if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
+ return err
+ }
+ reconciledBudgets = append(reconciledBudgets, budget)
+ }
}
- vk.BudgetID = &budget.ID
- vk.Budget = &budget
+ // Delete budgets that are no longer present
+ for _, existing := range vk.Budgets {
+ if !matchedIDs[existing.ID] {
+ if err := h.configStore.DeleteBudget(ctx, existing.ID, tx); err != nil {
+ return fmt.Errorf("failed to delete removed VK budget: %w", err)
+ }
+ }
}
+ vk.Budgets = reconciledBudgets
}
+
// Handle rate limit updates
if req.RateLimit != nil {
if isRateLimitRemovalRequest(req.RateLimit) {
@@ -864,21 +877,6 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
// Process new configs: create new ones and update existing ones
for _, pc := range req.ProviderConfigs {
if pc.ID == nil {
- // Validate budget if provided for new provider config
- if pc.Budget != nil {
- if pc.Budget.MaxLimit != nil && *pc.Budget.MaxLimit < 0 {
- return fmt.Errorf("provider config budget max_limit cannot be negative: %.2f", *pc.Budget.MaxLimit)
- }
- if pc.Budget.ResetDuration != nil {
- if _, err := configstoreTables.ParseDuration(*pc.Budget.ResetDuration); err != nil {
- return fmt.Errorf("invalid provider config budget reset duration format: %s", *pc.Budget.ResetDuration)
- }
- }
- // Both fields are required when creating new budget
- if pc.Budget.MaxLimit == nil || pc.Budget.ResetDuration == nil {
- return fmt.Errorf("both max_limit and reset_duration are required when creating a new provider budget")
- }
- }
if err := pc.AllowedModels.Validate(); err != nil {
return &badRequestError{err: fmt.Errorf("invalid allowed_models for provider %s: %w", pc.Provider, err)}
}
@@ -911,25 +909,6 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
AllowAllKeys: allowAllKeys,
Keys: keys,
}
- // Create budget for provider config if provided
- if pc.Budget != nil {
- pcCalAligned := pc.Budget.CalendarAligned != nil && *pc.Budget.CalendarAligned
- budget := configstoreTables.TableBudget{
- ID: uuid.NewString(),
- MaxLimit: *pc.Budget.MaxLimit,
- ResetDuration: *pc.Budget.ResetDuration,
- CalendarAligned: pcCalAligned,
- LastReset: budgetLastReset(pcCalAligned, *pc.Budget.ResetDuration),
- CurrentUsage: 0,
- }
- if err := validateBudget(&budget); err != nil {
- return err
- }
- if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
- return err
- }
- providerConfig.BudgetID = &budget.ID
- }
// Create rate limit for provider config if provided
if pc.RateLimit != nil {
rateLimit := configstoreTables.TableRateLimit{
@@ -952,6 +931,30 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
if err := h.configStore.CreateVirtualKeyProviderConfig(ctx, providerConfig, tx); err != nil {
return err
}
+ // Create multi-budgets for new provider config in update
+ if len(pc.Budgets) > 0 {
+ seenDurations := make(map[string]bool)
+ for _, b := range pc.Budgets {
+ if seenDurations[b.ResetDuration] {
+ return &badRequestError{err: fmt.Errorf("duplicate reset_duration in provider config budgets: %s", b.ResetDuration)}
+ }
+ seenDurations[b.ResetDuration] = true
+ budget := configstoreTables.TableBudget{
+ ID: uuid.NewString(),
+ MaxLimit: b.MaxLimit,
+ ResetDuration: b.ResetDuration,
+ LastReset: budgetLastReset(vk.CalendarAligned, b.ResetDuration),
+ CurrentUsage: 0,
+ ProviderConfigID: &providerConfig.ID,
+ }
+ if err := validateBudget(&budget); err != nil {
+ return err
+ }
+ if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
+ return err
+ }
+ }
+ }
} else {
// Update existing provider config
existing, ok := existingConfigsMap[*pc.ID]
@@ -987,68 +990,72 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
existing.AllowAllKeys = allowAllKeys
existing.Keys = keys
- // Handle budget updates for provider config
- if pc.Budget != nil {
- if isBudgetRemovalRequest(pc.Budget) {
- if existing.BudgetID != nil {
- providerBudgetIDsToDelete = append(providerBudgetIDsToDelete, *existing.BudgetID)
- existing.BudgetID = nil
- existing.Budget = nil
- }
- } else if existing.BudgetID != nil {
- // Update existing budget
- budget := configstoreTables.TableBudget{}
- if err := tx.First(&budget, "id = ?", *existing.BudgetID).Error; err != nil {
- return err
- }
- if pc.Budget.MaxLimit != nil {
- budget.MaxLimit = *pc.Budget.MaxLimit
- }
- if pc.Budget.ResetDuration != nil {
- budget.ResetDuration = *pc.Budget.ResetDuration
- }
- if pc.Budget.CalendarAligned != nil {
- wasCalendarAligned := budget.CalendarAligned
- budget.CalendarAligned = *pc.Budget.CalendarAligned
- if *pc.Budget.CalendarAligned && !wasCalendarAligned {
- budget.LastReset = configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, time.Now())
- budget.CurrentUsage = 0
+ // Handle multi-budget updates for existing provider config
+ if pc.Budgets != nil {
+ // Validate
+ seenDurations := make(map[string]bool)
+ for _, b := range pc.Budgets {
+ if b.MaxLimit < 0 {
+ return &badRequestError{err: fmt.Errorf("provider config budget max_limit cannot be negative: %.2f", b.MaxLimit)}
}
- }
- if err := validateBudget(&budget); err != nil {
- return err
- }
- if err := h.configStore.UpdateBudget(ctx, &budget, tx); err != nil {
- return err
- }
- } else {
- // Create new budget for existing provider config
- if pc.Budget.MaxLimit == nil || pc.Budget.ResetDuration == nil {
- return fmt.Errorf("both max_limit and reset_duration are required when creating a new provider budget")
+ if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil {
+ return &badRequestError{err: fmt.Errorf("invalid provider config budget reset duration format: %s", b.ResetDuration)}
}
- if *pc.Budget.MaxLimit < 0 {
- return fmt.Errorf("provider config budget max_limit cannot be negative: %.2f", *pc.Budget.MaxLimit)
+ if seenDurations[b.ResetDuration] {
+ return &badRequestError{err: fmt.Errorf("duplicate reset_duration in provider config budgets: %s", b.ResetDuration)}
}
- if _, err := configstoreTables.ParseDuration(*pc.Budget.ResetDuration); err != nil {
- return fmt.Errorf("invalid provider config budget reset duration format: %s", *pc.Budget.ResetDuration)
- }
- pcExistCalAligned := pc.Budget.CalendarAligned != nil && *pc.Budget.CalendarAligned
- budget := configstoreTables.TableBudget{
- ID: uuid.NewString(),
- MaxLimit: *pc.Budget.MaxLimit,
- ResetDuration: *pc.Budget.ResetDuration,
- CalendarAligned: pcExistCalAligned,
- LastReset: budgetLastReset(pcExistCalAligned, *pc.Budget.ResetDuration),
- CurrentUsage: 0,
+ seenDurations[b.ResetDuration] = true
}
- if err := validateBudget(&budget); err != nil {
- return err
+
+ // Build map of existing budgets by reset_duration for matching
+ pcExistingByDuration := make(map[string]configstoreTables.TableBudget)
+ for _, eb := range existing.Budgets {
+ pcExistingByDuration[eb.ResetDuration] = eb
}
- if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
- return err
+
+ // Reconcile: preserve existing budgets where possible
+ var pcReconciledBudgets []configstoreTables.TableBudget
+ pcMatchedIDs := make(map[string]bool)
+ for _, b := range pc.Budgets {
+ if eb, found := pcExistingByDuration[b.ResetDuration]; found {
+ // Budget with same duration exists — update max_limit, preserve usage
+ eb.MaxLimit = b.MaxLimit
+ if err := validateBudget(&eb); err != nil {
+ return err
+ }
+ if err := h.configStore.UpdateBudget(ctx, &eb, tx); err != nil {
+ return err
+ }
+ pcReconciledBudgets = append(pcReconciledBudgets, eb)
+ pcMatchedIDs[eb.ID] = true
+ } else {
+ // New budget duration — create fresh
+ budget := configstoreTables.TableBudget{
+ ID: uuid.NewString(),
+ MaxLimit: b.MaxLimit,
+ ResetDuration: b.ResetDuration,
+ LastReset: budgetLastReset(vk.CalendarAligned, b.ResetDuration),
+ CurrentUsage: 0,
+ ProviderConfigID: &existing.ID,
+ }
+ if err := validateBudget(&budget); err != nil {
+ return err
+ }
+ if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil {
+ return err
+ }
+ pcReconciledBudgets = append(pcReconciledBudgets, budget)
+ }
}
- existing.BudgetID = &budget.ID
+ // Delete budgets that are no longer present
+ for _, eb := range existing.Budgets {
+ if !pcMatchedIDs[eb.ID] {
+ if err := h.configStore.DeleteBudget(ctx, eb.ID, tx); err != nil {
+ return fmt.Errorf("failed to delete removed provider config budget: %w", err)
+ }
+ }
}
+ existing.Budgets = pcReconciledBudgets
}
// Handle rate limit updates for provider config
if pc.RateLimit != nil {
@@ -1179,11 +1186,6 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
}
}
- if budgetIDToDelete != "" {
- if err := tx.Delete(&configstoreTables.TableBudget{}, "id = ?", budgetIDToDelete).Error; err != nil {
- return err
- }
- }
if rateLimitIDToDelete != "" {
if err := tx.Delete(&configstoreTables.TableRateLimit{}, "id = ?", rateLimitIDToDelete).Error; err != nil {
return err
@@ -1399,8 +1401,7 @@ func (h *GovernanceHandler) createTeam(ctx *fasthttp.RequestCtx) {
ID: uuid.NewString(),
MaxLimit: req.Budget.MaxLimit,
ResetDuration: req.Budget.ResetDuration,
- CalendarAligned: req.Budget.CalendarAligned,
- LastReset: budgetLastReset(req.Budget.CalendarAligned, req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
@@ -1540,14 +1541,6 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) {
if req.Budget.ResetDuration != nil {
budget.ResetDuration = *req.Budget.ResetDuration
}
- if req.Budget.CalendarAligned != nil {
- wasCalendarAligned := budget.CalendarAligned
- budget.CalendarAligned = *req.Budget.CalendarAligned
- if *req.Budget.CalendarAligned && !wasCalendarAligned {
- budget.LastReset = configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, time.Now())
- budget.CurrentUsage = 0
- }
- }
if err := validateBudget(&budget); err != nil {
return err
}
@@ -1566,13 +1559,11 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) {
if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil {
return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration)
}
- teamCalAligned := req.Budget.CalendarAligned != nil && *req.Budget.CalendarAligned
budget := configstoreTables.TableBudget{
ID: uuid.NewString(),
MaxLimit: *req.Budget.MaxLimit,
ResetDuration: *req.Budget.ResetDuration,
- CalendarAligned: teamCalAligned,
- LastReset: budgetLastReset(teamCalAligned, *req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, *req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
@@ -1811,8 +1802,7 @@ func (h *GovernanceHandler) createCustomer(ctx *fasthttp.RequestCtx) {
ID: uuid.NewString(),
MaxLimit: req.Budget.MaxLimit,
ResetDuration: req.Budget.ResetDuration,
- CalendarAligned: req.Budget.CalendarAligned,
- LastReset: budgetLastReset(req.Budget.CalendarAligned, req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
@@ -1942,14 +1932,6 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) {
if req.Budget.ResetDuration != nil {
budget.ResetDuration = *req.Budget.ResetDuration
}
- if req.Budget.CalendarAligned != nil {
- wasCalendarAligned := budget.CalendarAligned
- budget.CalendarAligned = *req.Budget.CalendarAligned
- if *req.Budget.CalendarAligned && !wasCalendarAligned {
- budget.LastReset = configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, time.Now())
- budget.CurrentUsage = 0
- }
- }
if err := validateBudget(&budget); err != nil {
return err
}
@@ -1968,13 +1950,11 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) {
if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil {
return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration)
}
- custCalAligned := req.Budget.CalendarAligned != nil && *req.Budget.CalendarAligned
budget := configstoreTables.TableBudget{
ID: uuid.NewString(),
MaxLimit: *req.Budget.MaxLimit,
ResetDuration: *req.Budget.ResetDuration,
- CalendarAligned: custCalAligned,
- LastReset: budgetLastReset(custCalAligned, *req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, *req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
@@ -2199,9 +2179,6 @@ func validateBudget(budget *configstoreTables.TableBudget) error {
if _, err := configstoreTables.ParseDuration(budget.ResetDuration); err != nil {
return fmt.Errorf("invalid budget reset duration format: %s", budget.ResetDuration)
}
- if budget.CalendarAligned && !configstoreTables.IsCalendarAlignableDuration(budget.ResetDuration) {
- return fmt.Errorf("calendar_aligned is not supported for reset duration %q: only daily (d), weekly (w), monthly (M), and yearly (Y) periods support calendar alignment", budget.ResetDuration)
- }
return nil
}
@@ -2364,8 +2341,7 @@ func (h *GovernanceHandler) createModelConfig(ctx *fasthttp.RequestCtx) {
ID: uuid.NewString(),
MaxLimit: req.Budget.MaxLimit,
ResetDuration: req.Budget.ResetDuration,
- CalendarAligned: req.Budget.CalendarAligned,
- LastReset: budgetLastReset(req.Budget.CalendarAligned, req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
@@ -2470,14 +2446,6 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) {
if req.Budget.ResetDuration != nil {
budget.ResetDuration = *req.Budget.ResetDuration
}
- if req.Budget.CalendarAligned != nil {
- wasCalendarAligned := budget.CalendarAligned
- budget.CalendarAligned = *req.Budget.CalendarAligned
- if *req.Budget.CalendarAligned && !wasCalendarAligned {
- budget.LastReset = configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, time.Now())
- budget.CurrentUsage = 0
- }
- }
if err := validateBudget(&budget); err != nil {
return err
}
@@ -2496,13 +2464,11 @@ func (h *GovernanceHandler) updateModelConfig(ctx *fasthttp.RequestCtx) {
if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil {
return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration)
}
- mcCalAligned := req.Budget.CalendarAligned != nil && *req.Budget.CalendarAligned
budget := configstoreTables.TableBudget{
ID: uuid.NewString(),
MaxLimit: *req.Budget.MaxLimit,
ResetDuration: *req.Budget.ResetDuration,
- CalendarAligned: mcCalAligned,
- LastReset: budgetLastReset(mcCalAligned, *req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, *req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
@@ -2742,14 +2708,6 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) {
if req.Budget.ResetDuration != nil {
budget.ResetDuration = *req.Budget.ResetDuration
}
- if req.Budget.CalendarAligned != nil {
- wasCalendarAligned := budget.CalendarAligned
- budget.CalendarAligned = *req.Budget.CalendarAligned
- if *req.Budget.CalendarAligned && !wasCalendarAligned {
- budget.LastReset = configstoreTables.GetCalendarPeriodStart(budget.ResetDuration, time.Now())
- budget.CurrentUsage = 0
- }
- }
if err := validateBudget(&budget); err != nil {
return err
}
@@ -2762,13 +2720,11 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) {
if req.Budget.MaxLimit == nil || req.Budget.ResetDuration == nil {
return fmt.Errorf("both max_limit and reset_duration are required when creating a new budget")
}
- provCalAligned := req.Budget.CalendarAligned != nil && *req.Budget.CalendarAligned
budget := configstoreTables.TableBudget{
ID: uuid.NewString(),
MaxLimit: *req.Budget.MaxLimit,
ResetDuration: *req.Budget.ResetDuration,
- CalendarAligned: provCalAligned,
- LastReset: budgetLastReset(provCalAligned, *req.Budget.ResetDuration),
+ LastReset: budgetLastReset(false, *req.Budget.ResetDuration),
CurrentUsage: 0,
}
if err := validateBudget(&budget); err != nil {
diff --git a/transports/bifrost-http/handlers/governance_test.go b/transports/bifrost-http/handlers/governance_test.go
index 0fb9c7e40f..581a22e7b1 100644
--- a/transports/bifrost-http/handlers/governance_test.go
+++ b/transports/bifrost-http/handlers/governance_test.go
@@ -272,7 +272,7 @@ func TestCollectProviderConfigDeleteIDs(t *testing.T) {
{
name: "collects both IDs",
config: configstoreTables.TableVirtualKeyProviderConfig{
- BudgetID: &budgetID,
+ Budgets: []configstoreTables.TableBudget{{ID: budgetID}},
RateLimitID: &rateLimitID,
},
wantBudgetIDs: []string{budgetID},
@@ -281,7 +281,7 @@ func TestCollectProviderConfigDeleteIDs(t *testing.T) {
{
name: "appends to existing slices",
config: configstoreTables.TableVirtualKeyProviderConfig{
- BudgetID: &budgetID,
+ Budgets: []configstoreTables.TableBudget{{ID: budgetID}},
RateLimitID: &rateLimitID,
},
initialBudgetIDs: []string{"budget-0"},
diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go
index 19c8f3b7db..c1d39907b0 100644
--- a/transports/bifrost-http/lib/config.go
+++ b/transports/bifrost-http/lib/config.go
@@ -2237,7 +2237,6 @@ func reconcileVirtualKeyAssociations(
// Update existing provider config from file
existing.Weight = newPC.Weight
existing.AllowedModels = newPC.AllowedModels
- existing.BudgetID = newPC.BudgetID
existing.RateLimitID = newPC.RateLimitID
existing.Keys = newPC.Keys
if err := store.UpdateVirtualKeyProviderConfig(ctx, &existing, tx); err != nil {
diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go
index fe1e3bebee..e0edc1f923 100644
--- a/transports/bifrost-http/lib/config_test.go
+++ b/transports/bifrost-http/lib/config_test.go
@@ -240,8 +240,7 @@ End-to-end tests for virtual key provider configuration operations.
| TestSQLite_VKProviderConfig_KeyReference | VK provider config key references work |
| TestSQLite_VKProviderConfig_HashChangesOnKeyIDChange | Hash changes when key ID changes |
| TestSQLite_VKProviderConfig_WeightAndAllowedModels | Weight and allowed models handled correctly |
-| TestSQLite_VKProviderConfig_BudgetAndRateLimit | BudgetID/RateLimitID persisted correctly |
-| TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit | VK hash includes provider config budget/rate limit |
+| TestGenerateVirtualKeyHash_ProviderConfigRateLimit | VK hash includes provider config rate limit |
===================================================================================
SQLITE INTEGRATION TESTS - VK MCP CONFIGS
@@ -6347,7 +6346,6 @@ func TestProviderHashComparison_BedrockConfigChangedInFile(t *testing.T) {
func TestGenerateVirtualKeyHash(t *testing.T) {
// Create a virtual key
teamID := "team-1"
- budgetID := "budget-1"
vk1 := tables.TableVirtualKey{
ID: "vk-1",
Name: "test-vk",
@@ -6355,7 +6353,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
// Generate hash
@@ -6376,7 +6373,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
hash2, err := configstore.GenerateVirtualKeyHash(vk2)
@@ -6396,7 +6392,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
hash3, err := configstore.GenerateVirtualKeyHash(vk3)
@@ -6416,7 +6411,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_different", // Different value
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
hash4, err := configstore.GenerateVirtualKeyHash(vk4)
@@ -6436,7 +6430,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: false, // Different IsActive
TeamID: &teamID,
- BudgetID: &budgetID,
}
hash5, err := configstore.GenerateVirtualKeyHash(vk5)
@@ -6457,7 +6450,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &differentTeamID, // Different TeamID
- BudgetID: &budgetID,
}
hash6, err := configstore.GenerateVirtualKeyHash(vk6)
@@ -6477,7 +6469,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
hash7, err := configstore.GenerateVirtualKeyHash(vk7)
@@ -6498,7 +6489,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
CustomerID: &customerID, // CustomerID set
}
@@ -6520,7 +6510,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
CustomerID: &differentCustomerID, // Different CustomerID
}
@@ -6533,27 +6522,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
t.Error("Expected different hash for virtual keys with different CustomerID values")
}
- // Different BudgetID should produce different hash
- differentBudgetID := "budget-2"
- vk9 := tables.TableVirtualKey{
- ID: "vk-1",
- Name: "test-vk",
- Description: "Test virtual key",
- Value: "vk_abc123",
- IsActive: true,
- TeamID: &teamID,
- BudgetID: &differentBudgetID, // Different BudgetID
- }
-
- hash9, err := configstore.GenerateVirtualKeyHash(vk9)
- if err != nil {
- t.Fatalf("Failed to generate hash: %v", err)
- }
-
- if hash1 == hash9 {
- t.Error("Expected different hash for virtual keys with different BudgetID")
- }
-
// RateLimitID should produce different hash
rateLimitID := "ratelimit-1"
vk10 := tables.TableVirtualKey{
@@ -6563,7 +6531,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
RateLimitID: &rateLimitID, // RateLimitID set
}
@@ -6585,7 +6552,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
RateLimitID: &differentRateLimitID, // Different RateLimitID
}
@@ -6603,7 +6569,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) {
// TestGenerateVirtualKeyHash_WithProviderConfigs tests hash generation with provider configs
func TestGenerateVirtualKeyHash_WithProviderConfigs(t *testing.T) {
- budgetID := "budget-pc-1"
rateLimitID := "rl-pc-1"
// Virtual key with provider configs
@@ -6620,7 +6585,6 @@ func TestGenerateVirtualKeyHash_WithProviderConfigs(t *testing.T) {
Provider: "openai",
Weight: ptrFloat64(1.0),
AllowedModels: []string{"gpt-4", "gpt-3.5-turbo"},
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
Keys: []tables.TableKey{
{KeyID: "key-1", Name: "key-1"},
@@ -6653,7 +6617,6 @@ func TestGenerateVirtualKeyHash_WithProviderConfigs(t *testing.T) {
Provider: "anthropic", // Different provider
Weight: ptrFloat64(1.0),
AllowedModels: []string{"claude-3"},
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
},
},
@@ -6682,7 +6645,6 @@ func TestGenerateVirtualKeyHash_WithProviderConfigs(t *testing.T) {
Provider: "openai",
Weight: ptrFloat64(2.0), // Different weight
AllowedModels: []string{"gpt-4", "gpt-3.5-turbo"},
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
Keys: []tables.TableKey{
{KeyID: "key-1", Name: "key-1"},
@@ -6786,7 +6748,6 @@ func TestGenerateVirtualKeyHash_WithMCPConfigs(t *testing.T) {
// TestVirtualKeyHashComparison_MatchingHash tests that DB config is kept when hashes match
func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) {
teamID := "team-1"
- budgetID := "budget-1"
// Create a virtual key (simulating what's in config.json)
fileVK := tables.TableVirtualKey{
@@ -6796,7 +6757,6 @@ func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
// Generate file hash
@@ -6807,7 +6767,6 @@ func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) {
// Create DB virtual key with same content (simulating existing DB record)
dbTeamID := "team-1"
- dbBudgetID := "budget-1"
dbVK := tables.TableVirtualKey{
ID: "vk-1",
Name: "test-vk",
@@ -6815,7 +6774,6 @@ func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &dbTeamID,
- BudgetID: &dbBudgetID,
ConfigHash: fileHash, // Same hash as file
}
@@ -6840,7 +6798,6 @@ func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) {
// TestVirtualKeyHashComparison_DifferentHash tests that file config is used when hashes differ
func TestVirtualKeyHashComparison_DifferentHash(t *testing.T) {
teamID := "team-1"
- budgetID := "budget-1"
// Create DB virtual key with old config
dbVK := tables.TableVirtualKey{
@@ -6850,7 +6807,6 @@ func TestVirtualKeyHashComparison_DifferentHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
dbHash, err := configstore.GenerateVirtualKeyHash(dbVK)
@@ -6861,7 +6817,6 @@ func TestVirtualKeyHashComparison_DifferentHash(t *testing.T) {
// Create file virtual key with updated config
fileTeamID := "team-1"
- fileBudgetID := "budget-1"
fileVK := tables.TableVirtualKey{
ID: "vk-1",
Name: "new-name", // Updated name
@@ -6869,7 +6824,6 @@ func TestVirtualKeyHashComparison_DifferentHash(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &fileTeamID,
- BudgetID: &fileBudgetID,
}
fileHash, err := configstore.GenerateVirtualKeyHash(fileVK)
@@ -7041,26 +6995,6 @@ func TestVirtualKeyHashComparison_OptionalFieldsPresence(t *testing.T) {
t.Error("Expected different hash for team_id vs customer_id")
}
- // Virtual key with budget_id
- budgetID := "budget-1"
- vkWithBudget := tables.TableVirtualKey{
- ID: "vk-1",
- Name: "test-vk",
- Description: "",
- Value: "vk_abc123",
- IsActive: true,
- BudgetID: &budgetID,
- }
-
- hashWithBudget, err := configstore.GenerateVirtualKeyHash(vkWithBudget)
- if err != nil {
- t.Fatalf("Failed to generate hash: %v", err)
- }
-
- if hashNoOptional == hashWithBudget {
- t.Error("Expected different hash when budget_id is added")
- }
-
// Virtual key with rate_limit_id
rateLimitID := "rl-1"
vkWithRateLimit := tables.TableVirtualKey{
@@ -7087,7 +7021,6 @@ func TestVirtualKeyHashComparison_OptionalFieldsPresence(t *testing.T) {
// TestVirtualKeyHashComparison_FieldValueChanges tests hash changes when field values change
func TestVirtualKeyHashComparison_FieldValueChanges(t *testing.T) {
teamID := "team-1"
- budgetID := "budget-1"
// Base virtual key
baseVK := tables.TableVirtualKey{
@@ -7097,7 +7030,6 @@ func TestVirtualKeyHashComparison_FieldValueChanges(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
}
baseHash, err := configstore.GenerateVirtualKeyHash(baseVK)
@@ -7145,27 +7077,12 @@ func TestVirtualKeyHashComparison_FieldValueChanges(t *testing.T) {
t.Error("Expected different hash when TeamID value changes")
}
- // Change BudgetID value
- newBudgetID := "budget-2"
- vkChangedBudget := baseVK
- vkChangedBudget.BudgetID = &newBudgetID
-
- hashChangedBudget, err := configstore.GenerateVirtualKeyHash(vkChangedBudget)
- if err != nil {
- t.Fatalf("Failed to generate hash: %v", err)
- }
-
- if baseHash == hashChangedBudget {
- t.Error("Expected different hash when BudgetID value changes")
- }
-
t.Log("✓ Field value changes correctly detected in hash")
}
// TestVirtualKeyHashComparison_RoundTrip tests JSON → DB → same JSON produces no changes
func TestVirtualKeyHashComparison_RoundTrip(t *testing.T) {
teamID := "team-1"
- budgetID := "budget-1"
rateLimitID := "rl-1"
// Original config.json virtual key
@@ -7176,7 +7093,6 @@ func TestVirtualKeyHashComparison_RoundTrip(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &teamID,
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
{
@@ -7199,7 +7115,6 @@ func TestVirtualKeyHashComparison_RoundTrip(t *testing.T) {
// Same config.json on reload (simulating app restart)
reloadTeamID := "team-1"
- reloadBudgetID := "budget-1"
reloadRateLimitID := "rl-1"
reloadVK := tables.TableVirtualKey{
ID: "vk-1",
@@ -7208,7 +7123,6 @@ func TestVirtualKeyHashComparison_RoundTrip(t *testing.T) {
Value: "vk_abc123",
IsActive: true,
TeamID: &reloadTeamID,
- BudgetID: &reloadBudgetID,
RateLimitID: &reloadRateLimitID,
ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
{
@@ -8720,7 +8634,6 @@ func TestSQLite_FullLifecycle_InitialLoad(t *testing.T) {
Description: "Test virtual key 1",
Value: "vk_test123",
IsActive: true,
- BudgetID: &budgetID,
RateLimitID: &rateLimitID,
ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
{
@@ -10435,9 +10348,6 @@ func TestGenerateKeyHash_StableOrdering(t *testing.T) {
// TestGenerateVirtualKeyHash_StableProviderConfigOrdering verifies hash stability with different provider config orderings
func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
- budgetID1 := "budget-1"
- budgetID2 := "budget-2"
-
// VK with provider configs in order A
vkOrderA := tables.TableVirtualKey{
ID: "vk-1",
@@ -10452,7 +10362,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
Provider: "openai",
Weight: ptrFloat64(1.0),
AllowedModels: []string{"gpt-4"},
- BudgetID: &budgetID1,
},
{
ID: 2,
@@ -10460,7 +10369,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
Provider: "anthropic",
Weight: ptrFloat64(2.0),
AllowedModels: []string{"claude-3"},
- BudgetID: &budgetID2,
},
{
ID: 3,
@@ -10493,7 +10401,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
Provider: "anthropic",
Weight: ptrFloat64(2.0),
AllowedModels: []string{"claude-3"},
- BudgetID: &budgetID2,
},
{
ID: 1,
@@ -10501,7 +10408,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
Provider: "openai",
Weight: ptrFloat64(1.0),
AllowedModels: []string{"gpt-4"},
- BudgetID: &budgetID1,
},
},
}
@@ -10520,7 +10426,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
Provider: "anthropic",
Weight: ptrFloat64(2.0),
AllowedModels: []string{"claude-3"},
- BudgetID: &budgetID2,
},
{
ID: 1,
@@ -10528,7 +10433,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) {
Provider: "openai",
Weight: ptrFloat64(1.0),
AllowedModels: []string{"gpt-4"},
- BudgetID: &budgetID1,
},
{
ID: 3,
@@ -10941,8 +10845,6 @@ func TestGenerateVirtualKeyHash_StableToolsToExecuteOrdering(t *testing.T) {
// TestGenerateVirtualKeyHash_StableCombinedOrdering verifies hash stability with all nested orderings randomized
func TestGenerateVirtualKeyHash_StableCombinedOrdering(t *testing.T) {
- budgetID := "budget-1"
-
// VK with all elements in order A
vkOrderA := tables.TableVirtualKey{
ID: "vk-1",
@@ -10950,7 +10852,6 @@ func TestGenerateVirtualKeyHash_StableCombinedOrdering(t *testing.T) {
Description: "Test virtual key",
Value: "vk_abc123",
IsActive: true,
- BudgetID: &budgetID,
ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
{
ID: 1,
@@ -10993,7 +10894,6 @@ func TestGenerateVirtualKeyHash_StableCombinedOrdering(t *testing.T) {
Description: "Test virtual key",
Value: "vk_abc123",
IsActive: true,
- BudgetID: &budgetID,
ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
{
ID: 2,
@@ -14081,11 +13981,9 @@ func TestSQLite_Key_UseForBatchAPIChange_Detected(t *testing.T) {
}
}
-// TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit verifies that BudgetID and RateLimitID
-// in VK provider configs affect hash generation.
-func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
- budgetID1 := "budget-1"
- budgetID2 := "budget-2"
+// TestGenerateVirtualKeyHash_ProviderConfigRateLimit verifies that RateLimitID
+// in VK provider configs affects hash generation.
+func TestGenerateVirtualKeyHash_ProviderConfigRateLimit(t *testing.T) {
rateLimitID1 := "rate-limit-1"
rateLimitID2 := "rate-limit-2"
weight := 1.0
@@ -14096,34 +13994,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
vk2 tables.TableVirtualKey
expectEqual bool
}{
- {
- name: "different_budget_id_different_hash",
- vk1: tables.TableVirtualKey{
- ID: "vk-1",
- Name: "test-vk",
- IsActive: true,
- ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
- {
- Provider: "openai",
- Weight: &weight,
- BudgetID: &budgetID1,
- },
- },
- },
- vk2: tables.TableVirtualKey{
- ID: "vk-1",
- Name: "test-vk",
- IsActive: true,
- ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
- {
- Provider: "openai",
- Weight: &weight,
- BudgetID: &budgetID2,
- },
- },
- },
- expectEqual: false,
- },
{
name: "different_rate_limit_id_different_hash",
vk1: tables.TableVirtualKey{
@@ -14152,34 +14022,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
},
expectEqual: false,
},
- {
- name: "nil_vs_set_budget_id_different_hash",
- vk1: tables.TableVirtualKey{
- ID: "vk-1",
- Name: "test-vk",
- IsActive: true,
- ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
- {
- Provider: "openai",
- Weight: &weight,
- BudgetID: nil,
- },
- },
- },
- vk2: tables.TableVirtualKey{
- ID: "vk-1",
- Name: "test-vk",
- IsActive: true,
- ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
- {
- Provider: "openai",
- Weight: &weight,
- BudgetID: &budgetID1,
- },
- },
- },
- expectEqual: false,
- },
{
name: "nil_vs_set_rate_limit_id_different_hash",
vk1: tables.TableVirtualKey{
@@ -14209,7 +14051,7 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
expectEqual: false,
},
{
- name: "same_budget_and_rate_limit_same_hash",
+ name: "same_rate_limit_same_hash",
vk1: tables.TableVirtualKey{
ID: "vk-1",
Name: "test-vk",
@@ -14218,7 +14060,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
{
Provider: "openai",
Weight: &weight,
- BudgetID: &budgetID1,
RateLimitID: &rateLimitID1,
},
},
@@ -14231,7 +14072,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
{
Provider: "openai",
Weight: &weight,
- BudgetID: &budgetID1,
RateLimitID: &rateLimitID1,
},
},
@@ -14262,107 +14102,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) {
}
}
-// TestSQLite_VKProviderConfig_BudgetAndRateLimit verifies that BudgetID and RateLimitID
-// in VK provider configs are properly persisted and retrieved from SQLite.
-func TestSQLite_VKProviderConfig_BudgetAndRateLimit(t *testing.T) {
- initTestLogger()
- tempDir := createTempDir(t)
-
- budgetID := "budget-123"
- rateLimitID := "rate-limit-456"
- vkID := uuid.NewString()
- weight := 1.0
-
- // Create config with VK that has provider config with BudgetID and RateLimitID
- configData := makeConfigDataFullWithDir(
- nil,
- map[string]configstore.ProviderConfig{
- "openai": {
- Keys: []schemas.Key{
- {
- ID: uuid.NewString(),
- Name: "openai-key",
- Value: *schemas.NewEnvVar("sk-test"),
- Weight: 1,
- },
- },
- },
- },
- &configstore.GovernanceConfig{
- Budgets: []tables.TableBudget{
- {
- ID: budgetID,
- MaxLimit: 100.0,
- },
- },
- RateLimits: []tables.TableRateLimit{
- {
- ID: rateLimitID,
- RequestMaxLimit: int64Ptr(60),
- TokenMaxLimit: int64Ptr(10000),
- },
- },
- VirtualKeys: []tables.TableVirtualKey{
- {
- ID: vkID,
- Name: "test-vk",
- Value: "vk-test-value",
- IsActive: true,
- ProviderConfigs: []tables.TableVirtualKeyProviderConfig{
- {
- Provider: "openai",
- Weight: &weight,
- BudgetID: &budgetID,
- RateLimitID: &rateLimitID,
- },
- },
- },
- },
- },
- tempDir,
- )
-
- // Load config
- createConfigFile(t, tempDir, configData)
- config, err := LoadConfig(context.Background(), tempDir)
- if err != nil {
- t.Fatalf("LoadConfig failed: %v", err)
- }
- defer config.Close(context.Background())
-
- // Verify the governance config has the VK with provider configs
- if config.GovernanceConfig == nil {
- t.Fatal("Expected GovernanceConfig to exist")
- }
- if len(config.GovernanceConfig.VirtualKeys) == 0 {
- t.Fatal("Expected VirtualKeys in GovernanceConfig")
- }
-
- // Find the VK and verify provider config
- var foundVK *tables.TableVirtualKey
- for i := range config.GovernanceConfig.VirtualKeys {
- if config.GovernanceConfig.VirtualKeys[i].ID == vkID {
- foundVK = &config.GovernanceConfig.VirtualKeys[i]
- break
- }
- }
- if foundVK == nil {
- t.Fatalf("Virtual key %s not found in config", vkID)
- }
-
- if len(foundVK.ProviderConfigs) == 0 {
- t.Fatal("Expected VK to have provider configs")
- }
-
- pc := foundVK.ProviderConfigs[0]
- if pc.BudgetID == nil || *pc.BudgetID != budgetID {
- t.Errorf("Expected BudgetID=%s, got %v", budgetID, pc.BudgetID)
- }
- if pc.RateLimitID == nil || *pc.RateLimitID != rateLimitID {
- t.Errorf("Expected RateLimitID=%s, got %v", rateLimitID, pc.RateLimitID)
- }
-}
-
// intPtr is a helper to create a pointer to an int
func intPtr(i int) *int {
return &i
@@ -15378,9 +15117,11 @@ var excludedGoFields = map[string]map[string]bool{
},
// Table types have DB-specific fields
"tables.TableBudget": {
- "config_hash": true,
- "created_at": true,
- "updated_at": true,
+ "config_hash": true,
+ "created_at": true,
+ "updated_at": true,
+ "virtual_key_id": true, // Internal DB FK for multi-budget ownership
+ "provider_config_id": true, // Internal DB FK for multi-budget ownership
},
"tables.TableRateLimit": {
"config_hash": true,
@@ -15409,16 +15150,16 @@ var excludedGoFields = map[string]map[string]bool{
"config_hash": true,
"created_at": true,
"updated_at": true,
- "budget": true, // GORM relation
+ "budgets": true, // GORM relation (budgets have virtual_key_id FK)
"rate_limit": true, // GORM relation
"team": true, // GORM relation
"customer": true, // GORM relation
},
"tables.TableVirtualKeyProviderConfig": {
- "budget": true, // GORM relation
- "rate_limit": true, // GORM relation
+ "rate_limit": true, // GORM relation
"allow_all_keys": true, // Internal DB field; users configure via key_ids
- "keys": true, // GORM many2many relation; users configure via key_ids
+ "keys": true, // GORM many2many relation; users configure via key_ids
+ "budgets": true, // GORM relation (budgets have provider_config_id FK)
},
"tables.TableVirtualKeyMCPConfig": {
"mcp_client": true, // GORM relation
diff --git a/transports/config.schema.json b/transports/config.schema.json
index df417ddd47..bc8b267f15 100644
--- a/transports/config.schema.json
+++ b/transports/config.schema.json
@@ -165,7 +165,10 @@
},
"mcp_code_mode_binding_level": {
"type": "string",
- "enum": ["server", "tool"],
+ "enum": [
+ "server",
+ "tool"
+ ],
"description": "Code mode binding level for MCP tools"
},
"mcp_tool_sync_interval": {
@@ -291,6 +294,14 @@
"type": "boolean",
"description": "Snap resets to calendar boundaries (day/week/month/year start)",
"default": false
+ },
+ "virtual_key_id": {
+ "type": "string",
+ "description": "ID of the virtual key this budget belongs to (mutually exclusive with provider_config_id)"
+ },
+ "provider_config_id": {
+ "type": "integer",
+ "description": "ID of the provider config this budget belongs to (mutually exclusive with virtual_key_id)"
}
},
"required": [
@@ -457,6 +468,11 @@
"description": "Whether the virtual key is active",
"default": true
},
+ "calendar_aligned": {
+ "type": "boolean",
+ "description": "Snap all budget resets to calendar boundaries (day, week, month, year)",
+ "default": false
+ },
"team_id": {
"type": "string",
"description": "Associated team ID (mutually exclusive with customer_id)"
@@ -465,10 +481,6 @@
"type": "string",
"description": "Associated customer ID (mutually exclusive with team_id)"
},
- "budget_id": {
- "type": "string",
- "description": "Associated budget ID"
- },
"rate_limit_id": {
"type": "string",
"description": "Associated rate limit ID"
@@ -577,7 +589,9 @@
"description": "Capture raw request/response for internal logging only; strip from API responses returned to clients (default: false)"
}
},
- "required": ["name"]
+ "required": [
+ "name"
+ ]
}
}
},
@@ -932,7 +946,10 @@
},
"placement": {
"type": "string",
- "enum": ["pre_builtin", "post_builtin"],
+ "enum": [
+ "pre_builtin",
+ "post_builtin"
+ ],
"description": "Whether this plugin runs before or after built-in plugins. Default: post_builtin",
"optional": true,
"default": "post_builtin"
@@ -1012,10 +1029,15 @@
"description": "Password for basic authentication"
}
},
- "required": ["username", "password"]
+ "required": [
+ "username",
+ "password"
+ ]
}
},
- "required": ["push_gateway_url"]
+ "required": [
+ "push_gateway_url"
+ ]
}
},
"additionalProperties": false
@@ -1474,7 +1496,9 @@
"maximum": 1
}
},
- "required": ["weight"],
+ "required": [
+ "weight"
+ ],
"additionalProperties": false
},
"routing_rule": {
@@ -1521,12 +1545,17 @@
},
"scope": {
"type": "string",
- "enum": ["global", "team", "customer", "virtual_key"],
+ "enum": [
+ "global",
+ "team",
+ "customer",
+ "virtual_key"
+ ],
"description": "Rule scope level",
"default": "global"
},
"scope_id": {
- "type": ["string", "null"],
+ "type": "string",
"description": "Entity ID for non-global scopes (required for non-global scope)"
},
"priority": {
@@ -1540,8 +1569,37 @@
"additionalProperties": true
}
},
- "required": ["id", "name", "targets"],
- "additionalProperties": false
+ "required": [
+ "id",
+ "name",
+ "targets"
+ ],
+ "additionalProperties": false,
+ "if": {
+ "properties": {
+ "scope": {
+ "enum": [
+ "team",
+ "customer",
+ "virtual_key"
+ ]
+ }
+ },
+ "required": [
+ "scope"
+ ]
+ },
+ "then": {
+ "required": [
+ "scope_id"
+ ],
+ "properties": {
+ "scope_id": {
+ "type": "string",
+ "minLength": 1
+ }
+ }
+ }
},
"virtual_key_provider_config": {
"type": "object",
@@ -1560,7 +1618,10 @@
"description": "Provider name"
},
"weight": {
- "type": ["number", "null"],
+ "type": [
+ "number",
+ "null"
+ ],
"description": "Weight for load balancing (null opts out of weighted routing)",
"default": null
},
@@ -1571,10 +1632,6 @@
"type": "string"
}
},
- "budget_id": {
- "type": "string",
- "description": "Associated budget ID"
- },
"rate_limit_id": {
"type": "string",
"description": "Associated rate limit ID"
@@ -2486,7 +2543,11 @@
},
"auth_type": {
"type": "string",
- "enum": ["none", "headers", "oauth"],
+ "enum": [
+ "none",
+ "headers",
+ "oauth"
+ ],
"description": "Authentication type for MCP connection"
},
"oauth_config_id": {
@@ -2632,8 +2693,16 @@
}
},
"anyOf": [
- { "required": ["http_config"] },
- { "required": ["connection_string"] }
+ {
+ "required": [
+ "http_config"
+ ]
+ },
+ {
+ "required": [
+ "connection_string"
+ ]
+ }
]
},
{
@@ -2665,7 +2734,10 @@
},
"code_mode_binding_level": {
"type": "string",
- "enum": ["server", "tool"],
+ "enum": [
+ "server",
+ "tool"
+ ],
"description": "How tools are exposed in VFS for code execution"
},
"disable_auto_tool_inject": {
@@ -3148,7 +3220,11 @@
},
"cloud": {
"type": "string",
- "enum": ["commercial", "gcc-high", "dod"],
+ "enum": [
+ "commercial",
+ "gcc-high",
+ "dod"
+ ],
"default": "commercial",
"description": "Cloud environment: 'commercial' (default), 'gcc-high' for US Government GCC High, or 'dod' for Department of Defense"
},
@@ -3412,7 +3488,14 @@
"scope_kind": {
"type": "string",
"description": "Scope level for this override",
- "enum": ["global", "provider", "provider_key", "virtual_key", "virtual_key_provider", "virtual_key_provider_key"]
+ "enum": [
+ "global",
+ "provider",
+ "provider_key",
+ "virtual_key",
+ "virtual_key_provider",
+ "virtual_key_provider_key"
+ ]
},
"virtual_key_id": {
"type": "string",
@@ -3429,7 +3512,10 @@
"match_type": {
"type": "string",
"description": "How the pattern is matched against model names",
- "enum": ["exact", "wildcard"]
+ "enum": [
+ "exact",
+ "wildcard"
+ ]
},
"pattern": {
"type": "string",
@@ -3452,21 +3538,38 @@
"description": "Internal hash for change detection (auto-managed)"
}
},
- "required": ["id", "name", "scope_kind", "match_type", "pattern", "request_types"],
+ "required": [
+ "id",
+ "name",
+ "scope_kind",
+ "match_type",
+ "pattern",
+ "request_types"
+ ],
"additionalProperties": false
},
"pricing_override_match_type": {
"type": "string",
- "enum": ["exact", "wildcard"]
+ "enum": [
+ "exact",
+ "wildcard"
+ ]
},
"pricing_override_request_type": {
"type": "string",
"enum": [
- "chat_completion", "text_completion", "responses",
- "embedding", "rerank",
- "speech", "transcription",
- "image_generation", "image_variation", "image_edit",
- "video_generation", "video_remix"
+ "chat_completion",
+ "text_completion",
+ "responses",
+ "embedding",
+ "rerank",
+ "speech",
+ "transcription",
+ "image_generation",
+ "image_variation",
+ "image_edit",
+ "video_generation",
+ "video_remix"
]
},
"custom_provider_config": {
@@ -3485,65 +3588,167 @@
"type": "object",
"description": "Allowed request types for the custom provider",
"properties": {
- "list_models": { "type": "boolean" },
- "text_completion": { "type": "boolean" },
- "text_completion_stream": { "type": "boolean" },
- "chat_completion": { "type": "boolean" },
- "chat_completion_stream": { "type": "boolean" },
- "responses": { "type": "boolean" },
- "responses_stream": { "type": "boolean" },
- "count_tokens": { "type": "boolean" },
- "embedding": { "type": "boolean" },
- "rerank": { "type": "boolean" },
- "speech": { "type": "boolean" },
- "speech_stream": { "type": "boolean" },
- "transcription": { "type": "boolean" },
- "transcription_stream": { "type": "boolean" },
- "image_generation": { "type": "boolean" },
- "image_generation_stream": { "type": "boolean" },
- "image_edit": { "type": "boolean" },
- "image_edit_stream": { "type": "boolean" },
- "image_variation": { "type": "boolean" },
- "video_generation": { "type": "boolean" },
- "video_retrieve": { "type": "boolean" },
- "video_download": { "type": "boolean" },
- "video_delete": { "type": "boolean" },
- "video_list": { "type": "boolean" },
- "video_remix": { "type": "boolean" },
- "batch_create": { "type": "boolean" },
- "batch_list": { "type": "boolean" },
- "batch_retrieve": { "type": "boolean" },
- "batch_cancel": { "type": "boolean" },
- "batch_delete": { "type": "boolean" },
- "batch_results": { "type": "boolean" },
- "file_upload": { "type": "boolean" },
- "file_list": { "type": "boolean" },
- "file_retrieve": { "type": "boolean" },
- "file_delete": { "type": "boolean" },
- "file_content": { "type": "boolean" },
- "container_create": { "type": "boolean" },
- "container_list": { "type": "boolean" },
- "container_retrieve": { "type": "boolean" },
- "container_delete": { "type": "boolean" },
- "container_file_create": { "type": "boolean" },
- "container_file_list": { "type": "boolean" },
- "container_file_retrieve": { "type": "boolean" },
- "container_file_content": { "type": "boolean" },
- "container_file_delete": { "type": "boolean" },
- "passthrough": { "type": "boolean" },
- "passthrough_stream": { "type": "boolean" },
- "websocket_responses": { "type": "boolean" },
- "realtime": { "type": "boolean" }
+ "list_models": {
+ "type": "boolean"
+ },
+ "text_completion": {
+ "type": "boolean"
+ },
+ "text_completion_stream": {
+ "type": "boolean"
+ },
+ "chat_completion": {
+ "type": "boolean"
+ },
+ "chat_completion_stream": {
+ "type": "boolean"
+ },
+ "responses": {
+ "type": "boolean"
+ },
+ "responses_stream": {
+ "type": "boolean"
+ },
+ "count_tokens": {
+ "type": "boolean"
+ },
+ "embedding": {
+ "type": "boolean"
+ },
+ "rerank": {
+ "type": "boolean"
+ },
+ "speech": {
+ "type": "boolean"
+ },
+ "speech_stream": {
+ "type": "boolean"
+ },
+ "transcription": {
+ "type": "boolean"
+ },
+ "transcription_stream": {
+ "type": "boolean"
+ },
+ "image_generation": {
+ "type": "boolean"
+ },
+ "image_generation_stream": {
+ "type": "boolean"
+ },
+ "image_edit": {
+ "type": "boolean"
+ },
+ "image_edit_stream": {
+ "type": "boolean"
+ },
+ "image_variation": {
+ "type": "boolean"
+ },
+ "video_generation": {
+ "type": "boolean"
+ },
+ "video_retrieve": {
+ "type": "boolean"
+ },
+ "video_download": {
+ "type": "boolean"
+ },
+ "video_delete": {
+ "type": "boolean"
+ },
+ "video_list": {
+ "type": "boolean"
+ },
+ "video_remix": {
+ "type": "boolean"
+ },
+ "batch_create": {
+ "type": "boolean"
+ },
+ "batch_list": {
+ "type": "boolean"
+ },
+ "batch_retrieve": {
+ "type": "boolean"
+ },
+ "batch_cancel": {
+ "type": "boolean"
+ },
+ "batch_delete": {
+ "type": "boolean"
+ },
+ "batch_results": {
+ "type": "boolean"
+ },
+ "file_upload": {
+ "type": "boolean"
+ },
+ "file_list": {
+ "type": "boolean"
+ },
+ "file_retrieve": {
+ "type": "boolean"
+ },
+ "file_delete": {
+ "type": "boolean"
+ },
+ "file_content": {
+ "type": "boolean"
+ },
+ "container_create": {
+ "type": "boolean"
+ },
+ "container_list": {
+ "type": "boolean"
+ },
+ "container_retrieve": {
+ "type": "boolean"
+ },
+ "container_delete": {
+ "type": "boolean"
+ },
+ "container_file_create": {
+ "type": "boolean"
+ },
+ "container_file_list": {
+ "type": "boolean"
+ },
+ "container_file_retrieve": {
+ "type": "boolean"
+ },
+ "container_file_content": {
+ "type": "boolean"
+ },
+ "container_file_delete": {
+ "type": "boolean"
+ },
+ "passthrough": {
+ "type": "boolean"
+ },
+ "passthrough_stream": {
+ "type": "boolean"
+ },
+ "websocket_responses": {
+ "type": "boolean"
+ },
+ "realtime": {
+ "type": "boolean"
+ }
},
"additionalProperties": false
},
"request_path_overrides": {
"type": "object",
"description": "Mapping of request type to custom path overriding the default provider path",
- "additionalProperties": { "type": "string" }
+ "additionalProperties": {
+ "type": "string"
+ }
}
},
- "required": ["base_provider_type"],
+ "required": [
+ "base_provider_type"
+ ],
"additionalProperties": false
}
}
diff --git a/ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx b/ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
new file mode 100644
index 0000000000..e7747a2742
--- /dev/null
+++ b/ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
@@ -0,0 +1,17 @@
+import { ShieldCheck } from "lucide-react";
+import ContactUsView from "../views/contactUsView";
+
+export default function AccessProfilesIndexView() {
+ return (
+