feat: bringing provider governance to budget & limits

[roroghost17]

Summary

Provider-level governance (budget and rate limit) has been migrated from config_providers.budget_id / rate_limit_id into governance_model_configs as wildcard rows with scope='global', model_name='*', and provider=<name>. This makes governance_model_configs the single source of truth for all governance enforcement, eliminating the separate provider-governance enforcement path.

Changes

• Database migration (migrationMigrateProviderGovernanceToModelConfigs): Folds existing provider-level budget/rate-limit FK references into new (global, *, <provider>) model config rows, reusing the same budget/rate-limit rows. Provider FKs are then nulled out. The migration is idempotent and includes a rollback path.
• ModelConfigAllModels = "*" sentinel: Introduced as a named constant to represent "all models" in a model config row. The "*" sentinel is excluded from catalog-based model name normalization.
• collectModelConfigsFor: New helper that resolves all applicable model configs for a request across four tiers — exact model+provider, exact model (all providers), all models on this provider (*:provider), and all models on all providers (*:nil) — deduped by config ID. All budget/rate-limit check and usage-tracking paths now use this helper instead of duplicating two-tier lookup logic.
• modelConfigEntityKey: New helper that builds a stable entity key for a model config, replacing ad-hoc fmt.Sprintf strings scattered across check and update functions.
• GetProviderGovernanceModelConfigs: New store method that queries the wildcard model config rows backing provider governance, with budget and rate-limit preloads.
• DeleteProvider cleanup: When a provider is deleted, its associated wildcard model configs (and their owned budget/rate-limit rows) are now cleaned up as part of the transaction.
• Provider governance HTTP handlers: getProviderGovernance, updateProviderGovernance, and deleteProviderGovernance are rewritten to operate on wildcard model config rows instead of provider FK columns. The GET endpoint gains a from_memory query parameter to serve data from the in-memory governance store.
• UI: The model limits table renders "*" as "All Models". The model limit sheet exposes an allowAllOption on the model selector. RTK Query cache tags are cross-invalidated so that changes to model configs refresh the provider governance view and vice versa.

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 ./framework/configstore/... ./plugins/governance/...

# UI
cd ui
pnpm i
pnpm build

1. Start with a deployment that has providers with existing budget_id or rate_limit_id values. After the migration runs, verify those providers have budget_id = NULL and rate_limit_id = NULL, and that corresponding (global, *, <provider>) rows exist in governance_model_configs pointing to the same budget/rate-limit IDs.
2. Re-run the migration and confirm no duplicate wildcard rows are created (idempotency).
3. Make a request through a provider that had governance configured and confirm the budget/rate-limit is still enforced.
4. Use GET /api/governance/providers and GET /api/governance/providers?from_memory=true and confirm both return the expected provider governance data.
5. Use PUT and DELETE on /api/governance/providers/{name} and confirm the wildcard model config rows are created, updated, or removed accordingly, and that the Model Limits UI reflects the change without a manual refresh.
6. Delete a provider and confirm its wildcard model config and owned budget/rate-limit rows are removed.

Breaking changes

• Yes
• No

Provider governance is no longer stored in config_providers.budget_id / rate_limit_id. Any code or tooling that reads governance directly from the providers table will no longer find it there. All governance enforcement and management must go through governance_model_configs. The HTTP API surface is unchanged; the migration handles existing data automatically.

Security considerations

Budget and rate-limit rows are reused (not duplicated) during migration. The rollback path restores FK references to the provider table and removes the wildcard model config rows, leaving no orphaned governance rows in either direction.

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

Summary by CodeRabbit

• New Features

  • Governance policies can be configured for "All Models" on a specific provider, enabling provider-level budget and rate-limit controls.

• UI Improvements

  • Model limits table shows "All Models" for wildcard policies.
  • Model picker supports an "All" option for creating provider-scoped policies.
  • Changes now cause related views to refresh so updates appear immediately.

• Bug Fixes

  • Provider-scoped governance is applied and cleaned up reliably (including on provider removal).

• Tests

  • Added tests covering provider-scoped all-models governance and migrations.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 6c75d038-1218-4f08-af16-41341e44edc9

📥 Commits

Reviewing files that changed from the base of the PR and between f3ef789 and ea9049f.

📒 Files selected for processing (13)

• framework/configstore/migrations.go
• framework/configstore/migrations_test.go
• framework/configstore/rdb.go
• framework/configstore/rdb_test.go
• framework/configstore/store.go
• framework/configstore/tables/modelconfig.go
• plugins/governance/modelprovidergovernance_test.go
• plugins/governance/store.go
• transports/bifrost-http/handlers/governance.go
• transports/bifrost-http/lib/config_test.go
• ui/app/workspace/model-limits/views/modelLimitSheet.tsx
• ui/app/workspace/model-limits/views/modelLimitsTable.tsx
• ui/lib/store/apis/governanceApi.ts

---

📝 Walkthrough

Walkthrough

Provider-level governance (budget and rate-limit configuration) is migrated from provider row foreign keys to wildcard model configs, enabling provider-scoped "all models" governance policies. A database migration moves existing provider governance to wildcard entries, in-memory policy evaluation is refactored to support tiered config collection, HTTP endpoints transition to wildcard config persistence, and UI/cache invalidation ensure consistency.

Changes

Provider Governance Wildcard Model Config Migration

Layer / File(s)	Summary
Data Model & Interfaces framework/configstore/tables/modelconfig.go, framework/configstore/store.go	ModelConfigAllModels constant ("*") and ConfigStore.GetProviderGovernanceModelConfigs() interface method establish the contract for fetching provider-scoped wildcard governance configs.
Database Migration & Storage Layer framework/configstore/migrations.go, framework/configstore/migrations_test.go, framework/configstore/rdb.go, framework/configstore/rdb_test.go	Forward migration idempotently creates wildcard model configs from provider row governance FKs and nulls those FKs; rollback re-attaches FKs and deletes wildcard rows. RDBConfigStore.GetProviderGovernanceModelConfigs() fetches provider wildcard configs; DeleteProvider cleans up associated governance model configs and their owned budget/rate-limit rows. Tests validate migration correctness and provider deletion cleanup.
In-Memory Governance Policy Evaluation & Usage Tracking plugins/governance/store.go	collectModelConfigsFor centralizes tiered model-config selection across four tiers (exact model+provider, normalized model-only, provider-level wildcard, global wildcard) with deduplication. CheckModelBudget, CheckModelRateLimit, and scoped variants are refactored to evaluate budgets/limits across all applicable tiers. Usage-tracking methods now bump counters via collectModelConfigsFor, including the new provider-level "all models" tier. Model-name normalization is adjusted to preserve the wildcard sentinel.
Provider Governance HTTP Handler transports/bifrost-http/handlers/governance.go	getProviderGovernance reads provider governance from wildcard configs (with in-memory fallback), updateProviderGovernance creates/updates/deletes wildcard configs and their budget/rate-limit entities, and deleteProviderGovernance removes wildcard configs idempotently. All paths sync in-memory governance state after DB transactions.
Governance Tests & Mocks plugins/governance/modelprovidergovernance_test.go, transports/bifrost-http/lib/config_test.go, framework/configstore/migrations_test.go, framework/configstore/rdb_test.go	New tests verify provider-scoped all-models wildcard budget/rate-limit enforcement (including provider isolation and usage tracking). MockConfigStore is extended with GetProviderGovernanceModelConfigs() stub. Migration and deletion tests validate DB forward/rollback and cleanup.
UI Display & Cache Invalidation ui/app/workspace/model-limits/views/modelLimitSheet.tsx, ui/app/workspace/model-limits/views/modelLimitsTable.tsx, ui/lib/store/apis/governanceApi.ts	Model limits UI enables all-models selection and renders "*" as "All Models". RTK Query mutations invalidate cache tags: model config mutations invalidate ProviderGovernance cache, provider governance mutations invalidate ModelConfigs cache, ensuring cross-feature refreshes.

Sequence Diagram

sequenceDiagram
  participant Client
  participant Handler as HTTP Handler
  participant Storage as RDB ConfigStore
  participant InMem as In-Memory Governance
  Client->>Handler: POST /provider-governance (update)
  Handler->>Storage: Validate provider exists
  Handler->>Storage: Fetch or create provider wildcard model config
  Handler->>Storage: Upsert/delete wildcard model config + budget/rate-limit lifecycle
  Handler->>InMem: remove or reload model config (sync)
  Handler-->>Client: updated provider governance response
  Note over Client,InMem: Subsequent checks evaluate wildcard tier
  Client->>InMem: CheckModelBudget (model, provider)
  InMem->>InMem: collectModelConfigsFor (exact, normalized, provider-wildcard, global-wildcard)
  InMem-->>Client: budget/rate-limit decision

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• maximhq/bifrost#3937: Both PRs extend framework/configstore/migrations.go/migrations_test.go around the new governance_model_configs.scope columns—this PR adds a follow-on migration dependent on those scope changes.

Suggested reviewers

• danpiths
• akshaydeo

"A rabbit hopped through code tonight,
Budgets migrated by moonlit byte,
Wildcard rules now govern with glee,
All models bound to provider-tree,
Hooray — configs hop home, safe and bright!" 🐇

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 77.78% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: migrating provider governance to the budget and limits system via model configs.
Description check	✅ Passed	The description comprehensively covers the summary, detailed changes, type, affected areas, testing steps, breaking changes, and security considerations, following the template structure.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 05-27-feat_bringing_provider_governance_to_budget_limits

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

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

[CLAassistant]

All committers have signed the CLA.

[roroghost17]

• refactor: makes scope-level check methods extensible #3940 : 2 dependent PRs (#3960 , #3962 )
• feat: wires vk top-level and provider-level budgets from model configs table #3939
• feat: bringing provider governance to budget & limits #3938 👈 (View in Graphite)
• feat: adding scope to budget limits table #3937
• dev

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

[greptile-apps]

Confidence Score: 1/5

Not safe to merge — the governance package does not compile due to two call sites referencing a removed method, and the DeleteProvider cleanup leaves wildcard model-config rows orphaned in the database.

The governance plugin calls findModelOnlyConfig in two places (CollectApplicableGovernanceIDs and GetBudgetAndRateLimitStatus) but that method no longer exists anywhere in the codebase after being superseded by findScopedModelOnlyConfig. The package cannot be built. Separately, DeleteProvider in rdb.go declares mcIDs []string and immediately uses it in a WHERE id IN ? delete without ever appending the snapshotted IDs, so wildcard model-config rows survive provider deletion with dangling budget/rate-limit references already gone.

plugins/governance/store.go (undefined method at lines 2470 and 3540), framework/configstore/rdb.go (unpopulated mcIDs slice in DeleteProvider)

Important Files Changed

Filename	Overview
plugins/governance/store.go	Central refactor introducing collectModelConfigsFor (4-tier wildcard lookup), modelConfigEntityKey, and modelAndProvider. Two unreachable call sites to the now-deleted findModelOnlyConfig in CollectApplicableGovernanceIDs and GetBudgetAndRateLimitStatus cause a build failure.
framework/configstore/rdb.go	DeleteProvider loops to delete owned budget/rate-limit rows but the mcIDs slice is declared and immediately used in WHERE id IN ? without being populated, so wildcard model configs are never actually deleted from the DB. GetProviderGovernanceModelConfigs correctly preloads Budget/RateLimit.
transports/bifrost-http/handlers/governance.go	Provider governance handlers rewritten to operate on wildcard model configs; several previously-noted issues remain unresolved.
framework/configstore/migrations.go	New migration folds provider-level governance FKs into (global, *, provider) wildcard model-config rows and nulls provider FKs. Idempotent guard and rollback path look correct.
plugins/governance/modelprovidergovernance_test.go	Good new tests covering *:provider budget/rate-limit enforcement and cross-provider isolation. Tests will only run once the build failure in store.go is fixed.
framework/configstore/rdb_test.go	TestDeleteProvider_CleansUpProviderModelConfigs is added but will fail at the model-config deletion assertion due to the unpopulated mcIDs slice in DeleteProvider.
framework/configstore/tables/modelconfig.go	Adds ModelConfigAllModels = "*" sentinel constant with clear doc comment. Clean, no issues.
framework/configstore/store.go	Adds GetProviderGovernanceModelConfigs to the ConfigStore interface. Straightforward addition, no issues.
ui/lib/store/apis/governanceApi.ts	Adds cross-invalidation cache tags so model-limit changes refresh the provider governance view and vice versa. Correct.
ui/app/workspace/model-limits/views/modelLimitsTable.tsx	Renders "*" as "All Models" label. No data-testid removals, no issues.
ui/app/workspace/model-limits/views/modelLimitSheet.tsx	Adds allowAllOption to the model selector for creating provider-scoped * governance configs. One-line change, no issues.
transports/bifrost-http/lib/config_test.go	Adds stub GetProviderGovernanceModelConfigs to mock store to satisfy the updated interface. No issues.
framework/configstore/migrations_test.go	Migration test validates idempotency and FK cleanup. Well-structured and covers the intended scenarios.

_{Reviews (4): Last reviewed commit: "feat: bringing provider governance to bu..." | Re-trigger Greptile}

[roroghost17]

@coderabbitai full-review

[coderabbitai]

@roroghost17: I'll kick off a full review of this PR right away!

✅ Actions performed

Full review triggered.

[coderabbitai]

Actionable comments posted: 8

🤖 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 `@framework/configstore/migrations_test.go`:
- Around line 2441-2479: Add additional unit tests alongside
TestMigrationMigrateProviderGovernanceToModelConfigs to cover edge cases: create
a provider with only BudgetID (nil RateLimitID), a provider with only
RateLimitID (nil BudgetID), a provider with neither (ensure no wildcard config
created), and multiple providers (each gets its own wildcard config). For each
new test, call migrationMigrateProviderGovernanceToModelConfigs(ctx, db) and
assert expected TableModelConfig rows (query TableModelConfig by
scope/model_name/provider) and that provider TableProvider budget_id and
rate_limit_id are cleared as appropriate; also verify idempotency by re-running
the migration and checking no duplicates are created. Ensure tests reuse the
same setup pattern (setupTestDB, AutoMigrate for
tables.TableProvider/tables.TableModelConfig/tables.TableBudget/tables.TableRateLimit)
and unique IDs for budgets/rate-limits to avoid cross-test collisions.

In `@framework/configstore/migrations.go`:
- Around line 3983-4005: The rollback currently mutates and deletes canonical
provider-governance rows (see the tx query selecting tables.TableModelConfig
rows into wildcards and the subsequent updates to tables.TableProvider and
tx.Delete of mc.ID), which can overwrite live data; instead mark this migration
as non-rollbackable by removing these destructive Down steps and returning an
explicit irreversible error (e.g. return fmt.Errorf("irreversible migration:
provider-governance canonicalization cannot be rolled back")) from the
migration's Down/rollback handler so callers know the migration cannot be safely
reverted.

In `@framework/configstore/rdb.go`:
- Around line 1159-1176: The current deletion first snapshots
providerModelConfigs and then runs a second WHERE "provider = ?" delete, which
allows a race with concurrent CreateModelConfig/UpdateModelConfig; to fix,
delete exactly the snapped rows (use the collected providerModelConfigs' IDs in
a single DELETE WHERE id IN (...) so you only remove the rows you read) and
additionally serialize concurrent model-config writers by taking a row-level
lock on the provider in CreateModelConfig and UpdateModelConfig (use a FOR
UPDATE/Locking clause on the provider select before creating/updating model
configs) so no new provider-scoped rows can be inserted between your snapshot
and delete; reference providerModelConfigs, CreateModelConfig,
UpdateModelConfig, and the txDB delete call to locate the changes.

In `@framework/configstore/store.go`:
- Line 259: Clarify the method comment for GetProviderGovernanceModelConfigs to
state it returns only provider-specific wildcard model configs (i.e., entries
where provider != nil) and does not include global wildcard configs (provider ==
nil); update the docstring to explicitly say this so implementers know the
function should filter and return only provider-scoped "all models on a
provider" configurations (matching the downstream handler's expectations).

In `@plugins/governance/store.go`:
- Around line 1087-1126: GetBudgetAndRateLimitStatus and
CollectApplicableGovernanceIDs currently only check exact-model and model-only
rows so they miss provider-level wildcard configs ("*", provider) and global
wildcards ("*", nil); update both functions to call collectModelConfigsFor(ctx,
scope, scopeID, model, provider) and iterate the returned
[]*configstoreTables.TableModelConfig to build budgets, rate-limits and
governance ID lists (honoring the deduping/order already applied by
collectModelConfigsFor), handling provider==nil and preserving existing behavior
for model-only and exact-model entries. Ensure you replace any direct lookups
like findScopedModelOnlyConfig or gs.modelConfigs.Load for model entries inside
those methods with the collectModelConfigsFor call and use the returned slice to
compute status and attribution.
- Around line 1281-1299: The code currently overwrites
entityWiseRateLimits[modelConfigEntityKey(mc)] with a single-element slice,
losing multiple matching rate limits and causing CheckRateLimit to see only one
violation; change the logic in CheckModelRateLimit (and the similar block at
1327-1351) to append rate limits instead of replacing them: for each mc with
mc.RateLimitID, load the rateLimit and if non-nil do existing :=
entityWiseRateLimits[modelConfigEntityKey(mc)];
entityWiseRateLimits[modelConfigEntityKey(mc)] = append(existing, rateLimit) so
CheckRateLimit receives all matching limits and can correctly aggregate
violations into DecisionRateLimited when both tokens and requests are exceeded.

In `@transports/bifrost-http/handlers/governance.go`:
- Around line 3139-3144: Replace the inline removal check that examines only
req.RateLimit.TokenMaxLimit and req.RateLimit.RequestMaxLimit with a call to the
existing helper isRateLimitRemovalRequest to centralize semantics; update the
block that mutates mc.RateLimitID and mc.RateLimit so it runs when
isRateLimitRemovalRequest(req.RateLimit) returns true instead of the current
two-field check (ensure you keep the assignments rateLimitIDToDelete =
*mc.RateLimitID, mc.RateLimitID = nil, mc.RateLimit = nil), and verify this
behavioral change (reset-duration-only requests no longer count as removals) is
acceptable.

In `@ui/lib/store/apis/governanceApi.ts`:
- Line 550: Add a short explanatory comment above the invalidatesTags:
["ProviderGovernance"] occurrence in governanceApi.ts that mirrors the rationale
used at the other two spots (lines where ProviderGovernance is invalidated) —
mention the cross-entity relationship and why invalidating ProviderGovernance is
necessary for this endpoint; keep the phrasing and placement consistent with the
comments around the other invalidatesTags instances.

🪄 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: ASSERTIVE

Plan: Pro

Run ID: 4c379fe0-1870-41ba-bb5e-9abd8e53c73c

📥 Commits

Reviewing files that changed from the base of the PR and between eb27aa8 and f3ef789.

📒 Files selected for processing (13)

• framework/configstore/migrations.go
• framework/configstore/migrations_test.go
• framework/configstore/rdb.go
• framework/configstore/rdb_test.go
• framework/configstore/store.go
• framework/configstore/tables/modelconfig.go
• plugins/governance/modelprovidergovernance_test.go
• plugins/governance/store.go
• transports/bifrost-http/handlers/governance.go
• transports/bifrost-http/lib/config_test.go
• ui/app/workspace/model-limits/views/modelLimitSheet.tsx
• ui/app/workspace/model-limits/views/modelLimitsTable.tsx
• ui/lib/store/apis/governanceApi.ts

[akshaydeo]

Merge activity

• Jun 2, 1:06 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 2, 1:09 PM UTC: Graphite rebased this pull request as part of a merge.
• Jun 2, 1:10 PM UTC: @akshaydeo merged this pull request with Graphite.