refactor: makes scope-level check methods extensible - #3940
Merged
akshaydeo merged 1 commit intoJun 2, 2026
Conversation
Contributor
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (18)
📝 WalkthroughWalkthroughThis PR generalizes scoped model-config governance: it adds a runtime scope registry, converts VK-specific scoped APIs to generic (scope, scopeID) checks/updates, rewrites VK handler reconciliation to upsert VK-scoped model-configs, updates migrations and DB cleanup to bulk-delete associated records, adjusts resolver/tracker enforcement and tests, and makes the UI registry-driven. ChangesScoped Model Config Governance
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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 |
Contributor
Author
This stack of pull requests is managed by Graphite. Learn more about stacking. |
18 tasks
|
|
This was referenced May 31, 2026
roroghost17
force-pushed
the
06-01-refactor_makes_scope-level_check_methods_extensible
branch
2 times, most recently
from
June 1, 2026 10:04
41da4b0 to
c1eab5c
Compare
roroghost17
marked this pull request as ready for review
June 1, 2026 10:08
Contributor
Confidence Score: 3/5The refactor is broadly correct but both UpdateScopedModel*UsageInMemory methods no longer short-circuit on an empty model string, which can silently increment wildcard-config counters — both need the || model == "" guard restored before merging. Both UpdateScopedModelBudgetUsageInMemory and UpdateScopedModelRateLimitUsageInMemory dropped the || model == "" early-return that the old UpdateVirtualKeyScoped* methods had. Because collectModelConfigsFor always evaluates tiers 3 and 4 regardless of the model argument, passing an empty model matches and bumps every wildcard config for the given scope/scopeID. Production tracker call-sites are guarded today, but enterprise GovernanceStore implementations following the documented contract could trigger silent budget/rate-limit corruption. plugins/governance/store.go — UpdateScopedModelBudgetUsageInMemory and UpdateScopedModelRateLimitUsageInMemory both need the || model == "" guard restored. Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
Author
|
@coderabbitai full-review |
Contributor
|
✅ Actions performedFull review triggered. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
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 (4)
ui/app/workspace/model-limits/views/modelLimitSheet.tsx (1)
59-62:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winValidation logic must be registry-aware for enterprise scopes.
The schema hardcodes
"virtual_key"in the refinement, so it won't requirescopeIdfor enterprise scopes that also need a target (e.g.,"user"). The payload logic at lines 225-229 correctly uses the registry to decide whether to sendscope_id; validation should mirror that.🔧 Proposed fix to make validation registry-driven
+import { getModelLimitScope, getModelLimitScopes } from "`@/lib/registries/modelLimitScopes`"; + const formSchema = z .object({ modelName: z.string().min(1, "Model name is required"), ... }) - .refine((data) => data.scope !== "virtual_key" || !!data.scopeId, { - message: "Virtual key is required for the Virtual Key scope", + .refine((data) => { + const scopeEntry = getModelLimitScope(data.scope || "global"); + return !scopeEntry?.PickerComponent || !!data.scopeId; + }, { + message: "Scope target is required", path: ["scopeId"], });🤖 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/model-limits/views/modelLimitSheet.tsx` around lines 59 - 62, The current refinement hardcodes "virtual_key" and should instead consult the same registry logic used in the payload code to decide when scopeId is required; update the .refine check to ask the scope registry (the same object/utility used where the payload decides to include scope_id) whether the given data.scope requires a target/id and validate that !!data.scopeId when it does (use the same registry lookup function used in the payload logic to determine required scopes so validation stays in sync).plugins/governance/store.go (1)
1325-1364:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject unknown scope names instead of silently allowing them.
The old VK-only API could not be mistyped, but this string-based version will quietly return
DecisionAllowfor any non-empty, unregistered scope because the lookup just misses every key. That is a fail-open regression on a governance/rate-limit path. Please validatescopewithconfigstoreTables.IsValidModelConfigScope(scope)and return an error for unknown values; the twoUpdateScopedModel*methods should mirror the same guard.
As per coding guidelines: "Review budget, rate-limit, virtual key, and RBAC paths for fail-closed behavior where security is involved."🤖 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 `@plugins/governance/store.go` around lines 1325 - 1364, Validate the incoming scope string at the start of CheckScopedModelBudget and CheckScopedModelRateLimit by calling configstoreTables.IsValidModelConfigScope(scope) and return a non-nil error (not DecisionAllow) for unknown/invalid scope values; also add the same guard to the two UpdateScopedModel* methods so they reject unknown scope names instead of proceeding silently, ensuring all model-config budget/rate-limit paths fail-closed for invalid scopes.framework/configstore/migrations.go (1)
4083-4110:⚠️ Potential issue | 🟠 Major | ⚡ Quick winGuard
calendar_alignedwrites until the column exists.
triggerMigrationsrunsmigrationMigrateVirtualKeyGovernanceToModelConfigsbeforemigrationAddModelConfigCalendarAlignedColumn(lines ~840-845), butensureVKModelConfigpassescalendarAlignedinto thetables.TableModelConfig{ CalendarAligned: ... }create payload (lines ~4083-4107), and the call sites passvk.CalendarAligned(lines ~4144 and ~4173). On upgrade paths wheregovernance_model_configs.calendar_alignedhasn’t been added yet, this can fail with an unknown-column error.Add a guard before calling
ensureVKModelConfig(e.g.,tx.Migrator().HasColumn(&tables.TableModelConfig{}, "calendar_aligned")) and only write/enableCalendarAlignedwhen the column is present.🤖 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 `@framework/configstore/migrations.go` around lines 4083 - 4110, The code writes CalendarAligned into tables.TableModelConfig in ensureVKModelConfig which can run before the calendar_aligned column exists; before calling ensureVKModelConfig (from migrationMigrateVirtualKeyGovernanceToModelConfigs or its callers) check whether the column exists using tx.Migrator().HasColumn(&tables.TableModelConfig{}, "calendar_aligned") and only pass/assign CalendarAligned (or set the calendarAligned argument) when that check is true; alternatively, modify ensureVKModelConfig to accept a nil/absent flag and skip setting CalendarAligned on the created mc unless tx.Migrator().HasColumn reports the column present. Ensure references to ensureVKModelConfig, migrationMigrateVirtualKeyGovernanceToModelConfigs, tables.TableModelConfig, and tx.Migrator().HasColumn are used so the guard is applied where the create occurs.framework/configstore/rdb.go (1)
2952-2975:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDelete scoped model-config budgets before removing the configs.
This path only cleans up the legacy
BudgetIDfield. Any budgets owned throughmc.Budgetssurvive theDelete(&tables.TableModelConfig{})call, so deleting a virtual key leaks scoped model-budget rows.Proposed fix
var scopedModelConfigs []tables.TableModelConfig if err := txDB.WithContext(ctx).Preload("Budgets"). Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id). Find(&scopedModelConfigs).Error; err != nil { return err } + var scopedBudgetIDs []string + var scopedRateLimitIDs []string for _, mc := range scopedModelConfigs { + for i := range mc.Budgets { + scopedBudgetIDs = append(scopedBudgetIDs, mc.Budgets[i].ID) + } if mc.BudgetID != nil { - if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", *mc.BudgetID).Error; err != nil { - return err - } + scopedBudgetIDs = append(scopedBudgetIDs, *mc.BudgetID) } if mc.RateLimitID != nil { - if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *mc.RateLimitID).Error; err != nil { - return err - } + scopedRateLimitIDs = append(scopedRateLimitIDs, *mc.RateLimitID) } } if err := txDB.WithContext(ctx). Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id). Delete(&tables.TableModelConfig{}).Error; err != nil { return err } + if len(scopedBudgetIDs) > 0 { + if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id IN ?", scopedBudgetIDs).Error; err != nil { + return err + } + } + if len(scopedRateLimitIDs) > 0 { + if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id IN ?", scopedRateLimitIDs).Error; err != nil { + return err + } + }Based on learnings: budgets and rate limits have a 1:1 ownership with their parent entities and should be deleted together.
🤖 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 `@framework/configstore/rdb.go` around lines 2952 - 2975, The current cleanup only deletes the legacy BudgetID and RateLimitID fields but leaves any budgets in the preloaded mc.Budgets slice, leaking scoped model-budget rows; update the loop over scopedModelConfigs (scopedModelConfigs, TableModelConfig, Budgets, BudgetID, TableBudget, RateLimitID, TableRateLimit) to also remove all entries in mc.Budgets before deleting the model config—e.g., iterate mc.Budgets and delete each budget (or issue a bulk delete for those budget IDs) using txDB.WithContext(ctx).Delete(...), then continue deleting legacy BudgetID/RateLimitID as already done so that all owned budgets and rate-limits are removed prior to deleting the TableModelConfig rows.
🤖 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 `@plugins/governance/store.go`:
- Around line 1481-1483: The early return in
UpdateScopedModelBudgetUsageInMemory currently skips processing when model == ""
which prevents wildcard scoped entries from being matched; remove model == ""
from the early-return condition so the function still runs
collectModelConfigsFor(scope, scopeID, "", providerStr) and updates "*"
model/provider tiers when model is empty; apply the same change to the analogous
scoped rate-limit updater (the other function around lines 1502-1504) so both
UpdateScopedModelBudgetUsageInMemory and its rate-limit counterpart process
wildcard entries even when the incoming model string is empty.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 923-925: The handler is mutating shared in-memory virtual key
objects because it calls applyVKGovernanceFromModelConfigs on pointers from
data.VirtualKeys; to fix, deep-clone each TableVirtualKey (including its
Budgets, RateLimit/RateLimitID and provider-config slice/map tree) before
calling applyVKGovernanceFromModelConfigs so the hydration only affects the
response copy, not the shared GovernanceManager state; update the loop that uses
buildVKModelConfigIndex/virtualKeys to clone each vk (and do the same fix for
the similar block at lines 1232-1235) and ensure the clone preserves identity
fields but is a separate in-memory structure to avoid races.
---
Outside diff comments:
In `@framework/configstore/migrations.go`:
- Around line 4083-4110: The code writes CalendarAligned into
tables.TableModelConfig in ensureVKModelConfig which can run before the
calendar_aligned column exists; before calling ensureVKModelConfig (from
migrationMigrateVirtualKeyGovernanceToModelConfigs or its callers) check whether
the column exists using tx.Migrator().HasColumn(&tables.TableModelConfig{},
"calendar_aligned") and only pass/assign CalendarAligned (or set the
calendarAligned argument) when that check is true; alternatively, modify
ensureVKModelConfig to accept a nil/absent flag and skip setting CalendarAligned
on the created mc unless tx.Migrator().HasColumn reports the column present.
Ensure references to ensureVKModelConfig,
migrationMigrateVirtualKeyGovernanceToModelConfigs, tables.TableModelConfig, and
tx.Migrator().HasColumn are used so the guard is applied where the create
occurs.
In `@framework/configstore/rdb.go`:
- Around line 2952-2975: The current cleanup only deletes the legacy BudgetID
and RateLimitID fields but leaves any budgets in the preloaded mc.Budgets slice,
leaking scoped model-budget rows; update the loop over scopedModelConfigs
(scopedModelConfigs, TableModelConfig, Budgets, BudgetID, TableBudget,
RateLimitID, TableRateLimit) to also remove all entries in mc.Budgets before
deleting the model config—e.g., iterate mc.Budgets and delete each budget (or
issue a bulk delete for those budget IDs) using
txDB.WithContext(ctx).Delete(...), then continue deleting legacy
BudgetID/RateLimitID as already done so that all owned budgets and rate-limits
are removed prior to deleting the TableModelConfig rows.
In `@plugins/governance/store.go`:
- Around line 1325-1364: Validate the incoming scope string at the start of
CheckScopedModelBudget and CheckScopedModelRateLimit by calling
configstoreTables.IsValidModelConfigScope(scope) and return a non-nil error (not
DecisionAllow) for unknown/invalid scope values; also add the same guard to the
two UpdateScopedModel* methods so they reject unknown scope names instead of
proceeding silently, ensuring all model-config budget/rate-limit paths
fail-closed for invalid scopes.
In `@ui/app/workspace/model-limits/views/modelLimitSheet.tsx`:
- Around line 59-62: The current refinement hardcodes "virtual_key" and should
instead consult the same registry logic used in the payload code to decide when
scopeId is required; update the .refine check to ask the scope registry (the
same object/utility used where the payload decides to include scope_id) whether
the given data.scope requires a target/id and validate that !!data.scopeId when
it does (use the same registry lookup function used in the payload logic to
determine required scopes so validation stays in sync).
🪄 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: c4c34ca0-4d32-43ee-a919-9859473eb222
📒 Files selected for processing (15)
framework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/tables/modelconfig.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.gotransports/bifrost-http/handlers/governance.goui/app/_fallbacks/enterprise/lib/registrations/modelLimitScopes.tsui/app/workspace/model-limits/views/modelLimitSheet.tsxui/app/workspace/model-limits/views/modelLimitsTable.tsxui/lib/constants/governance.tsui/lib/registries/modelLimitScopes.tsxui/lib/store/apis/governanceApi.tsui/lib/utils/labels.ts
💤 Files with no reviewable changes (1)
- ui/lib/constants/governance.ts
roroghost17
force-pushed
the
06-01-refactor_makes_scope-level_check_methods_extensible
branch
2 times, most recently
from
June 1, 2026 14:28
963dbb9 to
c557c1a
Compare
This was referenced Jun 1, 2026
roroghost17
force-pushed
the
05-28-feat_wires_vk_top-level_and_provider-level_budgets_from_model_configs_table
branch
from
June 1, 2026 17:56
476a36b to
0afcab3
Compare
roroghost17
force-pushed
the
06-01-refactor_makes_scope-level_check_methods_extensible
branch
2 times, most recently
from
June 1, 2026 18:06
096e68f to
e63c607
Compare
This was referenced Jun 2, 2026
Contributor
Merge activity
|
akshaydeo
changed the base branch from
05-28-feat_wires_vk_top-level_and_provider-level_budgets_from_model_configs_table
to
graphite-base/3940
June 2, 2026 13:11
roroghost17
force-pushed
the
06-01-refactor_makes_scope-level_check_methods_extensible
branch
from
June 2, 2026 13:37
e63c607 to
4816c78
Compare
roroghost17
force-pushed
the
graphite-base/3940
branch
from
June 2, 2026 13:37
0afcab3 to
39828ee
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 17
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
transports/bifrost-http/handlers/governance.go (1)
2972-3029:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winReturn 400 for model-config validation failures.
These mutation paths run
validateBudget/validateRateLimitinside the transaction, but the outer handlers only translate failures to 500. Reachable client errors likemax_limit: 0or an invalid reset duration will therefore be reported as server faults even though nothing committed. Wrap validation failures inbadRequestErroror pre-validate before entering the transaction.One way to keep the current structure
}); err != nil { + var badReqErr *badRequestError + if errors.As(err, &badReqErr) { + SendError(ctx, 400, badReqErr.Error()) + return + } logger.Error("failed to create model config: %v", err) SendError(ctx, 500, fmt.Sprintf("Failed to create model config: %v", err)) return }Also wrap transaction-time validation failures the same way:
- if err := validateRateLimit(&rateLimit); err != nil { - return err + if err := validateRateLimit(&rateLimit); err != nil { + return &badRequestError{err: err} }Also applies to: 3061-3147
🤖 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.go` around lines 2972 - 3029, The handler is returning 500 for client-side validation failures originating from validateBudget/validateRateLimit inside h.configStore.ExecuteTransaction; change it to return 400 by either pre-validating all req.Budgets and req.RateLimit before calling ExecuteTransaction or by detecting validation errors returned from within the transaction and wrapping them in a badRequestError (or the project’s equivalent) so the outer error handling translates them to a 400 response; update both the model-config creation block (where validateRateLimit and validateBudget are called) and the analogous block handling the other mutation (the one referenced around the later similar code) to consistently wrap/translate validation errors to badRequestError before calling SendError.plugins/governance/store.go (1)
3128-3138:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve rate-limit reset timestamps here too.
UpdateModelConfigInMemorycarries the live counters forward but dropsTokenLastResetandRequestLastReset. Updating a model config with the same rate-limit ID can therefore shift the in-memory reset window and produce premature or delayed resets on the next check.Suggested fix
if clone.RateLimit != nil { clone.RateLimit.IsCalendarAligned = clone.CalendarAligned if existingRateLimitValue, exists := gs.rateLimits.Load(clone.RateLimit.ID); exists && existingRateLimitValue != nil { if erl, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && erl != nil { clone.RateLimit.TokenCurrentUsage = erl.TokenCurrentUsage clone.RateLimit.RequestCurrentUsage = erl.RequestCurrentUsage + clone.RateLimit.TokenLastReset = erl.TokenLastReset + clone.RateLimit.RequestLastReset = erl.RequestLastReset } } gs.rateLimits.Store(clone.RateLimit.ID, clone.RateLimit) }🤖 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 `@plugins/governance/store.go` around lines 3128 - 3138, When storing the rate limit in UpdateModelConfigInMemory (the block using clone.RateLimit and gs.rateLimits.Store), also preserve the last-reset timestamps from the in-memory entry: when retrieving existingRateLimitValue (cast to *configstoreTables.TableRateLimit as erl), copy erl.TokenLastReset and erl.RequestLastReset into clone.RateLimit.TokenLastReset and clone.RateLimit.RequestLastReset (with appropriate nil/type checks) before calling gs.rateLimits.Store, so updates with the same rate-limit ID keep the live reset windows intact.
🤖 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.go`:
- Around line 4208-4229: The code currently deletes the TableModelConfig row
(tx.Delete(&tables.TableModelConfig{}, "id = ?", mc.ID)) even when no provider
config (pcs) is found, which can orphan budgets; change the logic to avoid
deleting the model config unless budgets were successfully restored: after the
tx.Where(...) query, if len(pcs) == 0 return a non-rollbackable error (or simply
return an error indicating rollback cannot proceed) instead of proceeding to
tx.Delete, or move the tx.Delete call inside the if len(pcs) > 0 block so
deletion only happens when pcID and budget updates succeeded; reference
mc.ScopeID, mc.Provider, pcs, pcID, budgets, mc.RateLimitID, and mc.ID when
implementing this check.
- Around line 3881-3893: The new unique index creation must use PostgreSQL's
CONCURRENTLY option and avoid running inside a transaction to prevent long
locks: replace the plain migrator.CreateIndex call for
"idx_model_scope_provider" on modelConfig with a dialect check for Postgres
(e.g., DB.Dialector.Name() or migrator.DB.Dialector.Name()), and for Postgres
execute a raw "CREATE UNIQUE INDEX CONCURRENTLY ..." SQL statement via DB.Exec
(construct the index name and columns to match the struct tags) and handle
errors; ensure this path is not executed inside a transaction (CONCURRENTLY
cannot run in a transaction) and keep the existing migrator.CreateIndex fallback
for non-Postgres dialects, then continue to drop the old "idx_model_provider"
index as before.
- Around line 3933-3961: When a wildcard TableModelConfig (ModelConfigAllModels)
already exists for a provider you must merge the provider's governance FKs into
that row before clearing them on the TableProvider; modify the branch where
existing > 0 to locate the existing TableModelConfig (using the same
tx.Model(&tables.TableModelConfig{}).Where("scope = ? AND model_name = ? AND
provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels,
p.Name)) and Update that record to set/merge budget_id and rate_limit_id from
p.BudgetID and p.RateLimitID (e.g. only set if the target FK is nil or use
COALESCE-like logic), then continue to clear the provider FKs on TableProvider —
ensure you use the same tx instance and handle/return any errors from the merge
update before nulling provider fields.
In `@framework/configstore/rdb.go`:
- Around line 3003-3026: The code only collects legacy BudgetID/RateLimitID
fields and deletes TableModelConfig rows by scope (which may remove rows created
after the snapshot) causing owned budgets/rate-limits to leak; fix by (1)
loading the snapshot's has-many relations (Preload("Budgets") and
Preload("RateLimits") on the TableModelConfig query) and collect both legacy IDs
and the IDs from mc.Budgets and mc.RateLimits, (2) collect the snapshot
TableModelConfig IDs and delete budgets and rate-limits by their ModelConfigID
OR by collected budget/rate IDs (use TableBudget.ModelConfigID IN ? and
TableRateLimit.ModelConfigID IN ? plus any legacy id lists), and (3) delete
TableModelConfig rows by their specific IDs (WHERE id IN ?) instead of reusing
WHERE scope = ? AND scope_id = ? so you only remove the exact snapshot rows;
reference TableModelConfig (ID, BudgetID, RateLimitID, Budgets, ModelConfigID),
TableBudget, TableRateLimit, ModelConfigScopeVirtualKey, and txDB.WithContext
calls to locate the code.
- Around line 4344-4351: In RDBConfigStore.GetModelConfig, detect the invalid
call where scope != "global" and scopeID == nil and immediately return an error
instead of translating that into "scope_id IS NULL"; update the start of
GetModelConfig to validate scope and scopeID (using the scope and scopeID
parameters) and return a clear error when a non-global scope is missing
scope_id, otherwise proceed with building the query (keep the existing branch
that uses "scope_id IS NULL" only for global lookups); reference the function
name RDBConfigStore.GetModelConfig and the variables scope, scopeID, modelName
to locate and implement this check.
In `@framework/configstore/tables/modelconfig.go`:
- Around line 123-126: When persisting ModelConfig, trim and normalize the scope
id: if mc.Scope == ModelConfigScopeGlobal set mc.ScopeID = nil; otherwise set
mc.ScopeID to a pointer to strings.TrimSpace(*mc.ScopeID) and validate emptiness
on the trimmed value (i.e., use the trimmed string for the emptiness check and
assignment) so the stored ScopeID is canonical and will match runtime lookups;
update the logic around mc.ScopeID and the existing emptiness check to use the
trimmed value.
In `@plugins/governance/store.go`:
- Around line 1304-1313: The current code builds entityWiseRateLimits and calls
CheckRateLimit but CheckRateLimit (and its helpers) returns on the first
violated entry instead of aggregating all violations; update CheckRateLimit, and
any helpers it calls, to collect all violation results for every entry in
entityWiseRateLimits (as built from collectModelConfigsFor,
modelConfigEntityKey, LoadRateLimit) and only map the final decision to
DecisionTokenLimited or DecisionRequestLimited when exactly one violation
exists—otherwise return the aggregated rate_limited result (or equivalent
aggregate decision) so that multiple exceeded model-config tiers produce a
combined rate_limited outcome rather than order-dependent narrowing.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 3207-3219: The handler is silently ignoring additional budgets on
provider-level model-configs by only using mc.Budgets[0]; update validation and
handlers to reject multi-budget provider configs: in
modelConfigToProviderGovernance check if mc != nil && mc.Scope ==
configstoreTables.ModelConfigScopeGlobal && mc.ModelName ==
configstoreTables.ModelConfigAllModels && mc.Provider != nil and return false
(or an error upstream) if len(mc.Budgets) > 1; additionally add the same
validation to the PUT/POST code paths that create/update these provider-level
configs so attempts to create multi-budget rows are blocked (rather than
allowing extra budgets to survive PUTs unnoticed).
- Around line 751-770: hydrateVKListGovernance currently calls GetModelConfigs
and filters client-side; instead gather the virtual-key IDs from vks, call the
config store query that fetches only model configs with Scope ==
ModelConfigScopeVirtualKey and ModelName == ModelConfigAllModels for those scope
IDs (e.g. a method like GetModelConfigsByScope/ByScopeIDs or similar), handle
and wrap any error from that call, build the byKey map using
vkModelConfigIndexKey as before, and then call applyVKGovernanceFromModelConfigs
for each vk; update hydrateVKListGovernance to use the scoped query (passing ctx
for cancellation) rather than h.configStore.GetModelConfigs so only relevant
rows are loaded.
- Around line 102-119: NewGovernanceHandler must not register a process-global
resolver bound to the constructor-local configStore; remove the
RegisterScopeNameResolver call from NewGovernanceHandler and instead attach the
virtual-key resolver closure to the GovernanceHandler instance (e.g. add/assign
a field like virtualKeyResolver or scopeNameResolver on GovernanceHandler using
the same closure that calls configStore.GetVirtualKey). Update any code that
previously relied on the global resolver (e.g. resolveModelConfigScopeName
usage) to call the instance resolver on the handler. Ensure no global state is
written from NewGovernanceHandler so multiple handlers keep their own
configStore-bound lookup.
In `@transports/bifrost-http/lib/config_test.go`:
- Around line 969-978: The MockConfigStore.DeletePlugin implementation currently
performs in-memory filtering of m.plugins (embedding business logic in the
mock); change it to a simple no-op that does not implement filtering or
pagination—i.e., remove the loop and any mutation of m.plugins and simply return
nil, keeping the mock behavior minimal consistent with other mocks like
GetVirtualKeysPaginated and DeleteMCPClientConfig (real filtering should be
tested in SQLite integration tests).
- Around line 638-650: The MockConfigStore.DeleteMCPClientConfig currently
implements in-memory filtering logic; revert it to a simple no-op mock by
removing the filtering and mutation and just returning nil (preserve the early
nil check if desired), so the mock does not embed business logic—leave deletion
semantics to SQLite-backed integration tests (see createTestSQLiteConfigStore)
and keep MockConfigStore methods like DeleteMCPClientConfig (and similar methods
such as GetVirtualKeysPaginated) minimal.
In `@transports/config.schema.json`:
- Around line 696-704: Update the JSON Schema in transports/config.schema.json
to require "scope_id" whenever the "scope" property is not "global": add an
if/then/else (or oneOf) conditional around the existing properties so that when
"scope" has const "global" nothing extra is required, otherwise the schema
requires "scope_id"; target the existing "scope" and "scope_id" properties in
the schema and ensure uploads validating against the schema will fail if "scope"
!= "global" and "scope_id" is missing.
In `@ui/app/workspace/model-limits/views/modelLimitSheet.tsx`:
- Around line 118-122: When hydrating the form, preserve legacy single-budget
rows by checking modelConfig.budget in addition to modelConfig.budgets: if
modelConfig.budgets is empty/undefined but modelConfig.budget exists, include
that single budget (converted to the same shape: id, max_limit, reset_duration)
in the budgets array used to populate the form. Update the hydration logic where
budgets are created (the budgets: (modelConfig?.budgets ?? []).map(...) block
and the corresponding logic at lines ~149-155) to merge or fallback to
modelConfig.budget so the table (which still reads config.budget) continues to
show and save legacy rows.
- Around line 39-62: The current zod schema (formSchema) only requires scopeId
when scope === "virtual_key"; update the .refine on formSchema to require
scopeId for every non-"global" scope (i.e., validate that data.scope ===
"global" || !!data.scopeId) so it matches transports/config.schema.json
contract; locate the refine call that references scope and scopeId and change
the predicate and error message accordingly to enforce scope_id whenever scope
is not "global".
In `@ui/app/workspace/model-limits/views/modelLimitsTable.tsx`:
- Around line 326-344: The badge currently appears clickable for all scopes and
uses onClick only; update the JSX so you call getModelLimitScope(config.scope ??
"global")?.buildDeepLink?.(config.scope_id) first and only render the
interactive Badge/TooltipTrigger (with proper keyboard semantics—e.g., a Link or
a button with role and tabIndex) when buildDeepLink returns a target; for scopes
without a deep link render a non-interactive Badge (no onClick, no launch icon)
that still shows config.scope_name; ensure you reference getModelLimitScope,
buildDeepLink, navigate, Badge and TooltipTrigger when making the conditional
change.
In `@ui/lib/utils/labels.ts`:
- Around line 1-20: getScopeLabel calls getModelLimitScope but this file never
triggers the enterprise scope registrations, so consumers can see raw scopes;
fix by ensuring the enterprise model-limit registry is loaded before resolving
labels (e.g., add a side-effect import or call to the enterprise registry
bootstrap from this module so registrations run at module load time), making
sure to do this in the same file that defines getScopeLabel (referencing
getScopeLabel and getModelLimitScope) so callers like
ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsx get the enterprise
labels.
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 3128-3138: When storing the rate limit in
UpdateModelConfigInMemory (the block using clone.RateLimit and
gs.rateLimits.Store), also preserve the last-reset timestamps from the in-memory
entry: when retrieving existingRateLimitValue (cast to
*configstoreTables.TableRateLimit as erl), copy erl.TokenLastReset and
erl.RequestLastReset into clone.RateLimit.TokenLastReset and
clone.RateLimit.RequestLastReset (with appropriate nil/type checks) before
calling gs.rateLimits.Store, so updates with the same rate-limit ID keep the
live reset windows intact.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 2972-3029: The handler is returning 500 for client-side validation
failures originating from validateBudget/validateRateLimit inside
h.configStore.ExecuteTransaction; change it to return 400 by either
pre-validating all req.Budgets and req.RateLimit before calling
ExecuteTransaction or by detecting validation errors returned from within the
transaction and wrapping them in a badRequestError (or the project’s equivalent)
so the outer error handling translates them to a 400 response; update both the
model-config creation block (where validateRateLimit and validateBudget are
called) and the analogous block handling the other mutation (the one referenced
around the later similar code) to consistently wrap/translate validation errors
to badRequestError before calling SendError.
🪄 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: 8ff5d65f-f225-49e9-b2a2-ed986e5a30a5
📒 Files selected for processing (17)
framework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/lib/registrations/modelLimitScopes.tsui/app/workspace/model-limits/views/modelLimitSheet.tsxui/app/workspace/model-limits/views/modelLimitsTable.tsxui/lib/registries/modelLimitScopes.tsxui/lib/store/apis/governanceApi.tsui/lib/utils/labels.ts
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
🛑 Comments failed to post (17)
framework/configstore/migrations.go (3)
3881-3893:
⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftCreate the new unique index concurrently for Postgres
This migration creates a new unique index using
CreateIndexin the normal migration flow. On largegovernance_model_configs, that can hold heavy locks and block writes during upgrade.💡 Suggested approach
+ // For Postgres, run this migration outside a transaction and use CONCURRENTLY. + opts := *migrator.DefaultOptions + if db.Dialector.Name() == "postgres" { + opts.UseTransaction = false + } - m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + m := migrator.New(db, &opts, []*migrator.Migration{{ ... - if !migrator.HasIndex(modelConfig, "idx_model_scope_provider") { - if err := migrator.CreateIndex(modelConfig, "idx_model_scope_provider"); err != nil { - return fmt.Errorf("failed to create idx_model_scope_provider: %w", err) - } - } + if !migrator.HasIndex(modelConfig, "idx_model_scope_provider") { + if tx.Dialector.Name() == "postgres" { + if err := tx.Exec(` + CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_model_scope_provider + ON governance_model_configs (scope, scope_id, model_name, provider) + `).Error; err != nil { + return fmt.Errorf("failed to create idx_model_scope_provider concurrently: %w", err) + } + } else { + if err := migrator.CreateIndex(modelConfig, "idx_model_scope_provider"); err != nil { + return fmt.Errorf("failed to create idx_model_scope_provider: %w", err) + } + } + }As per coding guidelines: "When migrations are added or changed, verify they avoid deadlocks on large tables and create indexes concurrently."
🤖 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 `@framework/configstore/migrations.go` around lines 3881 - 3893, The new unique index creation must use PostgreSQL's CONCURRENTLY option and avoid running inside a transaction to prevent long locks: replace the plain migrator.CreateIndex call for "idx_model_scope_provider" on modelConfig with a dialect check for Postgres (e.g., DB.Dialector.Name() or migrator.DB.Dialector.Name()), and for Postgres execute a raw "CREATE UNIQUE INDEX CONCURRENTLY ..." SQL statement via DB.Exec (construct the index name and columns to match the struct tags) and handle errors; ensure this path is not executed inside a transaction (CONCURRENTLY cannot run in a transaction) and keep the existing migrator.CreateIndex fallback for non-Postgres dialects, then continue to drop the old "idx_model_provider" index as before.
3933-3961:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve provider governance when wildcard model-config already exists
The migration clears provider
budget_id/rate_limit_ideven when it only detects an existing wildcard row and never merges those FK values into that row. That can silently drop provider governance.💡 Suggested fix
- // Idempotency: skip if a global all-models row already exists for this provider. - var existing int64 - if err := tx.Model(&tables.TableModelConfig{}). - Where("scope = ? AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, p.Name). - Count(&existing).Error; err != nil { - return fmt.Errorf("failed to check existing wildcard config for provider %q: %w", p.Name, err) - } - if existing == 0 { + var existing []tables.TableModelConfig + if err := tx.Model(&tables.TableModelConfig{}). + Where("scope = ? AND scope_id IS NULL AND model_name = ? AND provider = ?", + tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, p.Name). + Limit(1). + Find(&existing).Error; err != nil { + return fmt.Errorf("failed to check existing wildcard config for provider %q: %w", p.Name, err) + } + if len(existing) == 0 { providerName := p.Name mc := tables.TableModelConfig{ ID: uuid.NewString(), ModelName: tables.ModelConfigAllModels, Provider: &providerName, Scope: tables.ModelConfigScopeGlobal, BudgetID: p.BudgetID, RateLimitID: p.RateLimitID, CreatedAt: now, UpdatedAt: now, } if err := tx.Create(&mc).Error; err != nil { return fmt.Errorf("failed to create wildcard model config for provider %q: %w", p.Name, err) } + } else { + if err := tx.Model(&tables.TableModelConfig{}). + Where("id = ?", existing[0].ID). + Updates(map[string]any{ + "budget_id": p.BudgetID, + "rate_limit_id": p.RateLimitID, + }).Error; err != nil { + return fmt.Errorf("failed to merge governance into wildcard model config for provider %q: %w", p.Name, err) + } }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.var existing []tables.TableModelConfig if err := tx.Model(&tables.TableModelConfig{}). Where("scope = ? AND scope_id IS NULL AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, p.Name). Limit(1). Find(&existing).Error; err != nil { return fmt.Errorf("failed to check existing wildcard config for provider %q: %w", p.Name, err) } if len(existing) == 0 { providerName := p.Name mc := tables.TableModelConfig{ ID: uuid.NewString(), ModelName: tables.ModelConfigAllModels, Provider: &providerName, Scope: tables.ModelConfigScopeGlobal, BudgetID: p.BudgetID, RateLimitID: p.RateLimitID, CreatedAt: now, UpdatedAt: now, } if err := tx.Create(&mc).Error; err != nil { return fmt.Errorf("failed to create wildcard model config for provider %q: %w", p.Name, err) } } else { if err := tx.Model(&tables.TableModelConfig{}). Where("id = ?", existing[0].ID). Updates(map[string]any{ "budget_id": p.BudgetID, "rate_limit_id": p.RateLimitID, }).Error; err != nil { return fmt.Errorf("failed to merge governance into wildcard model config for provider %q: %w", p.Name, err) } } // Detach governance from the provider (FK rows are reused by the model config above). if err := tx.Model(&tables.TableProvider{}).Where("name = ?", p.Name). Updates(map[string]any{"budget_id": nil, "rate_limit_id": nil}).Error; err != nil { return fmt.Errorf("failed to clear governance FKs for provider %q: %w", p.Name, err) }🤖 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 `@framework/configstore/migrations.go` around lines 3933 - 3961, When a wildcard TableModelConfig (ModelConfigAllModels) already exists for a provider you must merge the provider's governance FKs into that row before clearing them on the TableProvider; modify the branch where existing > 0 to locate the existing TableModelConfig (using the same tx.Model(&tables.TableModelConfig{}).Where("scope = ? AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, p.Name)) and Update that record to set/merge budget_id and rate_limit_id from p.BudgetID and p.RateLimitID (e.g. only set if the target FK is nil or use COALESCE-like logic), then continue to clear the provider FKs on TableProvider — ensure you use the same tx instance and handle/return any errors from the merge update before nulling provider fields.
4208-4229:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winRollback can orphan budgets when provider config is missing
If no provider config is found for
(virtual_key_id, provider), budgets are not restored, but the model-config row is still deleted. That can leave budget ownership dangling.💡 Suggested fix
- if len(pcs) > 0 { - pcID := pcs[0].ID - for _, b := range budgets { - if err := tx.Exec("UPDATE governance_budgets SET provider_config_id = ?, model_config_id = NULL WHERE id = ?", pcID, b.ID).Error; err != nil { - return fmt.Errorf("failed to restore provider-config budget %q: %w", b.ID, err) - } - } - if mc.RateLimitID != nil { - if err := tx.Exec("UPDATE governance_virtual_key_provider_configs SET rate_limit_id = ? WHERE id = ?", *mc.RateLimitID, pcID).Error; err != nil { - return fmt.Errorf("failed to restore provider-config rate limit: %w", err) - } - } - } + if len(pcs) == 0 { + return fmt.Errorf( + "cannot rollback model config %q: missing provider config for virtual_key_id=%q provider=%q", + mc.ID, *mc.ScopeID, *mc.Provider, + ) + } + pcID := pcs[0].ID + for _, b := range budgets { + if err := tx.Exec("UPDATE governance_budgets SET provider_config_id = ?, model_config_id = NULL WHERE id = ?", pcID, b.ID).Error; err != nil { + return fmt.Errorf("failed to restore provider-config budget %q: %w", b.ID, err) + } + } + if mc.RateLimitID != nil { + if err := tx.Exec("UPDATE governance_virtual_key_provider_configs SET rate_limit_id = ? WHERE id = ?", *mc.RateLimitID, pcID).Error; err != nil { + return fmt.Errorf("failed to restore provider-config rate limit: %w", err) + } + }As per coding guidelines: "If a migration cannot be rolled back, explicitly flag it as non-rollbackable."
🤖 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 `@framework/configstore/migrations.go` around lines 4208 - 4229, The code currently deletes the TableModelConfig row (tx.Delete(&tables.TableModelConfig{}, "id = ?", mc.ID)) even when no provider config (pcs) is found, which can orphan budgets; change the logic to avoid deleting the model config unless budgets were successfully restored: after the tx.Where(...) query, if len(pcs) == 0 return a non-rollbackable error (or simply return an error indicating rollback cannot proceed) instead of proceeding to tx.Delete, or move the tx.Delete call inside the if len(pcs) > 0 block so deletion only happens when pcID and budget updates succeeded; reference mc.ScopeID, mc.Provider, pcs, pcID, budgets, mc.RateLimitID, and mc.ID when implementing this check.framework/configstore/rdb.go (2)
3003-3026:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winDelete only the snapshotted VK-scoped model configs, and load their has-many budgets.
This path collects only the legacy
BudgetID, notTableModelConfig.Budgets, so active budgets owned viaModelConfigIDleak on virtual-key delete. It also deletes with a secondWHERE scope = ? AND scope_id = ?, which can remove rows created after the snapshot without deleting their owned budget/rate-limit rows.Suggested fix
- var scopedModelConfigs []tables.TableModelConfig - if err := txDB.WithContext(ctx). - Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id). - Find(&scopedModelConfigs).Error; err != nil { + var scopedModelConfigs []tables.TableModelConfig + if err := txDB.WithContext(ctx). + Preload("Budgets"). + Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id). + Find(&scopedModelConfigs).Error; err != nil { return err } + mcIDs := make([]string, 0, len(scopedModelConfigs)) budgetIDs := make([]string, 0, len(scopedModelConfigs)) rateLimitIDs := make([]string, 0, len(scopedModelConfigs)) for _, mc := range scopedModelConfigs { + mcIDs = append(mcIDs, mc.ID) + for i := range mc.Budgets { + budgetIDs = append(budgetIDs, mc.Budgets[i].ID) + } if mc.BudgetID != nil { budgetIDs = append(budgetIDs, *mc.BudgetID) } if mc.RateLimitID != nil { rateLimitIDs = append(rateLimitIDs, *mc.RateLimitID) } } - if err := txDB.WithContext(ctx). - Where("scope = ? AND scope_id = ?", tables.ModelConfigScopeVirtualKey, id). - Delete(&tables.TableModelConfig{}).Error; err != nil { - return err + if len(mcIDs) > 0 { + if err := txDB.WithContext(ctx). + Where("id IN ?", mcIDs). + Delete(&tables.TableModelConfig{}).Error; err != nil { + return err + } }Based on learnings, budgets and rate limits have a 1:1 ownership with their parent entities and should be deleted together.
🤖 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 `@framework/configstore/rdb.go` around lines 3003 - 3026, The code only collects legacy BudgetID/RateLimitID fields and deletes TableModelConfig rows by scope (which may remove rows created after the snapshot) causing owned budgets/rate-limits to leak; fix by (1) loading the snapshot's has-many relations (Preload("Budgets") and Preload("RateLimits") on the TableModelConfig query) and collect both legacy IDs and the IDs from mc.Budgets and mc.RateLimits, (2) collect the snapshot TableModelConfig IDs and delete budgets and rate-limits by their ModelConfigID OR by collected budget/rate IDs (use TableBudget.ModelConfigID IN ? and TableRateLimit.ModelConfigID IN ? plus any legacy id lists), and (3) delete TableModelConfig rows by their specific IDs (WHERE id IN ?) instead of reusing WHERE scope = ? AND scope_id = ? so you only remove the exact snapshot rows; reference TableModelConfig (ID, BudgetID, RateLimitID, Budgets, ModelConfigID), TableBudget, TableRateLimit, ModelConfigScopeVirtualKey, and txDB.WithContext calls to locate the code.
4344-4351:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winFail fast when a non-global scope lookup is missing
scope_id.This now translates
scope != global && scopeID == nilintoscope_id IS NULL, which violates the new scoped identity contract and turns caller bugs into silent misses.Suggested fix
func (s *RDBConfigStore) GetModelConfig(ctx context.Context, scope string, scopeID *string, modelName string, provider *string) (*tables.TableModelConfig, error) { var modelConfig tables.TableModelConfig + if scope != tables.ModelConfigScopeGlobal { + if scopeID == nil || strings.TrimSpace(*scopeID) == "" { + return nil, fmt.Errorf("scopeID is required for non-global scope %q", scope) + } + } query := s.DB().WithContext(ctx).Where("model_name = ?", modelName).Where("scope = ?", scope) if scopeID != nil { query = query.Where("scope_id = ?", *scopeID) } else { query = query.Where("scope_id IS NULL")As per coding guidelines,
transports/config.schema.jsonrequiresscope_idwhenscope != "global".📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.func (s *RDBConfigStore) GetModelConfig(ctx context.Context, scope string, scopeID *string, modelName string, provider *string) (*tables.TableModelConfig, error) { var modelConfig tables.TableModelConfig if scope != tables.ModelConfigScopeGlobal { if scopeID == nil || strings.TrimSpace(*scopeID) == "" { return nil, fmt.Errorf("scopeID is required for non-global scope %q", scope) } } query := s.DB().WithContext(ctx).Where("model_name = ?", modelName).Where("scope = ?", scope) if scopeID != nil { query = query.Where("scope_id = ?", *scopeID) } else { query = query.Where("scope_id IS NULL")🤖 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 `@framework/configstore/rdb.go` around lines 4344 - 4351, In RDBConfigStore.GetModelConfig, detect the invalid call where scope != "global" and scopeID == nil and immediately return an error instead of translating that into "scope_id IS NULL"; update the start of GetModelConfig to validate scope and scopeID (using the scope and scopeID parameters) and return a clear error when a non-global scope is missing scope_id, otherwise proceed with building the query (keep the existing branch that uses "scope_id IS NULL" only for global lookups); reference the function name RDBConfigStore.GetModelConfig and the variables scope, scopeID, modelName to locate and implement this check.framework/configstore/tables/modelconfig.go (1)
123-126:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winTrim
scope_idbefore persisting it.This only checks
strings.TrimSpace(*mc.ScopeID)for emptiness but keeps the untrimmed value in the row. A non-global config saved with whitespace-padded IDs will pass validation, then never match runtime lookups keyed by the canonical ID, so its scoped budgets/rate limits are silently skipped.As per coding guidelines, validate all untrusted input.
🤖 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 `@framework/configstore/tables/modelconfig.go` around lines 123 - 126, When persisting ModelConfig, trim and normalize the scope id: if mc.Scope == ModelConfigScopeGlobal set mc.ScopeID = nil; otherwise set mc.ScopeID to a pointer to strings.TrimSpace(*mc.ScopeID) and validate emptiness on the trimmed value (i.e., use the trimmed string for the emptiness check and assignment) so the stored ScopeID is canonical and will match runtime lookups; update the logic around mc.ScopeID and the existing emptiness check to use the trimmed value.plugins/governance/store.go (1)
1304-1313:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winAggregate model-config rate-limit violations before deciding.
These paths now evaluate multiple model-config tiers at once, but
CheckRateLimitstill returns on the first violated entry. When more than one matched model-config limit is exceeded, the result becomes order-dependent and can incorrectly narrow totoken_limited/request_limitedinstead of the required aggregatedrate_limited.Based on learnings,
CheckRateLimitand derived helpers must accumulate all violations and only narrow toDecisionTokenLimitedorDecisionRequestLimitedwhen exactly one violation exists.Also applies to: 1349-1358
🤖 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 `@plugins/governance/store.go` around lines 1304 - 1313, The current code builds entityWiseRateLimits and calls CheckRateLimit but CheckRateLimit (and its helpers) returns on the first violated entry instead of aggregating all violations; update CheckRateLimit, and any helpers it calls, to collect all violation results for every entry in entityWiseRateLimits (as built from collectModelConfigsFor, modelConfigEntityKey, LoadRateLimit) and only map the final decision to DecisionTokenLimited or DecisionRequestLimited when exactly one violation exists—otherwise return the aggregated rate_limited result (or equivalent aggregate decision) so that multiple exceeded model-config tiers produce a combined rate_limited outcome rather than order-dependent narrowing.transports/bifrost-http/handlers/governance.go (3)
102-119:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't bind a handler-local store into the global scope resolver registry.
NewGovernanceHandleroverwrites the process-widevirtual_keyresolver with a closure over its ownconfigStore. If another handler/test creates a second instance later, the first handler will start resolving scope names against the second store, so/model-configscan return wrong VK names or blanks. Keep the default VK lookup instance-local instead of registering it from the constructor.Suggested direction
func NewGovernanceHandler(manager GovernanceManager, configStore configstore.ConfigStore) (*GovernanceHandler, error) { if manager == nil { return nil, fmt.Errorf("governance manager is required") } if configStore == nil { return nil, fmt.Errorf("config store is required") } - RegisterScopeNameResolver(configstoreTables.ModelConfigScopeVirtualKey, func(ctx context.Context, scopeID string) (string, bool) { - vk, err := configStore.GetVirtualKey(ctx, scopeID) - if err != nil || vk == nil { - return "", false - } - return vk.Name, true - }) return &GovernanceHandler{ governanceManager: manager, configStore: configStore, }, nil }func (h *GovernanceHandler) resolveModelConfigScopeName(ctx context.Context, mc *configstoreTables.TableModelConfig, cache map[string]string) { if mc == nil || mc.Scope == "" || mc.ScopeID == nil { return } + if mc.Scope == configstoreTables.ModelConfigScopeVirtualKey { + if vk, err := h.configStore.GetVirtualKey(ctx, *mc.ScopeID); err == nil && vk != nil { + mc.ScopeName = vk.Name + } + return + } resolver, ok := lookupScopeNameResolver(mc.Scope) if !ok { return } ... }As per coding guidelines,
**/*.go: Apply standard Go review practices: clear ownership, small interfaces, explicit error handling and wrapping, context propagation and cancellation, bounded goroutines/channels, race-safe shared state, deterministic tests, and table-driven coverage for behavior changes.🤖 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.go` around lines 102 - 119, NewGovernanceHandler must not register a process-global resolver bound to the constructor-local configStore; remove the RegisterScopeNameResolver call from NewGovernanceHandler and instead attach the virtual-key resolver closure to the GovernanceHandler instance (e.g. add/assign a field like virtualKeyResolver or scopeNameResolver on GovernanceHandler using the same closure that calls configStore.GetVirtualKey). Update any code that previously relied on the global resolver (e.g. resolveModelConfigScopeName usage) to call the instance resolver on the handler. Ensure no global state is written from NewGovernanceHandler so multiple handlers keep their own configStore-bound lookup.
751-770:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winAvoid loading every model config to hydrate one VK page.
hydrateVKListGovernancecallsGetModelConfigs(ctx)and filters client-side. After this PR adds more non-VK scopes,/api/governance/virtual-keys?limit=...now scales with total model-config rows in the deployment, not the VKs in the current response. Please use the scope/scope_id query surface from this stack to fetch onlyscope=virtual_key,model_name="*"rows for the VK IDs being returned.As per coding guidelines,
**/*.go: Apply standard Go review practices: clear ownership, small interfaces, explicit error handling and wrapping, context propagation and cancellation, bounded goroutines/channels, race-safe shared state, deterministic tests, and table-driven coverage for behavior changes.🤖 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.go` around lines 751 - 770, hydrateVKListGovernance currently calls GetModelConfigs and filters client-side; instead gather the virtual-key IDs from vks, call the config store query that fetches only model configs with Scope == ModelConfigScopeVirtualKey and ModelName == ModelConfigAllModels for those scope IDs (e.g. a method like GetModelConfigsByScope/ByScopeIDs or similar), handle and wrap any error from that call, build the byKey map using vkModelConfigIndexKey as before, and then call applyVKGovernanceFromModelConfigs for each vk; update hydrateVKListGovernance to use the scoped query (passing ctx for cancellation) rather than h.configStore.GetModelConfigs so only relevant rows are loaded.
3207-3219:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winProvider governance can't safely ignore extra budgets.
These handlers now back
/api/governance/providerswith provider-level model-config rows, but they only read/updatemc.Budgets[0]. Generic model-config CRUD in this same file now allows multiple budgets on that row shape, so any additional budgets become invisible in/providersand survive PUTs unmanaged. Please rejectlen(mc.Budgets) > 1here or block multi-budget creation forscope=global, model_name="*", provider!=nilconfigs.Also applies to: 3301-3431
🤖 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.go` around lines 3207 - 3219, The handler is silently ignoring additional budgets on provider-level model-configs by only using mc.Budgets[0]; update validation and handlers to reject multi-budget provider configs: in modelConfigToProviderGovernance check if mc != nil && mc.Scope == configstoreTables.ModelConfigScopeGlobal && mc.ModelName == configstoreTables.ModelConfigAllModels && mc.Provider != nil and return false (or an error upstream) if len(mc.Budgets) > 1; additionally add the same validation to the PUT/POST code paths that create/update these provider-level configs so attempts to create multi-budget rows are blocked (rather than allowing extra budgets to survive PUTs unnoticed).transports/bifrost-http/lib/config_test.go (2)
638-650: 🧹 Nitpick | 🔵 Trivial | 💤 Low value
Mock filtering logic may violate simplicity principle.
The updated
DeleteMCPClientConfignow performs in-memory filtering to remove the client by ID. Based on learnings, MockConfigStore methods should remain simple, returning zero/nil values without embedding business logic like filtering. Consider whether this filtering is necessary for mock behavior or if SQLite-backed integration tests should validate deletion semantics instead.Based on learnings: "In tests under transports/bifrost-http/lib/config_test.go, keep MockConfigStore methods (e.g., GetVirtualKeysPaginated) as simple, returning zero/nil values. Do not embed business logic (filtering/pagination) in the mock. Cover behavior with SQLite-backed integration tests using createTestSQLiteConfigStore instead to validate end-to-end behavior."
🤖 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/lib/config_test.go` around lines 638 - 650, The MockConfigStore.DeleteMCPClientConfig currently implements in-memory filtering logic; revert it to a simple no-op mock by removing the filtering and mutation and just returning nil (preserve the early nil check if desired), so the mock does not embed business logic—leave deletion semantics to SQLite-backed integration tests (see createTestSQLiteConfigStore) and keep MockConfigStore methods like DeleteMCPClientConfig (and similar methods such as GetVirtualKeysPaginated) minimal.
969-978: 🧹 Nitpick | 🔵 Trivial | 💤 Low value
Mock filtering logic may violate simplicity principle.
Similar to
DeleteMCPClientConfig, this method now filters the plugins slice in-memory. Based on learnings, consider whether this filtering logic belongs in the mock or should be validated through SQLite integration tests instead.Based on learnings: "In tests under transports/bifrost-http/lib/config_test.go, keep MockConfigStore methods (e.g., GetVirtualKeysPaginated) as simple, returning zero/nil values. Do not embed business logic (filtering/pagination) in the mock."
🤖 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/lib/config_test.go` around lines 969 - 978, The MockConfigStore.DeletePlugin implementation currently performs in-memory filtering of m.plugins (embedding business logic in the mock); change it to a simple no-op that does not implement filtering or pagination—i.e., remove the loop and any mutation of m.plugins and simply return nil, keeping the mock behavior minimal consistent with other mocks like GetVirtualKeysPaginated and DeleteMCPClientConfig (real filtering should be tested in SQLite integration tests).transports/config.schema.json (1)
696-704:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winEnforce
scope_idfor non-global model-config scopes.This only documents the requirement; it does not validate it. Because config uploads are validated against this schema,
{ "scope": "virtual_key" }will still pass withoutscope_id, which breaks the new scoped model-config contract.Suggested schema fix
"model_configs": { "type": "array", "description": "Per-model rate limit and budget configurations", "items": { "type": "object", "properties": { "id": { "type": "string", "description": "Model config ID" }, "model_name": { "type": "string", "description": "Model name to apply the configuration to" }, "provider": { "type": "string", "description": "Optional provider name to scope this config" }, "scope": { "type": "string", "description": "Scope where this config applies: \"global\" (default) or \"virtual_key\"", "default": "global" }, "scope_id": { "type": "string", "description": "Target entity ID for non-global scopes (e.g. virtual key ID). Required when scope != \"global\"" }, "budget_id": { "type": "string", "description": "Budget ID to associate with this model" }, "rate_limit_id": { "type": "string", "description": "Rate limit ID to associate with this model" } }, "required": ["id", "model_name"], + "if": { + "properties": { + "scope": { + "not": { "const": "global" } + } + }, + "required": ["scope"] + }, + "then": { + "required": ["scope_id"], + "properties": { + "scope_id": { + "type": "string", + "minLength": 1 + } + } + }, "additionalProperties": false } },As per coding guidelines,
transports/config.schema.jsonis the source of truth for config fields, and the referenced schema requirement saysscope_idis required whenscope != "global".🤖 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/config.schema.json` around lines 696 - 704, Update the JSON Schema in transports/config.schema.json to require "scope_id" whenever the "scope" property is not "global": add an if/then/else (or oneOf) conditional around the existing properties so that when "scope" has const "global" nothing extra is required, otherwise the schema requires "scope_id"; target the existing "scope" and "scope_id" properties in the schema and ensure uploads validating against the schema will fail if "scope" != "global" and "scope_id" is missing.ui/app/workspace/model-limits/views/modelLimitSheet.tsx (2)
39-62:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate
scopeIdfor every non-global scope, not justvirtual_key.This refinement hardcodes one scope, so any downstream scope added through the registry can submit with an empty
scopeIdand only fail after the API call. That breaks the new extensibility path and diverges from the config contract.As per coding guidelines, `transports/config.schema.json` is the source of truth and requires `scope_id` whenever `scope != "global"`.Suggested fix
const formSchema = z .object({ modelName: z.string().min(1, "Model name is required"), provider: z.string().optional(), scope: z.string().optional(), scopeId: z.string().optional(), @@ - .refine((data) => data.scope !== "virtual_key" || !!data.scopeId, { - message: "Virtual key is required for the Virtual Key scope", + .refine((data) => { + const scope = data.scope || "global"; + return scope === "global" || !!data.scopeId; + }, { + message: "Scope target is required for non-global scopes", path: ["scopeId"], });🤖 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/model-limits/views/modelLimitSheet.tsx` around lines 39 - 62, The current zod schema (formSchema) only requires scopeId when scope === "virtual_key"; update the .refine on formSchema to require scopeId for every non-"global" scope (i.e., validate that data.scope === "global" || !!data.scopeId) so it matches transports/config.schema.json contract; locate the refine call that references scope and scopeId and change the predicate and error message accordingly to enforce scope_id whenever scope is not "global".
118-122:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve legacy single-budget rows when hydrating the form.
The sheet now reads only
modelConfig.budgets, but the table still guards forconfig.budget. If a row arrives with only the legacy field populated, editing it opens with no budget lines and saving will drop the existing budget.Suggested fix
+ const initialBudgets = (modelConfig?.budgets ?? (modelConfig?.budget ? [modelConfig.budget] : [])).map((b) => ({ + id: b.id, + max_limit: b.max_limit, + reset_duration: b.reset_duration, + })); + const form = useForm<FormData>({ @@ - budgets: (modelConfig?.budgets ?? []).map((b) => ({ - id: b.id, - max_limit: b.max_limit, - reset_duration: b.reset_duration, - })), + budgets: initialBudgets, @@ - budgets: (modelConfig.budgets ?? []).map((b) => ({ - id: b.id, - max_limit: b.max_limit, - reset_duration: b.reset_duration, - })), + budgets: initialBudgets,Also applies to: 149-155
🤖 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/model-limits/views/modelLimitSheet.tsx` around lines 118 - 122, When hydrating the form, preserve legacy single-budget rows by checking modelConfig.budget in addition to modelConfig.budgets: if modelConfig.budgets is empty/undefined but modelConfig.budget exists, include that single budget (converted to the same shape: id, max_limit, reset_duration) in the budgets array used to populate the form. Update the hydration logic where budgets are created (the budgets: (modelConfig?.budgets ?? []).map(...) block and the corresponding logic at lines ~149-155) to merge or fallback to modelConfig.budget so the table (which still reads config.budget) continues to show and save legacy rows.ui/app/workspace/model-limits/views/modelLimitsTable.tsx (1)
326-344:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winOnly render an interactive scope badge when a deep link exists, and use button/link semantics.
buildDeepLinkis optional in the registry, but this badge always looks clickable and shows the launch icon. For scopes without a registered deep link it becomes a no-op CTA, and even when a target exists theBadgeis not keyboard-focusable because it's wired withonClickonly.Suggested fix
- {config.scope !== "global" && config.scope_id && config.scope_name ? ( + {config.scope !== "global" && config.scope_id && config.scope_name ? ( <TooltipProvider> <Tooltip> <TooltipTrigger asChild> - <Badge - variant="secondary" - className="flex max-w-[160px] cursor-pointer items-center gap-1 hover:opacity-80" - data-testid={`model-limit-scope-target-${config.scope_id}`} - onClick={() => { - if (!config.scope_id) return; - const target = getModelLimitScope(config.scope ?? "global")?.buildDeepLink?.(config.scope_id); - if (target) navigate(target as never); - }} - > - <span className="truncate">{config.scope_name}</span> - <ArrowUpRight className="h-3 w-3 shrink-0" /> - </Badge> + {(() => { + const target = getModelLimitScope(config.scope ?? "global")?.buildDeepLink?.(config.scope_id); + const content = <span className="truncate">{config.scope_name}</span>; + return target ? ( + <button + type="button" + data-testid={`model-limit-scope-target-${config.scope_id}`} + onClick={() => navigate(target as never)} + > + <Badge variant="secondary" className="flex max-w-[160px] items-center gap-1 hover:opacity-80"> + {content} + <ArrowUpRight className="h-3 w-3 shrink-0" /> + </Badge> + </button> + ) : ( + <Badge variant="secondary" className="max-w-[160px]"> + {content} + </Badge> + ); + })()} </TooltipTrigger> <TooltipContent className="max-w-[320px] break-all">{config.scope_name}</TooltipContent> </Tooltip> </TooltipProvider>📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.{config.scope !== "global" && config.scope_id && config.scope_name ? ( <TooltipProvider> <Tooltip> <TooltipTrigger asChild> {(() => { const target = getModelLimitScope(config.scope ?? "global")?.buildDeepLink?.(config.scope_id); const content = <span className="truncate">{config.scope_name}</span>; return target ? ( <button type="button" data-testid={`model-limit-scope-target-${config.scope_id}`} onClick={() => navigate(target as never)} > <Badge variant="secondary" className="flex max-w-[160px] items-center gap-1 hover:opacity-80"> {content} <ArrowUpRight className="h-3 w-3 shrink-0" /> </Badge> </button> ) : ( <Badge variant="secondary" className="max-w-[160px]"> {content} </Badge> ); })()} </TooltipTrigger> <TooltipContent className="max-w-[320px] break-all">{config.scope_name}</TooltipContent> </Tooltip> </TooltipProvider> ) : ( /* rest of the code */ )}🤖 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/model-limits/views/modelLimitsTable.tsx` around lines 326 - 344, The badge currently appears clickable for all scopes and uses onClick only; update the JSX so you call getModelLimitScope(config.scope ?? "global")?.buildDeepLink?.(config.scope_id) first and only render the interactive Badge/TooltipTrigger (with proper keyboard semantics—e.g., a Link or a button with role and tabIndex) when buildDeepLink returns a target; for scopes without a deep link render a non-interactive Badge (no onClick, no launch icon) that still shows config.scope_name; ensure you reference getModelLimitScope, buildDeepLink, navigate, Badge and TooltipTrigger when making the conditional change.ui/lib/utils/labels.ts (1)
1-20:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winLoad enterprise scope registrations before resolving labels.
getScopeLabel()now depends on the registry, but this module never imports the enterprise registration side effect. That means callers outside the model-limit pages can still see raw scope values like"user"because only the OSS defaults are registered here.ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsxalready consumes this helper without loading the registry bootstrap first.🤖 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/utils/labels.ts` around lines 1 - 20, getScopeLabel calls getModelLimitScope but this file never triggers the enterprise scope registrations, so consumers can see raw scopes; fix by ensuring the enterprise model-limit registry is loaded before resolving labels (e.g., add a side-effect import or call to the enterprise registry bootstrap from this module so registrations run at module load time), making sure to do this in the same file that defines getScopeLabel (referencing getScopeLabel and getModelLimitScope) so callers like ui/app/workspace/routing-rules/views/routingRuleInfoSheet.tsx get the enterprise labels.
akshaydeo
force-pushed
the
06-01-refactor_makes_scope-level_check_methods_extensible
branch
from
June 2, 2026 14:03
4816c78 to
6a0b77c
Compare
akshaydeo
deleted the
06-01-refactor_makes_scope-level_check_methods_extensible
branch
June 2, 2026 14:05
akshaydeo
pushed a commit
that referenced
this pull request
Jun 2, 2026
## Summary This PR generalizes the virtual-key-scoped model governance system into a scope-agnostic framework, enabling any registered scope (e.g. `user`) to carry per-scope model rate limits and budgets — not just virtual keys. It also introduces a runtime scope registry on both the backend and frontend so downstream (enterprise) builds can extend the system without modifying OSS code. ## Changes - **`CheckVirtualKeyScopedModelBudget` / `CheckVirtualKeyScopedModelRateLimit`** and their `UpdateVirtualKeyScoped*` counterparts are replaced by `CheckScopedModelBudget`, `CheckScopedModelRateLimit`, `UpdateScopedModelBudgetUsageInMemory`, and `UpdateScopedModelRateLimitUsageInMemory`. These accept a `(scope, scopeID)` pair instead of a `*TableVirtualKey`, making them scope-agnostic. An empty scope or scopeID is a no-op. - **`ModelConfigScopeUser`** constant added to `tables/modelconfig.go`, along with a `RegisterModelConfigScope` function and a `sync.RWMutex`-guarded registry so downstream builds can add scopes at startup without forking the OSS validation logic. - **`EvaluateUserRequest`** in `resolver.go` and `UpdateUsage` in `tracker.go` now invoke the scoped model check/update paths for the `user` scope, mirroring the existing VK-scoped block. - **`DeleteProvider`** in `rdb.go` is refactored to batch-delete budgets and rate limits with `IN` clauses instead of one-by-one, and the model config row is deleted before its owned resources to avoid constraint issues. - **`DeleteVirtualKey`** removes the loop that deleted budgets via `ModelConfigID`; only the `BudgetID` foreign key path is retained. - Internal naming throughout `migrations.go`, `governance.go`, and `store.go` drops the "wildcard" terminology (`ensureVKWildcardModelConfig` → `ensureVKModelConfig`, `vkWildcardDesired` → `vkModelConfigDesired`, `upsertVKWildcard` → `reconcileVKModelConfig`, etc.) to reflect that these configs are not exclusively wildcard rows. - **`RegisterScopeNameResolver`** added to `handlers/governance.go` with a package-level `sync.RWMutex`-guarded map. `resolveModelConfigScopeName` now dispatches to the registered resolver for any scope rather than hard-coding the VK lookup. The VK resolver is wired automatically in `NewGovernanceHandler`. - **UI scope registry** (`ui/lib/registries/modelLimitScopes.tsx`) replaces the static `MODEL_LIMIT_SCOPES` constant. Each entry can declare a `PickerComponent` and a `buildDeepLink` function. The OSS build registers `global` and `virtual_key` at module load; enterprise builds extend the registry via the `@enterprise` alias side-effect import. - The Model Limit sheet's VK picker is replaced by a registry-driven `PickerComponent` render, and the deep-link navigation in the table is driven by `buildDeepLink`, so adding a new scope (e.g. `user`) requires no changes to OSS sheet or table code. - `invalidatesTags` for model config mutations now includes `"Users"` and `"UserGovernance"` (no-op in OSS; picked up by enterprise tag wiring). ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [x] UI (React) - [ ] Docs ## How to test ```sh # Core/Transports go test ./framework/configstore/... ./plugins/governance/... ./transports/bifrost-http/... # UI cd ui pnpm i pnpm build ``` Existing governance tests in `modelprovidergovernance_test.go` have been updated to call the new `CheckScopedModel*` / `UpdateScopedModel*` signatures and continue to cover the VK-scoped budget and rate-limit paths. ## Breaking changes - [x] Yes - [ ] No The `GovernanceStore` interface methods `CheckVirtualKeyScopedModelBudget`, `CheckVirtualKeyScopedModelRateLimit`, `UpdateVirtualKeyScopedModelBudgetUsageInMemory`, and `UpdateVirtualKeyScopedModelRateLimitUsageInMemory` are removed and replaced by their scope-agnostic equivalents. Any downstream implementation of `GovernanceStore` must be updated to implement `CheckScopedModelBudget`, `CheckScopedModelRateLimit`, `UpdateScopedModelBudgetUsageInMemory`, and `UpdateScopedModelRateLimitUsageInMemory`. ## Security considerations The new `RegisterModelConfigScope` and `RegisterScopeNameResolver` functions are intended to be called once at process startup before serving requests. No user-supplied input reaches the registry directly. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added user-scoped model budget and rate-limit enforcement. * Added calendar-aligned reset support for model limits. * Implemented dynamic model-limit scope registry with scope-specific pickers and deep-linking. * **Bug Fixes** * Improved bulk cleanup for provider- and virtual-key-scoped model configs. * Ensured consistent calendar-alignment when creating scoped model configs. * **Refactor** * Migrated virtual-key governance to model-config backed storage. * Generalized scoped model checks/usage to a (scope, scope_id) model for reuse. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
18 tasks
akshaydeo
pushed a commit
that referenced
this pull request
Jun 4, 2026
## Summary This PR generalizes the virtual-key-scoped model governance system into a scope-agnostic framework, enabling any registered scope (e.g. `user`) to carry per-scope model rate limits and budgets — not just virtual keys. It also introduces a runtime scope registry on both the backend and frontend so downstream (enterprise) builds can extend the system without modifying OSS code. ## Changes - **`CheckVirtualKeyScopedModelBudget` / `CheckVirtualKeyScopedModelRateLimit`** and their `UpdateVirtualKeyScoped*` counterparts are replaced by `CheckScopedModelBudget`, `CheckScopedModelRateLimit`, `UpdateScopedModelBudgetUsageInMemory`, and `UpdateScopedModelRateLimitUsageInMemory`. These accept a `(scope, scopeID)` pair instead of a `*TableVirtualKey`, making them scope-agnostic. An empty scope or scopeID is a no-op. - **`ModelConfigScopeUser`** constant added to `tables/modelconfig.go`, along with a `RegisterModelConfigScope` function and a `sync.RWMutex`-guarded registry so downstream builds can add scopes at startup without forking the OSS validation logic. - **`EvaluateUserRequest`** in `resolver.go` and `UpdateUsage` in `tracker.go` now invoke the scoped model check/update paths for the `user` scope, mirroring the existing VK-scoped block. - **`DeleteProvider`** in `rdb.go` is refactored to batch-delete budgets and rate limits with `IN` clauses instead of one-by-one, and the model config row is deleted before its owned resources to avoid constraint issues. - **`DeleteVirtualKey`** removes the loop that deleted budgets via `ModelConfigID`; only the `BudgetID` foreign key path is retained. - Internal naming throughout `migrations.go`, `governance.go`, and `store.go` drops the "wildcard" terminology (`ensureVKWildcardModelConfig` → `ensureVKModelConfig`, `vkWildcardDesired` → `vkModelConfigDesired`, `upsertVKWildcard` → `reconcileVKModelConfig`, etc.) to reflect that these configs are not exclusively wildcard rows. - **`RegisterScopeNameResolver`** added to `handlers/governance.go` with a package-level `sync.RWMutex`-guarded map. `resolveModelConfigScopeName` now dispatches to the registered resolver for any scope rather than hard-coding the VK lookup. The VK resolver is wired automatically in `NewGovernanceHandler`. - **UI scope registry** (`ui/lib/registries/modelLimitScopes.tsx`) replaces the static `MODEL_LIMIT_SCOPES` constant. Each entry can declare a `PickerComponent` and a `buildDeepLink` function. The OSS build registers `global` and `virtual_key` at module load; enterprise builds extend the registry via the `@enterprise` alias side-effect import. - The Model Limit sheet's VK picker is replaced by a registry-driven `PickerComponent` render, and the deep-link navigation in the table is driven by `buildDeepLink`, so adding a new scope (e.g. `user`) requires no changes to OSS sheet or table code. - `invalidatesTags` for model config mutations now includes `"Users"` and `"UserGovernance"` (no-op in OSS; picked up by enterprise tag wiring). ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [x] UI (React) - [ ] Docs ## How to test ```sh # Core/Transports go test ./framework/configstore/... ./plugins/governance/... ./transports/bifrost-http/... # UI cd ui pnpm i pnpm build ``` Existing governance tests in `modelprovidergovernance_test.go` have been updated to call the new `CheckScopedModel*` / `UpdateScopedModel*` signatures and continue to cover the VK-scoped budget and rate-limit paths. ## Breaking changes - [x] Yes - [ ] No The `GovernanceStore` interface methods `CheckVirtualKeyScopedModelBudget`, `CheckVirtualKeyScopedModelRateLimit`, `UpdateVirtualKeyScopedModelBudgetUsageInMemory`, and `UpdateVirtualKeyScopedModelRateLimitUsageInMemory` are removed and replaced by their scope-agnostic equivalents. Any downstream implementation of `GovernanceStore` must be updated to implement `CheckScopedModelBudget`, `CheckScopedModelRateLimit`, `UpdateScopedModelBudgetUsageInMemory`, and `UpdateScopedModelRateLimitUsageInMemory`. ## Security considerations The new `RegisterModelConfigScope` and `RegisterScopeNameResolver` functions are intended to be called once at process startup before serving requests. No user-supplied input reaches the registry directly. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added user-scoped model budget and rate-limit enforcement. * Added calendar-aligned reset support for model limits. * Implemented dynamic model-limit scope registry with scope-specific pickers and deep-linking. * **Bug Fixes** * Improved bulk cleanup for provider- and virtual-key-scoped model configs. * Ensured consistent calendar-alignment when creating scoped model configs. * **Refactor** * Migrated virtual-key governance to model-config backed storage. * Generalized scoped model checks/usage to a (scope, scope_id) model for reuse. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This was referenced Jun 5, 2026
akshaydeo
pushed a commit
that referenced
this pull request
Jun 7, 2026
## Summary This PR generalizes the virtual-key-scoped model governance system into a scope-agnostic framework, enabling any registered scope (e.g. `user`) to carry per-scope model rate limits and budgets — not just virtual keys. It also introduces a runtime scope registry on both the backend and frontend so downstream (enterprise) builds can extend the system without modifying OSS code. ## Changes - **`CheckVirtualKeyScopedModelBudget` / `CheckVirtualKeyScopedModelRateLimit`** and their `UpdateVirtualKeyScoped*` counterparts are replaced by `CheckScopedModelBudget`, `CheckScopedModelRateLimit`, `UpdateScopedModelBudgetUsageInMemory`, and `UpdateScopedModelRateLimitUsageInMemory`. These accept a `(scope, scopeID)` pair instead of a `*TableVirtualKey`, making them scope-agnostic. An empty scope or scopeID is a no-op. - **`ModelConfigScopeUser`** constant added to `tables/modelconfig.go`, along with a `RegisterModelConfigScope` function and a `sync.RWMutex`-guarded registry so downstream builds can add scopes at startup without forking the OSS validation logic. - **`EvaluateUserRequest`** in `resolver.go` and `UpdateUsage` in `tracker.go` now invoke the scoped model check/update paths for the `user` scope, mirroring the existing VK-scoped block. - **`DeleteProvider`** in `rdb.go` is refactored to batch-delete budgets and rate limits with `IN` clauses instead of one-by-one, and the model config row is deleted before its owned resources to avoid constraint issues. - **`DeleteVirtualKey`** removes the loop that deleted budgets via `ModelConfigID`; only the `BudgetID` foreign key path is retained. - Internal naming throughout `migrations.go`, `governance.go`, and `store.go` drops the "wildcard" terminology (`ensureVKWildcardModelConfig` → `ensureVKModelConfig`, `vkWildcardDesired` → `vkModelConfigDesired`, `upsertVKWildcard` → `reconcileVKModelConfig`, etc.) to reflect that these configs are not exclusively wildcard rows. - **`RegisterScopeNameResolver`** added to `handlers/governance.go` with a package-level `sync.RWMutex`-guarded map. `resolveModelConfigScopeName` now dispatches to the registered resolver for any scope rather than hard-coding the VK lookup. The VK resolver is wired automatically in `NewGovernanceHandler`. - **UI scope registry** (`ui/lib/registries/modelLimitScopes.tsx`) replaces the static `MODEL_LIMIT_SCOPES` constant. Each entry can declare a `PickerComponent` and a `buildDeepLink` function. The OSS build registers `global` and `virtual_key` at module load; enterprise builds extend the registry via the `@enterprise` alias side-effect import. - The Model Limit sheet's VK picker is replaced by a registry-driven `PickerComponent` render, and the deep-link navigation in the table is driven by `buildDeepLink`, so adding a new scope (e.g. `user`) requires no changes to OSS sheet or table code. - `invalidatesTags` for model config mutations now includes `"Users"` and `"UserGovernance"` (no-op in OSS; picked up by enterprise tag wiring). ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [x] UI (React) - [ ] Docs ## How to test ```sh # Core/Transports go test ./framework/configstore/... ./plugins/governance/... ./transports/bifrost-http/... # UI cd ui pnpm i pnpm build ``` Existing governance tests in `modelprovidergovernance_test.go` have been updated to call the new `CheckScopedModel*` / `UpdateScopedModel*` signatures and continue to cover the VK-scoped budget and rate-limit paths. ## Breaking changes - [x] Yes - [ ] No The `GovernanceStore` interface methods `CheckVirtualKeyScopedModelBudget`, `CheckVirtualKeyScopedModelRateLimit`, `UpdateVirtualKeyScopedModelBudgetUsageInMemory`, and `UpdateVirtualKeyScopedModelRateLimitUsageInMemory` are removed and replaced by their scope-agnostic equivalents. Any downstream implementation of `GovernanceStore` must be updated to implement `CheckScopedModelBudget`, `CheckScopedModelRateLimit`, `UpdateScopedModelBudgetUsageInMemory`, and `UpdateScopedModelRateLimitUsageInMemory`. ## Security considerations The new `RegisterModelConfigScope` and `RegisterScopeNameResolver` functions are intended to be called once at process startup before serving requests. No user-supplied input reaches the registry directly. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added user-scoped model budget and rate-limit enforcement. * Added calendar-aligned reset support for model limits. * Implemented dynamic model-limit scope registry with scope-specific pickers and deep-linking. * **Bug Fixes** * Improved bulk cleanup for provider- and virtual-key-scoped model configs. * Ensured consistent calendar-alignment when creating scoped model configs. * **Refactor** * Migrated virtual-key governance to model-config backed storage. * Generalized scoped model checks/usage to a (scope, scope_id) model for reuse. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR generalizes the virtual-key-scoped model governance system into a scope-agnostic framework, enabling any registered scope (e.g.
user) to carry per-scope model rate limits and budgets — not just virtual keys. It also introduces a runtime scope registry on both the backend and frontend so downstream (enterprise) builds can extend the system without modifying OSS code.Changes
CheckVirtualKeyScopedModelBudget/CheckVirtualKeyScopedModelRateLimitand theirUpdateVirtualKeyScoped*counterparts are replaced byCheckScopedModelBudget,CheckScopedModelRateLimit,UpdateScopedModelBudgetUsageInMemory, andUpdateScopedModelRateLimitUsageInMemory. These accept a(scope, scopeID)pair instead of a*TableVirtualKey, making them scope-agnostic. An empty scope or scopeID is a no-op.ModelConfigScopeUserconstant added totables/modelconfig.go, along with aRegisterModelConfigScopefunction and async.RWMutex-guarded registry so downstream builds can add scopes at startup without forking the OSS validation logic.EvaluateUserRequestinresolver.goandUpdateUsageintracker.gonow invoke the scoped model check/update paths for theuserscope, mirroring the existing VK-scoped block.DeleteProviderinrdb.gois refactored to batch-delete budgets and rate limits withINclauses instead of one-by-one, and the model config row is deleted before its owned resources to avoid constraint issues.DeleteVirtualKeyremoves the loop that deleted budgets viaModelConfigID; only theBudgetIDforeign key path is retained.migrations.go,governance.go, andstore.godrops the "wildcard" terminology (ensureVKWildcardModelConfig→ensureVKModelConfig,vkWildcardDesired→vkModelConfigDesired,upsertVKWildcard→reconcileVKModelConfig, etc.) to reflect that these configs are not exclusively wildcard rows.RegisterScopeNameResolveradded tohandlers/governance.gowith a package-levelsync.RWMutex-guarded map.resolveModelConfigScopeNamenow dispatches to the registered resolver for any scope rather than hard-coding the VK lookup. The VK resolver is wired automatically inNewGovernanceHandler.ui/lib/registries/modelLimitScopes.tsx) replaces the staticMODEL_LIMIT_SCOPESconstant. Each entry can declare aPickerComponentand abuildDeepLinkfunction. The OSS build registersglobalandvirtual_keyat module load; enterprise builds extend the registry via the@enterprisealias side-effect import.PickerComponentrender, and the deep-link navigation in the table is driven bybuildDeepLink, so adding a new scope (e.g.user) requires no changes to OSS sheet or table code.invalidatesTagsfor model config mutations now includes"Users"and"UserGovernance"(no-op in OSS; picked up by enterprise tag wiring).Type of change
Affected areas
How to test
Existing governance tests in
modelprovidergovernance_test.gohave been updated to call the newCheckScopedModel*/UpdateScopedModel*signatures and continue to cover the VK-scoped budget and rate-limit paths.Breaking changes
The
GovernanceStoreinterface methodsCheckVirtualKeyScopedModelBudget,CheckVirtualKeyScopedModelRateLimit,UpdateVirtualKeyScopedModelBudgetUsageInMemory, andUpdateVirtualKeyScopedModelRateLimitUsageInMemoryare removed and replaced by their scope-agnostic equivalents. Any downstream implementation ofGovernanceStoremust be updated to implementCheckScopedModelBudget,CheckScopedModelRateLimit,UpdateScopedModelBudgetUsageInMemory, andUpdateScopedModelRateLimitUsageInMemory.Security considerations
The new
RegisterModelConfigScopeandRegisterScopeNameResolverfunctions are intended to be called once at process startup before serving requests. No user-supplied input reaches the registry directly.Checklist
docs/contributing/README.mdand followed the guidelinesSummary by CodeRabbit
New Features
Bug Fixes
Refactor