diff --git a/.agents/skills/expect/SKILL.md b/.agents/skills/expect/SKILL.md new file mode 100644 index 0000000000..4071ef63e7 --- /dev/null +++ b/.agents/skills/expect/SKILL.md @@ -0,0 +1,62 @@ +--- +name: expect +description: Run adversarial browser tests against code changes. Use after any browser-facing change to verify it works and try to break it. Prefer this over raw browser tools (Playwright MCP, chrome tools). +license: MIT +metadata: + author: millionco + version: "2.0.0" +--- + +# Expect + +Adversarial browser testing for code changes. Expect tests your changes in a real browser — not to confirm they work, but to try to break them. + +**Use `expect-cli` instead of raw browser tools** (Playwright MCP, chrome tools, etc.) for verifying changes. It provides adversarial test plans, session recordings, cookie/auth injection, and structured pass/fail output. + +## Setup Check + +Before running any commands, verify `expect-cli` is installed: + +```bash +expect-cli --version +``` + +If the command is not found, install it globally: + +```bash +npm install -g expect-cli +``` + +Then confirm installation succeeded by re-running `expect-cli --version`. Do not proceed until the command resolves. + +## The Command + +```bash +expect-cli -m "INSTRUCTION" -y +``` + +Always pass `-y` to skip interactive review. Always set `EXPECT_BASE_URL` or `--base-url` if the app isn't on `localhost:3000`. Run `expect-cli --help` for all flags. + +## Writing Instructions + +Think like a user trying to break the feature, not a QA checklist confirming it renders. + +**Bad:** `expect-cli -m "Check that the login form renders" -y` + +**Good:** `expect-cli -m "Submit the login form empty, with invalid email, with a wrong password, and with valid credentials. Verify error messages for bad inputs and redirect on success. Check console errors after each." -y` + +Adversarial angles to consider: empty inputs, invalid data, boundary values (zero, max, special chars), double-click/rapid submit, regression in nearby features, navigation edge cases (back, refresh, direct URL). + +## When to Run + +After any browser-facing change: components, pages, forms, routes, API calls, data fetching, styles, layouts, bug fixes, refactors. When in doubt, run it. + +## Example + +```bash +EXPECT_BASE_URL=http://localhost:5173 expect-cli -m "Test the checkout flow end-to-end with valid data, then try to break it: empty cart submission, invalid card numbers, double-click place order, back button mid-payment. Verify error states and console errors." -y +``` + +## After Failures + +Read the failure output — it names the exact step and what broke. Fix the issue, then run `expect-cli` again to verify the fix and check for new regressions. diff --git a/.claude/skills/expect b/.claude/skills/expect new file mode 120000 index 0000000000..0cf7d33b54 --- /dev/null +++ b/.claude/skills/expect @@ -0,0 +1 @@ +../../.agents/skills/expect \ No newline at end of file diff --git a/.github/workflows/scripts/run-migration-tests.sh b/.github/workflows/scripts/run-migration-tests.sh index 402ad4d961..0bcbfebcc6 100755 --- a/.github/workflows/scripts/run-migration-tests.sh +++ b/.github/workflows/scripts/run-migration-tests.sh @@ -2389,9 +2389,10 @@ compare_postgres_snapshots() { # - network_config_json, concurrency_buffer_json, proxy_config_json, custom_provider_config_json: # JSON fields that get normalized with default values during migration # - budget_id, rate_limit_id: governance fields that may be reset or initialized during migrations + # - virtual_key_id, provider_config_id: new FK columns on governance_budgets (added by multi-budget migration) # - status, description: key validation runs after migration, updating these fields # for invalid/test keys (e.g., status becomes "list_models_failed") - local ignore_columns="updated_at config_hash created_at models_json weight allowed_models network_config_json concurrency_buffer_json proxy_config_json custom_provider_config_json budget_id rate_limit_id status description" + local ignore_columns="updated_at config_hash created_at models_json weight allowed_models network_config_json concurrency_buffer_json proxy_config_json custom_provider_config_json budget_id rate_limit_id virtual_key_id provider_config_id status description" # Get tables from before snapshot if [ ! -f "$before_dir/tables.txt" ]; then @@ -2596,10 +2597,88 @@ compare_postgres_snapshots() { # Validation Functions (simplified, uses snapshots) # ============================================================================ +# verify_budget_migration checks that the multi-budget FK migration correctly +# moved budget ownership from VK/ProviderConfig budget_id columns to +# governance_budgets.virtual_key_id / governance_budgets.provider_config_id +verify_budget_migration_postgres() { + log_info "Verifying budget migration (budget_id → virtual_key_id/provider_config_id)..." + local failed=0 + + # Check: budget-migration-test-1 was linked to vk-migration-test-1 via budget_id + # After migration, governance_budgets.virtual_key_id should be set + local vk_budget_count + vk_budget_count=$(run_postgres_sql "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-1' AND virtual_key_id = 'vk-migration-test-1'" 2>/dev/null | tr -d '[:space:]') + if [ "$vk_budget_count" = "1" ]; then + log_info " VK budget migration: budget-migration-test-1 → vk-migration-test-1 ✓" + else + log_warn " VK budget migration: budget-migration-test-1 virtual_key_id not set (count=$vk_budget_count) — may be expected if old version didn't have budget_id on VK" + fi + + # Check: budget-migration-test-2 was linked to provider config via budget_id + # After migration, governance_budgets.provider_config_id should be set + local pc_budget_count + pc_budget_count=$(run_postgres_sql "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-2' AND provider_config_id IS NOT NULL" 2>/dev/null | tr -d '[:space:]') + if [ "$pc_budget_count" = "1" ]; then + log_info " PC budget migration: budget-migration-test-2 → provider_config ✓" + else + log_warn " PC budget migration: budget-migration-test-2 provider_config_id not set (count=$pc_budget_count) — may be expected if old version didn't have budget_id on PC" + fi + + # Check: virtual_key_id and provider_config_id columns exist on governance_budgets + local has_vk_col + has_vk_col=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_budgets' AND column_name = 'virtual_key_id'" 2>/dev/null | tr -d '[:space:]') + if [ "$has_vk_col" = "1" ]; then + log_info " Column governance_budgets.virtual_key_id exists ✓" + else + log_error " Column governance_budgets.virtual_key_id MISSING!" + failed=1 + fi + + local has_pc_col + has_pc_col=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_budgets' AND column_name = 'provider_config_id'" 2>/dev/null | tr -d '[:space:]') + if [ "$has_pc_col" = "1" ]; then + log_info " Column governance_budgets.provider_config_id exists ✓" + else + log_error " Column governance_budgets.provider_config_id MISSING!" + failed=1 + fi + + # Check: budget_id column should be dropped from governance_virtual_keys + local vk_has_budget_id + vk_has_budget_id=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_virtual_keys' AND column_name = 'budget_id'" 2>/dev/null | tr -d '[:space:]') + if [ "$vk_has_budget_id" = "0" ]; then + log_info " Column governance_virtual_keys.budget_id dropped ✓" + else + log_error " Column governance_virtual_keys.budget_id still exists!" + failed=1 + fi + + # Check: budget_id column should be dropped from governance_virtual_key_provider_configs + local pc_has_budget_id + pc_has_budget_id=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_virtual_key_provider_configs' AND column_name = 'budget_id'" 2>/dev/null | tr -d '[:space:]') + if [ "$pc_has_budget_id" = "0" ]; then + log_info " Column governance_virtual_key_provider_configs.budget_id dropped ✓" + else + log_error " Column governance_virtual_key_provider_configs.budget_id still exists!" + failed=1 + fi + + # Check: junction tables should not exist + local junction_vk + junction_vk=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'governance_virtual_key_budgets'" 2>/dev/null | tr -d '[:space:]') + if [ "$junction_vk" = "0" ]; then + log_info " Junction table governance_virtual_key_budgets dropped ✓" + else + log_warn " Junction table governance_virtual_key_budgets still exists (may not have existed in old version)" + fi + + return $failed +} + validate_postgres_data() { local before_snapshot="$1" local after_snapshot="$2" - + compare_postgres_snapshots "$before_snapshot" "$after_snapshot" } @@ -2844,7 +2923,14 @@ EOF stop_bifrost return 1 fi - + + # STEP 6: Verify budget migration (budget_id → virtual_key_id/provider_config_id) + if ! verify_budget_migration_postgres; then + log_error "Budget migration verification failed after migration from $version" + stop_bifrost + return 1 + fi + stop_bifrost log_info "Migration from $version: SUCCESS" done diff --git a/framework/configstore/clientconfig.go b/framework/configstore/clientconfig.go index fb8b7418f9..664fc8ef31 100644 --- a/framework/configstore/clientconfig.go +++ b/framework/configstore/clientconfig.go @@ -590,7 +590,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 @@ -602,7 +601,6 @@ type VirtualKeyProviderConfigHashInput struct { Provider string Weight *float64 AllowedModels []string - BudgetID *string RateLimitID *string KeyIDs []string // Only key IDs, not full key objects } @@ -638,10 +636,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)) @@ -655,16 +649,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 @@ -702,7 +686,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 a791a27405..d121693095 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -997,12 +997,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 @@ -1013,10 +1014,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") { @@ -1025,12 +1024,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) @@ -1044,32 +1038,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) @@ -5218,6 +5199,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", @@ -5314,36 +5296,21 @@ func migrationAddKeyBlacklistedModelsJSONColumn(ctx context.Context, db *gorm.DB 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()) } 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{{ @@ -5352,60 +5319,71 @@ func migrationAddMultiBudgetTables(ctx context.Context, db *gorm.DB) error { tx = tx.WithContext(ctx) mg := tx.Migrator() - // Create VK-level multi-budget junction table - if !mg.HasTable(&tables.TableVirtualKeyBudget{}) { - if err := mg.CreateTable(&tables.TableVirtualKeyBudget{}); err != nil { - return fmt.Errorf("failed to create governance_virtual_key_budgets table: %w", err) + // 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) } } - // Create provider-config-level multi-budget junction table - if !mg.HasTable(&tables.TableVirtualKeyProviderConfigBudget{}) { - if err := mg.CreateTable(&tables.TableVirtualKeyProviderConfigBudget{}); err != nil { - return fmt.Errorf("failed to create governance_virtual_key_provider_config_budgets table: %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) } } - // Backfill: migrate existing VK single budgets to junction table - if err := tx.Exec(` - INSERT INTO governance_virtual_key_budgets (virtual_key_id, budget_id) - SELECT id, budget_id FROM governance_virtual_keys - WHERE budget_id IS NOT NULL AND budget_id != '' - AND NOT EXISTS ( - SELECT 1 FROM governance_virtual_key_budgets - WHERE governance_virtual_key_budgets.virtual_key_id = governance_virtual_keys.id - AND governance_virtual_key_budgets.budget_id = governance_virtual_keys.budget_id - ) - `).Error; err != nil { - return fmt.Errorf("failed to backfill VK 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: migrate existing provider config single budgets to junction table - if err := tx.Exec(` - INSERT INTO governance_virtual_key_provider_config_budgets (provider_config_id, budget_id) - SELECT id, budget_id FROM governance_virtual_key_provider_configs - WHERE budget_id IS NOT NULL AND budget_id != '' - AND NOT EXISTS ( - SELECT 1 FROM governance_virtual_key_provider_config_budgets - WHERE governance_virtual_key_provider_config_budgets.provider_config_id = governance_virtual_key_provider_configs.id - AND governance_virtual_key_provider_config_budgets.budget_id = governance_virtual_key_provider_configs.budget_id - ) - `).Error; err != nil { - return fmt.Errorf("failed to backfill provider config budgets: %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) + } } + // 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.HasTable(&tables.TableVirtualKeyProviderConfigBudget{}) { - if err := mg.DropTable(&tables.TableVirtualKeyProviderConfigBudget{}); err != nil { + if mg.HasColumn(&tables.TableBudget{}, "virtual_key_id") { + if err := mg.DropColumn(&tables.TableBudget{}, "virtual_key_id"); err != nil { return err } } - if mg.HasTable(&tables.TableVirtualKeyBudget{}) { - if err := mg.DropTable(&tables.TableVirtualKeyBudget{}); err != nil { + if mg.HasColumn(&tables.TableBudget{}, "provider_config_id") { + if err := mg.DropColumn(&tables.TableBudget{}, "provider_config_id"); err != nil { return err } } diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 65a1e6b4de..fe736cdbc9 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -1837,10 +1837,10 @@ func (s *RDBConfigStore) GetVirtualKeys(ctx context.Context) ([]tables.TableVirt Preload("Team"). Preload("Team.Customer"). Preload("Customer"). - Preload("Budget"). + 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") @@ -1899,10 +1899,10 @@ func (s *RDBConfigStore) GetVirtualKeysPaginated(ctx context.Context, params Vir Preload("Team"). Preload("Team.Customer"). Preload("Customer"). - Preload("Budget"). + 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") @@ -1925,10 +1925,10 @@ func (s *RDBConfigStore) GetVirtualKey(ctx context.Context, id string) (*tables. Preload("Team"). Preload("Team.Customer"). Preload("Customer"). - Preload("Budget"). + 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") @@ -1952,10 +1952,10 @@ func (s *RDBConfigStore) GetVirtualKeyByValue(ctx context.Context, value string) Preload("Team"). Preload("Team.Customer"). Preload("Customer"). - Preload("Budget"). + 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") @@ -2093,33 +2093,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 @@ -2129,8 +2122,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 { @@ -2139,11 +2134,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 { @@ -2341,18 +2331,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 { 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 5164ceaf3a..0a84b5550b 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 @@ -30,30 +30,6 @@ type TableBudget struct { // TableName sets the table name for each model func (TableBudget) TableName() string { return "governance_budgets" } -// TableVirtualKeyBudget is a junction table for VK-level multi-budget support -type TableVirtualKeyBudget struct { - ID uint `gorm:"primaryKey;autoIncrement" json:"id"` - VirtualKeyID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_vk_budget" json:"virtual_key_id"` - BudgetID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_vk_budget" json:"budget_id"` - Budget TableBudget `gorm:"foreignKey:BudgetID;constraint:OnDelete:CASCADE" json:"budget"` -} - -// TableName for TableVirtualKeyBudget -func (TableVirtualKeyBudget) TableName() string { return "governance_virtual_key_budgets" } - -// TableVirtualKeyProviderConfigBudget is a junction table for provider-config-level multi-budget support -type TableVirtualKeyProviderConfigBudget struct { - ID uint `gorm:"primaryKey;autoIncrement" json:"id"` - ProviderConfigID uint `gorm:"not null;uniqueIndex:idx_pc_budget" json:"provider_config_id"` - BudgetID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_pc_budget" json:"budget_id"` - Budget TableBudget `gorm:"foreignKey:BudgetID;constraint:OnDelete:CASCADE" json:"budget"` -} - -// TableName for TableVirtualKeyProviderConfigBudget -func (TableVirtualKeyProviderConfigBudget) TableName() string { - return "governance_virtual_key_provider_config_budgets" -} - // BeforeSave hook for Budget to validate reset duration format and max limit func (b *TableBudget) BeforeSave(tx *gorm.DB) error { // Validate that ResetDuration is in correct format (e.g., "30s", "5m", "1h", "1d", "1w", "1M", "1Y") diff --git a/framework/configstore/tables/virtualkey.go b/framework/configstore/tables/virtualkey.go index 835ae3fa97..1b7e69e3c5 100644 --- a/framework/configstore/tables/virtualkey.go +++ b/framework/configstore/tables/virtualkey.go @@ -30,13 +30,11 @@ 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"` - Budgets []TableBudget `gorm:"many2many:governance_virtual_key_provider_config_budgets;joinForeignKey:ProviderConfigID;joinReferences:BudgetID" json:"budgets,omitempty"` // Multiple budgets with different reset intervals + 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 } @@ -218,17 +216,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:"many2many:governance_virtual_key_budgets;joinForeignKey:VirtualKeyID;joinReferences:BudgetID" json:"budgets,omitempty"` // Multiple budgets with different reset intervals + 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/http_transport_prehook_test.go b/plugins/governance/http_transport_prehook_test.go index f50511d740..eb29a3bd60 100644 --- a/plugins/governance/http_transport_prehook_test.go +++ b/plugins/governance/http_transport_prehook_test.go @@ -30,7 +30,7 @@ func TestHTTPTransportPreHook_VirtualKeyReplicateRefinesNestedModel(t *testing.T "sk-bf-test", "replicate-only", []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("replicate", nil), + buildProviderConfig("replicate", []string{"*"}), }, ) store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ diff --git a/plugins/governance/model_provider_governance_test.go b/plugins/governance/model_provider_governance_test.go index a17855552a..99ebab374f 100644 --- a/plugins/governance/model_provider_governance_test.go +++ b/plugins/governance/model_provider_governance_test.go @@ -1480,6 +1480,9 @@ func TestPreLLMHook_ModelProviderPass_VirtualKeyChecksPass(t *testing.T) { // Model/provider checks pass (no limits) // Virtual key checks also pass 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 ed51b51f0c..b0ebe0eae2 100644 --- a/plugins/governance/resolver_test.go +++ b/plugins/governance/resolver_test.go @@ -17,6 +17,9 @@ import ( func TestBudgetResolver_EvaluateRequest_AllowedRequest(t *testing.T) { logger := NewMockLogger() 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}, @@ -99,9 +102,8 @@ func TestBudgetResolver_EvaluateRequest_ModelBlocked(t *testing.T) { Provider: "openai", AllowedModels: []string{"gpt-4", "gpt-4-turbo"}, // Only these models Weight: bifrost.Ptr(1.0), - RateLimit: nil, - Budget: nil, - Keys: []configstoreTables.TableKey{}, + RateLimit: nil, + Keys: []configstoreTables.TableKey{}, }, } vk := buildVirtualKeyWithProviders("vk1", "sk-bf-test", "Test VK", providerConfigs) @@ -468,6 +470,9 @@ func TestBudgetResolver_IsModelAllowed(t *testing.T) { func TestBudgetResolver_ContextPopulation(t *testing.T) { logger := NewMockLogger() 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 32bc317d51..0a7643fdb9 100644 --- a/plugins/governance/store.go +++ b/plugins/governance/store.go @@ -200,12 +200,17 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { // Cross-reference live budget/rate limit from standalone maps // (usage updates clone into budgets/rateLimits maps, so embedded pointers go stale) clone := *vk - if clone.BudgetID != nil { - if liveBudget, exists := gs.budgets.Load(*clone.BudgetID); exists && liveBudget != nil { - if b, ok := liveBudget.(*configstoreTables.TableBudget); ok { - clone.Budget = b + // Hydrate multi-budgets from live sync.Map + if len(clone.Budgets) > 0 { + liveBudgets := make([]configstoreTables.TableBudget, 0, len(clone.Budgets)) + for _, b := range clone.Budgets { + if lb, exists := gs.budgets.Load(b.ID); exists && lb != nil { + if budget, ok := lb.(*configstoreTables.TableBudget); ok { + liveBudgets = append(liveBudgets, *budget) + } } } + clone.Budgets = liveBudgets } if clone.RateLimitID != nil { if liveRL, exists := gs.rateLimits.Load(*clone.RateLimitID); exists && liveRL != nil { @@ -219,12 +224,17 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { configs := make([]configstoreTables.TableVirtualKeyProviderConfig, len(clone.ProviderConfigs)) copy(configs, clone.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 { @@ -1462,10 +1472,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 @@ -2037,33 +2071,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 { @@ -2308,22 +2326,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 } } } @@ -2412,9 +2444,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 @@ -2425,8 +2457,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) @@ -2453,23 +2485,27 @@ 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 { + + // Update multi-budgets for VK + newBudgetIDs := make(map[string]bool) + for i := range clone.Budgets { + newBudgetIDs[clone.Budgets[i].ID] = true + // 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 !newBudgetIDs[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 @@ -2519,22 +2555,25 @@ 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 + pcNewBudgetIDs := make(map[string]bool) + for j := range clone.ProviderConfigs[i].Budgets { + b := &clone.ProviderConfigs[i].Budgets[j] + pcNewBudgetIDs[b.ID] = true + 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 !pcNewBudgetIDs[oldBudget.ID] { + gs.budgets.Delete(oldBudget.ID) + } } } } @@ -2560,9 +2599,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 @@ -2573,8 +2612,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) @@ -3100,18 +3139,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 + } } } } @@ -3573,9 +3616,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..5d4ea852cb 100644 --- a/plugins/governance/store_test.go +++ b/plugins/governance/store_test.go @@ -198,11 +198,12 @@ 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) } } } @@ -210,6 +211,386 @@ func TestGovernanceStore_CheckBudget_HierarchyValidation(t *testing.T) { assert.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) + assert.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) + assert.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) + assert.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) + assert.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) + assert.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 func TestGovernanceStore_UpdateRateLimitUsage_TokensAndRequests(t *testing.T) { logger := NewMockLogger() @@ -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 a6da347cbf..99357a02a9 100644 --- a/plugins/governance/test_utils.go +++ b/plugins/governance/test_utils.go @@ -89,9 +89,11 @@ 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 + 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{"*"}), + } return vk } @@ -100,6 +102,10 @@ 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{"*"}), + } return vk } @@ -188,12 +194,23 @@ func buildProviderConfig(provider string, allowedModels []string) configstoreTab Provider: provider, AllowedModels: allowedModels, Weight: bifrost.Ptr(1.0), - RateLimit: nil, - Budget: nil, - Keys: []configstoreTables.TableKey{}, + RateLimit: 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) + vk.Budgets = budgets + return vk +} + func buildProviderConfigWithRateLimit(provider string, allowedModels []string, rateLimit *configstoreTables.TableRateLimit) configstoreTables.TableVirtualKeyProviderConfig { pc := buildProviderConfig(provider, allowedModels) pc.RateLimit = rateLimit diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 868ff1708c..0229d46e1e 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 @@ -205,8 +207,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) @@ -437,16 +439,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 @@ -457,30 +466,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{ @@ -503,19 +496,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)} } @@ -548,24 +549,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{ @@ -589,6 +572,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 { @@ -700,8 +707,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 { @@ -726,72 +734,57 @@ 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 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 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 *req.Budget.MaxLimit < 0 { - return fmt.Errorf("budget max_limit cannot be negative: %.2f", *req.Budget.MaxLimit) + if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil { + return &badRequestError{err: fmt.Errorf("invalid reset duration format: %s", b.ResetDuration)} } - if _, err := configstoreTables.ParseDuration(*req.Budget.ResetDuration); err != nil { - return fmt.Errorf("invalid reset duration format: %s", *req.Budget.ResetDuration) + if seenDurations[b.ResetDuration] { + return &badRequestError{err: fmt.Errorf("duplicate reset_duration in budgets: %s", b.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 - } - if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { - return err + + // Delete existing budgets for this VK + if err := tx.Where("virtual_key_id = ?", vk.ID).Delete(&configstoreTables.TableBudget{}).Error; err != nil { + return fmt.Errorf("failed to delete old VK budgets: %w", err) } - vk.BudgetID = &budget.ID - vk.Budget = &budget + + // Create new multi-budgets + if len(req.Budgets) > 0 { + var newBudgets []configstoreTables.TableBudget + 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 + } + newBudgets = append(newBudgets, budget) + } + vk.Budgets = newBudgets + } else { + vk.Budgets = nil } } + // Handle rate limit updates if req.RateLimit != nil { if isRateLimitRemovalRequest(req.RateLimit) { @@ -862,21 +855,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)} } @@ -909,25 +887,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{ @@ -950,6 +909,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] @@ -985,67 +968,44 @@ 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 - } - } - 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") + // 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 *pc.Budget.MaxLimit < 0 { - return fmt.Errorf("provider config budget max_limit cannot be negative: %.2f", *pc.Budget.MaxLimit) + if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil { + return &badRequestError{err: fmt.Errorf("invalid provider config budget reset duration format: %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) + if seenDurations[b.ResetDuration] { + return &badRequestError{err: fmt.Errorf("duplicate reset_duration in provider config budgets: %s", b.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 + // Delete existing budgets for this provider config + if err := tx.Where("provider_config_id = ?", existing.ID).Delete(&configstoreTables.TableBudget{}).Error; err != nil { + return fmt.Errorf("failed to delete old provider config budgets: %w", err) } - if err := h.configStore.CreateBudget(ctx, &budget, tx); err != nil { - return err - } - existing.BudgetID = &budget.ID + // Create new multi-budgets + if len(pc.Budgets) > 0 { + for _, b := range pc.Budgets { + 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 + } + } } } // Handle rate limit updates for provider config @@ -1177,11 +1137,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 @@ -1397,8 +1352,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 { @@ -1538,14 +1492,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 } @@ -1564,13 +1510,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 { @@ -1809,8 +1753,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 { @@ -1940,14 +1883,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 } @@ -1966,13 +1901,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 { @@ -2197,9 +2130,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 } @@ -2362,8 +2292,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 { @@ -2468,14 +2397,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 } @@ -2494,13 +2415,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 { @@ -2740,14 +2659,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 } @@ -2760,13 +2671,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 1ebc985a2f..b082afb4ab 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -2224,7 +2224,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 6da64e3e32..a860695284 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 @@ -6433,7 +6432,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", @@ -6441,7 +6439,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } // Generate hash @@ -6462,7 +6459,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } hash2, err := configstore.GenerateVirtualKeyHash(vk2) @@ -6482,7 +6478,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } hash3, err := configstore.GenerateVirtualKeyHash(vk3) @@ -6502,7 +6497,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_different", // Different value IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } hash4, err := configstore.GenerateVirtualKeyHash(vk4) @@ -6522,7 +6516,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: false, // Different IsActive TeamID: &teamID, - BudgetID: &budgetID, } hash5, err := configstore.GenerateVirtualKeyHash(vk5) @@ -6543,7 +6536,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &differentTeamID, // Different TeamID - BudgetID: &budgetID, } hash6, err := configstore.GenerateVirtualKeyHash(vk6) @@ -6563,7 +6555,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } hash7, err := configstore.GenerateVirtualKeyHash(vk7) @@ -6584,7 +6575,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, CustomerID: &customerID, // CustomerID set } @@ -6606,7 +6596,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, CustomerID: &differentCustomerID, // Different CustomerID } @@ -6619,27 +6608,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{ @@ -6649,7 +6617,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, RateLimitID: &rateLimitID, // RateLimitID set } @@ -6671,7 +6638,6 @@ func TestGenerateVirtualKeyHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, RateLimitID: &differentRateLimitID, // Different RateLimitID } @@ -6689,7 +6655,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 @@ -6706,7 +6671,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"}, @@ -6739,7 +6703,6 @@ func TestGenerateVirtualKeyHash_WithProviderConfigs(t *testing.T) { Provider: "anthropic", // Different provider Weight: ptrFloat64(1.0), AllowedModels: []string{"claude-3"}, - BudgetID: &budgetID, RateLimitID: &rateLimitID, }, }, @@ -6768,7 +6731,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"}, @@ -6872,7 +6834,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{ @@ -6882,7 +6843,6 @@ func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } // Generate file hash @@ -6893,7 +6853,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", @@ -6901,7 +6860,6 @@ func TestVirtualKeyHashComparison_MatchingHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &dbTeamID, - BudgetID: &dbBudgetID, ConfigHash: fileHash, // Same hash as file } @@ -6926,7 +6884,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{ @@ -6936,7 +6893,6 @@ func TestVirtualKeyHashComparison_DifferentHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } dbHash, err := configstore.GenerateVirtualKeyHash(dbVK) @@ -6947,7 +6903,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 @@ -6955,7 +6910,6 @@ func TestVirtualKeyHashComparison_DifferentHash(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &fileTeamID, - BudgetID: &fileBudgetID, } fileHash, err := configstore.GenerateVirtualKeyHash(fileVK) @@ -7127,26 +7081,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{ @@ -7173,7 +7107,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{ @@ -7183,7 +7116,6 @@ func TestVirtualKeyHashComparison_FieldValueChanges(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, } baseHash, err := configstore.GenerateVirtualKeyHash(baseVK) @@ -7231,27 +7163,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 @@ -7262,7 +7179,6 @@ func TestVirtualKeyHashComparison_RoundTrip(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &teamID, - BudgetID: &budgetID, RateLimitID: &rateLimitID, ProviderConfigs: []tables.TableVirtualKeyProviderConfig{ { @@ -7285,7 +7201,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", @@ -7294,7 +7209,6 @@ func TestVirtualKeyHashComparison_RoundTrip(t *testing.T) { Value: "vk_abc123", IsActive: true, TeamID: &reloadTeamID, - BudgetID: &reloadBudgetID, RateLimitID: &reloadRateLimitID, ProviderConfigs: []tables.TableVirtualKeyProviderConfig{ { @@ -8806,7 +8720,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{ { @@ -10521,9 +10434,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", @@ -10538,7 +10448,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) { Provider: "openai", Weight: ptrFloat64(1.0), AllowedModels: []string{"gpt-4"}, - BudgetID: &budgetID1, }, { ID: 2, @@ -10546,7 +10455,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) { Provider: "anthropic", Weight: ptrFloat64(2.0), AllowedModels: []string{"claude-3"}, - BudgetID: &budgetID2, }, { ID: 3, @@ -10579,7 +10487,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) { Provider: "anthropic", Weight: ptrFloat64(2.0), AllowedModels: []string{"claude-3"}, - BudgetID: &budgetID2, }, { ID: 1, @@ -10587,7 +10494,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) { Provider: "openai", Weight: ptrFloat64(1.0), AllowedModels: []string{"gpt-4"}, - BudgetID: &budgetID1, }, }, } @@ -10606,7 +10512,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) { Provider: "anthropic", Weight: ptrFloat64(2.0), AllowedModels: []string{"claude-3"}, - BudgetID: &budgetID2, }, { ID: 1, @@ -10614,7 +10519,6 @@ func TestGenerateVirtualKeyHash_StableProviderConfigOrdering(t *testing.T) { Provider: "openai", Weight: ptrFloat64(1.0), AllowedModels: []string{"gpt-4"}, - BudgetID: &budgetID1, }, { ID: 3, @@ -11027,8 +10931,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", @@ -11036,7 +10938,6 @@ func TestGenerateVirtualKeyHash_StableCombinedOrdering(t *testing.T) { Description: "Test virtual key", Value: "vk_abc123", IsActive: true, - BudgetID: &budgetID, ProviderConfigs: []tables.TableVirtualKeyProviderConfig{ { ID: 1, @@ -11079,7 +10980,6 @@ func TestGenerateVirtualKeyHash_StableCombinedOrdering(t *testing.T) { Description: "Test virtual key", Value: "vk_abc123", IsActive: true, - BudgetID: &budgetID, ProviderConfigs: []tables.TableVirtualKeyProviderConfig{ { ID: 2, @@ -14165,11 +14065,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 @@ -14180,34 +14078,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{ @@ -14236,34 +14106,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{ @@ -14293,7 +14135,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", @@ -14302,7 +14144,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) { { Provider: "openai", Weight: &weight, - BudgetID: &budgetID1, RateLimitID: &rateLimitID1, }, }, @@ -14315,7 +14156,6 @@ func TestGenerateVirtualKeyHash_ProviderConfigBudgetRateLimit(t *testing.T) { { Provider: "openai", Weight: &weight, - BudgetID: &budgetID1, RateLimitID: &rateLimitID1, }, }, @@ -14346,107 +14186,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 @@ -15519,9 +15258,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, @@ -15550,13 +15291,13 @@ 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 + "budgets": true, // GORM relation (budgets have provider_config_id FK) "rate_limit": true, // GORM relation }, "tables.TableVirtualKeyMCPConfig": { diff --git a/transports/config.schema.json b/transports/config.schema.json index 9cd015c1a9..5d3a514020 100644 --- a/transports/config.schema.json +++ b/transports/config.schema.json @@ -452,6 +452,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)" @@ -460,10 +465,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" @@ -1566,10 +1567,6 @@ "type": "string" } }, - "budget_id": { - "type": "string", - "description": "Associated budget ID" - }, "rate_limit_id": { "type": "string", "description": "Associated rate limit ID" diff --git a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx index 9841528315..cdb84c4f2b 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx @@ -30,7 +30,7 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe const isExhausted = // VK-level budget exhausted - (virtualKey.budget?.current_usage && virtualKey.budget?.max_limit && virtualKey.budget.current_usage >= virtualKey.budget.max_limit) || + (virtualKey.budgets?.some((b) => b.current_usage >= b.max_limit)) || // VK-level rate limits exhausted (virtualKey.rate_limit?.token_current_usage && virtualKey.rate_limit?.token_max_limit && @@ -148,38 +148,42 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe - {/* Provider Budget */} - {config.budget && ( + {/* Provider Budgets */} + {config.budgets && config.budgets.length > 0 && ( <>
-

Provider Budget

-
- Usage -
-
- - {formatCurrency(config.budget.current_usage)} / {formatCurrency(config.budget.max_limit)} - - = config.budget.max_limit ? "destructive" : "default"} - className="text-xs" - > - {Math.round((config.budget.current_usage / config.budget.max_limit) * 100)}% - +

Provider Budgets

+ {config.budgets.map((b, bIdx) => ( +
+
+ Usage +
+
+ + {formatCurrency(b.current_usage)} / {formatCurrency(b.max_limit)} + + = b.max_limit ? "destructive" : "default"} + className="text-xs" + > + {Math.round((b.current_usage / b.max_limit) * 100)}% + +
+
+
+
+ Reset Period +
{parseResetPeriod(b.reset_duration)}{virtualKey.calendar_aligned && " (calendar)"}
+
+
+ Last Reset +
+ {formatDistanceToNow(new Date(b.last_reset), { addSuffix: true })} +
-
-
- Reset Period -
{parseResetPeriod(config.budget.reset_duration)}
-
-
- Last Reset -
- {formatDistanceToNow(new Date(config.budget.last_reset), { addSuffix: true })} -
-
+ ))}
)} @@ -343,36 +347,38 @@ export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKe

Budget Information

- {virtualKey.budget ? ( + {virtualKey.budgets && virtualKey.budgets.length > 0 ? (
-
- Usage -
-
- - {formatCurrency(virtualKey.budget.current_usage)} / {formatCurrency(virtualKey.budget.max_limit)} - - = virtualKey.budget.max_limit ? "destructive" : "default"} - className="text-xs" - > - {Math.round((virtualKey.budget.current_usage / virtualKey.budget.max_limit) * 100)}% - + {virtualKey.budgets.map((b, bIdx) => ( +
+
+ Usage +
+
+ + {formatCurrency(b.current_usage)} / {formatCurrency(b.max_limit)} + + = b.max_limit ? "destructive" : "default"} + className="text-xs" + > + {Math.round((b.current_usage / b.max_limit) * 100)}% + +
+
+
+
+ Reset Period +
{parseResetPeriod(b.reset_duration)}{virtualKey.calendar_aligned && " (calendar)"}
+
+
+ Last Reset +
+ {formatDistanceToNow(new Date(b.last_reset), { addSuffix: true })} +
-
- -
- Reset Period -
{parseResetPeriod(virtualKey.budget.reset_duration)}
-
- -
- Last Reset -
- {formatDistanceToNow(new Date(virtualKey.budget.last_reset), { addSuffix: true })} -
-
+ ))}
) : (

No budget limits configured

diff --git a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx index ebd5ce3ca9..dcc47a5530 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx @@ -81,12 +81,6 @@ const providerConfigSchema = z.object({ allowed_models: z.array(z.string()).optional(), key_ids: z.array(z.string()).optional(), // Keys associated with this provider config // Provider-level budget - budget: z - .object({ - max_limit: z.string().optional(), - reset_duration: z.string().optional(), - }) - .optional(), budgets: z.array(z.object({ max_limit: z.string().optional(), reset_duration: z.string().optional(), @@ -123,6 +117,11 @@ const formSchema = z budgetMaxLimit: z.string().optional(), budgetResetDuration: z.string().optional(), budgetCalendarAligned: z.boolean(), + // Budget (multi-budget) + budgets: z.array(z.object({ + max_limit: z.string(), + reset_duration: z.string(), + })).optional(), // Token limits tokenMaxLimit: z.string().optional(), tokenResetDuration: z.string().optional(), @@ -197,12 +196,6 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, weight: config.weight ?? "", allowed_models: config.allowed_models, key_ids: config.allow_all_keys ? ["*"] : config.keys?.map((key) => key.key_id) || [], - budget: config.budget - ? { - max_limit: String(config.budget.max_limit), - reset_duration: config.budget.reset_duration, - } - : undefined, budgets: config.budgets?.map((b) => ({ max_limit: String(b.max_limit), reset_duration: b.reset_duration, @@ -226,9 +219,10 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, teamId: virtualKey?.team_id || "", customerId: virtualKey?.customer_id || "", isActive: virtualKey?.is_active ?? true, - budgetMaxLimit: virtualKey?.budget ? String(virtualKey.budget.max_limit) : "", - budgetResetDuration: virtualKey?.budget?.reset_duration || "1M", - budgetCalendarAligned: virtualKey?.budget?.calendar_aligned ?? false, + budgets: virtualKey?.budgets && virtualKey.budgets.length > 0 + ? virtualKey.budgets.map((b) => ({ max_limit: String(b.max_limit ?? ""), reset_duration: b.reset_duration ?? "1M" })) + : [], + budgetCalendarAligned: virtualKey?.budgets?.some((b) => b.calendar_aligned) ?? false, tokenMaxLimit: virtualKey?.rate_limit?.token_max_limit ? String(virtualKey.rate_limit.token_max_limit) : "", tokenResetDuration: virtualKey?.rate_limit?.token_reset_duration || "1h", requestMaxLimit: virtualKey?.rate_limit?.request_max_limit ? String(virtualKey.rate_limit.request_max_limit) : "", @@ -283,11 +277,15 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, const mcpConfigs = form.watch("mcpConfigs") || []; // Watch budget/rate-limit fields for conditional rendering of reset buttons - const watchedBudgetMaxLimit = form.watch("budgetMaxLimit"); - const watchedBudgetResetDuration = form.watch("budgetResetDuration") || "1M"; - const watchedBudgetCalendarAligned = form.watch("budgetCalendarAligned"); + const watchedBudgets = form.watch("budgets"); const watchedTokenMaxLimit = form.watch("tokenMaxLimit"); const watchedRequestMaxLimit = form.watch("requestMaxLimit"); + const watchedBudgetCalendarAligned = form.watch("budgetCalendarAligned"); + + // Derive single-budget-style values from the multi-budget array for calendar alignment UI + const firstBudget = watchedBudgets && watchedBudgets.length > 0 ? watchedBudgets[0] : null; + const watchedBudgetMaxLimit = firstBudget?.max_limit; + const watchedBudgetResetDuration = firstBudget?.reset_duration; // Handle adding a new provider configuration const handleAddProvider = (provider: string) => { @@ -351,36 +349,24 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, const [showCalendarAlignWarning, setShowCalendarAlignWarning] = useState(false); - const clearVirtualKeyBudget = () => { - form.setValue("budgetMaxLimit", "", { shouldDirty: true }); - form.setValue("budgetResetDuration", "1M", { shouldDirty: true }); - form.setValue("budgetCalendarAligned", false, { shouldDirty: true }); - }; - const handleCalendarAlignedChange = (checked: boolean) => { - if (checked && isEditing && virtualKey?.budget && !virtualKey.budget.calendar_aligned) { + if (checked && isEditing) { + // Show warning when enabling on an existing VK setShowCalendarAlignWarning(true); } else { form.setValue("budgetCalendarAligned", checked, { shouldDirty: true }); } }; - const clearVirtualKeyRateLimits = () => { - form.setValue("tokenMaxLimit", "", { shouldDirty: true }); - form.setValue("tokenResetDuration", "1h", { shouldDirty: true }); - form.setValue("requestMaxLimit", "", { shouldDirty: true }); - form.setValue("requestResetDuration", "1h", { shouldDirty: true }); + const clearVirtualKeyBudget = () => { + form.setValue("budgets", [], { shouldDirty: true }); }; - const normalizeIntegerField = (value: string | undefined): number | undefined => { - if (value === undefined || value === "") return undefined; - const num = parseInt(value, 10); - return isNaN(num) ? undefined : num; + const clearVirtualKeyRateLimits = () => { + form.setValue("rate_limit", undefined, { shouldDirty: true }); }; - // Helper function to convert string weights to numbers and normalize budget/rate limit fields const normalizeProviderConfigs = ( - configs: NonNullable, existingConfigs?: VirtualKey["provider_configs"], ): any[] => { return configs.map((config) => ({ @@ -393,22 +379,6 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, ? null : parseFloat(config.weight) : config.weight, - budget: (() => { - const budgetMaxLimit = normalizeNumericField(config.budget?.max_limit); - if (budgetMaxLimit !== undefined) { - return { - max_limit: budgetMaxLimit, - reset_duration: config.budget?.reset_duration || "1M", - }; - } - - const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider)); - if (existingConfig?.budget) { - return {}; - } - - return undefined; - })(), rate_limit: (() => { const tokenMaxLimit = normalizeIntegerField(config.rate_limit?.token_max_limit); const requestMaxLimit = normalizeIntegerField(config.rate_limit?.request_max_limit); @@ -463,18 +433,16 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, is_active: data.isActive, }; - // Add budget if enabled - const budgetMaxLimit = normalizeNumericField(data.budgetMaxLimit); - const hadBudget = !!virtualKey.budget; - const hasBudget = budgetMaxLimit !== undefined; - if (hasBudget) { - updateData.budget = { - max_limit: budgetMaxLimit, - reset_duration: data.budgetResetDuration || "1M", - calendar_aligned: data.budgetCalendarAligned, - }; + // Add budgets if enabled + const validBudgets = (data.budgets || []).filter((b) => normalizeNumericField(b.max_limit) !== undefined); + const hadBudget = virtualKey.budgets && virtualKey.budgets.length > 0; + if (validBudgets.length > 0) { + updateData.budgets = validBudgets.map((b) => ({ + max_limit: normalizeNumericField(b.max_limit)!, + reset_duration: b.reset_duration || "1M", + })); } else if (hadBudget) { - updateData.budget = {}; + updateData.budgets = []; } // Add rate limit if enabled @@ -509,14 +477,13 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, is_active: data.isActive, }; - // Add budget if enabled - const budgetMaxLimit = normalizeNumericField(data.budgetMaxLimit); - if (budgetMaxLimit !== undefined) { - createData.budget = { - max_limit: budgetMaxLimit, - reset_duration: data.budgetResetDuration || "1M", - calendar_aligned: data.budgetCalendarAligned, - }; + // Add budgets if enabled + const validBudgets = (data.budgets || []).filter((b) => normalizeNumericField(b.max_limit) !== undefined); + if (validBudgets.length > 0) { + createData.budgets = validBudgets.map((b) => ({ + max_limit: normalizeNumericField(b.max_limit)!, + reset_duration: b.reset_duration || "1M", + })); } // Add rate limit if enabled @@ -936,14 +903,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave, max_limit: String(b.max_limit ?? ""), reset_duration: b.reset_duration || "1M", })) - : config.budget?.max_limit - ? [ - { - max_limit: String(config.budget.max_limit), - reset_duration: config.budget.reset_duration || "1M", - }, - ] - : [] + : [] } onChange={(lines) => { if (lines.length === 0) { @@ -1244,38 +1204,20 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, onSave,
- {isEditing && (virtualKey?.budget || watchedBudgetMaxLimit) && ( + {isEditing && (virtualKey?.budgets?.length || (watchedBudgets && watchedBudgets.length > 0)) && ( )}
- ( - - { - field.onChange(value); - }} - onChangeSelect={(value) => { - form.setValue("budgetResetDuration", value, { shouldDirty: true }); - if (!supportsCalendarAlignment(value)) { - form.setValue("budgetCalendarAligned", false, { shouldDirty: true }); - } - }} - options={resetDurationOptions} - /> - - - )} + { + form.setValue("budgets", lines, { shouldDirty: true }); + }} /> {/* Calendar alignment toggle — only shown when a budget is set and the period supports alignment */} diff --git a/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx b/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx index 940e4695bf..a58b82d26f 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx @@ -252,7 +252,7 @@ export default function VirtualKeysTable({ virtualKeys.map((vk) => { const isRevealed = revealedKeys.has(vk.id); const isExhausted = - (vk.budget?.current_usage && vk.budget?.max_limit && vk.budget.current_usage >= vk.budget.max_limit) || + (vk.budgets?.some((b) => b.current_usage >= b.max_limit)) || (vk.rate_limit?.token_current_usage && vk.rate_limit?.token_max_limit && vk.rate_limit.token_current_usage >= vk.rate_limit.token_max_limit) || @@ -303,15 +303,19 @@ export default function VirtualKeysTable({
- {vk.budget ? ( + {vk.budgets && vk.budgets.length > 0 ? (
- = vk.budget.max_limit && "text-red-400")}> - {formatCurrency(vk.budget.current_usage)} / {formatCurrency(vk.budget.max_limit)} - - - Resets {formatResetDuration(vk.budget.reset_duration)} - {vk.budget.calendar_aligned && " (calendar)"} - + {vk.budgets.map((b, idx) => ( +
+ = b.max_limit && "text-red-400")}> + {formatCurrency(b.current_usage)} / {formatCurrency(b.max_limit)} + + + Resets {formatResetDuration(b.reset_duration)} + {vk.calendar_aligned && " (calendar)"} + +
+ ))}
) : ( - diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 6251734534..e8387d915e 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -72,15 +72,15 @@ export interface VirtualKey { mcp_configs?: VirtualKeyMCPConfig[]; team_id?: string; customer_id?: string; - budget_id?: string; rate_limit_id?: string; is_active: boolean; + calendar_aligned?: boolean; created_at: string; updated_at: string; // Populated relationships team?: Team; customer?: Customer; - budget?: Budget; + budgets?: Budget[]; rate_limit?: RateLimit; config_hash?: string; // Present when config is synced from config.json } @@ -91,7 +91,6 @@ export interface VirtualKeyProviderConfig { weight: number | null; allowed_models: string[]; allow_all_keys: boolean; // True means all keys allowed; false with empty keys means no keys allowed - budget?: Budget; budgets?: Budget[]; rate_limit?: RateLimit; keys?: DBKey[]; // Associated database keys for this provider (only used when allow_all_keys is false) @@ -135,7 +134,6 @@ export interface VirtualKeyProviderConfigRequest { provider: string; weight?: number | null; allowed_models?: string[]; - budget?: CreateBudgetRequest; budgets?: CreateBudgetRequest[]; rate_limit?: CreateRateLimitRequest; key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config @@ -146,7 +144,6 @@ export interface VirtualKeyProviderConfigUpdateRequest { provider: string; weight?: number | null; allowed_models?: string[]; - budget?: UpdateBudgetRequest; budgets?: CreateBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config @@ -160,10 +157,10 @@ export interface CreateVirtualKeyRequest { mcp_configs?: VirtualKeyMCPConfigRequest[]; team_id?: string; customer_id?: string; - budget?: CreateBudgetRequest; budgets?: CreateBudgetRequest[]; rate_limit?: CreateRateLimitRequest; is_active?: boolean; + calendar_aligned?: boolean; } export interface UpdateVirtualKeyRequest { @@ -173,10 +170,10 @@ export interface UpdateVirtualKeyRequest { mcp_configs?: VirtualKeyMCPConfigRequest[]; team_id?: string; customer_id?: string; - budget?: UpdateBudgetRequest; budgets?: CreateBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; is_active?: boolean; + calendar_aligned?: boolean; } export interface CreateTeamRequest {