feat: add per-model budgets/rate-limits to virtual key provider configs

[BearTS]

Summary

Adds support for per-model budgets and rate limits scoped to a specific provider config on a Virtual Key. Previously, budgets and rate limits could only be set at the VK top-level or per-provider tier (the "*" all-models wildcard). This change allows operators to define finer-grained spending and request caps for individual models within a provider config.

Changes

• Introduced VKProviderModelBudget as a serialization-only struct on TableVirtualKeyProviderConfig, hydrated from VK-scoped model configs at read time rather than stored directly in the provider config table.
• Added model_budgets to CreateVirtualKeyRequest and UpdateVirtualKeyRequest provider config entries, with dedicated vkModelBudgetRequest and vkModelBudgetUpdateRequest types.
• Extended vkModelConfigDesired with a modelName field and a reconcileModelBudgets flag. When model_budgets is supplied for a provider, the backend treats the list as the full desired set and prunes any per-model configs absent from it; when omitted, existing per-model configs are left untouched.
• Refactored syncVKGovernanceToModelConfigs to key the keep map by provider + "\x00" + modelTier, enabling both the "*" provider tier and concrete model tiers to coexist and be independently reconciled.
• Added buildVKCreateModelBudgets and buildVKUpdateModelBudgets helpers with validation (validateVKModelBudgetNames) enforcing uniqueness, non-empty names, non-wildcard model names, and a cap of 100 per-model groups per provider.
• Added buildVKModelBudgetsIndex to group per-model VK-scoped model configs by provider for bulk hydration, and updated applyVKGovernanceFromModelConfigs, hydrateVKGovernance, hydrateVKListGovernance, getVirtualKeys, getVirtualKey, and collectVKModelUsage to pass and consume this index.
• Replaced the per-provider N+1 GetModelConfig calls in hydrateVKGovernance with a single GetModelConfigsByScopeAndScopeIDs bulk load.
• Updated the UI type definitions (VirtualKeyModelBudget, VirtualKeyModelBudgetRequest) and provider config request interfaces to include model_budgets.
• Extended the VK create/edit sheet (virtualKeySheet.tsx) to pass model_budgets through form state, normalization, and the ProviderConfigEditor component (via a new showModelBudgets prop).
• Updated the VK details sheet (virtualKeyDetailsSheet.tsx) to render per-model budgets, token limits, and request limits for each provider config.
• Updated tests and the mock config store to implement GetModelConfigsByScopeAndScopeIDs and pass the new perModelByKey argument to applyVKGovernanceFromModelConfigs.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

# Core/Transports
go test ./...

# UI
cd ui
pnpm i
pnpm build

1. Create a Virtual Key with a provider config that includes model_budgets, e.g.:

   {
     "provider_configs": [{
       "provider": "openai",
       "model_budgets": [
         { "model_name": "gpt-4o", "budgets": [{ "max_limit": 10.0, "reset_duration": "30d" }] },
         { "model_name": "gpt-4o-mini", "rate_limit": { "request_max_limit": 100, "request_reset_duration": "1h" } }
       ]
     }]
   }

2. Fetch the VK and confirm model_budgets is populated on the provider config in the response.
3. Update the VK omitting model_budgets for a provider and confirm existing per-model configs are preserved.
4. Update the VK supplying an empty model_budgets: [] for a provider and confirm existing per-model configs are pruned.
5. Verify the VK details sheet renders per-model budget and rate-limit rows under the provider config section.

Breaking changes

• Yes
• No

Related issues

Security considerations

Per-model budget validation rejects the "*" wildcard as a model name to prevent accidental shadowing of the provider-level tier. Model name uniqueness is enforced per provider config to prevent ambiguous budget application.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added per-model budgets and token/request rate limits for virtual key provider configurations.
  • Virtual key creation and editing now support governance settings for individual models.
  • Virtual key details display model-specific budget usage, reset schedules, and rate-limit information.
  • Per-model governance is preserved and shown in virtual key and quota details.

• Bug Fixes

  • Improved synchronization and cleanup of model-specific governance settings.
  • Added validation for concrete model names, configuration limits, and governance values.
  • Added support for explicitly removing model-specific rate limits.

Walkthrough

Per-model budgets and rate limits are added to virtual-key provider configurations. Backend reconciliation and hydration support model-specific governance. The virtual-key UI accepts, normalizes, submits, and displays these settings.

Changes

Virtual-key model governance

Layer / File(s)	Summary
Governance contracts and payload validation framework/configstore/tables/virtualkey.go, transports/bifrost-http/handlers/governance.go, ui/lib/types/governance.ts	Provider configuration contracts represent per-model budgets and rate limits. Create and update payloads validate model names, limits, uniqueness, and explicit clearing behavior.
Model-tier governance reconciliation transports/bifrost-http/handlers/governance.go	Virtual-key synchronization creates, retains, updates, and prunes model-config records by provider and model tier.
Model-budget hydration and validation coverage transports/bifrost-http/handlers/governance.go, transports/bifrost-http/handlers/governance_test.go	Virtual-key and quota hydration reverse-map scoped model configs into provider ModelBudgets. Test mocks and mapper call sites support bulk loading.
Model-budget editing and presentation ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx, ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx	Forms load, edit, normalize, and submit model-specific governance. Details views render budget usage, overrides, reset metadata, and token/request limits.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant VirtualKeySheet
  participant GovernanceAPI
  participant syncVKGovernanceToModelConfigs
  participant ConfigStore
  participant VirtualKeyDetailsSheet

  VirtualKeySheet->>GovernanceAPI: Submit model_budgets
  GovernanceAPI->>syncVKGovernanceToModelConfigs: Reconcile provider and model tiers
  syncVKGovernanceToModelConfigs->>ConfigStore: Store model-config governance
  VirtualKeyDetailsSheet->>GovernanceAPI: Request virtual-key details
  GovernanceAPI->>ConfigStore: Load scoped model configs
  ConfigStore-->>GovernanceAPI: Return model budgets and rate limits
  GovernanceAPI-->>VirtualKeyDetailsSheet: Return hydrated model_budgets

Loading

Possibly related PRs

• maximhq/bifrost#5704: Updates the ProviderConfigCard model-budget count header used by this virtual-key UI.
• maximhq/bifrost#5898: Also modifies virtualKeyDetailsSheet.tsx for per-model budget overrides.
• maximhq/bifrost#6017: Adds E2E coverage for virtual-key provider-level budget overrides.

Suggested reviewers: akshaydeo, impoiler, pratham-mishra04

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely identifies the main feature: per-model budgets and rate limits for Virtual Key provider configurations.
Description check	✅ Passed	The description is detailed and covers the change, affected areas, testing steps, breaking changes, security, and checklist; UI screenshots are not included.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 07-31-feat_add_virtual_key_per_model_budget_creation_on_the_virtual_key_sheet

---

_{Comment @coderabbitai help to get the list of available commands.}

[BearTS]

• feat: replace per-budget BudgetOverrideDialog with unified BudgetOverrideManagerDialog grouping provider and model budgets #5898
• fix: show model budget count in provider card header summary #5704
• feat: add per-model budgets/rate-limits to virtual key provider configs #5703 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)

1167-1197: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Per-model budget edits bypass the budget-reset prompt and over-limit warning.

hasBudgetResetRelevantChanges and getBudgetUsageWarning (Lines 691-760) only walk config.budgets. Now that model_budgets are editable here, changing a model's limit or reset frequency saves without the "Reset budget usage?" dialog, and existing usage that already exceeds the new per-model limit is never surfaced — inconsistent with provider- and VK-level behavior.

Extending budgetSignature/findBudgetUsageWarning over each config's model_budgets (scope label ${providerLabel} · ${model_name}) would close the gap.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 1167 -
1197, Extend hasBudgetResetRelevantChanges and getBudgetUsageWarning to include
each provider config’s model_budgets alongside config.budgets. Reuse
budgetSignature and findBudgetUsageWarning for per-model budget entries,
labeling their scope as `${providerLabel} · ${model_name}`, so edits trigger the
reset dialog and existing usage over the new limit produces a warning.

🧹 Nitpick comments (2)

transports/bifrost-http/handlers/governance_test.go (1)

116-131: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add coverage for the new per-model governance behavior.

The mock/signature updates keep existing tests compiling, but the new logic is untested: validateVKModelBudgetNames (limit, empty, wildcard, duplicate), buildVKModelBudgetsIndex grouping/sorting, and — most importantly — the prune-vs-preserve branch in syncVKGovernanceToModelConfigs (model_budgets omitted leaves existing per-model configs untouched; supplied set prunes absent models). That last branch is the one most likely to silently delete a customer's per-model budgets.

As per coding guidelines, apply "table-driven coverage for behavior changes".

Want me to draft these table-driven cases?

Also applies to: 3183-3183, 3218-3218

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@transports/bifrost-http/handlers/governance_test.go` around lines 116 - 131,
Add table-driven tests covering validateVKModelBudgetNames for limit, empty,
wildcard, and duplicate inputs; buildVKModelBudgetsIndex grouping and sorting;
and both syncVKGovernanceToModelConfigs branches, preserving existing per-model
configs when model_budgets is omitted and pruning absent models when it is
supplied. Use the existing mockRotateConfigStore and governance test helpers,
and assert the resulting configurations explicitly.

Source: Coding guidelines

ui/lib/types/governance.ts (1)

165-193: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider an update-specific model-budget request type.

The backend update payload uses UpdateRateLimitRequest semantics for model_budgets[].rate_limit (nulls / {} to clear), while VirtualKeyProviderConfigUpdateRequest.model_budgets reuses VirtualKeyModelBudgetRequest with CreateRateLimitRequest. virtualKeySheet.tsx already emits the removal shape, so the type no longer describes the wire payload (it's only hidden because normalization returns any[]).

As per coding guidelines, "Avoid any; prefer strict typing, inference, and reusable shared types."

♻️ Suggested split

+export interface VirtualKeyModelBudgetUpdateRequest {
+	model_name: string;
+	budgets?: CreateBudgetRequest[];
+	rate_limit?: UpdateRateLimitRequest;
+}

 export interface VirtualKeyProviderConfigUpdateRequest {
 	...
-	model_budgets?: VirtualKeyModelBudgetRequest[]; // Full desired per-model set when provider_configs is supplied
+	model_budgets?: VirtualKeyModelBudgetUpdateRequest[]; // Full desired per-model set when provider_configs is supplied
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ui/lib/types/governance.ts` around lines 165 - 193, Introduce an
update-specific model-budget request type alongside
VirtualKeyModelBudgetRequest, using UpdateRateLimitRequest for rate_limit while
preserving the shared model_name and budgets fields. Change
VirtualKeyProviderConfigUpdateRequest.model_budgets to use the new update type,
keeping the create request on VirtualKeyModelBudgetRequest and ensuring the type
represents null or empty-object removal payloads emitted by virtualKeySheet.tsx.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 418-490: Add BudgetOverrideDialog to each model budget row in the
config.model_budgets.map rendering, using the budget’s identifier and
virtual-key/model scope required by the override endpoint. Reuse the existing
provider or virtual-key budget dialog integration and effective-limit/override
data so model budgets support viewing and editing overrides consistently.

In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 100-124: Update the model_name field in the model_budgets schema
to reject empty or whitespace-only values and provide a meaningful inline
validation message, while preserving valid model names and the existing
normalization flow in normalizeProviderConfigs.

---

Outside diff comments:
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 1167-1197: Extend hasBudgetResetRelevantChanges and
getBudgetUsageWarning to include each provider config’s model_budgets alongside
config.budgets. Reuse budgetSignature and findBudgetUsageWarning for per-model
budget entries, labeling their scope as `${providerLabel} · ${model_name}`, so
edits trigger the reset dialog and existing usage over the new limit produces a
warning.

---

Nitpick comments:
In `@transports/bifrost-http/handlers/governance_test.go`:
- Around line 116-131: Add table-driven tests covering
validateVKModelBudgetNames for limit, empty, wildcard, and duplicate inputs;
buildVKModelBudgetsIndex grouping and sorting; and both
syncVKGovernanceToModelConfigs branches, preserving existing per-model configs
when model_budgets is omitted and pruning absent models when it is supplied. Use
the existing mockRotateConfigStore and governance test helpers, and assert the
resulting configurations explicitly.

In `@ui/lib/types/governance.ts`:
- Around line 165-193: Introduce an update-specific model-budget request type
alongside VirtualKeyModelBudgetRequest, using UpdateRateLimitRequest for
rate_limit while preserving the shared model_name and budgets fields. Change
VirtualKeyProviderConfigUpdateRequest.model_budgets to use the new update type,
keeping the create request on VirtualKeyModelBudgetRequest and ensuring the type
represents null or empty-object removal payloads emitted by virtualKeySheet.tsx.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 3b394f8d-d7e8-45e7-924c-2b52a26e72cf

📥 Commits

Reviewing files that changed from the base of the PR and between 19beff9 and 3d4c5f5.

📒 Files selected for processing (6)

• framework/configstore/tables/virtualkey.go
• transports/bifrost-http/handlers/governance.go
• transports/bifrost-http/handlers/governance_test.go
• ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
• ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
• ui/lib/types/governance.ts

The merge-base changed after approval.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

The merge-base changed after approval.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

The base branch was changed.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@transports/bifrost-http/handlers/governance.go`:
- Line 880: In reconcileVKModelConfig, declare modelName by calling
d.modelNameOrAll() before the query that uses it, ensuring the existing
references compile correctly.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: f1100786-815c-4f8a-aae9-80e4f032a3a9

📥 Commits

Reviewing files that changed from the base of the PR and between 40c75ff and b2209df.

📒 Files selected for processing (2)

• transports/bifrost-http/handlers/governance.go
• ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx

🚧 Files skipped from review as they are similar to previous changes (1)

• ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx

[akshaydeo]

Merge activity

• Aug 11, 2:09 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Aug 11, 2:09 PM UTC: @akshaydeo merged this pull request with Graphite.