Conversation
|
|
Contributor
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (3)
📒 Files selected for processing (15)
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughMigrate from single-budget to multi-budget ownership: budgets become first-class rows linked by Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Frontend
participant API as Bifrost HTTP Handler
participant Store as Governance Store
participant DB as Database
rect rgba(60,179,113,0.5)
Client->>API: CreateVirtualKey {budgets[], calendar_aligned}
API->>API: Validate budgets (parse, duplicates)
API->>DB: INSERT TableVirtualKey (calendar_aligned)
loop per budget
API->>API: compute LastReset (vk.calendar_aligned + reset_duration)
API->>DB: INSERT TableBudget (virtual_key_id/provider_config_id, reset_duration, max_limit)
end
API->>Store: Refresh in-memory budgets
API->>Client: Return VirtualKey with budgets[]
end
rect rgba(70,130,180,0.5)
Client->>API: UpdateVirtualKey {budgets[]}
API->>DB: SELECT existing budgets for VirtualKey
API->>API: Match incoming budgets by reset_duration
alt match -> update
API->>DB: UPDATE TableBudget (preserve CurrentUsage)
else new -> insert
API->>DB: INSERT TableBudget
else missing -> delete
API->>DB: DELETE TableBudget
end
API->>Store: Reconcile in-memory budgets
API->>Client: Return updated VirtualKey with budgets[]
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
Author
This stack of pull requests is managed by Graphite. Learn more about stacking. |
Contributor
Confidence Score: 4/5Not safe to merge as-is: provider-level multi-budget creation via the UI is broken due to unconverted string max_limit values. One confirmed P1 defect: provider config budgets cannot be set from the UI because max_limit is sent as a string instead of a number, causing a JSON deserialization error in the Go API. The remaining issues are P2 (calendar alignment regression for teams/customers, and a performance concern in ResetExpiredBudgetsInMemory). The core backend logic, migration, and tests are solid, bringing the score to 4 rather than lower. ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (normalizeProviderConfigs must convert budgets[].max_limit to a number); transports/bifrost-http/handlers/governance.go (calendar alignment removal from teams/customers is unannounced) Important Files Changed
|
![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: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ui/lib/types/governance.ts (1)
160-173:⚠️ Potential issue | 🟠 MajorThe multi-budget request shape is lossy on update.
UpdateVirtualKeyRequest.budgetsisCreateBudgetRequest[], so callers have no way to send the IDs of existing budgets even thoughBudgetin this file is stateful (id,current_usage,last_reset). That forces the backend to guess which rows should be updated versus recreated, which can reset usage whenever a virtual key is edited. The same file also still models the read/provider-config side as singularbudget, so the new field is not fully round-trippable end-to-end.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/lib/types/governance.ts` around lines 160 - 173, UpdateVirtualKeyRequest.budgets currently uses CreateBudgetRequest[] which loses existing budget IDs and prevents proper updates; change the type to accept stateful budget update payloads (e.g., UpdateBudgetRequest[] or a union like (UpdateBudgetRequest | CreateBudgetRequest)[]) so callers can include ids for existing Budget entities and new budget data for creation, and ensure the read-side still exposes the singular budget field consistently (align UpdateVirtualKeyRequest.budget and budgets types with the Budget model that contains id/current_usage/last_reset) so updates are round-trippable without forcing recreation.
🤖 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 5317-5346: migrationAddBudgetCalendarAlignedColumn is adding a
"calendar_aligned" column but GORM needs the corresponding struct field on
tables.TableBudget to infer type and tags; add a CalendarAligned bool field to
the TableBudget struct (named exactly CalendarAligned) with appropriate gorm
struct tag (nullable/default/index as required by your schema) in the
TableBudget definition so mg.AddColumn(&tables.TableBudget{},
"calendar_aligned") can succeed, then re-run the migration.
In `@framework/configstore/tables/budget.go`:
- Around line 29-35: The join structs TableVirtualKeyBudget and
TableVirtualKeyProviderConfigBudget lack parent relation fields so GORM doesn't
create FK constraints to their parent tables; add a VirtualKey field to
TableVirtualKeyBudget with
`gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE"` referencing the
TableVirtualKey model, and add a VirtualKeyProviderConfig (or
TableVirtualKeyProviderConfig) field to TableVirtualKeyProviderConfigBudget with
`gorm:"foreignKey:ProviderConfigID;constraint:OnDelete:CASCADE"` referencing the
provider-config model so deletions cascade and no orphaned join rows remain.
In `@framework/configstore/tables/virtualkey.go`:
- Around line 39-40: The Multi-budget association on TableVirtualKey (Budgets
[]TableBudget) isn’t being persisted or loaded: update UpdateVirtualKey to
handle many-to-many budgets instead of the legacy singular budget_id selection,
change the provider-config association logic that currently only replaces Keys
to also replace the Budgets association on ProviderConfigs, and adjust
ListVirtualKeys to preload Budgets and ProviderConfigs.Budgets (not just
Budget/ProviderConfigs.Budget) so GETs return the stored Budgets; ensure you use
GORM’s Association Replace/Preload APIs for the TableVirtualKey.Budgets and
ProviderConfig.Budgets relationships when saving and loading (touch functions
UpdateVirtualKey, the provider-config save path, and ListVirtualKeys).
In `@ui/app/workspace/governance/access-profiles/page.tsx`:
- Around line 3-7: Add an explicit RBAC guard inside AccessProfilesPage: call
useRbac(RbacResource.AccessProfiles, RbacOperation.View) and store the result
(e.g., hasAccessProfilesAccess); if false, render NoPermissionView with
entity="access-profiles" instead of rendering AccessProfilesIndexView. This
mirrors the pattern used by other governance pages (e.g., virtual-keys, teams,
customers) and ensures page-level protection even though AccessProfilesIndexView
is from an external package.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 861-878: Add a data-testid to the MultiBudgetLines element (e.g.,
data-testid="vk-provider-budget-${index}"), replace the `(b: any)` cast with the
proper BudgetLineEntry type when mapping existing budgets, and fix the onChange
handler logic in the MultiBudgetLines usage: instead of writing only lines[0]
into config.budget, persist all lines to config.budgets (call
handleUpdateProviderConfig(index, "budgets", lines) when lines.length > 0) and
clear budgets (or set to undefined) when lines.length === 0; keep config.budget
only if you intentionally support a single-entry legacy field, otherwise remove
its usage to avoid confusion and ensure consistency between the lines prop and
what you persist.
In `@ui/components/ui/multiBudgetLines.tsx`:
- Around line 45-48: The new Add Budget button (rendered with Button and
onClick={addLine}) and the per-row delete icon button lack stable test selectors
and the delete button lacks an accessible name; add data-testid attributes
following the pattern data-testid="budget-<element>-<qualifier>" (e.g.,
data-testid="budget-add-button" on the Add Budget Button and
data-testid="budget-row-delete-<index|id>" on each row delete control) and
ensure the icon-only delete button has an accessible name (aria-label or
aria-labelledby, e.g., aria-label="Delete budget row") so both testing and
screen readers can reliably target remove actions; update the corresponding JSX
where addLine is used and where rows are rendered (the delete button element) to
include these attributes.
In `@ui/lib/store/apis/baseApi.ts`:
- Line 165: The AccessProfiles tag declared in the API tag types is never
applied to endpoints — update the relevant endpoints in the
governanceApi/enterprise API slice to use providesTags and invalidatesTags so
cache updates work: add providesTags: ["AccessProfiles"] to the list/fetch
endpoint (e.g., getAccessProfiles) and add invalidatesTags: ["AccessProfiles"]
to mutating endpoints (e.g., createAccessProfile, updateAccessProfile,
deleteAccessProfile); ensure names match the tag string "AccessProfiles" and
apply the annotations on the endpoint definitions inside the API slice so RTK
Query can drive cache refreshes.
---
Outside diff comments:
In `@ui/lib/types/governance.ts`:
- Around line 160-173: UpdateVirtualKeyRequest.budgets currently uses
CreateBudgetRequest[] which loses existing budget IDs and prevents proper
updates; change the type to accept stateful budget update payloads (e.g.,
UpdateBudgetRequest[] or a union like (UpdateBudgetRequest |
CreateBudgetRequest)[]) so callers can include ids for existing Budget entities
and new budget data for creation, and ensure the read-side still exposes the
singular budget field consistently (align UpdateVirtualKeyRequest.budget and
budgets types with the Budget model that contains id/current_usage/last_reset)
so updates are round-trippable without forcing recreation.
🪄 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: 8e9757e7-5d06-4b6d-8edd-39cae172ac05
📒 Files selected for processing (10)
framework/configstore/migrations.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/components/sidebar.tsxui/components/ui/multiBudgetLines.tsxui/lib/store/apis/baseApi.tsui/lib/types/governance.ts
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
c094b66 to
a40a014
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
356-408:⚠️ Potential issue | 🟡 MinorConvert
budgetsarraymax_limitvalues from string to number innormalizeProviderConfigs.The
budgetsarray is spread into the normalized config but itsmax_limitvalues remain as strings from the form. This mirrors the singularbudgetfield, which explicitly convertsmax_limitvianormalizeNumericField. Add similar normalization for thebudgetsarray:budgets: config.budgets?.map((b) => ({ max_limit: normalizeNumericField(b.max_limit), reset_duration: b.reset_duration, })) || undefined,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 356 - 408, The normalizeProviderConfigs function currently leaves config.budgets' max_limit values as strings; update normalizeProviderConfigs (the function named normalizeProviderConfigs) to map config.budgets and convert each b.max_limit using normalizeNumericField (like the singular budget logic does), returning budgets: config.budgets?.map(b => ({ max_limit: normalizeNumericField(b.max_limit), reset_duration: b.reset_duration })) || undefined so that budgets' max_limit become numeric or undefined and follow the same reset_duration fallback behavior as budget.
♻️ Duplicate comments (1)
framework/configstore/tables/budget.go (1)
29-51:⚠️ Potential issue | 🟠 MajorMissing parent-side foreign key constraints on junction tables.
Both
TableVirtualKeyBudgetandTableVirtualKeyProviderConfigBudgetdefineBudgetrelations withOnDelete:CASCADE, but lack parent relations forVirtualKeyIDandProviderConfigIDrespectively. GORM won't generate FK constraints from these ID columns to the parent tables (governance_virtual_keys,governance_virtual_key_provider_configs), leaving orphan junction rows when parents are deleted.🔧 Proposed fix to add parent relations
type TableVirtualKeyBudget struct { ID uint `gorm:"primaryKey;autoIncrement" json:"id"` VirtualKeyID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_vk_budget" json:"virtual_key_id"` BudgetID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_vk_budget" json:"budget_id"` + VirtualKey TableVirtualKey `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"-"` Budget TableBudget `gorm:"foreignKey:BudgetID;constraint:OnDelete:CASCADE" json:"budget"` } type TableVirtualKeyProviderConfigBudget struct { ID uint `gorm:"primaryKey;autoIncrement" json:"id"` ProviderConfigID uint `gorm:"not null;uniqueIndex:idx_pc_budget" json:"provider_config_id"` BudgetID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_pc_budget" json:"budget_id"` + ProviderConfig TableVirtualKeyProviderConfig `gorm:"foreignKey:ProviderConfigID;constraint:OnDelete:CASCADE" json:"-"` Budget TableBudget `gorm:"foreignKey:BudgetID;constraint:OnDelete:CASCADE" json:"budget"` }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/budget.go` around lines 29 - 51, The junction structs lack parent-side relationship fields so GORM won't create FKs for VirtualKeyID and ProviderConfigID; add relation fields to both TableVirtualKeyBudget and TableVirtualKeyProviderConfigBudget that reference their parent models (e.g., add a VirtualKey field of type TableVirtualKey on TableVirtualKeyBudget and a ProviderConfig field of type TableVirtualKeyProviderConfig on TableVirtualKeyProviderConfigBudget) and annotate them with gorm tags to point to the ID fields and enforce constraint: e.g. gorm:"foreignKey:VirtualKeyID;references:ID;constraint:OnDelete:CASCADE" and gorm:"foreignKey:ProviderConfigID;references:ID;constraint:OnDelete:CASCADE" so deletions of parents cascade and foreign keys are created.
🧹 Nitpick comments (1)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
902-941: Missingdata-testidattribute on MultiBudgetLines.Per coding guidelines, add a
data-testidattribute following the pattern<entity>-<element>-<qualifier>.🧪 Proposed fix
<MultiBudgetLines id={`providerBudget-${index}`} + data-testid={`vk-provider-budget-${index}`} label="Provider Budget"As per coding guidelines:
ui/**/*.{tsx,ts}: Add new interactive UI elements withdata-testidattributes following the pattern:data-testid="<entity>-<element>-<qualifier>".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 902 - 941, The MultiBudgetLines component instance is missing a data-testid; add a data-testid attribute following the project's pattern (entity-element-qualifier) to this instance (the MultiBudgetLines with id `providerBudget-${index}`) so tests can target it reliably; update the JSX for MultiBudgetLines to include something like data-testid="virtualkey-provider-budget" or a variant that includes the index/qualifier you prefer so it follows the `<entity>-<element>-<qualifier>` convention while leaving existing props and the onChange handling (handleUpdateProviderConfig) unchanged.
🤖 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/tables/virtualkey.go`:
- Around line 39-40: Update the comment on the Keys field to reflect
deny-by-default semantics: change the current "Empty means all keys allowed for
this provider" to explain that interpretation depends on the AllowAllKeys flag —
when AllowAllKeys is true, all keys are allowed; when AllowAllKeys is false and
Keys is empty, no keys are allowed (deny-by-default). Reference the Keys field
and the AllowAllKeys flag on the TableKey/TableVirtualKey provider config types
so the relationship is clear.
- Line 231: Update the GORM preloads to include the missing Budgets
associations: modify the query functions GetVirtualKeys, GetVirtualKey,
GetVirtualKeyByValue, and GetVirtualKeysPaginated to add Preload("Budgets") and
Preload("ProviderConfigs.Budgets") when loading TableVirtualKey so the
many-to-many TableBudget relations are eagerly loaded, and change
GetVirtualKeyProviderConfigs to use full preloading (e.g., Preload("Budgets")
and any nested Preload("...ProviderConfig.Budgets") patterns used) for
TableVirtualKeyProviderConfig to eliminate N+1 queries; locate these changes by
editing the functions named above and the structs TableVirtualKey and
TableVirtualKeyProviderConfig where the Budgets field is defined to ensure the
correct association names are used.
---
Outside diff comments:
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 356-408: The normalizeProviderConfigs function currently leaves
config.budgets' max_limit values as strings; update normalizeProviderConfigs
(the function named normalizeProviderConfigs) to map config.budgets and convert
each b.max_limit using normalizeNumericField (like the singular budget logic
does), returning budgets: config.budgets?.map(b => ({ max_limit:
normalizeNumericField(b.max_limit), reset_duration: b.reset_duration })) ||
undefined so that budgets' max_limit become numeric or undefined and follow the
same reset_duration fallback behavior as budget.
---
Duplicate comments:
In `@framework/configstore/tables/budget.go`:
- Around line 29-51: The junction structs lack parent-side relationship fields
so GORM won't create FKs for VirtualKeyID and ProviderConfigID; add relation
fields to both TableVirtualKeyBudget and TableVirtualKeyProviderConfigBudget
that reference their parent models (e.g., add a VirtualKey field of type
TableVirtualKey on TableVirtualKeyBudget and a ProviderConfig field of type
TableVirtualKeyProviderConfig on TableVirtualKeyProviderConfigBudget) and
annotate them with gorm tags to point to the ID fields and enforce constraint:
e.g. gorm:"foreignKey:VirtualKeyID;references:ID;constraint:OnDelete:CASCADE"
and gorm:"foreignKey:ProviderConfigID;references:ID;constraint:OnDelete:CASCADE"
so deletions of parents cascade and foreign keys are created.
---
Nitpick comments:
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 902-941: The MultiBudgetLines component instance is missing a
data-testid; add a data-testid attribute following the project's pattern
(entity-element-qualifier) to this instance (the MultiBudgetLines with id
`providerBudget-${index}`) so tests can target it reliably; update the JSX for
MultiBudgetLines to include something like
data-testid="virtualkey-provider-budget" or a variant that includes the
index/qualifier you prefer so it follows the `<entity>-<element>-<qualifier>`
convention while leaving existing props and the onChange handling
(handleUpdateProviderConfig) unchanged.
🪄 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: aca4f929-2220-4b2c-8def-491e39d4f405
⛔ Files ignored due to path filters (1)
ui/package-lock.jsonis excluded by!**/package-lock.json
📒 Files selected for processing (10)
framework/configstore/migrations.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/components/sidebar.tsxui/components/ui/multiBudgetLines.tsxui/lib/store/apis/baseApi.tsui/lib/types/governance.ts
✅ Files skipped from review due to trivial changes (2)
- ui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsx
- ui/app/workspace/governance/access-profiles/page.tsx
🚧 Files skipped from review as they are similar to previous changes (3)
- ui/lib/store/apis/baseApi.ts
- ui/components/ui/multiBudgetLines.tsx
- framework/configstore/migrations.go
12 tasks
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
a40a014 to
71634a1
Compare
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
71634a1 to
291045f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ui/lib/types/governance.ts (1)
164-177:⚠️ Potential issue | 🟠 MajorRequest and response types are now asymmetric for virtual-key budgets.
CreateVirtualKeyRequest/UpdateVirtualKeyRequestacceptbudgets, butVirtualKeystill only exposesbudget?: Budget.virtualKeySheet.tsxinitializes fromvirtualKey.budget, so opening an existing multi-budget key cannot round-trip the new array and a save will collapse back to the legacy single-budget shape.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/lib/types/governance.ts` around lines 164 - 177, The VirtualKey model is asymmetric with the requests: CreateVirtualKeyRequest/UpdateVirtualKeyRequest accept budgets (array) but VirtualKey only exposes budget (single), causing virtualKeySheet.tsx to lose multi-budget data on save; update the VirtualKey interface to include budgets?: Budget[] (while keeping budget?: Budget for backward compatibility) and adjust virtualKeySheet.tsx initialization logic to prefer converting budgets to the single budget field when necessary (e.g., use virtualKey.budgets?.[0] || virtualKey.budget) and ensure save handlers emit the budgets array back to the API so existing keys round-trip correctly.
♻️ Duplicate comments (1)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
930-969:⚠️ Potential issue | 🟡 MinorAdd a selector to the new multi-budget editor.
MultiBudgetLinesis a new interactive control, but it still has nodata-testid, unlike the rest of this sheet. Please plumb a test id through the component and set one here using the repo’s three-part pattern. As per coding guidelines:ui/**/*.{tsx,ts}: Add new interactive UI elements withdata-testidattributes following the pattern:data-testid="<entity>-<element>-<qualifier>".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 930 - 969, The MultiBudgetLines instance in virtualKeySheet.tsx is missing a data-testid; add a data-testid prop to this MultiBudgetLines call using the repo’s three-part pattern (e.g. data-testid={`virtualkey-provider-budget-${index}`}) and ensure the MultiBudgetLines component accepts and forwards a data-testid prop to its root interactive element (update the MultiBudgetLines props/type to include data-testid?: string and apply it to the outermost JSX element). Reference symbols: MultiBudgetLines (call site), data-testid prop, and the MultiBudgetLines component definition to wire the prop through.
🤖 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/tables/budget.go`:
- Around line 33-55: The join tables TableVirtualKeyBudget and
TableVirtualKeyProviderConfigBudget leave orphaned TableBudget rows because the
provider-config/virtual-key delete/replace path currently only clears the legacy
BudgetID; update the provider-config and virtual-key delete/replace logic to,
within the same DB transaction, remove the corresponding join-table entries and
delete the referenced TableBudget rows (by BudgetID) for those join entries (use
cascade semantics or explicit DELETEs) so that when budgets are replaced or the
parent is deleted no owned TableBudget remains; specifically modify the
provider-config delete/replace routine to query the related
TableVirtualKeyProviderConfigBudget and TableVirtualKeyBudget rows, collect
their BudgetID values, and delete those TableBudget records along with removing
the join rows in the same transaction.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 90-93: The form schema for provider "budgets" currently declares
max_limit/reset_duration as strings and drops calendar_aligned when building the
API payload; update the budgets zod schema (the budgets array object) to accept
max_limit as a number (or nullable/optional) and calendar_aligned as an optional
boolean, and then update normalizeProviderConfigs() to convert form string
values into the proper types before sending (parse max_limit into a number or
null, keep reset_duration as the expected string, and include calendar_aligned
if present) so multi-budget entries round-trip without losing calendar_aligned
or posting max_limit as a string.
- Around line 704-709: Replace the clickable <Trash2> SVG inside the accordion
trigger with a proper button element that contains the Trash2 icon, wire the
button's onClick to call handleRemoveProvider(index) and call
event.stopPropagation() at the start of that handler (or inline) to prevent the
click from bubbling to the accordion trigger; ensure the button has an
accessible name via aria-label (e.g., aria-label={`Delete provider ${index +
1}`}) and retains the same styling/classes (e.g., "hover:bg-accent/50
cursor-pointer rounded-sm p-2") and the existing data-testid
(`vk-delete-provider-${index}`) so keyboard focus, screen reader announcement,
and tests all work.
---
Outside diff comments:
In `@ui/lib/types/governance.ts`:
- Around line 164-177: The VirtualKey model is asymmetric with the requests:
CreateVirtualKeyRequest/UpdateVirtualKeyRequest accept budgets (array) but
VirtualKey only exposes budget (single), causing virtualKeySheet.tsx to lose
multi-budget data on save; update the VirtualKey interface to include budgets?:
Budget[] (while keeping budget?: Budget for backward compatibility) and adjust
virtualKeySheet.tsx initialization logic to prefer converting budgets to the
single budget field when necessary (e.g., use virtualKey.budgets?.[0] ||
virtualKey.budget) and ensure save handlers emit the budgets array back to the
API so existing keys round-trip correctly.
---
Duplicate comments:
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 930-969: The MultiBudgetLines instance in virtualKeySheet.tsx is
missing a data-testid; add a data-testid prop to this MultiBudgetLines call
using the repo’s three-part pattern (e.g.
data-testid={`virtualkey-provider-budget-${index}`}) and ensure the
MultiBudgetLines component accepts and forwards a data-testid prop to its root
interactive element (update the MultiBudgetLines props/type to include
data-testid?: string and apply it to the outermost JSX element). Reference
symbols: MultiBudgetLines (call site), data-testid prop, and the
MultiBudgetLines component definition to wire the prop through.
🪄 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: b746a15a-e188-4cf5-98aa-be52f83d23ee
⛔ Files ignored due to path filters (1)
ui/package-lock.jsonis excluded by!**/package-lock.json
📒 Files selected for processing (10)
framework/configstore/migrations.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/components/sidebar.tsxui/components/ui/multiBudgetLines.tsxui/lib/store/apis/baseApi.tsui/lib/types/governance.ts
✅ Files skipped from review due to trivial changes (3)
- ui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsx
- ui/app/workspace/governance/access-profiles/page.tsx
- ui/lib/store/apis/baseApi.ts
🚧 Files skipped from review as they are similar to previous changes (3)
- framework/configstore/migrations.go
- ui/components/ui/multiBudgetLines.tsx
- framework/configstore/tables/virtualkey.go
Contributor
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (6)
ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx (1)
31-40:⚠️ Potential issue | 🟠 MajorThe status badge still ignores exhausted provider budgets.
You now render provider-level budgets below, but
isExhaustedonly checks VK-level budgets and rate limits. A key with no top-level budget and a single provider config at its budget limit will still showActivehere even though governance rejects that provider. A shared exhaustion helper would keep this and the table badge in sync.Also applies to: 151-189
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx` around lines 31 - 40, isExhausted currently only checks VK-level budgets and rate limits (variable isExhausted) and therefore ignores provider-level budgets; update isExhausted to also consider provider configs' budgets and rate limits (e.g. iterate virtualKey.providers or virtualKey.provider_configs to detect any provider with current_usage >= max_limit or exhausted rate limits) and consolidate this logic into a shared helper (e.g. computeVirtualKeyExhaustion) so the badge here and the table badge logic (the other block that checks exhaustion) call the same function to remain in sync.transports/bifrost-http/lib/config_test.go (1)
15261-15300:⚠️ Potential issue | 🟠 MajorKeep the owner-FK exclusions, but don't hide
budgetsfrom schema sync.Excluding
virtual_key_id/provider_config_idmakes sense, but excludingtables.TableVirtualKey.budgetsandtables.TableVirtualKeyProviderConfig.budgetsdisablesTestConfigSchemaSyncfor the new multi-budget config surface. That letsconfig.schema.jsondrift without this test failing.♻️ Proposed fix
"tables.TableVirtualKey": { "config_hash": true, "created_at": true, "updated_at": true, - "budgets": true, // GORM relation (budgets have virtual_key_id FK) "rate_limit": true, // GORM relation "team": true, // GORM relation "customer": true, // GORM relation }, "tables.TableVirtualKeyProviderConfig": { - "budgets": true, // GORM relation (budgets have provider_config_id FK) "rate_limit": true, // GORM relation },🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/lib/config_test.go` around lines 15261 - 15300, The test exclusions currently hide tables.TableVirtualKey.budgets and tables.TableVirtualKeyProviderConfig.budgets which prevents TestConfigSchemaSync from exercising the new multi-budget config surface; keep the owner-FK exclusions (virtual_key_id and provider_config_id) but remove the entries that set "tables.TableVirtualKey": {"budgets": true} and "tables.TableVirtualKeyProviderConfig": {"budgets": true} from the exclusions map so the schema sync test validates budgets and prevents config.schema.json drift; ensure only the FK fields (virtual_key_id/provider_config_id) remain excluded.plugins/governance/store.go (2)
2528-2579:⚠️ Potential issue | 🟠 MajorRemoved provider configs still leave orphan budgets and rate limits behind.
This block only cleans child governance objects for provider configs that are still present in
clone.ProviderConfigs. If a provider config is deleted outright, its oldBudgetsandRateLimitstay in the side maps.DumpBudgets()later iterates the wholegs.budgetsmap, so those stale budgets can keep getting flushed after the provider config is gone.Suggested fix
if clone.ProviderConfigs != nil { // Create a map of existing provider configs by ID for fast lookup existingProviderConfigs := make(map[uint]configstoreTables.TableVirtualKeyProviderConfig) if existingVK.ProviderConfigs != nil { for _, existingPC := range existingVK.ProviderConfigs { existingProviderConfigs[existingPC.ID] = existingPC } } + liveProviderConfigIDs := make(map[uint]struct{}, len(clone.ProviderConfigs)) // Process each new/updated provider config for i, pc := range clone.ProviderConfigs { + liveProviderConfigIDs[pc.ID] = struct{}{} if pc.RateLimit != nil { // Preserve existing usage from memory when updating provider config rate limit if existingRateLimitValue, exists := gs.rateLimits.Load(pc.RateLimit.ID); exists && existingRateLimitValue != nil { if existingRateLimit, ok := existingRateLimitValue.(*configstoreTables.TableRateLimit); ok && existingRateLimit != nil { // Preserve current usage and last reset times from existing in-memory rate limit @@ if existingPC, exists := existingProviderConfigs[pc.ID]; exists { for _, oldBudget := range existingPC.Budgets { if !pcNewBudgetIDs[oldBudget.ID] { gs.budgets.Delete(oldBudget.ID) } } } } + for id, existingPC := range existingProviderConfigs { + if _, ok := liveProviderConfigIDs[id]; ok { + continue + } + for _, oldBudget := range existingPC.Budgets { + gs.budgets.Delete(oldBudget.ID) + } + if existingPC.RateLimit != nil { + gs.rateLimits.Delete(existingPC.RateLimit.ID) + } + } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2528 - 2579, The code only cleans child objects for provider configs that remain in clone.ProviderConfigs; to remove orphaned budgets/rate limits when a provider config is deleted, after building existingProviderConfigs (from existingVK.ProviderConfigs) compute which existing IDs are missing from clone.ProviderConfigs and for each missing existingPC delete its RateLimit from gs.rateLimits (if existingPC.RateLimit != nil) and delete each oldBudget.ID from gs.budgets; perform this cleanup before or right after the loop that processes clone.ProviderConfigs so removed provider configs no longer leave stale entries in gs.budgets/gs.rateLimits.
2058-2095:⚠️ Potential issue | 🟠 MajorPopulate VirtualKey and ProviderConfig budget relationships in
loadFromConfigMemory().The relationship hydration loop (lines 2058-2095) wires team/customer/rate-limit references but skips budget relationships. When
collectBudgetsFromHierarchy()later attempts to iteratevk.Budgetsandpc.Budgets, these slices remain empty if not pre-populated by the config file, causing config-memory deployments to skip budget enforcement entirely.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2058 - 2095, The virtual key hydration in loadFromConfigMemory populates Team/Customer/RateLimit but never wires budget references, so vk.Budgets and pc.Budgets remain empty and collectBudgetsFromHierarchy sees no budgets; fix by iterating the global budgets slice when hydrating each virtualKeys entry and each vk.ProviderConfigs entry (use symbols virtualKeys, vk, vk.ProviderConfigs, pc) and for each budget ID match assign the pointer into vk.Budgets and pc.Budgets just like rate limits are assigned; ensure you handle nil ID checks and break/continue as appropriate so budget relationships are populated before collectBudgetsFromHierarchy runs.transports/bifrost-http/handlers/governance.go (2)
118-123:⚠️ Potential issue | 🟡 MinorRemove unused
CalendarAlignedfield fromCreateBudgetRequestor clarify intent.The
CalendarAlignedfield onCreateBudgetRequest(line 122) is never read during budget creation. The code consistently uses the VK-levelCalendarAlignedviabudgetLastReset(vk.CalendarAligned, b.ResetDuration)at lines 506, 587, 770, 924, 998. If per-budget calendar alignment is not intended, remove this field to avoid confusing API consumers. If it is intended for future use, document that assumption.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 118 - 123, The CreateBudgetRequest struct contains an unused CalendarAligned field; either remove the CalendarAligned field from CreateBudgetRequest to avoid exposing a misleading per-budget option, or if per-budget control is intended, wire it into budget creation by passing req.CalendarAligned into budgetLastReset instead of vk.CalendarAligned (update call sites that currently use vk.CalendarAligned such as calls to budgetLastReset in the budget creation/update flows). Ensure the API docs/comments reflect the chosen behavior.
845-848:⚠️ Potential issue | 🟡 MinorRemove redundant budget and rate limit deletion logic—
DeleteVirtualKeyProviderConfigalready handles it.At line 1068,
collectProviderConfigDeleteIDs()collects budget IDs to delete manually at lines 1145–1148. However,DeleteVirtualKeyProviderConfig()(called at line 1073) already deletes all associated budgets and rate limits internally (seeframework/configstore/rdb.go:2336). The manual deletion code is dead—it collects from an unpopulatedconfig.Budgetsslice (no Preload) and attempts to delete records that were already deleted.Remove the
collectProviderConfigDeleteIDscall and the manual deletion loops; the provider config deletion handles all cleanup:for id := range existingConfigsMap { if !requestConfigsMap[id] { - providerBudgetIDsToDelete, providerRateLimitIDsToDelete = collectProviderConfigDeleteIDs( - existingConfigsMap[id], - providerBudgetIDsToDelete, - providerRateLimitIDsToDelete, - ) if err := h.configStore.DeleteVirtualKeyProviderConfig(ctx, id, tx); err != nil { return err } } }Then remove the unused
providerBudgetIDsToDeleteandproviderRateLimitIDsToDeletevariables and the manual deletion block at lines 1145–1148.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 845 - 848, Remove the redundant manual deletion logic: delete the call to collectProviderConfigDeleteIDs and drop the providerBudgetIDsToDelete and providerRateLimitIDsToDelete variables and their subsequent manual tx.Delete loops, because DeleteVirtualKeyProviderConfig already removes associated budgets and rate limits; also remove any reliance on config.Budgets (which isn’t Preloaded) in this delete path and keep only the call to DeleteVirtualKeyProviderConfig to perform the cleanup.
♻️ Duplicate comments (4)
framework/configstore/tables/virtualkey.go (1)
38-38:⚠️ Potential issue | 🟡 MinorFix the
Keyscomment to match the new deny-by-default behavior.Line 38 still says empty
Keysmeans “all keys allowed”, but Line 32 and the JSON parsing path now treat emptyKeyswithAllowAllKeys=falseas deny-all. Keeping the old comment here will keep sending future changes in the wrong direction.✏️ Suggested doc fix
- Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty means all keys allowed for this provider + Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Interpreted with AllowAllKeys: true => all keys allowed; false + empty => no keys allowed (deny-by-default)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` at line 38, Update the comment on the Keys field in the VirtualKey struct to reflect the new deny-by-default semantics: note that an empty Keys slice does NOT mean all keys are allowed—it's treated as deny-all unless AllowAllKeys (or equivalent flag) is set to true; change the existing comment text on Keys and mention AllowAllKeys to avoid future confusion with the JSON parsing path and logic in virtual key handling.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (3)
897-929:⚠️ Potential issue | 🟡 MinorExpose
data-testids for the new budget editors.Both
MultiBudgetLinesusages add new interactive UI without stable selectors. Please forward adata-testidprop throughMultiBudgetLinesand apply the repo’s<entity>-<element>-<qualifier>pattern here.As per coding guidelines,
ui/**/*.{tsx,ts}: Add new interactive UI elements withdata-testidattributes following the pattern:data-testid="<entity>-<element>-<qualifier>".Also applies to: 1214-1220
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 897 - 929, Add a stable data-testid to the new budget editors by passing a data-testid prop into both MultiBudgetLines usages (the one with id={`providerBudget-${index}`} and the other at lines ~1214-1220) using the repo pattern entity-element-qualifier (e.g. provider-budget-editor or provider-budget-lines-<qualifier>), and ensure the MultiBudgetLines component accepts and forwards that prop to the root interactive element; update MultiBudgetLines props and its root JSX to forward data-testid so the interactive inputs rendered by MultiBudgetLines receive the stable selector while leaving existing behavior in handleUpdateProviderConfig unchanged.
369-403:⚠️ Potential issue | 🔴 Critical
normalizeProviderConfigs()is broken and still emits invalid budget payloads.Line 372 references
configs, but that identifier is not defined in this scope, so the sheet will not compile. Even after fixing that, the helper still spreadsconfig.budgetsstraight into the request, which keepsmax_limitas the string produced byMultiBudgetLinesinstead of the numericCreateBudgetRequest.max_limitexpected by the API.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 369 - 403, The function normalizeProviderConfigs currently references an undefined variable configs and passes budget max_limit strings through to the API; fix by changing the signature to accept both the incoming configs and optional existingConfigs (e.g. normalizeProviderConfigs = (configs?: VirtualKey["provider_configs"], existingConfigs?: VirtualKey["provider_configs"]) => ...) or rename the parameter so the loop iterates over the actual configs argument, then when mapping each config convert its budgets array so each CreateBudgetRequest.max_limit is a numeric value (use normalizeIntegerField or parseFloat with Number.isNaN guard and coerce invalid to null) instead of leaving the string from MultiBudgetLines; keep the existing rate_limit normalization logic but reference the correct params (config.id/provider) when looking up existingConfig.
671-676:⚠️ Potential issue | 🟠 MajorMake the delete affordances accessible.
The provider delete control is still a clickable SVG inside the accordion trigger, so it is not keyboard-focusable and its click bubbles into the accordion. The MCP delete control is a real button, but because it is icon-only it still needs an accessible name. Please make both affordances labeled buttons and stop propagation before deleting the provider.
Also applies to: 1180-1188
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 671 - 676, The provider delete SVG needs to be converted into a proper labeled button and the click must stop propagation before calling handleRemoveProvider; replace the clickable <Trash2 ... /> with a <button type="button"> that contains the Trash2 icon, add onClick={(e) => { e.stopPropagation(); handleRemoveProvider(index); }}, include an accessible name via aria-label (e.g. aria-label={`Remove provider ${index + 1}`} or similar) and keep the data-testid `vk-delete-provider-${index}` and styling classes. Do the same for the MCP delete affordance: ensure its button uses type="button", has an explicit aria-label, and calls e.stopPropagation() before invoking its remove handler so the accordion trigger does not receive the event; reference the existing handleRemoveProvider and Trash2 identifiers when making these changes.
🧹 Nitpick comments (3)
plugins/governance/test_utils.go (1)
192-199: This default provider-config helper now models “no keys allowed.”
buildProviderConfigleavesAllowAllKeysat its zero value. WithKeysempty, that is a deny-all key config now, not a permissive default. The new VK helpers use this builder specifically to keep governance checks open, so they end up constructing fixtures that still publish an empty include-only key list downstream.🛠️ One way to keep the helper permissive
func buildProviderConfig(provider string, allowedModels []string) configstoreTables.TableVirtualKeyProviderConfig { return configstoreTables.TableVirtualKeyProviderConfig{ Provider: provider, AllowedModels: allowedModels, Weight: bifrost.Ptr(1.0), + AllowAllKeys: true, RateLimit: nil, Keys: []configstoreTables.TableKey{}, } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 192 - 199, buildProviderConfig currently returns a TableVirtualKeyProviderConfig with Keys empty and AllowAllKeys left false (zero), which creates a deny-all key config; change the fixture to be permissive by setting AllowAllKeys (on TableVirtualKeyProviderConfig) to true (e.g., using bifrost.Ptr(true) to match existing pointer style) so the VK helpers built by buildProviderConfig do not inadvertently block all keys while keeping Keys as an empty slice.plugins/governance/store_test.go (1)
511-551: This test doesn't hit the calendar-aligned reset behavior yet.It only proves the flag survives hydration. If the boundary-snapping logic regresses, this still passes; the valuable assertion here is to move
LastResetacross a day/month boundary, runResetExpiredBudgets*, and verify the reset happens on the aligned cutoff.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 511 - 551, The test TestGovernanceStore_MultiBudget_CalendarAligned only checks that vk.CalendarAligned survives hydration; update it to exercise the calendar-aligned reset by mutating the budgets' LastReset to a time just before a day/month boundary (for the daily budget set LastReset to >24h ago or previous calendar day, for the monthly budget set LastReset to a prior month), call the store's reset routine (e.g., ResetExpiredBudgets or the appropriate ResetExpiredBudgets* method on the governance store created via NewLocalGovernanceStore), then fetch the VK via GetVirtualKey and assert the budgets' CurrentUsage and/or LastReset reflect the calendar-aligned snap (i.e., usages reset and LastReset advanced to the aligned cutoff) rather than no-op; keep existing checks with CheckBudget and reuse vk, dailyBudget, monthlyBudget identifiers to locate code.transports/bifrost-http/lib/config_test.go (1)
14068-14187: Add a replacement hash case for multi-budget provider configs.This rename narrows the coverage to
RateLimitIDonly. In this stack, provider-config budgets are the new behavior, so this file still needs a case proving that changingProviderConfigs[i].Budgetschanges the VK hash; otherwise a budget-only reconciliation regression can slip through unnoticed.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/lib/config_test.go` around lines 14068 - 14187, Add a test case inside TestGenerateVirtualKeyHash_ProviderConfigRateLimit that verifies GenerateVirtualKeyHash considers ProviderConfigs[i].Budgets: create two tables.TableVirtualKey instances (same ID/Name/etc.) whose ProviderConfigs differ only by the Budgets field (e.g. one with nil/empty Budgets and the other with a non-empty budgets slice or different budget values) and assert hashes are different; also include a same-budgets subcase asserting hashes are equal. Reference GenerateVirtualKeyHash and tables.TableVirtualKeyProviderConfig.Budgets to locate where to modify/add the cases.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.agents/skills/expect/SKILL.md:
- Around line 38-39: Update the SKILL.md guidance to include an explicit
non-production safety guardrail: require testers to use only local/staging test
environments and test accounts (no real payment or side-effect systems) and make
it mandatory to set EXPECT_BASE_URL or pass --base-url to a non-production host
(never production domains); also add a short checkbox/confirmation step (e.g.,
“I confirm this run uses test data/accounts”) and a sentence instructing to run
expect-cli with -y only against non-production URLs. Reference EXPECT_BASE_URL,
--base-url, -y, and expect-cli in the new wording so the reader knows where to
configure the safety check.
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2611-2625: The backfill checks should fail hard and validate exact
linkage: replace the soft log_warn branches for vk_budget_count and
pc_budget_count with strict checks that log an error and exit non-zero on
failure; specifically, for the VK check use run_postgres_sql to assert
governance_budgets.virtual_key_id = 'vk-migration-test-1' (instead of just IS
NOT NULL) and for the PC check assert governance_budgets.provider_config_id =
'<expected_provider_config_id>' (replace with the real fixture id) rather than
IS NOT NULL, and on mismatch call log_error and exit 1; update references to
vk_budget_count, pc_budget_count, run_postgres_sql, log_error/log_info
accordingly.
- Around line 2392-2395: The global ignore_columns variable currently includes
virtual_key_id and provider_config_id which incorrectly suppresses diffs for all
tables; remove virtual_key_id and provider_config_id from the global
ignore_columns string and instead add a table-specific case where, when the
inspected table name equals "governance_budgets" (or within the code path that
builds per-table ignore lists), you append "virtual_key_id provider_config_id"
to that table's ignore list so only governance_budgets ignores those two columns
while other tables continue to be validated.
- Around line 2610-2640: The scalar queries in the migration tests are using
run_postgres_sql which returns psql formatted output (headers/footers), so tr -d
'[:space:]' doesn't yield plain "1"; add a new helper run_postgres_scalar() that
calls psql with -t -A to return bare scalar values and use that for the four
checks currently assigning vk_budget_count, pc_budget_count, has_vk_col, and
has_pc_col (replace those run_postgres_sql calls with run_postgres_scalar).
Ensure run_postgres_scalar preserves exit codes and stderr handling like
run_postgres_sql and document its use in the migration checks.
In `@plugins/governance/test_utils.go`:
- Around line 90-96: The builders (buildVirtualKeyWithBudget and
buildProviderConfigWithBudgets / buildProviderConfig) currently attach budgets
but do not set the owner foreign-key fields, so update these helpers to stamp
the appropriate owner IDs: set each TableBudget.VirtualKeyID when building
vk.Budgets in buildVirtualKeyWithBudget (use vk.ID) and set each
TableBudget.ProviderConfigID when building pc.Budgets in
buildProviderConfigWithBudgets (use pc.ID); implement this by either adding an
explicit owner-id parameter to the builder or by assigning the budget.OwnerID
fields immediately after the caller assigns the parent ID (e.g., after pc.ID or
vk.ID is set) so budgets and rate-limits have the correct 1:1 linkage to their
parent entities.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 222-225: The calendar-alignment toggle currently reads from
virtualKey.budgets (via budgetCalendarAligned) but the true field is
virtualKey.calendar_aligned and the create/update payloads never set
calendar_aligned, causing toggled state to be lost; update initializers to
prefer virtualKey.calendar_aligned (falling back to budgets.some(...) only if
undefined), wire the form state field budgetCalendarAligned to that value, and
include calendar_aligned in both the create and update payloads so user changes
persist (update the code paths that build payloads in the create/update handlers
and any other places setting budgetCalendarAligned such as the occurrences
around budgetCalendarAligned and virtualKey.budgets).
- Around line 365-367: clearVirtualKeyRateLimits currently only sets rate_limit
to undefined but does not update the individual controlled inputs, so update the
function (clearVirtualKeyRateLimits) to also call form.setValue for each
controlled field—tokenMaxLimit, tokenResetDuration, requestMaxLimit, and
requestResetDuration—setting them to undefined (or their empty/default values)
with { shouldDirty: true } so the visible inputs are cleared; keep the existing
rate_limit setValue as well to preserve the underlying model reset.
In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx`:
- Around line 313-315: The `(calendar)` suffix is using the virtual key-level
flag vk.calendar_aligned but must reflect whether the specific budget interval
is calendar-aligned; change the conditional to use the budget-level property
(e.g., b.calendar_aligned) when rendering the reset duration so only
calendar-aligned budgets show the suffix (update the span that renders Resets
{formatResetDuration(b.reset_duration)} to check b.calendar_aligned instead of
vk.calendar_aligned).
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 2528-2579: The code only cleans child objects for provider configs
that remain in clone.ProviderConfigs; to remove orphaned budgets/rate limits
when a provider config is deleted, after building existingProviderConfigs (from
existingVK.ProviderConfigs) compute which existing IDs are missing from
clone.ProviderConfigs and for each missing existingPC delete its RateLimit from
gs.rateLimits (if existingPC.RateLimit != nil) and delete each oldBudget.ID from
gs.budgets; perform this cleanup before or right after the loop that processes
clone.ProviderConfigs so removed provider configs no longer leave stale entries
in gs.budgets/gs.rateLimits.
- Around line 2058-2095: The virtual key hydration in loadFromConfigMemory
populates Team/Customer/RateLimit but never wires budget references, so
vk.Budgets and pc.Budgets remain empty and collectBudgetsFromHierarchy sees no
budgets; fix by iterating the global budgets slice when hydrating each
virtualKeys entry and each vk.ProviderConfigs entry (use symbols virtualKeys,
vk, vk.ProviderConfigs, pc) and for each budget ID match assign the pointer into
vk.Budgets and pc.Budgets just like rate limits are assigned; ensure you handle
nil ID checks and break/continue as appropriate so budget relationships are
populated before collectBudgetsFromHierarchy runs.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 118-123: The CreateBudgetRequest struct contains an unused
CalendarAligned field; either remove the CalendarAligned field from
CreateBudgetRequest to avoid exposing a misleading per-budget option, or if
per-budget control is intended, wire it into budget creation by passing
req.CalendarAligned into budgetLastReset instead of vk.CalendarAligned (update
call sites that currently use vk.CalendarAligned such as calls to
budgetLastReset in the budget creation/update flows). Ensure the API
docs/comments reflect the chosen behavior.
- Around line 845-848: Remove the redundant manual deletion logic: delete the
call to collectProviderConfigDeleteIDs and drop the providerBudgetIDsToDelete
and providerRateLimitIDsToDelete variables and their subsequent manual tx.Delete
loops, because DeleteVirtualKeyProviderConfig already removes associated budgets
and rate limits; also remove any reliance on config.Budgets (which isn’t
Preloaded) in this delete path and keep only the call to
DeleteVirtualKeyProviderConfig to perform the cleanup.
In `@transports/bifrost-http/lib/config_test.go`:
- Around line 15261-15300: The test exclusions currently hide
tables.TableVirtualKey.budgets and tables.TableVirtualKeyProviderConfig.budgets
which prevents TestConfigSchemaSync from exercising the new multi-budget config
surface; keep the owner-FK exclusions (virtual_key_id and provider_config_id)
but remove the entries that set "tables.TableVirtualKey": {"budgets": true} and
"tables.TableVirtualKeyProviderConfig": {"budgets": true} from the exclusions
map so the schema sync test validates budgets and prevents config.schema.json
drift; ensure only the FK fields (virtual_key_id/provider_config_id) remain
excluded.
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 31-40: isExhausted currently only checks VK-level budgets and rate
limits (variable isExhausted) and therefore ignores provider-level budgets;
update isExhausted to also consider provider configs' budgets and rate limits
(e.g. iterate virtualKey.providers or virtualKey.provider_configs to detect any
provider with current_usage >= max_limit or exhausted rate limits) and
consolidate this logic into a shared helper (e.g. computeVirtualKeyExhaustion)
so the badge here and the table badge logic (the other block that checks
exhaustion) call the same function to remain in sync.
---
Duplicate comments:
In `@framework/configstore/tables/virtualkey.go`:
- Line 38: Update the comment on the Keys field in the VirtualKey struct to
reflect the new deny-by-default semantics: note that an empty Keys slice does
NOT mean all keys are allowed—it's treated as deny-all unless AllowAllKeys (or
equivalent flag) is set to true; change the existing comment text on Keys and
mention AllowAllKeys to avoid future confusion with the JSON parsing path and
logic in virtual key handling.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 897-929: Add a stable data-testid to the new budget editors by
passing a data-testid prop into both MultiBudgetLines usages (the one with
id={`providerBudget-${index}`} and the other at lines ~1214-1220) using the repo
pattern entity-element-qualifier (e.g. provider-budget-editor or
provider-budget-lines-<qualifier>), and ensure the MultiBudgetLines component
accepts and forwards that prop to the root interactive element; update
MultiBudgetLines props and its root JSX to forward data-testid so the
interactive inputs rendered by MultiBudgetLines receive the stable selector
while leaving existing behavior in handleUpdateProviderConfig unchanged.
- Around line 369-403: The function normalizeProviderConfigs currently
references an undefined variable configs and passes budget max_limit strings
through to the API; fix by changing the signature to accept both the incoming
configs and optional existingConfigs (e.g. normalizeProviderConfigs = (configs?:
VirtualKey["provider_configs"], existingConfigs?:
VirtualKey["provider_configs"]) => ...) or rename the parameter so the loop
iterates over the actual configs argument, then when mapping each config convert
its budgets array so each CreateBudgetRequest.max_limit is a numeric value (use
normalizeIntegerField or parseFloat with Number.isNaN guard and coerce invalid
to null) instead of leaving the string from MultiBudgetLines; keep the existing
rate_limit normalization logic but reference the correct params
(config.id/provider) when looking up existingConfig.
- Around line 671-676: The provider delete SVG needs to be converted into a
proper labeled button and the click must stop propagation before calling
handleRemoveProvider; replace the clickable <Trash2 ... /> with a <button
type="button"> that contains the Trash2 icon, add onClick={(e) => {
e.stopPropagation(); handleRemoveProvider(index); }}, include an accessible name
via aria-label (e.g. aria-label={`Remove provider ${index + 1}`} or similar) and
keep the data-testid `vk-delete-provider-${index}` and styling classes. Do the
same for the MCP delete affordance: ensure its button uses type="button", has an
explicit aria-label, and calls e.stopPropagation() before invoking its remove
handler so the accordion trigger does not receive the event; reference the
existing handleRemoveProvider and Trash2 identifiers when making these changes.
---
Nitpick comments:
In `@plugins/governance/store_test.go`:
- Around line 511-551: The test TestGovernanceStore_MultiBudget_CalendarAligned
only checks that vk.CalendarAligned survives hydration; update it to exercise
the calendar-aligned reset by mutating the budgets' LastReset to a time just
before a day/month boundary (for the daily budget set LastReset to >24h ago or
previous calendar day, for the monthly budget set LastReset to a prior month),
call the store's reset routine (e.g., ResetExpiredBudgets or the appropriate
ResetExpiredBudgets* method on the governance store created via
NewLocalGovernanceStore), then fetch the VK via GetVirtualKey and assert the
budgets' CurrentUsage and/or LastReset reflect the calendar-aligned snap (i.e.,
usages reset and LastReset advanced to the aligned cutoff) rather than no-op;
keep existing checks with CheckBudget and reuse vk, dailyBudget, monthlyBudget
identifiers to locate code.
In `@plugins/governance/test_utils.go`:
- Around line 192-199: buildProviderConfig currently returns a
TableVirtualKeyProviderConfig with Keys empty and AllowAllKeys left false
(zero), which creates a deny-all key config; change the fixture to be permissive
by setting AllowAllKeys (on TableVirtualKeyProviderConfig) to true (e.g., using
bifrost.Ptr(true) to match existing pointer style) so the VK helpers built by
buildProviderConfig do not inadvertently block all keys while keeping Keys as an
empty slice.
In `@transports/bifrost-http/lib/config_test.go`:
- Around line 14068-14187: Add a test case inside
TestGenerateVirtualKeyHash_ProviderConfigRateLimit that verifies
GenerateVirtualKeyHash considers ProviderConfigs[i].Budgets: create two
tables.TableVirtualKey instances (same ID/Name/etc.) whose ProviderConfigs
differ only by the Budgets field (e.g. one with nil/empty Budgets and the other
with a non-empty budgets slice or different budget values) and assert hashes are
different; also include a same-budgets subcase asserting hashes are equal.
Reference GenerateVirtualKeyHash and
tables.TableVirtualKeyProviderConfig.Budgets to locate where to modify/add the
cases.
🪄 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: 47f04c6e-f6e8-4f7e-b2b0-689be2bead6a
📒 Files selected for processing (25)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (1)
- .claude/skills/expect
🚧 Files skipped from review as they are similar to previous changes (2)
- framework/configstore/tables/budget.go
- framework/configstore/migrations.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
5359439 to
2ab508f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
framework/configstore/rdb.go (1)
2020-2023:⚠️ Potential issue | 🟠 MajorPersist
calendar_alignedwhen updating an existing virtual key.This whitelist still selects
budget_idand never selectscalendar_aligned, so edits to the new VK-level flag won't be saved throughUpdateVirtualKey.🐛 Proposed fix
- Select("name", "description", "value", "is_active", "team_id", "customer_id", "budget_id", "rate_limit_id", "config_hash", "updated_at", "encryption_status", "value_hash"). + Select("name", "description", "value", "is_active", "team_id", "customer_id", "calendar_aligned", "rate_limit_id", "config_hash", "updated_at", "encryption_status", "value_hash").🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2020 - 2023, The update statement in UpdateVirtualKey doesn't persist the new VK-level flag because "calendar_aligned" is missing from the Select projection; update the Select call (the chain starting with txDB.WithContext(ctx).Select(...).Updates(virtualKey)) to include "calendar_aligned" alongside the other fields so changes to virtualKey.CalendarAligned are saved when Updates(virtualKey).Error executes.ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx (1)
31-40:⚠️ Potential issue | 🟠 MajorInclude provider-config budgets in the exhausted-state badge.
isExhaustedonly checksvirtualKey.budgetsand the VK-level rate limit. A key that relies only on provider-config budgets now still renders asActivehere, even when every provider budget is exhausted.Also applies to: 151-189
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx` around lines 31 - 40, The exhausted-state logic in isExhausted currently only checks virtualKey.budgets and VK-level rate limits, so provider-config budgets are missing; update the condition to also consider provider-config budgets (e.g., add || virtualKey.provider_budgets?.some(b => b.current_usage >= b.max_limit)) alongside the existing virtualKey.budgets check and similarly extend the same logic used later (lines ~151-189) where exhausted-state is computed, referencing the same symbols: isExhausted, virtualKey.budgets, virtualKey.provider_budgets, and virtualKey.rate_limit.plugins/governance/store.go (1)
2529-2579:⚠️ Potential issue | 🟠 MajorRemove orphaned children when a provider config disappears.
Line 2529 builds
existingProviderConfigs, but the loop only reconciles configs that still exist inclone.ProviderConfigs. If a provider config is removed from a VK, its old budgets and rate limit stay ings.budgets/gs.rateLimits, soGetGovernanceData,DumpBudgets, andDumpRateLimitskeep operating on objects that no longer belong to any VK.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2529 - 2579, Orphaned budgets and rate limits remain when a provider config is removed: after building existingProviderConfigs and reconciling clone.ProviderConfigs, detect provider config IDs that existed in existingVK.ProviderConfigs but are not present in clone.ProviderConfigs and remove their child entries from gs.budgets and gs.rateLimits. Concretely, build a set of new provider IDs from clone.ProviderConfigs (or remove matched entries from existingProviderConfigs as you process them) and then for each remaining existingPC in existingProviderConfigs that is not in the new set, call gs.budgets.Delete(oldBudget.ID) for each oldBudget in existingPC.Budgets and gs.rateLimits.Delete(existingPC.RateLimit.ID) if existingPC.RateLimit != nil, and ensure clone.ProviderConfigs does not reference them.transports/bifrost-http/handlers/governance.go (1)
96-115:⚠️ Potential issue | 🔴 CriticalPreserve budget counters across VK updates.
UpdateVirtualKeyRequestcarries budgets as create-only objects, so Lines 758-780 and 988-1007 delete/recreate rows and assign new IDs on every PUT. Because this endpoint receives full state, ordinary edits resend unchanged budgets and resetCurrentUsage/LastResetto zero. Reconcile existing budgets in place, or add stable budget IDs to the request, so unchanged budgets keep their usage history.Based on learnings,
updateVirtualKeyexpects a full payload from the frontend.Also applies to: 757-785, 987-1008
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 96 - 115, The updateVirtualKey handler currently deletes and recreates budgets (losing CurrentUsage/LastReset) because CreateBudgetRequest is treated create-only; change updateVirtualKey to reconcile budgets in-place: when handling UpdateVirtualKeyRequest (and the provider/ VK budget slices), fetch existing budgets for that virtual key, match incoming budgets by stable identifier (prefer an ID field on CreateBudgetRequest if present, otherwise a unique key like budget name/type), and for matches update mutable fields (limit, period) while preserving CurrentUsage and LastReset; create new DB rows only for truly new budgets and delete only budgets absent from the incoming full payload. Update CreateBudgetRequest/handler logic accordingly so budget IDs are accepted/used (or implement matching logic) and remove the delete-all/insert-all flow in updateVirtualKey to keep usage history.
🧹 Nitpick comments (4)
ui/lib/types/governance.ts (1)
137-176: Use a VK-specific budget payload type here.These request interfaces reuse
CreateBudgetRequest, which still carriescalendar_aligned. After moving alignment to the virtual-key level, that lets VK/provider-config payloads type-check with both top-levelcalendar_alignedand nestedbudgets[i].calendar_aligned, which is an easy way for contradictory form state to leak into requests.♻️ Proposed refactor
+type VirtualKeyBudgetRequest = Omit<CreateBudgetRequest, "calendar_aligned">; + export interface VirtualKeyProviderConfigRequest { provider: string; weight?: number | null; allowed_models?: string[]; - budgets?: CreateBudgetRequest[]; + budgets?: VirtualKeyBudgetRequest[]; rate_limit?: CreateRateLimitRequest; key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config } export interface VirtualKeyProviderConfigUpdateRequest { id?: number; provider: string; weight?: number | null; allowed_models?: string[]; - budgets?: CreateBudgetRequest[]; + budgets?: VirtualKeyBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; key_ids?: string[]; // List of DBKey UUIDs to associate with this provider config } export interface CreateVirtualKeyRequest { name: string; @@ - budgets?: CreateBudgetRequest[]; + budgets?: VirtualKeyBudgetRequest[]; rate_limit?: CreateRateLimitRequest; is_active?: boolean; calendar_aligned?: boolean; } export interface UpdateVirtualKeyRequest { name?: string; @@ - budgets?: CreateBudgetRequest[]; + budgets?: VirtualKeyBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; is_active?: boolean; calendar_aligned?: boolean; }Also applies to: 205-215
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/lib/types/governance.ts` around lines 137 - 176, The nested budget payloads for virtual-key requests should use a VK-specific budget type without calendar_aligned to prevent conflicting state; create a new type (e.g., CreateVirtualKeyBudgetRequest / UpdateVirtualKeyBudgetRequest) that omits calendar_aligned from CreateBudgetRequest/UpdateBudgetRequest and replace all uses of CreateBudgetRequest (and UpdateBudgetRequest if present) inside VirtualKeyProviderConfigRequest/VirtualKeyProviderConfigUpdateRequest, CreateVirtualKeyRequest, UpdateVirtualKeyRequest, and VirtualKeyMCPConfigRequest payloads (also update the occurrences mentioned around lines 205-215) so top-level calendar_aligned remains the single source of truth.plugins/governance/store_test.go (2)
261-264: Guarderr.Error()behindrequire.Error(orif assert.Error(...)).
assert.Errordoesn’t stop execution, so these follow-uperr.Error()calls will panic if the function unexpectedly succeeds. That makes the test fail with a nil dereference instead of the actual regression.🧪 Safer assertion pattern
- assert.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") - assert.Contains(t, err.Error(), "budget exceeded") + if assert.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") { + assert.Contains(t, err.Error(), "budget exceeded") + }Also applies to: 287-290, 337-340, 363-366, 394-396
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 261 - 264, The test currently calls store.CheckBudget(...) and uses assert.Error(t, err) then immediately calls err.Error(), which can panic if err is nil; change these to use require.Error(t, err) (or wrap the follow-up checks in if assert.Error(t, err) { ... }) so execution stops when no error is returned; update the CheckBudget call sites in this file (the occurrences around the current assert blocks and the other instances noted) to either use require.Error for the initial assertion or guard the subsequent assert.Contains(err.Error(), ...) inside an if assert.Error(...) block so err is never dereferenced when nil.
511-550: This doesn’t actually exercise calendar-aligned resets.The test only proves that
CalendarAlignedround-trips through the store and that an under-limit request passes. If the new VK-level behavior is supposed to reset1d/1Mbudgets on calendar boundaries, add a case that straddles midnight or month-end and assertCheckBudgetresets usage before evaluating.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 511 - 550, The test currently only verifies VK.CalendarAligned persists and that an under-limit request passes; add a subcase that exercises calendar-aligned resets by simulating a boundary crossing: create the same daily and monthly budgets but set their LastReset timestamps to just before the previous calendar boundary (e.g., yesterday for the 1d budget and last month for the 1M budget) or otherwise ensure they are older than the current calendar period, then call store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil) and assert that the budgets' CurrentUsage have been reset (or reduced according to reset logic) before evaluation; use buildVirtualKeyWithMultiBudgets, NewLocalGovernanceStore, CheckBudget and store.GetVirtualKey/GetBudget to locate and verify the post-check state.plugins/governance/store.go (1)
1475-1499: Avoid a full VK scan for every budget reset.Lines 1478-1498 do a full
gs.virtualKeysscan for every budget just to recoverCalendarAligned. With multi-budget support this makes the background reset pass O(budgets × virtualKeys × providerConfigs). A small lookup keyed byVirtualKeyID/ProviderConfigIDwould keep this path linear.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 1475 - 1499, The reset path currently does a full gs.virtualKeys.Range scan to find CalendarAligned for each budget (checking budget.VirtualKeyID or budget.ProviderConfigID against configstoreTables.TableVirtualKey and its ProviderConfigs), which is O(budgets×virtualKeys); add O(1) lookups instead: maintain a lookup map in the governance store (e.g., vkByID map[int64]*configstoreTables.TableVirtualKey and a providerConfigToVK map[ProviderConfigIDType]*configstoreTables.TableVirtualKey or vkID) that is populated/updated whenever virtual keys change, protect it with the same concurrency mechanism used for gs.virtualKeys (sync.RWMutex or sync.Map), and replace the gs.virtualKeys.Range calls in the budget reset logic with direct map lookups to read TableVirtualKey.CalendarAligned for the found VK (using budget.VirtualKeyID or resolving budget.ProviderConfigID via providerConfigToVK).
🤖 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 1004-1007: The ALTER TABLE adding
governance_virtual_key_provider_configs.budget_id must not be silently ignored
because migrationAddMultiBudgetTables later reads that column; modify the
tx.Exec error handling so genuine failures are returned (or cause the
transaction to roll back) instead of swallowed: replace the current empty catch
with logic that only ignores errors known to mean “IF NOT EXISTS not supported”
for the current dialect (e.g., detect SQLite via the DB/Dialector or match the
specific error), otherwise propagate or wrap and return the error from this
migration step (reference tx.Exec call that adds budget_id and the
migrationAddMultiBudgetTables reader).
- Around line 5329-5338: The migration currently adds columns VirtualKeyID and
ProviderConfigID to tables.TableBudget but misses creating the corresponding
indexes and FK constraints; after the backfill and before returning, check
mg.HasIndex for "idx_virtual_key_id" and "idx_provider_config_id" and call
mg.CreateIndex(&tables.TableBudget{}, "VirtualKeyID") and
mg.CreateIndex(&tables.TableBudget{}, "ProviderConfigID") if missing, and
similarly check mg.HasConstraint for the FK names (e.g. "fk_virtual_key_budgets"
and "fk_provider_config_budgets") and call
mg.CreateConstraint(&tables.TableBudget{}, "VirtualKey") and
mg.CreateConstraint(&tables.TableBudget{}, "ProviderConfig") to ensure the
foreign keys (with OnDelete:CASCADE as defined on the TableBudget struct) are
created on upgraded installs, returning formatted errors on failure.
In `@framework/configstore/rdb.go`:
- Around line 2099-2102: The delete path in the function handling single
provider-config removal (DeleteVirtualKey) must mirror the bulk-delete cleanup:
begin a DB transaction (use tx :=
db.WithContext(ctx).Begin()/tx.Commit()/tx.Rollback()), first delete the
join-table rows from governance_virtual_key_provider_config_keys for the target
pc.ID using tx.WithContext(ctx).Exec(...), then perform the subsequent deletes
(budgets, config entries, rate-limits and the provider-config row) using the
same tx so all-or-nothing semantics apply; also apply the identical
transactional join-table-first cleanup to the other delete block referenced
around the 2326-2342 region to ensure budgets and rate-limits (1:1 owned
resources) are removed inside the transaction.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 442-459: The request validation for virtual key budgets is
returning generic errors (from validateBudget) which surface as 500 because
createVirtualKey only maps badRequestError to HTTP 400; ensure all payload
validation failures (e.g., max_limit == 0, negative or duplicate reset_duration)
are converted into badRequestError before they propagate: either update
validateBudget to return badRequestError for any client-side validation issue,
or wrap/translate its returned errors to badRequestError in createVirtualKey
(and the top-level budget pre-check paths that call validateBudget) so that
SendError yields a 400 response; apply the same fix to the other
budget-validation sites mentioned (lines ~499-516 and ~575-597).
In `@transports/config.schema.json`:
- Around line 455-459: Add the new multi-budget field that replaces the removed
single budget reference by adding a "budget_ids" array property to both
governance.virtual_keys[*] and $defs.virtual_key_provider_config in the schema:
define "budget_ids" as an array whose items match the existing budget ID type
used elsewhere (string/ID), include a clear description like "List of budget IDs
this virtual key/provider applies to", and set a sensible default (empty array)
so config.json can declare VK/provider budgets; ensure the property shape
exactly matches what the transport/parser expects (same item schema and
constraints) and mirror the change in the other schema block mentioned in the
review.
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 151-186: Replace the raw division (b.current_usage / b.max_limit)
in the Provider Budgets block with a call to
calculateUsagePercentage(b.current_usage, b.max_limit), use that returned
percentage for the Badge label (Math.round(percent) + "%") and for the badge
variant (e.g., variant={percent >= 100 ? "destructive" : "default"}), and apply
the same change to the other occurrence around lines 350-381 so zero max_limit
no longer yields NaN/Infinity; note the relevant symbols:
calculateUsagePercentage, b.current_usage, b.max_limit, Badge, and
virtualKeyDetailsSheet.tsx.
---
Outside diff comments:
In `@framework/configstore/rdb.go`:
- Around line 2020-2023: The update statement in UpdateVirtualKey doesn't
persist the new VK-level flag because "calendar_aligned" is missing from the
Select projection; update the Select call (the chain starting with
txDB.WithContext(ctx).Select(...).Updates(virtualKey)) to include
"calendar_aligned" alongside the other fields so changes to
virtualKey.CalendarAligned are saved when Updates(virtualKey).Error executes.
In `@plugins/governance/store.go`:
- Around line 2529-2579: Orphaned budgets and rate limits remain when a provider
config is removed: after building existingProviderConfigs and reconciling
clone.ProviderConfigs, detect provider config IDs that existed in
existingVK.ProviderConfigs but are not present in clone.ProviderConfigs and
remove their child entries from gs.budgets and gs.rateLimits. Concretely, build
a set of new provider IDs from clone.ProviderConfigs (or remove matched entries
from existingProviderConfigs as you process them) and then for each remaining
existingPC in existingProviderConfigs that is not in the new set, call
gs.budgets.Delete(oldBudget.ID) for each oldBudget in existingPC.Budgets and
gs.rateLimits.Delete(existingPC.RateLimit.ID) if existingPC.RateLimit != nil,
and ensure clone.ProviderConfigs does not reference them.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 96-115: The updateVirtualKey handler currently deletes and
recreates budgets (losing CurrentUsage/LastReset) because CreateBudgetRequest is
treated create-only; change updateVirtualKey to reconcile budgets in-place: when
handling UpdateVirtualKeyRequest (and the provider/ VK budget slices), fetch
existing budgets for that virtual key, match incoming budgets by stable
identifier (prefer an ID field on CreateBudgetRequest if present, otherwise a
unique key like budget name/type), and for matches update mutable fields (limit,
period) while preserving CurrentUsage and LastReset; create new DB rows only for
truly new budgets and delete only budgets absent from the incoming full payload.
Update CreateBudgetRequest/handler logic accordingly so budget IDs are
accepted/used (or implement matching logic) and remove the delete-all/insert-all
flow in updateVirtualKey to keep usage history.
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 31-40: The exhausted-state logic in isExhausted currently only
checks virtualKey.budgets and VK-level rate limits, so provider-config budgets
are missing; update the condition to also consider provider-config budgets
(e.g., add || virtualKey.provider_budgets?.some(b => b.current_usage >=
b.max_limit)) alongside the existing virtualKey.budgets check and similarly
extend the same logic used later (lines ~151-189) where exhausted-state is
computed, referencing the same symbols: isExhausted, virtualKey.budgets,
virtualKey.provider_budgets, and virtualKey.rate_limit.
---
Nitpick comments:
In `@plugins/governance/store_test.go`:
- Around line 261-264: The test currently calls store.CheckBudget(...) and uses
assert.Error(t, err) then immediately calls err.Error(), which can panic if err
is nil; change these to use require.Error(t, err) (or wrap the follow-up checks
in if assert.Error(t, err) { ... }) so execution stops when no error is
returned; update the CheckBudget call sites in this file (the occurrences around
the current assert blocks and the other instances noted) to either use
require.Error for the initial assertion or guard the subsequent
assert.Contains(err.Error(), ...) inside an if assert.Error(...) block so err is
never dereferenced when nil.
- Around line 511-550: The test currently only verifies VK.CalendarAligned
persists and that an under-limit request passes; add a subcase that exercises
calendar-aligned resets by simulating a boundary crossing: create the same daily
and monthly budgets but set their LastReset timestamps to just before the
previous calendar boundary (e.g., yesterday for the 1d budget and last month for
the 1M budget) or otherwise ensure they are older than the current calendar
period, then call store.CheckBudget(context.Background(), vk,
&EvaluationRequest{Provider: schemas.OpenAI}, nil) and assert that the budgets'
CurrentUsage have been reset (or reduced according to reset logic) before
evaluation; use buildVirtualKeyWithMultiBudgets, NewLocalGovernanceStore,
CheckBudget and store.GetVirtualKey/GetBudget to locate and verify the
post-check state.
In `@plugins/governance/store.go`:
- Around line 1475-1499: The reset path currently does a full
gs.virtualKeys.Range scan to find CalendarAligned for each budget (checking
budget.VirtualKeyID or budget.ProviderConfigID against
configstoreTables.TableVirtualKey and its ProviderConfigs), which is
O(budgets×virtualKeys); add O(1) lookups instead: maintain a lookup map in the
governance store (e.g., vkByID map[int64]*configstoreTables.TableVirtualKey and
a providerConfigToVK map[ProviderConfigIDType]*configstoreTables.TableVirtualKey
or vkID) that is populated/updated whenever virtual keys change, protect it with
the same concurrency mechanism used for gs.virtualKeys (sync.RWMutex or
sync.Map), and replace the gs.virtualKeys.Range calls in the budget reset logic
with direct map lookups to read TableVirtualKey.CalendarAligned for the found VK
(using budget.VirtualKeyID or resolving budget.ProviderConfigID via
providerConfigToVK).
In `@ui/lib/types/governance.ts`:
- Around line 137-176: The nested budget payloads for virtual-key requests
should use a VK-specific budget type without calendar_aligned to prevent
conflicting state; create a new type (e.g., CreateVirtualKeyBudgetRequest /
UpdateVirtualKeyBudgetRequest) that omits calendar_aligned from
CreateBudgetRequest/UpdateBudgetRequest and replace all uses of
CreateBudgetRequest (and UpdateBudgetRequest if present) inside
VirtualKeyProviderConfigRequest/VirtualKeyProviderConfigUpdateRequest,
CreateVirtualKeyRequest, UpdateVirtualKeyRequest, and VirtualKeyMCPConfigRequest
payloads (also update the occurrences mentioned around lines 205-215) so
top-level calendar_aligned remains the single source of truth.
🪄 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: 1937babe-ee7d-4748-9e8e-9c2355d24a8b
📒 Files selected for processing (25)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (3)
- .claude/skills/expect
- .agents/skills/expect/SKILL.md
- ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
🚧 Files skipped from review as they are similar to previous changes (12)
- plugins/governance/http_transport_prehook_test.go
- plugins/governance/resolver.go
- transports/bifrost-http/handlers/governance_test.go
- plugins/governance/model_provider_governance_test.go
- ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
- framework/configstore/rdb_test.go
- framework/configstore/tables/budget.go
- plugins/governance/resolver_test.go
- plugins/governance/test_utils.go
- .github/workflows/scripts/run-migration-tests.sh
- framework/configstore/tables/virtualkey.go
- transports/bifrost-http/lib/config_test.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
2ab508f to
5cff823
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
plugins/governance/store.go (1)
3619-3635:⚠️ Potential issue | 🟠 MajorKeep routing budget status aligned with
CheckBudget.This new loop covers provider-config budgets, but the function still never folds in
vk.Budgetsor inherited team/customer budgets. A key exhausted only at those levels will still reportBudgetPercentUsed == 0here, so budget-aware routing can pick a path thatCheckBudgetimmediately rejects. ReusecollectBudgetsFromHierarchyhere so status and enforcement evaluate the same hierarchy.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 3619 - 3635, The provider-config loop currently only inspects pc.Budgets and gs.budgets, missing vk.Budgets and inherited team/customer budgets so BudgetPercentUsed can be zero while CheckBudget would later reject; replace the manual loop with a call to collectBudgetsFromHierarchy (the same function used by CheckBudget) to gather the full set of budgets for the given provider-config/key, then iterate that returned budget list (using budgetBaselines and gs.budgets lookups as before) to compute and set result.BudgetPercentUsed so status calculation and CheckBudget use the identical hierarchy.transports/bifrost-http/handlers/governance.go (1)
1350-1356:⚠️ Potential issue | 🟠 MajorDon’t silently ignore
calendar_alignedon non-VK budgets.All of these create/recreate paths now hardcode
budgetLastReset(false, ...), butCreateBudgetRequest/UpdateBudgetRequeststill acceptcalendar_aligned. Callers can keep sending that field and get a rolling budget with no validation error. Either keep honoring the flag here or reject it explicitly so the API surface does not advertise a setting this code ignores.Also applies to: 1513-1518, 1751-1757, 1904-1909, 2290-2296, 2418-2423, 2674-2679
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 1350 - 1356, The handlers currently ignore req.Budget.CalendarAligned and always call budgetLastReset(false,...), so add validation in the Create/Update budget paths (where req.Budget is processed) to explicitly reject calendar_aligned for non-VK budgets: check req.Budget.CalendarAligned (from CreateBudgetRequest/UpdateBudgetRequest) and if true return a 400/BadRequest with a clear message like "calendar_aligned is only supported for VK budgets" instead of silently using budgetLastReset(false,...); apply the same check in the other identical create/recreate spots that call budgetLastReset (the blocks invoking budgetLastReset(false,...)) so the API no longer advertises an ignored setting.
♻️ Duplicate comments (9)
framework/configstore/rdb.go (1)
2319-2350:⚠️ Potential issue | 🟠 MajorMake
DeleteVirtualKeyProviderConfigatomic and clear key-join rows first.This path performs multiple deletes, but when no external
txis passed it is not wrapped in a local transaction, and unlikeDeleteVirtualKey(Line 2100) it does not explicitly cleargovernance_virtual_key_provider_config_keysfirst. A failure mid-path can leave partial deletion state.Suggested transactional cleanup pattern
func (s *RDBConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } - // First fetch the provider config to get budget and rate limit IDs - var providerConfig tables.TableVirtualKeyProviderConfig - if err := txDB.WithContext(ctx).First(&providerConfig, "id = ?", id).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return ErrNotFound - } - return err - } - // Store the rate limit ID before deleting - rateLimitID := providerConfig.RateLimitID - // Delete budgets owned by this provider config - if err := txDB.WithContext(ctx).Where("provider_config_id = ?", id).Delete(&tables.TableBudget{}).Error; err != nil { - return err - } - // Delete the provider config - if err := txDB.WithContext(ctx).Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id).Error; err != nil { - return err - } - // Delete the rate limit if it exists - if rateLimitID != nil { - if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil { - return err - } - } - return nil + run := func(txDB *gorm.DB) error { + var providerConfig tables.TableVirtualKeyProviderConfig + if err := txDB.WithContext(ctx).First(&providerConfig, "id = ?", id).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ErrNotFound + } + return err + } + rateLimitID := providerConfig.RateLimitID + + if err := txDB.WithContext(ctx).Exec( + "DELETE FROM governance_virtual_key_provider_config_keys WHERE table_virtual_key_provider_config_id = ?", + id, + ).Error; err != nil { + return err + } + if err := txDB.WithContext(ctx).Where("provider_config_id = ?", id).Delete(&tables.TableBudget{}).Error; err != nil { + return err + } + if err := txDB.WithContext(ctx).Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id).Error; err != nil { + return err + } + if rateLimitID != nil { + if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil { + return err + } + } + return nil + } + + if len(tx) > 0 { + return run(tx[0]) + } + return s.db.WithContext(ctx).Transaction(run) }Based on learnings: In the Bifrost codebase, budgets and rate limits have a 1:1 ownership with their parent entities (ModelConfig, Provider, VirtualKey, Team, Customer). They are created along with the parent and deleted together.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2319 - 2350, DeleteVirtualKeyProviderConfig currently performs multiple deletions without ensuring atomicity or clearing the join table; update it to begin a local transaction when no external tx is provided, and perform deletions inside that tx in this order: delete rows from governance_virtual_key_provider_config_keys for provider_config_id=id, delete budgets (TableBudget) where provider_config_id=id, fetch providerConfig to obtain rateLimitID (TableVirtualKeyProviderConfig), delete the provider config (TableVirtualKeyProviderConfig) by id, then delete the rate limit (TableRateLimit) if rateLimitID is non-nil; finally commit the tx or rollback on any error—mirror the transactional pattern used by DeleteVirtualKey to locate where to implement the tx handling and cleanup.plugins/governance/test_utils.go (1)
204-208:⚠️ Potential issue | 🟠 MajorStamp
ProviderConfigIDin provider-config budget fixtures.At Line 204, budgets are attached to
pc.Budgetsbut owner FK (provider_config_id) is not set. That can let tests miss regressions in owner-scoped budget handling.Suggested fixture adjustment
-func buildProviderConfigWithBudgets(provider string, allowedModels []string, budgets []configstoreTables.TableBudget) configstoreTables.TableVirtualKeyProviderConfig { - pc := buildProviderConfig(provider, allowedModels) - pc.Budgets = budgets +func buildProviderConfigWithBudgets(providerConfigID uint, provider string, allowedModels []string, budgets []configstoreTables.TableBudget) configstoreTables.TableVirtualKeyProviderConfig { + pc := buildProviderConfig(provider, allowedModels) + pc.ID = providerConfigID + for i := range budgets { + budgets[i].ProviderConfigID = bifrost.Ptr(providerConfigID) + budgets[i].VirtualKeyID = nil + } + pc.Budgets = budgets return pc }Based on learnings: In the governance store architecture for plugins/governance/store.go, budgets and rate limits are never shared between virtual keys or other entities. Each entity has its own dedicated budget/rate limit instance.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 204 - 208, The test fixture helper buildProviderConfigWithBudgets attaches budgets to pc.Budgets but does not set the owner foreign-key, so budgets lack ProviderConfigID and tests can miss owner-scoped bugs; fix by iterating the provided budgets before assigning (in buildProviderConfigWithBudgets) and set each budget's ProviderConfigID (the TableBudget.ProviderConfigID field) to the created provider config's ID (pc.ProviderConfigID), then assign pc.Budgets = budgets; this ensures each TableBudget is stamped with the proper owner ID returned by buildProviderConfig.framework/configstore/tables/virtualkey.go (1)
38-38:⚠️ Potential issue | 🟡 MinorFix stale
Keyscomment to match current deny-by-default behavior.At Line 38, the comment says empty
Keysmeans all keys are allowed, but Line 32 and the unmarshal logic useAllowAllKeysfor that. EmptyKeyswithAllowAllKeys=falseshould mean no keys allowed.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` at line 38, Update the stale comment on the Keys field to reflect the deny-by-default behavior: mention that when Keys is empty the provider allows no keys unless AllowAllKeys is true; reference the Keys field and the AllowAllKeys boolean (and the unmarshal logic that sets AllowAllKeys) so the comment reads that empty Keys means "no keys allowed" unless AllowAllKeys==true (which explicitly allows all keys).ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx (1)
151-186:⚠️ Potential issue | 🟡 MinorReuse the safe percentage helper for the new budget badges.
Both new budget sections still do raw
b.current_usage / b.max_limit, so a zero-limit budget rendersNaN%/Infinity%.calculateUsagePercentageis already imported here and keeps these badges consistent with the rate-limit sections.Also applies to: 350-381
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx` around lines 151 - 186, Replace the raw b.current_usage / b.max_limit math in the Provider Budgets badge with the existing calculateUsagePercentage helper: compute const pct = calculateUsagePercentage(b.current_usage, b.max_limit) and use pct for the Badge label (e.g. `${pct}%`) and for the variant check (e.g. variant={pct >= 100 ? "destructive" : "default"}) inside the budgets map in virtualKeyDetailsSheet.tsx; apply the same change to the other duplicate budget section mentioned so zero or infinite limits render safely and consistently.transports/config.schema.json (1)
455-459:⚠️ Potential issue | 🟠 MajorExpose the replacement multi-budget field in the config schema.
After removing the singular VK/provider-config budget reference, this schema still has no replacement array field for virtual-key budgets. That means
config.jsonstill cannot declare the new multi-budget relationship even though the rest of this stack expects it. Please add the plural budget reference here and mirror it under$defs.virtual_key_provider_configusing the same item shape the config loader consumes. As per coding guidelines,transports/config.schema.json: Updatetransports/config.schema.jsonwhen adding config fields; it is the authoritative definition for allconfig.jsonfields.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/config.schema.json` around lines 455 - 459, Add a new plural budget reference array to the schema to replace the removed singular VK/provider-config budget: create a "budgets" (or the agreed plural name) property alongside "calendar_aligned" with "type": "array" and items matching the existing virtual-key budget item shape used by the config loader, and then mirror that same array property under $defs.virtual_key_provider_config so both the root config and the virtual key provider config accept the multi-budget structure; ensure the array item schema exactly matches the loader-consumed shape (field names/types) used elsewhere in the codebase.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (2)
1178-1184:⚠️ Potential issue | 🟠 MajorGive the MCP delete button an accessible name.
This control is icon-only, so assistive tech has no label to announce. Please add an
aria-label(oraria-labelledby) before shipping.💡 Suggested fix
<Button type="button" variant="ghost" size="sm" + aria-label={`Remove MCP client ${config.mcp_client_name}`} onClick={() => handleRemoveMCPClient(index)} data-testid={`vk-delete-mcp-${index}`} >🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 1178 - 1184, The icon-only delete Button for MCP clients lacks an accessible name; update the Button component (the one with onClick={() => handleRemoveMCPClient(index)} and data-testid={`vk-delete-mcp-${index}`}) to include an aria-label or aria-labelledby (e.g., aria-label="Remove MCP client" or include client-specific text) so screen readers can announce the control; ensure the label is unique/clear if multiple clients are listed (you can incorporate index or client name if available).
378-410:⚠️ Potential issue | 🔴 CriticalNormalize provider budgets before serializing the request.
normalizeProviderConfigs()still spreads the form object verbatim. That leavesconfig.budgetsas{ max_limit: string }rows fromMultiBudgetLines, and it also forwards the legacybudgetshadow field that the API no longer declares.provider_configs[].budgets[].max_limitis numeric in both the UI request types and the Go handler, so saving any provider budget will post strings and the backend unmarshal will reject the whole request. Build the payload explicitly here and normalize/filter provider budgets the same way the top-leveldata.budgetspath does.💡 Suggested direction
const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => { - return configs.map((config) => ({ - ...config, + return configs.map((config) => { + const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider)); + const validBudgets = (config.budgets ?? []).filter((b) => normalizeNumericField(b.max_limit) !== undefined); + + return { + id: config.id, + provider: config.provider, + allowed_models: config.allowed_models, + key_ids: config.key_ids, weight: config.weight === undefined || config.weight === null ? null : typeof config.weight === "string" ? Number.isNaN(parseFloat(config.weight)) ? null : parseFloat(config.weight) : config.weight, + budgets: + validBudgets.length > 0 + ? validBudgets.map((b) => ({ + max_limit: normalizeNumericField(b.max_limit)!, + reset_duration: b.reset_duration || "1M", + })) + : existingConfig?.budgets?.length + ? [] + : undefined, rate_limit: (() => { const tokenMaxLimit = normalizeNumericField(config.rate_limit?.token_max_limit); const requestMaxLimit = normalizeNumericField(config.rate_limit?.request_max_limit); const hasTokenMaxLimit = tokenMaxLimit !== undefined; const hasRequestMaxLimit = requestMaxLimit !== undefined; @@ - const existingConfig = existingConfigs?.find((item) => (config.id ? item.id === config.id : item.provider === config.provider)); if (existingConfig?.rate_limit) { return {}; } return undefined; })(), - })); + }; + }); };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 378 - 410, normalizeProviderConfigs currently forwards config objects verbatim which keeps provider_configs[].budgets as strings and still includes the legacy budget shadow field; update normalizeProviderConfigs so each returned config builds an explicit payload: omit any legacy `budget` field, and replace the spread budgets with a normalized `budgets` array constructed like the top-level data.budgets logic — map config.budgets (from MultiBudgetLines) and for each row use normalizeNumericField(row.max_limit) (or parseFloat with NaN -> undefined) to produce numeric max_limit (or null/undefined per top-level behavior), filter out invalid/empty budget rows, and ensure the resulting budgets elements match the API shape (numeric max_limit type) before returning the config object (keep existing rate_limit normalization and other fields intact).transports/bifrost-http/handlers/governance.go (1)
442-459:⚠️ Potential issue | 🟠 MajorReturn 400 for every VK/provider budget payload validation failure.
These branches still let plain
validateBudgeterrors escape.createVirtualKeyonly mapsbadRequestErrorto 400, andupdateVirtualKeydoes the same for wrapped errors, so client payloads likemax_limit == 0still come back as 500. The provider-config create paths also skip the earliermax_limit/reset_durationprecheck entirely. Please wrap everyvalidateBudgetfailure here asbadRequestError(or makevalidateBudgetreturn that type) so bad budget payloads stay in the 4xx path.Also applies to: 499-516, 575-597, 741-755, 763-779, 912-935, 971-1007
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 442 - 459, The budget validation currently lets validateBudget errors propagate as non-400s; update the validation here (the budgets loop in governance.go) so any failure from validateBudget is converted to a badRequestError (or change validateBudget to return a badRequestError type) and returned/sent as a 400; specifically, wrap or map validateBudget errors to badRequestError before calling SendError so createVirtualKey and updateVirtualKey receive a 4xx for payload issues, and apply the same wrapping to the other budget validation blocks referenced (around lines 499-516, 575-597, 741-755, 763-779, 912-935, 971-1007) to ensure all provider/config and VK create/update paths return 400 for invalid budget payloads.ui/components/ui/multiBudgetLines.tsx (1)
64-70:⚠️ Potential issue | 🟡 MinorUse distinct test IDs for each budget editor instance.
budget-add-btnandbudget-remove-${index}repeat for everyMultiBudgetLinesinstance, so the VK budget editor and each provider budget editor expose the same selectors. That makes these new controls brittle to target once multiple sections are rendered. Please scope add/remove/reset with one explicit qualifier per component instance instead of mixing a rawidfor reset and sharedbudget-*IDs for the rest.As per coding guidelines, "Add new interactive UI elements with
data-testidattributes following the pattern:data-testid=\"<entity>-<element>-<qualifier>\"."Also applies to: 99-106
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/components/ui/multiBudgetLines.tsx` around lines 64 - 70, The add/remove/reset buttons in MultiBudgetLines use duplicate test IDs (e.g. "budget-add-btn" and "budget-remove-${index}") across instances; update the data-testid attributes to include the component instance qualifier (use the component's id prop or another instance-specific prop) so each instance has unique selectors—e.g. change the add button in the MultiBudgetLines render (where addLine is used) to data-testid={`budget-add-btn-${id}`} and likewise scope the remove buttons (budget-remove-${index} → budget-remove-${id}-${index}) and the reset button (currently using id for only reset; ensure it follows the same pattern) so all interactive elements follow the pattern <entity>-<element>-<qualifier>.
🧹 Nitpick comments (1)
.github/workflows/scripts/run-migration-tests.sh (1)
2621-2626: Consider adding SQLite budget migration verification.
verify_budget_migration_postgres()validates the multi-budget FK migration on PostgreSQL, but there's no corresponding verification for SQLite. If the migration applies to both database types, consider adding averify_budget_migration_sqlite()function to ensure consistent test coverage.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/scripts/run-migration-tests.sh around lines 2621 - 2626, The test suite only has verify_budget_migration_postgres() to validate the budget_id → virtual_key_id/provider_config_id migration, so add a counterpart verify_budget_migration_sqlite() that mirrors the Postgres checks (same assertions, queries and failure reporting) but using SQLite connection/commands; update the test runner to call verify_budget_migration_sqlite() when running against SQLite (same place where verify_budget_migration_postgres() is invoked), and reuse any shared helpers to avoid duplication while ensuring the SQLite-specific SQL/connection handling is used.
🤖 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/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx`:
- Around line 7-13: The ContactUsView instance in accessProfilesIndexView.tsx is
missing the testIdPrefix prop which prevents its interactive buttons from
emitting stable data-testid attributes; add
testIdPrefix="access-profiles-contact" (or similar following the
<entity>-<element>-<qualifier> pattern) to the ContactUsView component so its
internal buttons render with predictable data-testid values; update the
ContactUsView invocation (the one with icon={<ShieldCheck .../>} and
title="Unlock access profiles for better performance") to include this prop.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 367-369: The clearVirtualKeyBudget function currently only clears
the "budgets" field; update it to also reset the alignment flag by calling
form.setValue("budgetCalendarAligned", false, { shouldDirty: true }) (or the
appropriate default) so that when budgets are cleared the budgetCalendarAligned
state is reset too; modify clearVirtualKeyBudget to set both "budgets" and
"budgetCalendarAligned" via form.setValue to ensure a clean state before further
edits.
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 3619-3635: The provider-config loop currently only inspects
pc.Budgets and gs.budgets, missing vk.Budgets and inherited team/customer
budgets so BudgetPercentUsed can be zero while CheckBudget would later reject;
replace the manual loop with a call to collectBudgetsFromHierarchy (the same
function used by CheckBudget) to gather the full set of budgets for the given
provider-config/key, then iterate that returned budget list (using
budgetBaselines and gs.budgets lookups as before) to compute and set
result.BudgetPercentUsed so status calculation and CheckBudget use the identical
hierarchy.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 1350-1356: The handlers currently ignore
req.Budget.CalendarAligned and always call budgetLastReset(false,...), so add
validation in the Create/Update budget paths (where req.Budget is processed) to
explicitly reject calendar_aligned for non-VK budgets: check
req.Budget.CalendarAligned (from CreateBudgetRequest/UpdateBudgetRequest) and if
true return a 400/BadRequest with a clear message like "calendar_aligned is only
supported for VK budgets" instead of silently using budgetLastReset(false,...);
apply the same check in the other identical create/recreate spots that call
budgetLastReset (the blocks invoking budgetLastReset(false,...)) so the API no
longer advertises an ignored setting.
---
Duplicate comments:
In `@framework/configstore/rdb.go`:
- Around line 2319-2350: DeleteVirtualKeyProviderConfig currently performs
multiple deletions without ensuring atomicity or clearing the join table; update
it to begin a local transaction when no external tx is provided, and perform
deletions inside that tx in this order: delete rows from
governance_virtual_key_provider_config_keys for provider_config_id=id, delete
budgets (TableBudget) where provider_config_id=id, fetch providerConfig to
obtain rateLimitID (TableVirtualKeyProviderConfig), delete the provider config
(TableVirtualKeyProviderConfig) by id, then delete the rate limit
(TableRateLimit) if rateLimitID is non-nil; finally commit the tx or rollback on
any error—mirror the transactional pattern used by DeleteVirtualKey to locate
where to implement the tx handling and cleanup.
In `@framework/configstore/tables/virtualkey.go`:
- Line 38: Update the stale comment on the Keys field to reflect the
deny-by-default behavior: mention that when Keys is empty the provider allows no
keys unless AllowAllKeys is true; reference the Keys field and the AllowAllKeys
boolean (and the unmarshal logic that sets AllowAllKeys) so the comment reads
that empty Keys means "no keys allowed" unless AllowAllKeys==true (which
explicitly allows all keys).
In `@plugins/governance/test_utils.go`:
- Around line 204-208: The test fixture helper buildProviderConfigWithBudgets
attaches budgets to pc.Budgets but does not set the owner foreign-key, so
budgets lack ProviderConfigID and tests can miss owner-scoped bugs; fix by
iterating the provided budgets before assigning (in
buildProviderConfigWithBudgets) and set each budget's ProviderConfigID (the
TableBudget.ProviderConfigID field) to the created provider config's ID
(pc.ProviderConfigID), then assign pc.Budgets = budgets; this ensures each
TableBudget is stamped with the proper owner ID returned by buildProviderConfig.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 442-459: The budget validation currently lets validateBudget
errors propagate as non-400s; update the validation here (the budgets loop in
governance.go) so any failure from validateBudget is converted to a
badRequestError (or change validateBudget to return a badRequestError type) and
returned/sent as a 400; specifically, wrap or map validateBudget errors to
badRequestError before calling SendError so createVirtualKey and
updateVirtualKey receive a 4xx for payload issues, and apply the same wrapping
to the other budget validation blocks referenced (around lines 499-516, 575-597,
741-755, 763-779, 912-935, 971-1007) to ensure all provider/config and VK
create/update paths return 400 for invalid budget payloads.
In `@transports/config.schema.json`:
- Around line 455-459: Add a new plural budget reference array to the schema to
replace the removed singular VK/provider-config budget: create a "budgets" (or
the agreed plural name) property alongside "calendar_aligned" with "type":
"array" and items matching the existing virtual-key budget item shape used by
the config loader, and then mirror that same array property under
$defs.virtual_key_provider_config so both the root config and the virtual key
provider config accept the multi-budget structure; ensure the array item schema
exactly matches the loader-consumed shape (field names/types) used elsewhere in
the codebase.
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 151-186: Replace the raw b.current_usage / b.max_limit math in the
Provider Budgets badge with the existing calculateUsagePercentage helper:
compute const pct = calculateUsagePercentage(b.current_usage, b.max_limit) and
use pct for the Badge label (e.g. `${pct}%`) and for the variant check (e.g.
variant={pct >= 100 ? "destructive" : "default"}) inside the budgets map in
virtualKeyDetailsSheet.tsx; apply the same change to the other duplicate budget
section mentioned so zero or infinite limits render safely and consistently.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 1178-1184: The icon-only delete Button for MCP clients lacks an
accessible name; update the Button component (the one with onClick={() =>
handleRemoveMCPClient(index)} and data-testid={`vk-delete-mcp-${index}`}) to
include an aria-label or aria-labelledby (e.g., aria-label="Remove MCP client"
or include client-specific text) so screen readers can announce the control;
ensure the label is unique/clear if multiple clients are listed (you can
incorporate index or client name if available).
- Around line 378-410: normalizeProviderConfigs currently forwards config
objects verbatim which keeps provider_configs[].budgets as strings and still
includes the legacy budget shadow field; update normalizeProviderConfigs so each
returned config builds an explicit payload: omit any legacy `budget` field, and
replace the spread budgets with a normalized `budgets` array constructed like
the top-level data.budgets logic — map config.budgets (from MultiBudgetLines)
and for each row use normalizeNumericField(row.max_limit) (or parseFloat with
NaN -> undefined) to produce numeric max_limit (or null/undefined per top-level
behavior), filter out invalid/empty budget rows, and ensure the resulting
budgets elements match the API shape (numeric max_limit type) before returning
the config object (keep existing rate_limit normalization and other fields
intact).
In `@ui/components/ui/multiBudgetLines.tsx`:
- Around line 64-70: The add/remove/reset buttons in MultiBudgetLines use
duplicate test IDs (e.g. "budget-add-btn" and "budget-remove-${index}") across
instances; update the data-testid attributes to include the component instance
qualifier (use the component's id prop or another instance-specific prop) so
each instance has unique selectors—e.g. change the add button in the
MultiBudgetLines render (where addLine is used) to
data-testid={`budget-add-btn-${id}`} and likewise scope the remove buttons
(budget-remove-${index} → budget-remove-${id}-${index}) and the reset button
(currently using id for only reset; ensure it follows the same pattern) so all
interactive elements follow the pattern <entity>-<element>-<qualifier>.
---
Nitpick comments:
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2621-2626: The test suite only has
verify_budget_migration_postgres() to validate the budget_id →
virtual_key_id/provider_config_id migration, so add a counterpart
verify_budget_migration_sqlite() that mirrors the Postgres checks (same
assertions, queries and failure reporting) but using SQLite connection/commands;
update the test runner to call verify_budget_migration_sqlite() when running
against SQLite (same place where verify_budget_migration_postgres() is invoked),
and reuse any shared helpers to avoid duplication while ensuring the
SQLite-specific SQL/connection handling is used.
🪄 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: 99168762-075b-4949-9b86-54dd763cf140
📒 Files selected for processing (28)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/ui/multiBudgetLines.tsxui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (4)
- .claude/skills/expect
- framework/configstore/rdb_test.go
- .agents/skills/expect/SKILL.md
- plugins/governance/resolver.go
🚧 Files skipped from review as they are similar to previous changes (9)
- plugins/governance/model_provider_governance_test.go
- plugins/governance/http_transport_prehook_test.go
- transports/bifrost-http/handlers/governance_test.go
- ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
- framework/configstore/tables/budget.go
- plugins/governance/resolver_test.go
- plugins/governance/store_test.go
- framework/configstore/migrations.go
- transports/bifrost-http/lib/config_test.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
5cff823 to
a2077df
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/governance/store.go (1)
2080-2095:⚠️ Potential issue | 🟠 MajorConfig-memory path is not wiring VK/provider-config budgets.
loadFromConfigMemorynow populates provider-config rate limits only. If budgets are supplied via top-levelconfig.Budgets(same pattern used for model/provider hydration),vk.Budgets/pc.Budgetsstay empty, and VK budget checks can be skipped.💡 Proposed fix (rebuild budget ownership maps in memory-mode)
// Load budgets budgets := config.Budgets + + // Build ownership maps for multi-budgets + budgetsByVKID := make(map[string][]configstoreTables.TableBudget) + budgetsByProviderConfigID := make(map[uint][]configstoreTables.TableBudget) + for i := range budgets { + b := budgets[i] + if b.VirtualKeyID != nil { + budgetsByVKID[*b.VirtualKeyID] = append(budgetsByVKID[*b.VirtualKeyID], b) + } + if b.ProviderConfigID != nil { + budgetsByProviderConfigID[*b.ProviderConfigID] = append(budgetsByProviderConfigID[*b.ProviderConfigID], b) + } + } @@ for i := range virtualKeys { vk := &virtualKeys[i] + if len(vk.Budgets) == 0 { + vk.Budgets = append(vk.Budgets, budgetsByVKID[vk.ID]...) + } @@ if vk.ProviderConfigs != nil { for j := range vk.ProviderConfigs { pc := &vk.ProviderConfigs[j] + if len(pc.Budgets) == 0 { + pc.Budgets = append(pc.Budgets, budgetsByProviderConfigID[pc.ID]...) + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2080 - 2095, loadFromConfigMemory currently wires only rate limits for vk.ProviderConfigs and ignores top-level config.Budgets, leaving vk.Budgets and pc.Budgets empty; fix this by rebuilding budget ownership maps in memory-mode: iterate over config.Budgets, for each budget inspect its Owner/OwnerType and attach the budget to the matching virtual kernel (vk.Budgets) or provider-config (pc.Budgets) similar to the existing model/provider hydration pattern, ensuring you initialize vk.Budgets and pc.Budgets slices/maps before appending and use the same ID matching approach as for RateLimitID (compare budget.Owner/OwnerType to vk.ID and pc.ID) so VK budget checks see the budgets.
♻️ Duplicate comments (5)
framework/configstore/rdb.go (1)
2334-2342:⚠️ Potential issue | 🟠 MajorKeep standalone provider-config deletes atomic.
This helper now removes owned budgets, the provider-config row, and then the rate limit, but it only becomes atomic when a caller passes
tx. A failure mid-sequence leaves partially deleted governance state.Based on learnings: budgets and rate limits have a 1:1 ownership with their parent entities and are deleted together.🔒 Proposed transaction wrapper
func (s *RDBConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db + if len(tx) == 0 { + return s.db.WithContext(ctx).Transaction(func(inner *gorm.DB) error { + return s.DeleteVirtualKeyProviderConfig(ctx, id, inner) + }) } + txDB := tx[0] + // First fetch the provider config to get budget and rate limit IDs var providerConfig tables.TableVirtualKeyProviderConfig🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2334 - 2342, The sequence that deletes budgets, the provider-config row, and the rate limit must be performed inside a DB transaction in this helper so it is atomic even when the caller did not pass one; modify the code (use the existing txDB, ctx and the stored rateLimitID variable) to start a transaction (e.g., tx := txDB.Begin()/txDB.WithContext(ctx).BeginTx or use txDB.Transaction(fn)), run the Where(...).Delete(&tables.TableBudget{}), Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id) and the subsequent rate-limit deletion using that tx variable instead of txDB, and commit or rollback on error so partial deletes cannot be left behind. Ensure all operations reference the stored rateLimitID and return tx errors as before.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (3)
367-369:⚠️ Potential issue | 🟡 MinorReset the alignment toggle when clearing VK budgets.
Clearing the rows currently leaves
budgetCalendarAligneduntouched. If the user adds a new budget in the same edit session, the old alignment state is reused unexpectedly.💡 Suggested fix
const clearVirtualKeyBudget = () => { form.setValue("budgets", [], { shouldDirty: true }); + form.setValue("budgetCalendarAligned", false, { shouldDirty: true }); };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 367 - 369, The clearVirtualKeyBudget function only clears the "budgets" field but does not reset the alignment toggle; update clearVirtualKeyBudget to also reset the budgetCalendarAligned form field (or state) so the toggle returns to its default when budgets are cleared — locate clearVirtualKeyBudget and call form.setValue("budgetCalendarAligned", <defaultValue>, { shouldDirty: true }) (or the equivalent state setter) immediately after clearing "budgets" to ensure new budgets in the same edit session do not inherit the old alignment.
84-91:⚠️ Potential issue | 🔴 CriticalProvider budgets are still posted in form-string shape.
normalizeProviderConfigs()convertsweightandrate_limit, but it still forwardsconfig.budgetsexactly as stored in the form. Since this schema keepsbudgets[*].max_limitas a string, provider multi-budget saves still serialize"max_limit":"..."instead of a number, and the Go handler will reject the request during JSON decoding.💡 Suggested fix
const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => { return configs.map((config) => ({ ...config, + budgets: + config.budgets === undefined + ? undefined + : config.budgets + .filter((b) => normalizeNumericField(b.max_limit) !== undefined) + .map((b) => ({ + max_limit: normalizeNumericField(b.max_limit)!, + reset_duration: b.reset_duration || "1M", + })), weight: config.weight === undefined || config.weight === null ? nullAlso applies to: 378-410
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 84 - 91, normalizeProviderConfigs currently converts weight and rate_limit but forwards config.budgets as form strings, causing budgets[*].max_limit and budgets[*].reset_duration to be sent as strings; update normalizeProviderConfigs to map over config.budgets (if present) and coerce max_limit and reset_duration to numbers (e.g., parseInt/Number or null when empty) before returning the normalized config so the outgoing JSON matches the Go handler's expected numeric types; reference the normalizeProviderConfigs function and the config.budgets/budgets[*].max_limit and budgets[*].reset_duration fields when making the change.
905-937:⚠️ Potential issue | 🟡 MinorAdd
data-testidhooks to the new budget editors.Both
MultiBudgetLinesinstances are new interactive controls, but neither usage exposes adata-testid. IfMultiBudgetLinesdoes not forward arbitrary DOM props yet, this should be added there too.As per coding guidelines,
ui/**/*.{tsx,ts}:Add new interactive UI elements with data-testid attributes following the pattern: data-testid="<entity>-<element>-<qualifier>".Also applies to: 1200-1209
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 905 - 937, Add a data-testid to the MultiBudgetLines usages and ensure the MultiBudgetLines component forwards arbitrary DOM props (including data-testid) to its root element; specifically, update the two MultiBudgetLines instances (currently using id={`providerBudget-${index}`}) to include a data-testid following the pattern data-testid="provider-budget-editor-<qualifier>" (e.g., including the index) and, if MultiBudgetLines does not already accept/rest props, modify its props signature to accept ...rest and spread those onto the top-level DOM node so data-testid is actually rendered; keep existing behavior in handleUpdateProviderConfig unchanged.transports/bifrost-http/handlers/governance.go (1)
499-516:⚠️ Potential issue | 🟠 MajorWrap multi-budget validation failures as bad requests.
These paths still
return errdirectly aftervalidateBudget(&budget).validateBudget()rejectsmax_limit == 0and bad durations, so invalid client payloads can still bubble out as 500s from both create and update VK flows instead of 400s.💡 Suggested fix
if err := validateBudget(&budget); err != nil { - return err + return &badRequestError{err: err} }Apply the same wrapping at each VK/provider multi-budget call site.
Also applies to: 575-597, 762-785, 912-935, 994-1007
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 499 - 516, The validateBudget call returns client-side validation errors (e.g., zero max_limit or bad duration) but the handler currently returns them directly causing 500s; change the validateBudget error path so validation failures are converted into an explicit 400 Bad Request before returning (e.g., after if err := validateBudget(&budget); err != nil { return wrapAsBadRequest(err) }), then proceed to call h.configStore.CreateBudget; apply this same pattern for each multi-budget site (the blocks that call validateBudget and h.configStore.CreateBudget around req.Budgets / vk creation/update).
🧹 Nitpick comments (1)
plugins/governance/store.go (1)
1475-1499: Consider indexing ownership once per reset cycle.
ResetExpiredBudgetsInMemorycurrently scans all virtual keys (and nested provider configs) for each budget to resolveCalendarAligned. This becomes O(B×VK×PC). Precomputing{virtualKeyID -> calendarAligned}and{providerConfigID -> calendarAligned}once per call would keep resets linear.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 1475 - 1499, The ResetExpiredBudgetsInMemory function currently scans gs.virtualKeys for every budget; instead, before iterating budgets build two lookup maps once per call — e.g. map[VirtualKeyID]bool and map[ProviderConfigID]bool — by ranging gs.virtualKeys and inspecting each *configstoreTables.TableVirtualKey and its ProviderConfigs to set CalendarAligned, then replace the inner gs.virtualKeys.Range scans with O(1) lookups of those maps when resolving calendarAligned for a budget (use budget.VirtualKeyID and budget.ProviderConfigID as keys); this keeps the algorithm linear and preserves the existing TableVirtualKey/ProviderConfigs logic.
🤖 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 5299-5306: The migration currently no-ops; implement the Migrate
function in migrationAddBudgetCalendarAlignedColumn to backfill legacy calendar
alignment: query governance_budgets for rows where calendar_aligned = true and
set governance_virtual_keys.calendar_aligned = true for the corresponding
virtual key records (including budgets that are direct-owned or
provider-config-owned), using the tx *gorm.DB transaction; ensure you update
only VKs that exist and run the update idempotently (use joins or IN-subquery
against governance_budgets.vk_id/governance_budgets.virtual_key_id as
appropriate) so that after migration reads can safely switch to
governance_virtual_keys.calendar_aligned.
In `@framework/configstore/tables/budget.go`:
- Around line 18-20: Add a GORM BeforeSave hook on the Budget model (implement
BeforeSave(tx *gorm.DB) error) that validates ownership: return a non-nil error
if both VirtualKeyID and ProviderConfigID are non‑nil (reject budgets that point
at both owners), but allow both to be nil (team/customer/provider/model-config
budgets remain valid); update the Budget struct's BeforeSave method to perform
this check and return a clear error message when the guard fails.
In `@plugins/governance/store.go`:
- Around line 2490-2507: The delete logic can remove budgets that were moved
between VK-level and provider-config-level within the same update because it
only checks newBudgetIDs (VK) or pcNewBudgetIDs (per-PC) separately; change the
code in the update routine around gs.budgets handling (where clone.Budgets are
stored and later deleted, and similarly in the per-PC block at lines ~2571-2578)
to first collect a single combined set (e.g., incomingBudgetIDs) of all budget
IDs encountered in this update from both clone.Budgets and any provider-config
budgets, use that combined set when deciding which existingVK.Budgets (and
per-PC budgets) to Delete from gs.budgets so only truly removed IDs are deleted.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 740-785: The PUT handler is deleting and recreating all vk.Budgets
even when unchanged, which resets CurrentUsage/LastReset; instead, in
updateVirtualKey compare req.Budgets to existing vk.Budgets and preserve
matching rows in-place: for each req.Budgets entry look up a matching existing
configstoreTables.TableBudget by ResetDuration and MaxLimit (keep ID,
CurrentUsage, LastReset), only create a new TableBudget (via
h.configStore.CreateBudget) when no match exists or when MaxLimit/ResetDuration
changed, and only zero CurrentUsage and recalculate LastReset (using
budgetLastReset) when the budget is being replaced or when calendar alignment is
being enabled (vk.CalendarAligned toggles true). After reconciling, delete any
existing TableBudget rows not present in the reconciled set and assign
vk.Budgets to the final slice; keep validateBudget/budgetLastReset usage for
new/forced-reset budgets.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 291-294: The component currently derives calendar_aligned and
other UI state from the first element (firstBudget) of watchedBudgets, which
prevents keys whose first row is non-alignable (or only provider-config budgets)
from showing or submitting the VK-wide calendar_aligned flag; update the logic
in VirtualKeySheet so that calendar_aligned is computed across all
watchedBudgets (e.g., any budget with calendar_aligned true or any budget that
supports alignment) rather than only firstBudget, show the calendar alignment
control when any budget is alignable, and when submitting, write the VK-level
calendar_aligned flag (not just per-first-budget values); also replace usages of
watchedBudgetMaxLimit and watchedBudgetResetDuration to derive from a chosen
alignable budget or from aggregated/default values across watchedBudgets so
later alignable rows aren’t ignored.
In `@ui/components/ui/multiBudgetLines.tsx`:
- Around line 64-70: The test IDs in MultiBudgetLines collide across instances;
update data-testid attributes to scope them with the component's id prop (e.g.,
use `${id}-budget-add-btn` for the add button and `${id}-budget-remove-${index}`
for remove buttons) in the JSX around the addLine handler and the remove
button(s) (also update the occurrences around lines where budget-remove and
budget-add are defined, including the places noted at 100-101); ensure all
interactive elements in this component follow the pattern
data-testid={`${id}-<entity>-<element>-<qualifier>`} so selectors remain unique.
- Around line 41-45: The addLine function currently falls back to the hardcoded
"1M" when no unused duration is found; change the fallback so the new row's
reset_duration is always one of the select options by using available?.value ??
options[0]?.value ?? "" instead of "1M". Update the line that builds the new row
in addLine (which references lines, options, available, and calls onChange) to
use that fallback so newly added rows never contain a value not present in
options.
In `@ui/lib/types/governance.ts`:
- Around line 94-95: The frontend is using Budget/CreateBudgetRequest which
include calendar_aligned (a field the backend no longer persists), so replace
those usages with a provider-config-specific budget type that omits
calendar_aligned; add a new type (e.g., ProviderConfigBudget or TableBudgetLite)
matching the backend-serialised shape used by
TableBudget/TableVirtualKeyProviderConfig (no calendar_aligned), update the
budgets?: Budget[] property to budgets?: ProviderConfigBudget[] and replace
other CreateBudgetRequest references (see the block around lines 133-149) to use
the new provider-config budget type so the frontend won’t send a field the Go
backend cannot round-trip.
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 2080-2095: loadFromConfigMemory currently wires only rate limits
for vk.ProviderConfigs and ignores top-level config.Budgets, leaving vk.Budgets
and pc.Budgets empty; fix this by rebuilding budget ownership maps in
memory-mode: iterate over config.Budgets, for each budget inspect its
Owner/OwnerType and attach the budget to the matching virtual kernel
(vk.Budgets) or provider-config (pc.Budgets) similar to the existing
model/provider hydration pattern, ensuring you initialize vk.Budgets and
pc.Budgets slices/maps before appending and use the same ID matching approach as
for RateLimitID (compare budget.Owner/OwnerType to vk.ID and pc.ID) so VK budget
checks see the budgets.
---
Duplicate comments:
In `@framework/configstore/rdb.go`:
- Around line 2334-2342: The sequence that deletes budgets, the provider-config
row, and the rate limit must be performed inside a DB transaction in this helper
so it is atomic even when the caller did not pass one; modify the code (use the
existing txDB, ctx and the stored rateLimitID variable) to start a transaction
(e.g., tx := txDB.Begin()/txDB.WithContext(ctx).BeginTx or use
txDB.Transaction(fn)), run the Where(...).Delete(&tables.TableBudget{}),
Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id) and the subsequent
rate-limit deletion using that tx variable instead of txDB, and commit or
rollback on error so partial deletes cannot be left behind. Ensure all
operations reference the stored rateLimitID and return tx errors as before.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 499-516: The validateBudget call returns client-side validation
errors (e.g., zero max_limit or bad duration) but the handler currently returns
them directly causing 500s; change the validateBudget error path so validation
failures are converted into an explicit 400 Bad Request before returning (e.g.,
after if err := validateBudget(&budget); err != nil { return
wrapAsBadRequest(err) }), then proceed to call h.configStore.CreateBudget; apply
this same pattern for each multi-budget site (the blocks that call
validateBudget and h.configStore.CreateBudget around req.Budgets / vk
creation/update).
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 367-369: The clearVirtualKeyBudget function only clears the
"budgets" field but does not reset the alignment toggle; update
clearVirtualKeyBudget to also reset the budgetCalendarAligned form field (or
state) so the toggle returns to its default when budgets are cleared — locate
clearVirtualKeyBudget and call form.setValue("budgetCalendarAligned",
<defaultValue>, { shouldDirty: true }) (or the equivalent state setter)
immediately after clearing "budgets" to ensure new budgets in the same edit
session do not inherit the old alignment.
- Around line 84-91: normalizeProviderConfigs currently converts weight and
rate_limit but forwards config.budgets as form strings, causing
budgets[*].max_limit and budgets[*].reset_duration to be sent as strings; update
normalizeProviderConfigs to map over config.budgets (if present) and coerce
max_limit and reset_duration to numbers (e.g., parseInt/Number or null when
empty) before returning the normalized config so the outgoing JSON matches the
Go handler's expected numeric types; reference the normalizeProviderConfigs
function and the config.budgets/budgets[*].max_limit and
budgets[*].reset_duration fields when making the change.
- Around line 905-937: Add a data-testid to the MultiBudgetLines usages and
ensure the MultiBudgetLines component forwards arbitrary DOM props (including
data-testid) to its root element; specifically, update the two MultiBudgetLines
instances (currently using id={`providerBudget-${index}`}) to include a
data-testid following the pattern
data-testid="provider-budget-editor-<qualifier>" (e.g., including the index)
and, if MultiBudgetLines does not already accept/rest props, modify its props
signature to accept ...rest and spread those onto the top-level DOM node so
data-testid is actually rendered; keep existing behavior in
handleUpdateProviderConfig unchanged.
---
Nitpick comments:
In `@plugins/governance/store.go`:
- Around line 1475-1499: The ResetExpiredBudgetsInMemory function currently
scans gs.virtualKeys for every budget; instead, before iterating budgets build
two lookup maps once per call — e.g. map[VirtualKeyID]bool and
map[ProviderConfigID]bool — by ranging gs.virtualKeys and inspecting each
*configstoreTables.TableVirtualKey and its ProviderConfigs to set
CalendarAligned, then replace the inner gs.virtualKeys.Range scans with O(1)
lookups of those maps when resolving calendarAligned for a budget (use
budget.VirtualKeyID and budget.ProviderConfigID as keys); this keeps the
algorithm linear and preserves the existing TableVirtualKey/ProviderConfigs
logic.
🪄 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: d85973c7-5dfe-4fbb-b0bc-21fc0ce4645c
📒 Files selected for processing (28)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/ui/multiBudgetLines.tsxui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (4)
- .claude/skills/expect
- ui/app/workspace/governance/access-profiles/page.tsx
- .agents/skills/expect/SKILL.md
- .github/workflows/scripts/run-migration-tests.sh
🚧 Files skipped from review as they are similar to previous changes (12)
- plugins/governance/http_transport_prehook_test.go
- plugins/governance/model_provider_governance_test.go
- plugins/governance/resolver.go
- framework/configstore/rdb_test.go
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- plugins/governance/resolver_test.go
- transports/config.schema.json
- plugins/governance/test_utils.go
- transports/bifrost-http/handlers/governance_test.go
- ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
- plugins/governance/store_test.go
- transports/bifrost-http/lib/config_test.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
a2077df to
a4f5278
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
framework/configstore/rdb.go (1)
2318-2350:⚠️ Potential issue | 🔴 CriticalAdd join table cleanup in
DeleteVirtualKeyProviderConfig.The join table
governance_virtual_key_provider_config_keyshas no FK constraint tags defined on its columns, so database-level CASCADE does not apply. When this method deletes a provider config directly (not viaDeleteVirtualKey), orphaned join table entries remain. Add explicit cleanup before deleting the provider config:Suggested fix
// Delete the keys join table entries if err := txDB.WithContext(ctx).Exec("DELETE FROM governance_virtual_key_provider_config_keys WHERE table_virtual_key_provider_config_id = ?", id).Error; err != nil { return err } // Delete the provider config if err := txDB.WithContext(ctx).Delete(&tables.TableVirtualKeyProviderConfig{}, "id = ?", id).Error; err != nil { return err }This mirrors the cleanup already performed in
DeleteVirtualKey(line 2100-2101).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2318 - 2350, DeleteVirtualKeyProviderConfig currently deletes budgets, the provider config, and rate limits but misses cleaning up the join table governance_virtual_key_provider_config_keys, leaving orphaned entries; before deleting the provider config in DeleteVirtualKeyProviderConfig add an explicit deletion of join rows (same approach used in DeleteVirtualKey) by executing a DELETE FROM governance_virtual_key_provider_config_keys WHERE table_virtual_key_provider_config_id = ? via txDB.WithContext(ctx).Exec and return any error, then proceed to delete the provider config and rate limit as existing.transports/bifrost-http/handlers/governance.go (1)
863-873:⚠️ Potential issue | 🔴 CriticalAdd missing Preload("Budgets") to provider config query during reconciliation.
At line 866, the query fetches existing provider configs without preloading their budgets:
tx.Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs). The reconciliation logic at lines 1009–1012 expectsexisting.Budgetsto be populated, building a map to match budgets byResetDuration. Without the preload, this slice is nil/empty, causing:
- All submitted budgets to be treated as new (creating duplicates)
- The deletion loop at lines 1048–1055 to iterate an empty slice (orphaned budgets remain)
Add
.Preload("Budgets")to match the pattern used inGetVirtualKey()and other read paths:Fix: preload budgets for provider config reconciliation
// Get existing provider configs for comparison var existingConfigs []configstoreTables.TableVirtualKeyProviderConfig - if err := tx.Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs).Error; err != nil { + if err := tx.Where("virtual_key_id = ?", vk.ID).Preload("Budgets").Find(&existingConfigs).Error; err != nil { return err }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 863 - 873, The provider-config reconciliation loads existingConfigs without their Budgets, so existingConfigs[i].Budgets is nil and budget comparisons fail; update the DB query that builds existingConfigs (the tx.Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs) call) to preload budgets (e.g., use tx.Preload("Budgets").Where(...).Find(&existingConfigs)) so that TableVirtualKeyProviderConfig entries have their Budgets populated for the later reconciliation and deletion logic that inspects existing.Budgets and matches by ResetDuration.
♻️ Duplicate comments (2)
framework/configstore/migrations.go (1)
5395-5408:⚠️ Potential issue | 🟠 MajorBackfill
calendar_alignedthrough provider configs too.This update only promotes legacy rows where
governance_budgets.virtual_key_idis set. Provider-config-owned budgets were just migrated intoprovider_config_idabove, so any oldcalendar_aligned=trueon those rows is lost permanently when Line 5408 drops the source column.🛠️ Suggested fix
- if err := tx.Exec(` - UPDATE governance_virtual_keys SET calendar_aligned = true - WHERE id IN ( - SELECT DISTINCT virtual_key_id FROM governance_budgets - WHERE calendar_aligned = true AND virtual_key_id IS NOT NULL - ) AND calendar_aligned = false - `).Error; err != nil { + if err := tx.Exec(` + UPDATE governance_virtual_keys SET calendar_aligned = true + WHERE id IN ( + SELECT DISTINCT b.virtual_key_id + FROM governance_budgets b + WHERE b.calendar_aligned = true AND b.virtual_key_id IS NOT NULL + UNION + SELECT DISTINCT pc.virtual_key_id + FROM governance_budgets b + JOIN governance_virtual_key_provider_configs pc + ON pc.id = b.provider_config_id + WHERE b.calendar_aligned = true AND pc.virtual_key_id IS NOT NULL + ) AND calendar_aligned = false + `).Error; err != nil {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 5395 - 5408, The migration only backfills calendar_aligned for budgets with governance_budgets.virtual_key_id set and then drops the column, losing calendar_aligned=true on budgets owned by provider configs; add a second backfill that finds virtual keys via the provider config relationship and sets governance_virtual_keys.calendar_aligned = true for those VKs where a budget had calendar_aligned = true and the VK field is currently false. Concretely, after the existing UPDATE that uses governance_budgets.virtual_key_id, add an UPDATE that joins governance_budgets (where provider_config_id IS NOT NULL AND calendar_aligned = true) to the provider configs table (the table that links provider_config_id -> virtual_key_id) and updates governance_virtual_keys.calendar_aligned for the resolved virtual_key_id, then only after both backfills drop governance_budgets.calendar_aligned; reference symbols: migrations.go backfill block, tables.TableBudget, governance_budgets.calendar_aligned, governance_budgets.provider_config_id, governance_virtual_keys.calendar_aligned, and the provider-configs table/column that maps provider_config_id -> virtual_key_id.transports/bifrost-http/handlers/governance.go (1)
510-512:⚠️ Potential issue | 🟠 Major
validateBudgeterrors return HTTP 500 instead of 400 for client validation failures.The
validateBudgetfunction returns a plain error, but the error handling at lines 629-636 only mapsbadRequestErrorto HTTP 400. Validation failures (e.g.,max_limit == 0) will surface as 500. Either wrap the error or addmax_limit == 0to the pre-validation.Also applies to: lines 591-593, 770-772, 788-789, 948-950, 1021-1023, 1039-1041
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 510 - 512, validateBudget returns generic errors which are being treated as internal (500) by the HTTP handlers; wrap or convert validation failures into a badRequestError so they map to HTTP 400. Update validateBudget to return a typed error (e.g., badRequestError) for client validation problems (like max_limit == 0) or, where validateBudget is called (the handlers invoking validateBudget), detect validation failures and wrap them with badRequestError before returning; ensure all call sites of validateBudget (the budget-related handlers that call validateBudget) consistently return badRequestError on validation failure so the existing error-to-status mapping will produce HTTP 400.
🧹 Nitpick comments (3)
plugins/governance/test_utils.go (1)
210-218: Consider adding a default provider config for consistency.
buildVirtualKeyWithBudgetadds a default"openai"provider config (lines 96-98) so the resolver doesn't block at provider check, butbuildVirtualKeyWithMultiBudgetsomits this. If tests usingbuildVirtualKeyWithMultiBudgetsalso require the resolver to pass provider checks, they may fail unexpectedly.♻️ Suggested fix
func buildVirtualKeyWithMultiBudgets(id, value, name string, budgets []configstoreTables.TableBudget) *configstoreTables.TableVirtualKey { vk := buildVirtualKey(id, value, name, true) for i := range budgets { vkID := id budgets[i].VirtualKeyID = &vkID } vk.Budgets = budgets + // Add a default provider config so the resolver doesn't block at provider check + vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{ + buildProviderConfig("openai", []string{"*"}), + } return vk }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 210 - 218, buildVirtualKeyWithMultiBudgets misses adding the default provider config that buildVirtualKeyWithBudget sets, which can cause resolver provider checks to block in tests; update buildVirtualKeyWithMultiBudgets to add the same default provider config (e.g., the "openai" provider entry added by buildVirtualKeyWithBudget) to the returned VirtualKey before returning, and ensure VirtualKey.Budgets is still set and each TableBudget.VirtualKeyID is assigned just as currently done so tests behave consistently with buildVirtualKeyWithBudget.transports/bifrost-http/handlers/governance.go (2)
118-123:CreateBudgetRequest.CalendarAlignedfield appears unused.The
CalendarAlignedfield onCreateBudgetRequest(line 122) is no longer used—budgetLastResetalways reads from the VK-levelCalendarAlignedinstead. Consider removing this field to avoid confusing API consumers, or update the comment to note it's deprecated/ignored.♻️ Proposed cleanup
// CreateBudgetRequest represents the request body for creating a budget type CreateBudgetRequest struct { MaxLimit float64 `json:"max_limit" validate:"required"` // Maximum budget in dollars ResetDuration string `json:"reset_duration" validate:"required"` // e.g., "30s", "5m", "1h", "1d", "1w", "1M" - CalendarAligned bool `json:"calendar_aligned,omitempty"` // Snap resets to calendar boundaries (day/week/month/year) + // Deprecated: CalendarAligned is now controlled at the VirtualKey level; this field is ignored. + CalendarAligned bool `json:"calendar_aligned,omitempty"` }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 118 - 123, The CreateBudgetRequest.CalendarAligned field is unused because budgetLastReset always reads the VK-level CalendarAligned; either remove CalendarAligned from the CreateBudgetRequest struct or mark it deprecated/ignored in its comment and API docs. Locate the CreateBudgetRequest type and update its definition: if removing, delete the CalendarAligned field and any JSON tag; if keeping for backward compatibility, change the comment to “deprecated: ignored—VK-level CalendarAligned is used” and ensure no code reads CreateBudgetRequest.CalendarAligned anywhere (verify usages tied to budgetLastReset). Also update any API docs or tests that reference CreateBudgetRequest.CalendarAligned.
737-739: TogglingCalendarAligneddoes not recalculateLastResetfor existing budgets.When
CalendarAlignedchanges (especially false → true), existing matched budgets retain their original non-alignedLastResetvalues, while any new budgets in the same request get calendar-aligned values. This creates inconsistent reset timing within the same virtual key.If preserving usage is intentional, consider adding a comment documenting this behavior. Otherwise, consider recalculating
LastResetfor existing budgets whenCalendarAlignedis enabled:💡 Optional fix to snap existing budgets to calendar boundaries
if existing, found := existingByDuration[b.ResetDuration]; found { // Budget with same duration exists — update max_limit, preserve usage existing.MaxLimit = b.MaxLimit + // If calendar alignment was just enabled, snap LastReset to calendar boundary + if req.CalendarAligned != nil && *req.CalendarAligned && !vk.CalendarAligned { + existing.LastReset = budgetLastReset(true, existing.ResetDuration) + existing.CurrentUsage = 0 + } if err := validateBudget(&existing); err != nil {Also applies to: 767-777
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 737 - 739, The code currently sets vk.CalendarAligned from req.CalendarAligned but does not adjust existing budgets' reset timestamps, causing mixed LastReset values; when req.CalendarAligned changes to true you should iterate the existing budgets on the virtual key (e.g. vk.Budgets / each Budget in vk) and recalculate Budget.LastReset to the appropriate calendar-aligned boundary (use or add a helper like computeAlignedLastReset(now, budget.Window/Interval) to determine the start of the current calendar window) so all budgets share aligned LastReset; if you intend to preserve legacy timestamps instead, add a clear comment next to the vk.CalendarAligned assignment (and the similar block at the other location) documenting that existing Budget.LastReset values are intentionally left unchanged.
🤖 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/tables/virtualkey.go`:
- Around line 37-38: TableVirtualKeyProviderConfig and TableVirtualKey lost
backwards compatibility when the singular Budget field and legacy
budget.calendar_aligned were removed; add a temporary JSON/unmarshal
compatibility shim that for one release checks for the old singular "budget"
field and legacy "calendar_aligned" and, if present and Budgets is empty, lifts
that singular budget into Budgets[0] and applies CalendarAligned into that
TableBudget instance. Implement the shim as a custom UnmarshalJSON (or
equivalent decode path) on TableVirtualKeyProviderConfig and TableVirtualKey
that merges old payloads into the new Budgets []TableBudget shape, and mirror
the same logic for the TableBudget calendar_aligned migration (see the
TableBudget type and the related code around the other occurrence lines
referenced 222-228) so older payloads deserialize correctly.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 932-955: Add a pre-validation inside the pc.Budgets loop (where
duplicate reset_duration is checked) to reject budgets with non-positive
max_limit before calling validateBudget/CreateBudget: inspect b.MaxLimit and if
it's <= 0 return a badRequestError (similar to the duplicate-reset-duration
case) with a clear message like "invalid max_limit in provider config budgets:
%v"; keep using the same loop and error type so client validation errors
consistently return 400 (references: pc.Budgets, b.MaxLimit, validateBudget,
badRequestError, h.configStore.CreateBudget).
In `@transports/bifrost-http/lib/config_test.go`:
- Line 243: Add an end-to-end SQLite integration test that exercises
provider-config budget persistence and reload/reconciliation: replace or
complement heavy MockConfigStore logic with simple mocks and use
createTestSQLiteConfigStore to persist a VirtualKey with ProviderConfigs that
include Tables.TableVirtualKeyProviderConfig.Budgets, then assert that
GetVirtualKey preloads ProviderConfigs.Budgets and returns the written budgets;
next mutate the underlying config (simulate file change), run the
reload/reconcile path, and assert budgets are reconciled/removed as expected.
Keep MockConfigStore methods minimal and focused, reuse
createTestSQLiteConfigStore for write/read round-trips, and ensure the test
covers write->read->reload->delete behavior to prevent regressions in the
multi-budget code paths.
---
Outside diff comments:
In `@framework/configstore/rdb.go`:
- Around line 2318-2350: DeleteVirtualKeyProviderConfig currently deletes
budgets, the provider config, and rate limits but misses cleaning up the join
table governance_virtual_key_provider_config_keys, leaving orphaned entries;
before deleting the provider config in DeleteVirtualKeyProviderConfig add an
explicit deletion of join rows (same approach used in DeleteVirtualKey) by
executing a DELETE FROM governance_virtual_key_provider_config_keys WHERE
table_virtual_key_provider_config_id = ? via txDB.WithContext(ctx).Exec and
return any error, then proceed to delete the provider config and rate limit as
existing.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 863-873: The provider-config reconciliation loads existingConfigs
without their Budgets, so existingConfigs[i].Budgets is nil and budget
comparisons fail; update the DB query that builds existingConfigs (the
tx.Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs) call) to preload
budgets (e.g., use tx.Preload("Budgets").Where(...).Find(&existingConfigs)) so
that TableVirtualKeyProviderConfig entries have their Budgets populated for the
later reconciliation and deletion logic that inspects existing.Budgets and
matches by ResetDuration.
---
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 5395-5408: The migration only backfills calendar_aligned for
budgets with governance_budgets.virtual_key_id set and then drops the column,
losing calendar_aligned=true on budgets owned by provider configs; add a second
backfill that finds virtual keys via the provider config relationship and sets
governance_virtual_keys.calendar_aligned = true for those VKs where a budget had
calendar_aligned = true and the VK field is currently false. Concretely, after
the existing UPDATE that uses governance_budgets.virtual_key_id, add an UPDATE
that joins governance_budgets (where provider_config_id IS NOT NULL AND
calendar_aligned = true) to the provider configs table (the table that links
provider_config_id -> virtual_key_id) and updates
governance_virtual_keys.calendar_aligned for the resolved virtual_key_id, then
only after both backfills drop governance_budgets.calendar_aligned; reference
symbols: migrations.go backfill block, tables.TableBudget,
governance_budgets.calendar_aligned, governance_budgets.provider_config_id,
governance_virtual_keys.calendar_aligned, and the provider-configs table/column
that maps provider_config_id -> virtual_key_id.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 510-512: validateBudget returns generic errors which are being
treated as internal (500) by the HTTP handlers; wrap or convert validation
failures into a badRequestError so they map to HTTP 400. Update validateBudget
to return a typed error (e.g., badRequestError) for client validation problems
(like max_limit == 0) or, where validateBudget is called (the handlers invoking
validateBudget), detect validation failures and wrap them with badRequestError
before returning; ensure all call sites of validateBudget (the budget-related
handlers that call validateBudget) consistently return badRequestError on
validation failure so the existing error-to-status mapping will produce HTTP
400.
---
Nitpick comments:
In `@plugins/governance/test_utils.go`:
- Around line 210-218: buildVirtualKeyWithMultiBudgets misses adding the default
provider config that buildVirtualKeyWithBudget sets, which can cause resolver
provider checks to block in tests; update buildVirtualKeyWithMultiBudgets to add
the same default provider config (e.g., the "openai" provider entry added by
buildVirtualKeyWithBudget) to the returned VirtualKey before returning, and
ensure VirtualKey.Budgets is still set and each TableBudget.VirtualKeyID is
assigned just as currently done so tests behave consistently with
buildVirtualKeyWithBudget.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 118-123: The CreateBudgetRequest.CalendarAligned field is unused
because budgetLastReset always reads the VK-level CalendarAligned; either remove
CalendarAligned from the CreateBudgetRequest struct or mark it
deprecated/ignored in its comment and API docs. Locate the CreateBudgetRequest
type and update its definition: if removing, delete the CalendarAligned field
and any JSON tag; if keeping for backward compatibility, change the comment to
“deprecated: ignored—VK-level CalendarAligned is used” and ensure no code reads
CreateBudgetRequest.CalendarAligned anywhere (verify usages tied to
budgetLastReset). Also update any API docs or tests that reference
CreateBudgetRequest.CalendarAligned.
- Around line 737-739: The code currently sets vk.CalendarAligned from
req.CalendarAligned but does not adjust existing budgets' reset timestamps,
causing mixed LastReset values; when req.CalendarAligned changes to true you
should iterate the existing budgets on the virtual key (e.g. vk.Budgets / each
Budget in vk) and recalculate Budget.LastReset to the appropriate
calendar-aligned boundary (use or add a helper like computeAlignedLastReset(now,
budget.Window/Interval) to determine the start of the current calendar window)
so all budgets share aligned LastReset; if you intend to preserve legacy
timestamps instead, add a clear comment next to the vk.CalendarAligned
assignment (and the similar block at the other location) documenting that
existing Budget.LastReset values are intentionally left unchanged.
🪄 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: 863575c2-1266-4317-b5f2-f9c181c5c013
📒 Files selected for processing (28)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/ui/multiBudgetLines.tsxui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- framework/configstore/clientconfig.go
- transports/bifrost-http/lib/config.go
✅ Files skipped from review due to trivial changes (5)
- .claude/skills/expect
- plugins/governance/resolver.go
- transports/bifrost-http/handlers/governance_test.go
- .agents/skills/expect/SKILL.md
- ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
🚧 Files skipped from review as they are similar to previous changes (10)
- plugins/governance/http_transport_prehook_test.go
- ui/app/workspace/governance/access-profiles/page.tsx
- plugins/governance/model_provider_governance_test.go
- plugins/governance/resolver_test.go
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- framework/configstore/tables/budget.go
- ui/components/ui/multiBudgetLines.tsx
- framework/configstore/rdb_test.go
- .github/workflows/scripts/run-migration-tests.sh
- plugins/governance/store.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
a4f5278 to
413d855
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (4)
framework/configstore/migrations.go (1)
5395-5406:⚠️ Potential issue | 🟠 MajorBackfill misses calendar-aligned budgets owned via provider configs.
On Lines 5399-5403, the VK backfill only reads
governance_budgets.virtual_key_id. Legacy provider-config budgets (governance_virtual_key_provider_configs.budget_id) can havecalendar_aligned = truewithoutvirtual_key_idpopulated, so their parent virtual keys stayfalseafter migration.💡 Suggested fix
- if err := tx.Exec(` - UPDATE governance_virtual_keys SET calendar_aligned = true - WHERE id IN ( - SELECT DISTINCT virtual_key_id FROM governance_budgets - WHERE calendar_aligned = true AND virtual_key_id IS NOT NULL - ) AND calendar_aligned = false - `).Error; err != nil { + if err := tx.Exec(` + UPDATE governance_virtual_keys SET calendar_aligned = true + WHERE id IN ( + SELECT DISTINCT COALESCE(b.virtual_key_id, pc.virtual_key_id) AS vk_id + FROM governance_budgets b + LEFT JOIN governance_virtual_key_provider_configs pc + ON pc.id = b.provider_config_id + WHERE b.calendar_aligned = true + AND COALESCE(b.virtual_key_id, pc.virtual_key_id) IS NOT NULL + ) AND calendar_aligned = false + `).Error; err != nil { return fmt.Errorf("failed to backfill calendar_aligned from budgets to virtual keys: %w", err) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 5395 - 5406, The backfill only updates governance_virtual_keys based on governance_budgets.virtual_key_id and therefore misses budgets referenced from governance_virtual_key_provider_configs.budget_id; update the migration in the block that checks mg.HasColumn(&tables.TableBudget{}, "calendar_aligned") to also mark VKs whose IDs are parents of budgets referenced by governance_virtual_key_provider_configs where calendar_aligned = true — e.g., run an additional UPDATE (or extend the WHERE ... IN subquery) that selects DISTINCT provider_config.virtual_key_id FROM governance_virtual_key_provider_configs JOIN governance_budgets ON governance_budgets.id = governance_virtual_key_provider_configs.budget_id WHERE governance_budgets.calendar_aligned = true and provider_config.virtual_key_id IS NOT NULL, ensuring governance_virtual_keys.calendar_aligned is set to true for those VKs as well.transports/bifrost-http/handlers/governance.go (2)
863-868:⚠️ Potential issue | 🔴 CriticalPreload provider-config budgets before reconciling them.
existingConfigsis loaded withoutPreload("Budgets"), soexisting.Budgetsis empty here. Every PUT with provider budgets then treats them as new, creates fresh rows withCurrentUsage: 0, and never deletes the old ones.GetVirtualKeyalready preloadsProviderConfigs.Budgetsinframework/configstore/rdb.go:1922-1945; this tx-local requery drops that association.🛠️ Minimal fix
- if err := tx.Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs).Error; err != nil { + if err := tx.Where("virtual_key_id = ?", vk.ID). + Preload("Budgets"). + Find(&existingConfigs).Error; err != nil { return err }Based on learnings,
updateVirtualKey (PUT /api/governance/virtual-keys/{vk_id})expects a full payload from the frontend.Also applies to: 991-1057
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 863 - 868, The re-query that populates existingConfigs in the updateVirtualKey flow is missing Preload("Budgets"), so existingConfigs' Budgets slices are empty and budgets are recreated; change the tx.Where(...).Find(&existingConfigs) call to preload budgets (e.g., tx.Preload("Budgets").Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs).Error) so the reconciler sees existing ProviderConfig.Budgets and can update/delete them correctly; mirror the preload used by GetVirtualKey (ProviderConfigs.Budgets) for both places noted (around existingConfigs and the similar block at lines ~991-1057).
442-459:⚠️ Potential issue | 🟠 MajorReturn 400 for all invalid budget payloads.
These paths still only classify some client errors up front.
max_limit == 0at VK/provider-config level, and malformedreset_durationin the new provider-config budget paths, still fall through tovalidateBudget, which bubbles out as a generic error and turns a bad request into HTTP 500. WrapvalidateBudgethere inbadRequestError, or tighten every loop to precheck<= 0plusParseDuration(...)before creating/updating the budget.Also applies to: 575-597, 740-755, 932-955, 991-1005
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 442 - 459, The code currently lets some invalid budget payloads (e.g., MaxLimit == 0 or malformed ResetDuration) reach validateBudget and become 500s; update the pre-validation logic in the request handlers that iterate req.Budgets (the block using seenDurations, b.MaxLimit and configstoreTables.ParseDuration) to treat MaxLimit <= 0 as a 400 and to ParseDuration(b.ResetDuration) before creating/updating budgets, or alternatively wrap any call to validateBudget with badRequestError so validation failures are returned as HTTP 400; ensure references are to req.Budgets, b.MaxLimit, b.ResetDuration, validateBudget and badRequestError and use SendError(…,400,…) for these client errors.framework/configstore/tables/virtualkey.go (1)
37-38:⚠️ Potential issue | 🟠 MajorKeep a legacy single-budget decoder for one release.
TableVirtualKeyProviderConfigis still a config/JSON boundary type — it already has a customUnmarshalJSONforkey_ids. Removing the singularbudgetfield here and onTableVirtualKeymeans older single-budget payloads now deserialize to zero budgets, and legacybudget.calendar_alignedhas no path toTableVirtualKey.CalendarAligned. That silently drops governance settings on config-file upgrades. Add a temporary compatibility shim that liftsbudgetintoBudgets[0]and copies the old alignment flag during unmarshal.Also applies to: 222-228
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` around lines 37 - 38, Add a one-release compatibility shim in the custom UnmarshalJSON for TableVirtualKeyProviderConfig (and similarly for TableVirtualKey) that detects the legacy singular "budget" payload and, if present and Budgets is empty, creates Budgets with that single budget as Budgets[0]; also copy the legacy budget.calendar_aligned (or legacy field) into the new TableVirtualKey.CalendarAligned field when unmarshalling so the old alignment flag is preserved. Implement this logic inside the existing UnmarshalJSON function(s) that already handle "key_ids", ensure it only runs when Budgets is empty and the legacy "budget" key is present, and mark the shim as temporary to be removed after one release.
🧹 Nitpick comments (4)
ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx (1)
175-178: Use the budget row as the first source for the(calendar)suffix.These rows are rendered per budget, but they still read
virtualKey.calendar_aligned.virtualKeysTable.tsxalready readsb.calendar_aligned, so the details sheet can drift if that flag is only hydrated on the budget objects.b.calendar_aligned ?? virtualKey.calendar_alignedwould keep both shapes working during the rollout.♻️ Suggested change
-<div className="col-span-2 text-sm">{parseResetPeriod(b.reset_duration)}{virtualKey.calendar_aligned && " (calendar)"}</div> +<div className="col-span-2 text-sm"> + {parseResetPeriod(b.reset_duration)} + {(b.calendar_aligned ?? virtualKey.calendar_aligned) && " (calendar)"} +</div>Apply the same change in both budget sections.
Also applies to: 370-373
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx` around lines 175 - 178, The Reset Period row uses virtualKey.calendar_aligned but should prefer the budget's flag to avoid drift; update the JSX in virtualKeyDetailsSheet.tsx where you render reset period for each budget (the div showing parseResetPeriod(b.reset_duration) and the calendar suffix) to use b.calendar_aligned ?? virtualKey.calendar_aligned instead of virtualKey.calendar_aligned, and apply the same change in the other budget section rendering (the other Reset Period block around the second budget render).plugins/governance/test_utils.go (1)
92-93: Usebifrost.Ptr(id)here for consistency.These are the new pointer-assignment sites in this file, so it's worth keeping them aligned with the repo's standard helper instead of taking the address of a temporary local.
Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability.
Also applies to: 213-214
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 92 - 93, Replace the local temporary address-take for VirtualKeyID with the repo helper: instead of creating vkID := id and assigning &vkID to budget.VirtualKeyID, use bifrost.Ptr(id) to produce the pointer (update the assignment site involving vkID and budget.VirtualKeyID). Also scan the same file for other places that take addresses of temporaries (the similar sites noted in the review) and switch them to bifrost.Ptr(...) for consistency.plugins/governance/store_test.go (2)
553-592: Cover provider-config budgets in the create/delete test too.
CreateVirtualKeyInMemoryandDeleteVirtualKeyInMemoryboth have separatepc.Budgetsloops, but this test only asserts the top-levelvk.Budgetspath. A regression in provider-config budget insertion or cleanup would still pass here.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 553 - 592, The test currently only checks vk.Budgets but must also cover provider-config budgets: after calling CreateVirtualKeyInMemory(vk) assert that budgets referenced from each vk.ProviderConfigs[*].Budgets are present in store.budgets and that the retrieved virtual key (GetVirtualKey("sk-bf-test")) contains the provider configs with their Budgets preserved; after calling DeleteVirtualKeyInMemory("vk1") assert that those same provider-config budget keys are removed from store.budgets and that GetVirtualKey no longer finds the key. Update the test to use the existing buildProviderConfig/buildVirtualKeyWithMultiBudgets helpers to include provider-config budgets and add presence checks for pc.Budgets before delete and absence checks after delete, referencing CreateVirtualKeyInMemory, DeleteVirtualKeyInMemory, vk.ProviderConfigs and pc.Budgets to locate the code to change.
511-550: This doesn't actually test theCalendarAlignedreset behavior.The assertions only prove that the flag round-trips onto the VK and that
CheckBudgetstill passes under limit. The new behavior lives inResetExpiredBudgetsInMemory, so this would still pass if the calendar-period reset logic regressed. Please driveLastResetacross a calendar boundary and assert the reset result directly; ideally add a provider-config budget case too, since that branch resolvesCalendarAlignedthroughProviderConfigID.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 511 - 550, Test currently only verifies the CalendarAligned flag round-trips and CheckBudget passes, but doesn't exercise the calendar-boundary reset logic in ResetExpiredBudgetsInMemory; update TestGovernanceStore_MultiBudget_CalendarAligned to drive LastReset across a calendar boundary and assert budgets are reset accordingly. Specifically, set a budget's LastReset to a timestamp just before the previous calendar period (e.g., yesterday for daily or previous month for monthly), call store.ResetExpiredBudgetsInMemory(context.Background()), then fetch the virtual key via store.GetVirtualKey("sk-bf-test") or the budget entries and assert CurrentUsage was zeroed (or reset as expected). Also add a case where the budget is resolved via a ProviderConfig (use ProviderConfigID in buildVirtualKeyWithMultiBudgets / buildProviderConfig) to verify ProviderConfig-level CalendarAligned resolution triggers the same reset behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2655-2731: The SQLite migration lane is missing the
budget-migration verification: add the same checks currently in
verify_budget_migration_postgres() (column existence, budget backfill for
budget-migration-test-1/2, dropped budget_id columns, and junction table
absence) to the SQLite path or invoke that verification from the SQLite flow;
update run_sqlite_migration_tests() (or validate_sqlite_data()) to run the
budget migration assertions (or factor out a shared verify_budget_migration()
used by both run_postgres_migration_tests() and the SQLite lane) so
virtual_key_id/provider_config_id backfill and legacy budget_id cleanup are
validated for SQLite as well.
---
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 5395-5406: The backfill only updates governance_virtual_keys based
on governance_budgets.virtual_key_id and therefore misses budgets referenced
from governance_virtual_key_provider_configs.budget_id; update the migration in
the block that checks mg.HasColumn(&tables.TableBudget{}, "calendar_aligned") to
also mark VKs whose IDs are parents of budgets referenced by
governance_virtual_key_provider_configs where calendar_aligned = true — e.g.,
run an additional UPDATE (or extend the WHERE ... IN subquery) that selects
DISTINCT provider_config.virtual_key_id FROM
governance_virtual_key_provider_configs JOIN governance_budgets ON
governance_budgets.id = governance_virtual_key_provider_configs.budget_id WHERE
governance_budgets.calendar_aligned = true and provider_config.virtual_key_id IS
NOT NULL, ensuring governance_virtual_keys.calendar_aligned is set to true for
those VKs as well.
In `@framework/configstore/tables/virtualkey.go`:
- Around line 37-38: Add a one-release compatibility shim in the custom
UnmarshalJSON for TableVirtualKeyProviderConfig (and similarly for
TableVirtualKey) that detects the legacy singular "budget" payload and, if
present and Budgets is empty, creates Budgets with that single budget as
Budgets[0]; also copy the legacy budget.calendar_aligned (or legacy field) into
the new TableVirtualKey.CalendarAligned field when unmarshalling so the old
alignment flag is preserved. Implement this logic inside the existing
UnmarshalJSON function(s) that already handle "key_ids", ensure it only runs
when Budgets is empty and the legacy "budget" key is present, and mark the shim
as temporary to be removed after one release.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 863-868: The re-query that populates existingConfigs in the
updateVirtualKey flow is missing Preload("Budgets"), so existingConfigs' Budgets
slices are empty and budgets are recreated; change the
tx.Where(...).Find(&existingConfigs) call to preload budgets (e.g.,
tx.Preload("Budgets").Where("virtual_key_id = ?",
vk.ID).Find(&existingConfigs).Error) so the reconciler sees existing
ProviderConfig.Budgets and can update/delete them correctly; mirror the preload
used by GetVirtualKey (ProviderConfigs.Budgets) for both places noted (around
existingConfigs and the similar block at lines ~991-1057).
- Around line 442-459: The code currently lets some invalid budget payloads
(e.g., MaxLimit == 0 or malformed ResetDuration) reach validateBudget and become
500s; update the pre-validation logic in the request handlers that iterate
req.Budgets (the block using seenDurations, b.MaxLimit and
configstoreTables.ParseDuration) to treat MaxLimit <= 0 as a 400 and to
ParseDuration(b.ResetDuration) before creating/updating budgets, or
alternatively wrap any call to validateBudget with badRequestError so validation
failures are returned as HTTP 400; ensure references are to req.Budgets,
b.MaxLimit, b.ResetDuration, validateBudget and badRequestError and use
SendError(…,400,…) for these client errors.
---
Nitpick comments:
In `@plugins/governance/store_test.go`:
- Around line 553-592: The test currently only checks vk.Budgets but must also
cover provider-config budgets: after calling CreateVirtualKeyInMemory(vk) assert
that budgets referenced from each vk.ProviderConfigs[*].Budgets are present in
store.budgets and that the retrieved virtual key (GetVirtualKey("sk-bf-test"))
contains the provider configs with their Budgets preserved; after calling
DeleteVirtualKeyInMemory("vk1") assert that those same provider-config budget
keys are removed from store.budgets and that GetVirtualKey no longer finds the
key. Update the test to use the existing
buildProviderConfig/buildVirtualKeyWithMultiBudgets helpers to include
provider-config budgets and add presence checks for pc.Budgets before delete and
absence checks after delete, referencing CreateVirtualKeyInMemory,
DeleteVirtualKeyInMemory, vk.ProviderConfigs and pc.Budgets to locate the code
to change.
- Around line 511-550: Test currently only verifies the CalendarAligned flag
round-trips and CheckBudget passes, but doesn't exercise the calendar-boundary
reset logic in ResetExpiredBudgetsInMemory; update
TestGovernanceStore_MultiBudget_CalendarAligned to drive LastReset across a
calendar boundary and assert budgets are reset accordingly. Specifically, set a
budget's LastReset to a timestamp just before the previous calendar period
(e.g., yesterday for daily or previous month for monthly), call
store.ResetExpiredBudgetsInMemory(context.Background()), then fetch the virtual
key via store.GetVirtualKey("sk-bf-test") or the budget entries and assert
CurrentUsage was zeroed (or reset as expected). Also add a case where the budget
is resolved via a ProviderConfig (use ProviderConfigID in
buildVirtualKeyWithMultiBudgets / buildProviderConfig) to verify
ProviderConfig-level CalendarAligned resolution triggers the same reset
behavior.
In `@plugins/governance/test_utils.go`:
- Around line 92-93: Replace the local temporary address-take for VirtualKeyID
with the repo helper: instead of creating vkID := id and assigning &vkID to
budget.VirtualKeyID, use bifrost.Ptr(id) to produce the pointer (update the
assignment site involving vkID and budget.VirtualKeyID). Also scan the same file
for other places that take addresses of temporaries (the similar sites noted in
the review) and switch them to bifrost.Ptr(...) for consistency.
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 175-178: The Reset Period row uses virtualKey.calendar_aligned but
should prefer the budget's flag to avoid drift; update the JSX in
virtualKeyDetailsSheet.tsx where you render reset period for each budget (the
div showing parseResetPeriod(b.reset_duration) and the calendar suffix) to use
b.calendar_aligned ?? virtualKey.calendar_aligned instead of
virtualKey.calendar_aligned, and apply the same change in the other budget
section rendering (the other Reset Period block around the second budget
render).
🪄 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: 68294cc9-915b-4ec0-87cc-14c7bf0763f8
📒 Files selected for processing (31)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/sidebar.tsxui/components/ui/multiBudgetLines.tsxui/lib/store/apis/baseApi.tsui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (8)
- .claude/skills/expect
- transports/bifrost-http/handlers/governance_test.go
- ui/lib/store/apis/baseApi.ts
- ui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsx
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
- .agents/skills/expect/SKILL.md
- transports/config.schema.json
🚧 Files skipped from review as they are similar to previous changes (12)
- plugins/governance/resolver.go
- plugins/governance/resolver_test.go
- plugins/governance/model_provider_governance_test.go
- plugins/governance/http_transport_prehook_test.go
- ui/app/workspace/governance/access-profiles/page.tsx
- framework/configstore/rdb_test.go
- framework/configstore/tables/budget.go
- ui/lib/types/governance.ts
- ui/components/ui/multiBudgetLines.tsx
- framework/configstore/rdb.go
- ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
- plugins/governance/store.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
413d855 to
1994f2a
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/governance/store.go (1)
2080-2095:⚠️ Potential issue | 🟠 MajorAdd budget hydration for provider configs in config-memory mode.
Provider configs can have multiple budgets (with different reset intervals per
TableBudget.ProviderConfigID), but theloadFromConfigMemoryfunction only hydrates provider-config rate limits. Budgets must also be populated by matchingTableBudget.ProviderConfigIDwith each provider config, matching the pattern used for model configs and providers. In config-memory mode, budget checks can silently miss provider-config limits without this.Lines 2080–2095
// Populate provider config relationships with rate limits if vk.ProviderConfigs != nil { for j := range vk.ProviderConfigs { pc := &vk.ProviderConfigs[j] // Populate rate limit if pc.RateLimitID != nil { for k := range rateLimits { if rateLimits[k].ID == *pc.RateLimitID { pc.RateLimit = &rateLimits[k] break } } } } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2080 - 2095, In loadFromConfigMemory, after the existing ProviderConfigs loop that hydrates pc.RateLimit, also populate provider-config budgets by iterating the in-memory TableBudget slice (e.g., variable named budgets or tableBudgets) and for each pc (ProviderConfigs[j]) append or set any TableBudget whose ProviderConfigID matches pc.ID (same pattern used for model configs/providers). Ensure you check vk.ProviderConfigs != nil and the budgets slice != nil, and attach budget entries to the ProviderConfig (e.g., pc.Budgets or equivalent field) rather than replacing existing fields so multiple budgets per provider config are preserved.
♻️ Duplicate comments (5)
.github/workflows/scripts/run-migration-tests.sh (1)
2982-2987:⚠️ Potential issue | 🟠 MajorAdd the same budget-migration assertions to the SQLite lane.
This new step only runs in
run_postgres_migration_tests().run_sqlite_migration_tests()still stops after the coarse row-count checks invalidate_sqlite_data(), so a SQLite regression invirtual_key_id/provider_config_idbackfill orbudget_idcleanup can pass unnoticed.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/scripts/run-migration-tests.sh around lines 2982 - 2987, The SQLite migration lane is missing the detailed budget backfill/cleanup assertions that exist in run_postgres_migration_tests; add the same verification step (call the budget verification helper used by Postgres or create a verify_budget_migration_sqlite wrapper) into run_sqlite_migration_tests immediately after validate_sqlite_data (and before stopping the test harness), and on failure log the same error, call stop_bifrost, and return a non-zero exit so SQLite regressions in virtual_key_id/provider_config_id or budget_id cleanup are caught; reference verify_budget_migration_postgres, run_postgres_migration_tests, run_sqlite_migration_tests, validate_sqlite_data, and stop_bifrost to locate where to insert the call.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (2)
291-294:⚠️ Potential issue | 🟠 MajorTreat provider budgets as part of VK-wide
calendar_aligned.The toggle and both submit paths still only look at top-level
budgets. A virtual key that only uses provider-config budgets can't render, change, or persistcalendar_aligned, even though this flag now lives on the virtual key and applies to all budgets under that key.🧭 Fold provider budgets into the VK-wide alignment logic
-const hasAnyAlignableBudget = - watchedBudgets && watchedBudgets.length > 0 && watchedBudgets.some((b) => b.max_limit && supportsCalendarAlignment(b.reset_duration || "1M")); +const allBudgetLines = [ + ...(watchedBudgets ?? []), + ...providerConfigs.flatMap((config) => config.budgets ?? []), +]; +const hasAnyAlignableBudget = allBudgetLines.some( + (b) => b.max_limit && supportsCalendarAlignment(b.reset_duration || "1M"), +);Also gate
createData.calendar_aligned/updateData.calendar_alignedon any VK-owned budget being present, not justvalidBudgets.length > 0.Also applies to: 443-455, 489-497, 1213-1230
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 291 - 294, The calendar alignment toggle and submit logic only consider top-level budgets and ignore provider-config (watched) budgets, so update the alignment checks to include provider budgets: change the hasAnyAlignableBudget computation to fold provider budgets (watchedBudgets) into the same logic used for top-level budgets (use both budgets and watchedBudgets when checking max_limit and supportsCalendarAlignment), and ensure createData.calendar_aligned and updateData.calendar_aligned are set/gated when any VK-owned budget exists (not just when validBudgets.length > 0) by treating provider budgets as VK-owned in the conditions that build createData/updateData; update references to watchedBudgets, budgets, hasAnyAlignableBudget, createData.calendar_aligned, updateData.calendar_aligned, and validBudgets accordingly.
378-410:⚠️ Potential issue | 🔴 CriticalNormalize provider budgets before building the payload.
normalizeProviderConfigs()still spreadsconfig.budgetsback into the request unchanged. Those rows are stored in form state as strings, so any save on a key that already has provider budgets re-posts"max_limit": "..."instead of a number, and blank rows are never stripped. Theany[]return type hides the mismatch from TypeScript, but the API still gets the wrong shape.🔧 Normalize the provider-budget rows here
const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => { return configs.map((config) => ({ ...config, + budgets: + config.budgets === undefined + ? undefined + : config.budgets + .map((budget) => { + const maxLimit = normalizeNumericField(budget.max_limit); + if (maxLimit === undefined) return undefined; + return { + max_limit: maxLimit, + reset_duration: budget.reset_duration || "1M", + }; + }) + .filter( + ( + budget, + ): budget is { + max_limit: number; + reset_duration: string; + } => budget !== undefined, + ), weight:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 378 - 410, normalizeProviderConfigs is currently passing config.budgets through untouched (and typed as any[]), which leaves string values and blank rows in the payload; update normalizeProviderConfigs to explicitly normalize config.budgets for each config: filter out empty budget rows, coerce numeric fields (e.g., max_limit) from strings to numbers (or null when blank/NaN) and preserve other fields, and return the normalized budgets in place of config.budgets; also tighten the return typing from any[] to a proper shape if possible so TypeScript will catch mismatches.framework/configstore/tables/virtualkey.go (1)
37-38:⚠️ Potential issue | 🟠 MajorKeep a legacy single-budget decode path.
These structs still sit on JSON/config decode paths, but the old singular
budgetfield no longer has anywhere to land. Older single-budget payloads will now unmarshal to zero budgets before any DB backfill can help, and legacy virtual-keybudget.calendar_alignedis also lost. Add a temporary compatibility shim that lifts the old shape intoBudgetsandCalendarAligned.Also applies to: 218-228
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` around lines 37 - 38, Add a temporary compatibility shim that detects the legacy singular "budget" JSON shape and lifts it into the new Budgets slice and preserves legacy calendar-aligned flag: implement a custom UnmarshalJSON on the struct that owns Budgets (the one in virtualkey.go that declares Budgets []TableBudget and Keys []TableKey) which first attempts normal unmarshalling, then checks if the input contains a top-level "budget" object (or "budget":{"calendar_aligned":...}) and if so converts that object into a single TableBudget appended to Budgets and sets the struct's CalendarAligned field accordingly; ensure the UnmarshalJSON falls back to the standard decode for current shapes and add the same shim to the other equivalent struct mentioned in the comment (the block around lines 218-228).framework/configstore/migrations.go (1)
5395-5404:⚠️ Potential issue | 🟠 MajorBackfill misses provider-config-owned legacy budgets when propagating
calendar_aligned.The current update only reads
governance_budgets.virtual_key_id. Budgets backfilled viaprovider_config_idwon’t propagatecalendar_aligned=trueto their owning virtual key, so upgraded data can end up inconsistent.💡 Suggested fix
- if err := tx.Exec(` - UPDATE governance_virtual_keys SET calendar_aligned = true - WHERE id IN ( - SELECT DISTINCT virtual_key_id FROM governance_budgets - WHERE calendar_aligned = true AND virtual_key_id IS NOT NULL - ) AND calendar_aligned = false - `).Error; err != nil { + if err := tx.Exec(` + UPDATE governance_virtual_keys + SET calendar_aligned = true + WHERE id IN ( + SELECT DISTINCT COALESCE(b.virtual_key_id, pc.virtual_key_id) + FROM governance_budgets b + LEFT JOIN governance_virtual_key_provider_configs pc + ON pc.id = b.provider_config_id + WHERE b.calendar_aligned = true + AND COALESCE(b.virtual_key_id, pc.virtual_key_id) IS NOT NULL + ) AND (calendar_aligned = false OR calendar_aligned IS NULL) + `).Error; err != nil { return fmt.Errorf("failed to backfill calendar_aligned from budgets to virtual keys: %w", err) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 5395 - 5404, The backfill only propagates calendar_aligned from governance_budgets.virtual_key_id and thus misses budgets owned via provider_config_id; update the tx.Exec SQL used in the mg.HasColumn check so it also considers budgets that map to a VK through their provider_config: join governance_provider_configs on governance_budgets.provider_config_id = governance_provider_configs.id (or select governance_provider_configs.virtual_key_id) and include those virtual_key_ids in the UPDATE predicate so governance_virtual_keys.calendar_aligned is set true for VKs referenced either directly by governance_budgets.virtual_key_id or indirectly via governance_provider_configs.virtual_key_id.
🧹 Nitpick comments (3)
docs/openapi/openapi.yaml (1)
551-555: Add a top-levelUserstag entry for API docs consistency.These new paths reference operations tagged as
Users, butopenapi.yamldoes not currently define aUserstag in the globaltagssection. Adding it improves grouping/description quality in generated docs.Suggested diff
- name: Cache description: Cache management endpoints + - name: Users + description: User management endpoints🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/openapi/openapi.yaml` around lines 551 - 555, The OpenAPI spec is missing a top-level tag for "Users" even though paths /api/users and /api/users/{id} use that tag; add a new entry named "Users" to the global tags array in openapi.yaml (the document's top-level tags section) with a short description (e.g., "User management operations") so generated docs properly group the operations; update the tags list where other tags are defined so the "Users" tag matches the exact tag string used in the operations.ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx (1)
306-307: Use a stable key for each budget row.
idxwill reshuffle DOM state when budgets are added, removed, or reordered. These rows already have a stableb.id.♻️ Proposed change
- {vk.budgets.map((b, idx) => ( - <div key={idx} className="flex flex-col"> + {vk.budgets.map((b) => ( + <div key={b.id} className="flex flex-col">🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx` around lines 306 - 307, The budget list mapping in VirtualKeysTable uses the array index (idx) as the React key which can cause DOM/state reshuffles; change the key on the mapped element in the vk.budgets.map callback to use the stable budget identifier (b.id) instead of idx (e.g., key={b.id} or a stringified form) so each budget row keeps a stable identity when budgets are added, removed, or reordered.plugins/governance/store.go (1)
1475-1502: Avoid O(budgets × virtualKeys × providerConfigs) scans during budget reset checks.
ResetExpiredBudgetsInMemorynow scansgs.virtualKeysfor each budget to infercalendarAligned. This will degrade noticeably at scale. Precomputing ownership indexes once per reset pass would keep this near O(budgets + virtualKeys + providerConfigs).♻️ Proposed refactor sketch
func (gs *LocalGovernanceStore) ResetExpiredBudgetsInMemory(ctx context.Context) []*configstoreTables.TableBudget { now := time.Now() var resetBudgets []*configstoreTables.TableBudget + + // Build lookup tables once for this pass. + vkCalendarByID := make(map[string]bool) + vkCalendarByProviderConfigID := make(map[uint]bool) + gs.virtualKeys.Range(func(_, v interface{}) bool { + if vk, ok := v.(*configstoreTables.TableVirtualKey); ok && vk != nil { + vkCalendarByID[vk.ID] = vk.CalendarAligned + for _, pc := range vk.ProviderConfigs { + vkCalendarByProviderConfigID[pc.ID] = vk.CalendarAligned + } + } + return true + }) gs.budgets.Range(func(key, value interface{}) bool { budget, ok := value.(*configstoreTables.TableBudget) if !ok || budget == nil { return true } @@ - calendarAligned := false - if budget.VirtualKeyID != nil { - gs.virtualKeys.Range(func(_, v interface{}) bool { ... }) - } else if budget.ProviderConfigID != nil { - gs.virtualKeys.Range(func(_, v interface{}) bool { ... }) - } + calendarAligned := false + if budget.VirtualKeyID != nil { + calendarAligned = vkCalendarByID[*budget.VirtualKeyID] + } else if budget.ProviderConfigID != nil { + calendarAligned = vkCalendarByProviderConfigID[*budget.ProviderConfigID] + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 1475 - 1502, ResetExpiredBudgetsInMemory currently iterates gs.virtualKeys for every budget to determine calendarAligned (using budget.VirtualKeyID, budget.ProviderConfigID and TableVirtualKey.ProviderConfigs), which is O(budgets × virtualKeys × providerConfigs); instead, during the start of the reset pass build one-time lookup maps from virtual key ID -> CalendarAligned and from providerConfig ID -> CalendarAligned (scan gs.virtualKeys once, record vk.ID -> vk.CalendarAligned and for each pc in vk.ProviderConfigs record pc.ID -> vk.CalendarAligned), then in the per-budget loop consult these maps (falling back to false if not found) to decide calendar alignment; this keeps complexity near O(budgets + virtualKeys + providerConfigs) and preserves the existing logic in ResetExpiredBudgetsInMemory.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/openapi/paths/management/users.yaml`:
- Around line 9-20: Update the OpenAPI parameter schemas for the query params
named "page" and "limit" to encode their bounds: set "page" schema to have
"minimum": 1 (keeping type: integer and default: 1) and set "limit" schema to
have "minimum": 1 and "maximum": 100 (keeping type: integer and default: 20) so
validators and clients enforce the 1-based page and max-100 limit constraints.
In `@plugins/governance/store_test.go`:
- Around line 261-264: Replace assert.Error checks that are followed by
dereferencing err with require.Error to stop the test immediately if err is nil;
specifically update calls around store.CheckBudget(...) and similar assertions
(e.g., the cases at lines referenced by the reviewer: the blocks using
store.CheckBudget with EvaluationRequest and subsequent assert.Contains on
err.Error()) so they call require.Error(t, err, ...) before doing
assert.Contains(t, err.Error(), ...). Do the same replacement for the other
occurrences noted (the blocks around the ranges 287-290, 337-340, 363-366,
394-396) to avoid panics when err is unexpectedly nil.
In `@plugins/governance/test_utils.go`:
- Around line 210-217: buildVirtualKeyWithMultiBudgets skips seeding the default
provider config so resolver tests short-circuit before exercising multi-budget
logic; modify buildVirtualKeyWithMultiBudgets to mirror
buildVirtualKeyWithBudget/buildVirtualKeyWithRateLimit by seeding/assigning the
same default provider config (e.g., call the same helper or set the
ProviderConfigID and Provider fields used by those builders) before returning,
and ensure each budgets[i].VirtualKeyID is set as currently done so the
multi-budget fixture lives on the same resolver path as the other builders.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 1179-1185: The MCP delete Button is an icon-only control and needs
an accessible name; update the Button (the one calling handleRemoveMCPClient) to
include an aria-label (or aria-labelledby) such as aria-label={`Remove MCP
client ${index + 1}`} or, if a client name is available, aria-label={`Remove
${client.name}`}, ensuring the attribute is added alongside the existing props
(type, variant, size, onClick, data-testid) so screen readers can announce the
control.
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 2080-2095: In loadFromConfigMemory, after the existing
ProviderConfigs loop that hydrates pc.RateLimit, also populate provider-config
budgets by iterating the in-memory TableBudget slice (e.g., variable named
budgets or tableBudgets) and for each pc (ProviderConfigs[j]) append or set any
TableBudget whose ProviderConfigID matches pc.ID (same pattern used for model
configs/providers). Ensure you check vk.ProviderConfigs != nil and the budgets
slice != nil, and attach budget entries to the ProviderConfig (e.g., pc.Budgets
or equivalent field) rather than replacing existing fields so multiple budgets
per provider config are preserved.
---
Duplicate comments:
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2982-2987: The SQLite migration lane is missing the detailed
budget backfill/cleanup assertions that exist in run_postgres_migration_tests;
add the same verification step (call the budget verification helper used by
Postgres or create a verify_budget_migration_sqlite wrapper) into
run_sqlite_migration_tests immediately after validate_sqlite_data (and before
stopping the test harness), and on failure log the same error, call
stop_bifrost, and return a non-zero exit so SQLite regressions in
virtual_key_id/provider_config_id or budget_id cleanup are caught; reference
verify_budget_migration_postgres, run_postgres_migration_tests,
run_sqlite_migration_tests, validate_sqlite_data, and stop_bifrost to locate
where to insert the call.
In `@framework/configstore/migrations.go`:
- Around line 5395-5404: The backfill only propagates calendar_aligned from
governance_budgets.virtual_key_id and thus misses budgets owned via
provider_config_id; update the tx.Exec SQL used in the mg.HasColumn check so it
also considers budgets that map to a VK through their provider_config: join
governance_provider_configs on governance_budgets.provider_config_id =
governance_provider_configs.id (or select
governance_provider_configs.virtual_key_id) and include those virtual_key_ids in
the UPDATE predicate so governance_virtual_keys.calendar_aligned is set true for
VKs referenced either directly by governance_budgets.virtual_key_id or
indirectly via governance_provider_configs.virtual_key_id.
In `@framework/configstore/tables/virtualkey.go`:
- Around line 37-38: Add a temporary compatibility shim that detects the legacy
singular "budget" JSON shape and lifts it into the new Budgets slice and
preserves legacy calendar-aligned flag: implement a custom UnmarshalJSON on the
struct that owns Budgets (the one in virtualkey.go that declares Budgets
[]TableBudget and Keys []TableKey) which first attempts normal unmarshalling,
then checks if the input contains a top-level "budget" object (or
"budget":{"calendar_aligned":...}) and if so converts that object into a single
TableBudget appended to Budgets and sets the struct's CalendarAligned field
accordingly; ensure the UnmarshalJSON falls back to the standard decode for
current shapes and add the same shim to the other equivalent struct mentioned in
the comment (the block around lines 218-228).
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 291-294: The calendar alignment toggle and submit logic only
consider top-level budgets and ignore provider-config (watched) budgets, so
update the alignment checks to include provider budgets: change the
hasAnyAlignableBudget computation to fold provider budgets (watchedBudgets) into
the same logic used for top-level budgets (use both budgets and watchedBudgets
when checking max_limit and supportsCalendarAlignment), and ensure
createData.calendar_aligned and updateData.calendar_aligned are set/gated when
any VK-owned budget exists (not just when validBudgets.length > 0) by treating
provider budgets as VK-owned in the conditions that build createData/updateData;
update references to watchedBudgets, budgets, hasAnyAlignableBudget,
createData.calendar_aligned, updateData.calendar_aligned, and validBudgets
accordingly.
- Around line 378-410: normalizeProviderConfigs is currently passing
config.budgets through untouched (and typed as any[]), which leaves string
values and blank rows in the payload; update normalizeProviderConfigs to
explicitly normalize config.budgets for each config: filter out empty budget
rows, coerce numeric fields (e.g., max_limit) from strings to numbers (or null
when blank/NaN) and preserve other fields, and return the normalized budgets in
place of config.budgets; also tighten the return typing from any[] to a proper
shape if possible so TypeScript will catch mismatches.
---
Nitpick comments:
In `@docs/openapi/openapi.yaml`:
- Around line 551-555: The OpenAPI spec is missing a top-level tag for "Users"
even though paths /api/users and /api/users/{id} use that tag; add a new entry
named "Users" to the global tags array in openapi.yaml (the document's top-level
tags section) with a short description (e.g., "User management operations") so
generated docs properly group the operations; update the tags list where other
tags are defined so the "Users" tag matches the exact tag string used in the
operations.
In `@plugins/governance/store.go`:
- Around line 1475-1502: ResetExpiredBudgetsInMemory currently iterates
gs.virtualKeys for every budget to determine calendarAligned (using
budget.VirtualKeyID, budget.ProviderConfigID and
TableVirtualKey.ProviderConfigs), which is O(budgets × virtualKeys ×
providerConfigs); instead, during the start of the reset pass build one-time
lookup maps from virtual key ID -> CalendarAligned and from providerConfig ID ->
CalendarAligned (scan gs.virtualKeys once, record vk.ID -> vk.CalendarAligned
and for each pc in vk.ProviderConfigs record pc.ID -> vk.CalendarAligned), then
in the per-budget loop consult these maps (falling back to false if not found)
to decide calendar alignment; this keeps complexity near O(budgets + virtualKeys
+ providerConfigs) and preserves the existing logic in
ResetExpiredBudgetsInMemory.
In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx`:
- Around line 306-307: The budget list mapping in VirtualKeysTable uses the
array index (idx) as the React key which can cause DOM/state reshuffles; change
the key on the mapped element in the vk.budgets.map callback to use the stable
budget identifier (b.id) instead of idx (e.g., key={b.id} or a stringified form)
so each budget row keeps a stable identity when budgets are added, removed, or
reordered.
🪄 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: 1990e10f-8dce-4e25-9cbb-304640bc9f50
📒 Files selected for processing (31)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shdocs/openapi/openapi.yamldocs/openapi/paths/management/users.yamldocs/openapi/schemas/management/users.yamlframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/ui/multiBudgetLines.tsxui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (6)
- .claude/skills/expect
- plugins/governance/model_provider_governance_test.go
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- .agents/skills/expect/SKILL.md
- transports/config.schema.json
- docs/openapi/schemas/management/users.yaml
🚧 Files skipped from review as they are similar to previous changes (12)
- plugins/governance/resolver.go
- plugins/governance/http_transport_prehook_test.go
- transports/bifrost-http/handlers/governance_test.go
- ui/app/workspace/governance/access-profiles/page.tsx
- plugins/governance/resolver_test.go
- framework/configstore/tables/budget.go
- framework/configstore/rdb_test.go
- ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
- ui/components/ui/multiBudgetLines.tsx
- framework/configstore/rdb.go
- transports/bifrost-http/lib/config_test.go
- transports/bifrost-http/handlers/governance.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
1994f2a to
de37270
Compare
11 tasks
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (8)
.github/workflows/scripts/run-migration-tests.sh (1)
2982-2987:⚠️ Potential issue | 🟠 MajorSQLite still skips the multi-budget migration checks.
Step 6 only runs in the PostgreSQL lane. The SQLite path still stops at
validate_sqlite_data(), which only verifies that a couple of tables are non-empty, so regressions invirtual_key_id/provider_config_idbackfill or legacybudget_idcleanup can still pass unnoticed.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/scripts/run-migration-tests.sh around lines 2982 - 2987, The SQLite test lane is skipping budget migration checks; after validate_sqlite_data() you should run the same budget verification used for Postgres rather than stopping early—either call a shared verification function (e.g., rename verify_budget_migration_postgres to verify_budget_migration and invoke it from both lanes) or add a new verify_budget_migration_sqlite and call it after validate_sqlite_data(); ensure the script logs failures with log_error and calls stop_bifrost on failure (same behavior as the PostgreSQL branch).framework/configstore/rdb.go (1)
2334-2347:⚠️ Potential issue | 🟠 MajorMake standalone provider-config deletes atomic.
When
DeleteVirtualKeyProviderConfigruns without a suppliedtx, the budget delete, provider-config delete, and rate-limit delete execute as separate statements. If a later step fails, this leaves a partially deleted provider config and orphaned owned resources. MirrorDeleteVirtualKeyby opening a local transaction whentxis empty.Based on learnings: budgets and rate limits have a 1:1 ownership with their parent entities and are deleted together.🧩 Suggested wrapper
func (s *RDBConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } + if len(tx) == 0 { + return s.db.WithContext(ctx).Transaction(func(inner *gorm.DB) error { + return s.DeleteVirtualKeyProviderConfig(ctx, id, inner) + }) + } + txDB := tx[0]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2334 - 2347, The DeleteVirtualKeyProviderConfig implementation performs multiple separate deletes (TableBudget, TableVirtualKeyProviderConfig, TableRateLimit) when called without a transaction, risking partial deletes; modify DeleteVirtualKeyProviderConfig to mirror DeleteVirtualKey by creating a local transaction when the incoming tx is nil (e.g., open txDB = db.Begin()/db.WithContext(ctx).Begin()), run the three deletes inside that transaction using txDB (same variables: providerConfig, rateLimitID, TableBudget, TableVirtualKeyProviderConfig, TableRateLimit), and commit/rollback the local transaction on success/error so the budget, provider-config, and rate-limit deletions are atomic.plugins/governance/store_test.go (1)
261-263:⚠️ Potential issue | 🟡 MinorUse
require.Errorbefore callingerr.Error().
assert.Errorkeeps the test running, so a nilerrhere panics onerr.Error()and hides the actual regression.🧪 Safer assertion pattern
err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil) -assert.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") +require.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") assert.Contains(t, err.Error(), "budget exceeded")Also applies to: 287-289, 337-339, 363-365, 394-396
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 261 - 263, Replace usages of assert.Error followed immediately by err.Error() with require.Error to stop the test on a nil error; specifically change the calls around store.CheckBudget(...) that assert.Error(t, err, ...) before using err.Error() to require.Error(t, err, ...) so the test fails fast and avoids nil dereference — apply this correction for the occurrences that call store.CheckBudget with vk and &EvaluationRequest{Provider: schemas.OpenAI} (and the similar assertions at the other reported locations) so each time you check err.Error() you have first required a non-nil err.transports/bifrost-http/handlers/governance.go (2)
864-873:⚠️ Potential issue | 🔴 CriticalPreload provider-config budgets before reconciling them.
Find(&existingConfigs)only loads the provider-config rows;existing.Budgetsis empty here. Because this PUT path receives the full provider-config payload, every ordinary save withpc.Budgetspopulated will miss all matches, recreate fresh budget rows withCurrentUsage: 0, and leave the old ones behind.🐛 Proposed fix
- if err := tx.Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs).Error; err != nil { + if err := tx.Preload("Budgets").Where("virtual_key_id = ?", vk.ID).Find(&existingConfigs).Error; err != nil { return err }#!/bin/bash set -euo pipefail # Verify that TableVirtualKeyProviderConfig has a Budgets association, # and compare the current update path with any preload-based load paths. rg -n -C3 'type TableVirtualKeyProviderConfig struct|Budgets \[\].*TableBudget' framework/configstore transports/bifrost-http/handlers rg -n -C3 'Where\("virtual_key_id = \?", vk.ID\)\.Find\(&existingConfigs\)|Preload\("Budgets"\)|func .*GetVirtualKey\(' framework/configstore transports/bifrost-http/handlersBased on learnings,
updateVirtualKey (PUT /api/governance/virtual-keys/{vk_id}) expects a full payload from the frontend.Also applies to: 1008-1056
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 864 - 873, The existing load of provider configs uses tx.Where(...).Find(&existingConfigs) which does not populate the TableVirtualKeyProviderConfig.Budgets association, causing budget rows to be treated as missing and recreated; change the query that builds existingConfigs (used to create existingConfigsMap) to preload the Budgets association (e.g., use tx.Preload("Budgets").Where(...).Find(&existingConfigs)) so each config's Budgets slice is populated before reconciliation; apply the same Preload("Budgets") fix to the other similar load in the updateVirtualKey PUT path (the similar block around lines 1008-1056).
442-459:⚠️ Potential issue | 🟠 MajorReturn 400 for all invalid multi-budget payloads.
max_limit == 0still slips past the prechecks here, and the provider-config create branches still let malformed budgets fall through tovalidateBudget(). Those failures are plainerrors, socreateVirtualKey()/updateVirtualKey()still return 500 for client validation mistakes instead of 400.Also applies to: 575-597, 740-755, 932-955, 991-1006
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance.go` around lines 442 - 459, The budget validation allows max_limit == 0 and lets validation errors escalate to 500; update the checks to reject non-positive limits by changing the condition to b.MaxLimit <= 0 and return a 400 error message for zero or negative values, ensure ParseDuration errors already produce 400, and modify callers (createVirtualKey, updateVirtualKey) so any error returned from validateBudget (or similar budget validation functions) is converted into a client 400 response via SendError instead of letting it propagate to a 500; apply the same fix to the other budget-validation blocks noted (the similar blocks around the other line ranges) so all malformed multi-budget payloads consistently return 400.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (3)
291-293:⚠️ Potential issue | 🟠 Major
calendar_alignedis still scoped to top-level budgets only.The new visibility check now scans all top-level VK budget rows, but the toggle, reset path, and both payload builders still ignore
providerConfigs[*].budgets. A key with only provider budgets still can't manage the VK-wide flag, and clearing top-level rows forces it off even when provider budgets remain.Also applies to: 366-369, 443-455, 489-497, 1201-1230
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 291 - 293, The visibility and handling of the VK-wide calendar_aligned flag currently only inspects top-level budgets; update all places that determine visibility and mutate or build payloads to include providerConfigs[*].budgets as well — specifically extend the hasAnyAlignableBudget check (and analogous checks at lines ~366-369, 443-455, 489-497, 1201-1230) to iterate both top-level watchedBudgets and each providerConfig.budgets array, and update the toggle handler, reset path, and both payload builder functions to read and set calendar_aligned by considering provider budgets (e.g., aggregate supportsCalendarAlignment(...) and max_limit checks across providerConfigs[*].budgets) so keys with only provider budgets correctly show, preserve, and clear the VK-wide flag.
1179-1185:⚠️ Potential issue | 🟡 MinorAdd an accessible name to the MCP delete button.
This is still icon-only, so assistive tech gets an unlabeled button.
♿ Minimal fix
<Button type="button" variant="ghost" size="sm" + aria-label={`Remove MCP client ${config.mcp_client_name}`} onClick={() => handleRemoveMCPClient(index)} data-testid={`vk-delete-mcp-${index}`} >🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 1179 - 1185, The MCP delete Button (rendered with onClick={handleRemoveMCPClient} and data-testid={`vk-delete-mcp-${index}`}) is icon-only and missing an accessible name; add an accessible label (e.g., aria-label={`Delete MCP client ${index + 1}`} or aria-label="Delete MCP client") or include visually hidden text inside the Button so screen readers can identify it, ensuring the change is applied to the Button component where handleRemoveMCPClient(index) is used.
378-410:⚠️ Potential issue | 🔴 CriticalNormalize/filter provider budgets before serializing them.
normalizeProviderConfigs()convertsweightandrate_limit, but it forwardsconfig.budgetsunchanged. Because provider budget rows are stored as form strings and can include blank entries, create/update can still post string/emptymax_limitvalues instead of numeric budgets.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 378 - 410, normalizeProviderConfigs currently leaves config.budgets untouched so string/empty budget rows get serialized; update normalizeProviderConfigs (used with providerConfigs and existingConfigs) to map and filter config.budgets: for each budget entry run its max_limit through normalizeNumericField (same helper used for rate_limit fields) converting numeric strings to numbers and invalid/empty values to null/undefined as appropriate, and drop empty/blank budget rows before returning the provider config object so creates/updates only send normalized numeric budget values; reference normalizeProviderConfigs, config.budgets, normalizeNumericField and VirtualKey["provider_configs"] when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2721-2728: The assertion for governance_virtual_key_budgets is
inverted: it uses run_postgres_scalar to set junction_vk and then treats "0"
(missing) as success; instead, flip the condition so a non-zero count is treated
as the expected (success) case and zero is a warning. Update the if-statement
that checks junction_vk (the variable set by run_postgres_scalar) to test for !=
"0" and swap the log calls (use log_info for the table existing and log_warn for
it being missing) so messages correctly reflect the expected presence of
governance_virtual_key_budgets created by the migration.
- Around line 2701-2718: compare_postgres_snapshots() is currently treating the
intentional drops of governance_virtual_keys.budget_id and
governance_virtual_key_provider_configs.budget_id as unexpected removals,
causing validation to fail before your later checks run; update
compare_postgres_snapshots() to ignore these specific column removals (e.g., add
a small allowlist/exception check for table names "governance_virtual_keys" and
"governance_virtual_key_provider_configs" with column "budget_id") so the
snapshot comparator does not flag them as failures, then re-run
validate_postgres_data() flow to confirm the later verifier can assert the
columns were dropped.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 208-214: The current mapping for rate_limit uses truthiness so
numeric zeros get treated as absent causing stored 0 limits to become undefined;
change the checks for token_max_limit and request_max_limit to explicit
null/undefined checks (e.g. config.rate_limit.token_max_limit !== undefined &&
config.rate_limit.token_max_limit !== null) so 0 is preserved and still
converted to String(...) for round-tripping, updating the mapping where
rate_limit: { token_max_limit: ..., request_max_limit: ... } is constructed in
virtualKeySheet.tsx.
- Around line 1201-1204: The data-testid on the MultiBudgetLines component uses
a two-part value "vk-budget"; update it to the three-part convention
(entity-element-qualifier) - for example change data-testid="vk-budget" to
data-testid="vk-budget-lines" on the MultiBudgetLines element (id="vkBudget") so
the selector pattern `entity-element-qualifier` is respected for interactive UI
elements.
- Around line 917-936: The two back-to-back calls to handleUpdateProviderConfig
(updating "budget" then "budgets") race because each clones the same
providerConfigs snapshot; instead compute the legacy budget object and the
budgets array once (e.g., const nextBudget = lines.length ? { max_limit:
lines[0].max_limit, reset_duration: lines[0].reset_duration } : undefined; const
nextBudgets = lines.length ? lines.map(l => ({ max_limit: l.max_limit,
reset_duration: l.reset_duration })) : []) and then perform a single state
update so both fields are written together—either by calling a single
handleUpdateProviderConfig that accepts a partial update object (e.g.,
handleUpdateProviderConfig(index, { budget: nextBudget, budgets: nextBudgets }))
or by invoking the underlying set function once with the merged provider config;
reference handleUpdateProviderConfig, "budget", and "budgets" when making the
change.
In `@ui/lib/types/governance.ts`:
- Around line 163-166: Create a VK-specific budget type by adding
VirtualKeyBudgetRequest = Omit<CreateBudgetRequest, "calendar_aligned"> and
replace any property types that currently use CreateBudgetRequest[] for the
virtual-key payload (the budgets?: field in the governance type and the
analogous array usage later in the file) with VirtualKeyBudgetRequest[] so
callers can no longer pass budgets[i].calendar_aligned; leave the top-level
calendar_aligned boolean on the virtual-key object unchanged.
---
Duplicate comments:
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2982-2987: The SQLite test lane is skipping budget migration
checks; after validate_sqlite_data() you should run the same budget verification
used for Postgres rather than stopping early—either call a shared verification
function (e.g., rename verify_budget_migration_postgres to
verify_budget_migration and invoke it from both lanes) or add a new
verify_budget_migration_sqlite and call it after validate_sqlite_data(); ensure
the script logs failures with log_error and calls stop_bifrost on failure (same
behavior as the PostgreSQL branch).
In `@framework/configstore/rdb.go`:
- Around line 2334-2347: The DeleteVirtualKeyProviderConfig implementation
performs multiple separate deletes (TableBudget, TableVirtualKeyProviderConfig,
TableRateLimit) when called without a transaction, risking partial deletes;
modify DeleteVirtualKeyProviderConfig to mirror DeleteVirtualKey by creating a
local transaction when the incoming tx is nil (e.g., open txDB =
db.Begin()/db.WithContext(ctx).Begin()), run the three deletes inside that
transaction using txDB (same variables: providerConfig, rateLimitID,
TableBudget, TableVirtualKeyProviderConfig, TableRateLimit), and commit/rollback
the local transaction on success/error so the budget, provider-config, and
rate-limit deletions are atomic.
In `@plugins/governance/store_test.go`:
- Around line 261-263: Replace usages of assert.Error followed immediately by
err.Error() with require.Error to stop the test on a nil error; specifically
change the calls around store.CheckBudget(...) that assert.Error(t, err, ...)
before using err.Error() to require.Error(t, err, ...) so the test fails fast
and avoids nil dereference — apply this correction for the occurrences that call
store.CheckBudget with vk and &EvaluationRequest{Provider: schemas.OpenAI} (and
the similar assertions at the other reported locations) so each time you check
err.Error() you have first required a non-nil err.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 864-873: The existing load of provider configs uses
tx.Where(...).Find(&existingConfigs) which does not populate the
TableVirtualKeyProviderConfig.Budgets association, causing budget rows to be
treated as missing and recreated; change the query that builds existingConfigs
(used to create existingConfigsMap) to preload the Budgets association (e.g.,
use tx.Preload("Budgets").Where(...).Find(&existingConfigs)) so each config's
Budgets slice is populated before reconciliation; apply the same
Preload("Budgets") fix to the other similar load in the updateVirtualKey PUT
path (the similar block around lines 1008-1056).
- Around line 442-459: The budget validation allows max_limit == 0 and lets
validation errors escalate to 500; update the checks to reject non-positive
limits by changing the condition to b.MaxLimit <= 0 and return a 400 error
message for zero or negative values, ensure ParseDuration errors already produce
400, and modify callers (createVirtualKey, updateVirtualKey) so any error
returned from validateBudget (or similar budget validation functions) is
converted into a client 400 response via SendError instead of letting it
propagate to a 500; apply the same fix to the other budget-validation blocks
noted (the similar blocks around the other line ranges) so all malformed
multi-budget payloads consistently return 400.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 291-293: The visibility and handling of the VK-wide
calendar_aligned flag currently only inspects top-level budgets; update all
places that determine visibility and mutate or build payloads to include
providerConfigs[*].budgets as well — specifically extend the
hasAnyAlignableBudget check (and analogous checks at lines ~366-369, 443-455,
489-497, 1201-1230) to iterate both top-level watchedBudgets and each
providerConfig.budgets array, and update the toggle handler, reset path, and
both payload builder functions to read and set calendar_aligned by considering
provider budgets (e.g., aggregate supportsCalendarAlignment(...) and max_limit
checks across providerConfigs[*].budgets) so keys with only provider budgets
correctly show, preserve, and clear the VK-wide flag.
- Around line 1179-1185: The MCP delete Button (rendered with
onClick={handleRemoveMCPClient} and data-testid={`vk-delete-mcp-${index}`}) is
icon-only and missing an accessible name; add an accessible label (e.g.,
aria-label={`Delete MCP client ${index + 1}`} or aria-label="Delete MCP client")
or include visually hidden text inside the Button so screen readers can identify
it, ensuring the change is applied to the Button component where
handleRemoveMCPClient(index) is used.
- Around line 378-410: normalizeProviderConfigs currently leaves config.budgets
untouched so string/empty budget rows get serialized; update
normalizeProviderConfigs (used with providerConfigs and existingConfigs) to map
and filter config.budgets: for each budget entry run its max_limit through
normalizeNumericField (same helper used for rate_limit fields) converting
numeric strings to numbers and invalid/empty values to null/undefined as
appropriate, and drop empty/blank budget rows before returning the provider
config object so creates/updates only send normalized numeric budget values;
reference normalizeProviderConfigs, config.budgets, normalizeNumericField and
VirtualKey["provider_configs"] when making the change.
🪄 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: 0b1b97ce-c299-4b4b-936f-8e4fbe34da09
⛔ Files ignored due to path filters (10)
docs/media/user-provisioning/zitadel-add-role.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-add-user-select-key.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-auth-method.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-namne.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-uri.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-role-assignemnt.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-select-project.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-token-config.pngis excluded by!**/*.pngui/package-lock.jsonis excluded by!**/package-lock.json
📒 Files selected for processing (32)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shdocs/openapi/openapi.yamldocs/openapi/paths/management/users.yamldocs/openapi/schemas/management/users.yamlframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/http_transport_prehook_test.goplugins/governance/model_provider_governance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/ui/multiBudgetLines.tsxui/lib/types/governance.tsui/package.json
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (10)
- .claude/skills/expect
- plugins/governance/model_provider_governance_test.go
- docs/openapi/openapi.yaml
- framework/configstore/rdb_test.go
- docs/openapi/paths/management/users.yaml
- framework/configstore/tables/virtualkey.go
- .agents/skills/expect/SKILL.md
- docs/openapi/schemas/management/users.yaml
- transports/bifrost-http/lib/config_test.go
- framework/configstore/migrations.go
🚧 Files skipped from review as they are similar to previous changes (11)
- plugins/governance/http_transport_prehook_test.go
- ui/app/workspace/governance/access-profiles/page.tsx
- plugins/governance/resolver_test.go
- transports/bifrost-http/handlers/governance_test.go
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
- framework/configstore/tables/budget.go
- plugins/governance/test_utils.go
- ui/components/ui/multiBudgetLines.tsx
- plugins/governance/store.go
- transports/config.schema.json
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
de37270 to
3e8d874
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (11)
.github/workflows/scripts/run-migration-tests.sh (2)
2624-2627:⚠️ Potential issue | 🟠 MajorAllow the legacy
budget_idremovals in the snapshot comparator.Step 6 now asserts that
governance_virtual_keys.budget_idandgovernance_virtual_key_provider_configs.budget_idwere dropped, but Step 5 still treats those removals as unexpected column loss. Any migration from a pre-multi-budget tag will fail incompare_postgres_snapshots()before the new verifier runs.🧩 Suggested fix
- local dropped_columns="enable_governance" + local dropped_columns="enable_governance" + if [ "$table" = "governance_virtual_keys" ] || [ "$table" = "governance_virtual_key_provider_configs" ]; then + dropped_columns="$dropped_columns budget_id" + fi # provider, model (dropped from routing_rules only in v1.4.12) if [ "$table" = "routing_rules" ]; then dropped_columns="$dropped_columns provider model" fiAlso applies to: 2674-2680
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/scripts/run-migration-tests.sh around lines 2624 - 2627, The snapshot comparator is flagging legacy removals of budget_id as unexpected; update the comparator used in compare_postgres_snapshots() to treat budget_id as an allowed/ignored dropped column by adding "budget_id" to the ignore_columns list (the local variable named ignore_columns in the script) so removals of governance_virtual_keys.budget_id and governance_virtual_key_provider_configs.budget_id are not considered failures during snapshot comparison.
2837-2913:⚠️ Potential issue | 🟠 MajorBudget-migration verification still never runs for SQLite.
The new multi-budget assertions are only added to the PostgreSQL flow.
validate_sqlite_data()still just checks table counts, so a SQLite regression invirtual_key_id/provider_config_idbackfill or legacybudget_idcleanup would stay green.Also applies to: 3164-3169
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/scripts/run-migration-tests.sh around lines 2837 - 2913, The SQLite validation path is missing the new multi-budget assertions added in verify_budget_migration_postgres(); update validate_sqlite_data (or add a new verify_budget_migration_sqlite and call it from the SQLite flow) to run the same checks: ensure governance_budgets.virtual_key_id/provider_config_id existence and proper values for the test rows, confirm governance_virtual_keys and governance_virtual_key_provider_configs no longer have budget_id, and ensure junction table governance_virtual_key_budgets is dropped (or warn if present); mirror the queries and logging logic from verify_budget_migration_postgres() so SQLite migrations are validated the same way.framework/configstore/rdb.go (1)
2327-2339:⚠️ Potential issue | 🟠 MajorMake single provider-config deletion match the transactional bulk-delete path.
This path still skips
governance_virtual_key_provider_config_keyscleanup and performs the budget/config/rate-limit deletes outside a transaction. If a later delete fails, you can leave orphaned key joins or a partially deleted provider-config state, unlikeDeleteVirtualKeyabove which already cleans this up in-order.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2327 - 2339, The single-item provider-config delete path must mirror the transactional bulk-delete in DeleteVirtualKey: run the join-table cleanup, budget, provider-config and rate-limit deletions inside the same txDB transaction and in the same order to avoid orphaned joins or partial deletes. Modify the code that uses txDB.WithContext(ctx).Delete for tables.TableBudget, tables.TableVirtualKeyProviderConfig{}, and tables.TableRateLimit{} to also delete rows from the governance_virtual_key_provider_config_keys join table (using the provider config id) and ensure all four deletes execute on the same txDB transaction context (ctx/txDB) so any failure rolls back the entire operation.plugins/governance/store_test.go (1)
261-263:⚠️ Potential issue | 🟡 MinorUse
require.Errorbefore dereferencingerr.
assert.Errorkeeps the test running. If one of these calls unexpectedly returnsnil, the nexterr.Error()panics and hides the real failure.Safer assertion pattern
err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil) -assert.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") +require.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") assert.Contains(t, err.Error(), "budget exceeded")Also applies to: 287-289, 337-339, 363-365, 394-396
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 261 - 263, Replace assert.Error calls that are followed by dereferencing err (e.g., the call to store.CheckBudget(ctx, vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil) followed by assert.Contains(t, err.Error(), ...)) with require.Error so the test stops when err is nil; do this for the instances around the CheckBudget assertions (current snippet and the other occurrences mentioned at lines ~287-289, ~337-339, ~363-365, ~394-396) so you don’t call err.Error() on a nil error and hide the real failure.framework/configstore/tables/virtualkey.go (1)
33-38:⚠️ Potential issue | 🟠 MajorKeep a one-release decoder shim for legacy single-budget JSON.
These types are still JSON decode boundaries, but the single-budget shape (
budget_id/budget) has been removed without a lift intoBudgets. Older config/API payloads now deserialize to zero budgets, and legacybudget.calendar_alignedalso has no path intoTableVirtualKey.CalendarAligned. DB migrations do not cover file-backed configs here.Also applies to: 46-77, 219-228
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` around lines 33 - 38, The JSON decoder currently drops legacy single-budget fields (budget_id / budget and budget.calendar_aligned) causing Budgets on TableVirtualKey to be empty and CalendarAligned lost; add a one-release decoder shim by implementing a custom UnmarshalJSON for TableVirtualKey that first attempts normal decoding, then checks for legacy keys (budget_id or budget object with calendar_aligned) and, if present, converts them into a single-entry Budgets slice and sets TableVirtualKey.CalendarAligned accordingly; ensure the shim only affects JSON decoding paths (not DB migrations) and remove after the release.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (4)
1201-1204:⚠️ Potential issue | 🟡 MinorUse a 3-part
data-testidfor VK budget control.Line 1203 uses
vk-budget(2-part). Use<entity>-<element>-<qualifier>format for consistency.Suggested fix
- data-testid="vk-budget" + data-testid="vk-budget-lines"As per coding guidelines,
ui/**/*.{tsx,ts}: Add new interactive UI elements withdata-testidattributes following the patterndata-testid="<entity>-<element>-<qualifier>".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 1201 - 1204, The MultiBudgetLines instance uses a 2-part data-testid "vk-budget"; update it to a 3-part test id following the <entity>-<element>-<qualifier> convention (e.g., change data-testid on the MultiBudgetLines with id "vkBudget" to "virtual-key-budget-control" or "virtual-key-budget-config") so it matches the ui/**/*.{tsx,ts} guideline for new interactive UI elements.
1179-1185:⚠️ Potential issue | 🟡 MinorAdd an accessible name to the MCP remove button.
Line 1179 is icon-only and lacks
aria-label, so assistive tech announces an unlabeled control.Suggested fix
<Button type="button" variant="ghost" size="sm" + aria-label={`Remove MCP client ${config.mcp_client_name || index + 1}`} onClick={() => handleRemoveMCPClient(index)} data-testid={`vk-delete-mcp-${index}`} >🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 1179 - 1185, The icon-only MCP remove Button (rendered with onClick={handleRemoveMCPClient} and data-testid={`vk-delete-mcp-${index}`}) lacks an accessible name; add an aria-label or equivalent accessibleName (for example aria-label={`Remove MCP client ${index}`} or a meaningful string) to the Button component so screen readers announce its purpose while retaining the existing onClick handler and data-testid.
208-214:⚠️ Potential issue | 🟠 MajorPreserve zero-valued provider rate limits in defaults.
Line 210 and Line 212 use truthiness checks, so
0becomesundefinedand is lost on edit/save round-trip.Suggested fix
rate_limit: config.rate_limit ? { - token_max_limit: config.rate_limit.token_max_limit ? String(config.rate_limit.token_max_limit) : undefined, + token_max_limit: + config.rate_limit.token_max_limit === undefined || config.rate_limit.token_max_limit === null + ? undefined + : String(config.rate_limit.token_max_limit), token_reset_duration: config.rate_limit.token_reset_duration, - request_max_limit: config.rate_limit.request_max_limit ? String(config.rate_limit.request_max_limit) : undefined, + request_max_limit: + config.rate_limit.request_max_limit === undefined || config.rate_limit.request_max_limit === null + ? undefined + : String(config.rate_limit.request_max_limit), request_reset_duration: config.rate_limit.request_reset_duration, } : undefined,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 208 - 214, In virtualKeySheet.tsx the mapping that builds rate_limit currently uses truthiness checks for token_max_limit and request_max_limit so numeric zero is treated as absent; in the object construction for rate_limit (the lines assigning token_max_limit and request_max_limit) change the checks to explicitly test for undefined/null (e.g. use !== undefined or != null) and then String(...) the value, so a value of 0 is preserved as "0" instead of becoming undefined during the edit/save round-trip.
917-936:⚠️ Potential issue | 🟠 MajorMake provider
budget+budgetsupdate atomic.Line 919/Line 920 and Line 923/Line 928 call
handleUpdateProviderConfigtwice from the sameproviderConfigssnapshot; the second write can overwrite the first. Keep the intentional dual-write, but do it in onesetValue.Suggested fix
onChange={(lines) => { - if (lines.length === 0) { - handleUpdateProviderConfig(index, "budget", undefined); - handleUpdateProviderConfig(index, "budgets", []); - } else { - // Legacy single budget field = first line - handleUpdateProviderConfig(index, "budget", { - max_limit: lines[0].max_limit, - reset_duration: lines[0].reset_duration, - }); - // All lines in budgets array - handleUpdateProviderConfig( - index, - "budgets", - lines.map((l) => ({ - max_limit: l.max_limit, - reset_duration: l.reset_duration, - })), - ); - } + const updatedConfigs = [...providerConfigs]; + updatedConfigs[index] = { + ...updatedConfigs[index], + budget: + lines.length > 0 + ? { + max_limit: lines[0].max_limit, + reset_duration: lines[0].reset_duration, + } + : undefined, + budgets: lines.map((l) => ({ + max_limit: l.max_limit, + reset_duration: l.reset_duration, + })), + }; + form.setValue("providerConfigs", updatedConfigs, { shouldDirty: true }); }}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 917 - 936, Multiple sequential calls to handleUpdateProviderConfig (for "budget" and "budgets") use the same providerConfigs snapshot and can overwrite each other; instead build a single update object and call handleUpdateProviderConfig once for the given index to set both fields atomically. In the onChange handler (referencing handleUpdateProviderConfig and index) compute the newBudget (or undefined) and newBudgets array from lines, then call handleUpdateProviderConfig(index, { budget: newBudget, budgets: newBudgets }) in a single invocation so both fields are updated together.ui/lib/types/governance.ts (1)
156-179:⚠️ Potential issue | 🟠 MajorDrop per-budget
calendar_alignedfrom VK requests too.Lines 163 and 176 still use
CreateBudgetRequest[], so callers can sendbudgets[i].calendar_alignedalongside the new top-levelcalendar_aligned. That keeps two sources of truth in the payload even though provider-config budgets were already split out.🧩 Proposed type split
+export type VirtualKeyBudgetRequest = Omit<CreateBudgetRequest, "calendar_aligned">; + export interface CreateVirtualKeyRequest { name: string; description?: string; provider_configs?: VirtualKeyProviderConfigRequest[]; mcp_configs?: VirtualKeyMCPConfigRequest[]; team_id?: string; customer_id?: string; - budgets?: CreateBudgetRequest[]; + budgets?: VirtualKeyBudgetRequest[]; rate_limit?: CreateRateLimitRequest; is_active?: boolean; calendar_aligned?: boolean; } export interface UpdateVirtualKeyRequest { name?: string; description?: string; provider_configs?: VirtualKeyProviderConfigUpdateRequest[]; mcp_configs?: VirtualKeyMCPConfigRequest[]; team_id?: string; customer_id?: string; - budgets?: CreateBudgetRequest[]; + budgets?: VirtualKeyBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; is_active?: boolean; calendar_aligned?: boolean; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/lib/types/governance.ts` around lines 156 - 179, Change the budgets element type on CreateVirtualKeyRequest and UpdateVirtualKeyRequest so callers cannot set per-budget calendar_aligned; replace CreateBudgetRequest[] with a budget type that omits calendar_aligned (e.g., CreateBudgetNoCalendarRequest[] or a mapped Omit<CreateBudgetRequest,'calendar_aligned'>[]) and update any imports/exports accordingly; adjust the symbols CreateVirtualKeyRequest and UpdateVirtualKeyRequest to reference the new budget type so only the top-level calendar_aligned remains the source of truth.framework/configstore/migrations.go (1)
5544-5558:⚠️ Potential issue | 🟠 MajorBackfill
calendar_alignedthrough provider-config budgets too.Line 5550 only promotes budgets that already have
virtual_key_idset. Budgets migrated throughgovernance_virtual_key_provider_configs.budget_idonly getprovider_config_idon Lines 5529-5542, so any legacycalendar_aligned = trueon those rows is lost instead of flowing up to the parent virtual key.🛠️ Suggested fix
if mg.HasColumn(&tables.TableBudget{}, "calendar_aligned") { if err := tx.Exec(` UPDATE governance_virtual_keys SET calendar_aligned = true WHERE id IN ( - SELECT DISTINCT virtual_key_id FROM governance_budgets - WHERE calendar_aligned = true AND virtual_key_id IS NOT NULL + SELECT DISTINCT b.virtual_key_id + FROM governance_budgets b + WHERE b.calendar_aligned = true AND b.virtual_key_id IS NOT NULL + UNION + SELECT DISTINCT pc.virtual_key_id + FROM governance_budgets b + JOIN governance_virtual_key_provider_configs pc ON pc.id = b.provider_config_id + WHERE b.calendar_aligned = true AND b.provider_config_id IS NOT NULL ) AND calendar_aligned = false `).Error; err != nil {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 5544 - 5558, The backfill only updates governance_virtual_keys for governance_budgets rows that already have virtual_key_id set, missing budgets linked via governance_virtual_key_provider_configs; update the tx.Exec backfill to also propagate calendar_aligned=true from governance_budgets -> governance_virtual_key_provider_configs (budget_id -> provider_config_id) -> governance_provider_configs (id -> virtual_key_id) and set governance_virtual_keys.calendar_aligned=true for those virtual_key_ids (only where calendar_aligned is false), then drop the legacy governance_budgets.calendar_aligned column as before; reference the existing mg.HasColumn check and the tx.Exec block around governance_budgets, governance_virtual_key_provider_configs, governance_provider_configs and governance_virtual_keys.
🧹 Nitpick comments (1)
plugins/governance/test_utils.go (1)
92-94: Usebifrost.Ptr(id)for these owner pointers.The new fixture helpers switched to taking the address of locals. The repo convention in test utilities is to use
bifrost.Ptr(...)instead, which also removes the temporary variable.♻️ Suggested cleanup
func buildVirtualKeyWithBudget(id, value, name string, budget *configstoreTables.TableBudget) *configstoreTables.TableVirtualKey { vk := buildVirtualKey(id, value, name, true) - vkID := id - budget.VirtualKeyID = &vkID + budget.VirtualKeyID = bifrost.Ptr(id) vk.Budgets = []configstoreTables.TableBudget{*budget} @@ func buildVirtualKeyWithMultiBudgets(id, value, name string, budgets []configstoreTables.TableBudget) *configstoreTables.TableVirtualKey { vk := buildVirtualKey(id, value, name, true) for i := range budgets { - vkID := id - budgets[i].VirtualKeyID = &vkID + budgets[i].VirtualKeyID = bifrost.Ptr(id) } vk.Budgets = budgets return vk }Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability.
Also applies to: 212-215
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 92 - 94, Replace local-address pointer creation with bifrost.Ptr calls: instead of creating vkID := id and setting budget.VirtualKeyID = &vkID, call budget.VirtualKeyID = bifrost.Ptr(id); similarly, when building vk.Budgets set it using the budget value directly (no temp pointer) and replace any other &local patterns (notably the occurrences around lines referenced for vk/budget at 212-215) with bifrost.Ptr(...) to follow the repo test-utility convention.
🤖 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/rdb.go`:
- Around line 1854-1862: preloadVirtualKeyBaseRelations currently only preloads
the singular "Budget" association but the virtual-key struct now exposes Budgets
[]TableBudget; update the Preload chain in function
preloadVirtualKeyBaseRelations to also Preload("Budgets") (in addition to the
existing Preload("Budget")) so VK-level plural budgets are loaded for callers
like GetVirtualKeys/GetVirtualKey/GetVirtualKeyByValue; keep the existing
ProviderConfigs.Budgets preload as-is.
In `@plugins/governance/store.go`:
- Around line 1543-1567: The code only checks budget.VirtualKeyID and
budget.ProviderConfigID to infer calendar alignment, so config-memory budgets
with both owner FKs nil fall through and get treated as non-calendar-aligned;
update the logic in the calendarAligned resolution to add a fallback branch
that, when both budget.VirtualKeyID and budget.ProviderConfigID are nil,
iterates gs.virtualKeys (same Range used currently) and inspects each
*configstoreTables.TableVirtualKey for membership of this budget (e.g., check
vk.Budgets or the collection that lists associated budgets for a matching
budget.ID), and if found set calendarAligned = vk.CalendarAligned and stop the
Range; keep existing checks for VirtualKeyID and ProviderConfigID unchanged.
- Around line 200-212: The clone of vk is being hydrated with live budgets but
never written back, so the updated Budgets are dropped; after constructing
liveBudgets and setting clone.Budgets, assign the updated budgets back to the
original vk (e.g., set vk.Budgets = clone.Budgets or copy clone into *vk) so
GetGovernanceData() sees the hydrated values; locate the clone := *vk block and
update it to persist the modified budgets into vk after the hydration loop that
uses gs.budgets and configstoreTables.TableBudget.
In `@transports/config.schema.json`:
- Around line 1546-1563: Update the JSON Schema for the routing rule to enforce
that scope_id must be a non-empty string for any non-global scope: add a
conditional (if/then) on the "scope" property so that when "scope" is not
"global" (i.e., "team", "customer", "virtual_key") the schema requires
"scope_id" and constrains it to type "string" with minimum length > 0 (disallow
null and empty string); apply the same conditional enforcement for the other
occurrence referenced (the second "scope_id" block) so both schema locations
validate non-global scopes consistently.
- Around line 298-305: Add a JSON Schema constraint to enforce exclusivity
between virtual_key_id and provider_config_id: update
transports/config.schema.json around the budget object (the properties
virtual_key_id and provider_config_id) to reject documents that contain both
keys by adding a "not": {"required": ["virtual_key_id","provider_config_id"]}
(or equivalently a oneOf that forbids both) at the same schema level so the
validator fails when both virtual_key_id and provider_config_id are present.
---
Duplicate comments:
In @.github/workflows/scripts/run-migration-tests.sh:
- Around line 2624-2627: The snapshot comparator is flagging legacy removals of
budget_id as unexpected; update the comparator used in
compare_postgres_snapshots() to treat budget_id as an allowed/ignored dropped
column by adding "budget_id" to the ignore_columns list (the local variable
named ignore_columns in the script) so removals of
governance_virtual_keys.budget_id and
governance_virtual_key_provider_configs.budget_id are not considered failures
during snapshot comparison.
- Around line 2837-2913: The SQLite validation path is missing the new
multi-budget assertions added in verify_budget_migration_postgres(); update
validate_sqlite_data (or add a new verify_budget_migration_sqlite and call it
from the SQLite flow) to run the same checks: ensure
governance_budgets.virtual_key_id/provider_config_id existence and proper values
for the test rows, confirm governance_virtual_keys and
governance_virtual_key_provider_configs no longer have budget_id, and ensure
junction table governance_virtual_key_budgets is dropped (or warn if present);
mirror the queries and logging logic from verify_budget_migration_postgres() so
SQLite migrations are validated the same way.
In `@framework/configstore/migrations.go`:
- Around line 5544-5558: The backfill only updates governance_virtual_keys for
governance_budgets rows that already have virtual_key_id set, missing budgets
linked via governance_virtual_key_provider_configs; update the tx.Exec backfill
to also propagate calendar_aligned=true from governance_budgets ->
governance_virtual_key_provider_configs (budget_id -> provider_config_id) ->
governance_provider_configs (id -> virtual_key_id) and set
governance_virtual_keys.calendar_aligned=true for those virtual_key_ids (only
where calendar_aligned is false), then drop the legacy
governance_budgets.calendar_aligned column as before; reference the existing
mg.HasColumn check and the tx.Exec block around governance_budgets,
governance_virtual_key_provider_configs, governance_provider_configs and
governance_virtual_keys.
In `@framework/configstore/rdb.go`:
- Around line 2327-2339: The single-item provider-config delete path must mirror
the transactional bulk-delete in DeleteVirtualKey: run the join-table cleanup,
budget, provider-config and rate-limit deletions inside the same txDB
transaction and in the same order to avoid orphaned joins or partial deletes.
Modify the code that uses txDB.WithContext(ctx).Delete for tables.TableBudget,
tables.TableVirtualKeyProviderConfig{}, and tables.TableRateLimit{} to also
delete rows from the governance_virtual_key_provider_config_keys join table
(using the provider config id) and ensure all four deletes execute on the same
txDB transaction context (ctx/txDB) so any failure rolls back the entire
operation.
In `@framework/configstore/tables/virtualkey.go`:
- Around line 33-38: The JSON decoder currently drops legacy single-budget
fields (budget_id / budget and budget.calendar_aligned) causing Budgets on
TableVirtualKey to be empty and CalendarAligned lost; add a one-release decoder
shim by implementing a custom UnmarshalJSON for TableVirtualKey that first
attempts normal decoding, then checks for legacy keys (budget_id or budget
object with calendar_aligned) and, if present, converts them into a single-entry
Budgets slice and sets TableVirtualKey.CalendarAligned accordingly; ensure the
shim only affects JSON decoding paths (not DB migrations) and remove after the
release.
In `@plugins/governance/store_test.go`:
- Around line 261-263: Replace assert.Error calls that are followed by
dereferencing err (e.g., the call to store.CheckBudget(ctx, vk,
&EvaluationRequest{Provider: schemas.OpenAI}, nil) followed by
assert.Contains(t, err.Error(), ...)) with require.Error so the test stops when
err is nil; do this for the instances around the CheckBudget assertions (current
snippet and the other occurrences mentioned at lines ~287-289, ~337-339,
~363-365, ~394-396) so you don’t call err.Error() on a nil error and hide the
real failure.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 1201-1204: The MultiBudgetLines instance uses a 2-part data-testid
"vk-budget"; update it to a 3-part test id following the
<entity>-<element>-<qualifier> convention (e.g., change data-testid on the
MultiBudgetLines with id "vkBudget" to "virtual-key-budget-control" or
"virtual-key-budget-config") so it matches the ui/**/*.{tsx,ts} guideline for
new interactive UI elements.
- Around line 1179-1185: The icon-only MCP remove Button (rendered with
onClick={handleRemoveMCPClient} and data-testid={`vk-delete-mcp-${index}`})
lacks an accessible name; add an aria-label or equivalent accessibleName (for
example aria-label={`Remove MCP client ${index}`} or a meaningful string) to the
Button component so screen readers announce its purpose while retaining the
existing onClick handler and data-testid.
- Around line 208-214: In virtualKeySheet.tsx the mapping that builds rate_limit
currently uses truthiness checks for token_max_limit and request_max_limit so
numeric zero is treated as absent; in the object construction for rate_limit
(the lines assigning token_max_limit and request_max_limit) change the checks to
explicitly test for undefined/null (e.g. use !== undefined or != null) and then
String(...) the value, so a value of 0 is preserved as "0" instead of becoming
undefined during the edit/save round-trip.
- Around line 917-936: Multiple sequential calls to handleUpdateProviderConfig
(for "budget" and "budgets") use the same providerConfigs snapshot and can
overwrite each other; instead build a single update object and call
handleUpdateProviderConfig once for the given index to set both fields
atomically. In the onChange handler (referencing handleUpdateProviderConfig and
index) compute the newBudget (or undefined) and newBudgets array from lines,
then call handleUpdateProviderConfig(index, { budget: newBudget, budgets:
newBudgets }) in a single invocation so both fields are updated together.
In `@ui/lib/types/governance.ts`:
- Around line 156-179: Change the budgets element type on
CreateVirtualKeyRequest and UpdateVirtualKeyRequest so callers cannot set
per-budget calendar_aligned; replace CreateBudgetRequest[] with a budget type
that omits calendar_aligned (e.g., CreateBudgetNoCalendarRequest[] or a mapped
Omit<CreateBudgetRequest,'calendar_aligned'>[]) and update any imports/exports
accordingly; adjust the symbols CreateVirtualKeyRequest and
UpdateVirtualKeyRequest to reference the new budget type so only the top-level
calendar_aligned remains the source of truth.
---
Nitpick comments:
In `@plugins/governance/test_utils.go`:
- Around line 92-94: Replace local-address pointer creation with bifrost.Ptr
calls: instead of creating vkID := id and setting budget.VirtualKeyID = &vkID,
call budget.VirtualKeyID = bifrost.Ptr(id); similarly, when building vk.Budgets
set it using the budget value directly (no temp pointer) and replace any other
&local patterns (notably the occurrences around lines referenced for vk/budget
at 212-215) with bifrost.Ptr(...) to follow the repo test-utility convention.
🪄 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: 63ec2414-f4c7-443d-876c-d94c61a91a1b
⛔ Files ignored due to path filters (9)
docs/media/user-provisioning/zitadel-add-role.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-add-user-select-key.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-auth-method.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-namne.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-uri.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-role-assignemnt.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-select-project.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-token-config.pngis excluded by!**/*.png
📒 Files selected for processing (33)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shdocs/openapi/openapi.yamldocs/openapi/paths/management/users.yamldocs/openapi/schemas/management/users.yamlframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/sidebar.tsxui/components/ui/multiBudgetLines.tsxui/lib/store/apis/baseApi.tsui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (8)
- .claude/skills/expect
- ui/lib/store/apis/baseApi.ts
- docs/openapi/openapi.yaml
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- docs/openapi/paths/management/users.yaml
- .agents/skills/expect/SKILL.md
- docs/openapi/schemas/management/users.yaml
- ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
🚧 Files skipped from review as they are similar to previous changes (10)
- ui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsx
- transports/bifrost-http/handlers/governance_test.go
- ui/app/workspace/governance/access-profiles/page.tsx
- ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
- plugins/governance/resolver_test.go
- framework/configstore/tables/budget.go
- framework/configstore/rdb_test.go
- ui/components/ui/multiBudgetLines.tsx
- transports/bifrost-http/lib/config_test.go
- transports/bifrost-http/handlers/governance.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
3e8d874 to
5f285e3
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (10)
ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx (1)
311-313:⚠️ Potential issue | 🟡 MinorRender the calendar suffix from
vk.calendar_aligned.This stack moved
calendar_alignedto the virtual key, sob.calendar_alignedstays falsy and the"(calendar)"label disappears for every budget row.🔧 Suggested fix
- {b.calendar_aligned && " (calendar)"} + {vk.calendar_aligned && " (calendar)"}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx` around lines 311 - 313, The calendar suffix is rendering from the wrong object—budget rows check b.calendar_aligned which is now moved to the virtual key; update the render in virtualKeysTable.tsx to use vk.calendar_aligned instead of b.calendar_aligned (the surrounding JSX that displays Resets {formatResetDuration(b.reset_duration)} should keep b.reset_duration but switch the conditional to vk.calendar_aligned so the "(calendar)" label appears correctly).framework/configstore/rdb.go (2)
1861-1868:⚠️ Potential issue | 🟠 MajorPreload VK-level
Budgetshere as well.Callers using
preloadVirtualKeyBaseRelations()now readvirtualKey.Budgets, but this chain still only pulls the singularBudgetrelation plus provider-config budgets.GetVirtualKeys(),GetVirtualKey(), andGetVirtualKeyByValue()will otherwise return empty VK budgets after the multi-budget migration.🔧 Suggested fix
return db. + Preload("Budgets"). Preload("Budget"). Preload("RateLimit"). Preload("ProviderConfigs"). Preload("ProviderConfigs.Budgets").🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 1861 - 1868, preloadVirtualKeyBaseRelations currently preloads the singular Budget and provider-config budgets but not the VirtualKey-level Budgets, causing VirtualKey.Budgets to be empty after the multi-budget migration; update the preload chain in preloadVirtualKeyBaseRelations to include Preload("Budgets") (alongside the existing Preload("Budget") and Preload("ProviderConfigs.Budgets")) so callers like GetVirtualKeys, GetVirtualKey, and GetVirtualKeyByValue receive the VK-level Budgets populated.
2319-2347:⚠️ Potential issue | 🟠 MajorMirror the bulk VK delete cleanup when removing one provider config.
This path still deletes budgets/config/rate-limits without a local transaction and never clears
governance_virtual_key_provider_config_keys. If the join rows remain or a later delete fails, you can leave a partially deleted provider-config state.🔧 Suggested fix
func (s *RDBConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } + if len(tx) == 0 { + return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + return s.DeleteVirtualKeyProviderConfig(ctx, id, tx) + }) + } + txDB := tx[0] // First fetch the provider config to get budget and rate limit IDs var providerConfig tables.TableVirtualKeyProviderConfig if err := txDB.WithContext(ctx).First(&providerConfig, "id = ?", id).Error; err != nil { ... } + if err := txDB.WithContext(ctx).Exec( + "DELETE FROM governance_virtual_key_provider_config_keys WHERE table_virtual_key_provider_config_id = ?", + id, + ).Error; err != nil { + return err + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2319 - 2347, DeleteVirtualKeyProviderConfig currently performs multiple deletes (budgets, provider config, rate limit) without a local transaction and never removes join rows from governance_virtual_key_provider_config_keys; change it to start a transaction when no tx is provided (using txDB = s.db.Begin()/txDB.Commit()/txDB.Rollback()), run all deletes inside that transaction, ensure you Delete join rows from governance_virtual_key_provider_config_keys (e.g. txDB.Where("provider_config_id = ?", id).Delete(&tables.TableVirtualKeyProviderConfigKeys{} or raw table name) before deleting TableVirtualKeyProviderConfig, and only commit if all operations (deleting TableBudget, governance_virtual_key_provider_config_keys, TableVirtualKeyProviderConfig, and TableRateLimit when rateLimitID != nil) succeed, rolling back on any error.plugins/governance/test_utils.go (1)
210-217:⚠️ Potential issue | 🟠 MajorSeed the default provider config in
buildVirtualKeyWithMultiBudgets().Unlike
buildVirtualKeyWithBudget()andbuildVirtualKeyWithRateLimit(), this helper still returns a VK with no provider configs. Resolver tests that use it can short-circuit at provider selection instead of exercising the multi-budget path.🧪 Suggested fix
func buildVirtualKeyWithMultiBudgets(id, value, name string, budgets []configstoreTables.TableBudget) *configstoreTables.TableVirtualKey { vk := buildVirtualKey(id, value, name, true) for i := range budgets { vkID := id budgets[i].VirtualKeyID = &vkID } vk.Budgets = budgets + vk.ProviderConfigs = []configstoreTables.TableVirtualKeyProviderConfig{ + buildProviderConfig("openai", []string{"*"}), + } return vk }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 210 - 217, buildVirtualKeyWithMultiBudgets returns a TableVirtualKey with budgets but no provider configs, causing resolver tests to short-circuit at provider selection; modify buildVirtualKeyWithMultiBudgets to seed the default provider config the same way buildVirtualKeyWithBudget() and buildVirtualKeyWithRateLimit() do (i.e., create and attach a default provider config entry to the TableVirtualKey.ProviderConfigs or equivalent field after calling buildVirtualKey and before returning) and ensure each TableBudget.VirtualKeyID is set to the vk ID as currently done.framework/configstore/tables/virtualkey.go (2)
37-38:⚠️ Potential issue | 🟠 MajorKeep a one-release legacy JSON decode path for singular
budgetpayloads.Removing singular budget fields at this JSON boundary can silently drop legacy single-budget payloads unless they are lifted into
Budgets. Add a temporary compatibility shim so old config shapes still deserialize correctly during transition.🧩 Suggested compatibility shim (pattern)
func (pc *TableVirtualKeyProviderConfig) UnmarshalJSON(data []byte) error { type Alias TableVirtualKeyProviderConfig + type legacyBudget struct { + TableBudget + } type TempProviderConfig struct { Alias - KeyIDs []string `json:"key_ids"` + KeyIDs []string `json:"key_ids"` + Budget *legacyBudget `json:"budget"` // legacy single-budget shape } var temp TempProviderConfig if err := json.Unmarshal(data, &temp); err != nil { return err } *pc = TableVirtualKeyProviderConfig(temp.Alias) + if len(pc.Budgets) == 0 && temp.Budget != nil { + pc.Budgets = []TableBudget{temp.Budget.TableBudget} + } // existing key_ids conversion... return nil }// Add similar shim on TableVirtualKey: // - read legacy "budget" // - if Budgets is empty, lift into Budgets[0] // - map legacy budget.calendar_aligned to vk.CalendarAligned when presentAlso applies to: 222-223
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` around lines 37 - 38, Add a temporary JSON-compat shim on TableVirtualKey by implementing a custom UnmarshalJSON that reads legacy singular "budget" payloads: if the incoming JSON has "budget" and the struct's Budgets slice is empty, lift that legacy budget into Budgets[0]; also map legacy budget.calendar_aligned to the TableVirtualKey.CalendarAligned field when present. Ensure the shim only runs during unmarshalling and preserves normal behavior for the current "budgets" array shape so old single-budget config shapes deserialize correctly into Budgets.
38-38:⚠️ Potential issue | 🟡 MinorFix contradictory
Keyssemantics comment.The comment says empty
keysallows all, but code/documented behavior is deny-by-default unlessAllowAllKeysis true.✏️ Suggested comment fix
- Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty means all keys allowed for this provider + Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty with AllowAllKeys=false means no keys allowed (deny-by-default)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/tables/virtualkey.go` at line 38, The comment on the Keys field is incorrect: update the comment for Keys to reflect that absence of entries is treated as deny-by-default unless the AllowAllKeys boolean is true; specifically, in the struct containing Keys and AllowAllKeys (the provider config in virtualkey.go) change the comment from "Empty means all keys allowed for this provider" to something like "Empty means no keys are allowed unless AllowAllKeys is true" so the semantics of Keys and AllowAllKeys are accurate and unambiguous.plugins/governance/store_test.go (1)
261-264:⚠️ Potential issue | 🟡 MinorUse
require.Errorbefore callingerr.Error()in these failure-path assertions.These blocks can panic and hide the real failure if
erris unexpectedly nil. Replaceassert.Errorwithrequire.Errorbefore dereferencingerr.🛠️ Suggested fix pattern
err = store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil) -assert.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") +require.Error(t, err, "Should fail when hourly budget is exceeded even though daily is fine") assert.Contains(t, err.Error(), "budget exceeded")Also applies to: 287-290, 337-340, 363-366, 394-396
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 261 - 264, The test is dereferencing err (err.Error()) after using assert.Error which can be nil and cause a panic; replace assert.Error(t, err, ...) with require.Error(t, err, ...) before any err.Error() usage in the failing-path assertions (e.g., the block calling store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider: schemas.OpenAI}, nil) and the other similar blocks at the indicated locations) so the test stops immediately on nil err; update all instances (the one shown and the ones at 287-290, 337-340, 363-366, 394-396) to use require.Error and keep the subsequent assert.Contains(t, err.Error(), "budget exceeded") as-is.transports/bifrost-http/lib/config_test.go (1)
243-243:⚠️ Potential issue | 🟠 MajorRestore SQLite round-trip coverage for provider-config budgets.
Line 243 now points this section to a hash-unit test, but the multi-budget provider-config SQLite write/read/reload path is still not covered here. This remains an important regression risk for the stacked ownership migration.
Based on learnings: In tests under transports/bifrost-http/lib/config_test.go, keep MockConfigStore methods simple and cover behavior with SQLite-backed integration tests using createTestSQLiteConfigStore instead to validate end-to-end behavior.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/lib/config_test.go` at line 243, The unit test TestGenerateVirtualKeyHash_ProviderConfigRateLimit was changed to exercise hashing only, but you should restore end-to-end SQLite round-trip coverage for provider-config budgets: simplify MockConfigStore implementations (avoid simulating persistence) and add/instrument an integration test that uses createTestSQLiteConfigStore to write a multi-budget provider-config, reload it from SQLite, and assert the budgets and stacked ownership are preserved; locate MockConfigStore and TestGenerateVirtualKeyHash_ProviderConfigRateLimit in transports/bifrost-http/lib/config_test.go and replace the mocked persistence checks with a run that uses createTestSQLiteConfigStore to validate the full write/read/reload path.plugins/governance/store.go (2)
1543-1567:⚠️ Potential issue | 🟠 MajorCalendar-aligned resets can be skipped for config-memory budgets with nil owner FKs.
At lines 1546 and 1554, reset mode is inferred only from
budget.VirtualKeyID/budget.ProviderConfigID. In config-memory mode (loaded vialoadFromConfigMemory), those owner fields can be nil because they're not explicitly populated during relationship wiring. This causes a false rolling-reset path even when the owning VK is calendar-aligned.🛡️ Proposed fix (fallback owner resolution by budget membership)
calendarAligned := false + ownerResolved := false if budget.VirtualKeyID != nil { gs.virtualKeys.Range(func(_, v interface{}) bool { if vk, ok := v.(*configstoreTables.TableVirtualKey); ok && vk != nil && vk.ID == *budget.VirtualKeyID { calendarAligned = vk.CalendarAligned + ownerResolved = true return false // stop } return true }) } else if budget.ProviderConfigID != nil { // Provider config budgets: look up the VK that owns this provider config gs.virtualKeys.Range(func(_, v interface{}) bool { if vk, ok := v.(*configstoreTables.TableVirtualKey); ok && vk != nil { for _, pc := range vk.ProviderConfigs { if pc.ID == *budget.ProviderConfigID { calendarAligned = vk.CalendarAligned + ownerResolved = true return false // stop } } } return true }) } + // Fallback for config-memory budgets that may not have owner FK fields set + if !ownerResolved { + gs.virtualKeys.Range(func(_, v interface{}) bool { + vk, ok := v.(*configstoreTables.TableVirtualKey) + if !ok || vk == nil { + return true + } + for _, b := range vk.Budgets { + if b.ID == budget.ID { + calendarAligned = vk.CalendarAligned + return false + } + } + for _, pc := range vk.ProviderConfigs { + for _, b := range pc.Budgets { + if b.ID == budget.ID { + calendarAligned = vk.CalendarAligned + return false + } + } + } + return true + }) + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 1543 - 1567, The current lookup only checks budget.VirtualKeyID and budget.ProviderConfigID, which fails for config-memory budgets where those FKs are nil; update the gs.virtualKeys.Range scan (using configstoreTables.TableVirtualKey and vk.ProviderConfigs) to add a fallback path when both owner FKs are nil: iterate all vk entries and check membership (e.g., check any vk.ProviderConfigs for a matching budget.Owner relationship and/or any vk.Budgets or budget list field for an entry matching budget.ID) and if found set calendarAligned = vk.CalendarAligned and stop; keep using gs.virtualKeys.Range and the same vk.CalendarAligned assignment so calendar-aligned state is resolved even for config-memory-loaded budgets.
198-212:⚠️ Potential issue | 🟠 MajorHydrated VK budgets are dropped due to local copy not being written back.
At line 200,
clone := *vkcreates a local copy for budget hydration, butclone.Budgetsis never assigned back tovk. The rate limit (line 216) and provider configs (line 244) are correctly written back tovk, but VK-level budgets are not, causingGetGovernanceData()to return stale VK-level budget usage.🐛 Proposed fix
- clone := *vk - // Hydrate multi-budgets from live sync.Map - if len(clone.Budgets) > 0 { - liveBudgets := make([]configstoreTables.TableBudget, 0, len(clone.Budgets)) - for _, b := range clone.Budgets { + // Hydrate multi-budgets from live sync.Map + if len(vk.Budgets) > 0 { + liveBudgets := make([]configstoreTables.TableBudget, 0, len(vk.Budgets)) + for _, b := range vk.Budgets { if lb, exists := gs.budgets.Load(b.ID); exists && lb != nil { if budget, ok := lb.(*configstoreTables.TableBudget); ok { liveBudgets = append(liveBudgets, *budget) } } } - clone.Budgets = liveBudgets + vk.Budgets = liveBudgets }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 198 - 212, The local copy created by "clone := *vk" is hydrated with live budgets but never written back to the original VK, so VK-level budgets remain stale; after populating clone.Budgets from gs.budgets, assign the updated budgets back to the original value (e.g., update vk.Budgets or write the modified clone back into *vk) so the hydrated budgets are persisted; ensure you modify the same code block that updates rate limits and provider configs (the area handling clone, gs.budgets, and subsequent writes) so VK-level budgets are treated the same as rateLimits/provider configs.
🧹 Nitpick comments (1)
plugins/governance/test_utils.go (1)
92-93: Preferbifrost.Ptr(id)for these owner ID pointers.It keeps the test builders consistent with the rest of the repo and removes the throwaway local variable.
♻️ Suggested cleanup
func buildVirtualKeyWithBudget(id, value, name string, budget *configstoreTables.TableBudget) *configstoreTables.TableVirtualKey { vk := buildVirtualKey(id, value, name, true) - vkID := id - budget.VirtualKeyID = &vkID + budget.VirtualKeyID = bifrost.Ptr(id) vk.Budgets = []configstoreTables.TableBudget{*budget}for i := range budgets { - vkID := id - budgets[i].VirtualKeyID = &vkID + budgets[i].VirtualKeyID = bifrost.Ptr(id) }Based on learnings: prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically.
Also applies to: 213-214
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/test_utils.go` around lines 92 - 93, Replace the throwaway local pointer creation for owner IDs with the repository-standard helper: instead of creating vkID := id and assigning budget.VirtualKeyID = &vkID, call bifrost.Ptr(id) and assign budget.VirtualKeyID = bifrost.Ptr(id); apply the same change for the other occurrence referenced at lines 213-214 so all owner ID pointer assignments use bifrost.Ptr(...) (referencing symbols: id, vkID, budget.VirtualKeyID, bifrost.Ptr).
🤖 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/lib/types/governance.ts`:
- Around line 163-166: CreateBudgetRequest still defines calendar_aligned?:
boolean which allows callers to set budgets[i].calendar_aligned alongside the
top-level calendar_aligned on CreateVirtualKeyRequest/UpdateVirtualKeyRequest,
creating conflicting sources of truth; update the types so that the budgets
field on CreateVirtualKeyRequest and UpdateVirtualKeyRequest uses a budget type
that omits calendar_aligned (e.g., a new CreateBudgetForVK or by using
Omit<CreateBudgetRequest,'calendar_aligned'>) and remove calendar_aligned from
the per-budget shape, ensuring only the VK-level calendar_aligned remains, and
adjust any references to CreateBudgetRequest in the file accordingly (symbols:
CreateBudgetRequest, CreateVirtualKeyRequest, UpdateVirtualKeyRequest, budgets,
calendar_aligned).
---
Duplicate comments:
In `@framework/configstore/rdb.go`:
- Around line 1861-1868: preloadVirtualKeyBaseRelations currently preloads the
singular Budget and provider-config budgets but not the VirtualKey-level
Budgets, causing VirtualKey.Budgets to be empty after the multi-budget
migration; update the preload chain in preloadVirtualKeyBaseRelations to include
Preload("Budgets") (alongside the existing Preload("Budget") and
Preload("ProviderConfigs.Budgets")) so callers like GetVirtualKeys,
GetVirtualKey, and GetVirtualKeyByValue receive the VK-level Budgets populated.
- Around line 2319-2347: DeleteVirtualKeyProviderConfig currently performs
multiple deletes (budgets, provider config, rate limit) without a local
transaction and never removes join rows from
governance_virtual_key_provider_config_keys; change it to start a transaction
when no tx is provided (using txDB =
s.db.Begin()/txDB.Commit()/txDB.Rollback()), run all deletes inside that
transaction, ensure you Delete join rows from
governance_virtual_key_provider_config_keys (e.g. txDB.Where("provider_config_id
= ?", id).Delete(&tables.TableVirtualKeyProviderConfigKeys{} or raw table name)
before deleting TableVirtualKeyProviderConfig, and only commit if all operations
(deleting TableBudget, governance_virtual_key_provider_config_keys,
TableVirtualKeyProviderConfig, and TableRateLimit when rateLimitID != nil)
succeed, rolling back on any error.
In `@framework/configstore/tables/virtualkey.go`:
- Around line 37-38: Add a temporary JSON-compat shim on TableVirtualKey by
implementing a custom UnmarshalJSON that reads legacy singular "budget"
payloads: if the incoming JSON has "budget" and the struct's Budgets slice is
empty, lift that legacy budget into Budgets[0]; also map legacy
budget.calendar_aligned to the TableVirtualKey.CalendarAligned field when
present. Ensure the shim only runs during unmarshalling and preserves normal
behavior for the current "budgets" array shape so old single-budget config
shapes deserialize correctly into Budgets.
- Line 38: The comment on the Keys field is incorrect: update the comment for
Keys to reflect that absence of entries is treated as deny-by-default unless the
AllowAllKeys boolean is true; specifically, in the struct containing Keys and
AllowAllKeys (the provider config in virtualkey.go) change the comment from
"Empty means all keys allowed for this provider" to something like "Empty means
no keys are allowed unless AllowAllKeys is true" so the semantics of Keys and
AllowAllKeys are accurate and unambiguous.
In `@plugins/governance/store_test.go`:
- Around line 261-264: The test is dereferencing err (err.Error()) after using
assert.Error which can be nil and cause a panic; replace assert.Error(t, err,
...) with require.Error(t, err, ...) before any err.Error() usage in the
failing-path assertions (e.g., the block calling
store.CheckBudget(context.Background(), vk, &EvaluationRequest{Provider:
schemas.OpenAI}, nil) and the other similar blocks at the indicated locations)
so the test stops immediately on nil err; update all instances (the one shown
and the ones at 287-290, 337-340, 363-366, 394-396) to use require.Error and
keep the subsequent assert.Contains(t, err.Error(), "budget exceeded") as-is.
In `@plugins/governance/store.go`:
- Around line 1543-1567: The current lookup only checks budget.VirtualKeyID and
budget.ProviderConfigID, which fails for config-memory budgets where those FKs
are nil; update the gs.virtualKeys.Range scan (using
configstoreTables.TableVirtualKey and vk.ProviderConfigs) to add a fallback path
when both owner FKs are nil: iterate all vk entries and check membership (e.g.,
check any vk.ProviderConfigs for a matching budget.Owner relationship and/or any
vk.Budgets or budget list field for an entry matching budget.ID) and if found
set calendarAligned = vk.CalendarAligned and stop; keep using
gs.virtualKeys.Range and the same vk.CalendarAligned assignment so
calendar-aligned state is resolved even for config-memory-loaded budgets.
- Around line 198-212: The local copy created by "clone := *vk" is hydrated with
live budgets but never written back to the original VK, so VK-level budgets
remain stale; after populating clone.Budgets from gs.budgets, assign the updated
budgets back to the original value (e.g., update vk.Budgets or write the
modified clone back into *vk) so the hydrated budgets are persisted; ensure you
modify the same code block that updates rate limits and provider configs (the
area handling clone, gs.budgets, and subsequent writes) so VK-level budgets are
treated the same as rateLimits/provider configs.
In `@plugins/governance/test_utils.go`:
- Around line 210-217: buildVirtualKeyWithMultiBudgets returns a TableVirtualKey
with budgets but no provider configs, causing resolver tests to short-circuit at
provider selection; modify buildVirtualKeyWithMultiBudgets to seed the default
provider config the same way buildVirtualKeyWithBudget() and
buildVirtualKeyWithRateLimit() do (i.e., create and attach a default provider
config entry to the TableVirtualKey.ProviderConfigs or equivalent field after
calling buildVirtualKey and before returning) and ensure each
TableBudget.VirtualKeyID is set to the vk ID as currently done.
In `@transports/bifrost-http/lib/config_test.go`:
- Line 243: The unit test TestGenerateVirtualKeyHash_ProviderConfigRateLimit was
changed to exercise hashing only, but you should restore end-to-end SQLite
round-trip coverage for provider-config budgets: simplify MockConfigStore
implementations (avoid simulating persistence) and add/instrument an integration
test that uses createTestSQLiteConfigStore to write a multi-budget
provider-config, reload it from SQLite, and assert the budgets and stacked
ownership are preserved; locate MockConfigStore and
TestGenerateVirtualKeyHash_ProviderConfigRateLimit in
transports/bifrost-http/lib/config_test.go and replace the mocked persistence
checks with a run that uses createTestSQLiteConfigStore to validate the full
write/read/reload path.
In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx`:
- Around line 311-313: The calendar suffix is rendering from the wrong
object—budget rows check b.calendar_aligned which is now moved to the virtual
key; update the render in virtualKeysTable.tsx to use vk.calendar_aligned
instead of b.calendar_aligned (the surrounding JSX that displays Resets
{formatResetDuration(b.reset_duration)} should keep b.reset_duration but switch
the conditional to vk.calendar_aligned so the "(calendar)" label appears
correctly).
---
Nitpick comments:
In `@plugins/governance/test_utils.go`:
- Around line 92-93: Replace the throwaway local pointer creation for owner IDs
with the repository-standard helper: instead of creating vkID := id and
assigning budget.VirtualKeyID = &vkID, call bifrost.Ptr(id) and assign
budget.VirtualKeyID = bifrost.Ptr(id); apply the same change for the other
occurrence referenced at lines 213-214 so all owner ID pointer assignments use
bifrost.Ptr(...) (referencing symbols: id, vkID, budget.VirtualKeyID,
bifrost.Ptr).
🪄 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: 02dd8641-185c-44cb-a107-2b71cc6232f5
⛔ Files ignored due to path filters (9)
docs/media/user-provisioning/zitadel-add-role.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-add-user-select-key.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-auth-method.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-namne.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app-uri.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-create-app.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-role-assignemnt.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-select-project.pngis excluded by!**/*.pngdocs/media/user-provisioning/zitadel-token-config.pngis excluded by!**/*.png
📒 Files selected for processing (33)
.agents/skills/expect/SKILL.md.claude/skills/expect.github/workflows/scripts/run-migration-tests.shdocs/openapi/openapi.yamldocs/openapi/paths/management/users.yamldocs/openapi/schemas/management/users.yamlframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/tables/budget.goframework/configstore/tables/virtualkey.goplugins/governance/modelprovidergovernance_test.goplugins/governance/resolver.goplugins/governance/resolver_test.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/governance_test.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/config.schema.jsonui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsxui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsxui/app/workspace/governance/access-profiles/page.tsxui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/components/sidebar.tsxui/components/ui/multiBudgetLines.tsxui/lib/store/apis/baseApi.tsui/lib/types/governance.ts
💤 Files with no reviewable changes (2)
- transports/bifrost-http/lib/config.go
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (10)
- .claude/skills/expect
- ui/lib/store/apis/baseApi.ts
- ui/app/_fallbacks/enterprise/lib/contexts/rbacContext.tsx
- ui/app/_fallbacks/enterprise/components/access-profiles/accessProfilesIndexView.tsx
- docs/openapi/openapi.yaml
- docs/openapi/schemas/management/users.yaml
- docs/openapi/paths/management/users.yaml
- .agents/skills/expect/SKILL.md
- .github/workflows/scripts/run-migration-tests.sh
- transports/bifrost-http/handlers/governance.go
🚧 Files skipped from review as they are similar to previous changes (10)
- plugins/governance/modelprovidergovernance_test.go
- plugins/governance/resolver.go
- transports/bifrost-http/handlers/governance_test.go
- ui/app/workspace/governance/access-profiles/page.tsx
- framework/configstore/tables/budget.go
- ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
- framework/configstore/rdb_test.go
- transports/config.schema.json
- ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
- framework/configstore/migrations.go
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/enterprise/setting-up-okta.mdx (1)
209-209:⚠️ Potential issue | 🟡 MinorIncorrect step reference after renumbering.
The note references "Steps 4-7" but this includes Step 6 (Assign Users) and Step 7 (API token), which doesn't align with the intent. The note seems to describe skipping the optional role/group claim configuration — likely should reference "Steps 3-5" (Custom Role Attribute, Role Claim, and Groups).
📝 Proposed fix
-Role claims are available only when you configure custom claims on your authorization server. Ensure you add role claims to your chosen authorization server (for example, `/oauth2/default`) to enable RBAC. If you skipped Steps 4-7, the first user to sign in automatically receives the **Admin** role and can manage RBAC for all subsequent users through the Bifrost dashboard. +Role claims are available only when you configure custom claims on your authorization server. Ensure you add role claims to your chosen authorization server (for example, `/oauth2/default`) to enable RBAC. If you skipped Steps 3-5, the first user to sign in automatically receives the **Admin** role and can manage RBAC for all subsequent users through the Bifrost dashboard.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/enterprise/setting-up-okta.mdx` at line 209, Update the incorrect step reference in the note that currently says "Steps 4-7" to the correct range "Steps 3-5" so the sentence about skipping optional role/group claim configuration correctly points to the Custom Role Attribute, Role Claim, and Groups steps; locate the sentence containing "Steps 4-7" in the docs/enterprise/setting-up-okta.mdx content and replace that fragment with "Steps 3-5" (ensure surrounding wording remains unchanged).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/enterprise/setting-up-okta.mdx`:
- Line 77: Fix the subject-verb agreement in the sentence "Roles takes
precedence over groups in role assignment" by changing "takes" to "take" so it
reads "Roles take precedence over groups in role assignment"; update the
sentence in the docs content where that phrase appears.
- Around line 214-229: Step 7 tells users to create an API token but Step 8 and
the Configuration Reference never surface that token; add a clear configuration
field for it. Update the Step 8 configuration table (the "Step 8" section) to
include an "Okta API Token" (or similarly named) entry, and add the matching
entry in the "Configuration Reference" documentation with its expected key/name,
description, and usage (e.g., used for bulk user and team sync); alternatively,
if the token is consumed elsewhere, add a note in Step 7 pointing users exactly
to the config key or location where the token should be pasted.
In `@ui/app/_fallbacks/enterprise/components/user-groups/businessUnitsView.tsx`:
- Around line 7-13: The ContactUsView instance is missing the testIdPrefix prop,
so add testIdPrefix="business-units" to the ContactUsView JSX in
businessUnitsView.tsx (the ContactUsView component is the unique symbol to
update) so the fallback renders stable data-testid hooks (following the pattern
data-testid="<entity>-<element>-<qualifier>") for the business-units screen;
ensure the prop name is exactly testIdPrefix and that ContactUsView's internal
controls will use that prefix.
In `@ui/app/_fallbacks/enterprise/components/user-groups/teamsView.tsx`:
- Around line 7-13: The ContactUsView instance lacks the testIdPrefix prop;
update the ContactUsView call to include a stable prefix (e.g. pass
testIdPrefix="teams-contact" or similar) so its internal elements render
predictable data-testid attributes; ensure the chosen prefix follows the project
pattern (entity-element-qualifier) so E2E selectors like
data-testid="teams-contact-readme-link" or
data-testid="teams-contact-submit-button" will be produced.
---
Outside diff comments:
In `@docs/enterprise/setting-up-okta.mdx`:
- Line 209: Update the incorrect step reference in the note that currently says
"Steps 4-7" to the correct range "Steps 3-5" so the sentence about skipping
optional role/group claim configuration correctly points to the Custom Role
Attribute, Role Claim, and Groups steps; locate the sentence containing "Steps
4-7" in the docs/enterprise/setting-up-okta.mdx content and replace that
fragment with "Steps 3-5" (ensure surrounding wording remains unchanged).
🪄 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: 7b56020a-7e86-42ae-bdcd-a942a4022754
⛔ Files ignored due to path filters (3)
docs/media/user-provisioning/okta-api-token-created.pngis excluded by!**/*.pngdocs/media/user-provisioning/okta-create-token-form.pngis excluded by!**/*.pngdocs/media/user-provisioning/okta-tokens-screen.pngis excluded by!**/*.png
📒 Files selected for processing (9)
docs/enterprise/setting-up-okta.mdxframework/configstore/rdb.goplugins/governance/test_utils.goui/app/_fallbacks/enterprise/components/user-groups/businessUnitsView.tsxui/app/_fallbacks/enterprise/components/user-groups/teamsView.tsxui/app/workspace/dashboard/components/charts/modelFilterSelect.tsxui/app/workspace/governance/business-units/page.tsxui/app/workspace/governance/teams/page.tsxui/components/sidebar.tsx
✅ Files skipped from review due to trivial changes (1)
- ui/app/workspace/dashboard/components/charts/modelFilterSelect.tsx
🚧 Files skipped from review as they are similar to previous changes (2)
- ui/components/sidebar.tsx
- plugins/governance/test_utils.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
f227d90 to
d45bac7
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 (2)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
384-417:⚠️ Potential issue | 🟠 MajorProvider config budgets need
max_limitconverted from string to number.The function normalizes
weightandrate_limitfields but spreadsconfig.budgetsunchanged. Since the form schema storesmax_limitas a string but the API expects a number, provider config budgets will fail validation or be rejected by the backend.VK-level budgets are correctly converted in
onSubmit(lines 453–456), but the same conversion is missing here.🔧 Proposed fix
const normalizeProviderConfigs = (configs: typeof providerConfigs, existingConfigs?: VirtualKey["provider_configs"]): any[] => { return configs.map((config) => ({ ...config, + budgets: config.budgets + ?.map((b) => { + const maxLimit = b.max_limit ? parseFloat(b.max_limit) : undefined; + if (maxLimit === undefined || isNaN(maxLimit)) return null; + return { + max_limit: maxLimit, + reset_duration: b.reset_duration || "1M", + }; + }) + .filter((b): b is NonNullable<typeof b> => b !== null), weight: config.weight === undefined || config.weight === null ? null // ... rest of weight handling🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 384 - 417, The normalizeProviderConfigs function is leaving config.budgets unchanged so budget.max_limit remains a string; update normalizeProviderConfigs to normalize each config.budgets entry (e.g., map config.budgets) and convert budget.max_limit to a numeric value (use the existing normalizeNumericField helper or the same parseFloat/Number logic used for weight) returning null/undefined consistently when conversion fails, keeping other budget fields intact; reference the normalizeProviderConfigs function, the config.budgets array, and the normalizeNumericField helper (and mirror the VK-level conversion done in onSubmit) to ensure the API receives numbers for max_limit.framework/configstore/rdb.go (1)
2015-2018:⚠️ Potential issue | 🔴 CriticalPersist
calendar_alignedinUpdateVirtualKey.Line 2017 still uses the old VK update whitelist: it selects
budget_idbut omitscalendar_aligned. Existing virtual keys will never save changes to the new VK-level reset setting.🛠️ Suggested fix
- Select("name", "description", "value", "is_active", "team_id", "customer_id", "budget_id", "rate_limit_id", "config_hash", "updated_at", "encryption_status", "value_hash"). + Select("name", "description", "value", "is_active", "team_id", "customer_id", "calendar_aligned", "rate_limit_id", "config_hash", "updated_at", "encryption_status", "value_hash").🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2015 - 2018, The update call for virtual keys is missing the new field `calendar_aligned`, so changes to the VK-level reset setting are never persisted; in the UpdateVirtualKey path (around the txDB.WithContext(...).Select(...).Updates(virtualKey) call), add "calendar_aligned" to the Select whitelist alongside "budget_id" (and the other selected fields) so the Updates(virtualKey) will persist that field from the virtualKey struct (virtualKey.ID is already assigned from existing.ID).
♻️ Duplicate comments (2)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
1182-1190:⚠️ Potential issue | 🟡 MinorAdd
aria-labelto the MCP delete button for accessibility.This is an icon-only button without an accessible name. Screen readers will announce it as an unlabeled button.
♿ Proposed fix
<Button type="button" variant="ghost" size="sm" + aria-label={`Remove MCP client ${config.mcp_client_name}`} onClick={() => handleRemoveMCPClient(index)} data-testid={`vk-delete-mcp-${index}`} > <Trash2 className="h-4 w-4" /> </Button>🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 1182 - 1190, The delete button is an icon-only control (Button with Trash2) and lacks an accessible name; add an aria-label to the Button used for deleting MCP clients so screen readers can announce its purpose (e.g., aria-label describing the action and index like "Delete MCP client 1"). Update the Button JSX (the element using onClick={handleRemoveMCPClient} and data-testid={`vk-delete-mcp-${index}`}) to include the aria-label, keeping the existing event handler and test id intact.framework/configstore/rdb.go (1)
2315-2345:⚠️ Potential issue | 🟠 MajorMake standalone provider-config deletes atomic too.
Line 2333, Line 2337, and Line 2342 execute a multi-step delete without a local transaction when
txisn't passed. If the later rate-limit delete fails, the provider config is already gone. This helper also skips the join-table cleanup thatDeleteVirtualKeydoes first.Based on learnings, budgets and rate limits have a 1:1 ownership with their parent entities and are deleted together.🛠️ Suggested fix
func (s *RDBConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) error { - var txDB *gorm.DB - if len(tx) > 0 { - txDB = tx[0] - } else { - txDB = s.db - } + if len(tx) == 0 { + return s.db.WithContext(ctx).Transaction(func(innerTx *gorm.DB) error { + return s.DeleteVirtualKeyProviderConfig(ctx, id, innerTx) + }) + } + txDB := tx[0] + // First fetch the provider config to get budget and rate limit IDs var providerConfig tables.TableVirtualKeyProviderConfig if err := txDB.WithContext(ctx).First(&providerConfig, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } + + if err := txDB.WithContext(ctx).Exec( + "DELETE FROM governance_virtual_key_provider_config_keys WHERE table_virtual_key_provider_config_id = ?", + id, + ).Error; err != nil { + return err + } + // Store the rate limit ID before deleting rateLimitID := providerConfig.RateLimitID🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2315 - 2345, DeleteVirtualKeyProviderConfig performs multiple related deletes (budgets, provider config, rate limit) without a local transaction when no tx is provided, risking partial deletes and skipping join-table cleanup; modify the function to start a local transaction (use s.db.Begin()/txDB := tx[0] if provided) and run all operations (fetch providerConfig, delete join-table/budgets like DeleteVirtualKey does, delete provider config, delete rate limit) using that single txDB, and ensure you Commit on success and Rollback on any error so the multi-step delete is atomic and includes the same join-table cleanup as DeleteVirtualKey.
🧹 Nitpick comments (2)
plugins/governance/store_test.go (2)
201-207: Fail fast instead of silently skipping budget-overage setup.The nested
ifchain can no-op if budgets are missing/mismatched, which obscures root cause when the test fails later. Preferrequire.*guards before mutating usage.♻️ Suggested test-hardening diff
- if len(vk.Budgets) > 0 { - budgetID := vk.Budgets[0].ID - if budgetValue, exists := store.budgets.Load(budgetID); exists && budgetValue != nil { - if budget, ok := budgetValue.(*configstoreTables.TableBudget); ok && budget != nil { - budget.CurrentUsage = 100.0 - store.budgets.Store(budgetID, budget) - } - } - } + require.NotEmpty(t, vk.Budgets, "VK must have at least one budget") + budgetID := vk.Budgets[0].ID + budgetValue, exists := store.budgets.Load(budgetID) + require.True(t, exists, "VK budget must exist in in-memory budget store") + budget, ok := budgetValue.(*configstoreTables.TableBudget) + require.True(t, ok && budget != nil, "VK budget entry must be a non-nil *TableBudget") + budget.CurrentUsage = 100.0 + store.budgets.Store(budgetID, budget)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 201 - 207, The test silently skips budget setup because the nested ifs can no-op; replace them with hard assertions (require.*) so failures are explicit: assert vk.Budgets has length > 0 (e.g., require.Len or require.NotEmpty on vk.Budgets), then require that store.budgets.Load(budgetID) returns exists==true and a non-nil value, and require the value is a *configstoreTables.TableBudget before mutating budget.CurrentUsage and calling store.budgets.Store; update the block around vk.Budgets, budgetID, store.budgets.Load and the type assertion to use these require checks.
553-592: Extend create/delete lifecycle assertions to provider-config budgets.This test currently verifies VK budgets only. Given the stack’s multi-budget scope, include provider-config budgets in the fixture and assert they are also inserted/removed to cover the full in-memory cleanup path.
As per coding guidelines
**: always check the stack if there is one for the current PR and review changes in light of the whole stack.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store_test.go` around lines 553 - 592, Update TestGovernanceStore_MultiBudget_InMemoryCreateAndDelete to also include and assert provider-config budgets: modify the vk fixture (vk.ProviderConfigs built via buildProviderConfig) to include provider-level budgets in addition to b1/b2, call store.CreateVirtualKeyInMemory(vk) and assert those provider-config budget entries exist in store.budgets, verify they appear in retrieved.Budgets from store.GetVirtualKey("sk-bf-test") as appropriate, then after store.DeleteVirtualKeyInMemory("vk1") assert the provider-config budget entries are removed from store.budgets and that GetVirtualKey no longer finds the VK; use the existing helpers buildProviderConfig, CreateVirtualKeyInMemory, DeleteVirtualKeyInMemory and GetVirtualKey to locate the code to change.
🤖 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/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 84-91: The build fails because providerConfigSchema defines
budgets (array) but the MultiBudgetLines onChange handler assigns a legacy
budget property; edit the onChange in the MultiBudgetLines handler to stop
setting the singular "budget" key and instead update the "budgets" array
(matching providerConfigSchema's z.array of objects with max_limit and
reset_duration), ensuring the object shape matches the schema and removing any
legacy "budget" assignment in the state/props update.
---
Outside diff comments:
In `@framework/configstore/rdb.go`:
- Around line 2015-2018: The update call for virtual keys is missing the new
field `calendar_aligned`, so changes to the VK-level reset setting are never
persisted; in the UpdateVirtualKey path (around the
txDB.WithContext(...).Select(...).Updates(virtualKey) call), add
"calendar_aligned" to the Select whitelist alongside "budget_id" (and the other
selected fields) so the Updates(virtualKey) will persist that field from the
virtualKey struct (virtualKey.ID is already assigned from existing.ID).
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 384-417: The normalizeProviderConfigs function is leaving
config.budgets unchanged so budget.max_limit remains a string; update
normalizeProviderConfigs to normalize each config.budgets entry (e.g., map
config.budgets) and convert budget.max_limit to a numeric value (use the
existing normalizeNumericField helper or the same parseFloat/Number logic used
for weight) returning null/undefined consistently when conversion fails, keeping
other budget fields intact; reference the normalizeProviderConfigs function, the
config.budgets array, and the normalizeNumericField helper (and mirror the
VK-level conversion done in onSubmit) to ensure the API receives numbers for
max_limit.
---
Duplicate comments:
In `@framework/configstore/rdb.go`:
- Around line 2315-2345: DeleteVirtualKeyProviderConfig performs multiple
related deletes (budgets, provider config, rate limit) without a local
transaction when no tx is provided, risking partial deletes and skipping
join-table cleanup; modify the function to start a local transaction (use
s.db.Begin()/txDB := tx[0] if provided) and run all operations (fetch
providerConfig, delete join-table/budgets like DeleteVirtualKey does, delete
provider config, delete rate limit) using that single txDB, and ensure you
Commit on success and Rollback on any error so the multi-step delete is atomic
and includes the same join-table cleanup as DeleteVirtualKey.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 1182-1190: The delete button is an icon-only control (Button with
Trash2) and lacks an accessible name; add an aria-label to the Button used for
deleting MCP clients so screen readers can announce its purpose (e.g.,
aria-label describing the action and index like "Delete MCP client 1"). Update
the Button JSX (the element using onClick={handleRemoveMCPClient} and
data-testid={`vk-delete-mcp-${index}`}) to include the aria-label, keeping the
existing event handler and test id intact.
---
Nitpick comments:
In `@plugins/governance/store_test.go`:
- Around line 201-207: The test silently skips budget setup because the nested
ifs can no-op; replace them with hard assertions (require.*) so failures are
explicit: assert vk.Budgets has length > 0 (e.g., require.Len or
require.NotEmpty on vk.Budgets), then require that store.budgets.Load(budgetID)
returns exists==true and a non-nil value, and require the value is a
*configstoreTables.TableBudget before mutating budget.CurrentUsage and calling
store.budgets.Store; update the block around vk.Budgets, budgetID,
store.budgets.Load and the type assertion to use these require checks.
- Around line 553-592: Update
TestGovernanceStore_MultiBudget_InMemoryCreateAndDelete to also include and
assert provider-config budgets: modify the vk fixture (vk.ProviderConfigs built
via buildProviderConfig) to include provider-level budgets in addition to b1/b2,
call store.CreateVirtualKeyInMemory(vk) and assert those provider-config budget
entries exist in store.budgets, verify they appear in retrieved.Budgets from
store.GetVirtualKey("sk-bf-test") as appropriate, then after
store.DeleteVirtualKeyInMemory("vk1") assert the provider-config budget entries
are removed from store.budgets and that GetVirtualKey no longer finds the VK;
use the existing helpers buildProviderConfig, CreateVirtualKeyInMemory,
DeleteVirtualKeyInMemory and GetVirtualKey to locate the code to change.
🪄 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: a472f834-4e49-4b20-8c50-db07c571813b
⛔ Files ignored due to path filters (3)
docs/media/user-provisioning/okta-api-token-created.pngis excluded by!**/*.pngdocs/media/user-provisioning/okta-create-token-form.pngis excluded by!**/*.pngdocs/media/user-provisioning/okta-tokens-screen.pngis excluded by!**/*.png
📒 Files selected for processing (14)
docs/enterprise/setting-up-okta.mdxdocs/openapi/paths/management/users.yamlframework/configstore/rdb.goplugins/governance/store.goplugins/governance/store_test.goplugins/governance/test_utils.goui/app/_fallbacks/enterprise/components/user-groups/businessUnitsView.tsxui/app/_fallbacks/enterprise/components/user-groups/teamsView.tsxui/app/workspace/dashboard/components/charts/modelFilterSelect.tsxui/app/workspace/governance/business-units/page.tsxui/app/workspace/governance/teams/page.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/components/sidebar.tsxui/lib/types/governance.ts
✅ Files skipped from review due to trivial changes (4)
- ui/app/workspace/governance/business-units/page.tsx
- ui/app/_fallbacks/enterprise/components/user-groups/businessUnitsView.tsx
- docs/openapi/paths/management/users.yaml
- ui/app/workspace/dashboard/components/charts/modelFilterSelect.tsx
🚧 Files skipped from review as they are similar to previous changes (5)
- ui/app/_fallbacks/enterprise/components/user-groups/teamsView.tsx
- plugins/governance/test_utils.go
- docs/enterprise/setting-up-okta.mdx
- ui/components/sidebar.tsx
- plugins/governance/store.go
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
d45bac7 to
53c6cf2
Compare
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
53c6cf2 to
5262d83
Compare
Contributor
Author
Merge activity
|
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
Adds multi-budget support to virtual keys and provider configurations, allowing multiple budget limits with different reset intervals. Also introduces a new Access Profiles feature for enhanced role-based access control and includes the
expectskill for adversarial browser testing.Changes
governance_budgetstablecalendar_alignedfrom budget-level to VK-level setting that applies to all budgets under that virtual keybudget_idrelationships to newvirtual_key_id/provider_config_idforeign keysMultiBudgetLinesUI component for managing multiple budget configurations with duplicate detectionAccessProfilesfor managing role-based access profilesexpectskill for adversarial browser testing with structured test plans and session recordingsbudget_idfields from provider configsType of change
Affected areas
How to test
Validate database migrations and multi-budget functionality:
Test the new multi-budget feature by:
Screenshots/Recordings
UI changes include a new multi-budget configuration interface that allows adding/removing multiple budget lines with different reset intervals, replacing the single budget input field. Virtual key details now show all active budgets with their individual usage and reset schedules.
Breaking changes
The changes are backward compatible - existing single budget configurations are automatically migrated to the new multi-budget system via database migration that preserves all existing budget relationships.
Related issues
Enables more granular budget control for virtual keys and provider configurations with multiple time-based limits (e.g., hourly + daily + monthly budgets on the same VK).
Security considerations
New Access Profiles feature introduces additional RBAC controls. Database migrations include proper foreign key constraints and cascade deletion rules to maintain data integrity. The
expectskill provides secure adversarial testing capabilities for browser-facing changes.Checklist
docs/contributing/README.mdand followed the guidelines