feat: add virtual key provider configs to budget/rate limit status and sort results#1535
Merged
Pratham-Mishra04 merged 1 commit intomainfrom Feb 4, 2026
Merged
Conversation
Contributor
|
Caution Review failedThe pull request is closed. 📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR makes governance budget/rate-limit checks virtual-key aware, moves VK-specific routing actions to occur after routing decisions, enforces deterministic ordering by Changes
Sequence DiagramsequenceDiagram
participant Client
participant GovPlugin as Governance Plugin
participant GovStore as Governance Store
participant ConfigStore as ConfigStore/DB
Client->>GovPlugin: Request (model, provider, optional virtualKey)
GovPlugin->>GovPlugin: Evaluate routing rules (routing program)
GovPlugin->>GovStore: GetBudgetAndRateLimitStatus(model, provider, virtualKey)
GovStore->>ConfigStore: Fetch model/provider configs (ordered by created_at)
ConfigStore-->>GovStore: Return configs
alt virtualKey present
GovStore->>ConfigStore: Fetch virtual-key specific configs
ConfigStore-->>GovStore: Return VK configs
end
GovStore->>GovStore: Aggregate model/provider/VK baselines and limits
GovStore-->>GovPlugin: Return BudgetAndRateLimitStatus
alt virtualKey present
GovPlugin->>GovPlugin: Perform VK load-balancing and inject MCP headers
end
GovPlugin-->>Client: Return routed response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
This was referenced Feb 3, 2026
Collaborator
Author
18 tasks
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@plugins/governance/store.go`:
- Around line 2915-2964: The code currently reads rate-limit data directly from
pc.RateLimit which can be stale; instead, when iterating vk.ProviderConfigs in
the block handling provider == string(provider), look up the canonical rate
limit from gs.rateLimits using pc.RateLimit.ID (or pc.RateLimitID/RateLimitID
field) and use that object for TokenCurrentUsage/TokenMaxLimit and
RequestCurrentUsage/RequestMaxLimit calculations, falling back to pc.RateLimit
only if the lookup misses or is nil; update references in the
tokenPercent/requestPercent calculations to use this fetched rateLimit and keep
the existing baseline and result logic unchanged (refer to
UpdateVirtualKeyRateLimitUsageInMemory, gs.rateLimits, pc.RateLimit,
ProviderConfigs, and vk).
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 2522-2532: When applying the update, compute the resulting scope
and scope_id (apply incoming req.Scope and req.ScopeID on top of existing
rule.Scope and rule.ScopeID), then validate: if the resulting scope != "global"
ensure the resulting scope_id is non-nil and non-empty and return a 400 error on
violation; if the resulting scope == "global" set rule.ScopeID = nil, otherwise
set rule.ScopeID to the validated value (use req.Scope, req.ScopeID, rule.Scope,
rule.ScopeID to locate the fields). Ensure you perform this check before
persisting or calling the DB so invalid inputs return 400 rather than causing a
DB error.
🧹 Nitpick comments (2)
framework/configstore/rdb.go (1)
1287-1303: Add a stable tiebreaker to CreatedAt ordering.Ordering only by
created_atcan still be nondeterministic when timestamps collide (e.g., batch inserts). Consider adding a secondary key likeidfor stable ordering; same idea applies to the other new Order clauses in this file. Also, is this PR part of a stack? If yes, please share the stack context so I can review across it.🔧 Suggested tweak
- Order("created_at ASC"). + Order("created_at ASC, id ASC").As per coding guidelines: 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).
transports/bifrost-http/handlers/governance.go (1)
294-296: Consider a deterministic tie-breaker when CreatedAt matches.Map iteration order is random; if multiple virtual keys share the same
CreatedAt, ordering can still vary. AddingIDas a secondary key makes the in-memory ordering fully stable.🔧 Suggested tweak
- sort.Slice(virtualKeys, func(i, j int) bool { - return virtualKeys[i].CreatedAt.Before(virtualKeys[j].CreatedAt) - }) + sort.Slice(virtualKeys, func(i, j int) bool { + if virtualKeys[i].CreatedAt.Equal(virtualKeys[j].CreatedAt) { + return virtualKeys[i].ID < virtualKeys[j].ID + } + return virtualKeys[i].CreatedAt.Before(virtualKeys[j].CreatedAt) + })
Pratham-Mishra04
force-pushed
the
02-03-fix_minor_routing_rules_fixes
branch
from
e05b657 to
6ec3119
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@plugins/governance/main.go`:
- Around line 349-365: The code assumes req.Headers is non-nil and will panic
when assigning headers; change how p.addMCPIncludeTools is used so it safely
merges headers: call p.addMCPIncludeTools with the existing req.Headers (or nil)
and virtualKey, check the returned headers map, initialize req.Headers if it's
nil, then copy/merge each returned header into req.Headers (instead of directly
assigning into a possibly nil map); keep the existing error handling and
preserve needsMarshal = true. Use the symbols virtualKey, p.addMCPIncludeTools,
req.Headers, headers, p.loadBalanceProvider, and needsMarshal to locate and
update the logic.
danpiths
changed the base branch from
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
to
graphite-base/1535
Pratham-Mishra04
force-pushed
the
graphite-base/1535
branch
from
7beaffe to
d98d5dc
Compare
Pratham-Mishra04
force-pushed
the
02-03-fix_minor_routing_rules_fixes
branch
from
6ec3119 to
1aa4a9c
Compare
Pratham-Mishra04
changed the base branch from
graphite-base/1535
to
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
Contributor
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)
plugins/governance/main.go (1)
340-371:⚠️ Potential issue | 🟠 MajorRouting decision can be overridden for /genai after reordering.
For
/genairequests,loadBalanceProviderreads the model from the URL path (unchanged byapplyRoutingRules) and can overwrite a routing decision that was just applied viactx.SetValue("model", ...). This makes routing rules ineffective for genai requests. Consider skipping load balancing when a routing decision has already been made (or using the routed model when available).🔧 Suggested guard
- if virtualKey != nil { - //2. Load balance provider - payload, err = p.loadBalanceProvider(ctx, req, payload, virtualKey) - if err != nil { - return nil, err - } + if virtualKey != nil { + //2. Load balance provider (skip when routing rules already selected provider/model) + if routingDecision == nil { + payload, err = p.loadBalanceProvider(ctx, req, payload, virtualKey) + if err != nil { + return nil, err + } + } //3. Add MCP tools headers, err := p.addMCPIncludeTools(nil, virtualKey) if err != nil { p.logger.Error("failed to add MCP include tools: %v", err) return nil, nil } for header, value := range headers { req.Headers[header] = value } needsMarshal = true }
🤖 Fix all issues with AI agents
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 10-11: The current use of sort.Slice to order virtual keys by
CreatedAt is unstable when CreatedAt ties; replace the unstable sort with a
stable comparison that deterministically breaks ties (e.g., use sort.SliceStable
or add a secondary comparison on the unique ID field) wherever ordering is
performed (the sort.Slice calls that compare CreatedAt, including the one around
lines 10-11 and the similar block at 281-283); update the comparator to first
compare CreatedAt and if equal compare the key ID (or another unique, stable
field) so responses are deterministic even when CreatedAt is identical or zero.
Collaborator
Author
Merge activity
|
Pratham-Mishra04
force-pushed
the
02-03-fix_minor_routing_rules_fixes
branch
from
1aa4a9c to
f5e9f3f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Added consistent ordering to database queries and enhanced virtual key provider-specific rate limits and budgets in the routing system.
Changes
Order("created_at ASC")to database queries for virtual keys, teams, customers, rate limits, and budgets to ensure consistent ordering in API responsesGetBudgetAndRateLimitStatusfunction to consider virtual key level provider-specific rate limits and budgetsType of change
Affected areas
How to test
Test the ordering of items in the UI to ensure they appear in a consistent order:
Test virtual key provider-specific rate limits by creating a virtual key with provider configs and verifying the rate limits are applied correctly.
Breaking changes
Related issues
Addresses inconsistent ordering in API responses and improves virtual key provider-specific rate limiting.
Security considerations
No security implications.
Checklist