Conversation
Contributor
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 4 minutes and 41 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (27)
📝 WalkthroughWalkthroughThread request contexts through governance/store APIs, change budget/rate-limit checks to return (Decision, error), reorder and centralize VK/customer/team evaluation flow, add customer/team evaluators, introduce CalendarAligned on budgets/rate-limits with a backfill migration and tests, and update many call sites and tests to new context and return shapes. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Plugin as GovernancePlugin
participant Resolver as BudgetResolver
participant Store as GovernanceStore
participant Tracker as UsageTracker
Client->>Plugin: EvaluateGovernanceRequest(ctx, req)
Plugin->>Store: GetVirtualKey(ctx, vkValue)
Store-->>Plugin: VirtualKey
Plugin->>Resolver: EvaluateVirtualKeyRequest(ctx, vk, req)
Resolver->>Store: CheckVirtualKeyRateLimit(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver->>Store: CheckVirtualKeyBudget(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver-->>Plugin: VKDecision
Plugin->>Resolver: EvaluateCustomerRequest(ctx, customerID, req)
Resolver->>Store: CheckCustomerRateLimit(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver->>Store: CheckCustomerBudget(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver-->>Plugin: CustomerDecision
Plugin->>Resolver: EvaluateTeamRequest(ctx, teamID, req)
Resolver->>Store: CheckTeamRateLimit(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver->>Store: CheckTeamBudget(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver-->>Plugin: TeamDecision
Plugin->>Resolver: EvaluateModelAndProviderRequest(ctx,...)
Resolver->>Store: CheckProviderRateLimit(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver->>Store: CheckProviderBudget(ctx,...)
Store-->>Resolver: (Decision, error)
Resolver-->>Plugin: ProviderDecision
Plugin->>Tracker: UpdateUsage(ctx, usage)
Tracker->>Store: GetVirtualKey(ctx, vkValue)
Store-->>Tracker: VirtualKey
Tracker-->>Plugin: UpdateConfirmed
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 4❌ Failed checks (2 warnings, 2 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
|
Contributor
Author
This stack of pull requests is managed by Graphite. Learn more about stacking. |
This was referenced Apr 19, 2026
Merged
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
34c5133 to
bc994b4
Compare
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
bc994b4 to
2f2566b
Compare
Contributor
Confidence Score: 5/5Safe to merge; all remaining findings are P2 style/polish items All previously flagged P0/P1 issues (empty team-budget block, index-out-of-bounds panic, missing rate-limit reset guard, division-by-zero) are resolved. Three remaining observations are P2: a stale comment in the store interface, an InitFromStore lock inconsistency vs Init, and a startup-reset gap for team/customer rate limits (the 10-second periodic worker closes it within one tick). plugins/governance/tracker.go (startup reset gap for team/customer rate limits) Important Files Changed
Reviews (11): Last reviewed commit: "team budgets" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![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: 12
🧹 Nitpick comments (3)
plugins/governance/store.go (3)
2592-2616: Inconsistent behavior:UpdateTeamGovernanceInMemoryreturns silently if entry doesn't exist.Other
Update*InMemorymethods (e.g.,UpdateTeamInMemoryat line 2400,UpdateCustomerInMemoryat line 2509) fall through to their correspondingCreate*InMemorymethod when the entry doesn't exist. This method silently returns, which could lead to lost updates if the caller expects upsert semantics.Consider aligning with the established pattern:
♻️ Proposed fix for consistent upsert behavior
func (gs *LocalGovernanceStore) UpdateTeamGovernanceInMemory(ctx context.Context, teamID string, budget *configstoreTables.TableBudget, rateLimit *configstoreTables.TableRateLimit) { if teamID == "" { return } value, exists := gs.teamGovernances.Load(teamID) if !exists || value == nil { - return + gs.CreateTeamGovernanceInMemory(ctx, teamID, budget, rateLimit) + return }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2592 - 2616, The UpdateTeamGovernanceInMemory currently returns silently when no existing entry is found; change it to perform an upsert like other methods by invoking CreateTeamGovernanceInMemory(ctx, teamID, budget, rateLimit) whenever teamGovernances.Load indicates missing or nil value or the type assertion for *TeamGovernance fails, and otherwise apply the existing update logic (setting tg.BudgetID/tg.RateLimitID and storing to gs.budgets/gs.rateLimits before storing tg back into gs.teamGovernances).
621-621: Consider joining violations slice for cleaner error message.Passing a
[]stringslice to%swill produce output like[violation1 violation2]. Usingstrings.Joinwould produce a cleaner message.✨ Proposed improvement
- return decision, fmt.Errorf("rate limit violated for %s: %s", entity, violations) + return decision, fmt.Errorf("rate limit violated for %s: %s", entity, strings.Join(violations, "; "))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` at line 621, The error message currently interpolates the []string violations slice directly (return decision, fmt.Errorf("rate limit violated for %s: %s", entity, violations)); change it to join the slice into a single string (e.g., msg := strings.Join(violations, ", ")) and use that in the fmt.Errorf call (fmt.Errorf("rate limit violated for %s: %s", entity, msg)); also add the strings import if missing. Ensure you update the return site that references decision, entity, and violations.
73-77: RemoveBusinessUnitGovernanceor add a comment explaining its purpose.This struct is defined at lines 72-75 but is not used anywhere in the codebase. If this is reserved for future use or part of a larger architectural plan, clarify with a comment. Otherwise, remove the dead code.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 73 - 77, Remove the unused struct BusinessUnitGovernance (and its fields BudgetID and RateLimitID) if it’s dead code; otherwise, keep it and add a concise comment above the type explaining its reserved purpose and planned usage (e.g., "reserved for future in-memory budget and rate limit tracking per business unit") so its intent is clear. Locate the type declaration named BusinessUnitGovernance and either delete the whole block or replace it with the explanatory comment plus the struct definition to document why it remains.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@plugins/governance/main.go`:
- Around line 1098-1116: The team-level branch is empty and the current logic
sets skipVKBudgetLimit = evaluationRequest.UserID != "" then calls
p.resolver.EvaluateVirtualKeyRequest(...) which—per
plugins/governance/resolver.go—causes all hierarchy checks to be skipped for
user-authenticated requests, bypassing team/customer budgets; fix by either
calling p.resolver.EvaluateTeamRequest(ctx, evaluationRequest.TeamID,
evaluationRequest) (or equivalent team evaluation function) before invoking
EvaluateUserRequest/EvaluateVirtualKeyRequest so team limits run, or change the
skipVKBudgetLimit boolean into a more specific flag (e.g.,
skipOnlyVKLocalChecks) and pass that to EvaluateVirtualKeyRequest so it skips
only VK-local checks while preserving team and customer evaluations.
In `@plugins/governance/store.go`:
- Line 2607: Fix the typo in the inline comment that reads "Updateing budgets
map" to "Updating budgets map" so comments are spelled correctly; locate the
comment in plugins/governance/store.go near the budgets update logic (search for
the string "Updateing budgets map") and update it to "Updating budgets map".
- Around line 2118-2127: The function collectRateLimitIDsFromMemory initializes
rateLimitIDs as an empty slice but assigns by index inside the nested loop which
causes an index out of range panic; change the inner loop in
collectRateLimitIDsFromMemory to append each rateLimit.ID to rateLimitIDs
(similar to collectBudgetIDsFromMemory) instead of assigning by index, keeping
the call to collectRateLimitsFromHierarchy and preserving vk/provider usage.
- Around line 2053-2057: The condition is checking the wrong map key: replace
the check of entityWiseBudgets["VK"] with entityWiseBudgets["Team"] so the
nil-check and initialization align with the subsequent append into
entityWiseBudgets["Team"]; update the nil check around entityWiseBudgets in the
same block (where budget is appended and seen[budget.ID] set) to initialize
entityWiseBudgets["Team"] = []*configstoreTables.TableBudget{} when nil.
- Line 800: The fmt.Sprintf call that constructs key is passing two arguments
but the format string "Model:%s" has only one verb, causing provider to be
ignored; update the call in the code that sets key (the fmt.Sprintf call
referencing mc.ModelName and provider) to include a second format verb (e.g.,
"Model:%s:%s") if you intend to include provider in the key, or remove the
provider argument if it should not be included, ensuring the final key variable
correctly reflects mc.ModelName and optionally provider.
- Around line 2070-2073: The conditional is checking the wrong map key: change
the if condition from entityWiseBudgets["VK"] to entityWiseBudgets["Customer"]
so you initialize the correct slice before appending; locate the block using the
entityWiseBudgets variable and the TableBudget append (the lines that read
entityWiseBudgets["Customer"] = []*configstoreTables.TableBudget{} and
entityWiseBudgets["Customer"] = append(..., budget)) and update the key in the
nil check to "Customer".
- Around line 857-860: The error text for enterprise-only user governance is
inconsistent; update the LocalGovernanceStore.CheckUserBudget method and the
other governance methods that return "user budgets are not supported" and "user
rate limits are not suported" to use a single standardized message (fixing the
typo "suported") such as "user budgets and rate limits are not supported in the
community plan"; specifically replace the error returned in CheckUserBudget
(method CheckUserBudget on LocalGovernanceStore) and the other methods that
currently return the two other strings so they all return errors.New("user
budgets and rate limits are not supported in the community plan").
In `@ui/app/workspace/config/views/mcpView.tsx`:
- Around line 219-226: The Tool Sync Interval numeric input is missing a test
selector; update the Input component (id="mcp-tool-sync-interval") to include a
data-testid following the pattern data-testid="mcp-tool-sync-interval-input" (or
similar pattern used elsewhere) so tests can target it; locate the Input in
mcpView.tsx that uses value={localValues.mcp_tool_sync_interval} and
onChange={(e) => handleToolSyncIntervalChange(e.target.value)} and add the
data-testid attribute to that component.
- Around line 307-315: Replace the accidental double-dot occurrences in the VFS
example strings: change "subtract..py", "GET_CHANNELS..py", "SEARCH_VIDEOS..py",
"get_forecast..py" to "subtract.py", "GET_CHANNELS.py", "SEARCH_VIDEOS.py",
"get_forecast.py" and update the explanatory paragraph "Individual ..py file for
each tool" to "Individual .py file for each tool"; these strings appear in the
JSX elements (the <div> entries and the <p className="text-muted-foreground mt-3
text-xs">) in mcpView.tsx.
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 176-218: The current truthy checks (e.g., formData.budgetMaxLimit,
formData.tokenMaxLimit, formData.requestMaxLimit) skip validation when the value
is 0; change them to explicit null/undefined checks (e.g.,
formData.budgetMaxLimit != null) so zero values are validated against
budgetMaxLimitNum and token/requestMaxLimitNum with Validator.minValue and
Validator.required; update the three conditional spreads around
Validator.minValue/Validator.required to use != null (or typeof !== "undefined")
instead of truthiness so 0 cannot bypass the minValue rules.
- Around line 102-108: initialState is only created once on mount via
useState(createInitialState(team)), so when the team prop changes the baseline
never updates and isDirty compares against stale data; fix by deriving/resetting
the baseline whenever team changes—either replace the static initialState
useState with a memoized value (useMemo(() => createInitialState(team), [team]))
or run an effect (useEffect) that calls createInitialState(team) and resets
formData to the new baseline with isDirty: false; update references to
initialState and setFormData so the form always compares against the current
team baseline.
In `@ui/components/ui/badge.tsx`:
- Around line 12-14: The default badge variant currently keeps text-primary
while switching to bg-primary/90 on hover, reducing contrast; update the class
string for the "default" variant (the object key named default in
ui/components/ui/badge.tsx) so the hover state uses primary-foreground for text
(mirror the outline variant pattern) — i.e., replace or augment the hover
classes in the default variant's class list that contains "bg-primary/90" and
"text-primary" so the hover text class becomes "primary-foreground" (or
"text-primary-foreground") to ensure sufficient contrast.
---
Nitpick comments:
In `@plugins/governance/store.go`:
- Around line 2592-2616: The UpdateTeamGovernanceInMemory currently returns
silently when no existing entry is found; change it to perform an upsert like
other methods by invoking CreateTeamGovernanceInMemory(ctx, teamID, budget,
rateLimit) whenever teamGovernances.Load indicates missing or nil value or the
type assertion for *TeamGovernance fails, and otherwise apply the existing
update logic (setting tg.BudgetID/tg.RateLimitID and storing to
gs.budgets/gs.rateLimits before storing tg back into gs.teamGovernances).
- Line 621: The error message currently interpolates the []string violations
slice directly (return decision, fmt.Errorf("rate limit violated for %s: %s",
entity, violations)); change it to join the slice into a single string (e.g.,
msg := strings.Join(violations, ", ")) and use that in the fmt.Errorf call
(fmt.Errorf("rate limit violated for %s: %s", entity, msg)); also add the
strings import if missing. Ensure you update the return site that references
decision, entity, and violations.
- Around line 73-77: Remove the unused struct BusinessUnitGovernance (and its
fields BudgetID and RateLimitID) if it’s dead code; otherwise, keep it and add a
concise comment above the type explaining its reserved purpose and planned usage
(e.g., "reserved for future in-memory budget and rate limit tracking per
business unit") so its intent is clear. Locate the type declaration named
BusinessUnitGovernance and either delete the whole block or replace it with the
explanatory comment plus the struct definition to document why it remains.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 278c6bdd-aa15-4805-be7e-6ce28ffa8fa1
📒 Files selected for processing (14)
plugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (4)
ui/app/workspace/governance/views/teamDialog.tsx (2)
102-108:⚠️ Potential issue | 🟠 MajorReset form baseline when
teamchanges to avoid stale dirty-state tracking.Lines 102-108 initialize
initialStateonly once; ifteamchanges without remount,isDirtycomparisons (Lines 131-157) use the old baseline.Suggested patch
- const [initialState] = useState<Omit<TeamFormData, "isDirty">>( - createInitialState(team), - ); + const initialState = useMemo<Omit<TeamFormData, "isDirty">>( + () => createInitialState(team), + [team], + ); const [formData, setFormData] = useState<TeamFormData>({ ...initialState, isDirty: false, }); + + useEffect(() => { + setFormData({ ...initialState, isDirty: false }); + }, [initialState]);Also applies to: 131-157
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/governance/views/teamDialog.tsx` around lines 102 - 108, The initialState is captured once by useState(createInitialState(team)) causing stale baseline when team updates; update the logic so initialState is derived from team on every change (e.g., replace the one-time useState with useMemo(() => createInitialState(team), [team]) or update initialState inside a useEffect that runs when team changes), and when team changes reset formData via setFormData({...newInitialState, isDirty: false}) so the isDirty comparisons in the formData/isDirty tracking code (functions/variables: initialState, createInitialState, formData, setFormData, TeamFormData, isDirty) use the fresh baseline.
176-218:⚠️ Potential issue | 🟠 MajorReplace numeric truthy checks with nullish checks so
0cannot bypass logic.Lines 176/191/206 skip validators when value is
0, but submit paths still include0in payload. Line 422 has the same pattern and hides the calendar-alignment UI for0.Suggested patch
- ...(formData.budgetMaxLimit + ...(formData.budgetMaxLimit != null ? [ Validator.minValue( budgetMaxLimitNum || 0, 0.01, "Budget max limit must be greater than $0.01", ), Validator.required( formData.budgetResetDuration, "Budget reset duration is required", ), ] : []), @@ - ...(formData.tokenMaxLimit + ...(formData.tokenMaxLimit != null ? [ Validator.minValue( tokenMaxLimitNum || 0, 1, "Token max limit must be at least 1", ), Validator.required( formData.tokenResetDuration, "Token reset duration is required", ), ] : []), @@ - ...(formData.requestMaxLimit + ...(formData.requestMaxLimit != null ? [ Validator.minValue( requestMaxLimitNum || 0, 1, "Request max limit must be at least 1", ), Validator.required( formData.requestResetDuration, "Request reset duration is required", ), ] : []), @@ - {formData.budgetMaxLimit && + {formData.budgetMaxLimit != null && supportsCalendarAlignment(formData.budgetResetDuration) && (Also applies to: 422-423
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/governance/views/teamDialog.tsx` around lines 176 - 218, The validators and UI toggles currently use truthy checks on formData.budgetMaxLimit, formData.tokenMaxLimit, and formData.requestMaxLimit which cause 0 to be treated as "absent" and skip validation/UI; change those conditions to explicit nullish checks (e.g., check against null/undefined) so 0 is accepted and validators like Validator.minValue and Validator.required (used with budgetMaxLimitNum, tokenMaxLimitNum, requestMaxLimitNum and the corresponding reset duration fields) still run; also update the duplicate check that controls the calendar-alignment UI (the same budgetMaxLimit truthy check) to the same nullish check to ensure the UI shows when value is 0.ui/components/ui/badge.tsx (1)
13-13:⚠️ Potential issue | 🟠 MajorFix unresolved hover contrast in the
defaultbadge variant.Line 13 still keeps
text-primarywhile hover switches tobg-primary/90, which causes low-contrast text on hover.Suggested patch
default: - "border-transparent bg-primary/10 border-primary/50 text-primary [a&]:hover:bg-primary/90", + "border-transparent bg-primary/10 border-primary/50 text-primary [a&]:hover:bg-primary/90 [a&]:hover:text-primary-foreground",🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/components/ui/badge.tsx` at line 13, The default badge variant's class string ("border-transparent bg-primary/10 border-primary/50 text-primary [a&]:hover:bg-primary/90") leaves text-primary unchanged on hover causing low contrast; update that class in ui/components/ui/badge.tsx (the default variant definition) to also set an appropriate hover text color (for example add "[a&]:hover:text-white" or "[a&]:hover:text-primary-foreground") so the text contrasts with the hover background bg-primary/90.plugins/governance/store.go (1)
2116-2124:⚠️ Potential issue | 🔴 CriticalAppend rate-limit IDs instead of indexing into an empty slice.
rateLimitIDsstarts at length 0, sorateLimitIDs[i] = rateLimit.IDpanics as soon as any category contains a rate limit. That breaksUpdateVirtualKeyRateLimitUsageInMemoryfor governed keys.#!/bin/bash # Verify the collector initializes an empty slice and then writes by index. rg -n -C3 'func \(gs \*LocalGovernanceStore\) collectRateLimitIDsFromMemory|rateLimitIDs := \[\]string\{\}|rateLimitIDs\[i\] =' plugins/governance/store.go🐛 Proposed fix
func (gs *LocalGovernanceStore) collectRateLimitIDsFromMemory(ctx context.Context, vk *configstoreTables.TableVirtualKey, provider schemas.ModelProvider) []string { rateLimitsWithCategories := gs.collectRateLimitsFromHierarchy(ctx, vk, provider) rateLimitIDs := []string{} for _, rateLimits := range rateLimitsWithCategories { - for i, rateLimit := range rateLimits { - rateLimitIDs[i] = rateLimit.ID + for _, rateLimit := range rateLimits { + rateLimitIDs = append(rateLimitIDs, rateLimit.ID) } } return rateLimitIDs }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2116 - 2124, collectRateLimitIDsFromMemory is writing into an empty slice by index which panics; change the logic to append IDs instead of assigning by index (or preallocate the slice using known lengths). Locate collectRateLimitIDsFromMemory which calls collectRateLimitsFromHierarchy and replace the inner assignment rateLimitIDs[i] = rateLimit.ID with an append (rateLimitIDs = append(rateLimitIDs, rateLimit.ID)) so UpdateVirtualKeyRateLimitUsageInMemory can safely iterate governed keys.
🧹 Nitpick comments (1)
ui/app/workspace/governance/views/teamDialog.tsx (1)
343-343: Drop emptyclassNameon<form>.Line 343 has
className=""; removing it reduces noise with no behavior change.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/governance/views/teamDialog.tsx` at line 343, Remove the empty className attribute from the form element in teamDialog.tsx — locate the JSX <form onSubmit={handleSubmit} className=""> (around the team dialog render where handleSubmit is used) and delete the className="" so the element becomes simply <form onSubmit={handleSubmit}> to eliminate needless noise.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@plugins/governance/resolver.go`:
- Around line 145-146: The inline comment above the CheckTeamBudget call in
EvaluateTeamRequest is incorrect ("user-level"); update that comment to
accurately reflect the operation (e.g., "Check team-level budget") and scan
nearby comments in EvaluateTeamRequest for any other copy-paste references to
user-level logic to correct them to team-level where appropriate (target
symbols: EvaluateTeamRequest, CheckTeamBudget).
- Around line 368-371: The code uses a misspelled variable name `decsion` when
assigning the result of r.store.CheckVirtualKeyBudget(ctx, vk, request, nil);
rename `decsion` to `decision` wherever used (including the return constructing
&EvaluationResult{Decision: decision, ...}) so the variable matches the expected
identifier; ensure references in that conditional branch and any subsequent uses
match the corrected `decision` name and compile.
- Around line 129-158: The empty conditional that intends to run team-level
checks should call BudgetResolver.EvaluateTeamRequest and handle its result:
inside the empty block where you check "if result.Decision == DecisionAllow"
invoke the resolver (e.g., r.BudgetResolver.EvaluateTeamRequest(ctx, teamID,
request)), then if that returns a non-allow Decision propagate that
EvaluationResult (or merge its Decision/Reason into the overall result) and
short-circuit; if you truly intend this to be scaffolding instead of wired
behavior, replace the empty block with a clear TODO comment referencing the PR
number (e.g., "// TODO: Wire up team-level evaluation in PR#<n>") so intent is
explicit.
In `@plugins/governance/store.go`:
- Around line 855-857: The stubs CheckUserBudget (and the similar CheckOrgBudget
stub at the other location) currently return DecisionAllow with a non-nil error;
change them to return a denial decision (e.g., DecisionDeny) alongside the
existing error so callers cannot treat an unsupported enterprise-only check as
allowed. Locate LocalGovernanceStore.CheckUserBudget (and the CheckOrgBudget
stub) and replace the DecisionAllow return value with the appropriate denial
constant (DecisionDeny) while keeping the error text unchanged.
- Around line 553-656: The checks in CheckRateLimit and CheckBudget iterate a
map (entityWiseRateLimits / entityWiseBudgets) which loses the hierarchy order
produced by collectRateLimitsFromHierarchy / collectBudgetsFromHierarchy; change
the flow so the checker preserves that order: have the collectors return an
ordered slice of (entityKey, items) or provide a stable ordered entity list
(Provider→VK→Team→Customer) and iterate that slice in CheckRateLimit and
CheckBudget instead of using "for entity, ... := range ..."; alternatively, if
changing collectors is hard, build a deterministically ordered list of entity
keys using your precedence constants and iterate that list to look up entries in
entityWiseRateLimits/entityWiseBudgets before performing the existing checks and
returning decisions (ensuring DecisionTokenLimited/DecisionRequestLimited
selection remains deterministic).
---
Duplicate comments:
In `@plugins/governance/store.go`:
- Around line 2116-2124: collectRateLimitIDsFromMemory is writing into an empty
slice by index which panics; change the logic to append IDs instead of assigning
by index (or preallocate the slice using known lengths). Locate
collectRateLimitIDsFromMemory which calls collectRateLimitsFromHierarchy and
replace the inner assignment rateLimitIDs[i] = rateLimit.ID with an append
(rateLimitIDs = append(rateLimitIDs, rateLimit.ID)) so
UpdateVirtualKeyRateLimitUsageInMemory can safely iterate governed keys.
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 102-108: The initialState is captured once by
useState(createInitialState(team)) causing stale baseline when team updates;
update the logic so initialState is derived from team on every change (e.g.,
replace the one-time useState with useMemo(() => createInitialState(team),
[team]) or update initialState inside a useEffect that runs when team changes),
and when team changes reset formData via setFormData({...newInitialState,
isDirty: false}) so the isDirty comparisons in the formData/isDirty tracking
code (functions/variables: initialState, createInitialState, formData,
setFormData, TeamFormData, isDirty) use the fresh baseline.
- Around line 176-218: The validators and UI toggles currently use truthy checks
on formData.budgetMaxLimit, formData.tokenMaxLimit, and formData.requestMaxLimit
which cause 0 to be treated as "absent" and skip validation/UI; change those
conditions to explicit nullish checks (e.g., check against null/undefined) so 0
is accepted and validators like Validator.minValue and Validator.required (used
with budgetMaxLimitNum, tokenMaxLimitNum, requestMaxLimitNum and the
corresponding reset duration fields) still run; also update the duplicate check
that controls the calendar-alignment UI (the same budgetMaxLimit truthy check)
to the same nullish check to ensure the UI shows when value is 0.
In `@ui/components/ui/badge.tsx`:
- Line 13: The default badge variant's class string ("border-transparent
bg-primary/10 border-primary/50 text-primary [a&]:hover:bg-primary/90") leaves
text-primary unchanged on hover causing low contrast; update that class in
ui/components/ui/badge.tsx (the default variant definition) to also set an
appropriate hover text color (for example add "[a&]:hover:text-white" or
"[a&]:hover:text-primary-foreground") so the text contrasts with the hover
background bg-primary/90.
---
Nitpick comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Line 343: Remove the empty className attribute from the form element in
teamDialog.tsx — locate the JSX <form onSubmit={handleSubmit} className="">
(around the team dialog render where handleSubmit is used) and delete the
className="" so the element becomes simply <form onSubmit={handleSubmit}> to
eliminate needless noise.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ee865dd0-4651-444c-9466-e32b9dd83a6a
📒 Files selected for processing (14)
plugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (3)
- plugins/governance/tracker.go
- plugins/governance/tracker_test.go
- plugins/governance/resolver_test.go
🚧 Files skipped from review as they are similar to previous changes (4)
- plugins/governance/main.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/utils.go
- plugins/governance/modelprovidergovernance_test.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
2f2566b to
f9bfbad
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@framework/configstore/migrations.go`:
- Around line 6592-6599: The migration currently aborts when
tx.First(&rateLimit, virtualKeys[i].RateLimitID) returns gorm.ErrRecordNotFound;
update the code around virtualKeys[i].RateLimitID handling so that if err ==
gorm.ErrRecordNotFound you skip this virtual key (do not return an error) and
continue, but still return on any other non-nil err; keep the subsequent update
of rateLimit.CalendarAligned and tx.Save(&rateLimit) unchanged for found rows
and ensure you still propagate unexpected tx.Save errors; use the existing
symbols virtualKeys, RateLimitID, tables.TableRateLimit, tx.First, tx.Save and
gorm.ErrRecordNotFound to locate the change.
- Around line 6586-6589: The query chain currently calls tx.Find(&virtualKeys)
first which executes immediately and prevents Preload/Where/Select from
applying; change the call order so filters and preloads run before the finisher
by invoking tx.Preload("budgets").Where("calendar_aligned = ?",
true).Select("rate_limit_id").Find(&virtualKeys) and then check .Error; update
the code that references virtualKeys/tx in the migration to use this reordered
chain.
In `@plugins/governance/resolver.go`:
- Around line 91-96: The Check*RateLimit branches (e.g.,
r.store.CheckProviderRateLimit) only treat failures when err != nil and ignore
non-allow Decisions returned with a nil error, so update each rate-limit check
(provider, model, team, user, vector/ VK checks referenced by
CheckModelRateLimit, CheckTeamRateLimit, CheckUserRateLimit,
CheckVectorDBRateLimit or similar) to treat any non-allow Decision
(DecisionRateLimited/DecisionTokenLimited/DecisionRequestLimited) as a blocking
EvaluationResult just like isProviderRateLimitViolated does; when constructing
the EvaluationResult.Reason use a nil-safe format that prefers the decision
reason if err is nil (avoid calling err.Error() when err == nil) and include the
decision name and any provided reason for clarity.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5608b6a0-f2a9-477c-8e37-1df2ff7266a6
📒 Files selected for processing (18)
framework/configstore/migrations.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (6)
- framework/configstore/tables/ratelimit.go
- ui/components/ui/badge.tsx
- plugins/governance/routing_test.go
- framework/configstore/tables/budget.go
- framework/configstore/tables/virtualkey.go
- plugins/governance/routing.go
🚧 Files skipped from review as they are similar to previous changes (9)
- plugins/governance/tracker.go
- plugins/governance/tracker_test.go
- plugins/governance/resolver_test.go
- plugins/governance/modelprovidergovernance_test.go
- ui/app/workspace/governance/views/teamDialog.tsx
- plugins/governance/main.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/store_test.go
- plugins/governance/store.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
f9bfbad to
092fc8d
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (3)
framework/configstore/migrations.go (2)
6592-6596:⚠️ Potential issue | 🟠 MajorSkip orphaned rate-limit rows instead of failing startup.
A stale
rate_limit_idon a virtual key currently aborts the whole migration. These references are intentionally not FK-enforced here, sogorm.ErrRecordNotFoundshould be treated as a skip and only unexpected DB errors should fail the upgrade.Suggested fix
- if err := tx.First(&rateLimit, virtualKeys[i].RateLimitID).Error; err != nil { + if err := tx.First(&rateLimit, virtualKeys[i].RateLimitID).Error; err != nil { + if err == gorm.ErrRecordNotFound { + continue + } return fmt.Errorf("failed to load rate limit for virtual key %s: %w", virtualKeys[i].ID, err) }Based on learnings, FK enforcement in
framework/configstore/migrations.gois limited toTableVirtualKeyProviderConfig;TableVirtualKeyrate-limit references are intentionally not protected by a DB constraint.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 6592 - 6596, The loop that loads rate limits for virtual keys using tx.First(&rateLimit, virtualKeys[i].RateLimitID) should not abort startup when the DB row is missing; update the error handling so that if err == gorm.ErrRecordNotFound you simply continue (skip the orphaned rate-limit reference) but for any other non-nil err you return the wrapped error (e.g., return fmt.Errorf("failed to load rate limit for virtual key %s: %w", virtualKeys[i].ID, err)); this change affects the code around virtualKeys[i].RateLimitID, TableRateLimit/rateLimit and tx.First calls.
6586-6588:⚠️ Potential issue | 🔴 CriticalFix the GORM query before this migration lands.
Findis the finisher here, so theWhere/Preload/Selectnever affect the executed query. On top of that, the preload path should use the Go field nameBudgets, and the parentidneeds to be selected for that preload to work. As written, this migration can scan all virtual keys and skip the budget backfill entirely.Suggested fix
- if err := tx.Find(&virtualKeys).Preload("budgets").Where("calendar_aligned = ?", true).Select("rate_limit_id").Error; err != nil { + if err := tx. + Preload("Budgets"). + Where("calendar_aligned = ?", true). + Select("id", "rate_limit_id"). + Find(&virtualKeys).Error; err != nil {In GORM v2, is `Find` a finisher method that executes immediately, must `Preload` association paths use Go struct field names like `Budgets`, and when preloading a has-many association after `Select`, does the parent query need to include the primary key column?🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 6586 - 6588, The GORM call currently calls tx.Find(...) first so Where/Preload/Select are ignored and Preload uses the wrong path; change the query to build filters then execute: use tx.Where("calendar_aligned = ?", true).Select("id, rate_limit_id").Preload("Budgets").Find(&virtualKeys). Ensure the Preload uses the Go field name "Budgets" and include the parent primary key ("id") in Select so the has-many preload can map child rows to parents for tables.TableVirtualKey.plugins/governance/resolver.go (1)
91-95:⚠️ Potential issue | 🔴 CriticalBlock on non-error rate-limit decisions in the resolver path.
These branches still only stop on
err != nil. The store API now returns(Decision, error), and this same file already treatsDecisionRateLimited/DecisionTokenLimited/DecisionRequestLimitedas blockers elsewhere, so adecision != allow, err == nilresult will still fall through here and incorrectly allow the request.Representative fix
- if decision, err := r.store.CheckProviderRateLimit(ctx, request, nil, nil); err != nil { + if decision, err := r.store.CheckProviderRateLimit(ctx, request, nil, nil); err != nil || isRateLimitViolation(decision) { + reason := "provider-level rate limit exceeded" + if err != nil { + reason = fmt.Sprintf("%s: %s", reason, err.Error()) + } return &EvaluationResult{ Decision: decision, - Reason: fmt.Sprintf("Provider-level rate limit check failed: %s", err.Error()), + Reason: reason, } }Apply the same pattern to the model/customer/team/user/VK rate-limit branches.
Also applies to: 107-110, 138-142, 168-172, 203-207, 372-389
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 91 - 95, The current resolver branches only check for err != nil when calling r.store.CheckProviderRateLimit (and the analogous model/customer/team/user/VK calls), so non-allow Decisions like DecisionRateLimited / DecisionTokenLimited / DecisionRequestLimited still fall through; update each branch (e.g., the r.store.CheckProviderRateLimit call and the other rate-limit calls at the noted spots) to block when the returned Decision is not DecisionAllow or when err != nil by returning an EvaluationResult with that Decision and a Reason (include err.Error() when err != nil), mirroring the pattern used elsewhere in this file where DecisionRateLimited/DecisionTokenLimited/DecisionRequestLimited are treated as blockers.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 569-606: The render checks for the token usage row use truthiness
(team?.rate_limit?.token_max_limit) which hides rows when the limit is 0; change
those guards to explicit nullish checks (e.g. team?.rate_limit?.token_max_limit
!= null) and do the same for the request usage row so rows render for zero
values; keep the existing percent/fallback logic for Badge and ensure references
like team.rate_limit.token_current_usage, team.rate_limit.token_max_limit and
team.rate_limit.token_last_reset are used unchanged.
In `@ui/components/ui/badge.tsx`:
- Around line 20-21: The success badge variant in ui/components/ui/badge.tsx
sets a dark hover background but leaves text-black, making linked badges hard to
read on hover; update the 'success' variant string to add an anchor-hover
foreground rule (e.g., add [a&]:hover:text-white) so links inside the Badge
switch to a readable color on hover while preserving text-black by default.
---
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 6592-6596: The loop that loads rate limits for virtual keys using
tx.First(&rateLimit, virtualKeys[i].RateLimitID) should not abort startup when
the DB row is missing; update the error handling so that if err ==
gorm.ErrRecordNotFound you simply continue (skip the orphaned rate-limit
reference) but for any other non-nil err you return the wrapped error (e.g.,
return fmt.Errorf("failed to load rate limit for virtual key %s: %w",
virtualKeys[i].ID, err)); this change affects the code around
virtualKeys[i].RateLimitID, TableRateLimit/rateLimit and tx.First calls.
- Around line 6586-6588: The GORM call currently calls tx.Find(...) first so
Where/Preload/Select are ignored and Preload uses the wrong path; change the
query to build filters then execute: use tx.Where("calendar_aligned = ?",
true).Select("id, rate_limit_id").Preload("Budgets").Find(&virtualKeys). Ensure
the Preload uses the Go field name "Budgets" and include the parent primary key
("id") in Select so the has-many preload can map child rows to parents for
tables.TableVirtualKey.
In `@plugins/governance/resolver.go`:
- Around line 91-95: The current resolver branches only check for err != nil
when calling r.store.CheckProviderRateLimit (and the analogous
model/customer/team/user/VK calls), so non-allow Decisions like
DecisionRateLimited / DecisionTokenLimited / DecisionRequestLimited still fall
through; update each branch (e.g., the r.store.CheckProviderRateLimit call and
the other rate-limit calls at the noted spots) to block when the returned
Decision is not DecisionAllow or when err != nil by returning an
EvaluationResult with that Decision and a Reason (include err.Error() when err
!= nil), mirroring the pattern used elsewhere in this file where
DecisionRateLimited/DecisionTokenLimited/DecisionRequestLimited are treated as
blockers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2e5cb50d-9136-47ee-a0f1-7e5217d201ba
📒 Files selected for processing (18)
framework/configstore/migrations.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (1)
- plugins/governance/tracker_test.go
🚧 Files skipped from review as they are similar to previous changes (10)
- framework/configstore/tables/budget.go
- framework/configstore/tables/ratelimit.go
- plugins/governance/resolver_test.go
- plugins/governance/utils.go
- plugins/governance/tracker.go
- plugins/governance/routing_test.go
- plugins/governance/store_test.go
- framework/configstore/tables/virtualkey.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/modelprovidergovernance_test.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
092fc8d to
7468b5e
Compare
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/resolver.go (1)
411-452:⚠️ Potential issue | 🟡 Minor
isProviderBudgetViolatedonly checkserr != nil, inconsistent with other paths.Lines 416 and 425 return
trueonly whenerr != nil, but don't check for non-allow decisions whenerr == nil. This is inconsistent with the hardened pattern applied elsewhere (e.g.,isProviderRateLimitViolatedat line 437 correctly uses|| isRateLimitViolation(decision)).🔧 Proposed fix
func (r *BudgetResolver) isProviderBudgetViolated(ctx context.Context, vk *configstoreTables.TableVirtualKey, config configstoreTables.TableVirtualKeyProviderConfig) bool { request := &EvaluationRequest{Provider: schemas.ModelProvider(config.Provider)} // 1. Check global provider-level budget first - if _, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil { + if decision, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil || isBudgetViolation(decision) { r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) return true } // 2. Check VK-level provider config budget if len(config.Budgets) == 0 { return false } - if _, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil { + if decision, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil || isBudgetViolation(decision) { r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, err.Error())) return true } return false }Note: The debug log calls
err.Error()which will panic iferr == nilbutisBudgetViolation(decision)is true. Consider usingreasonFromErr(err, decision)for nil-safety:- r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) + r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, reasonFromErr(err, decision)))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 411 - 452, The isProviderBudgetViolated function currently only treats err != nil as a violation; update it to follow the hardened pattern used by isProviderRateLimitViolated by checking both the returned decision and the error (i.e., use something like: if decision, err := r.store.CheckProviderBudget(...); err != nil || isBudgetViolation(decision) { ... }) for the global provider check and similarly for CheckVirtualKeyBudget; also replace direct err.Error() in the debug logs with a nil-safe helper (e.g., reasonFromErr(err, decision)) to avoid panics when err is nil but decision indicates a violation; target functions/symbols: isProviderBudgetViolated, CheckProviderBudget, CheckVirtualKeyBudget, isBudgetViolation, reasonFromErr.
♻️ Duplicate comments (1)
ui/app/workspace/governance/views/teamDialog.tsx (1)
569-606:⚠️ Potential issue | 🟡 MinorUse nullish guards for zero-valued rate-limit rows too.
These two conditions still rely on truthiness, so a persisted
0limit hides the row completely and the new> 0 ? ... : 0fallback never executes.Suggested fix
- {team?.rate_limit?.token_max_limit && ( + {team?.rate_limit?.token_max_limit != null && ( @@ - {team?.rate_limit?.request_max_limit && ( + {team?.rate_limit?.request_max_limit != null && (🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/governance/views/teamDialog.tsx` around lines 569 - 606, The conditional rendering uses truthiness (team?.rate_limit?.token_max_limit) which hides legitimate zero values; change the checks for both token_max_limit and request_max_limit to explicit nullish checks (e.g., typeof team?.rate_limit?.token_max_limit === "number" or team?.rate_limit?.token_max_limit != null) so a persisted 0 will still render the row (you can keep the existing > 0 checks inside the Badge/percent calculation to handle zero behavior); apply the same nullish guard pattern for request_max_limit as well.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 250-253: The update payload is sending an empty string for
`customer_id` when the user selects “None”, unlike the create path which
normalizes that to `undefined`; modify the construction of `updateData:
UpdateTeamRequest` so `customer_id` uses the same normalization as create (e.g.,
set `customer_id` to `undefined` when `formData.customerId` is an empty string)
to ensure “remove customer” edits are handled consistently for the `updateData`
object.
---
Outside diff comments:
In `@plugins/governance/resolver.go`:
- Around line 411-452: The isProviderBudgetViolated function currently only
treats err != nil as a violation; update it to follow the hardened pattern used
by isProviderRateLimitViolated by checking both the returned decision and the
error (i.e., use something like: if decision, err :=
r.store.CheckProviderBudget(...); err != nil || isBudgetViolation(decision) {
... }) for the global provider check and similarly for CheckVirtualKeyBudget;
also replace direct err.Error() in the debug logs with a nil-safe helper (e.g.,
reasonFromErr(err, decision)) to avoid panics when err is nil but decision
indicates a violation; target functions/symbols: isProviderBudgetViolated,
CheckProviderBudget, CheckVirtualKeyBudget, isBudgetViolation, reasonFromErr.
---
Duplicate comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 569-606: The conditional rendering uses truthiness
(team?.rate_limit?.token_max_limit) which hides legitimate zero values; change
the checks for both token_max_limit and request_max_limit to explicit nullish
checks (e.g., typeof team?.rate_limit?.token_max_limit === "number" or
team?.rate_limit?.token_max_limit != null) so a persisted 0 will still render
the row (you can keep the existing > 0 checks inside the Badge/percent
calculation to handle zero behavior); apply the same nullish guard pattern for
request_max_limit as well.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a4b6aa7f-d32b-4c95-bc1d-776770c5db75
📒 Files selected for processing (18)
framework/configstore/migrations.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (4)
- framework/configstore/tables/ratelimit.go
- plugins/governance/tracker_test.go
- plugins/governance/store_test.go
- framework/configstore/migrations.go
🚧 Files skipped from review as they are similar to previous changes (8)
- framework/configstore/tables/budget.go
- plugins/governance/routing.go
- plugins/governance/utils.go
- framework/configstore/tables/virtualkey.go
- ui/components/ui/badge.tsx
- plugins/governance/modelprovidergovernance_test.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/routing_test.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
7468b5e to
37d5f3a
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
plugins/governance/resolver.go (1)
416-425:⚠️ Potential issue | 🔴 CriticalTreat decision-only budget blocks as violations here too.
Lines 416 and 425 still only gate on
err != nil. After this PR’s store-contract change,CheckProviderBudgetandCheckVirtualKeyBudgetcan reject withDecisionBudgetExceededand a nil error, so this helper can keep an over-budget provider eligible for routing.🔧 Proposed fix
- if _, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) + if decision, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, reasonFromErr(err, decision))) return true } @@ - if _, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, err.Error())) + if decision, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, reasonFromErr(err, decision))) return true }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 416 - 425, The helper currently only treats non-nil errors from r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget as budget violations, but the store can return a DecisionBudgetExceeded result with a nil error; update the calls to capture the first return value (e.g., provRes, vkRes) and consider the budget exceeded if err != nil OR the returned status equals DecisionBudgetExceeded. Concretely: change the two calls to assign the first return value, then replace the condition `if _, err := r.store.CheckProviderBudget(...); err != nil {` and the similar CheckVirtualKeyBudget check so they return true when `err != nil || provRes == store.DecisionBudgetExceeded` (and likewise for `vkRes`) while keeping the existing log for provider (include config.Provider and the error when present).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@framework/logstore/asyncjob.go`:
- Line 37: The GovernanceStore interface signature for GetVirtualKey changed to
include a context.Context parameter (GetVirtualKey(ctx context.Context, vkValue
string)), so update every implementation and mock/test double that still defines
GetVirtualKey(vkValue string); locate concrete types implementing
GovernanceStore and test fakes/mocks and change their method signatures to match
the new signature, propagate the ctx parameter through any call sites (e.g.,
places invoking GetVirtualKey) and ensure any uses of the vkValue argument
remain unchanged while passing the appropriate context value.
In `@ui/components/ui/badge.tsx`:
- Around line 16-17: The destructive badge variant uses a hardcoded text-black
which conflicts with dark:bg-destructive/60 and the dark hover background;
update the destructive variant string in ui/components/ui/badge.tsx (the
destructive key) to use a responsive foreground such as replacing text-black
with text-black dark:text-white and add a hover state like [a&]:hover:text-white
or dark:[a&]:hover:text-white so the text becomes white in dark mode and on
hover against the darker destructive background.
---
Duplicate comments:
In `@plugins/governance/resolver.go`:
- Around line 416-425: The helper currently only treats non-nil errors from
r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget as budget
violations, but the store can return a DecisionBudgetExceeded result with a nil
error; update the calls to capture the first return value (e.g., provRes, vkRes)
and consider the budget exceeded if err != nil OR the returned status equals
DecisionBudgetExceeded. Concretely: change the two calls to assign the first
return value, then replace the condition `if _, err :=
r.store.CheckProviderBudget(...); err != nil {` and the similar
CheckVirtualKeyBudget check so they return true when `err != nil || provRes ==
store.DecisionBudgetExceeded` (and likewise for `vkRes`) while keeping the
existing log for provider (include config.Provider and the error when present).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8b8cb097-add5-4bd5-b317-31a68b31c205
📒 Files selected for processing (26)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (6)
- plugins/governance/tracker_test.go
- framework/configstore/tables/ratelimit.go
- framework/configstore/tables/virtualkey.go
- plugins/governance/store_test.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/routing_test.go
🚧 Files skipped from review as they are similar to previous changes (7)
- plugins/governance/resolver_test.go
- framework/configstore/tables/budget.go
- plugins/governance/modelprovidergovernance_test.go
- plugins/governance/routing.go
- framework/configstore/migrations.go
- ui/app/workspace/governance/views/teamDialog.tsx
- plugins/governance/utils.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
2 times, most recently
from
e8a123d to
73874c4
Compare
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
plugins/governance/resolver.go (1)
415-426:⚠️ Potential issue | 🔴 CriticalHandle decision-only budget violations here too.
These two branches still ignore the returned
Decision. IfCheckProviderBudgetorCheckVirtualKeyBudgetreturnsDecisionBudgetExceededwitherr == nil, this helper returnsfalseand keeps the provider config eligible even though governance already denied it.🔧 Proposed fix
// 1. Check global provider-level budget first - if _, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) + if decision, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, reasonFromErr(err, decision))) return true } @@ - if _, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, err.Error())) + if decision, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, reasonFromErr(err, decision))) return true }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 415 - 426, The budget-check branches call r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget but only treat non-nil err as a violation; update both call sites in the resolver helper to inspect the returned Decision value and treat DecisionBudgetExceeded as a budget violation (return true) even when err == nil — i.e., after calling CheckProviderBudget/CheckVirtualKeyBudget, check if decision == DecisionBudgetExceeded and log the VK/provider with r.logger.Debug before returning true; keep existing err-based logging/returns intact for other error cases.
🧹 Nitpick comments (1)
framework/logstore/asyncjob.go (1)
142-148: Consider standardizing update payload map types for consistency.This file now mixes
map[string]anyandmap[string]interface{}. Unifying on one form (preferablymap[string]any) would improve readability.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/logstore/asyncjob.go` around lines 142 - 148, The update payload in calls to e.logstore.UpdateAsyncJob and similar places mixes map[string]any and map[string]interface{}; standardize them to map[string]any for consistency. Locate usages such as the payload passed in the UpdateAsyncJob call (the map literal with "status", "status_code", "error", "completed_at", "expires_at") and replace any map[string]interface{} declarations or literals with map[string]any so all update payloads use the same type.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@plugins/governance/resolver.go`:
- Around line 415-426: The budget-check branches call
r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget but only treat
non-nil err as a violation; update both call sites in the resolver helper to
inspect the returned Decision value and treat DecisionBudgetExceeded as a budget
violation (return true) even when err == nil — i.e., after calling
CheckProviderBudget/CheckVirtualKeyBudget, check if decision ==
DecisionBudgetExceeded and log the VK/provider with r.logger.Debug before
returning true; keep existing err-based logging/returns intact for other error
cases.
---
Nitpick comments:
In `@framework/logstore/asyncjob.go`:
- Around line 142-148: The update payload in calls to e.logstore.UpdateAsyncJob
and similar places mixes map[string]any and map[string]interface{}; standardize
them to map[string]any for consistency. Locate usages such as the payload passed
in the UpdateAsyncJob call (the map literal with "status", "status_code",
"error", "completed_at", "expires_at") and replace any map[string]interface{}
declarations or literals with map[string]any so all update payloads use the same
type.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 1684acc9-aee8-4142-a2e1-1dad84dbd4a2
📒 Files selected for processing (27)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goframework/logstore/asyncjob_test.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (6)
- framework/configstore/tables/budget.go
- plugins/governance/tracker.go
- ui/app/workspace/config/views/mcpView.tsx
- framework/configstore/tables/virtualkey.go
- framework/configstore/migrations.go
- ui/app/workspace/governance/views/teamDialog.tsx
🚧 Files skipped from review as they are similar to previous changes (12)
- transports/bifrost-http/handlers/governance_test.go
- plugins/governance/tracker_test.go
- transports/bifrost-http/handlers/mcp.go
- plugins/governance/main.go
- framework/configstore/migrations_test.go
- framework/configstore/tables/ratelimit.go
- plugins/governance/modelprovidergovernance_test.go
- transports/bifrost-http/handlers/pricing_override_test.go
- plugins/governance/routing_test.go
- plugins/governance/utils.go
- plugins/governance/store_test.go
- transports/bifrost-http/lib/config_test.go
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
plugins/governance/resolver.go (1)
416-426:⚠️ Potential issue | 🟠 MajorHandle decision-only budget violations in provider filtering.
isProviderBudgetViolatedstill treats budget failures aserr != nilonly. If a store returnsDecisionBudgetExceededwith a nil error, this helper will keep that provider in the load-balancing pool and route traffic to an over-budget backend. The sibling rate-limit helper already uses the correcterr != nil || is*Violation(decision)guard.🔧 Suggested fix
- if _, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) + if decision, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, reasonFromErr(err, decision))) return true } @@ - if _, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, err.Error())) + if decision, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, reasonFromErr(err, decision))) return true }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 416 - 426, isProviderBudgetViolated currently treats a budget hit only when err != nil; change it to also detect decision-only violations by capturing the decision return from r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget and returning true when either err != nil OR the decision indicates a budget violation (e.g., DecisionBudgetExceeded). Update the two checks that call r.store.CheckProviderBudget(...) and r.store.CheckVirtualKeyBudget(...) to use the same pattern as the sibling rate-limit helper: assign (decision, err) from the store call and return true when err != nil || decision indicates a budget-exceeded result.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@framework/configstore/migrations.go`:
- Around line 6587-6616: The migration is updating all VK-owned budgets and rate
limits unconditionally; change it to only backfill when the source virtual key
was previously calendar-aligned. In the loop over virtualKeys (the
[]tables.TableVirtualKey slice), check virtualKeys[i].CalendarAligned (the
TableVirtualKey.CalendarAligned flag) and skip both the RateLimit update and the
Budgets loop when that flag is false; otherwise proceed to load and tx.Save the
rate limit and to set CalendarAligned=true and tx.Save each
virtualKeys[i].Budgets[j] as before. Ensure you still handle nil RateLimitID and
gorm.ErrRecordNotFound the same way.
In `@plugins/governance/main.go`:
- Around line 1109-1125: The hierarchy derivation only checks FK pointer fields
(hierarchyVK.CustomerID, hierarchyVK.TeamID, hierarchyVK.Team.CustomerID) so VKs
that were loaded with relations (vk.Customer.ID / vk.Team.ID) but without those
FK fields are skipped; update the customer and team ID extraction in the block
around result = p.resolver.EvaluateCustomerRequest(...) and the subsequent
team-level check to fall back to nested relation IDs (e.g., use
hierarchyVK.Customer.ID when hierarchyVK.CustomerID is nil and use
hierarchyVK.Team.ID when hierarchyVK.TeamID is nil, and similarly check
hierarchyVK.Team.Customer.ID when hierarchyVK.Team.CustomerID is nil) before
calling EvaluateCustomerRequest and EvaluateTeamRequest so both relation-loaded
and FK-populated VKs are enforced consistently.
In `@plugins/governance/modelprovidergovernance_test.go`:
- Around line 155-158: The test is asserting the wrong concrete decision:
CheckProviderRateLimit returns the first violated concrete decision (token
check) when both counters are at max, so update the assertion to expect
DecisionTokenLimited instead of DecisionRateLimited in the call that invokes
CheckProviderRateLimit with EvaluationRequest{Provider: schemas.OpenAI};
specifically replace the assert on decision (currently asserting
DecisionRateLimited) with assert.Equal(t, DecisionTokenLimited, decision) and
keep the error contains "rate limit" assertion as-is.
---
Duplicate comments:
In `@plugins/governance/resolver.go`:
- Around line 416-426: isProviderBudgetViolated currently treats a budget hit
only when err != nil; change it to also detect decision-only violations by
capturing the decision return from r.store.CheckProviderBudget and
r.store.CheckVirtualKeyBudget and returning true when either err != nil OR the
decision indicates a budget violation (e.g., DecisionBudgetExceeded). Update the
two checks that call r.store.CheckProviderBudget(...) and
r.store.CheckVirtualKeyBudget(...) to use the same pattern as the sibling
rate-limit helper: assign (decision, err) from the store call and return true
when err != nil || decision indicates a budget-exceeded result.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b519d348-5069-44f6-be39-81c214dd95f9
📒 Files selected for processing (27)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goframework/logstore/asyncjob_test.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
🚧 Files skipped from review as they are similar to previous changes (19)
- plugins/governance/resolver_test.go
- transports/bifrost-http/handlers/governance_test.go
- framework/configstore/tables/ratelimit.go
- plugins/governance/tracker.go
- plugins/governance/tracker_test.go
- framework/configstore/tables/budget.go
- framework/configstore/tables/virtualkey.go
- plugins/governance/utils.go
- transports/bifrost-http/handlers/pricing_override_test.go
- transports/bifrost-http/handlers/mcp.go
- framework/logstore/asyncjob.go
- transports/bifrost-http/handlers/governance.go
- framework/configstore/migrations_test.go
- ui/components/ui/badge.tsx
- plugins/governance/routing_test.go
- plugins/governance/store_test.go
- transports/bifrost-http/lib/config_test.go
- ui/app/workspace/governance/views/teamDialog.tsx
- ui/app/workspace/config/views/mcpView.tsx
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
73874c4 to
a54bb32
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
ui/app/workspace/governance/views/teamDialog.tsx (1)
569-606:⚠️ Potential issue | 🟡 MinorUse nullish checks for zero-valued rate-limit rows.
Line 569 and Line 606 still rely on truthiness, so stored
0limits hide the row entirely and the> 0 ? ... : 0fallback never runs. This matches the previously reported issue and looks like a regression in the current file state.Suggested fix
- {team?.rate_limit?.token_max_limit && ( + {team?.rate_limit?.token_max_limit != null && ( @@ - {team?.rate_limit?.request_max_limit && ( + {team?.rate_limit?.request_max_limit != null && (Also applies to: 606-642
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/governance/views/teamDialog.tsx` around lines 569 - 606, The token and request limit rows currently use truthy checks (team?.rate_limit?.token_max_limit and similar) which hide rows when the stored limit is 0; change those conditionals to nullish/defined checks (e.g., check that team?.rate_limit?.token_max_limit is not null/undefined) so a zero value still renders, and update any ternary guards (the > 0 ? ... : 0 logic around token_max_limit and request_max_limit) to rely on the same nullish test (use explicit != null or typeof !== "undefined") when deciding whether to compute percentages and show "Last Reset" for token_last_reset and request_last_reset in the TeamDialog component where token_current_usage, token_max_limit, request_current_usage, and request_max_limit are referenced.
🧹 Nitpick comments (2)
plugins/governance/modelprovidergovernance_test.go (1)
1998-2002: Strengthen token-only assertions with exact decision checks.Both tests set up token-limit-only violations, so asserting the concrete decision improves regression signal versus
NotEqual(DecisionAllow).♻️ Suggested test tightening
- assert.NotEqual(t, DecisionAllow, decision) + assert.Equal(t, DecisionTokenLimited, decision)Also applies to: 2051-2055
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/modelprovidergovernance_test.go` around lines 1998 - 2002, The test currently asserts token-limit violations only by checking errResult and that decision is not DecisionAllow; change those to assert the specific deny decision instead: replace assert.NotEqual(t, DecisionAllow, decision) with assert.Equal(t, DecisionDeny, decision) (or the project's explicit token-deny constant if different) for the CheckModelRateLimit call in the block using decision, errResult := store.CheckModelRateLimit(...), and make the same replacement for the analogous assertion at lines 2051-2055 so the test asserts the exact denial decision rather than just NotEqual(DecisionAllow).transports/bifrost-http/server/server.go (1)
699-699: Consider using passedctxinstead ofcontext.Background().The function receives
ctx context.Contextas a parameter but usescontext.Background()for theGetClientConfigcall. This prevents cancellation and deadline propagation from the caller's context.If this is intentional (e.g., to ensure config reload completes regardless of request timeouts), consider documenting the rationale. Otherwise, use the passed
ctxfor consistency.🔧 Suggested fix
- config, err := s.Config.ConfigStore.GetClientConfig(context.Background()) + config, err := s.Config.ConfigStore.GetClientConfig(ctx)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/server/server.go` at line 699, The call to s.Config.ConfigStore.GetClientConfig currently uses context.Background(), which ignores cancellation/deadline propagation from the function parameter ctx; change the call to pass the received ctx parameter instead (i.e., replace context.Background() with ctx) so GetClientConfig respects caller cancellation, or if you intentionally need an independent context, add a comment explaining the rationale; update the invocation in server.go where s.Config.ConfigStore.GetClientConfig(...) is called and ensure any upstream callers still behave correctly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 110-126: The transient showCalendarAlignWarning boolean isn’t
reset when switching teams, so add resetting it in the same effect that
reinitializes the form: inside the useEffect that calls createInitialState(team)
and then setInitialState and setFormData, also call
setShowCalendarAlignWarning(false) so the warning state does not carry over
between teams; update the effect containing createInitialState, setInitialState,
setFormData to include setShowCalendarAlignWarning(false).
---
Duplicate comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 569-606: The token and request limit rows currently use truthy
checks (team?.rate_limit?.token_max_limit and similar) which hide rows when the
stored limit is 0; change those conditionals to nullish/defined checks (e.g.,
check that team?.rate_limit?.token_max_limit is not null/undefined) so a zero
value still renders, and update any ternary guards (the > 0 ? ... : 0 logic
around token_max_limit and request_max_limit) to rely on the same nullish test
(use explicit != null or typeof !== "undefined") when deciding whether to
compute percentages and show "Last Reset" for token_last_reset and
request_last_reset in the TeamDialog component where token_current_usage,
token_max_limit, request_current_usage, and request_max_limit are referenced.
---
Nitpick comments:
In `@plugins/governance/modelprovidergovernance_test.go`:
- Around line 1998-2002: The test currently asserts token-limit violations only
by checking errResult and that decision is not DecisionAllow; change those to
assert the specific deny decision instead: replace assert.NotEqual(t,
DecisionAllow, decision) with assert.Equal(t, DecisionDeny, decision) (or the
project's explicit token-deny constant if different) for the CheckModelRateLimit
call in the block using decision, errResult := store.CheckModelRateLimit(...),
and make the same replacement for the analogous assertion at lines 2051-2055 so
the test asserts the exact denial decision rather than just
NotEqual(DecisionAllow).
In `@transports/bifrost-http/server/server.go`:
- Line 699: The call to s.Config.ConfigStore.GetClientConfig currently uses
context.Background(), which ignores cancellation/deadline propagation from the
function parameter ctx; change the call to pass the received ctx parameter
instead (i.e., replace context.Background() with ctx) so GetClientConfig
respects caller cancellation, or if you intentionally need an independent
context, add a comment explaining the rationale; update the invocation in
server.go where s.Config.ConfigStore.GetClientConfig(...) is called and ensure
any upstream callers still behave correctly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a9c7f419-c446-429d-893c-160269fd7021
📒 Files selected for processing (27)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goframework/logstore/asyncjob_test.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (3)
- framework/configstore/tables/ratelimit.go
- plugins/governance/routing_test.go
- framework/configstore/tables/virtualkey.go
🚧 Files skipped from review as they are similar to previous changes (12)
- plugins/governance/tracker_test.go
- transports/bifrost-http/handlers/mcp.go
- plugins/governance/tracker.go
- plugins/governance/resolver_test.go
- transports/bifrost-http/handlers/governance_test.go
- plugins/governance/utils.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/routing.go
- framework/logstore/asyncjob.go
- transports/bifrost-http/handlers/pricing_override_test.go
- plugins/governance/store_test.go
- plugins/governance/resolver.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
a54bb32 to
108770f
Compare
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/resolver.go (1)
412-429:⚠️ Potential issue | 🟠 MajorHonor decision-only budget denials in provider filtering.
This helper still blocks only on
err != nil, even thoughCheckProviderBudgetandCheckVirtualKeyBudgetnow return meaningfulDecisionvalues. If either call returnsDecisionBudgetExceeded, nil,loadBalanceProviderwill keep routing traffic to a provider that should already be excluded.🔧 Proposed fix
func (r *BudgetResolver) isProviderBudgetViolated(ctx context.Context, vk *configstoreTables.TableVirtualKey, config configstoreTables.TableVirtualKeyProviderConfig) bool { request := &EvaluationRequest{Provider: schemas.ModelProvider(config.Provider)} // 1. Check global provider-level budget first - if _, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) + if decision, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, reasonFromErr(err, decision))) return true } // 2. Check VK-level provider config budget if len(config.Budgets) == 0 { return false } - if _, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil { - r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, err.Error())) + if decision, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil || isBudgetViolation(decision) { + r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, reasonFromErr(err, decision))) return true } return false }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 412 - 429, The isProviderBudgetViolated helper currently only treats non-nil errors as violations; update it to inspect the Decision returned by r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget as well: call them as (decision, err) := r.store.CheckProviderBudget(...) and if decision == DecisionBudgetExceeded return true; if err != nil log and return true; do the same for r.store.CheckVirtualKeyBudget (use vk and config) so DecisionBudgetExceeded is honored when filtering providers in loadBalanceProvider.
♻️ Duplicate comments (1)
ui/app/workspace/governance/views/teamDialog.tsx (1)
570-643:⚠️ Potential issue | 🟡 MinorUse nullish guards for the token/request usage rows too.
Line 570 and Line 607 still rely on truthiness, so a stored
0limit hides the entire row and bypasses the new> 0 ? ... : 0fallback you added below. Switching both guards to!= nullkeeps the usage section consistent with the updated badge logic.Suggested fix
- {team?.rate_limit?.token_max_limit && ( + {team?.rate_limit?.token_max_limit != null && ( @@ - {team?.rate_limit?.request_max_limit && ( + {team?.rate_limit?.request_max_limit != null && (🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/governance/views/teamDialog.tsx` around lines 570 - 643, The token and request usage sections currently hide when a limit is 0 because they use truthy checks; replace the conditional guards team?.rate_limit?.token_max_limit and team?.rate_limit?.request_max_limit with nullish checks (e.g., team?.rate_limit?.token_max_limit != null and team?.rate_limit?.request_max_limit != null) so the rows render for zero values and keep the existing Badge fallback logic intact.
🧹 Nitpick comments (2)
framework/configstore/migrations.go (1)
6585-6623: Prefer set-based updates for this backfill.Loading every virtual key and
Save-ing each child row one by one will make startup cost grow with tenant size. A couple ofUPDATE ... WHERE EXISTS (...)statements would make the migration much cheaper and more resilient on large installs.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 6585 - 6623, The current migration loads all TableVirtualKey rows and calls tx.Save per-rate-limit and per-budget which scales poorly; replace the per-row loads/saves with set-based UPDATEs using tx.Exec: 1) UPDATE table_rate_limits SET calendar_aligned = TRUE WHERE id IN (SELECT rate_limit_id FROM table_virtual_keys WHERE calendar_aligned = TRUE AND rate_limit_id IS NOT NULL) to mirror the existing "skip stale rate_limit_id" behavior; 2) UPDATE table_budgets SET calendar_aligned = TRUE WHERE virtual_key_id IN (SELECT id FROM table_virtual_keys WHERE calendar_aligned = TRUE) to update all budgets in one statement; run both via tx.Exec and return errors on Exec failure, and remove the tx.Preload/Find + Save loops (symbols: TableVirtualKey, RateLimitID, TableRateLimit, Budgets, CalendarAligned, tx.Preload/Find, tx.Save).framework/configstore/migrations_test.go (1)
2187-2253: Preferbifrost.Ptr()over local pointer helper for consistency.
strPtrcan be removed and callsites can use the project-wide pointer helper directly.♻️ Proposed cleanup
- insertBudgetRaw(t, db, "budget-stale-vk", strPtr("vk-stale")) + insertBudgetRaw(t, db, "budget-stale-vk", bifrost.Ptr("vk-stale")) ... -func strPtr(s string) *string { return &s }Based on learnings: In this repository, prefer
bifrost.Ptr()for pointer creation across code paths, including tests.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations_test.go` around lines 2187 - 2253, Replace the local helper strPtr with the project-wide helper bifrost.Ptr: update all callsites that use strPtr (e.g., the insertBudgetRaw call that passes strPtr("vk-stale")) to call bifrost.Ptr("...") instead, then remove the strPtr function declaration at the bottom of the file; ensure imports include the package that exposes bifrost.Ptr if not already referenced.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@plugins/governance/main.go`:
- Around line 1101-1152: The code validates hierarchy (customer/team/user)
before the VirtualKey is filtered, which can let revoked or disallowed VKs
produce hierarchy 402/429 instead of a VK 403; call
p.resolver.EvaluateVirtualKeyRequest(ctx, evaluationRequest.VirtualKey,
evaluationRequest.Provider, evaluationRequest.Model, requestType, false) as an
early short-circuit right after resolving hierarchyVK (or immediately after
reading evaluationRequest.VirtualKey) and if the returned result.Decision !=
DecisionAllow return that result; only when the VK validation passes continue to
run the existing customer/team/user budget checks (the current uses of
hierarchyVK, EvaluateCustomerRequest, EvaluateTeamRequest, EvaluateUserRequest
should remain unchanged).
---
Outside diff comments:
In `@plugins/governance/resolver.go`:
- Around line 412-429: The isProviderBudgetViolated helper currently only treats
non-nil errors as violations; update it to inspect the Decision returned by
r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget as well: call them
as (decision, err) := r.store.CheckProviderBudget(...) and if decision ==
DecisionBudgetExceeded return true; if err != nil log and return true; do the
same for r.store.CheckVirtualKeyBudget (use vk and config) so
DecisionBudgetExceeded is honored when filtering providers in
loadBalanceProvider.
---
Duplicate comments:
In `@ui/app/workspace/governance/views/teamDialog.tsx`:
- Around line 570-643: The token and request usage sections currently hide when
a limit is 0 because they use truthy checks; replace the conditional guards
team?.rate_limit?.token_max_limit and team?.rate_limit?.request_max_limit with
nullish checks (e.g., team?.rate_limit?.token_max_limit != null and
team?.rate_limit?.request_max_limit != null) so the rows render for zero values
and keep the existing Badge fallback logic intact.
---
Nitpick comments:
In `@framework/configstore/migrations_test.go`:
- Around line 2187-2253: Replace the local helper strPtr with the project-wide
helper bifrost.Ptr: update all callsites that use strPtr (e.g., the
insertBudgetRaw call that passes strPtr("vk-stale")) to call bifrost.Ptr("...")
instead, then remove the strPtr function declaration at the bottom of the file;
ensure imports include the package that exposes bifrost.Ptr if not already
referenced.
In `@framework/configstore/migrations.go`:
- Around line 6585-6623: The current migration loads all TableVirtualKey rows
and calls tx.Save per-rate-limit and per-budget which scales poorly; replace the
per-row loads/saves with set-based UPDATEs using tx.Exec: 1) UPDATE
table_rate_limits SET calendar_aligned = TRUE WHERE id IN (SELECT rate_limit_id
FROM table_virtual_keys WHERE calendar_aligned = TRUE AND rate_limit_id IS NOT
NULL) to mirror the existing "skip stale rate_limit_id" behavior; 2) UPDATE
table_budgets SET calendar_aligned = TRUE WHERE virtual_key_id IN (SELECT id
FROM table_virtual_keys WHERE calendar_aligned = TRUE) to update all budgets in
one statement; run both via tx.Exec and return errors on Exec failure, and
remove the tx.Preload/Find + Save loops (symbols: TableVirtualKey, RateLimitID,
TableRateLimit, Budgets, CalendarAligned, tx.Preload/Find, tx.Save).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ffb514db-5773-4189-9939-37568d0d7197
📒 Files selected for processing (27)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goframework/logstore/asyncjob_test.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (6)
- plugins/governance/routing_test.go
- framework/configstore/tables/ratelimit.go
- framework/configstore/tables/virtualkey.go
- transports/bifrost-http/handlers/governance_test.go
- plugins/governance/modelprovidergovernance_test.go
- transports/bifrost-http/lib/config_test.go
🚧 Files skipped from review as they are similar to previous changes (12)
- transports/bifrost-http/handlers/mcp.go
- plugins/governance/resolver_test.go
- plugins/governance/tracker.go
- plugins/governance/tracker_test.go
- framework/logstore/asyncjob.go
- transports/bifrost-http/handlers/pricing_override_test.go
- transports/bifrost-http/server/server.go
- framework/logstore/asyncjob_test.go
- ui/app/workspace/config/views/mcpView.tsx
- ui/components/ui/badge.tsx
- plugins/governance/store_test.go
- framework/configstore/tables/budget.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
108770f to
46c8fde
Compare
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/resolver.go (1)
412-430:⚠️ Potential issue | 🟠 MajorInconsistent decision handling — budget decisions are ignored.
isProviderBudgetViolateddiscards thedecisionreturn value with_and only checkserr != nil. This is inconsistent with the hardening applied to all other call sites (e.g., lines 98, 115, 146, 176, 398) whereisBudgetViolation(decision)is also checked. If the store returnsDecisionBudgetExceededwith anilerror, this function would incorrectly returnfalse, allowing requests to bypass provider budget governance.🔧 Proposed fix
func (r *BudgetResolver) isProviderBudgetViolated(ctx context.Context, vk *configstoreTables.TableVirtualKey, config configstoreTables.TableVirtualKeyProviderConfig) bool { request := &EvaluationRequest{Provider: schemas.ModelProvider(config.Provider)} // 1. Check global provider-level budget first - if _, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil { + if decision, err := r.store.CheckProviderBudget(ctx, request, nil); err != nil || isBudgetViolation(decision) { r.logger.Debug(fmt.Sprintf("Global provider budget exceeded for provider %s: %s", config.Provider, err.Error())) return true } // 2. Check VK-level provider config budget if len(config.Budgets) == 0 { return false } - if _, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil { + if decision, err := r.store.CheckVirtualKeyBudget(ctx, vk, request, nil); err != nil || isBudgetViolation(decision) { r.logger.Debug(fmt.Sprintf("VK provider config budget exceeded for VK %s: %s", vk.ID, err.Error())) return true } return false }Note: The debug log formatting will also need adjustment to handle
err == nilcases (usereasonFromErr(err, decision)instead oferr.Error()).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/resolver.go` around lines 412 - 430, The function isProviderBudgetViolated currently ignores the returned decision from r.store.CheckProviderBudget and r.store.CheckVirtualKeyBudget; change it to capture both (decision, err) and return true when isBudgetViolation(decision) || err != nil; update the debug logs to call reasonFromErr(err, decision) instead of err.Error() so nil errors with DecisionBudgetExceeded are logged correctly; ensure you reference the decision handling for both CheckProviderBudget and CheckVirtualKeyBudget and cover the DecisionBudgetExceeded case.
🧹 Nitpick comments (1)
framework/configstore/migrations_test.go (1)
2253-2253: Usebifrost.Ptrin helper for repository consistency.
strPtrworks, but this codebase prefersbifrost.Ptr(...)over&value, including test helpers.♻️ Suggested refactor
-func strPtr(s string) *string { return &s } +func strPtr(s string) *string { return bifrost.Ptr(s) }Based on learnings: in this repository, prefer using
bifrost.Ptr()over address-of (&) consistently across all code paths, including tests.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations_test.go` at line 2253, Replace the local helper strPtr(s string) *string that returns &s with a call to the repository helper bifrost.Ptr to keep consistency; locate the strPtr function in migrations_test.go and remove or replace its usages so tests call bifrost.Ptr("...") instead of strPtr(...) (or inline bifrost.Ptr where strPtr is referenced).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@transports/bifrost-http/server/server.go`:
- Around line 643-649: The governance plugin selection must be done from the
incoming request context, not the server bootstrap context: change
GetGovernanceData to resolve the plugin name from ctx (read the
BifrostContextKeyGovernancePluginName request value) and call
lib.FindPluginAs[governance.BaseGovernancePlugin](s.Config, pluginNameFromCtx)
(or add/use a s.getGovernancePlugin(ctx) helper) so the returned plugin and its
store come from the request-scoped override; apply the same pattern to other
callers that use s.getGovernancePlugin() so they accept ctx and use the request
value rather than s.getGovernancePluginName().
---
Outside diff comments:
In `@plugins/governance/resolver.go`:
- Around line 412-430: The function isProviderBudgetViolated currently ignores
the returned decision from r.store.CheckProviderBudget and
r.store.CheckVirtualKeyBudget; change it to capture both (decision, err) and
return true when isBudgetViolation(decision) || err != nil; update the debug
logs to call reasonFromErr(err, decision) instead of err.Error() so nil errors
with DecisionBudgetExceeded are logged correctly; ensure you reference the
decision handling for both CheckProviderBudget and CheckVirtualKeyBudget and
cover the DecisionBudgetExceeded case.
---
Nitpick comments:
In `@framework/configstore/migrations_test.go`:
- Line 2253: Replace the local helper strPtr(s string) *string that returns &s
with a call to the repository helper bifrost.Ptr to keep consistency; locate the
strPtr function in migrations_test.go and remove or replace its usages so tests
call bifrost.Ptr("...") instead of strPtr(...) (or inline bifrost.Ptr where
strPtr is referenced).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 21575033-2006-414c-b645-01e58c7a439f
📒 Files selected for processing (27)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goframework/logstore/asyncjob_test.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (5)
- plugins/governance/routing.go
- framework/configstore/tables/ratelimit.go
- ui/app/workspace/config/views/mcpView.tsx
- transports/bifrost-http/handlers/governance.go
- transports/bifrost-http/lib/config_test.go
🚧 Files skipped from review as they are similar to previous changes (15)
- plugins/governance/tracker_test.go
- transports/bifrost-http/handlers/governance_test.go
- plugins/governance/resolver_test.go
- framework/configstore/tables/budget.go
- transports/bifrost-http/handlers/pricing_override_test.go
- transports/bifrost-http/handlers/mcp.go
- plugins/governance/utils.go
- framework/configstore/tables/virtualkey.go
- plugins/governance/tracker.go
- ui/components/ui/badge.tsx
- framework/logstore/asyncjob_test.go
- plugins/governance/routing_test.go
- ui/app/workspace/governance/views/teamDialog.tsx
- framework/configstore/migrations.go
- plugins/governance/store_test.go
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
46c8fde to
ead875f
Compare
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
transports/bifrost-http/server/server.go (1)
618-633:⚠️ Potential issue | 🟠 MajorDo not fail provider deletion after the provider is already removed.
This path deletes the provider from the client and config first, then unconditionally resolves the governance plugin. If governance is disabled or lookup fails, the API returns an error even though the deletion already happened, which makes the operation appear failed and turns retries into confusing
not foundcases.Suggested fix
err = s.Config.RemoveProvider(ctx, provider) if err != nil && !errors.Is(err, lib.ErrNotFound) { logger.Error("failed to remove provider from config: %v. Client and config may be out of sync, please restart bifrost", err) return fmt.Errorf("failed to remove provider from config: %w. Client and config may be out of sync, please restart bifrost", err) } - governancePlugin, err := s.getGovernancePlugin() - if err != nil { - return err - } - governancePlugin.GetGovernanceStore().DeleteProviderInMemory(ctx, string(provider)) + if s.Config.IsPluginLoaded(s.getGovernancePluginName()) { + governancePlugin, err := s.getGovernancePlugin() + if err != nil { + logger.Warn("governance plugin found but failed to get: %v", err) + } else { + governancePlugin.GetGovernanceStore().DeleteProviderInMemory(ctx, string(provider)) + } + } if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("pricing manager not found") }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/server/server.go` around lines 618 - 633, RemoveProvider currently returns an error if getGovernancePlugin fails after the provider was already removed from the client/config; change RemoveProvider so that after s.Client.RemoveProvider and s.Config.RemoveProvider complete (including "not found" cases), calls to getGovernancePlugin and governancePlugin.GetGovernanceStore().DeleteProviderInMemory are best-effort only: if getGovernancePlugin returns an error, log it and return nil (do not propagate), and if DeleteProviderInMemory fails, log the error but do not return it (ignore not-found cases similarly). Keep existing error handling for s.Client.RemoveProvider and s.Config.RemoveProvider but ensure governance lookup/cleanup cannot cause the overall deletion to fail.
🧹 Nitpick comments (2)
framework/configstore/migrations_test.go (1)
2253-2253: Usebifrost.Ptrinstead of address-of helper for consistency.
strPtrworks, but repository convention prefersbifrost.Ptr(...)across code paths (including tests).♻️ Optional cleanup
-func strPtr(s string) *string { return &s } +func strPtr(s string) *string { return bifrost.Ptr(s) }Based on learnings: in maximhq/bifrost, prefer
bifrost.Ptr()over&valuefor pointer creation in all code paths, including tests.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations_test.go` at line 2253, Replace the local helper strPtr(s string) *string with the repository-standard bifrost.Ptr usage: find and remove the strPtr function and update all call sites in this file (e.g., where strPtr(...) is used in migrations_test.go) to call bifrost.Ptr("...") instead so tests use bifrost.Ptr for pointer creation consistently with the rest of the codebase.plugins/governance/main.go (1)
1062-1120: Consider consolidating redundant VK lookups.The virtual key is looked up three times for the same
evaluationRequest.VirtualKey:
- Line 1064 — existence check
- Line 1109 —
hierarchyVKresolution- Inside
EvaluateVirtualKeyRequestat line 1119 (seeresolver.go:239)While
sync.Map.Loadis fast, consolidating#1and#2would reduce redundancy. For example, the existence check at line 1064 could usehierarchyVK != nilafter moving the resolution earlier:// Resolve VK once, use for both existence check and hierarchy extraction var hierarchyVK *configstoreTables.TableVirtualKey if evaluationRequest.VirtualKey != "" { vk, exists := p.store.GetVirtualKey(ctx, evaluationRequest.VirtualKey) if !exists { return nil, &schemas.BifrostError{...} // 401 not found } hierarchyVK = vk }The third lookup inside
EvaluateVirtualKeyRequestis harder to eliminate since it also sets governance context values — this could be a follow-up if you refactor the resolver interface.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/main.go` around lines 1062 - 1120, The code redundantly looks up the same virtual key multiple times; resolve the VK once early and reuse it: call p.store.GetVirtualKey(ctx, evaluationRequest.VirtualKey) once to both validate existence (returning the 401 virtual_key_not_found error if missing) and assign hierarchyVK, then use hierarchyVK (and its existence) instead of re-calling GetVirtualKey for the initial check; keep the remaining lookup inside resolver.EvaluateVirtualKeyRequest for now (since it sets governance context) but pass evaluationRequest.VirtualKey and hierarchyVK as needed so the initial existence check and hierarchy extraction reuse the single resolved TableVirtualKey.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@transports/bifrost-http/server/server.go`:
- Around line 618-633: RemoveProvider currently returns an error if
getGovernancePlugin fails after the provider was already removed from the
client/config; change RemoveProvider so that after s.Client.RemoveProvider and
s.Config.RemoveProvider complete (including "not found" cases), calls to
getGovernancePlugin and
governancePlugin.GetGovernanceStore().DeleteProviderInMemory are best-effort
only: if getGovernancePlugin returns an error, log it and return nil (do not
propagate), and if DeleteProviderInMemory fails, log the error but do not return
it (ignore not-found cases similarly). Keep existing error handling for
s.Client.RemoveProvider and s.Config.RemoveProvider but ensure governance
lookup/cleanup cannot cause the overall deletion to fail.
---
Nitpick comments:
In `@framework/configstore/migrations_test.go`:
- Line 2253: Replace the local helper strPtr(s string) *string with the
repository-standard bifrost.Ptr usage: find and remove the strPtr function and
update all call sites in this file (e.g., where strPtr(...) is used in
migrations_test.go) to call bifrost.Ptr("...") instead so tests use bifrost.Ptr
for pointer creation consistently with the rest of the codebase.
In `@plugins/governance/main.go`:
- Around line 1062-1120: The code redundantly looks up the same virtual key
multiple times; resolve the VK once early and reuse it: call
p.store.GetVirtualKey(ctx, evaluationRequest.VirtualKey) once to both validate
existence (returning the 401 virtual_key_not_found error if missing) and assign
hierarchyVK, then use hierarchyVK (and its existence) instead of re-calling
GetVirtualKey for the initial check; keep the remaining lookup inside
resolver.EvaluateVirtualKeyRequest for now (since it sets governance context)
but pass evaluationRequest.VirtualKey and hierarchyVK as needed so the initial
existence check and hierarchy extraction reuse the single resolved
TableVirtualKey.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 50c1dc24-9241-4012-9a14-993419e67211
📒 Files selected for processing (27)
framework/configstore/migrations.goframework/configstore/migrations_test.goframework/configstore/tables/budget.goframework/configstore/tables/ratelimit.goframework/configstore/tables/virtualkey.goframework/logstore/asyncjob.goframework/logstore/asyncjob_test.goplugins/governance/main.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/routing.goplugins/governance/routing_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/tracker.goplugins/governance/tracker_test.goplugins/governance/utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/pricing_override_test.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/governance/views/teamDialog.tsxui/components/ui/badge.tsx
✅ Files skipped from review due to trivial changes (4)
- transports/bifrost-http/handlers/mcp.go
- framework/configstore/tables/ratelimit.go
- plugins/governance/store_test.go
- transports/bifrost-http/lib/config_test.go
🚧 Files skipped from review as they are similar to previous changes (14)
- plugins/governance/resolver_test.go
- transports/bifrost-http/handlers/governance_test.go
- plugins/governance/tracker.go
- framework/logstore/asyncjob.go
- ui/app/workspace/config/views/mcpView.tsx
- plugins/governance/tracker_test.go
- framework/configstore/tables/virtualkey.go
- transports/bifrost-http/handlers/governance.go
- framework/configstore/migrations.go
- plugins/governance/resolver.go
- plugins/governance/routing.go
- framework/logstore/asyncjob_test.go
- plugins/governance/modelprovidergovernance_test.go
- ui/app/workspace/governance/views/teamDialog.tsx
akshaydeo
force-pushed
the
threadsafty_for_plugin_timing_detectors
branch
from
7a5168a to
fe043ae
Compare
akshaydeo
force-pushed
the
04-19-team_budgets
branch
from
ead875f to
b963068
Compare
18 tasks
This was referenced Apr 19, 2026
18 tasks
Contributor
Author
Merge activity
|
akshaydeo
changed the base branch from
threadsafty_for_plugin_timing_detectors
to
graphite-base/2836
akshaydeo
dismissed
coderabbitai[bot]’s stale review
The base branch was changed.
dominictayloruk
pushed a commit
to dominictayloruk/bifrost
that referenced
this pull request
Apr 21, 2026
## Summary Refactors GovernanceStore for reusability ## Changes - Removes unnecessary governance objects - Removes dead code ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [ ] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [ ] Docs ## How to test Describe the steps to validate this change. Include commands and expected outcomes. ```sh # Core/Transports go version go test ./... # UI cd ui pnpm i || npm i pnpm test || npm test pnpm build || npm run build ``` If adding new configs or environment variables, document them here. ## Screenshots/Recordings If UI changes, add before/after screenshots or short clips. ## Breaking changes - [ ] Yes - [x] No If yes, describe impact and migration instructions. ## Related issues Link related issues and discussions. Example: Closes maximhq#123 ## Security considerations Note any security implications (auth, secrets, PII, sandboxing, etc.). ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [x] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [x] I verified the CI pipeline passes locally if applicable
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
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Refactors GovernanceStore for reusability
Changes
Type of change
Affected areas
How to test
Describe the steps to validate this change. Include commands and expected outcomes.
If adding new configs or environment variables, document them here.
Screenshots/Recordings
If UI changes, add before/after screenshots or short clips.
Breaking changes
If yes, describe impact and migration instructions.
Related issues
Link related issues and discussions. Example: Closes #123
Security considerations
Note any security implications (auth, secrets, PII, sandboxing, etc.).
Checklist
docs/contributing/README.mdand followed the guidelines