feat: add support for per-model and per-provider level budgeting and rate limiting in governance plugin #1120
Conversation
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
|
Warning Rate limit exceeded@danpiths has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 11 minutes and 34 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughAdds provider- and model-level governance: new ORM table for model configs, provider governance fields, migrations, configstore CRUD and RDB persistence, in-memory provider/model caching and checks, resolver split for model/provider vs virtual-key evaluation, and reordered hook/tracker usage updates. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant GovernancePlugin
participant BudgetResolver
participant LocalStore
participant Tracker
rect rgb(240,248,255)
Client->>GovernancePlugin: PreHook(req: provider, model, optional VK)
GovernancePlugin->>BudgetResolver: EvaluateModelAndProviderRequest(provider, model, requestID)
BudgetResolver->>LocalStore: CheckProviderRateLimit/Budget (if provider)
LocalStore-->>BudgetResolver: decision
BudgetResolver->>LocalStore: CheckModelRateLimit/Budget (if model)
LocalStore-->>BudgetResolver: decision
BudgetResolver-->>GovernancePlugin: decision
end
alt virtual key present
rect rgb(245,255,240)
GovernancePlugin->>BudgetResolver: EvaluateVirtualKeyRequest(VK, provider, model, requestID)
BudgetResolver->>LocalStore: VK-level checks (rate-limit, budget)
LocalStore-->>BudgetResolver: decision
BudgetResolver-->>GovernancePlugin: final decision (may override)
end
end
GovernancePlugin-->>Client: PreHook response
Client->>GovernancePlugin: PostHook(response with model)
GovernancePlugin->>Tracker: Schedule updates (provider+model first, then VK if present)
Tracker->>LocalStore: UpdateProviderAndModelRateLimitUsageInMemory(...)
LocalStore-->>Tracker: ack
Tracker->>LocalStore: UpdateProviderAndModelBudgetUsageInMemory(...)
LocalStore-->>Tracker: ack
alt VK present
Tracker->>LocalStore: UpdateVirtualKeyRateLimitUsageInMemory/UpdateVirtualKeyBudgetUsageInMemory
LocalStore-->>Tracker: ack
end
Tracker-->>GovernancePlugin: done
GovernancePlugin-->>Client: PostHook complete
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
Comment |
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
bc0a584 to
0ed2f53
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
d145550 to
543a580
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
plugins/governance/tracker.go (1)
98-110: Consider guarding model-level updates when model name is empty.Unlike provider-level updates which check
if update.Provider != "", model-level updates don't verify thatupdate.Modelis non-empty. If a request could have an empty model name, this would cause unnecessary store calls.- // 3. Update model-level rate limit usage (after provider-level) - // This applies even when virtual keys are disabled or not present - if err := t.store.UpdateModelRateLimitUsage(ctx, update.Model, &update.Provider, update.TokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { - t.logger.Error("failed to update model-level rate limit usage for model %s: %v", update.Model, err) - } - - // 4. Update model-level budget usage (after provider-level) - // This applies even when virtual keys are disabled or not present - if shouldUpdateBudget && update.Cost > 0 { - if err := t.store.UpdateModelBudgetUsage(ctx, update.Model, &update.Provider, update.Cost); err != nil { - t.logger.Error("failed to update model-level budget usage for model %s: %v", update.Model, err) - } - } + // 3. Update model-level rate limit usage (after provider-level) + // This applies even when virtual keys are disabled or not present + if update.Model != "" { + if err := t.store.UpdateModelRateLimitUsage(ctx, update.Model, &update.Provider, update.TokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { + t.logger.Error("failed to update model-level rate limit usage for model %s: %v", update.Model, err) + } + + // 4. Update model-level budget usage (after provider-level) + // This applies even when virtual keys are disabled or not present + if shouldUpdateBudget && update.Cost > 0 { + if err := t.store.UpdateModelBudgetUsage(ctx, update.Model, &update.Provider, update.Cost); err != nil { + t.logger.Error("failed to update model-level budget usage for model %s: %v", update.Model, err) + } + } + }framework/configstore/tables/modelconfig.go (1)
36-59: Good validation, but inconsistent with TableProvider validation.The validation uses
strings.TrimSpace()which is more thorough than the simple== ""check inTableProvider.BeforeSave. Consider aligning the validation approach across both tables for consistency.In
framework/configstore/tables/provider.go, consider updating the validation to usestrings.TrimSpace()for consistency:// Validate governance fields -if p.BudgetID != nil && *p.BudgetID == "" { +if p.BudgetID != nil && strings.TrimSpace(*p.BudgetID) == "" { return fmt.Errorf("budget_id cannot be an empty string") } -if p.RateLimitID != nil && *p.RateLimitID == "" { +if p.RateLimitID != nil && strings.TrimSpace(*p.RateLimitID) == "" { return fmt.Errorf("rate_limit_id cannot be an empty string") }This would require adding
"strings"to the imports inprovider.go.framework/configstore/rdb.go (1)
1890-1904: Consider batch update for UpdateModelConfigs.The current implementation iterates and saves each model config individually. While this follows the existing pattern (e.g.,
UpdateBudgets,UpdateRateLimits), for large updates this could be inefficient.For now this is acceptable given the expected cardinality of model configs, but consider batching if performance becomes a concern.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.go(1 hunks)framework/configstore/migrations.go(2 hunks)framework/configstore/rdb.go(4 hunks)framework/configstore/store.go(2 hunks)framework/configstore/tables/modelconfig.go(1 hunks)framework/configstore/tables/provider.go(2 hunks)plugins/governance/main.go(5 hunks)plugins/governance/resolver.go(2 hunks)plugins/governance/store.go(8 hunks)plugins/governance/tracker.go(2 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goplugins/governance/tracker.goframework/configstore/migrations.goplugins/governance/main.goplugins/governance/resolver.goframework/configstore/tables/provider.goframework/configstore/store.goframework/configstore/rdb.goplugins/governance/store.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goplugins/governance/tracker.goframework/configstore/migrations.goplugins/governance/main.goplugins/governance/resolver.goframework/configstore/tables/provider.goframework/configstore/store.goframework/configstore/rdb.goplugins/governance/store.go
🧬 Code graph analysis (8)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
framework/configstore/clientconfig.go (3)
framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(282-309)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (2)
core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(282-309)
framework/configstore/tables/provider.go (2)
framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-32)TableRateLimit(35-35)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
plugins/governance/store.go (6)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/utils.go (1)
ParseDuration(9-43)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-32)TableRateLimit(35-35)core/schemas/provider.go (1)
Provider(282-309)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (15)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (19)
framework/configstore/clientconfig.go (1)
705-714: LGTM! New governance config fields are well-structured.The addition of
ModelConfigsandProvidersfields aligns with the hierarchical governance model introduced in this PR. The JSON tags follow the existing naming convention.framework/configstore/tables/provider.go (2)
41-47: LGTM! Provider governance fields follow established patterns.The governance fields and relationships mirror the structure used in
TableVirtualKeyand the newTableModelConfig, maintaining consistency across the codebase.
90-97: LGTM! Validation is consistent with other governance tables.The empty string validation for
BudgetIDandRateLimitIDmatches the pattern used inTableModelConfig.BeforeSave, ensuring consistent data integrity checks.plugins/governance/tracker.go (2)
82-96: LGTM! Provider-level governance updates are correctly prioritized.The implementation correctly applies provider-level rate limit and budget updates before model-level and virtual key checks, ensuring hierarchical governance is enforced even when virtual keys are disabled.
112-123: LGTM! Virtual key handling with early return is clean.The early return when
VirtualKeyis empty prevents unnecessary lookups and clearly separates the non-VK path from VK-specific updates. The debug log provides useful traceability.framework/configstore/migrations.go (3)
104-109: LGTM! Migrations are correctly integrated into the migration chain.The new migrations are placed at the end of the migration sequence, ensuring they run after all prerequisite tables (like
governance_budgetsandgovernance_rate_limits) exist.
1687-1715: LGTM! Model config table migration is well-structured.The migration correctly:
- Checks for table existence before creation
- Uses GORM's migrator for portable schema operations
- Provides a complete rollback function
1717-1769: LGTM! Provider governance columns migration follows established patterns.The migration safely adds
budget_idandrate_limit_idcolumns with existence checks. The rollback correctly reverses the operation.Note: Foreign key constraints and indexes are handled by GORM via the struct tags in
TableProvider, so explicit constraint creation isn't needed here.framework/configstore/tables/modelconfig.go (1)
11-29: LGTM! TableModelConfig struct is well-designed.The struct correctly implements:
- Composite unique index on
(ModelName, Provider)allowing provider-specific or provider-agnostic model configs- Nullable
Providerfield for model-wide governance without provider specificity- Consistent relationship patterns with cascade delete behavior
plugins/governance/main.go (1)
441-465: Hierarchical governance evaluation logic looks correct.The new flow properly:
- Extracts provider/model from request fields first
- Checks VK mandatory requirement and short-circuits if missing
- Evaluates model/provider constraints before VK constraints
- Only evaluates VK constraints if model/provider checks pass and VK is present
This ensures provider/model-level governance applies even when virtual keys are disabled or not present.
framework/configstore/store.go (1)
36-36: Interface extensions are well-structured.The new methods follow the established patterns in the interface:
GetProviderscomplementsGetProvidersConfigfor governance-specific data- Model config CRUD follows the same signature patterns as existing entities (Budget, RateLimit, VirtualKey)
The API is consistent and follows the existing conventions.
Also applies to: 118-126
plugins/governance/resolver.go (2)
77-125: Provider and model evaluation hierarchy is correctly implemented.The evaluation order follows a sensible hierarchy:
- Provider-level rate limits (most granular for provider)
- Provider-level budgets
- Model-level rate limits
- Model-level budgets
The use of a provider pointer for model checks allows proper scoping when a provider is specified.
127-213: VK evaluation properly handles context enrichment and hierarchical checks.The method correctly:
- Sets governance metadata in context for downstream use
- Validates VK state before proceeding
- Checks provider and model filtering
- Performs rate limit and budget hierarchy checks
- Populates include-only-keys for provider matching
framework/configstore/rdb.go (3)
660-670: GetProviders correctly preloads governance relationships.The method follows the established pattern for loading entities with relationships, preloading Budget and RateLimit for governance enforcement.
1832-1848: GetModelConfig handles optional provider correctly.The query properly distinguishes between:
provider = ?when provider is specifiedprovider IS NULLwhen provider is nil (model-only config)This aligns with the unique index
idx_model_provideron the table.
1981-1990: GovernanceConfig correctly includes new fields.The return struct properly includes
ModelConfigsandProviders, and the nil check condition at line 1954 is updated to consider the new fields.plugins/governance/store.go (3)
20-26: In-memory store properly extended with new maps.The addition of
modelConfigsandproviderssync.Map fields follows the established pattern for the other governance entities.
1246-1304: Config memory loading correctly populates relationships.The relationship population for modelConfigs and providers mirrors the existing pattern for virtualKeys, teams, and customers. The nested loops find matching Budget and RateLimit records by ID.
1410-1429: Model configs and providers correctly keyed in memory.The key format is well-designed:
- Model-only configs:
modelName- Model+provider configs:
modelName:provider- Providers:
providerNameThis allows efficient lookups for both specific and general configurations.
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/governance/resolver.go (1)
282-297: Early return on token reset expiry skips request limit check.When token rate limit is expired (lines 283-289), the function returns
DecisionAllowimmediately. However, the request rate limit (lines 290-297) may still be valid and should be checked. The same pattern exists inCheckModelRateLimitin store.go.The logic should continue checking request limits even if token limits are expired:
// Check if rate limit needs reset (in-memory check) + tokenExpired := false if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { // Token rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + tokenExpired = true } } } + requestExpired := false if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { // Request rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + requestExpired = true } } } + // If both are expired, allow the request + if tokenExpired && requestExpired { + return nil, DecisionAllow + }Then update the token/request limit checks to skip only when their respective flags are set.
🧹 Nitpick comments (3)
framework/configstore/tables/provider.go (1)
90-98: Minor inconsistency: validation differs fromTableModelConfig.This validation uses direct empty string comparison, while
TableModelConfig.BeforeSaveusesstrings.TrimSpace(). This means whitespace-only strings like" "would pass validation here but fail inTableModelConfig.Consider aligning for consistency:
// Validate governance fields - if p.BudgetID != nil && *p.BudgetID == "" { + if p.BudgetID != nil && strings.TrimSpace(*p.BudgetID) == "" { return fmt.Errorf("budget_id cannot be an empty string") } - if p.RateLimitID != nil && *p.RateLimitID == "" { + if p.RateLimitID != nil && strings.TrimSpace(*p.RateLimitID) == "" { return fmt.Errorf("rate_limit_id cannot be an empty string") }This would also require adding
"strings"to the imports.plugins/governance/main.go (1)
605-605: Unnamed parameters_, _reduce readability.The
isCacheReadandisBatchparameters are unnamed (_, _), which makes the function signature harder to understand. Consider keeping the names for documentation purposes even if unused, or add a comment explaining why they're ignored.-func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, _, _, isFinalChunk bool) { +func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, isCacheRead, isBatch, isFinalChunk bool) { + // Note: isCacheRead and isBatch are currently unused but preserved for future cache/batch-aware usage trackingframework/configstore/rdb.go (1)
1890-1904: Consider batch update optimization forUpdateModelConfigs.The current implementation iterates and calls
Saveindividually, which may be slow for large batches. Consider using batch updates if performance becomes a concern.For future optimization, consider using GORM's batch update capabilities or a single transaction with prepared statements if this method is called with large slices.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.go(1 hunks)framework/configstore/migrations.go(2 hunks)framework/configstore/rdb.go(4 hunks)framework/configstore/store.go(2 hunks)framework/configstore/tables/modelconfig.go(1 hunks)framework/configstore/tables/provider.go(2 hunks)plugins/governance/main.go(5 hunks)plugins/governance/resolver.go(2 hunks)plugins/governance/store.go(8 hunks)plugins/governance/tracker.go(2 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/tables/provider.goframework/configstore/tables/modelconfig.goplugins/governance/tracker.goframework/configstore/migrations.goplugins/governance/resolver.goframework/configstore/store.goframework/configstore/rdb.goframework/configstore/clientconfig.goplugins/governance/store.goplugins/governance/main.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/tables/provider.goframework/configstore/tables/modelconfig.goplugins/governance/tracker.goframework/configstore/migrations.goplugins/governance/resolver.goframework/configstore/store.goframework/configstore/rdb.goframework/configstore/clientconfig.goplugins/governance/store.goplugins/governance/main.go
🧬 Code graph analysis (8)
framework/configstore/tables/provider.go (2)
framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-32)TableRateLimit(35-35)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(282-309)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (3)
core/schemas/context.go (1)
BifrostContext(30-40)core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(282-309)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
framework/configstore/clientconfig.go (4)
framework/configstore/tables/customer.go (2)
TableCustomer(6-22)TableCustomer(25-25)framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/store.go (7)
framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-32)TableRateLimit(35-35)core/schemas/provider.go (1)
Provider(282-309)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(198-213)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (18)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (20)
framework/configstore/migrations.go (1)
1687-1715: LGTM - Model config table migration follows established patterns.The migration correctly creates the
governance_model_configstable with proper idempotency checks usingHasTable. The rollback is straightforward.framework/configstore/tables/provider.go (1)
41-47: LGTM - Provider governance fields properly defined.The governance fields follow the established patterns:
- Nullable foreign keys with
*stringtype- Proper indexes for query performance
- CASCADE delete to clean up when budgets/rate limits are removed
framework/configstore/tables/modelconfig.go (2)
11-29: LGTM - Well-structured model config table.The table design is sound:
- Composite unique index on
(model_name, provider)correctly allows the same model to have different configs per provider- Nullable
Providerenables model-only configs (provider-agnostic)- Foreign key relationships with CASCADE delete for proper cleanup
- ConfigHash for reconciliation support
36-59: LGTM - Thorough validation in BeforeSave.The validation correctly:
- Requires non-empty
ModelName- Prevents empty-string values for optional fields when provided
- Uses
strings.TrimSpaceconsistently to catch whitespace-only valuesplugins/governance/tracker.go (2)
138-144: LGTM - Clean helper extraction.The
updateBudgetHierarchyhelper cleanly encapsulates the VK→Team→Customer budget update flow with consistent error logging.
82-116: LGTM - Hierarchical usage tracking correctly implemented.All store methods exist with correct signatures. The update sequence properly follows the priority order:
- Provider-level rate limit
- Provider-level budget
- Model-level rate limit
- Model-level budget
- VK-level (only if VirtualKey is present)
Good design: errors are logged but don't block subsequent updates, provider checks gate provider-specific operations, and the early return when no VirtualKey ensures provider/model updates still occur.
framework/configstore/clientconfig.go (1)
705-714: LGTM - GovernanceConfig properly extended.The new
ModelConfigsandProvidersfields align with the new governance tables and follow the existing pattern for governance data types. Both fields are correctly fetched from the database inGetGovernanceConfigand properly populated in the return statement. JSON tags are consistent with codebase conventions.plugins/governance/main.go (2)
441-466: Governance evaluation hierarchy looks correct.The two-stage evaluation (model/provider first, then VK) is well-structured. The flow correctly short-circuits VK evaluation when model/provider checks fail, and only overwrites the result when both checks pass and a VK exists.
563-573: Usage tracking now proceeds without VK, confirmed tracker handles empty VK gracefully.The tracker correctly handles empty virtual keys by skipping only VK-level tracking while always updating provider-level and model-level usage (rate limits and budgets). Explicit checks in
UpdateUsage(line 53) confirm this graceful degradation pattern, with comments confirming provider and model-level updates apply "even when virtual keys are disabled or not present." This design reasonably supports governance without requiring VKs.framework/configstore/store.go (2)
118-126: Interface additions for ModelConfig CRUD are well-structured.The new methods follow the existing patterns in the interface (Get/Create/Update/Delete with optional transaction parameter). The
GetModelConfigwith optional provider parameter allows flexible lookups.
36-36:GetProvidersaddition completes the provider governance support.This method returns providers with their governance relationships (Budget/RateLimit), enabling the new provider-level governance checks.
plugins/governance/resolver.go (2)
77-125:EvaluateModelAndProviderRequestfollows correct priority order.Provider-level checks (rate limits and budgets) are evaluated before model-level checks, which aligns with the hierarchical governance design. The method correctly handles empty provider/model cases.
185-193: InlineEvaluationRequestconstruction for budget check is appropriate.Creating the
EvaluationRequestlocally for the budget hierarchy check avoids exposing implementation details while reusing existing budget checking logic.framework/configstore/rdb.go (3)
660-670:GetProvidersimplementation correctly preloads governance relationships.The preloading of Budget and RateLimit ensures complete provider data is available for governance checks without N+1 queries.
1832-1848:GetModelConfiglookup logic handles NULL provider correctly.The conditional query correctly differentiates between model-only configs (provider IS NULL) and model+provider configs, using the composite unique index.
1943-1948: Model configs and providers loading added toGetGovernanceConfig.The additions correctly populate the governance configuration with the new model configs and providers data.
plugins/governance/store.go (4)
20-26: New in-memory maps for modelConfigs and providers are well-documented.The key format documentation ("modelName" or "modelName:provider" for modelConfigs, provider name for providers) is clear and consistent with the implementation.
1246-1304: Relationship population for modelConfigs and providers inloadFromConfigMemoryis thorough.The nested loops correctly wire up Budget and RateLimit relationships for both modelConfigs (lines 1252-1277) and providers (lines 1279-1304) before the in-memory rebuild.
1371-1429:rebuildInMemoryStructurescorrectly handles new maps.The function properly clears and rebuilds all sync.Maps including the new modelConfigs and providers maps. The key format logic (lines 1416-1420) correctly distinguishes between model-only and model+provider configs.
46-55: Interface additions for provider/model governance are comprehensive.The new methods cover all necessary operations: checking and updating budgets and rate limits for both providers and models. The method signatures are consistent with existing patterns.
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
543a580 to
fc68a13
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
plugins/governance/main.go (1)
563-573: Usage tracking is still skipped when model is empty despite comments claiming “always process”The comment says “Always process usage tracking (with or without virtual key)” and that the tracker will handle empty VK by falling back to provider/model usage, but the
if model != ""guard meanspostHookWorkeris never invoked whenmodelis empty, so:
- No provider‑level or model‑level usage is updated for such requests.
- This contradicts the documented behavior and the new provider/model‑level tracking added in this PR.
To align behavior with the comment and with
UsageTracker.UpdateUsage, either:
- Remove the
if model != ""conditional sopostHookWorkeris always scheduled (still passing an empty model/provider when not available), or- Keep the guard but update the comments to clearly state that usage tracking is skipped when the model cannot be resolved.
Given earlier review feedback on this exact pattern, I’d strongly recommend removing the guard so provider/model tracking can run even when VK is empty and model resolution is flaky.
Also applies to: 593-601, 641-656
plugins/governance/store.go (1)
206-262: Rate-limit expiry logic can skip checking non-expired limits (provider/model) and under-enforce constraintsIn both provider- and model-level rate-limit checks:
CheckProviderRateLimitreturns(nil, DecisionAllow)as soon as either the token or request reset window is expired (lines 304–319), which:
- Completely skips checking the other limit type for that provider, even if it is still in-window and over its max.
CheckModelRateLimitusescontinuewhen a token or request reset window is expired (lines 497–512), which:
- Skips the entire model config (both token and request limits) whenever either side is expired.
This means that if, for example, the token window has expired but the request window has not (and the request count is already over its limit), the request limit is never evaluated and the request is incorrectly allowed.
A safer pattern is to:
- Track expiry per limit type (e.g.,
tokenExpired,requestExpired), and:
- Skip only the specific limit that has expired.
- Still evaluate the other limit for the same provider/model config.
- Only treat the entire rate limit as “effectively reset” if both token and request windows are expired.
This matches how you want per-dimension rate limits to behave and avoids silently disabling one dimension due to the other expiring.
Also applies to: 264-364, 366-439, 441-558
🧹 Nitpick comments (4)
plugins/governance/tracker.go (1)
73-81: Hierarchical provider/model updates are correct; optional small guard on empty modelThe new ordering (provider RL → provider budget → model RL → model budget → VK) matches the intended hierarchy and will update usage even when no VK is present. The only minor nit is that
UpdateModelRateLimitUsage/UpdateModelBudgetUsageare called even whenupdate.Model == "", which currently results in a no-op but still does the map lookups. You could optionally guard those calls onupdate.Model != ""to avoid unnecessary work, but it’s not functionally required.Also applies to: 82-111, 112-116
framework/configstore/rdb.go (1)
1820-1915: ModelConfig CRUD and governance config aggregation look coherentThe new ModelConfig methods (
GetModelConfigs,GetModelConfig,GetModelConfigByID,Create/Update/UpdateModelConfigs/DeleteModelConfig) correctly:
- Preload
Budget/RateLimitso the in-memory store can rely on relationships.- Use
(model_name, provider)as the lookup, matching theuniqueIndex:idx_model_providerdefinition.- Respect existing error translation patterns (
ErrNotFound,parseGormError).
GetGovernanceConfignow also loadsmodelConfigsandprovidersand includes them in the returnedGovernanceConfig, with the presence check updated so an entirely empty governance DB still returnsnil. This all lines up with the extended public config shape.Also applies to: 1918-1990
plugins/governance/store.go (2)
20-27: In-memory maps for providers/model configs are well-keyed; consider exposing them in GovernanceDataAdding
modelConfigsandproviderssync.Maps (keyed asmodelName/modelName:providerand provider name) and wiring them throughrebuildInMemoryStructuresmatches how the resolver and tracker look these up and keeps reads lock‑free. If you expect to debug or introspect provider/model governance state viaGetGovernanceData, you might also want to surface modelConfigs/providers there, but that’s optional and doesn’t affect behavior.Also applies to: 35-41, 1423-1433, 1464-1483
674-759: Update*Usage clones should avoid mutating the loaded struct to reduce race riskThe new provider/model
Update*Usagemethods follow the existing pattern of:if cached, ok := gs.budgets.Load(...); ok { cached.CurrentUsage = 0 cached.LastReset = now clone := *cached clone.CurrentUsage += cost gs.budgets.Store(..., &clone) }and similarly for rate limits. This mutates the struct pointed to by the
sync.Mapbefore cloning, so any concurrent goroutine that has also loaded the same pointer can observe partially updated fields, which is a classic data race.A safer pattern (still lock‑free at the map level) is:
if v, ok := gs.budgets.Load(id); ok { orig := v.(*configstoreTables.TableBudget) clone := *orig // mutate clone only if clone.ResetDuration != "" { ... use clone.LastReset ... } clone.CurrentUsage += cost gs.budgets.Store(id, &clone) }and the analogous change for rate limits. That way, the object returned from
Loadis never mutated after publication; only new clones are written back.Given the same pattern exists in older paths, this is more of a refactor than a blocker, but it would make the concurrency story materially cleaner.
Also applies to: 761-870
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.go(1 hunks)framework/configstore/migrations.go(2 hunks)framework/configstore/rdb.go(4 hunks)framework/configstore/store.go(2 hunks)framework/configstore/tables/modelconfig.go(1 hunks)framework/configstore/tables/provider.go(2 hunks)plugins/governance/main.go(5 hunks)plugins/governance/resolver.go(2 hunks)plugins/governance/store.go(8 hunks)plugins/governance/tracker.go(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- framework/configstore/tables/modelconfig.go
- framework/configstore/store.go
- framework/configstore/migrations.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/tracker.goframework/configstore/clientconfig.goplugins/governance/store.goframework/configstore/rdb.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/tracker.goframework/configstore/clientconfig.goplugins/governance/store.goframework/configstore/rdb.go
🧬 Code graph analysis (7)
framework/configstore/tables/provider.go (3)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-32)TableRateLimit(35-35)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(198-213)
plugins/governance/resolver.go (2)
core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(282-309)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(282-309)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
framework/configstore/clientconfig.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/store.go (10)
ui/lib/types/config.ts (1)
ModelProvider(172-175)plugins/governance/resolver.go (5)
Decision(15-15)DecisionAllow(18-18)DecisionRateLimited(21-21)DecisionTokenLimited(23-23)DecisionRequestLimited(24-24)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/budget.go (2)
TableBudget(11-24)TableBudget(27-27)framework/configstore/tables/utils.go (1)
ParseDuration(9-43)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-32)TableRateLimit(35-35)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)core/schemas/provider.go (1)
Provider(282-309)
framework/configstore/rdb.go (4)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(705-714)AuthConfig(695-700)ui/lib/types/config.ts (1)
AuthConfig(233-238)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (17)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (6)
framework/configstore/tables/provider.go (1)
41-47: Provider-level governance fields and validation look consistentThe added
BudgetID/RateLimitIDfields and relationships mirror existing governance patterns (e.g., on virtual keys and model configs), and theBeforeSavechecks correctly reject empty-string IDs while allowing nil. No functional issues spotted here.Also applies to: 91-97
framework/configstore/clientconfig.go (1)
705-713: GovernanceConfig extension to include model configs/providers is soundAdding
ModelConfigsandProviderstoGovernanceConfigis a backward‑compatible extension and matches the UI side types; existing consumers using the old fields should continue to work untouched.plugins/governance/main.go (1)
441-465: PreHook hierarchy (model/provider first, then VK) aligns with new resolver APIExtracting
provider, modelviaGetRequestFieldsand first callingEvaluateModelAndProviderRequest, thenEvaluateVirtualKeyRequestonly when VK is present and model/provider checks passed, matches the intended evaluation order and keeps VK checks scoped to already‑allowed requests.plugins/governance/resolver.go (1)
77-125: Split model/provider vs VK evaluation is clean and matches decision API
EvaluateModelAndProviderRequestandEvaluateVirtualKeyRequestclearly separate concerns: provider/model limits are checked first, VK validation/filtering and VK‑hierarchy rate‑limit/budget checks only run when applicable, and decisions are mapped consistently toDecision*values used byPreHook. The parameterization with explicitprovider,model, andvirtualKeyValuemakes the flow easier to reason about than the previous monolithic request struct.Also applies to: 128-213
framework/configstore/rdb.go (1)
660-670: GetProviders wiring matches new provider-level governance fields
GetProviderspreloadingBudgetandRateLimitensures the in-memory governance store can hydrate provider-level governance without extra queries. Returning an empty slice (rather than ErrNotFound) when no providers exist is fine and consistent with the rest of the store.plugins/governance/store.go (1)
1260-1271: Database/config loading for modelConfigs/providers is consistent with in-memory store expectations
loadFromDatabasenow pullsmodelConfigsviaGetModelConfigsandprovidersviaGetProviders, both of which preloadBudgetandRateLimit.loadFromConfigMemorymirrors that by wiringModelConfigs/Providersfrom the in-memoryGovernanceConfigand manually populating theBudget/RateLimitrelationships before callingrebuildInMemoryStructures.This keeps the in-memory
modelConfigs/providersmaps fully hydrated regardless of whether you boot from DB or config memory, and matches how the resolver and tracker query them.Also applies to: 1299-1357, 1417-1483
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
fc68a13 to
d0af6d9
Compare
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
0ed2f53 to
b9764ab
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (3)
plugins/governance/main.go (1)
578-588: Clarify the PostHook usage-tracking comment relative to themodel != ""guard.The comment states that usage tracking is “always” processed and that the tracker handles empty VKs, but the
if model != ""guard still skips tracking entirely when the model is empty—so provider-level updates also don’t run in that case. If that behavior is intentional (e.g., non-model endpoints), it’d be clearer to soften the comment to “for modelled requests” or similar.plugins/governance/store.go (2)
337-352: Early return on expired limit skips checking the other limit type.If the token reset duration is expired (lines 337-343), the method returns
DecisionAllowimmediately without checking whether the request limit is still active and potentially exceeded. The same issue occurs for expired request limits (lines 345-352).This means if only the token limit is expired but the request limit is still active and exceeded, the request would incorrectly be allowed.
🔎 Fix: Track expiry per limit type
// Check if rate limit needs reset (in-memory check) + tokenExpired := false + requestExpired := false + if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - // Token rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + tokenExpired = true } } } if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - // Request rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + requestExpired = true } } }Then conditionally skip only the specific expired check when validating limits:
// Token limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { + if !tokenExpired && rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { duration := "unknown" if rateLimit.TokenResetDuration != nil { duration = *rateLimit.TokenResetDuration } violations = append(violations, fmt.Sprintf("token limit exceeded (%d/%d, resets every %s)", rateLimit.TokenCurrentUsage+tokensBaseline, *rateLimit.TokenMaxLimit, duration)) } // Request limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { + if !requestExpired && rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { duration := "unknown" if rateLimit.RequestResetDuration != nil { duration = *rateLimit.RequestResetDuration } violations = append(violations, fmt.Sprintf("request limit exceeded (%d/%d, resets every %s)", rateLimit.RequestCurrentUsage+requestsBaseline, *rateLimit.RequestMaxLimit, duration)) }
529-545: Early continue on expired limit skips checking the other limit type for this model config.When either the token reset duration is expired (lines 530-536) or the request reset duration is expired (lines 538-544), the code uses
continueto skip the entire model config. This means if only one limit type is expired while the other is still active and exceeded, the request would incorrectly be allowed.🔎 Fix: Track expiry per limit type
// Check if rate limit needs reset (in-memory check) + tokenExpired := false + requestExpired := false + if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - // Token rate limit expired but hasn't been reset yet - skip check - continue + tokenExpired = true } } } if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - // Request rate limit expired but hasn't been reset yet - skip check - continue + requestExpired = true } } }Then conditionally skip only the specific expired check when validating limits:
// Token limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { + if !tokenExpired && rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { duration := "unknown" if rateLimit.TokenResetDuration != nil { duration = *rateLimit.TokenResetDuration } violations = append(violations, fmt.Sprintf("token limit exceeded (%d/%d, resets every %s)", rateLimit.TokenCurrentUsage+tokensBaseline, *rateLimit.TokenMaxLimit, duration)) } // Request limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { + if !requestExpired && rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { duration := "unknown" if rateLimit.RequestResetDuration != nil { duration = *rateLimit.RequestResetDuration } violations = append(violations, fmt.Sprintf("request limit exceeded (%d/%d, resets every %s)", rateLimit.RequestCurrentUsage+requestsBaseline, *rateLimit.RequestMaxLimit, duration)) }
🧹 Nitpick comments (4)
framework/configstore/migrations.go (1)
128-133: Migrations for model configs and provider governance look correct; consider normalizing column-name usage.The new migrations are idempotent and follow the existing style:
migrationAddModelConfigTablesafely creates/dropsgovernance_model_configs, andmigrationAddProviderGovernanceColumnsaddsbudget_id/rate_limit_idplus explicit indexes and clean rollback. That’s good.Two minor polish points:
- In
migrationAddProviderGovernanceColumns,HasColumnuses"budget_id"/"rate_limit_id"whileAddColumn/DropColumnuse"BudgetID"/"RateLimitID". GORM will resolve these correctly via schema metadata, but for consistency with the rest of this file (which usually passes snake_case column names) you might want to stick to one convention, e.g. always"budget_id"/"rate_limit_id"forHasColumn/AddColumn/DropColumn.- If you want DB-level referential integrity for provider governance similar to
migrationAddProviderConfigBudgetRateLimit, you could also addCreateConstraint/DropConstraintfor the providerBudget/RateLimitFKs here; right now only indexes are enforced.Also applies to: 2090-2197
plugins/governance/tracker.go (1)
72-123: Hierarchical provider/model/VK usage updates look sound; optional guard for empty model.The new order (provider-level → model-level → VK-level) and the streaming gates (
shouldUpdateTokens,shouldUpdateRequests,shouldUpdateBudget) are consistent with the governance design and allow VK‑less requests to still update provider/model usage.Given
UpdateUsageis currently only called from a path wheremodelis non-empty, this is correct. For extra robustness against future callers, you might consider short‑circuiting the model-level updates whenupdate.Model == ""to make the contract self‑contained, but it’s not required by the current call graph.plugins/governance/main.go (1)
606-671: Consider removing unused cache/batch flags from postHookWorker’s signature.
postHookWorkernow ignoresisCacheReadandisBatchvia blank identifiers in the parameter list but the parameters are still threaded through PostHook. If they’re no longer needed for tracking semantics, simplifying the signature and call to only passisFinalChunkwould reduce noise.Possible simplification
-func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, _, _, isFinalChunk bool) { +func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, isFinalChunk bool) { @@ - p.postHookWorker(result, provider, model, requestType, virtualKey, requestID, isCacheRead, isBatch, bifrost.IsFinalChunk(ctx)) + p.postHookWorker(result, provider, model, requestType, virtualKey, requestID, bifrost.IsFinalChunk(ctx))framework/configstore/rdb.go (1)
1930-2025: ModelConfig CRUD API is consistent with existing store patterns.The new methods (
GetModelConfigs,GetModelConfig,GetModelConfigByID,CreateModelConfig,UpdateModelConfig,UpdateModelConfigs,DeleteModelConfig) follow the existing RDBConfigStore conventions:
- Preloading
Budget/RateLimitwhere appropriate.- Distinguishing provider-scoped vs global configs via the optional
providerparameter (provider IS NULLvs exact match).- Using
parseGormErroron create/update paths.Only small note:
DeleteModelConfigwill not surfaceErrNotFoundwhen no rows match (sinceDeletedoesn’t returnErrRecordNotFound), but that’s consistent with several other delete helpers here.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (3)
- framework/configstore/tables/modelconfig.go
- framework/configstore/tables/provider.go
- framework/configstore/store.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/migrations.goplugins/governance/tracker.goframework/configstore/clientconfig.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goframework/configstore/rdb.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/migrations.goplugins/governance/tracker.goframework/configstore/clientconfig.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goframework/configstore/rdb.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.go
🧬 Code graph analysis (5)
framework/configstore/migrations.go (3)
framework/migrator/migrator.go (3)
New(131-149)DefaultOptions(100-106)Migration(62-69)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
framework/configstore/clientconfig.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (4)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)
framework/configstore/rdb.go (3)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (1)
GovernanceConfig(720-729)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (13)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (6)
framework/configstore/clientconfig.go (1)
720-729: GovernanceConfig extension for model configs and providers is coherent.Adding
ModelConfigsandProvidershere lines up with the new RDB getters andGetGovernanceConfigwiring; the JSON schema remains straightforward for consumers.plugins/governance/main.go (1)
456-480: PreHook’s split between provider/model and VK evaluation is consistent and ordered correctly.Extracting
(provider, model)viaGetRequestFieldsand runningEvaluateModelAndProviderRequestfirst, then layeringEvaluateVirtualKeyRequestonly when the first decision isallowand a VK is present, gives a clean hierarchy: global provider/model governance first, then VK-specific constraints. This matches the new resolver APIs and keeps VK-optional flows working.plugins/governance/resolver.go (2)
77-125: Provider/model evaluation path is well-factored and safely gated.
EvaluateModelAndProviderRequestcleanly handles provider-level and model-level rate limits/budgets, only invoking the relevant checks whenprovider/modelare non-empty and returning precise decisions/reasons. TheproviderPtrindirection for model-level checks is a reasonable way to keep the API compatible with model configs that may or may not be provider-scoped.
127-213: VK evaluation correctly reuses shared helpers and explicit provider/model parameters.
EvaluateVirtualKeyRequestnow:
- Validates VK existence/activeness and seeds context with VK/team/customer metadata.
- Uses
isProviderAllowed/isModelAllowedbased on the explicitproviderandmodelarguments.- Runs VK-centric rate limit checks via
checkRateLimitHierarchyand hierarchical budget checks viacheckBudgetHierarchy.- Finally sets
bf-governance-include-only-keysfor matching provider configs.The separation from the provider/model evaluation path in
PreHookis clear and preserves previous VK behavior while allowing provider/model governance to stand on its own.framework/configstore/rdb.go (2)
701-711: Provider fetch with governance preloads looks good.
GetProviderscorrectly returns all providers withBudgetandRateLimitpreloaded, and treats “no providers” as an empty slice without error, which matches howGetGovernanceConfighandles emptiness. TheErrNotFoundbranch will effectively never trigger forFind, but it’s harmless.
2034-2100: GovernanceConfig now exposes model configs and providers coherently.
GetGovernanceConfignow loadsmodelConfigsandproviders, includes them in the “no config” check, and returns them inGovernanceConfigalongside virtual keys, teams, customers, budgets, rate limits, and auth. This lines up with the extendedGovernanceConfigstruct and gives the governance plugin enough data to build model- and provider-level governance state.
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (2)
plugins/governance/store.go (2)
337-352: Early return on expired reset duration skips subsequent limit checks.If the token reset duration is expired (lines 337-343), the method returns
DecisionAllowimmediately, skipping the request limit checks entirely (lines 345-352 and lines 364-381). Similarly, if the request reset duration is expired (lines 345-352), it returns early.This means if only the token limit is expired but the request limit is still active and exceeded, the request would incorrectly be allowed.
🔎 Recommended fix
Apply the same fix pattern used in the VK-level
CheckRateLimitmethod (lines 609-633):- // Check if rate limit needs reset (in-memory check) + // Determine token and request expiration independently + tokenExpired := false + requestExpired := false + + // Check if token reset duration is expired if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - // Token rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + tokenExpired = true } } } + + // Check if request reset duration is expired if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - // Request rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + requestExpired = true } } }Then modify the limit checks to use the flags:
// Token limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { + if !tokenExpired && rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { // ... violation handling } // Request limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { + if !requestExpired && rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { // ... violation handling }Based on learnings, this is the same issue identified in previous reviews for similar rate limit checks.
529-545: Early continue on expired reset duration skips subsequent limit checks for model config.When either token or request reset duration is expired, the code uses
continue(lines 534, 542), which skips checking the other limit type for that model config. If only the token limit is expired but the request limit is still active and exceeded, the request would incorrectly be allowed for that model configuration.🔎 Recommended fix
Apply the same fix pattern used in the VK-level
CheckRateLimitmethod:for i, mc := range modelConfigsToCheck { // ... rate limit lookup ... - // Check if rate limit needs reset (in-memory check) + // Determine token and request expiration independently + tokenExpired := false + requestExpired := false + + // Check if token reset duration is expired if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - // Token rate limit expired but hasn't been reset yet - skip check - continue + tokenExpired = true } } } + + // Check if request reset duration is expired if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - // Request rate limit expired but hasn't been reset yet - skip check - continue + requestExpired = true } } }Then conditionally check only non-expired limits:
// Token limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { + if !tokenExpired && rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { // ... violation handling } // Request limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { + if !requestExpired && rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { // ... violation handling }Based on learnings, this is the same issue identified in previous reviews for similar rate limit checks.
🧹 Nitpick comments (2)
plugins/governance/main.go (1)
620-620: Consider adding a brief comment explaining unused parameters.The unnamed parameters
_, _forisCacheReadandisBatchare intentionally unused. A brief inline comment would clarify this is deliberate and prevent future confusion about whether these should be utilized.🔎 Suggested clarification
-func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, _, _, isFinalChunk bool) { +func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, _, _ /* isCacheRead, isBatch: reserved for future use */, isFinalChunk bool) {framework/configstore/migrations.go (1)
2120-2156: Consider adding FK constraints for referential integrity.The migration correctly adds columns and explicit indexes (addressing the past review concern). However, unlike
migrationAddProviderConfigBudgetRateLimit(lines 716-725), this migration doesn't create FK constraints forBudgetandRateLimitrelationships. This could lead to orphaned references if budgets/rate limits are deleted without CASCADE.🔎 Suggested addition after index creation
// Create index for rate_limit_id if !migrator.HasIndex(provider, "idx_provider_rate_limit") { if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_rate_limit ON config_providers (rate_limit_id)").Error; err != nil { return fmt.Errorf("failed to create rate_limit_id index: %w", err) } } } + // Create FK constraints (dialect-agnostic) for referential integrity + if !migrator.HasConstraint(provider, "Budget") { + if err := migrator.CreateConstraint(provider, "Budget"); err != nil { + return fmt.Errorf("failed to create Budget FK constraint: %w", err) + } + } + if !migrator.HasConstraint(provider, "RateLimit") { + if err := migrator.CreateConstraint(provider, "RateLimit"); err != nil { + return fmt.Errorf("failed to create RateLimit FK constraint: %w", err) + } + } + return nilAnd in the Rollback, add constraint drops before index drops:
Rollback: func(tx *gorm.DB) error { tx = tx.WithContext(ctx) migrator := tx.Migrator() provider := &tables.TableProvider{} + // Drop FK constraints first + if migrator.HasConstraint(provider, "RateLimit") { + if err := migrator.DropConstraint(provider, "RateLimit"); err != nil { + return fmt.Errorf("failed to drop RateLimit FK constraint: %w", err) + } + } + if migrator.HasConstraint(provider, "Budget") { + if err := migrator.DropConstraint(provider, "Budget"); err != nil { + return fmt.Errorf("failed to drop Budget FK constraint: %w", err) + } + } + // Drop indexes first if migrator.HasIndex(provider, "idx_provider_rate_limit") {
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (1)
- framework/configstore/tables/provider.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/migrations.goplugins/governance/tracker.goplugins/governance/resolver.goframework/configstore/clientconfig.goframework/configstore/tables/modelconfig.goframework/configstore/store.goplugins/governance/main.goframework/configstore/rdb.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/migrations.goplugins/governance/tracker.goplugins/governance/resolver.goframework/configstore/clientconfig.goframework/configstore/tables/modelconfig.goframework/configstore/store.goplugins/governance/main.goframework/configstore/rdb.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/tracker.goplugins/governance/resolver.goplugins/governance/main.goplugins/governance/store.go
🧬 Code graph analysis (8)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
plugins/governance/resolver.go (2)
core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)
framework/configstore/clientconfig.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
framework/configstore/store.go (1)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
framework/configstore/rdb.go (3)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(720-729)AuthConfig(710-715)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (26)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (18)
framework/configstore/tables/modelconfig.go (2)
12-29: LGTM! Well-structured model with appropriate validations.The
TableModelConfigstruct is well-designed with proper GORM tags, a composite unique index on(ModelName, Provider)to prevent duplicates, and appropriate foreign key relationships. The nullableProviderfield correctly allows for model-level configs that apply across all providers.
36-59: LGTM! Proper validation in BeforeSave hook.The validation correctly ensures that when optional pointer fields are provided, they contain meaningful values rather than empty strings. This prevents data integrity issues.
plugins/governance/tracker.go (2)
81-115: LGTM! Well-structured hierarchical usage tracking.The multi-stage update sequence correctly prioritizes provider-level updates first, followed by model-level, and finally VK-level. This ensures usage is tracked at all governance levels even when a virtual key is absent. The error handling approach (log and continue) is appropriate here to prevent one failed update from blocking others.
117-139: LGTM! Clean VK-conditional flow.The VK handling is properly isolated and only executes when a virtual key is present. The graceful handling of missing VKs (logging and returning without error) aligns with the VK-optional design introduced in this PR.
plugins/governance/main.go (2)
473-481: LGTM! Proper hierarchical evaluation order.The evaluation flow correctly implements the hierarchical governance model: model/provider checks are evaluated first via
EvaluateModelAndProviderRequest, and VK-specific checks only run if the initial checks pass and a VK is present. This ensures provider/model-level governance applies universally regardless of VK presence.
578-588: LGTM! Usage tracking with VK-optional support.The implementation correctly processes usage tracking for any request with a non-empty model, regardless of whether a virtual key is present. The comments accurately describe that the tracker handles empty virtual keys gracefully by updating only provider-level and model-level usage. This addresses the previous review concern about misleading comments.
plugins/governance/resolver.go (2)
77-125: LGTM! Clean implementation of provider/model governance checks.The
EvaluateModelAndProviderRequestmethod correctly implements the hierarchical check order: provider-level rate limits and budgets are checked first, followed by model-level checks. The optional provider pointer handling (lines 100-103) is appropriate for cases where provider may not be specified.
127-213: LGTM! Comprehensive VK evaluation with proper context enrichment.The
EvaluateVirtualKeyRequestmethod correctly handles the full VK evaluation flow: validation, provider/model filtering, rate limit checks, and budget hierarchy checks. The context enrichment (lines 139-152) properly propagates VK, team, and customer metadata for downstream observability.framework/configstore/migrations.go (1)
2090-2118: LGTM! Clean migration for model config table.The migration correctly creates the
governance_model_configstable using theTableModelConfigstruct definition. GORM will create the necessary indexes defined in the struct tags duringCreateTable. The rollback properly drops the table.framework/configstore/clientconfig.go (1)
720-729: LGTM - GovernanceConfig extended with model and provider fields.The additions of
ModelConfigsandProvidersfields are well-placed and properly typed, aligning with the PR's goal of supporting per-model and per-provider governance.framework/configstore/store.go (2)
36-36: LGTM - GetProviders method added to interface.The method signature is consistent with existing patterns and correctly returns a slice of providers.
118-126: LGTM - Model config CRUD methods added to interface.The method signatures follow established patterns and provide comprehensive CRUD operations for model configurations, including flexible query options (by ID, by model name with optional provider).
framework/configstore/rdb.go (3)
701-711: LGTM - GetProviders implementation follows established patterns.The method correctly preloads governance relationships (Budget and RateLimit) and handles errors consistently with other retrieval methods in the store.
1930-2025: LGTM - Model config CRUD implementations are well-structured.All methods follow established patterns, correctly preload relationships, handle optional transactions, and include proper error handling. The
GetModelConfigmethod appropriately handles both model-only and model+provider query scenarios.
2027-2101: LGTM - GetGovernanceConfig correctly integrated with model configs and providers.The method now loads and returns model configurations and providers alongside existing governance entities, with the presence check properly updated to include the new fields.
plugins/governance/store.go (3)
18-33: LGTM - In-memory maps added for model configs and providers.The new
modelConfigsandproviderssync.Map fields are well-documented with clear key format comments, following the established pattern for other governance entities.
52-92: LGTM - GovernanceStore interface extended with provider and model level operations.The new methods provide hierarchical governance checks and updates (model → provider → VK), with clear naming and consistent signatures.
1322-1576: LGTM - Loading and rebuilding logic correctly integrated.The database loading, config memory loading, and in-memory structure rebuilding all properly handle the new model configs and providers. The relationship population in
loadFromConfigMemorycorrectly links budgets and rate limits, and the key formats for the sync.Maps match the documented conventions.
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
b9764ab to
f6b85fa
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
2 times, most recently
from
79247f6 to
98c8540
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (6)
plugins/governance/store.go (6)
337-352: Critical: Early return skips independent limit checks.Lines 337-344 return
DecisionAllowif the token reset duration is expired, and lines 345-352 do the same for the request reset duration. This skips checking the other limit type entirely.For example, if the token limit is expired but the request limit is still active and exceeded, the request would incorrectly be allowed.
🔎 Recommended fix
Track expiration per-limit independently:
// Check if rate limit needs reset (in-memory check) + tokenExpired := false if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - // Token rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + tokenExpired = true } } } + requestExpired := false if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - // Request rate limit expired but hasn't been reset yet - skip check - return nil, DecisionAllow + requestExpired = true } } } + // If both are expired, allow request + if tokenExpired && requestExpired { + return nil, DecisionAllow + } // ... then check token limits only if !tokenExpired, request limits only if !requestExpired - // Token limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { + if !tokenExpired && rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { // ... violations } - // Request limits - check if total usage (local + remote baseline) exceeds limit - if rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { + if !requestExpired && rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { // ... violations }
529-545: Critical: Continue skips independent limit checks for model config.Lines 530-536 use
continuewhen the token reset duration is expired, and lines 538-544 do the same for the request reset duration. This skips checking the other limit type for that model config entirely.If either the token OR request limit is expired, the entire rate limit check for that model config is skipped via
continue.🔎 Recommended fix
Track expiration per-limit independently and only skip the specific expired check:
for i, mc := range modelConfigsToCheck { // ... load rateLimit ... // Check if rate limit needs reset (in-memory check) + tokenExpired := false if rateLimit.TokenResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.TokenResetDuration); err == nil { if time.Since(rateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - // Token rate limit expired but hasn't been reset yet - skip check - continue + tokenExpired = true } } } + requestExpired := false if rateLimit.RequestResetDuration != nil { if duration, err := configstoreTables.ParseDuration(*rateLimit.RequestResetDuration); err == nil { if time.Since(rateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - // Request rate limit expired but hasn't been reset yet - skip check - continue + requestExpired = true } } } + // If both are expired, skip this model config + if tokenExpired && requestExpired { + continue + } // ... then check token limits only if !tokenExpired, request limits only if !requestExpired - if rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { + if !tokenExpired && rateLimit.TokenMaxLimit != nil && rateLimit.TokenCurrentUsage+tokensBaseline >= *rateLimit.TokenMaxLimit { // ... violations } - if rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { + if !requestExpired && rateLimit.RequestMaxLimit != nil && rateLimit.RequestCurrentUsage+requestsBaseline >= *rateLimit.RequestMaxLimit { // ... violations } }
742-760: Critical data race: modifying shared budget before cloning.Lines 746-750 directly mutate
cachedBudget(settingCurrentUsageandLastReset) before creating a clone at line 754. This creates a data race where concurrent goroutines may read the shared budget object while it's being modified.While
sync.Mapprovides atomic Load/Store operations, it does NOT protect the values themselves from concurrent modification.🔎 Recommended fix
Clone FIRST, then operate on the clone (compare with
UpdateBudgetUsageInMemoryat lines 696-717 which does this correctly):if cachedBudgetValue, exists := gs.budgets.Load(*providerTable.BudgetID); exists && cachedBudgetValue != nil { if cachedBudget, ok := cachedBudgetValue.(*configstoreTables.TableBudget); ok && cachedBudget != nil { + // Clone FIRST to avoid race conditions + clone := *cachedBudget now := time.Now() + // Check if budget needs reset (in-memory check) - if cachedBudget.ResetDuration != "" { - if duration, err := configstoreTables.ParseDuration(cachedBudget.ResetDuration); err == nil { - if now.Sub(cachedBudget.LastReset).Round(time.Millisecond) >= duration { - cachedBudget.CurrentUsage = 0 - cachedBudget.LastReset = now + if clone.ResetDuration != "" { + if duration, err := configstoreTables.ParseDuration(clone.ResetDuration); err == nil { + if now.Sub(clone.LastReset).Round(time.Millisecond) >= duration { + clone.CurrentUsage = 0 + clone.LastReset = now } } } - clone := *cachedBudget + clone.CurrentUsage += cost gs.budgets.Store(*providerTable.BudgetID, &clone) } }
779-810: Critical data race: modifying shared rate limit before cloning.Lines 783-787 and 791-795 directly mutate
cachedRateLimit(resettingTokenCurrentUsage,TokenLastReset,RequestCurrentUsage,RequestLastReset) before creating a clone at line 799. This creates a data race with concurrent reads.🔎 Recommended fix
Clone FIRST, then operate on the clone:
if cachedRateLimitValue, exists := gs.rateLimits.Load(*providerTable.RateLimitID); exists && cachedRateLimitValue != nil { if cachedRateLimit, ok := cachedRateLimitValue.(*configstoreTables.TableRateLimit); ok && cachedRateLimit != nil { + // Clone FIRST to avoid race conditions + clone := *cachedRateLimit now := time.Now() + // Check if rate limit needs reset (in-memory check) - if cachedRateLimit.TokenResetDuration != nil { - if duration, err := configstoreTables.ParseDuration(*cachedRateLimit.TokenResetDuration); err == nil { - if now.Sub(cachedRateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - cachedRateLimit.TokenCurrentUsage = 0 - cachedRateLimit.TokenLastReset = now + if clone.TokenResetDuration != nil { + if duration, err := configstoreTables.ParseDuration(*clone.TokenResetDuration); err == nil { + if now.Sub(clone.TokenLastReset).Round(time.Millisecond) >= duration { + clone.TokenCurrentUsage = 0 + clone.TokenLastReset = now } } } - if cachedRateLimit.RequestResetDuration != nil { - if duration, err := configstoreTables.ParseDuration(*cachedRateLimit.RequestResetDuration); err == nil { - if now.Sub(cachedRateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - cachedRateLimit.RequestCurrentUsage = 0 - cachedRateLimit.RequestLastReset = now + if clone.RequestResetDuration != nil { + if duration, err := configstoreTables.ParseDuration(*clone.RequestResetDuration); err == nil { + if now.Sub(clone.RequestLastReset).Round(time.Millisecond) >= duration { + clone.RequestCurrentUsage = 0 + clone.RequestLastReset = now } } } - clone := *cachedRateLimit + if shouldUpdateTokens { clone.TokenCurrentUsage += tokensUsed } if shouldUpdateRequests { clone.RequestCurrentUsage += 1 } gs.rateLimits.Store(*providerTable.RateLimitID, &clone) } }
841-859: Critical data race: modifying shared budget before cloning (model-level).Lines 844-848 directly mutate
cachedBudgetbefore creating a clone at line 852. This has the same data race asUpdateProviderBudgetUsage.🔎 Recommended fix
Clone FIRST:
if cachedBudgetValue, exists := gs.budgets.Load(*mc.BudgetID); exists && cachedBudgetValue != nil { if cachedBudget, ok := cachedBudgetValue.(*configstoreTables.TableBudget); ok && cachedBudget != nil { + // Clone FIRST to avoid race conditions + clone := *cachedBudget + now := time.Now() + // Check if budget needs reset (in-memory check) - if cachedBudget.ResetDuration != "" { - if duration, err := configstoreTables.ParseDuration(cachedBudget.ResetDuration); err == nil { - if now.Sub(cachedBudget.LastReset).Round(time.Millisecond) >= duration { - cachedBudget.CurrentUsage = 0 - cachedBudget.LastReset = now + if clone.ResetDuration != "" { + if duration, err := configstoreTables.ParseDuration(clone.ResetDuration); err == nil { + if now.Sub(clone.LastReset).Round(time.Millisecond) >= duration { + clone.CurrentUsage = 0 + clone.LastReset = now } } } - clone := *cachedBudget + clone.CurrentUsage += cost gs.budgets.Store(*mc.BudgetID, &clone) } }
890-921: Critical data race: modifying shared rate limit before cloning (model-level).Lines 893-897 and 901-905 directly mutate
cachedRateLimitbefore creating a clone at line 909. This has the same data race asUpdateProviderRateLimitUsage.🔎 Recommended fix
Clone FIRST:
if cachedRateLimitValue, exists := gs.rateLimits.Load(*mc.RateLimitID); exists && cachedRateLimitValue != nil { if cachedRateLimit, ok := cachedRateLimitValue.(*configstoreTables.TableRateLimit); ok && cachedRateLimit != nil { + // Clone FIRST to avoid race conditions + clone := *cachedRateLimit + now := time.Now() + // Check if rate limit needs reset (in-memory check) - if cachedRateLimit.TokenResetDuration != nil { - if duration, err := configstoreTables.ParseDuration(*cachedRateLimit.TokenResetDuration); err == nil { - if now.Sub(cachedRateLimit.TokenLastReset).Round(time.Millisecond) >= duration { - cachedRateLimit.TokenCurrentUsage = 0 - cachedRateLimit.TokenLastReset = now + if clone.TokenResetDuration != nil { + if duration, err := configstoreTables.ParseDuration(*clone.TokenResetDuration); err == nil { + if now.Sub(clone.TokenLastReset).Round(time.Millisecond) >= duration { + clone.TokenCurrentUsage = 0 + clone.TokenLastReset = now } } } - if cachedRateLimit.RequestResetDuration != nil { - if duration, err := configstoreTables.ParseDuration(*cachedRateLimit.RequestResetDuration); err == nil { - if now.Sub(cachedRateLimit.RequestLastReset).Round(time.Millisecond) >= duration { - cachedRateLimit.RequestCurrentUsage = 0 - cachedRateLimit.RequestLastReset = now + if clone.RequestResetDuration != nil { + if duration, err := configstoreTables.ParseDuration(*clone.RequestResetDuration); err == nil { + if now.Sub(clone.RequestLastReset).Round(time.Millisecond) >= duration { + clone.RequestCurrentUsage = 0 + clone.RequestLastReset = now } } } - clone := *cachedRateLimit + if shouldUpdateTokens { clone.TokenCurrentUsage += tokensUsed } if shouldUpdateRequests { clone.RequestCurrentUsage += 1 } gs.rateLimits.Store(*mc.RateLimitID, &clone) } }
🧹 Nitpick comments (1)
plugins/governance/main.go (1)
620-620: Clarify unused isCacheRead and isBatch parameters.The
postHookWorkersignature shows two unnamed parameters (_, _) at positions whereisCacheReadandisBatchwere previously used (based on lines 565-576 where these values are extracted from context). This suggests these parameters are no longer needed in the function body.Consider removing these unused parameters entirely from the signature rather than keeping them as unnamed placeholders, which would make the API cleaner and the intent more explicit.
🔎 Proposed refactor to remove unused parameters
-func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, _, _, isFinalChunk bool) { +func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provider schemas.ModelProvider, model string, requestType schemas.RequestType, virtualKey, requestID string, isFinalChunk bool) {And update the call site on line 586:
- p.postHookWorker(result, provider, model, requestType, virtualKey, requestID, isCacheRead, isBatch, bifrost.IsFinalChunk(ctx)) + p.postHookWorker(result, provider, model, requestType, virtualKey, requestID, bifrost.IsFinalChunk(ctx))
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/migrations.goplugins/governance/resolver.goframework/configstore/store.goframework/configstore/rdb.goframework/configstore/tables/modelconfig.goplugins/governance/main.goframework/configstore/tables/provider.goplugins/governance/store.goplugins/governance/tracker.goframework/configstore/clientconfig.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/migrations.goplugins/governance/resolver.goframework/configstore/store.goframework/configstore/rdb.goframework/configstore/tables/modelconfig.goplugins/governance/main.goframework/configstore/tables/provider.goplugins/governance/store.goplugins/governance/tracker.goframework/configstore/clientconfig.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/resolver.goplugins/governance/main.goplugins/governance/store.goplugins/governance/tracker.go
🧬 Code graph analysis (9)
framework/configstore/migrations.go (3)
framework/migrator/migrator.go (3)
New(131-149)DefaultOptions(100-106)Migration(62-69)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
framework/configstore/rdb.go (6)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(720-729)AuthConfig(710-715)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)ui/lib/types/config.ts (1)
AuthConfig(259-264)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
framework/configstore/tables/provider.go (3)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
plugins/governance/tracker.go (2)
core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)
framework/configstore/clientconfig.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
98c8540 to
de01994
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
plugins/governance/store.go (1)
1416-1468: Optional: Consider optimizing nested loops for relationship population.The nested loops for populating Budget and RateLimit relationships (lines 1417-1441 for model configs, 1444-1468 for providers) are O(n*m). For typical governance setups this is acceptable, but if the number of budgets/rate limits grows large, consider building index maps first (e.g.,
budgetsByID := make(map[string]*configstoreTables.TableBudget)) and using direct lookups.🔎 Example optimization pattern
+ // Build index maps for faster lookups + budgetsByID := make(map[string]*configstoreTables.TableBudget) + for i := range budgets { + budgetsByID[budgets[i].ID] = &budgets[i] + } + rateLimitsByID := make(map[string]*configstoreTables.TableRateLimit) + for i := range rateLimits { + rateLimitsByID[rateLimits[i].ID] = &rateLimits[i] + } + // Populate model configs with their relationships (Budget and RateLimit) for i := range modelConfigs { mc := &modelConfigs[i] // Populate budget if mc.BudgetID != nil { - for j := range budgets { - if budgets[j].ID == *mc.BudgetID { - mc.Budget = &budgets[j] - break - } - } + if budget, exists := budgetsByID[*mc.BudgetID]; exists { + mc.Budget = budget + } } // Populate rate limit if mc.RateLimitID != nil { - for j := range rateLimits { - if rateLimits[j].ID == *mc.RateLimitID { - mc.RateLimit = &rateLimits[j] - break - } - } + if rateLimit, exists := rateLimitsByID[*mc.RateLimitID]; exists { + mc.RateLimit = rateLimit + } } modelConfigs[i] = *mc }Apply similar pattern to provider relationship population (lines 1444-1468).
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (2)
- framework/configstore/tables/provider.go
- framework/configstore/migrations.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goplugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goframework/configstore/rdb.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goplugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goframework/configstore/rdb.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.go
🧬 Code graph analysis (6)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
framework/configstore/clientconfig.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
plugins/governance/resolver.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)
framework/configstore/rdb.go (3)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(720-729)AuthConfig(710-715)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (20)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (8)
framework/configstore/store.go (1)
36-36: LGTM! New interface methods follow existing patterns.The new GetProviders and ModelConfig CRUD methods are well-designed and consistent with the existing governance interface methods. The signatures properly support optional transaction parameters and follow the established naming conventions.
Also applies to: 118-125
framework/configstore/tables/modelconfig.go (1)
11-59: LGTM! Well-designed model with proper validation.The TableModelConfig struct is well-designed with:
- Appropriate use of nullable Provider field to support both global and provider-specific model configs
- Unique index on (ModelName, Provider) preventing duplicate configurations
- Defensive BeforeSave validation using strings.TrimSpace to catch empty/whitespace-only values
- CASCADE deletion for related budgets and rate limits
plugins/governance/tracker.go (1)
81-122: LGTM! Multi-stage update sequence properly implements hierarchical governance.The new update order (provider → model → VK) correctly implements the hierarchical governance model. The error handling appropriately logs failures but continues with subsequent updates, which is suitable for usage tracking where availability is more important than strict consistency.
Note: Updates are non-transactional, meaning partial failures are possible (e.g., provider update succeeds but model update fails). This is acceptable for usage tracking as it's designed to be best-effort and non-blocking.
plugins/governance/resolver.go (1)
77-213: LGTM! Clean refactoring implements hierarchical governance evaluation.The split into EvaluateModelAndProviderRequest and EvaluateVirtualKeyRequest properly implements the new hierarchical governance model:
- Provider-level checks execute first
- Model-level checks execute second
- VK-specific checks execute last (only if VK exists)
This separation makes the evaluation flow explicit and maintainable while preserving all existing validation logic.
plugins/governance/main.go (1)
473-481: LGTM! Properly implements two-stage governance evaluation.The PreHook correctly orchestrates the new evaluation flow:
- First evaluates model/provider constraints (applies to all requests)
- Then evaluates VK constraints if VK exists and first stage passed
The PostHook appropriately schedules usage tracking when a model is present, regardless of VK presence. The tracker handles empty VKs gracefully by updating only provider/model usage in that case.
Also applies to: 578-588
framework/configstore/rdb.go (2)
701-711: LGTM! RDB implementations follow established patterns.The new data access methods are well-implemented:
- GetProviders properly preloads Budget and RateLimit relationships
- ModelConfig CRUD methods follow existing conventions for transaction handling, error handling, and preloading
- GetModelConfig correctly handles the nullable Provider field with appropriate WHERE clauses
All implementations are consistent with the existing codebase patterns.
Also applies to: 1930-2025
2034-2099: LGTM! GetGovernanceConfig properly extended for new governance surface.The method correctly fetches and returns the new ModelConfigs and Providers fields as part of the governance configuration. The presence check on line 2064 is appropriately updated to include these new fields.
framework/configstore/clientconfig.go (1)
720-729: LGTM! GovernanceConfig appropriately extended.The addition of ModelConfigs and Providers fields properly extends the governance configuration structure to support the new per-model and per-provider governance features. The fields are correctly typed and annotated for JSON serialization.
| } | ||
|
|
||
| // 4. Check rate limits hierarchy (Provider level first, then VK level) | ||
| if rateLimitResult := r.checkRateLimitHierarchy(ctx, vk, string(evaluationRequest.Provider), evaluationRequest.Model, evaluationRequest.RequestID); rateLimitResult != nil { |
There was a problem hiding this comment.
lets have evaluationRequest here in checkRateLimitHierarchy as well so there's symmetry between this and checkBudgetHierarchy
There was a problem hiding this comment.
i virtually recreate the evaluationRequest for checkBudgetHierarchy, so thought might as well just directly pass it in.
since we've split EvaluateRequest into EvaluateModelAndProviderRequest and EvaluateVirtualKeyRequest we no longer take the whole evaluationRequest and only take the required params, i.e., provider, model and virtualKeyValue along with the reqId
There was a problem hiding this comment.
this is also done now, you can check and resolve
plugins/governance/store.go
Outdated
| // Check if rate limit needs reset (in-memory check) - operate on clone | ||
| if clone.TokenResetDuration != nil { | ||
| if duration, err := configstoreTables.ParseDuration(*clone.TokenResetDuration); err == nil { | ||
| if now.Sub(clone.TokenLastReset).Round(time.Millisecond) >= duration { |
There was a problem hiding this comment.
can remove .Round(time.Millisecond)
There was a problem hiding this comment.
this is done now. you can check and resolve
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
f6b85fa to
654de5d
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
de01994 to
4b1d30d
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
plugins/governance/tracker.go (1)
72-110: Align model-level usage tracking with resolver's provider pointer semantics.Provider-level sequencing looks good, but model-level calls currently pass
&update.Providerunconditionally. Whenupdate.Provideris empty, this passes a pointer to an empty string rather than nil, causing the store to attempt a lookup for keys like"model:"which will fail:if err := t.store.UpdateModelRateLimitUsage(ctx, update.Model, &update.Provider, update.TokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { ... } if shouldUpdateBudget && update.Cost > 0 { if err := t.store.UpdateModelBudgetUsage(ctx, update.Model, &update.Provider, update.Cost); err != nil { ... } }The resolver's
EvaluateModelAndProviderRequestintentionally buildsproviderPtronly whenprovider != ""and passes nil otherwise. While the store's fallback logic ensures model-only configs are still checked, the inconsistency causes unnecessary failed lookups and conceptual confusion.Mirror the resolver pattern:
var providerPtr *schemas.ModelProvider if update.Provider != "" { providerPtr = &update.Provider } if err := t.store.UpdateModelRateLimitUsage(ctx, update.Model, providerPtr, update.TokensUsed, shouldUpdateTokens, shouldUpdateRequests); err != nil { ... } if shouldUpdateBudget && update.Cost > 0 { if err := t.store.UpdateModelBudgetUsage(ctx, update.Model, providerPtr, update.Cost); err != nil { ... } }plugins/governance/store.go (3)
2167-2228: Missing budget reference updates for model configs and providers.The
updateBudgetReferencesmethod updates budget references in VKs, teams, and customers, but it doesn't update model configs and providers that may also reference the reset budget. This causes stale budget references in themodelConfigsandproviderssync.Maps when budgets are reset.🔎 Proposed fix
Add the following after line 2227 to update model configs and providers:
return true // continue }) + + // Update model configs that reference this budget + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.BudgetID != nil && *mc.BudgetID == budgetID { + clone := *mc + clone.Budget = resetBudget + gs.modelConfigs.Store(key, &clone) + } + return true // continue + }) + + // Update providers that reference this budget + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.BudgetID != nil && *provider.BudgetID == budgetID { + clone := *provider + clone.Budget = resetBudget + gs.providers.Store(key, &clone) + } + return true // continue + }) }
2230-2263: Missing rate limit reference updates for model configs and providers.The
updateRateLimitReferencesmethod updates rate limit references in VKs and their provider configs, but it doesn't update model configs and providers that may also reference the reset rate limit. This causes stale rate limit references in themodelConfigsandproviderssync.Maps when rate limits are reset.🔎 Proposed fix
Add the following after line 2262 to update model configs and providers:
return true // continue }) + + // Update model configs that reference this rate limit + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.RateLimitID != nil && *mc.RateLimitID == rateLimitID { + clone := *mc + clone.RateLimit = resetRateLimit + gs.modelConfigs.Store(key, &clone) + } + return true // continue + }) + + // Update providers that reference this rate limit + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.RateLimitID != nil && *provider.RateLimitID == rateLimitID { + clone := *provider + clone.RateLimit = resetRateLimit + gs.providers.Store(key, &clone) + } + return true // continue + }) }
1184-1267: Model and provider rate limit usage not dumped to database.The
DumpRateLimitsmethod only collects rate limit IDs from virtual keys (lines 1186-1202), excluding model config and provider rate limits. Since model and provider rate limits are updated in memory viaUpdateModelRateLimitUsageandUpdateProviderRateLimitUsage, their usage data will never be persisted to the database. This causes data loss on restart and prevents usage sync across nodes in multi-node deployments.🔎 Proposed fix
Extend the rate limit ID collection to include model configs and providers:
// Collect unique rate limit IDs from virtual keys rateLimitIDs := make(map[string]bool) gs.virtualKeys.Range(func(key, value interface{}) bool { vk, ok := value.(*configstoreTables.TableVirtualKey) if !ok || vk == nil { return true // continue } if vk.RateLimitID != nil { rateLimitIDs[*vk.RateLimitID] = true } if vk.ProviderConfigs != nil { for _, pc := range vk.ProviderConfigs { if pc.RateLimitID != nil { rateLimitIDs[*pc.RateLimitID] = true } } } return true // continue }) + + // Collect rate limit IDs from model configs + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.RateLimitID != nil { + rateLimitIDs[*mc.RateLimitID] = true + } + return true // continue + }) + + // Collect rate limit IDs from providers + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.RateLimitID != nil { + rateLimitIDs[*provider.RateLimitID] = true + } + return true // continue + })
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (2)
- framework/configstore/store.go
- framework/configstore/tables/modelconfig.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/migrations.goplugins/governance/tracker.goframework/configstore/tables/provider.goframework/configstore/clientconfig.goplugins/governance/resolver.goplugins/governance/store.goframework/configstore/rdb.goplugins/governance/main.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/migrations.goplugins/governance/tracker.goframework/configstore/tables/provider.goframework/configstore/clientconfig.goplugins/governance/resolver.goplugins/governance/store.goframework/configstore/rdb.goplugins/governance/main.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/tracker.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/main.go
🧬 Code graph analysis (6)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
framework/configstore/tables/provider.go (2)
framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
framework/configstore/clientconfig.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)
framework/configstore/rdb.go (4)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(720-729)AuthConfig(710-715)ui/lib/types/config.ts (1)
AuthConfig(259-264)
plugins/governance/main.go (2)
plugins/governance/resolver.go (1)
Decision(15-15)core/utils.go (1)
IsFinalChunk(212-227)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (24)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (6)
framework/configstore/migrations.go (1)
128-133: Model-config table and provider governance migration look consistent with existing patterns.The new migrations are wired into
triggerMigrationsin order, create/drop thegovernance_model_configstable safely, and add/removebudget_id/rate_limit_idcolumns plus indexes onconfig_providerswith guarded Rollback logic. This aligns with the existing migration style in this file and should behave correctly on both fresh and existing schemas.Also applies to: 2090-2118, 2120-2197
framework/configstore/clientconfig.go (1)
720-728: GovernanceConfig extension is coherent with the new governance surface.Adding
ModelConfigsandProvidershere matches the new RDB getters and migrations and is backward compatible for JSON clients that ignore unknown fields.framework/configstore/tables/provider.go (1)
41-48: Provider-level governance fields and validation are well-shaped.The new
BudgetID/RateLimitIDfields, indexes, and relationships line up with the migrations and RDB getters, and theBeforeSavechecks correctly prevent empty-string IDs when the pointers are set, avoiding subtly broken FKs.Also applies to: 91-97
plugins/governance/main.go (1)
456-480: Pre/Post hook restructuring cleanly separates model/provider and VK responsibilities.
PreHooknow evaluates provider/model governance viaEvaluateModelAndProviderRequestbefore (optionally) layeringEvaluateVirtualKeyRequestwhen a VK is present, which matches the new hierarchical design.PostHookandpostHookWorkercorrectly treat the virtual key as optional for usage tracking, only requiring a non-emptymodelto spawn the worker and passing the (possibly empty) VK through so the tracker can decide what to update.This wiring looks consistent with the resolver and tracker changes.
Also applies to: 578-587, 606-671
plugins/governance/resolver.go (1)
77-125: Resolver split into model/provider and VK paths is logically sound and matches the hierarchy.The new
EvaluateModelAndProviderRequestenforces provider-level rate limits and budgets before model-level checks, using aproviderPtronly when a provider is present, and returns a clearDecisionAllowresult on success.EvaluateVirtualKeyRequestthen reuses the explicitproviderandmodelfor VK-level filtering, rate-limit hierarchy, and budget checks via a locally constructedEvaluationRequest. This keeps provider/model and VK concerns cleanly separated and consistent with how PreHook and the tracker now call into the resolver.Also applies to: 128-191, 250-289
framework/configstore/rdb.go (1)
701-711: Provider and model-config persistence APIs integrate cleanly with existing governance storage.
GetProvidersmirrors other list getters and preloadsBudget/RateLimit, which fits the new provider-level governance fields.- The model-config CRUD methods (
GetModelConfigs,GetModelConfig,GetModelConfigByID,Create/Update/UpdateModelConfigs,DeleteModelConfig) follow the same transaction and error-handling conventions as budgets/rate limits;GetModelConfig’sprovider == nilvsprovider != nilbranching correctly distinguishes global vs provider-specific configs at the SQL level.GetGovernanceConfignow includesmodelConfigsandprovidersin both the emptiness check and the returnedGovernanceConfig, so installations that only configure model/provider-level governance still get a non-nil config.Everything here is consistent with the broader governance model introduced in this stack.
Also applies to: 1930-2025, 2034-2058, 2064-2099
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (2)
- framework/configstore/migrations.go
- framework/configstore/tables/provider.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/clientconfig.goplugins/governance/resolver.goplugins/governance/tracker.goframework/configstore/tables/modelconfig.goframework/configstore/rdb.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/clientconfig.goplugins/governance/resolver.goplugins/governance/tracker.goframework/configstore/tables/modelconfig.goframework/configstore/rdb.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/resolver.goplugins/governance/tracker.goplugins/governance/main.goplugins/governance/store.go
🧬 Code graph analysis (6)
framework/configstore/clientconfig.go (8)
framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/team.go (2)
TableTeam(12-38)TableTeam(41-41)framework/configstore/tables/customer.go (2)
TableCustomer(6-22)TableCustomer(25-25)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)ui/app/workspace/providers/page.tsx (1)
Providers(29-248)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)
plugins/governance/tracker.go (3)
core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)ui/lib/types/governance.ts (1)
VirtualKey(61-80)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
plugins/governance/store.go (6)
framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/utils.go (1)
ParseDuration(9-43)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (33)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (16)
framework/configstore/clientconfig.go (1)
720-729: LGTM! GovernanceConfig extended to support provider/model-level governance.The new
ModelConfigsandProvidersfields integrate cleanly with the existing governance data model and align with the related ORM changes in tables/modelconfig.go and tables/provider.go.framework/configstore/tables/modelconfig.go (2)
11-29: LGTM! TableModelConfig structure is well-designed.The unique index on
(ModelName, Provider)correctly enforces that each model-provider combination is configured once, while allowing provider-agnostic configs when Provider is NULL.
36-59: LGTM! Validation logic is comprehensive and consistent.The BeforeSave hook properly validates that all optional fields contain meaningful values (non-empty after trimming) if they are set, preventing database inconsistencies.
plugins/governance/tracker.go (2)
81-109: LGTM! Hierarchical usage tracking correctly implements provider → model → VK order.The new update sequence properly decouples provider/model-level tracking from virtual key presence, enabling governance even when VKs are disabled or absent. Error handling is resilient, logging failures but continuing with subsequent updates.
111-122: LGTM! Virtual key handling is now properly optional.The deferred VK lookup and graceful handling of empty/missing VKs correctly supports the new VK-agnostic governance model while preserving VK-level tracking when a VK is present.
framework/configstore/store.go (1)
36-36: LGTM! New CRUD methods follow established patterns.The GetProviders and model config methods integrate cleanly with the existing ConfigStore interface. The GetModelConfig signature correctly uses
provider *stringto query the composite unique key (model_name, provider).Also applies to: 118-125
plugins/governance/main.go (3)
456-481: LGTM! Two-step evaluation correctly implements hierarchical governance.The PreHook now properly evaluates model/provider constraints first (universally), then conditionally evaluates VK-specific constraints. This aligns with the PR's objective of adding per-model and per-provider governance that applies even without virtual keys.
578-588: LGTM! Usage tracking correctly handles optional virtual keys.The goroutine dispatches tracking when a model is present, passing the virtualKey (which may be empty) to the tracker. The tracker's UpdateUsage method (reviewed earlier) properly handles empty VKs by only updating provider/model-level usage, aligning with the new VK-agnostic governance model.
607-672: LGTM! postHookWorker correctly supports VK-optional tracking.The worker properly constructs UsageUpdate with the virtualKey field (which may be empty) and delegates to the tracker, which handles empty VKs by updating only provider/model-level usage. Comments accurately describe this behavior.
plugins/governance/resolver.go (2)
77-125: LGTM! Provider/model evaluation correctly implements hierarchical checks.The method properly evaluates provider-level constraints before model-level constraints, with clear early returns on violations. The conditional checks (
provider != "",model != "") correctly handle cases where either may be absent.
127-213: LGTM! Virtual key evaluation properly validates and applies VK-specific constraints.The method correctly validates VK existence/status, applies provider/model filtering using explicit parameters, checks VK-level rate limits and budgets, and computes include-only-keys for the matching provider. The logic is sound and well-structured.
framework/configstore/rdb.go (3)
701-711: LGTM! GetProviders correctly preloads governance relationships.The method properly retrieves providers with their associated budgets and rate limits, following the same pattern as other governance entity retrieval methods.
1930-2025: LGTM! Model config CRUD methods are well-implemented.All methods follow established patterns with proper preloading, error handling, and transaction support. The GetModelConfig method (lines 1942-1958) correctly handles the composite unique key by querying with both modelName and provider (including NULL handling).
2027-2101: LGTM! GetGovernanceConfig correctly integrates model configs and providers.The method properly fetches, checks presence, and constructs GovernanceConfig with the new ModelConfigs and Providers fields, maintaining consistency with existing entity handling.
plugins/governance/store.go (2)
735-772: LGTM: Race condition fixes properly applied.All the new Update* methods (UpdateProviderBudgetUsage, UpdateProviderRateLimitUsage, UpdateModelBudgetUsage, UpdateModelRateLimitUsage) correctly follow the clone-first pattern:
- Clone immediately after loading from sync.Map
- Perform all mutations on the clone
- Store the clone back
This prevents the race conditions that were flagged in previous reviews.
Also applies to: 774-824, 826-875, 877-939
297-402: LGTM: Rate limit expiry checks correctly use per-limit flags.Both CheckProviderRateLimit and CheckModelRateLimit properly track token and request expiry independently (lines 338-355, 536-553) and skip only the specific expired check (lines 368, 379, 566, 577). This ensures that if only one limit type is expired, the other is still properly evaluated.
This addresses the early return issue flagged in previous reviews and matches the pattern used in VK-level CheckRateLimit.
Also applies to: 479-601
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
4b1d30d to
8171345
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
plugins/governance/store.go (2)
2264-2325: Add reference updates for model configs and providers in budget reset.The
updateBudgetReferencesfunction updates VKs, teams, and customers when a budget is reset, but does not update model configs or providers that reference the same budget. This means model configs and providers may hold stale budget pointers after a reset.Based on learnings, budgets have 1:1 relationships with their parent entities, but the reset updates should propagate to all entities holding references to ensure consistency.
🔎 Add model config and provider updates
After updating customers (line 2324), add:
+ // Update model configs that reference this budget + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.BudgetID != nil && *mc.BudgetID == budgetID { + clone := *mc + clone.Budget = resetBudget + gs.modelConfigs.Store(key, &clone) + } + return true // continue + }) + + // Update providers that reference this budget + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.BudgetID != nil && *provider.BudgetID == budgetID { + clone := *provider + clone.Budget = resetBudget + gs.providers.Store(key, &clone) + } + return true // continue + })
2327-2360: Add reference updates for model configs and providers in rate limit reset.The
updateRateLimitReferencesfunction updates VKs when a rate limit is reset, but does not update model configs or providers. This creates the same stale pointer issue as with budgets.🔎 Add model config and provider updates
After updating VKs (line 2359), add:
+ // Update model configs that reference this rate limit + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.RateLimitID != nil && *mc.RateLimitID == rateLimitID { + clone := *mc + clone.RateLimit = resetRateLimit + gs.modelConfigs.Store(key, &clone) + } + return true // continue + }) + + // Update providers that reference this rate limit + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.RateLimitID != nil && *provider.RateLimitID == rateLimitID { + clone := *provider + clone.RateLimit = resetRateLimit + gs.providers.Store(key, &clone) + } + return true // continue + })
🧹 Nitpick comments (6)
framework/configstore/migrations.go (1)
2120-2197: Provider governance columns and indexes are correct; consider minor consistency cleanupThe migration adds
budget_id/rate_limit_idwith matching indexes and rollback drops the indexes before dropping the columns, which is aligned with how the virtual‑key provider config migration handles similar fields.Two small nits you might consider:
- Use either field names (
BudgetID,RateLimitID) or column names (budget_id,rate_limit_id) consistently acrossHasColumn/AddColumn/DropColumnfor readability.- If you later decide to enforce DB‑level FK constraints between providers and budgets/rate limits (similar to VK relationships), this is the logical place to add
CreateConstraint/DropConstraintaround these columns.plugins/governance/main.go (1)
578-588: postHookWorker async usage tracking is fine; unused cache/batch flags could be cleaned upThe PostHook path:
- Always schedules usage tracking when
model != "", regardless of VK presence.- Passes
virtualKey(possibly empty) intopostHookWorker, which then hands an appropriateUsageUpdateto the tracker.Inside
postHookWorkerthe last three bools are now_, _, isFinalChunk, soisCacheReadandisBatchare intentionally unused. That’s correct if cache/batch no longer affect governance, but the parameter docs still describe them as meaningful.If you want to reduce confusion later, consider either:
- Dropping those parameters entirely (and updating the call site), or
- Updating the comment to state explicitly that cache/batch flags are currently ignored.
Also applies to: 608-620, 669-671
framework/configstore/store.go (1)
36-38: ConfigStore interface extensions match new governance needs; check other implementationsAdding:
- Provider accessors (
GetProviders,GetProviderByName),- Targeted usage updaters (
UpdateBudgetUsage,UpdateRateLimitUsage),- ModelConfig CRUD APIs
aligns well with the new provider/model-level governance features. Please just confirm any other
ConfigStoreimplementations (e.g., test fakes or non-RDB stores) have been updated to satisfy the expanded interface so builds don’t break elsewhere.Also applies to: 118-128
framework/configstore/rdb.go (2)
577-600: Cascade deletion of budgets/rate limits is consistent with 1:1 invariants; consider centralizing/transactionalizingThese blocks consistently:
- Capture associated
BudgetID/RateLimitIDbefore deleting the owning row (provider, VK, VK provider config, team, customer, model config).- Then delete the owner and finally delete the budget/rate-limit rows referenced.
This matches the invariant that budgets and rate limits are 1:1 with their parent entities and not shared, and it keeps the DB free of orphans. Based on learnings, that’s the intended model.
Two optional improvements you might consider over time:
- Wrap
DeleteProviderandDeleteVirtualKeyProviderConfigin explicit transactions like the VK/team/customer/model-config paths, so owner and child deletions are atomic.- Factor the repeated “capture FK IDs then delete budget/rate-limit” pattern into a small helper to reduce duplication and keep future changes (e.g., adding audit logging) consistent.
Also applies to: 1445-1483, 1590-1617, 1748-1761, 1840-1853, 2119-2152
718-740: Provider governance getters look good
GetProvidersandGetProviderByName:
- Preload
BudgetandRateLimit, which is exactly what the governance layer will need.- Return
ErrNotFoundonly for the single-row lookup, and an empty slice for the multi-row case.The extra
ErrRecordNotFoundcheck afterFindis redundant but harmless, sinceFindreturns nil error with an empty slice on no rows.plugins/governance/store.go (1)
1217-1283: Consider collecting budget IDs from entities for consistency with rate limits.
DumpRateLimits(lines 1094-1215) collects rate limit IDs from virtual keys, model configs, and providers before dumping, ensuring only actively-used rate limits are persisted. However,DumpBudgetsiterates over all budgets ings.budgetswithout filtering by entity references.While both approaches work, aligning
DumpBudgetsto collect budget IDs from entities (likeDumpRateLimitsdoes) would prevent dumping orphaned budgets and ensure consistency.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (3)
- framework/configstore/tables/provider.go
- framework/configstore/tables/modelconfig.go
- framework/configstore/clientconfig.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/tracker.goframework/configstore/migrations.goplugins/governance/resolver.goframework/configstore/rdb.goplugins/governance/main.goframework/configstore/store.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/tracker.goframework/configstore/migrations.goplugins/governance/resolver.goframework/configstore/rdb.goplugins/governance/main.goframework/configstore/store.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/tracker.goplugins/governance/resolver.goplugins/governance/main.goplugins/governance/store.go
🧬 Code graph analysis (6)
plugins/governance/tracker.go (3)
core/schemas/models.go (1)
Model(109-129)core/schemas/provider.go (1)
Provider(315-362)ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/resolver.go (3)
core/schemas/bifrost.go (1)
ModelProvider(32-32)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
plugins/governance/store.go (6)
framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/utils.go (1)
ParseDuration(9-43)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (23)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (6)
framework/configstore/migrations.go (1)
128-133: ModelConfig migration wiring looks consistent and safe
migrationAddModelConfigTableis registered late intriggerMigrations, after budgets and rate limits exist, and uses the sameCreateTable/DropTablepattern as existing migrations. No ordering or rollback issues stand out.Also applies to: 2090-2118
plugins/governance/tracker.go (1)
69-123: Usage update hierarchy (provider/model → VK) is well-structuredThe refactored
UpdateUsagenow:
- Updates provider- and model-level rate limits and budgets unconditionally (subject to streaming flags and
Cost > 0), regardless of VK presence.- Only then looks up and, if found, applies VK-level rate limit and budget updates.
- Logs and continues on provider/model update failures, keeping VK updates independent.
The comments about handling empty virtual keys (provider/model only) match the actual behavior.
Also applies to: 656-672
plugins/governance/main.go (1)
453-480: PreHook’s model/provider → VK evaluation pipeline is coherentExtracting
(provider, model, _)from the request and:
- First calling
EvaluateModelAndProviderRequest, then- Only if allowed and a VK is present, calling
EvaluateVirtualKeyRequestgives a clear, ordered governance flow and keeps VK-free requests governed at provider/model level only. The virtual-key-mandatory check remains orthogonal and still gates missing VKs when configured.
plugins/governance/resolver.go (1)
77-125: Resolver split between provider/model and VK paths is logically sound
EvaluateModelAndProviderRequestcleanly sequences:
- Provider rate limits, then provider budgets.
- Model rate limits, then model budgets (optionally scoped by provider).
EvaluateVirtualKeyRequest:
- Validates VK existence/activeness and enriches the context with VK/team/customer IDs.
- Applies provider and model filters using VK provider configs.
- Delegates VK-centric rate limits to
checkRateLimitHierarchyand VK→Team→Customer budgets tocheckBudgetHierarchy.checkRateLimitHierarchynow takes anEvaluationRequestand returns a structuredEvaluationResultwith the first matching rate-limit info.This lines up with the PreHook flow and the hierarchical governance design.
Also applies to: 179-195, 252-290
framework/configstore/rdb.go (2)
2000-2029: Targeted usage updates with SkipHooks are appropriate
UpdateBudgetUsageandUpdateRateLimitUsage:
- Use
Session(&gorm.Session{SkipHooks: true})to bypass validation hooks, which is correct when only updating usage counters.- Narrow the UPDATE to:
current_usagefor budgets.token_current_usage/request_current_usagefor rate limits.This minimizes the chance of clobbering config fields while still benefiting from
parseGormErrorfor uniform error handling.
2031-2152: ModelConfig CRUD and governance exposure are coherent
GetModelConfigs/GetModelConfig/GetModelConfigByID:
- Preload
BudgetandRateLimitso callers can inspect governance details.- For
GetModelConfig, the(model_name, provider)query correctly distinguishes provider-scoped configs vs global (provider IS NULL) configs.CreateModelConfig,UpdateModelConfig, andUpdateModelConfigsfollow the same patterns as existing budget/rate-limit CRUD, includingparseGormErrorusage.DeleteModelConfigmirrors the other cascade paths by deleting any associated budget/rate limit after removing the model config row.GetGovernanceConfignow also loadsmodelConfigsandprovidersand returns them via theGovernanceConfigstruct; the “all empty → nil” check was updated to include these slices.This gives the governance layer full visibility into model- and provider-level governance while keeping the RDB store behavior consistent with the rest of the API.
Also applies to: 2161-2185, 2219-2227
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
654de5d to
be49a44
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
8171345 to
360b6a9
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/configstore/store.go (1)
560-602: Be explicit about 1:1 ownership assumptions when deleting provider/model budgets and rate limits.
DeleteProviderandDeleteModelConfignow:
- capture
BudgetID/RateLimitID,- delete the parent row, and
- unconditionally delete the referenced budget and rate-limit rows.
This mirrors existing VK/team/customer behavior, but it relies on a strict 1:1 ownership invariant for budgets and rate limits. If a budget or rate limit were ever shared (e.g., reused across a VK and a provider, or across multiple model configs), these deletes would silently remove it for all parents.
If the design guarantee is indeed “budgets/rate limits are never shared across entities”, consider:
- documenting that invariant near these methods (and/or in the table structs), and
- adding tests that exercise delete flows with multiple potential references to ensure we don’t accidentally introduce sharing later.
Based on learnings, this keeps cascade deletes safe as the model evolves.
Also applies to: 2118-2151
🧹 Nitpick comments (3)
framework/configstore/migrations.go (1)
128-133: Align column/field naming for provider governance migration.The new migrations are wired correctly and rollback covers both indexes and columns. One small robustness nit: in
migrationAddProviderGovernanceColumnsyou mix DB column names and struct field names:
HasColumn(provider, "budget_id")/HasColumn(provider, "rate_limit_id")AddColumn(provider, "BudgetID")/DropColumn(provider, "RateLimitID")Elsewhere in this file you consistently use the same snake_case string for
HasColumn/AddColumn/DropColumn. It would be safer and more predictable to stick to one convention here as well (e.g., always"budget_id"/"rate_limit_id"), so future dialect or GORM changes don’t surprise you.This is non-blocking but worth tightening for consistency.
Also applies to: 2120-2197
plugins/governance/main.go (1)
578-588: Asynchronous usage tracking with optional VK is consistent and safe.The PostHook path now:
- skips tracking when
model == "",- always passes
virtualKey(possibly empty) intopostHookWorker, and- constructs a
UsageUpdatewhere VK-specific work is gated onVirtualKey != ""in the tracker.This matches the comments about handling empty VKs gracefully and keeps request processing non-blocking via the wait group. The two unused bools in
postHookWorker’s signature are fine as placeholders; you can drop them later if you decide they won’t come back.Also applies to: 606-620, 656-672
framework/configstore/rdb.go (1)
717-739: RDB implementations for providers/model-configs are consistent with the schema.The new
GetProviders/GetProviderByName, usage-only update helpers, and full ModelConfig CRUD are all idiomatic GORM and line up with theTableProviderandTableModelConfigdefinitions. IncludingModelConfigsandProvidersinGetGovernanceConfig’s payload cleanly exposes the new governance dimensions to callers.The extra ErrRecordNotFound checks after
Findcalls are unnecessary but benign; you can leave them as-is.Also applies to: 1999-2028, 2030-2058, 2060-2070, 2072-2151, 2153-2225
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (2)
- framework/configstore/clientconfig.go
- framework/configstore/tables/provider.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/tables/modelconfig.goplugins/governance/store.goplugins/governance/tracker.goframework/configstore/migrations.goframework/configstore/store.goframework/configstore/rdb.goplugins/governance/main.goplugins/governance/resolver.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/tables/modelconfig.goplugins/governance/store.goplugins/governance/tracker.goframework/configstore/migrations.goframework/configstore/store.goframework/configstore/rdb.goplugins/governance/main.goplugins/governance/resolver.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/store.goplugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.go
🧬 Code graph analysis (7)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
plugins/governance/tracker.go (1)
ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
framework/configstore/rdb.go (5)
framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(720-729)AuthConfig(710-715)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
plugins/governance/resolver.go (4)
core/schemas/bifrost.go (1)
ModelProvider(32-32)ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (10)
plugins/governance/main.go (1)
456-481: PreHook’s hierarchical evaluation order looks sound.Extracting
provider, modelviareq.GetRequestFields()and then:
- running
EvaluateModelAndProviderRequestunconditionally, and- only if that passes and a VK is present, running
EvaluateVirtualKeyRequestgives you a clean hierarchy (model/provider first, then VK) and keeps the decision mapping (403/429/402) centralized in one switch. This is a good fit with the new resolver APIs.
plugins/governance/tracker.go (1)
72-123: UsageUpdate flow correctly prioritizes provider/model before VK.The new
UpdateUsagephases:
- Provider/model rate-limit update
- Provider/model budget update (when
Cost > 0)- Optional VK lookup and VK-specific rate-limit/budget updates
ensure that non-VK governance data is always updated, and VK-specific work is strictly additive and isolated behind the VK lookup. Debug-only logging for VK-not-found and per-step failures keeps this non-disruptive.
This structure matches the intended “model/provider first, then VK” semantics and is a solid improvement.
framework/configstore/tables/modelconfig.go (1)
11-59: ModelConfig schema and validation look appropriate.The
TableModelConfigdefinition andBeforeSavehook provide the right invariants (non-empty model name; no empty-but-present IDs/providers) and align with the new migrations and CRUD methods. Relationships and indexes are also consistent with how budgets and rate limits are modeled elsewhere.plugins/governance/resolver.go (2)
77-125: Provider/model evaluation is well-factored and ordered correctly.
EvaluateModelAndProviderRequestchecking:
- provider rate-limit → provider budget
- model rate-limit → model budget
before returning
DecisionAllowgives a predictable, top‑down hierarchy and makes it easy to see where a denial originates (provider vs model). The error messages are specific enough to surface which layer failed.
128-193: VK evaluation and hierarchy checks are coherent and match the data model.
EvaluateVirtualKeyRequest:
- Validates VK existence/active state and stamps VK/team/customer context into
ctx.- Applies provider and model filters via
isProviderAllowed/isModelAllowed, correctly treating empty configs as “allow all”.- Uses
checkRateLimitHierarchyandcheckBudgetHierarchywith anEvaluationRequest, returning detailedEvaluationResults that carry VK and rate-limit/budget info where relevant.This structure keeps VK-specific policy clearly separated from provider/model policy while still exposing enough context for callers to understand denials.
Also applies to: 251-290, 294-363
framework/configstore/store.go (1)
36-38: Store surface and usage-only update helpers align with hierarchical governance.The ConfigStore interface extensions plus their RDB implementations cover:
- direct provider access (
GetProviders,GetProviderByName) with budgets and rate limits preloaded, and- narrow, hook-free
UpdateBudgetUsage/UpdateRateLimitUsagehelpers that only touch usage fields.This is a good fit for the in‑memory tracker’s periodic dump/reset logic and avoids accidental drift in config fields when persisting usage.
Also applies to: 118-128, 1999-2028, 2030-2058, 2060-2070, 2072-2151
plugins/governance/store.go (4)
242-405: LGTM! Provider-level governance checks implemented correctly.The CheckProviderBudget and CheckProviderRateLimit methods follow established patterns from VK-level checks, with proper nil handling, expiration checks, and the flag-based approach for independently tracking token vs request limit expiration. This addresses the early-return concerns from previous reviews.
407-604: LGTM! Model-level governance checks properly support both specific and global configurations.The CheckModelBudget and CheckModelRateLimit methods correctly implement hierarchical checking: model+provider configs (when provider is specified) followed by model-only global configs. The flag-based expiration tracking and nil provider handling are both correct.
738-863: LGTM! Usage update methods correctly implement clone-before-mutate pattern.The UpdateProviderAndModelBudgetUsageInMemory and UpdateProviderAndModelRateLimitUsageInMemory methods use inline helper functions to avoid duplication while maintaining the critical clone-before-mutate pattern that prevents race conditions. The hierarchical update order (provider → model+provider → model-only) is consistent with the check methods.
2110-2260: LGTM! Model config and provider CRUD operations properly implemented.The UpdateModelConfigInMemory, DeleteModelConfigInMemory, UpdateProviderInMemory, and DeleteProviderInMemory methods follow established patterns from VK/Team/Customer CRUD operations. The cascade deletion of associated budgets and rate limits aligns with the 1:1 relationship documented in learnings, and the key format logic is consistent throughout the codebase.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
plugins/governance/store.go (2)
2264-2325: Update model configs and providers that reference the reset budget.The
updateBudgetReferencesfunction updates VKs, teams, and customers that reference a reset budget, but does not update model configs or providers. Since model configs and providers can also have budgets (as seen inUpdateModelConfigInMemoryandUpdateProviderInMemory), they need to be updated as well to avoid stale references.🔎 Proposed fix to update model configs and providers
func (gs *LocalGovernanceStore) updateBudgetReferences(resetBudget *configstoreTables.TableBudget) { budgetID := resetBudget.ID // Update VKs that reference this budget gs.virtualKeys.Range(func(key, value interface{}) bool { vk, ok := value.(*configstoreTables.TableVirtualKey) if !ok || vk == nil { return true // continue } needsUpdate := false clone := *vk // Check VK-level budget if vk.BudgetID != nil && *vk.BudgetID == budgetID { clone.Budget = resetBudget needsUpdate = true } // Check provider config budgets if vk.ProviderConfigs != nil { for i, pc := range clone.ProviderConfigs { if pc.BudgetID != nil && *pc.BudgetID == budgetID { clone.ProviderConfigs[i].Budget = resetBudget needsUpdate = true } } } if needsUpdate { gs.virtualKeys.Store(key, &clone) } return true // continue }) // Update teams that reference this budget gs.teams.Range(func(key, value interface{}) bool { team, ok := value.(*configstoreTables.TableTeam) if !ok || team == nil { return true // continue } if team.BudgetID != nil && *team.BudgetID == budgetID { clone := *team clone.Budget = resetBudget gs.teams.Store(key, &clone) } return true // continue }) // Update customers that reference this budget gs.customers.Range(func(key, value interface{}) bool { customer, ok := value.(*configstoreTables.TableCustomer) if !ok || customer == nil { return true // continue } if customer.BudgetID != nil && *customer.BudgetID == budgetID { clone := *customer clone.Budget = resetBudget gs.customers.Store(key, &clone) } return true // continue }) + + // Update model configs that reference this budget + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.BudgetID != nil && *mc.BudgetID == budgetID { + clone := *mc + clone.Budget = resetBudget + gs.modelConfigs.Store(key, &clone) + } + return true // continue + }) + + // Update providers that reference this budget + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.BudgetID != nil && *provider.BudgetID == budgetID { + clone := *provider + clone.Budget = resetBudget + gs.providers.Store(key, &clone) + } + return true // continue + }) }
2327-2360: Update model configs and providers that reference the reset rate limit.The
updateRateLimitReferencesfunction updates VKs that reference a reset rate limit, but does not update model configs or providers. Since model configs and providers can also have rate limits (as seen inUpdateModelConfigInMemoryandUpdateProviderInMemory), they need to be updated as well to avoid stale references.🔎 Proposed fix to update model configs and providers
func (gs *LocalGovernanceStore) updateRateLimitReferences(resetRateLimit *configstoreTables.TableRateLimit) { rateLimitID := resetRateLimit.ID // Update VKs that reference this rate limit gs.virtualKeys.Range(func(key, value interface{}) bool { vk, ok := value.(*configstoreTables.TableVirtualKey) if !ok || vk == nil { return true // continue } needsUpdate := false clone := *vk // Check VK-level rate limit if vk.RateLimitID != nil && *vk.RateLimitID == rateLimitID { clone.RateLimit = resetRateLimit needsUpdate = true } // Check provider config rate limits if vk.ProviderConfigs != nil { for i, pc := range clone.ProviderConfigs { if pc.RateLimitID != nil && *pc.RateLimitID == rateLimitID { clone.ProviderConfigs[i].RateLimit = resetRateLimit needsUpdate = true } } } if needsUpdate { gs.virtualKeys.Store(key, &clone) } return true // continue }) + + // Update model configs that reference this rate limit + gs.modelConfigs.Range(func(key, value interface{}) bool { + mc, ok := value.(*configstoreTables.TableModelConfig) + if !ok || mc == nil { + return true // continue + } + if mc.RateLimitID != nil && *mc.RateLimitID == rateLimitID { + clone := *mc + clone.RateLimit = resetRateLimit + gs.modelConfigs.Store(key, &clone) + } + return true // continue + }) + + // Update providers that reference this rate limit + gs.providers.Range(func(key, value interface{}) bool { + provider, ok := value.(*configstoreTables.TableProvider) + if !ok || provider == nil { + return true // continue + } + if provider.RateLimitID != nil && *provider.RateLimitID == rateLimitID { + clone := *provider + clone.RateLimit = resetRateLimit + gs.providers.Store(key, &clone) + } + return true // continue + }) }
🧹 Nitpick comments (4)
framework/configstore/migrations.go (1)
2090-2197: ModelConfig table + provider governance columns migration follow existing patterns.The
governance_model_configstable creation and thebudget_id/rate_limit_idadditions onconfig_providers(with matching indexes and rollback) are consistent with earlier governance migrations and should play nicely with the stacked changes.framework/configstore/rdb.go (3)
560-601: Confirm provider budgets/rate limits remain 1:1 before hard-deleting them.
DeleteProvidernow deletes the provider’sBudgetIDandRateLimitIDafter removing the provider. This is correct only if those budget/rate‑limit rows are never shared with any other entity (including other providers or model configs). Based on learnings, budgets/rate limits are already modeled as 1:1 for VKs/teams/customers; please confirm the same invariant holds for provider‑level governance and keep tests enforcing that so we don’t accidentally delete shared limits in the future.
1999-2151: Usage-only updates and ModelConfig CRUD align with the new hierarchy.
UpdateBudgetUsage/UpdateRateLimitUsagecorrectly narrow updates to usage fields and skip hooks, which is appropriate for periodic dumps from the in‑memory tracker.- ModelConfig getters and CRUD (including cascading deletes of attached budget/rate‑limit rows) mirror the existing VK/team/customer patterns and should integrate cleanly with the governance store.
If you ever need stronger guarantees when an ID is wrong (e.g., treat “0 rows updated” as not found), you could add an explicit
RowsAffected == 0check in the usage updaters, but the current approach is acceptable if callers already validate IDs.
2117-2151: Same 1:1 invariant consideration for model-config budgets and rate limits.
DeleteModelConfigalso hard-deletes its associatedBudgetID/RateLimitIDafter removing the config. As with providers, this assumes those rows are not shared with any other entity. It would be good to explicitly document and test that ModelConfig budgets/rate‑limits are 1:1 with their owning config so future reuse/sharing doesn’t accidentally break this behavior.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (3)
- framework/configstore/tables/provider.go
- framework/configstore/clientconfig.go
- framework/configstore/tables/modelconfig.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/resolver.goplugins/governance/main.goplugins/governance/tracker.goframework/configstore/rdb.goframework/configstore/migrations.goframework/configstore/store.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/resolver.goplugins/governance/main.goplugins/governance/tracker.goframework/configstore/rdb.goframework/configstore/migrations.goframework/configstore/store.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/resolver.goplugins/governance/main.goplugins/governance/tracker.goplugins/governance/store.go
🧬 Code graph analysis (5)
plugins/governance/resolver.go (5)
core/schemas/bifrost.go (1)
ModelProvider(32-32)core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
plugins/governance/tracker.go (1)
ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
framework/configstore/store.go (1)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (10)
framework/configstore/store.go (2)
36-37: Provider accessors fit the new governance model surface.Adding
GetProviders/GetProviderByNamehere matches the RDB implementation and cleanly exposes provider-level governance metadata (including budgets/rate limits) to callers.
118-128: Usage-only updaters and ModelConfig CRUD are well-shaped for hierarchical governance.
UpdateBudgetUsage/UpdateRateLimitUsagegive narrow, usage-focused writes, and the ModelConfig CRUD methods mirror existing budget/rate‑limit/virtual‑key patterns, which should make wiring in the governance store straightforward.framework/configstore/migrations.go (1)
128-133: New migrations are registered in a safe, idempotent order.Hooking
migrationAddModelConfigTableandmigrationAddProviderGovernanceColumnsat the end oftriggerMigrationsmatches the existing sequencing and avoids interfering with older schema steps.plugins/governance/tracker.go (1)
69-123: Three‑phase provider/model → VK usage update flow looks solid.Updating provider/model rate limits and budgets unconditionally, then layering VK rate‑limit/budget updates only when a VK exists (and is found), matches the new governance hierarchy and preserves accounting even for non‑VK traffic.
framework/configstore/rdb.go (2)
717-739: Provider queries with governance preloads are consistent with the new interface.
GetProviders/GetProviderByNamepreloadingBudgetandRateLimitgive callers everything needed for provider‑level governance decisions without extra round trips.
2160-2225: Extending GovernanceConfig to include ModelConfigs and Providers is coherent.Fetching
modelConfigsandprovidersalongside virtual keys, teams, customers, budgets, and rate limits, and exposing them viaGovernanceConfiggives the governance plugin a complete snapshot of all hierarchy layers, which matches the stacked resolver/tracker changes.plugins/governance/main.go (3)
456-482: PreHook’s split between model/provider and VK evaluation is well-structured.Pulling
provider, modelviaGetRequestFields, runningEvaluateModelAndProviderRequestfirst, and only then layeringEvaluateVirtualKeyRequestwhen a VK is present (and non-mandatory VKs are respected) gives a clear hierarchy and keeps VK logic from masking provider/model-level rejections.
578-588: Usage tracking trigger matches the tracker’s expectations.Conditioning the
postHookWorkergoroutine onmodel != ""while always passing throughvirtualKey(possibly empty) lines up with the tracker’s behavior: provider/model usage is still updated even when VK is absent, and the earlier comments about VK being optional refer specifically to VK presence, not the model guard.
606-672: postHookWorker’s handling of empty virtual keys and streaming is consistent.The worker now explicitly documents and supports empty
virtualKeyby only feeding that through toUsageUpdate, leaving the tracker to apply provider/model-only accounting when VK is missing, and it correctly defers cost and token aggregation to the final streaming chunk.plugins/governance/resolver.go (1)
77-290: Hierarchical resolver flow (provider/model → VK) is coherent with the plugin hooks.
EvaluateModelAndProviderRequestcleanly applies provider-level then model-level rate‑limit and budget checks, independent of VK presence.EvaluateVirtualKeyRequestthen validates VK, enforces provider/model allowlists, and delegates VK rate‑limit and hierarchical budget checks viacheckRateLimitHierarchy/checkBudgetHierarchyusing a sharedEvaluationRequest.- Surfacing
RateLimitInfofrom matching provider configs (falling back to VK-level) gives useful context to callers without complicating the main decision API.Overall this matches how PreHook now sequences model/provider and VK decisions.
akshaydeo
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
2 times, most recently
from
ee064c2 to
51556a1
Compare
akshaydeo
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
360b6a9 to
c33348c
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
framework/configstore/tables/provider.go (1)
41-47: Provider-level governance fields and validation align with 1:1 budget/rate-limit design.BudgetID/RateLimitID plus the CASCADE relations and non-empty checks are consistent with treating provider budgets/rate limits as dedicated children of a provider. If you expect user input here, you might optionally mirror
TableModelConfigand usestrings.TrimSpacein the validations to reject whitespace-only IDs, but that’s not critical.Based on learnings, budgets/rate limits are intended to be 1:1 with their parent entity.
Also applies to: 91-97
framework/configstore/migrations.go (1)
128-133: New migrations for model configs and provider governance columns are well-shaped and rollback-safe.
migrationAddModelConfigTablecleanly creates/dropsgovernance_model_configs.migrationAddProviderGovernanceColumnscorrectly guardsAddColumn/index creation withHasColumn/HasIndexand mirrors this in rollback.If you later need DB-level referential guarantees like
migrationAddProviderConfigBudgetRateLimitprovides, you could addCreateConstraintcalls for the provider→budget/rate-limit FKs, but current behavior is consistent with the rest of the file.Also applies to: 2090-2197
plugins/governance/resolver.go (1)
127-214: Virtual-key evaluation and hierarchy checks integrate cleanly with the new model/provider path.
EvaluateVirtualKeyRequest:
- Validates VK existence/active state and decorates context with VK/team/customer IDs.
- Enforces provider and model allowlists via
isProviderAllowed/isModelAllowed.- Delegates VK rate-limit and budget decisions to
checkRateLimitHierarchyandcheckBudgetHierarchy.- Sets include-only-keys for the matched provider config.
checkRateLimitHierarchy/checkBudgetHierarchysignatures are now symmetric viaEvaluationRequest, which keeps parameter passing tidy.Only tiny nit: inside
checkRateLimitHierarchyyou can drop theschemas.ModelProvider(request.Provider)cast sincerequest.Provideris already of that type, but that’s purely cosmetic.Also applies to: 251-290
framework/configstore/rdb.go (1)
577-599: Consider wrapping provider + governance deletions in a single transaction.
DeleteProvidernow:
- Reads
BudgetID/RateLimitIDfrom the provider.- Deletes the provider row.
- Then deletes the associated budget and rate-limit rows.
Unlike
DeleteVirtualKey,DeleteTeam,DeleteCustomer, andDeleteModelConfig, this flow is not enclosed in a transaction when notxis passed, so a failure in the second phase could leave orphaned budgets or rate limits. Wrapping the whole sequence in a transaction (for the default path) would make provider deletion atomic and consistent with the other governance delete helpers.Based on learnings, budgets and rate limits are intended to be 1:1 with their parent entities, so keeping these deletes atomic helps preserve that invariant.
plugins/governance/store.go (1)
2149-2157: Document key immutability or handle key changes in UpdateModelConfigInMemory.
UpdateModelConfigInMemorystores model configs under composite keys (modelNameormodelName:provider). IfModelNameorProviderchanges between updates, the method creates a new map entry but leaves the old key in place, potentially causing stale entries.For example:
- Initial:
"gpt-4:openai" -> {ID: "mc1", ModelName: "gpt-4", Provider: "openai"}- After update: Both
"gpt-4:openai"(stale) and"gpt-4:anthropic"existDeleteModelConfigInMemory("mc1")would only remove one entryConsider one of the following approaches:
- Document: Add a comment stating that
ModelNameandProviderare immutable identifiers; changes requireDelete+Createflow- Validate: Add a check that rejects updates where
ModelNameorProviderdiffers from the existing entry- Auto-cleanup: Before storing under the new key, iterate to find and delete any existing entry with the same ID
🔎 Example: Document the immutability assumption
// UpdateModelConfigInMemory adds or updates a model config in the in-memory store (lock-free) // Preserves existing usage values when updating budgets and rate limits // Returns the updated model config with potentially modified usage values +// Note: ModelName and Provider are treated as immutable identifiers. To change either field, +// use DeleteModelConfigInMemory followed by UpdateModelConfigInMemory with the new values. func (gs *LocalGovernanceStore) UpdateModelConfigInMemory(mc *configstoreTables.TableModelConfig) *configstoreTables.TableModelConfig {
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (1)
- framework/configstore/clientconfig.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
framework/configstore/migrations.goframework/configstore/tables/modelconfig.goplugins/governance/resolver.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/tracker.goframework/configstore/store.goframework/configstore/rdb.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/configstore/migrations.goframework/configstore/tables/modelconfig.goplugins/governance/resolver.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/tracker.goframework/configstore/store.goframework/configstore/rdb.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/resolver.goplugins/governance/main.goplugins/governance/tracker.goplugins/governance/store.go
🧬 Code graph analysis (7)
framework/configstore/migrations.go (3)
framework/migrator/migrator.go (3)
New(131-149)DefaultOptions(100-106)Migration(62-69)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
framework/configstore/tables/provider.go (3)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
plugins/governance/tracker.go (1)
ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)
framework/configstore/store.go (2)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)
framework/configstore/rdb.go (5)
framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(720-729)AuthConfig(710-715)ui/lib/types/config.ts (1)
AuthConfig(259-264)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (10)
plugins/governance/main.go (3)
456-482: PreHook ordering cleanly enforces model/provider before virtual-key governance.Using
req.GetRequestFields()and then callingEvaluateModelAndProviderRequestbeforeEvaluateVirtualKeyRequestgives a clear, hierarchical decision flow and ensures provider/model limits apply even when VKs are disabled or absent. Looks good.
578-588: PostHook usage tracking gating on non-empty model is consistent with tracker contract.Conditioning the goroutine on
model != ""while always passingvirtualKey(possibly empty) intopostHookWorkermatches the tracker’s behavior of updating provider/model usage unconditionally and VK usage only when present. Concurrency handling viawg.Add/Donealso looks correct.
608-671: postHookWorker usage aggregation and tracker call look correct.The worker correctly:
- Derives
success, streaming flags, tokens across response types, and cost viamodelCatalog.- Constructs
UsageUpdatewith VK, provider, model, and streaming metadata.- Delegates to
tracker.UpdateUsage(which already handles empty VKs and streaming semantics).No issues from the governance/usage perspective.
framework/configstore/tables/modelconfig.go (1)
11-29: Model-config ORM and validation match the intended governance semantics.The
TableModelConfigshape (unique model/provider pair, optional per-model budget/rate-limit) plus the BeforeSave checks for non-empty IDs/provider give you a solid base for per-model governance while preventing half-baked rows. This looks good.Also applies to: 36-59
plugins/governance/tracker.go (1)
69-123: Hierarchical usage tracking sequencing is correct and robust.UpdateUsage now:
- Always applies provider/model rate-limit and budget updates (subject to streaming flags) regardless of VK presence.
- Only performs VK-level rate-limit/budget updates when a VK string is provided and found.
- Treats missing VKs as a soft condition with a debug log.
This cleanly decouples provider/model governance from VK presence while preserving VK semantics.
plugins/governance/resolver.go (1)
77-125: Model/provider evaluation flow is clearly separated and ordered.
EvaluateModelAndProviderRequestcleanly enforces provider rate-limit/budget first, then model rate-limit/budget, and returns a simpleDecisionAllowresult for the happy path. This makes PreHook’s hierarchical decision tree much easier to follow.framework/configstore/store.go (1)
30-38: ConfigStore interface extensions are coherent and well-grouped.Adding
GetProviders/GetProviderByName, usage-only updaters, and full model-config CRUD under the existing CRUD sections keeps the interface organized and matches the new RDB implementation signatures.Also applies to: 118-128
framework/configstore/rdb.go (1)
717-739: Provider, usage, and model-config datastore extensions look correct and align with governance needs.
GetProviders/GetProviderByName:
- Preload
BudgetandRateLimit, which is exactly what governance views and API surfaces need.UpdateBudgetUsage/UpdateRateLimitUsage:
- Use
SkipHooksand targetedUpdate/Updatesto mutate only usage counters, avoiding interfering with validation logic.- Model-config CRUD:
- Supports list/get/create/update/bulk-update/delete for
TableModelConfig.DeleteModelConfigcorrectly collects attached BudgetID/RateLimitID and removes the children inside a transaction, mirroring the 1:1 governance design.GetGovernanceConfig:
- Now includes
modelConfigsandprovidersin both the “any config present?” check and the returnedGovernanceConfig, which keeps the API surface in sync with the new governance entities.These look consistent with the rest of the config store and the tracker/resolver expectations.
Also applies to: 1999-2028, 2030-2151, 2153-2226
plugins/governance/store.go (2)
339-358: Well done fixing the early-return issue from previous reviews.The expired-limit handling now correctly tracks
tokenLimitExpiredandrequestLimitExpiredindependently, allowing both token and request limits to be evaluated even when one has expired. This addresses the feedback from previous reviews that flagged early returns skipping subsequent checks.Also applies to: 537-556
746-759: Well done fixing the race conditions from previous reviews.The usage update helpers now clone first before mutating, preventing data races where concurrent reads could see partial updates. This follows the correct pattern established in
UpdateBudgetUsageInMemoryand addresses the critical race condition issues flagged in previous reviews.Also applies to: 803-829
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
c33348c to
5a2256c
Compare
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
51556a1 to
c6e01b1
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
5a2256c to
509aaee
Compare
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
2 times, most recently
from
60af306 to
d4e7365
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
509aaee to
b72e947
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
plugins/governance/resolver.go (1)
301-405: Fix error formatting inCheckProviderRateLimit(passing[]stringto%s).The logic and decision mapping are correct, but the error construction uses
%swith a[]string, which will render as a%!s(MISSING=...)formatting error instead of a readable message.Recommend joining the violations, as done in
CheckRateLimit:Proposed fix for `CheckProviderRateLimit` error message
- if len(violations) > 0 { - // Determine specific violation type - decision := DecisionRateLimited // Default to general rate limited decision - if len(violations) == 1 { - if strings.Contains(violations[0], "token") { - decision = DecisionTokenLimited // More specific violation type - } else if strings.Contains(violations[0], "request") { - decision = DecisionRequestLimited // More specific violation type - } - } - return fmt.Errorf("rate limit violated for %s: %s", providerKey, violations), decision - } + if len(violations) > 0 { + // Determine specific violation type + decision := DecisionRateLimited // Default to general rate limited decision + if len(violations) == 1 { + if strings.Contains(violations[0], "token") { + decision = DecisionTokenLimited // More specific violation type + } else if strings.Contains(violations[0], "request") { + decision = DecisionRequestLimited // More specific violation type + } + } + msg := strings.Join(violations, "; ") + return fmt.Errorf("rate limit violated for %s: %s", providerKey, msg), decision + }framework/configstore/rdb.go (1)
562-604: Cascade deletes for provider budgets/rate limits rely on 1:1 ownership invariants.
DeleteProvidernow deletes the associatedTableBudgetandTableRateLimitunconditionally after removing the provider. This is fine as long as each budget/rate limit is owned exclusively by a single provider and never shared across entities, which matches the current governance design. If that invariant ever changes (e.g., shared budgets), this delete path would need additional reference checks or constraints.
🧹 Nitpick comments (2)
framework/configstore/rdb.go (1)
2033-2154: ModelConfig CRUD and cascades are consistent with the provider behavior but share the same 1:1 assumption.The new Get*/Create/Update/DeleteModelConfig methods look correct and mirror existing governance patterns.
DeleteModelConfigalso cascades deletes to itsBudgetandRateLimit. As with providers, this assumes those child records are not referenced elsewhere; tests enforcing the 1:1 relationship are important to keep this safe over time.plugins/governance/store.go (1)
2110-2260: ModelConfig/Provider in-memory CRUD preserves usage while updating limits; minor opportunity to keep reference helpers in sync.The
UpdateModelConfigInMemory/DeleteModelConfigInMemoryandUpdateProviderInMemory/DeleteProviderInMemorymethods:
- Clone the incoming structs.
- Use
checkAndUpdateBudget/checkAndUpdateRateLimitto preserve or reset usage as appropriate.- Keep the central
budgets/rateLimitsmaps as the canonical source of usage state.One optional improvement:
updateBudgetReferencesandupdateRateLimitReferencescurrently refresh VK, team, and customer pointers only. If future code starts relying onTableModelConfig.Budget/RateLimitorTableProvider.Budget/RateLimitpointers after resets, mirroring those updates there as well would keep all in-memory relationships strictly consistent, at the cost of a bit more bookkeeping.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (5)
- framework/configstore/clientconfig.go
- framework/configstore/migrations.go
- framework/configstore/tables/provider.go
- framework/configstore/store.go
- framework/configstore/tables/modelconfig.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goframework/configstore/rdb.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goframework/configstore/rdb.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/tracker.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.go
🧬 Code graph analysis (3)
plugins/governance/tracker.go (3)
core/schemas/models.go (1)
Model(109-129)core/schemas/provider.go (1)
Provider(315-362)ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)
plugins/governance/main.go (2)
plugins/governance/resolver.go (2)
Decision(15-15)DecisionAllow(18-18)core/utils.go (1)
IsFinalChunk(212-227)
framework/configstore/rdb.go (8)
framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(727-736)AuthConfig(717-722)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/team.go (2)
TableTeam(12-38)TableTeam(41-41)ui/lib/types/config.ts (1)
AuthConfig(259-264)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (9)
plugins/governance/tracker.go (1)
69-123: Hierarchical usage update flow inUpdateUsagelooks sound and matches the new model/provider semantics.Provider/model rate-limit and budget updates are applied first and unconditionally, with VK-level updates layered on only when a VK is present and resolvable. Streaming flags and cost/tokens guards are used consistently, and errors are logged without disrupting subsequent updates. No changes needed here.
plugins/governance/main.go (2)
473-515: PreHook’s reordered evaluation (model/provider first, then VK) is consistent and avoids redundant VK checks.Using
EvaluateModelAndProviderRequestas a first gate, then layeringEvaluateVirtualKeyRequestwhen a VK is present, gives clear precedence and isolates VK-specific failures. The VK‑mandatory path is cleanly handled viaisVkMandatory. No changes needed here.
573-612: PostHook usage tracking behavior is coherent with the tracker’s VK‑optional design.Restricting tracking to cases where
model != ""while always passingvirtualKey(possibly empty) intopostHookWorkermatches the tracker’s behavior of updating provider/model usage first and handling VK only when present. Streaming final‑chunk and success semantics are preserved.framework/configstore/rdb.go (1)
2157-2229: GovernanceConfig now correctly surfaces model configs and providers.Including
ModelConfigsandProvidersinGetGovernanceConfigaligns the backend payload with the extended governance surface (and the UI types), while still short‑circuiting tonilwhen all slices are empty. The directFindqueries here are sufficient since budgets and rate limits are also returned explicitly in the config.plugins/governance/store.go (5)
52-95: GovernanceStore interface extensions are cohesive and keep concerns nicely separated.The added provider/model-level check methods and in-memory CRUD/update hooks for model configs and providers fit well alongside the existing VK/team/customer APIs. Naming and signatures are consistent, and the contract around errors-as-policy-violations remains clear.
242-604: Provider/model-level budget and rate-limit checks align with existing VK hierarchy semantics.
CheckProviderBudget,CheckProviderRateLimit,CheckModelBudget, andCheckModelRateLimitall:
- Treat missing config or missing budget/rate-limit as “allow”.
- Respect reset durations by skipping only the expired side (tokens vs requests) instead of short‑circuiting the entire check.
- Incorporate per‑ID baselines, mirroring VK‑level behavior.
The use of cloned
TableRateLimit/TableBudgetvalues elsewhere ensures these reads are race‑free. This block looks good.
697-912: In-memory usage updates for VK/provider/model budgets and rate limits are concurrency-safe and correctly scoped.The new
UpdateVirtualKeyBudgetUsageInMemory,UpdateProviderAndModelBudgetUsageInMemory,UpdateProviderAndModelRateLimitUsageInMemory, andUpdateVirtualKeyRateLimitUsageInMemory:
- Clone map entries before mutation, avoiding the prior sync.Map value races.
- Respect reset durations on the cloned structs before incrementing usage.
- Apply cost/tokens/requests updates to all relevant scopes (VK hierarchy, provider, and model configs) without double‑mutating shared structs.
This matches the intended per-model and per-provider accounting, with safe in-memory updates.
1088-1150: ExtendingDumpRateLimitsto include model-config and provider rate limits is correct and complete.By collecting rate limit IDs from virtual keys, model configs, and providers,
DumpRateLimitsnow persists usage for all governance dimensions. Baseline handling and the deadlock-tolerant transaction pattern remain consistent with the existing VK behavior.
1288-1544: Database load and in-memory rebuild now correctly bring model configs and providers into the cache.
loadFromDatabase/loadFromConfigMemoryandrebuildInMemoryStructures:
- Fetch model configs and providers from the ConfigStore and wire Budget/RateLimit relationships.
- Populate the new
modelConfigsandprovidersmaps with the expected keys (modelNameandmodelName:provider, provider name).This provides the necessary backing for the new provider/model governance checks.
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
b72e947 to
d7c5209
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
plugins/governance/resolver.go (1)
127-195: Consider normalizing rate-limit check return signatures for readability
EvaluateVirtualKeyRequestandcheckRateLimitHierarchynow use:if decision, err := r.store.CheckRateLimit(...); err != nil { ... }while
EvaluateModelAndProviderRequestuses:if err, decision := r.store.CheckProviderRateLimit(...); err != nil { ... }So
CheckRateLimitis effectively(Decision, error)andCheckProviderRateLimitis(error, Decision). This compiles, but the differing orders add cognitive load when reading and maintaining the code.If it’s feasible, consider standardizing all
Check*RateLimithelpers to a single convention (e.g.,func(...) (Decision, error)), then adjust the call sites to match. That will make future extensions of the hierarchy (e.g., model-config-level checks) less error‑prone.Also applies to: 251-290
framework/configstore/rdb.go (3)
579-582: Provider delete now cascades to its budget/rate limit – keep 1:1 invariant in mindCapturing
budgetID/rateLimitIDbefore provider deletion and then deleting the correspondingTableBudget/TableRateLimitrows matches how VKs, teams, and customers are handled, and aligns with the documented 1:1 relationship between parents and their budgets/rate limits (no sharing).Given that, this is reasonable, but it does bake in the invariant that provider-level budgets/rate limits are not shared with any other entity. It would be good to:
- Keep tests/assertions enforcing that a given budget/rate-limit ID is referenced by at most one owner, and
- Revisit this cascade if you ever introduce shared budgets/limits.
Also applies to: 590-601
2002-2031: Usage-only updaters correctly bypass hooks but rely on hook semantics
UpdateBudgetUsageandUpdateRateLimitUsage:
- Scope the update by
id.- Use
Session(&gorm.Session{SkipHooks: true})and update only the usage counters (current_usage,token_current_usage,request_current_usage).- Avoid touching config fields like limits or reset durations.
This is exactly what you want for high-frequency usage writes, assuming there are no
BeforeSave/AfterSavehooks that maintain invariants purely on the usage fields (e.g., clamping or side-effect logging). If such hooks are added later, these helpers will need revisiting.If you add or already have GORM hooks on
TableBudget/TableRateLimit, please verify they don’t assumeSkipHooksis off for usage updates in GORM v1.31.1.
2033-2154: ModelConfig CRUD mirrors other governance entities; cascades assume 1:1 ownershipThe new model-config methods provide a complete persistence surface:
GetModelConfigs/GetModelConfig/GetModelConfigByIDwithBudget/RateLimitpreloads.CreateModelConfig,UpdateModelConfig,UpdateModelConfigsdelegating to GORM withparseGormError.DeleteModelConfigtransaction that:
- Fetches the config to capture
BudgetID/RateLimitID.- Deletes the config.
- Deletes the associated budget and rate limit rows if present.
This is consistent with how VK provider-configs and other governance owners are handled and again encodes the invariant that these budgets/rate limits are not shared across entities. As with providers, keeping tests that assert non-sharing will make these cascades safe long term.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/tables/provider.goplugins/governance/main.goplugins/governance/resolver.goplugins/governance/store.goplugins/governance/tracker.go
🚧 Files skipped from review as they are similar to previous changes (4)
- framework/configstore/tables/provider.go
- framework/configstore/tables/modelconfig.go
- framework/configstore/clientconfig.go
- framework/configstore/store.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/resolver.goframework/configstore/migrations.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/main.goplugins/governance/store.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/resolver.goframework/configstore/migrations.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/main.goplugins/governance/store.go
📚 Learning: 2025-12-22T10:50:40.990Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.
Applied to files:
plugins/governance/resolver.goplugins/governance/tracker.goplugins/governance/main.goplugins/governance/store.go
🧬 Code graph analysis (6)
plugins/governance/resolver.go (5)
core/schemas/bifrost.go (1)
ModelProvider(33-33)core/schemas/provider.go (1)
Provider(315-362)core/schemas/models.go (1)
Model(109-129)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
framework/configstore/migrations.go (2)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
plugins/governance/tracker.go (3)
core/schemas/models.go (1)
Model(109-129)core/schemas/provider.go (1)
Provider(315-362)ui/lib/types/governance.ts (2)
VirtualKey(61-80)RateLimit(13-25)
framework/configstore/rdb.go (3)
framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(727-736)AuthConfig(717-722)
plugins/governance/main.go (4)
examples/plugins/hello-world/main.go (1)
HTTPTransportMiddleware(19-27)core/schemas/plugin.go (1)
BifrostHTTPMiddleware(38-38)plugins/governance/resolver.go (1)
Decision(15-15)core/utils.go (1)
IsFinalChunk(212-227)
plugins/governance/store.go (6)
plugins/governance/resolver.go (5)
Decision(15-15)DecisionAllow(18-18)DecisionRateLimited(21-21)DecisionTokenLimited(23-23)DecisionRequestLimited(24-24)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(95-122)TableVirtualKey(125-125)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)framework/configstore/tables/budget.go (2)
TableBudget(11-27)TableBudget(30-30)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (11)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (13)
framework/configstore/migrations.go (2)
131-136: Wiring new migrations at the end oftriggerMigrationslooks correctAppending
migrationAddModelConfigTableandmigrationAddProviderGovernanceColumnsat the tail preserves existing migration ordering while ensuring prerequisites (core tables, budgets, rate limits) already exist. No issues here.
2124-2152: ModelConfig table migration is consistent with existing patterns
migrationAddModelConfigTablefollows the establishedCreateTable/DropTablepattern used elsewhere and delegates indexes/FKs to theTableModelConfigschema tags. This aligns with the new governance_model_configs data model.plugins/governance/tracker.go (1)
72-106: Provider/model usage updates before VK-level logic match the new hierarchyThe refactored
UpdateUsagenow:
- Skips failed requests early (with a debug log).
- Always updates provider- and model-level rate limits and budgets for successful calls, regardless of VK presence.
- Only then looks up the VK and applies VK/Team/Customer updates if a VK value exists and is found.
That matches the intended “model → provider → VK” hierarchy and ensures provider/model governance still works when virtual keys are disabled or absent. This looks sound assuming the
GovernanceStoremethods gracefully handle cases where provider or model are non-empty and VK is missing.Also applies to: 108-123
plugins/governance/resolver.go (1)
77-125: Model/provider evaluation split is clear and ordered correctly
EvaluateModelAndProviderRequestcleanly enforces:
- Provider-level rate limit and budget checks when
provider != "".- Model-level rate limit and budget checks when
model != ""(optionally scoped by provider).Returning a plain
DecisionAllowwhen these pass keeps PreHook’s control flow simple, and the separation from VK logic inEvaluateVirtualKeyRequestmatches the new hierarchical design. This looks good; just ensure all call sites (e.g., the plugin PreHook) invoke this method before any VK-specific evaluation, as you’re doing inmain.go.plugins/governance/main.go (3)
42-49: HTTP middleware interface aligns plugin with core transport patternsSwitching
BaseGovernancePluginto exposeHTTPTransportMiddleware() schemas.BifrostHTTPMiddlewarematches the coreBifrostHTTPMiddlewareabstraction and the concreteGovernancePlugin.HTTPTransportMiddlewareimplementation below. This keeps the plugin’s transport integration consistent with other plugins and the core examples.
477-503: PreHook correctly layers provider/model checks before VK-specific evaluationThe updated PreHook flow:
- Extracts
providerandmodelviareq.GetRequestFields().- Enforces VK-mandatory behavior (if configured).
- Runs
resolver.EvaluateModelAndProviderRequestunconditionally for governance, even when no VK is present.- Only if that passes and
virtualKeyValueis non-empty, runsEvaluateVirtualKeyRequestto apply VK/team/customer hierarchy.This ordering ensures provider/model limits are always applied first, with VK governance layered on top when available, and keeps the VK-required path behavior unchanged. The decision-mapping code that follows remains compatible with the new Decision values.
599-609: PostHook usage tracking is now VK-optional but still model-dependentIn PostHook:
- You always compute
virtualKey,requestID,provider, andmodelfrom context/response.- You only spawn
postHookWorkerwhenmodel != "", which matches the tracker’s expectation that provider/model must be known for usage accounting.virtualKey(possibly empty) is passed through; the worker builds aUsageUpdatethat the tracker uses to:
- Always update provider/model rate limits and budgets, and
- Optionally update VK hierarchy when
VirtualKeyis non-empty and exists.The worker’s comment and the note before
UpdateUsagemake it clear that empty VK values are handled by only touching provider/model usage. This is consistent with the new hierarchical tracking behavior.Also applies to: 630-641, 691-693
framework/configstore/rdb.go (2)
719-742: Provider getters preload governance fields for resolver/tracker use
GetProvidersandGetProviderByNamecorrectly:
- Use the
config_providerstable viatables.TableProvider.- Preload
BudgetandRateLimitso higher layers don’t need extra queries for governance metadata.- Map “not found” to the existing
ErrNotFoundsentinel for the single-provider path.Using
Find+ empty-slice semantics forGetProvidersis appropriate; theErrRecordNotFoundcheck is effectively dead but harmless.
2163-2165: GovernanceConfig now fully reflects model/provider-level governance
GetGovernanceConfighas been extended to:
- Load
modelConfigsandprovidersalongside existing virtual keys, teams, customers, budgets, and rate limits.- Treat any non-empty slice of these (including modelConfigs/providers) or governanceConfigs as sufficient to consider governance “present” (the all-zero-length check includes them).
- Populate the returned
GovernanceConfigwith the newModelConfigsandProvidersfields.This ensures callers (including the UI) see the new per-model and per-provider governance surfaces without needing separate queries, and avoids returning
nilwhen only model/provider-level config exists.Also applies to: 2182-2187, 2193-2193, 2221-2228
plugins/governance/store.go (4)
18-604: LGTM: Solid implementation of provider and model-level governance checks.The new provider and model-level governance methods are well-implemented:
Past issues resolved: The flag-based approach for expired limits (lines 341-358, 539-556) correctly avoids the early return problem identified in previous reviews. Now if only the token limit is expired, the request limit is still checked.
Consistent patterns: The implementation follows the same patterns as VK-level checks—collecting applicable configs, checking expiry per limit type, failing fast on first violation.
Hierarchical checking: Model+provider configs are checked before model-only configs (lines 419-436, 499-516), ensuring more specific configurations take precedence while still enforcing all applicable constraints.
Proper nil handling: Appropriate nil checks throughout (e.g., lines 252-260, 314-324).
Based on coding guidelines: This review considers the PR as part of the stack (#1170, #1169, #1121, #1120, #1119, #888), and the governance evaluation hierarchy (model → provider → virtual key) mentioned in the PR objectives should be verified in the resolver implementation.
738-863: LGTM: Race conditions properly addressed with clone-first pattern.The new provider and model usage update methods correctly implement the clone-first pattern:
- Lines 746-747, 803-804: Budgets and rate limits are cloned immediately after retrieval and before any mutations
- All reset checks and usage increments operate on the clone
- The clone is stored back atomically
This addresses the critical race conditions identified in previous reviews where the original cached object was being mutated before cloning, causing concurrent readers to see partial updates.
The inline helper functions (updateBudget and updateRateLimit) provide clean code reuse while maintaining the correct concurrency semantics.
1287-1544: LGTM: Data loading and structure rebuilding correctly extended.The loading and rebuilding logic properly handles the new model configs and providers:
Consistent key format (lines 1527-1535): Provider-specific configs use
"modelName:provider", global configs use"modelName". This matches the retrieval logic in check methods.Relationship population (lines 1364-1416): When loading from config memory (config.json), relationships are manually wired for both model configs and providers, ensuring budgets and rate limits are properly linked.
Complete rebuild (lines 1490-1543): All sync.Maps are properly cleared and rebuilt with the new entities.
The implementation maintains consistency with the existing patterns for VKs, teams, and customers.
2110-2260: LGTM: Model config and provider CRUD methods follow established patterns.The new CRUD methods are well-implemented:
Usage preservation (lines 2129, 2143, 2212, 2226): The
checkAndUpdateBudgetandcheckAndUpdateRateLimithelpers preserve existing usage values when updating, preventing loss of accumulated usage data.Cascade deletes (lines 2177-2184, 2248-2255): Associated budgets and rate limits are deleted when the parent entity is removed, which is correct given the 1:1 relationships noted in learnings.
Key format handling (lines 2149-2157, 2233): Storage correctly uses provider-specific or global keys based on whether the Provider field is set.
The implementation is consistent with existing VK, team, and customer CRUD methods, maintaining uniformity across the governance store.
Based on learnings: The 1:1 relationship assumption (budgets and rate limits are not shared) is maintained throughout the implementation.
…rate limiting in governance plugin
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
d7c5209 to
d6c6b20
Compare
Summary
Added model-level and provider-level governance capabilities to enhance the
governance system with hierarchical controls that work even without virtual
keys.
Changes
ModelConfigtable to store model-specific governance settingsProvidertablevirtual key constraints
virtual key is present
Type of change
Affected areas
How to test
check PR #1121
Breaking changes
Security considerations
This change enhances security by providing more granular control over model and
provider usage, allowing for better resource protection even when virtual keys
are not used.
Checklist
docs/contributing/README.mdand followed the guidelines