Conversation
📝 WalkthroughSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughAdds model-config and provider governance: new ORM model and migrations; extends configstore RDB and interface with provider/model CRUD and usage updates; expands in-memory governance store and checks; splits resolver into model/provider vs virtual-key flows; updates hooks and tracker sequencing. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant GovernancePlugin
participant BudgetResolver
participant LocalStore
participant Tracker
rect rgba(240,248,255,0.5)
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 rgba(245,255,240,0.5)
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 Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
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.
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
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
b232bfb to
7cbddba
Compare
danpiths
force-pushed
the
graphite-base/1120
branch
from
b0ed2b6 to
074ece1
Compare
danpiths
changed the base branch from
graphite-base/1120
to
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
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 (2)
framework/configstore/rdb.go (1)
608-630: Budgets and rate limits can be shared across entities—deleting a provider will orphan other references.Lines 608–628: The code deletes the provider's
BudgetIDandRateLimitIDoutright. However, the same Budget/RateLimit IDs can be referenced by other entities (virtual keys, provider configs, model configs). The data model defines CASCADE on the referencing side only, so deleting one entity's budget/rate limit will break other entities that reference the same ID. This pattern is replicated acrossDeleteVirtualKey,DeleteVirtualKeyProviderConfig, andDeleteModelConfigas well. Either enforce that each Budget/RateLimit is exclusively owned by a single entity, or add reference counting/constraint validation before deletion.framework/configstore/store.go (1)
37-125: MockConfigStore is incomplete and will cause compilation failures.The
RDBConfigStoreimplementation inframework/configstore/rdb.gohas been properly updated with all new interface methods. However, theMockConfigStoreintransports/bifrost-http/lib/config_test.gois missing 12 methods:
GetProvidersGetProviderByNameUpsertPluginUpdateBudgetUsageUpdateRateLimitUsageGetModelConfigsGetModelConfigGetModelConfigByIDCreateModelConfigUpdateModelConfigUpdateModelConfigsDeleteModelConfigThese missing stub methods must be added to MockConfigStore to satisfy the interface contract, even if not currently tested. Add minimal implementations (e.g., returning nil/empty values or panicking) to prevent interface assertion failures.
🤖 Fix all issues with AI agents
In `@framework/configstore/rdb.go`:
- Around line 2131-2164: DeleteModelConfig currently deletes associated
Budget/RateLimit blindly; update it (and mirror the same pattern in
DeleteProvider, DeleteVirtualKey, DeleteVirtualKeyProviderConfig, DeleteTeam,
DeleteCustomer) to first check whether the budgetID/rateLimitID are referenced
elsewhere and only delete when no other references exist. Concretely, inside the
transaction after reading budgetID/rateLimitID, run count queries (using
tx.Model(&tables.TableModelConfig)/tx.Model(&tables.TableProvider)/tx.Model(&tables.TableVirtualKey)/etc.
with Where("budget_id = ?" or "rate_limit_id = ?") excluding the current record
when appropriate) to get reference counts; only call
tx.Delete(&tables.TableBudget{}, "id = ?", id) or
tx.Delete(&tables.TableRateLimit{}, "id = ?", id) if the respective count is
zero, otherwise skip deletion. Keep error handling consistent (use
s.parseGormError where used) and perform these checks/deletes within the same
transaction so DeleteModelConfig, DeleteProvider, DeleteVirtualKey,
DeleteVirtualKeyProviderConfig, DeleteTeam, and DeleteCustomer all enforce
exclusive-deletion semantics.
- Around line 2014-2043: Both UpdateBudgetUsage and UpdateRateLimitUsage must
treat 0 rows affected as a not-found error; after the Update/Updates call check
result.RowsAffected and if it is 0 return
s.parseGormError(gorm.ErrRecordNotFound) (or the store's ErrNotFound via
parseGormError) so callers receive ErrNotFound instead of a silent success;
implement this check in UpdateBudgetUsage and UpdateRateLimitUsage immediately
after verifying result.Error.
In `@framework/configstore/tables/modelconfig.go`:
- Around line 13-17: The unique index idx_model_provider allows multiple rows
with provider=NULL which breaks GetModelConfig(..., provider=nil); change
Provider from a nullable *string to a non-nullable string sentinel and enforce
NOT NULL/default "" so the unique index covers (model_name, provider)
deterministically: update the ModelConfig struct field (Provider) type and GORM
tag to not null with a default empty string (remove omitempty), adjust the
BeforeSave hook and any callers to stop using nil (normalize to ""), and add a
migration to alter the column to NOT NULL DEFAULT '' so the unique index
prevents duplicate global configs.
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
2175-2245: Consider preloading governance relationships for providers/modelConfigs.In
GetGovernanceConfig,modelConfigsandprovidersare loaded withoutPreload("Budget")/Preload("RateLimit"), unlikeGetProviders/GetModelConfigs. If consumers expect nested budgets/rate limits in the governance payload, consider preloading here or explicitly documenting that only IDs are provided.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/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.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goframework/configstore/rdb.go
🧠 Learnings (5)
📚 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/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2025-12-29T11:54:55.836Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1153
File: framework/configstore/rdb.go:2221-2246
Timestamp: 2025-12-29T11:54:55.836Z
Learning: In Go reviews, do not flag range-over-int patterns like for i := range n as compile-time errors, assuming Go 1.22+ semantics. Only flag actual range-capable values (slices, arrays, maps, channels, strings) and other compile-time issues. This applies to all Go files across the repository.
Applied to files:
framework/configstore/migrations.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2026-01-14T04:40:11.480Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1312
File: framework/modelcatalog/pricing.go:276-426
Timestamp: 2026-01-14T04:40:11.480Z
Learning: In the Bifrost codebase, ImageUsage and other usage types guarantee that TotalTokens is populated (computed as InputTokens + OutputTokens if providers don’t supply TotalTokens). Reviewers can rely on this invariant and should not assume TotalTokens may be missing when input/output tokens exist. When implementing tiering logic or token-based decisions, you can safely use TotalTokens without extra null/zero guards, provided you’re in a context where InputTokens and OutputTokens are present. If a branch might discard tokens, ensure the invariant is preserved or add explicit checks only where the inputs are confirmed to be valid.
Applied to files:
framework/configstore/migrations.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2026-01-14T13:30:28.760Z
Learnt from: Radheshg04
Repo: maximhq/bifrost PR: 1326
File: plugins/semanticcache/test_utils.go:545-559
Timestamp: 2026-01-14T13:30:28.760Z
Learning: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.
Applied to files:
framework/configstore/migrations.goframework/configstore/store.goframework/configstore/tables/modelconfig.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2026-01-14T21:52:10.307Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1316
File: framework/configstore/clientconfig.go:276-280
Timestamp: 2026-01-14T21:52:10.307Z
Learning: In framework/configstore/clientconfig.go, when implementing GenerateKeyHash, prefix hashed keys with distinct identifiers: use 'env:' for environment-sourced values and 'val:' for literal values, mirroring the existing 'GenerateMCPClientHash' pattern to prevent hash collisions. This is a file-specific guideline for this component; no broad pattern is implied beyond this file.
Applied to files:
framework/configstore/clientconfig.go
🧬 Code graph analysis (4)
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 (6)
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)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)ui/lib/types/config.ts (1)
AuthConfig(278-283)
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/errors.go (1)
ErrNotFound(9-9)framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/clientconfig.go (2)
GovernanceConfig(759-768)AuthConfig(749-754)
⏰ 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
🔇 Additional comments (8)
framework/configstore/tables/modelconfig.go (2)
31-34: Table mapping looks good.
36-58: Validation hook is solid. Trimming + non-empty checks are a nice guardrail.framework/configstore/clientconfig.go (1)
759-767: GovernanceConfig extension looks clean.framework/configstore/migrations.go (3)
144-149: Migration wiring looks good.
2336-2413: Provider governance column migration is well-structured (indexes + rollback).
2306-2334: No action needed. GORM v1.31 automatically creates indexes from struct tags duringCreateTable. TheTableModelConfigdefinition already includes named index tags (uniqueIndex:idx_model_provider,index:idx_model_config_budget,index:idx_model_config_rate_limit), so these indexes will be created without explicitCREATE INDEXstatements. The migration follows GORM best practices.Likely an incorrect or invalid review comment.
framework/configstore/rdb.go (2)
718-740: Provider getters with preloaded governance look good.
2045-2129: Model config CRUD additions look good.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@framework/configstore/rdb.go`:
- Around line 2028-2043: In UpdateRateLimitUsage, after executing the Updates
call check result.RowsAffected and if result.Error is nil and RowsAffected == 0
return ErrNotFound so callers can distinguish a missing rate-limit ID; keep the
existing error path by still returning s.parseGormError(result.Error) when
result.Error != nil and only treat zero RowsAffected as ErrNotFound (avoid
changing SkipHooks/session logic or update fields).
- Around line 2014-2026: In UpdateBudgetUsage, check result.RowsAffected after
the update to detect when no row was updated and return ErrNotFound in that
case; specifically, after calling
s.db.WithContext(ctx).Session(...).Model(&tables.TableBudget{}).Where("id = ?",
id).Update(...), if result.Error != nil return s.parseGormError(result.Error),
else if result.RowsAffected == 0 return ErrNotFound, otherwise return nil —
update the function to inspect RowsAffected and use the existing parseGormError
and ErrNotFound symbol to distinguish a missing budget ID from a successful
no-op update.
In `@framework/configstore/tables/modelconfig.go`:
- Around line 13-17: The unique index on ModelConfig (fields ModelName and
Provider) allows multiple NULL providers for the same model; modify the schema
or logic to enforce a single default provider: either add a DB partial unique
index on the table (e.g., UNIQUE ON model_name WHERE provider IS NOT NULL and
separately ensure a single default row per model), or replace nullable Provider
with a sentinel string (e.g., "default") and update GetModelConfig to query that
sentinel, or add application-level checks in CreateModelConfig and
UpdateModelConfig to reject creating/updating a second NULL/default provider for
the same ModelName; locate the struct fields (Provider, ModelName) in
modelconfig.go and update the migration/index creation or adjust
CreateModelConfig/UpdateModelConfig and GetModelConfig to implement the chosen
approach.
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
608-630: Wrap provider + budget/rate-limit deletions in a transaction.Without a transaction, provider deletion could succeed while budget/rate-limit cleanup fails, leaving orphans. Consider wrapping these steps in a tx when one isn’t supplied.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
framework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.go
🚧 Files skipped from review as they are similar to previous changes (1)
- 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/modelconfig.goframework/configstore/store.goframework/configstore/clientconfig.goframework/configstore/rdb.go
🧠 Learnings (5)
📚 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/store.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2025-12-29T11:54:55.836Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1153
File: framework/configstore/rdb.go:2221-2246
Timestamp: 2025-12-29T11:54:55.836Z
Learning: In Go reviews, do not flag range-over-int patterns like for i := range n as compile-time errors, assuming Go 1.22+ semantics. Only flag actual range-capable values (slices, arrays, maps, channels, strings) and other compile-time issues. This applies to all Go files across the repository.
Applied to files:
framework/configstore/tables/modelconfig.goframework/configstore/store.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2026-01-14T04:40:11.480Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1312
File: framework/modelcatalog/pricing.go:276-426
Timestamp: 2026-01-14T04:40:11.480Z
Learning: In the Bifrost codebase, ImageUsage and other usage types guarantee that TotalTokens is populated (computed as InputTokens + OutputTokens if providers don’t supply TotalTokens). Reviewers can rely on this invariant and should not assume TotalTokens may be missing when input/output tokens exist. When implementing tiering logic or token-based decisions, you can safely use TotalTokens without extra null/zero guards, provided you’re in a context where InputTokens and OutputTokens are present. If a branch might discard tokens, ensure the invariant is preserved or add explicit checks only where the inputs are confirmed to be valid.
Applied to files:
framework/configstore/tables/modelconfig.goframework/configstore/store.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2026-01-14T13:30:28.760Z
Learnt from: Radheshg04
Repo: maximhq/bifrost PR: 1326
File: plugins/semanticcache/test_utils.go:545-559
Timestamp: 2026-01-14T13:30:28.760Z
Learning: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.
Applied to files:
framework/configstore/tables/modelconfig.goframework/configstore/store.goframework/configstore/clientconfig.goframework/configstore/rdb.go
📚 Learning: 2026-01-14T21:52:10.307Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1316
File: framework/configstore/clientconfig.go:276-280
Timestamp: 2026-01-14T21:52:10.307Z
Learning: In framework/configstore/clientconfig.go, when implementing GenerateKeyHash, prefix hashed keys with distinct identifiers: use 'env:' for environment-sourced values and 'val:' for literal values, mirroring the existing 'GenerateMCPClientHash' pattern to prevent hash collisions. This is a file-specific guideline for this component; no broad pattern is implied beyond this file.
Applied to files:
framework/configstore/clientconfig.go
🧬 Code graph analysis (2)
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 (3)
framework/configstore/tables/modelconfig.go (2)
TableModelConfig(12-29)TableModelConfig(32-34)framework/configstore/tables/provider.go (2)
TableProvider(15-52)TableProvider(55-55)ui/lib/types/config.ts (1)
AuthConfig(278-283)
⏰ 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). (10)
- 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/store.go (2)
37-38: LGTM: provider retrieval interface additions.These methods align with the new provider-level governance APIs.
115-125: LGTM: governance usage + model config CRUD surface.Interface extensions look consistent with the new model/provider governance flows.
framework/configstore/clientconfig.go (1)
760-767: LGTM: GovernanceConfig includes model/provider slices.Struct update cleanly exposes the new governance data in JSON.
framework/configstore/tables/modelconfig.go (1)
37-56: Validation hook looks solid.Checks cover empty fields for model, budget, rate limit, and provider.
framework/configstore/rdb.go (9)
718-740: LGTM: provider fetch with governance preloads.Preloading Budget/RateLimit here matches the provider-level governance usage.
2045-2055: LGTM: model-config list retrieval.
2057-2073: LGTM: model-config lookup by model/provider.
2075-2084: LGTM: model-config lookup by ID.
2087-2099: LGTM: create model-config.
2101-2113: LGTM: update model-config.
2115-2129: LGTM: batch update model-configs.
2131-2165: LGTM: model-config delete with budget/rate-limit cleanup.
2175-2245: Confirm whether GovernanceConfig expects preloaded provider/model relationships.GetGovernanceConfig returns providers/modelConfigs without Preload("Budget"/"RateLimit"). If downstream relies on those relationship pointers (instead of linking by ID with budgets/rateLimits slices), consider preloading here.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
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
7cbddba to
c478008
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/configstore/rdb.go (1)
608-630: Wrap provider + budget/rate-limit deletes in a transaction when no tx is supplied.Lines 608-630 delete the provider and then its budget/rate-limit without a transaction when
txisn't supplied. Although CASCADE constraints are defined in the schema, the manual deletion of budget and rate limit requires atomicity—a failure after provider deletion can leave orphan governance records. Follow the pattern used inDeleteTeamandDeleteVirtualKeyby wrapping the operation ins.db.WithContext(ctx).Transaction()when no tx is passed.🐛 Suggested pattern
func (s *RDBConfigStore) DeleteProvider(ctx context.Context, provider schemas.ModelProvider, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } + var txDB *gorm.DB + var localTx bool + if len(tx) > 0 { + txDB = tx[0] + } else { + txDB = s.db.Begin() + localTx = true + } // Find the existing provider var dbProvider tables.TableProvider if err := txDB.WithContext(ctx).Where("name = ?", string(provider)).First(&dbProvider).Error; err != nil { + if localTx { + txDB.Rollback() + } if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } @@ if err := txDB.WithContext(ctx).Delete(&dbProvider).Error; err != nil { + if localTx { + txDB.Rollback() + } if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } @@ if budgetID != nil { if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", *budgetID).Error; err != nil { + if localTx { + txDB.Rollback() + } return err } } @@ if rateLimitID != nil { if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil { + if localTx { + txDB.Rollback() + } return err } } - - return nil + if localTx { + return txDB.Commit().Error + } + return nil }
🤖 Fix all issues with AI agents
In `@plugins/governance/store.go`:
- Around line 2007-2036: When rehydrating clone.Budget and clone.RateLimit you
only copy current usage but not reset timestamps or DB-tracking fields, which
allows stale LastReset/TokenLastReset and LastDB* fields to overwrite in-memory
state; update the budget block (gs.budgets handling) to, when
existingBudgetValue is a *configstoreTables.TableBudget, copy eb.CurrentUsage
AND eb.LastReset and any eb.LastDB* fields into clone.Budget (e.g.,
clone.Budget.LastReset = eb.LastReset; clone.Budget.LastDBReset =
eb.LastDBReset), and update the rate limit block (gs.rateLimits handling) to,
when existingRateLimitValue is a *configstoreTables.TableRateLimit, copy
erl.TokenCurrentUsage, erl.RequestCurrentUsage AND
erl.TokenLastReset/RequestLastReset and any erl.TokenLastDB*/RequestLastDB* into
clone.RateLimit before calling gs.budgets.Store/gs.rateLimits.Store; apply the
same preservation logic for both provider-specific and global model config paths
where clone is stored in gs.modelConfigs.
- Around line 418-435: The model config gating is wrongly checking mc.Budget !=
nil which skips configs that only have BudgetID set; update both checks inside
the gs.modelConfigs.Load branches to test mc.BudgetID (same pattern as
CheckModelRateLimit uses RateLimitID) so configs with a BudgetID but without the
preloaded Budget relation are still considered—keep the rest of the logic
(appending to modelConfigsToCheck and budgetNames) unchanged and ensure you
reference configstoreTables.TableModelConfig's BudgetID field.
♻️ Duplicate comments (1)
framework/configstore/tables/modelconfig.go (1)
12-29: Struct design and relationships look correct.The
TableModelConfigstruct properly models model-specific governance with:
- Appropriate field types and GORM tags
- Correct cascade delete behavior for Budget/RateLimit relationships (aligns with 1:1 ownership pattern per learnings)
- ConfigHash for change detection during config sync
However, the nullable
Providerfield in the unique indexidx_model_provider(lines 14-15) still has the known issue where SQL semantics allow multiple rows with(model_name, provider=NULL)since NULL values aren't considered equal in unique constraints. This was flagged in previous reviews.
🧹 Nitpick comments (3)
framework/configstore/tables/provider.go (1)
90-98: Consider usingstrings.TrimSpacefor consistency withTableModelConfig.The validation here checks for empty strings directly (
*p.BudgetID == ""), whileTableModelConfig.BeforeSaveusesstrings.TrimSpace(*mc.BudgetID) == "". This inconsistency means whitespace-only strings like" "would pass validation here but fail inTableModelConfig.♻️ Suggested fix 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") }Note: You'll need to add
"strings"to the imports.plugins/governance/main.go (1)
705-705: Consider removing unused parametersisCacheReadandisBatch.The function signature uses unnamed parameters (
_, _) forisCacheReadandisBatch, indicating they're intentionally ignored. These values are extracted from context inPostHook(lines 650-660) but never used inpostHookWorker. If these are for future use, consider documenting that intent; otherwise, they could be removed to simplify the interface.plugins/governance/resolver.go (1)
103-106: Usebifrost.Ptrfor optional provider pointer creation.Line 103-106 uses
&provider. Repo convention prefersbifrost.Ptr(...)for consistency.♻️ Suggested change
import ( "context" "fmt" "slices" "time" + bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" "github.com/maximhq/bifrost/framework/modelcatalog" ) @@ var providerPtr *schemas.ModelProvider if provider != "" { - providerPtr = &provider + providerPtr = bifrost.Ptr(provider) }Based on learnings, prefer
bifrost.Ptr(...)over&valuefor pointer creation.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
ui/package-lock.jsonis excluded by!**/package-lock.json
📒 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/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/clientconfig.goframework/configstore/tables/provider.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/resolver.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.goframework/configstore/tables/modelconfig.go
🧠 Learnings (7)
📚 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.goframework/configstore/tables/provider.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/resolver.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.goframework/configstore/tables/modelconfig.go
📚 Learning: 2025-12-29T11:54:55.836Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1153
File: framework/configstore/rdb.go:2221-2246
Timestamp: 2025-12-29T11:54:55.836Z
Learning: In Go reviews, do not flag range-over-int patterns like for i := range n as compile-time errors, assuming Go 1.22+ semantics. Only flag actual range-capable values (slices, arrays, maps, channels, strings) and other compile-time issues. This applies to all Go files across the repository.
Applied to files:
framework/configstore/clientconfig.goframework/configstore/tables/provider.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/resolver.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.goframework/configstore/tables/modelconfig.go
📚 Learning: 2026-01-14T04:40:11.480Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1312
File: framework/modelcatalog/pricing.go:276-426
Timestamp: 2026-01-14T04:40:11.480Z
Learning: In the Bifrost codebase, ImageUsage and other usage types guarantee that TotalTokens is populated (computed as InputTokens + OutputTokens if providers don’t supply TotalTokens). Reviewers can rely on this invariant and should not assume TotalTokens may be missing when input/output tokens exist. When implementing tiering logic or token-based decisions, you can safely use TotalTokens without extra null/zero guards, provided you’re in a context where InputTokens and OutputTokens are present. If a branch might discard tokens, ensure the invariant is preserved or add explicit checks only where the inputs are confirmed to be valid.
Applied to files:
framework/configstore/clientconfig.goframework/configstore/tables/provider.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/resolver.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.goframework/configstore/tables/modelconfig.go
📚 Learning: 2026-01-14T13:30:28.760Z
Learnt from: Radheshg04
Repo: maximhq/bifrost PR: 1326
File: plugins/semanticcache/test_utils.go:545-559
Timestamp: 2026-01-14T13:30:28.760Z
Learning: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.
Applied to files:
framework/configstore/clientconfig.goframework/configstore/tables/provider.goplugins/governance/tracker.goframework/configstore/rdb.goplugins/governance/resolver.goframework/configstore/store.goplugins/governance/main.goplugins/governance/store.goframework/configstore/tables/modelconfig.go
📚 Learning: 2026-01-14T21:52:10.307Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1316
File: framework/configstore/clientconfig.go:276-280
Timestamp: 2026-01-14T21:52:10.307Z
Learning: In framework/configstore/clientconfig.go, when implementing GenerateKeyHash, prefix hashed keys with distinct identifiers: use 'env:' for environment-sourced values and 'val:' for literal values, mirroring the existing 'GenerateMCPClientHash' pattern to prevent hash collisions. This is a file-specific guideline for this component; no broad pattern is implied beyond this file.
Applied to files:
framework/configstore/clientconfig.go
📚 Learning: 2026-01-16T14:35:57.601Z
Learnt from: danpiths
Repo: maximhq/bifrost PR: 1120
File: framework/configstore/rdb.go:2131-2164
Timestamp: 2026-01-16T14:35:57.601Z
Learning: In the Bifrost codebase, budgets and rate limits have a 1:1 ownership with their parent entities (ModelConfig, Provider, VirtualKey, Team, Customer). They are created along with the parent and deleted together. There is no API to share a budget or rate limit across multiple entities. Therefore review checks should not rely on cross-entity references or pre-deletion reference validation; ensure deletions cascade through the parent and avoid patterns that imply sharing budgets/rates. This guideline applies to all Go files under framework/configstore.
Applied to files:
framework/configstore/clientconfig.goframework/configstore/tables/provider.goframework/configstore/rdb.goframework/configstore/store.goframework/configstore/tables/modelconfig.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)
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/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/resolver.go (4)
core/schemas/context.go (1)
BifrostContext(32-43)core/schemas/bifrost.go (1)
ModelProvider(30-30)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(148-175)TableVirtualKey(178-178)framework/configstore/tables/ratelimit.go (2)
TableRateLimit(11-36)TableRateLimit(39-39)
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 (10)
core/schemas/bifrost.go (1)
ModelProvider(30-30)ui/lib/types/config.ts (1)
ModelProvider(216-220)plugins/governance/resolver.go (2)
Decision(16-16)EvaluationRequest(31-36)framework/configstore/tables/virtualkey.go (2)
TableVirtualKey(148-175)TableVirtualKey(178-178)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)ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)core/schemas/provider.go (1)
Provider(323-378)
framework/configstore/tables/modelconfig.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
⏰ 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). (1)
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (16)
framework/configstore/tables/modelconfig.go (1)
36-59: LGTM!The
BeforeSavevalidation correctly ensures:
ModelNameis never empty- Optional pointer fields (
BudgetID,RateLimitID,Provider) are validated only when non-nil- Consistent use of
strings.TrimSpacefor whitespace handlingframework/configstore/tables/provider.go (1)
41-48: LGTM!The governance fields are correctly added with:
- Appropriate nullable types and indexes
- Cascade delete behavior matching the 1:1 ownership pattern
framework/configstore/clientconfig.go (1)
759-768: LGTM!The
GovernanceConfigstruct is correctly extended with:
ModelConfigsfor model-level governance settingsProvidersfor provider-level governance settingsThese additions enable the hierarchical governance evaluation (model → provider → virtual key) introduced in this PR.
plugins/governance/main.go (2)
537-566: LGTM!The PreHook correctly implements hierarchical governance evaluation:
- Extracts provider/model early for model/provider-level checks
- Validates VK requirement when mandatory
- Evaluates model/provider constraints first (applies universally)
- Only evaluates VK-specific constraints if model/provider checks passed AND VK exists
This enables governance enforcement even when virtual keys are disabled, supporting the PR's goal of per-model and per-provider budgeting/rate limiting.
663-673: LGTM!The PostHook correctly processes usage tracking:
- Spawns worker when model is present (model is required for meaningful tracking)
- VirtualKey can be empty string; tracker handles this gracefully
- Comments accurately describe the behavior
plugins/governance/tracker.go (1)
68-123: LGTM!The
UpdateUsagemethod correctly implements hierarchical usage tracking:
- Provider/model updates (lines 81-93): Run unconditionally, enabling governance even without virtual keys
- Early return (lines 95-99): Cleanly exits when no VK is present
- VK-specific updates (lines 108-122): Only run when VK exists and is found
Error handling appropriately logs failures but continues processing to prevent cascading failures.
framework/configstore/store.go (1)
37-125: Verify all ConfigStore implementations are updated for the new interface.Line 37-125 expands
ConfigStorewith provider/model CRUD and usage update APIs. Please confirm all concrete stores and mocks in the stack (SQLite/Postgres, test fakes) implement these methods so downstream PRs compile cleanly.plugins/governance/resolver.go (2)
80-128: Verify all call sites are updated for the split evaluation API.Line 80-128 introduces
EvaluateModelAndProviderRequestand removes the old single entry point. Please confirm all callers across the stack (main/tracker/interceptor) now invoke the split methods to avoid integration drift.
130-217: LGTM: VK path and hierarchy helpers are cleanly structured.The explicit provider/model filtering and the new hierarchy helpers make the VK path easy to follow and align with the new governance flow.
Also applies to: 262-300
framework/configstore/rdb.go (1)
718-740: LGTM: provider/model CRUD and usage-only updates align with governance needs.The new CRUD surface, usage-only update helpers, and governance config population are consistent and keep DB writes minimal.
Also applies to: 2014-2049, 2051-2171, 2181-2251
plugins/governance/store.go (6)
20-27: LGTM: in-memory caches and governance store contract cover providers/model configs.The additional caches and interface methods make the provider/model governance flow explicit and consistent.
Also applies to: 55-95
1319-1333: LGTM: load paths wire model configs/providers into memory cleanly.Both DB and config-memory loaders populate model configs/providers and rebuild the maps coherently.
Also applies to: 1358-1478, 1483-1543
242-298: LGTM: provider/model rate-limit checks handle expiries per limit.Per-limit expiry flags keep token vs request enforcement consistent across levels.
Also applies to: 300-405, 482-604
697-736: LGTM: usage updates clone before mutation and cover provider/model scopes.The clone-first pattern keeps
sync.Mapsafe while extending usage tracking to provider/model budgets and rate limits.Also applies to: 738-863, 865-912
1108-1150: LGTM: dump now includes model/provider rate limits for persistence.Including model/provider rate limits in the dump aligns persistence with the new hierarchy.
2041-2070: LGTM: delete paths clean up associated budgets/rate limits.Deletion logic consistently cleans related governance rows.
Also applies to: 2110-2132
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
074ece1 to
dbdfd7e
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
c478008 to
5eeeda1
Compare
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
dbdfd7e to
a63e698
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
5eeeda1 to
e61e6d3
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@framework/configstore/rdb.go`:
- Around line 718-728: The GetProviders method on RDBConfigStore contains an
unreachable check for gorm.ErrRecordNotFound after using GORM's Find (which
returns an empty slice instead of ErrRecordNotFound); remove the dead branch
that checks errors.Is(err, gorm.ErrRecordNotFound) and simply return the error
if non-nil (mirroring GetBudgets/GetRateLimits), i.e., in the GetProviders
function eliminate the ErrNotFound handling and keep the direct err return path
so callers receive the actual DB error or an empty providers slice.
- Around line 2051-2061: The ErrNotFound check in RDBConfigStore.GetModelConfigs
is unreachable because s.db.WithContext(...).Preload(...).Find(&modelConfigs)
returns no error for empty results; remove the errors.Is(err,
gorm.ErrRecordNotFound) branch and simply return modelConfigs (which may be an
empty slice) with nil error when Find succeeds, keeping the method signature and
otherwise-preserved error handling; update the function around GetModelConfigs
and the
s.db.WithContext(...).Preload("Budget").Preload("RateLimit").Find(&modelConfigs).Error
usage accordingly.
♻️ Duplicate comments (2)
plugins/governance/store.go (2)
446-463: Budget gating can be skipped if Budget isn’t preloaded.
CheckModelBudgetfilters onmc.Budget != nil, which will silently skip a budget if onlyBudgetIDis set (e.g., when relationships aren’t preloaded). That can allow policy bypass. Prefer gating onBudgetIDinstead, mirroring the rate‑limit path.✅ Suggested fix
- if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.Budget != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.BudgetID != nil { modelConfigsToCheck = append(modelConfigsToCheck, mc) budgetNames = append(budgetNames, fmt.Sprintf("Model:%s:Provider:%s", model, string(*provider))) } @@ - if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.Budget != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.BudgetID != nil { modelConfigsToCheck = append(modelConfigsToCheck, mc) budgetNames = append(budgetNames, fmt.Sprintf("Model:%s", model)) }
2109-2127: Preserve reset timestamps when rehydrating model/provider budgets & rate limits.Only
CurrentUsagefields are preserved right now. IfLastReset/TokenLastReset/RequestLastReset(and DB-tracking fields) are overwritten by stale values from config, limits can be incorrectly treated as unexpired.✅ Suggested fix
if clone.Budget != nil { if existingBudgetValue, exists := gs.budgets.Load(clone.Budget.ID); exists && existingBudgetValue != nil { if eb, ok := existingBudgetValue.(*configstoreTables.TableBudget); ok && eb != nil { clone.Budget.CurrentUsage = eb.CurrentUsage + clone.Budget.LastReset = eb.LastReset + clone.Budget.LastDBUsage = eb.LastDBUsage } } gs.budgets.Store(clone.Budget.ID, clone.Budget) } @@ if clone.RateLimit != nil { if existingRateLimitValue, exists := gs.rateLimits.Load(clone.RateLimit.ID); exists && existingRateLimitValue != nil { if erl, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && erl != nil { clone.RateLimit.TokenCurrentUsage = erl.TokenCurrentUsage clone.RateLimit.RequestCurrentUsage = erl.RequestCurrentUsage + clone.RateLimit.TokenLastReset = erl.TokenLastReset + clone.RateLimit.RequestLastReset = erl.RequestLastReset + clone.RateLimit.LastDBTokenUsage = erl.LastDBTokenUsage + clone.RateLimit.LastDBRequestUsage = erl.LastDBRequestUsage } } gs.rateLimits.Store(clone.RateLimit.ID, clone.RateLimit) }Also applies to: 2185-2204
🧹 Nitpick comments (2)
plugins/governance/store.go (2)
55-64: Standardize Check return ordering for consistency.*
CheckRateLimitreturns(Decision, error)while provider/model variants return(error, Decision). Consider aligning the return order across all Check* methods to reduce call‑site mistakes and follow Go’s “error last” convention.
433-439: Preferbifrost.Ptrfor provider pointers to match repo conventions.This keeps pointer creation consistent with the rest of the codebase.
♻️ Suggested refactor
import ( "context" "fmt" "strings" "sync" "time" + bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" "gorm.io/gorm" ) @@ - if request.Provider != "" { - provider = &request.Provider - } + if request.Provider != "" { + provider = bifrost.Ptr(request.Provider) + } @@ - if request.Provider != "" { - provider = &request.Provider - } + if request.Provider != "" { + provider = bifrost.Ptr(request.Provider) + }Based on learnings, prefer
bifrost.Ptr()over address-of for pointer creation.Also applies to: 523-529
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/configstore/rdb.go (1)
608-630: Wrap the multi-step provider deletion in a transaction to ensure atomicity.If the provider delete succeeds but a subsequent budget or rate-limit delete fails, you can leave orphan records. Apply the
localTxpattern already used elsewhere in this file (e.g.,UpdatePlugin), wrapping budget and rate-limit deletions in a transaction when notxis supplied, with proper rollback on error.💡 Suggested pattern
func (s *RDBConfigStore) DeleteProvider(ctx context.Context, provider schemas.ModelProvider, tx ...*gorm.DB) error { var txDB *gorm.DB + var localTx bool if len(tx) > 0 { txDB = tx[0] + localTx = false } else { - txDB = s.db + txDB = s.db.Begin() + localTx = true } // ... existing Find and Delete logic ... + if localTx { + if err != nil { + txDB.Rollback() + return err + } + return txDB.Commit().Error + } return nil }
🤖 Fix all issues with AI agents
In `@framework/configstore/migrations.go`:
- Around line 2356-2379: The migration currently only creates indexes
(idx_provider_budget, idx_provider_rate_limit) inside the blocks that add
columns, so if a previous run added the column but crashed before creating the
index the index never gets created; update the logic in the migration function
to always attempt to create the indexes after ensuring the columns exist: keep
the existing checks using migrator.HasColumn and migrator.AddColumn for
budget_id and rate_limit_id, but move the tx.Exec("CREATE INDEX IF NOT EXISTS
... ON config_providers (...)") and the migrator.HasIndex guards outside those
HasColumn blocks (or run them unconditionally) so idx_provider_budget and
idx_provider_rate_limit are created idempotently regardless of whether the
column was just added or already present.
♻️ Duplicate comments (2)
plugins/governance/store.go (2)
446-463: Avoid skipping model budgets when onlyBudgetIDis set.
The gating onmc.Budget != nilcan skip checks if the relation wasn’t preloaded butBudgetIDexists. Switching toBudgetID != nilkeeps behavior consistent with rate-limit gating.💡 Suggested adjustment
- if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.Budget != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.BudgetID != nil {
2109-2127: Preserve reset timestamps and DB-tracking fields when rehydrating budgets/rate limits.
Only copying usage counters can overwrite in-memory reset state and cause incorrect enforcement after config reloads. Carry overLastReset/TokenLastReset/RequestLastResetandLastDB*fields as well.💡 Suggested preservation
if eb, ok := existingBudgetValue.(*configstoreTables.TableBudget); ok && eb != nil { clone.Budget.CurrentUsage = eb.CurrentUsage + clone.Budget.LastReset = eb.LastReset + clone.Budget.LastDBUsage = eb.LastDBUsage } ... if erl, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && erl != nil { clone.RateLimit.TokenCurrentUsage = erl.TokenCurrentUsage clone.RateLimit.RequestCurrentUsage = erl.RequestCurrentUsage + clone.RateLimit.TokenLastReset = erl.TokenLastReset + clone.RateLimit.RequestLastReset = erl.RequestLastReset + clone.RateLimit.LastDBTokenUsage = erl.LastDBTokenUsage + clone.RateLimit.LastDBRequestUsage = erl.LastDBRequestUsage }Also applies to: 2195-2204
🧹 Nitpick comments (1)
plugins/governance/store.go (1)
433-439: Usebifrost.Ptrfor pointer creation for consistency.
The code currently uses&request.Provider; preferbifrost.Ptr(request.Provider)to match repo style.♻️ Suggested tweak
- if request.Provider != "" { - provider = &request.Provider - } + if request.Provider != "" { + provider = bifrost.Ptr(request.Provider) + }Based on learnings, use
bifrost.Ptr()for pointer creation consistency.Also applies to: 523-529
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/configstore/rdb.go (1)
608-630: Wrap provider + budget/rate‑limit deletes in an explicit transaction for consistency when notxis supplied.The current code relies on CASCADE constraints to ensure atomicity, which is safe. However, for consistency with transaction patterns elsewhere in the codebase (e.g.,
CreateMCPClientConfig,UpdateFrameworkConfig) and to make the transactional intent explicit, wrap the entire operation in a transaction when no external transaction is passed.Suggested approach
func (s *RDBConfigStore) DeleteProvider(ctx context.Context, provider schemas.ModelProvider, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } + doDelete := func(txDB *gorm.DB) error { + // existing body (find → delete provider → delete budget/rate limit) + return nil + } + if len(tx) == 0 { + return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + return doDelete(tx) + }) + } + return doDelete(tx[0]) - // existing body... }
🤖 Fix all issues with AI agents
In `@framework/configstore/tables/provider.go`:
- Around line 91-97: The validation for governance fields (p.BudgetID and
p.RateLimitID) currently only rejects exact empty strings; update the checks in
the provider model to trim whitespace and reject values that are empty after
trimming (i.e., treat whitespace-only IDs as invalid) for parity with
TableModelConfig—use strings.TrimSpace on *p.BudgetID and *p.RateLimitID, and
return the same fmt.Errorf("budget_id cannot be an empty string") /
fmt.Errorf("rate_limit_id cannot be an empty string") if the trimmed value is
empty.
In `@plugins/governance/store.go`:
- Around line 2109-2127: The code currently copies only usage counters from
existing in-memory entries into clone.Budget and clone.RateLimit but overwrites
their reset timestamp and DB-related reset metadata; update the merge logic in
the block that handles clone.Budget (and the analogous block for clone.RateLimit
around lines noted) to also preserve in-memory reset fields instead of replacing
them with DB values — when existingBudgetValue.(*configstoreTables.TableBudget)
is present copy eb.CurrentUsage into clone.Budget.CurrentUsage but also copy
eb.LastReset and any eb.LastDB* reset fields (e.g., LastDBReset or similar) back
into clone.Budget; likewise for
existingRateLimitValue.(*configstoreTables.TableRateLimit) copy
erl.TokenCurrentUsage and erl.RequestCurrentUsage into clone.RateLimit while
preserving erl.TokenLastReset, erl.RequestLastReset and any LastDB* fields by
assigning them into clone.RateLimit rather than letting the DB values overwrite
them.
🧹 Nitpick comments (1)
plugins/governance/store.go (1)
432-439: Preferbifrost.Ptr()for provider pointer creation.This repo’s convention is to avoid
&valuein favor ofbifrost.Ptr(value)for consistency.Suggested refactor
+import ( + bifrost "github.com/maximhq/bifrost/core" +) ... - provider = &request.Provider + provider = bifrost.Ptr(request.Provider)Based on learnings, use
bifrost.Ptr(...)consistently.Also applies to: 523-529
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
e61e6d3 to
eabf9e0
Compare
danpiths
force-pushed
the
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
branch
from
a63e698 to
8a05992
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/rdb.go (1)
608-630: Wrap provider deletion + budget/rate-limit cleanup in a transaction.
When notxis passed, the provider delete and budget/rate-limit deletes can partially succeed. That can leave orphan budgets or mismatched state.🐛 Proposed fix (transaction when no tx is supplied)
func (s *RDBConfigStore) DeleteProvider(ctx context.Context, provider schemas.ModelProvider, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } + if len(tx) == 0 { + return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + return s.DeleteProvider(ctx, provider, tx) + }) + } + txDB := tx[0] // Find the existing provider var dbProvider tables.TableProvider if err := txDB.WithContext(ctx).Where("name = ?", string(provider)).First(&dbProvider).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } @@ return nil }
♻️ Duplicate comments (5)
framework/configstore/tables/provider.go (1)
91-97: Reject whitespace-only governance IDs.Right now only
""is rejected; whitespace-only values can still slip through. Trim spaces to align with other governance validators.🛠️ Suggested fix
import ( "encoding/json" "fmt" + "strings" "time" @@ - 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") }framework/configstore/migrations.go (1)
2356-2380: Make index creation idempotent even when columns already exist.If a previous run added the column but failed before creating the index, reruns will skip index creation. Create indexes outside the
HasColumnblocks.🛠️ Suggested fix
// Add budget_id column if it doesn't exist if !migrator.HasColumn(provider, "budget_id") { if err := migrator.AddColumn(provider, "budget_id"); err != nil { return fmt.Errorf("failed to add budget_id column: %w", err) } - // Create index for budget_id - if !migrator.HasIndex(provider, "idx_provider_budget") { - if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_budget ON config_providers (budget_id)").Error; err != nil { - return fmt.Errorf("failed to create budget_id index: %w", err) - } - } } + // Create index for budget_id (also covers reruns where column already exists) + if !migrator.HasIndex(provider, "idx_provider_budget") { + if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_budget ON config_providers (budget_id)").Error; err != nil { + return fmt.Errorf("failed to create budget_id index: %w", err) + } + } // Add rate_limit_id column if it doesn't exist if !migrator.HasColumn(provider, "rate_limit_id") { if err := migrator.AddColumn(provider, "rate_limit_id"); err != nil { return fmt.Errorf("failed to add rate_limit_id column: %w", err) } - // 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 index for rate_limit_id (also covers reruns where column already exists) + 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) + } + }framework/configstore/rdb.go (1)
718-725: Duplicate:Find()doesn’t return ErrRecordNotFound.
TheErrNotFoundbranches afterFind()are unreachable. This matches earlier review feedback.Please confirm with GORM docs:
GORM Find ErrRecordNotFound behavior for empty resultsAlso applies to: 2051-2057
plugins/governance/store.go (2)
446-462: Duplicate: gate onBudgetID, notBudget, in model-config checks.
If the relation isn’t preloaded,mc.Budgetis nil and the check is skipped even whenBudgetIDis set. This matches prior review feedback.
2109-2125: Duplicate: preserve reset timestamps/LastDB when rehydrating budgets & limits.*
This matches earlier feedback—only usage counters are preserved today.Also applies to: 2185-2202
🧹 Nitpick comments (2)
plugins/governance/resolver.go (1)
263-266: Clarify “provider-level” wording in this helper comment.
checkRateLimitHierarchyis operating on VK provider configs (and VK), not global provider-level rate limits. Tweaking the comment would avoid confusion.plugins/governance/store.go (1)
432-439: Preferbifrost.Ptr(...)for pointer creation.
Use the shared helper instead of&valuefor consistency.Based on learnings, prefer `bifrost.Ptr()` for pointer creation.♻️ Suggested change
import ( "context" "fmt" "strings" "sync" "time" + bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" "gorm.io/gorm" ) @@ if request != nil { model = request.Model if request.Provider != "" { - provider = &request.Provider + provider = bifrost.Ptr(request.Provider) } } @@ if request != nil { model = request.Model if request.Provider != "" { - provider = &request.Provider + provider = bifrost.Ptr(request.Provider) } }Also applies to: 523-529
…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
eabf9e0 to
d166de7
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@framework/configstore/migrations.go`:
- Around line 2347-2417: The migrationAddProviderGovernanceColumns migration
currently only adds budget_id and rate_limit_id columns and indexes but does not
create FK constraints, so update the Migrate function to explicitly create
foreign key constraints (with ON DELETE CASCADE) referencing the appropriate
tables after adding the columns and indexes (e.g., create CONSTRAINT
fk_provider_budget FOREIGN KEY(budget_id) REFERENCES config_budgets(id) ON
DELETE CASCADE and similarly for rate_limit_id), and update the Rollback
function to drop those constraints before dropping indexes/columns (DROP
CONSTRAINT IF EXISTS fk_provider_budget, fk_provider_rate_limit). Locate
TableProvider, budget_id, rate_limit_id, idx_provider_budget,
idx_provider_rate_limit and add/create/drop the named constraints to ensure
parity with fresh installs.
♻️ Duplicate comments (2)
plugins/governance/store.go (2)
446-463: Don’t gate model budget checks on preloaded Budget relations.Line 450 and Line 460 check
mc.Budget != nil, which can skip enforcement when onlyBudgetIDis present (e.g., relationship not preloaded). Gate onBudgetIDinstead so model budgets are always considered.✅ Suggested fix
- if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.Budget != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.BudgetID != nil { modelConfigsToCheck = append(modelConfigsToCheck, mc) budgetNames = append(budgetNames, fmt.Sprintf("Model:%s:Provider:%s", model, string(*provider))) }- if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.Budget != nil { + if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil && mc.BudgetID != nil { modelConfigsToCheck = append(modelConfigsToCheck, mc) budgetNames = append(budgetNames, fmt.Sprintf("Model:%s", model)) }
2109-2127: Preserve reset timestamps and DB counters when rehydrating model/provider budgets & limits.Only copying
CurrentUsage(and token/request usage) risks overwriting in-memory reset metadata with stale DB values, which can skew expiry decisions after config refresh. PreserveLastResetandLastDB*fields for budgets and rate limits.✅ Suggested fix
if clone.Budget != nil { if existingBudgetValue, exists := gs.budgets.Load(clone.Budget.ID); exists && existingBudgetValue != nil { if eb, ok := existingBudgetValue.(*configstoreTables.TableBudget); ok && eb != nil { clone.Budget.CurrentUsage = eb.CurrentUsage + clone.Budget.LastReset = eb.LastReset + clone.Budget.LastDBUsage = eb.LastDBUsage } } gs.budgets.Store(clone.Budget.ID, clone.Budget) } @@ if clone.RateLimit != nil { if existingRateLimitValue, exists := gs.rateLimits.Load(clone.RateLimit.ID); exists && existingRateLimitValue != nil { if erl, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && erl != nil { clone.RateLimit.TokenCurrentUsage = erl.TokenCurrentUsage clone.RateLimit.RequestCurrentUsage = erl.RequestCurrentUsage + clone.RateLimit.TokenLastReset = erl.TokenLastReset + clone.RateLimit.RequestLastReset = erl.RequestLastReset + clone.RateLimit.LastDBTokenUsage = erl.LastDBTokenUsage + clone.RateLimit.LastDBRequestUsage = erl.LastDBRequestUsage } } gs.rateLimits.Store(clone.RateLimit.ID, clone.RateLimit) }Also applies to: 2185-2204
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
2115-2129: Consider transaction wrapper for UpdateModelConfigs to ensure atomicity.If any
Savefails mid-iteration, the database will be left in a partially updated state. Consider wrapping in a transaction when no externaltxis provided, similar to other batch operations.♻️ Proposed refactor for atomicity
func (s *RDBConfigStore) UpdateModelConfigs(ctx context.Context, modelConfigs []*tables.TableModelConfig, tx ...*gorm.DB) error { var txDB *gorm.DB if len(tx) > 0 { txDB = tx[0] - } else { - txDB = s.db - } - for _, mc := range modelConfigs { - if err := txDB.WithContext(ctx).Save(mc).Error; err != nil { - return s.parseGormError(err) + for _, mc := range modelConfigs { + if err := txDB.WithContext(ctx).Save(mc).Error; err != nil { + return s.parseGormError(err) + } } + return nil } - return nil + return s.db.WithContext(ctx).Transaction(func(localTx *gorm.DB) error { + for _, mc := range modelConfigs { + if err := localTx.Save(mc).Error; err != nil { + return s.parseGormError(err) + } + } + return nil + }) }
Merge activity
|
akshaydeo
changed the base branch from
feature/12-12-_nix_support_adds_nix_flake_support_for_bifrost
to
graphite-base/1120
akshaydeo
deleted the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
3 participants
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