chore: adds test for governance plugin's model configs and providers budgeting and rate limiting #1121
Conversation
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
📝 WalkthroughSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughAdds test helpers and a comprehensive governance test suite for provider/model budgets, rate limits, hooks, and virtual-key flows; updates the BudgetResolver API from EvaluateRequest(...) to EvaluateVirtualKeyRequest(virtualKey, provider, model, requestID) and updates tests accordingly. Changes
Sequence Diagram(s)(Skipped — changes are test additions and an API signature update; no new multi-component runtime control flow requiring visualization.) Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: Organization UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (3)
🧰 Additional context used📓 Path-based instructions (1)**⚙️ CodeRabbit configuration file
Files:
🧠 Learnings (2)📚 Learning: 2025-12-09T17:07:42.007ZApplied to files:
📚 Learning: 2025-12-22T10:50:40.990ZApplied to files:
🧬 Code graph analysis (2)plugins/governance/fixtures_test.go (1)
plugins/governance/resolver_test.go (1)
⏰ 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)
🔇 Additional comments (3)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
plugins/governance/model_provider_governance_test.go (2)
1714-1717: Consider more robust async wait pattern.The
time.Sleep(200 * time.Millisecond)pattern for waiting on async processing can lead to flaky tests under CI load. If the async operations expose any synchronization primitives (channels, wait groups), consider using those instead. Otherwise, a polling loop with timeout would be more resilient.
724-727: Standardize loop syntax across test file for consistency.The file mixes loop styles: lines 534–548 use traditional
for i := 0; i < 500; i++syntax while lines 724–727, 735–738, 764–767, and 775–778 usefor range 500syntax. Since the project targets Go 1.24.3, thefor range Nsyntax is fully supported. Standardize on one style throughout this file.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.go(1 hunks)plugins/governance/model_provider_governance_test.go(1 hunks)plugins/governance/resolver_test.go(15 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:
plugins/governance/resolver_test.goplugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.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:
plugins/governance/resolver_test.goplugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.go
🧬 Code graph analysis (2)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
plugins/governance/model_provider_governance_test.go (5)
framework/configstore/clientconfig.go (1)
GovernanceConfig(705-714)core/schemas/bifrost.go (2)
OpenAI(35-35)Azure(36-36)plugins/governance/resolver.go (4)
DecisionAllow(18-18)DecisionTokenLimited(23-23)DecisionRequestLimited(24-24)DecisionRateLimited(21-21)plugins/governance/main.go (1)
InitFromStore(163-209)core/schemas/chatcompletions.go (1)
BifrostChatRequest(12-19)
⏰ 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). (4)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (5)
plugins/governance/fixtures_test.go (1)
222-261: LGTM! Well-structured test helpers.The new helper functions (
buildModelConfig,buildProviderWithGovernance,boolPtr) follow the established patterns in this file and provide clean abstractions for constructing test data. The consistent use of conditional attachment for optional Budget/RateLimit fields is appropriate.plugins/governance/resolver_test.go (1)
28-28: API migration toEvaluateVirtualKeyRequestlooks correct.The tests have been properly updated to use the new
EvaluateVirtualKeyRequest(ctx, virtualKey, provider, model, requestID)signature. The migration is consistent across all test cases, and the test assertions remain valid.plugins/governance/model_provider_governance_test.go (3)
1-13: Comprehensive test suite with good organization.The test file is well-structured with clear section headers and covers a wide range of governance scenarios including provider/model budgets, rate limits, and VK interactions. The test helpers from
fixtures_test.goare used effectively.
1091-1116: Good PreHook integration test coverage.The PreHook tests properly verify the short-circuit behavior when governance limits are exceeded. The test setup with
boolPtr(false)forIsVkMandatoryand the assertion onshortCircuit.Error.Error.Messageeffectively validates the expected behavior.
1667-1733: Good PostHook integration test with appropriate documentation.The test correctly acknowledges the limitation that "Without model catalog, cost will be 0" and focuses on verifying the flow works correctly. The PreHook → PostHook → PreHook pattern effectively validates the usage tracking integration.
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-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
f952287 to
311c5f9
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
plugins/governance/model_provider_governance_test.go (2)
1716-1716: Consider using synchronization instead of fixed sleep for async tests.The PostHook tests use
time.Sleep(200 * time.Millisecond)to wait for async processing. While this is a common pattern, it can be flaky under load or on slower CI runners. If there's a way to synchronize on completion (e.g., a channel or condition variable in the tracker), that would be more reliable.That said, for unit tests with a 200ms buffer, this is typically acceptable.
Also applies to: 1785-1785, 1852-1852, 1921-1921
534-548: Minor style inconsistency in loop syntax (optional modernization).The file uses both traditional
for i := 0; i < 500; i++loops (lines 534-537, 545-548) and the Go 1.22+for range 500syntax (lines 724, 735, 764, 775). While both are valid, the range syntax is now the preferred form for iteration in Go. Consider adopting a consistent style throughout the file by preferringfor range Nonce Go 1.22 is the project's minimum version.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- plugins/governance/fixtures_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.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_test.goplugins/governance/model_provider_governance_test.go
🧬 Code graph analysis (2)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
plugins/governance/model_provider_governance_test.go (6)
framework/configstore/clientconfig.go (1)
GovernanceConfig(720-729)core/schemas/bifrost.go (2)
OpenAI(35-35)Azure(36-36)plugins/governance/resolver.go (6)
DecisionAllow(18-18)DecisionTokenLimited(23-23)DecisionRequestLimited(24-24)DecisionRateLimited(21-21)NewBudgetResolver(70-75)DecisionBudgetExceeded(22-22)core/schemas/context.go (2)
NewBifrostContext(46-64)NoDeadline(10-10)plugins/governance/main.go (1)
InitFromStore(178-224)core/schemas/chatcompletions.go (1)
BifrostChatRequest(12-19)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (2)
plugins/governance/model_provider_governance_test.go (1)
1-13: Comprehensive test coverage for model and provider governance.This test file provides excellent coverage for the new governance features:
- Store-level tests for provider/model budget and rate limit checks
- Update operation tests for usage tracking
- Resolver tests for
EvaluateModelAndProviderRequest- End-to-end PreHook/PostHook integration tests
- Virtual key fallback scenarios
The tests are well-organized with clear section headers and cover both success and failure paths.
plugins/governance/resolver_test.go (1)
28-28: API migration toEvaluateVirtualKeyRequestlooks correct.The migration from
EvaluateRequesttoEvaluateVirtualKeyRequestis consistent across all test cases. The new signature(ctx, virtualKey, provider, model, requestID)is properly used with:
- Virtual key strings matching the test setup
schemas.OpenAIas the provider- Appropriate model names
- Request IDs where needed (empty string when not required)
Also applies to: 43-43, 61-61, 85-85, 117-117, 139-139, 162-162, 200-200, 222-222, 249-249, 284-284, 295-295, 317-317, 340-340, 478-478
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
plugins/governance/model_provider_governance_test.go (1)
1716-1716: Consider replacingtime.Sleepwith proper synchronization.Using
time.Sleepto wait for async processing can lead to flaky tests:
- 200ms may be insufficient in CI or under load
- Tests will be slower than necessary when processing completes faster
- No guarantee that processing is complete after the sleep
Consider alternatives:
- Use channels or wait groups to signal completion
- Add a test-only synchronous mode to the plugin
- Implement a retry mechanism with timeout
Example: Using retry with timeout
- // Wait for async processing to complete - time.Sleep(200 * time.Millisecond) + // Wait for async processing to complete with retry + require.Eventually(t, func() bool { + _, sc, _ := plugin.PreHook(ctx2, req2) + return sc != nil && strings.Contains(sc.Error.Error.Message, "token limit exceeded") + }, 1*time.Second, 50*time.Millisecond, "Expected token limit to be exceeded after update")Also applies to: 1785-1785, 1852-1852, 1921-1921
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- plugins/governance/fixtures_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.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_test.goplugins/governance/model_provider_governance_test.go
🧬 Code graph analysis (1)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
⏰ 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). (9)
- GitHub 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)
plugins/governance/model_provider_governance_test.go (7)
19-84: LGTM! Comprehensive provider budget test coverage.The test cases cover all critical scenarios: no config, no budget, within limits, exceeded, and with baseline calculations. The test structure is clear and assertions are descriptive.
90-174: LGTM! Complete provider rate limit test coverage.The tests properly verify token limits, request limits, and the combined case. The decision hierarchy is correctly tested (DecisionRateLimited when both are exceeded).
180-291: LGTM! Thorough model budget test coverage.The tests excellently cover the nuanced interactions between model-only and model+provider configurations, as well as provider-specific filtering. The test at lines 254-273 correctly verifies that both model-only and model+provider budgets are checked when both exist.
297-439: LGTM! Complete model rate limit test coverage.The tests mirror the comprehensive structure of the model budget tests, properly covering token limits, request limits, and the interaction between model-only and model+provider rate limit configurations.
445-785: LGTM! Comprehensive update operation test coverage.The tests properly verify that usage updates are correctly tracked and reflected in subsequent checks. The loop-based tests (lines 534-548, 724-785) effectively simulate incremental usage accumulation.
791-1085: LGTM! Thorough resolver integration tests.The resolver tests effectively verify the integration of provider and model checks, including the correct evaluation order (provider checks before model checks) and proper handling of provider-specific configurations.
1091-1661: LGTM! Comprehensive PreHook integration tests.The end-to-end PreHook tests thoroughly cover the interaction between model/provider governance and virtual key governance, including proper short-circuit behavior and the cascade of checks (provider → model → VK).
plugins/governance/resolver_test.go (1)
28-28: LGTM! Consistent API migration to EvaluateVirtualKeyRequest.All test calls have been correctly updated to use the new explicit parameter-based API. The changes are consistent throughout the file and improve clarity by making required parameters explicit rather than struct-based.
The mixture of specific request IDs (e.g., "req-123") and empty strings suggests requestID is optional, which is reasonable for different test scenarios.
Also applies to: 43-43, 61-61, 85-85, 117-117, 139-139, 162-162, 200-200, 222-222, 249-249, 284-284, 295-295, 317-317, 340-340, 478-478
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
311c5f9 to
41cd83b
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
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
2 times, most recently
from
69b5088 to
13222e0
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
98c8540 to
de01994
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
13222e0 to
4d4f403
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
de01994 to
4b1d30d
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
plugins/governance/fixtures_test.go (1)
259-261: Simple pointer helper - verify for duplication.The
boolPtrhelper is correctly implemented and useful for tests. However, verify whether a similar utility already exists in the codebase (e.g., in atestutilpackage or similar) to avoid duplication.plugins/governance/model_provider_governance_test.go (1)
1667-1937: PostHook tests verify usage tracking - consider async handling improvement.The PostHook integration tests effectively verify that usage updates are applied and reflected in subsequent requests. However, the tests use
time.Sleep(200 * time.Millisecond)at lines 1716, 1785, 1852, and 1921 to wait for async processing to complete.This approach may be flaky in CI environments with variable performance. Consider implementing one of these alternatives:
- Add a synchronous flush method to the store for testing purposes
- Use polling with a timeout to verify the update completed
- Make the async behavior configurable to allow synchronous updates in tests
Alternative approaches for async testing
Option 1: Add test helper for synchronous flush
// In store implementation (test-only) func (s *LocalGovernanceStore) FlushPendingUpdates() { // Wait for all pending async operations to complete } // In tests plugin.PostHook(ctx1, result1, nil) store.FlushPendingUpdates() // Instead of time.SleepOption 2: Polling with timeout
// Replace time.Sleep with polling require.Eventually(t, func() bool { _, decision := store.CheckProviderRateLimit(...) return decision == DecisionTokenLimited }, 1*time.Second, 50*time.Millisecond, "Usage should be updated")
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/fixtures_test.goplugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/fixtures_test.goplugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.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/fixtures_test.goplugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.go
🧬 Code graph analysis (3)
plugins/governance/fixtures_test.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
plugins/governance/model_provider_governance_test.go (9)
plugins/governance/fixtures_test.go (1)
NewMockLogger(24-32)plugins/governance/store.go (1)
NewLocalGovernanceStore(95-114)framework/configstore/clientconfig.go (1)
GovernanceConfig(720-729)core/schemas/bifrost.go (7)
OpenAI(35-35)Azure(36-36)BifrostRequest(165-185)BifrostContextKeyVirtualKey(121-121)BifrostContextKeyRequestID(123-123)BifrostResponse(327-347)BifrostResponseExtraFields(395-406)plugins/governance/resolver.go (6)
DecisionAllow(18-18)DecisionTokenLimited(23-23)DecisionRequestLimited(24-24)DecisionRateLimited(21-21)NewBudgetResolver(70-75)DecisionBudgetExceeded(22-22)core/schemas/context.go (1)
NewBifrostContext(46-64)core/schemas/chatcompletions.go (3)
BifrostChatRequest(12-19)BifrostChatResponse(27-42)BifrostLLMUsage(845-852)core/schemas/models.go (1)
Model(109-129)examples/plugins/hello-world/main.go (2)
PreHook(24-30)PostHook(32-39)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (9)
plugins/governance/fixtures_test.go (2)
223-240: LGTM! Test helper correctly wires model config relationships.The helper properly constructs a
TableModelConfigwith optional Budget and RateLimit associations, correctly setting both the entity pointer and the corresponding ID field when non-nil. This follows the same pattern as existing helpers likebuildVirtualKeyWithBudget.
242-257: LGTM! Provider builder follows established patterns.The helper correctly constructs a
TableProviderwith optional governance settings, maintaining proper relationships between the provider and its budget/rate limit configurations.plugins/governance/resolver_test.go (1)
28-490: API migration looks correct - verify empty requestID behavior.The migration from
EvaluateRequest(ctx, &EvaluationRequest{...})toEvaluateVirtualKeyRequest(ctx, virtualKey, provider, model, requestID)is consistently applied across all test cases. However, many tests pass an empty string""for therequestIDparameter (e.g., lines 43, 61, 85, 117, etc.).Please verify that passing an empty
requestIDis the expected behavior for these test scenarios and doesn't bypass any important request tracking logic.plugins/governance/model_provider_governance_test.go (6)
19-84: Comprehensive provider budget test coverage.The provider budget tests cover all major scenarios including no config, no budget, within limits, exceeded, and baseline checks. Test isolation is properly maintained with fresh store creation for each test case.
90-174: Well-structured provider rate limit tests.The test suite comprehensively covers provider rate limiting scenarios including token limits, request limits, and combined limits. The use of decision enums for assertions provides clear verification of expected outcomes.
180-439: Excellent coverage of model-level governance.The model budget and rate limit tests thoroughly cover various scenarios including model-only configurations, model+provider configurations, and provider-specific filtering. The tests correctly verify that provider-specific model configs don't affect requests using different providers (e.g., lines 275-291, 403-439).
445-785: Update operation tests verify full cycle.The tests for update operations (budget usage and rate limit usage) correctly verify that updates are applied and subsequent checks reflect the updated values. This ensures the governance tracking works end-to-end.
791-1085: Resolver tests verify correct evaluation order.The
EvaluateModelAndProviderRequesttests comprehensively verify the resolver's evaluation logic, including the correct ordering of checks (provider-level before model-level) and proper handling of various configuration scenarios. Tests at lines 893-915 and 917-963 specifically verify that provider-level failures are caught before model-level checks.
1091-1624: PreHook integration tests cover end-to-end flows.The PreHook integration tests comprehensively verify the plugin's behavior in real-world scenarios, including requests without virtual keys and with VK fallback. The tests correctly verify short-circuit behavior when governance checks fail (e.g., lines 1114-1116, 1141-1143).
The VK fallback tests (lines 1385-1624) effectively demonstrate the layered governance model where model/provider checks happen first, followed by VK-specific checks.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
plugins/governance/model_provider_governance_test.go (1)
1716-1716: Consider replacing sleep-based synchronization with deterministic waits.The PostHook tests use
time.Sleep(200 * time.Millisecond)to wait for async processing to complete. This approach can lead to flaky tests if the processing takes longer than expected or introduces unnecessary delays in CI.🔧 Consider using synchronization primitives or test hooks
Options to improve test reliability:
- Add a synchronous test mode for the governance plugin that processes updates synchronously
- Use channels or wait groups to signal when async processing is complete
- Implement a test helper that polls for the expected state with a timeout
- Add a flush/drain method to the plugin that waits for pending operations
Example polling approach:
// Helper function func waitForBudgetUpdate(t *testing.T, store *LocalGovernanceStore, provider schemas.ModelProvider, expectedExceeded bool) { timeout := time.After(500 * time.Millisecond) ticker := time.NewTicker(10 * time.Millisecond) defer ticker.Stop() for { select { case <-timeout: t.Fatal("Timeout waiting for budget update") case <-ticker.C: err := store.CheckProviderBudget(context.Background(), provider, nil) if (err != nil) == expectedExceeded { return } } } }Also applies to: 1785-1785, 1852-1852, 1921-1921
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/fixtures_test.goplugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/fixtures_test.goplugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.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/fixtures_test.goplugins/governance/resolver_test.goplugins/governance/model_provider_governance_test.go
🧬 Code graph analysis (1)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
⏰ 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 (5)
plugins/governance/resolver_test.go (1)
28-28: LGTM! API migration is consistent.The test updates correctly migrate from the old
EvaluateRequestAPI to the newEvaluateVirtualKeyRequestAPI signature. The virtual key, provider, model, and requestID parameters are properly passed in all test cases.Also applies to: 43-43, 61-61, 85-85, 117-117, 139-139, 162-162, 200-200, 222-222, 249-249, 284-284, 295-295, 317-317, 340-340, 478-478
plugins/governance/fixtures_test.go (3)
223-240: LGTM! Helper follows existing patterns.The
buildModelConfighelper function is well-structured and consistent with existing test fixture builders. It properly handles optional budget and rate limit parameters, setting both the object reference and the ID field when non-nil.
242-257: LGTM! Helper follows existing patterns.The
buildProviderWithGovernancehelper function mirrors the pattern used inbuildModelConfigand existing helpers, providing a clean way to construct test providers with optional governance settings.
259-261: LGTM! Standard Go helper pattern.The
boolPtrhelper is a standard Go pattern for creating boolean pointers, useful for optional boolean fields in structs.plugins/governance/model_provider_governance_test.go (1)
1-1937: LGTM! Comprehensive test coverage for governance features.This test file provides excellent coverage of the new model and provider governance capabilities, including:
- Provider and model budget/rate limit checking
- Update operations for budgets and rate limits
- Resolver evaluation logic
- PreHook and PostHook integration
- Edge cases like provider-specific configs and virtual key fallback scenarios
The test organization is clear with well-defined sections, and test names are descriptive.
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
4b1d30d to
8171345
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
4d4f403 to
7167ba6
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
plugins/governance/model_provider_governance_test.go (1)
1352-1379: Consider shorter test names for readability.Some test names are quite long (e.g.,
TestPreHook_ProviderSpecificModelRateLimit_DifferentProvider_Passes_NoVirtualKey_RequestLimit). While descriptive, consider abbreviating or restructuring for better readability. For example:
TestPreHook_NoVK_ProviderSpecificModelRL_DifferentProvider_RequestLimitThis is a minor style consideration and doesn't affect test functionality.
Also applies to: 1067-1085, 422-439
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/model_provider_governance_test.goplugins/governance/fixtures_test.goplugins/governance/resolver_test.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/model_provider_governance_test.goplugins/governance/fixtures_test.goplugins/governance/resolver_test.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/model_provider_governance_test.goplugins/governance/fixtures_test.goplugins/governance/resolver_test.go
🧬 Code graph analysis (3)
plugins/governance/model_provider_governance_test.go (4)
plugins/governance/store.go (1)
NewLocalGovernanceStore(98-117)core/schemas/bifrost.go (4)
OpenAI(35-35)Azure(36-36)BifrostContextKeyVirtualKey(121-121)BifrostContextKeyRequestID(123-123)plugins/governance/resolver.go (5)
DecisionAllow(18-18)DecisionTokenLimited(23-23)DecisionRequestLimited(24-24)DecisionRateLimited(21-21)NewBudgetResolver(70-75)core/schemas/context.go (1)
NewBifrostContext(46-64)
plugins/governance/fixtures_test.go (1)
ui/lib/types/governance.ts (2)
Budget(5-11)RateLimit(13-25)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
⏰ 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 (6)
plugins/governance/fixtures_test.go (3)
223-240: LGTM! Well-structured test helper.The
buildModelConfigfunction follows the existing fixture builder patterns and correctly handles optional Budget and RateLimit associations.
242-257: LGTM! Consistent helper implementation.The
buildProviderWithGovernancefunction mirrors the structure ofbuildModelConfigand provides clear fixture setup for provider governance testing.
259-261: LGTM! Standard Go testing utility.The
boolPtrhelper is a common Go idiom for creating boolean pointers in tests.plugins/governance/resolver_test.go (1)
28-28: LGTM! Consistent API migration.All test invocations have been updated to use the new
EvaluateVirtualKeyRequestAPI signature with explicit positional parameters. The changes are mechanical and preserve the original test coverage and intent.Also applies to: 43-43, 61-61, 85-85, 117-117, 139-139, 162-162, 200-200, 222-222, 249-249, 284-284, 295-295, 317-317, 340-340, 478-478
plugins/governance/model_provider_governance_test.go (2)
1-13: LGTM! Excellent test organization.The test file is well-structured with clear section headers organizing tests into logical categories: Store tests (budgets and rate limits), Resolver tests, PreHook integration, and PostHook integration. This makes the comprehensive test suite highly navigable.
Also applies to: 15-17, 86-88, 176-178, 293-295, 441-443, 557-559, 633-635, 787-789, 1087-1089, 1381-1383, 1663-1665
19-785: LGTM! Comprehensive test coverage.The test suite provides excellent coverage across multiple dimensions:
- Store tests: All CRUD operations for provider/model budgets and rate limits, including edge cases
- Resolver tests: Request evaluation with various governance constraint combinations and ordering verification
- PreHook integration: Both standalone model/provider governance and interaction with virtual key governance
- PostHook integration: Usage tracking and propagation to subsequent requests
The progressive structure from unit → integration → end-to-end tests makes the suite robust and maintainable.
Also applies to: 791-1085, 1091-1661, 1667-1937
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
plugins/governance/model_provider_governance_test.go (1)
1716-1716: Consider replacing sleep-based synchronization with more reliable approach.The tests use
time.Sleep(200 * time.Millisecond)to wait for async processing to complete. This is a common test anti-pattern that can lead to flaky tests in CI environments, especially under load:
- 200ms might be insufficient in resource-constrained environments
- Tests will always wait the full duration even if processing completes faster
- No guarantee that processing is actually complete
Consider alternatives like:
- Exposing a synchronous test mode for the async operations
- Using channels to signal completion
- Polling with a reasonable timeout
Also applies to: 1785-1785, 1852-1852, 1921-1921
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
plugins/governance/fixtures_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- plugins/governance/fixtures_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/governance/model_provider_governance_test.goplugins/governance/resolver_test.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/model_provider_governance_test.goplugins/governance/resolver_test.go
🧬 Code graph analysis (1)
plugins/governance/resolver_test.go (1)
core/schemas/bifrost.go (1)
OpenAI(35-35)
⏰ 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 (2)
plugins/governance/resolver_test.go (1)
28-28: LGTM! API migration is consistent and complete.The update from
EvaluateRequestwith a structured request object toEvaluateVirtualKeyRequestwith explicit parameters (virtualKey, provider, model, requestID) has been applied consistently across all test cases. The new API signature is more explicit and easier to use.Also applies to: 43-43, 61-61, 85-85, 117-117, 139-139, 162-162, 200-200, 222-222, 249-249, 284-284, 295-295, 317-317, 340-340, 478-478
plugins/governance/model_provider_governance_test.go (1)
1-1937: Excellent comprehensive test coverage for governance features.This new test file provides thorough coverage of the model and provider governance functionality with ~80+ test cases organized into clear sections:
- Store-level tests for budget and rate limit checking
- Store-level tests for usage updates
- Resolver-level evaluation tests
- PreHook integration tests (with and without virtual keys)
- PostHook integration tests (usage tracking)
The tests are well-structured, properly isolated, and cover numerous edge cases including provider-specific configurations, cross-provider scenarios, and evaluation order verification.
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
7167ba6 to
ed45bf7
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
8171345 to
360b6a9
Compare
akshaydeo
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
360b6a9 to
c33348c
Compare
akshaydeo
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
ed45bf7 to
85fa076
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
85fa076 to
ff67ef4
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
5a2256c to
509aaee
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
ff67ef4 to
b128177
Compare
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
509aaee to
b72e947
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
b128177 to
0ba0d94
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
b72e947 to
d7c5209
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
0ba0d94 to
ffb1435
Compare
…budgeting and rate limiting
danpiths
force-pushed
the
feature/12-17-feat_add_support_for_per-model_and_per-provider_level_budgeting_and_rate_limiting_in_governance_plugin
branch
from
d7c5209 to
d6c6b20
Compare
danpiths
force-pushed
the
feature/12-17-chore_adds_test_for_governance_plugin_s_model_configs_and_providers_budgeting_and_rate_limiting
branch
from
ffb1435 to
03f1cc1
Compare
Summary
Added comprehensive model and provider governance capabilities to the governance
plugin, enabling budget and rate limit enforcement at both the model and
provider levels.
Changes
settings
system
Type of change
Affected areas
How to test
Breaking changes
Related issues
none
Security considerations
none
Checklist
docs/contributing/README.mdand followed the guidelines