add multiple budgets for virtual keys#2365
Closed
akshaydeo wants to merge 1 commit intographite-base/2365from
Closed
add multiple budgets for virtual keys#2365akshaydeo wants to merge 1 commit intographite-base/2365from
akshaydeo wants to merge 1 commit intographite-base/2365from
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 selected for processing (25)
📝 WalkthroughSummary by CodeRabbit
WalkthroughReplaces singular budget links with owned budgets: budgets now carry Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant HTTP as HTTP Handler
participant DB as RDB (GORM)
participant Store as Governance Store
participant InMem as gs.budgets
Client->>HTTP: Create/Update VirtualKey with Budgets[]
HTTP->>HTTP: validate budgets (limits, reset_duration, duplicates)
HTTP->>DB: DELETE FROM governance_budgets WHERE virtual_key_id/provider_config_id = ?
HTTP->>DB: INSERT TableBudget rows with virtual_key_id / provider_config_id
DB-->>HTTP: persisted budgets
HTTP->>Store: notify Create/UpdateVirtualKeyInMemory(vk with Budgets)
Store->>InMem: upsert each budget (preserve CurrentUsage/LastReset)
Client->>HTTP: Request using VirtualKey
HTTP->>Store: Evaluate request -> iterate budgets & check usage
Store->>InMem: read CurrentUsage/LastReset per budget ID
InMem-->>Store: budget statuses
Store-->>HTTP: decision (allow / budget exceeded)
HTTP-->>Client: response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate 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 — multiple P1 regressions in the UI layer (calendar_aligned never sent, broken Reset button) and a data-loss risk in the migration, on top of previously flagged unresolved issues (compile errors, budget usage reset). The Go backend schema changes and in-memory store refactor are clean and correct. However, the UI has at least three new P1 bugs introduced in this PR, plus the migration misses a backfill that would silently zero-out calendar-alignment settings for existing users. Combined with the still-open issues from prior review rounds, the PR is not yet merge-ready. ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx needs the most attention (three distinct P1 bugs). framework/configstore/migrations.go needs a calendar_aligned backfill. transports/bifrost-http/handlers/governance.go still has unresolved budget-usage-reset issues from prior review. 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: 5
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)
2558-2628:⚠️ Potential issue | 🟠 MajorDropped provider configs leave multi-budgets behind in
gs.budgets.This reconciliation only cleans
Budgetsfor provider configs that still exist inclone.ProviderConfigs. If a provider config is removed entirely, its multi-budget IDs stay in the global cache, soGetGovernanceData().Budgets,DumpBudgets, and the reset sweep keep processing orphaned budgets.Possible cleanup pass
if clone.ProviderConfigs != nil { // Create a map of existing provider configs by ID for fast lookup existingProviderConfigs := make(map[uint]configstoreTables.TableVirtualKeyProviderConfig) + seenProviderConfigIDs := make(map[uint]struct{}) if existingVK.ProviderConfigs != nil { for _, existingPC := range existingVK.ProviderConfigs { existingProviderConfigs[existingPC.ID] = existingPC } } @@ // Process each new/updated provider config for i, pc := range clone.ProviderConfigs { + if pc.ID != 0 { + seenProviderConfigIDs[pc.ID] = struct{}{} + } ... } + + for id, existingPC := range existingProviderConfigs { + if _, ok := seenProviderConfigIDs[id]; ok { + continue + } + for _, oldBudget := range existingPC.Budgets { + gs.budgets.Delete(oldBudget.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 2558 - 2628, The reconciliation misses removing multi-budgets when a ProviderConfig is deleted; after building existingProviderConfigs and processing clone.ProviderConfigs, add a cleanup pass that finds provider configs present in existingProviderConfigs but not in clone.ProviderConfigs and deletes their associated budget IDs and rate limit IDs from the in-memory caches; specifically, use the same identifiers used above (existingProviderConfigs, clone.ProviderConfigs, gs.budgets.Delete and gs.rateLimits.Delete) to iterate removed existingPC.Budgets (delete each oldBudget.ID), delete existingPC.Budget.ID if non-nil, and delete existingPC.RateLimit.ID if non-nil so orphaned budgets/rate-limits are removed from the global caches.
🤖 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 1840-1845: The delete paths still only remove the legacy BudgetID
row and leave many-to-many Budget links, so update the deletion code in
DeleteVirtualKey and DeleteVirtualKeyProviderConfig (and the provider-config
cleanup loop) to remove all associated Budgets and RateLimit entries created
with the parent: explicitly delete or dissociate entries from the join table for
the Budgets association (the Budgets many-to-many relation) in addition to
clearing BudgetID, and ensure RateLimit rows related to the
parent/provider-config are deleted as well; use the already-preloaded "Budgets"
and "RateLimit" relations to iterate and delete/dissociate each related record
so no Budget rows remain after the parent is removed.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 444-462: The new budgets[] pre-checks allow MaxLimit==0 to slip
through and later cause a 500; update each budgets validation block (the one
using req.Budgets, seenDurations, configstoreTables.ParseDuration and SendError)
to reject zero limits up front by either calling validateBudget and treating its
error as a badRequestError before starting the transaction, or explicitly change
the MaxLimit check to <= 0 (SendError 400) instead of < 0; apply the same change
to the other similar blocks mentioned so all budget branches consistently return
400 for zero or invalid budgets.
- Around line 516-545: When req.Budgets is present we must always derive or
clear the legacy vk.BudgetID from that list rather than only setting it when
nil; update the block that builds budgetsToAssociate (using validateBudget,
h.configStore.CreateBudget, and tx.Model(&vk).Association("Budgets").Replace) to
set vk.BudgetID = &budgetsToAssociate[0].ID when len(budgetsToAssociate) > 0 and
to set vk.BudgetID = nil (and persist via tx.Model(&vk).Update("budget_id",
vk.BudgetID)) when budgetsToAssociate is empty, so an empty budgets:[] clears
the legacy field; apply the same change to the other analogous handlers
referenced (the blocks around the vk.BudgetID logic at the other locations).
- Around line 1090-1123: When deleting provider configs (the later
!requestConfigsMap[id] cleanup), also remove their associated multi-budget rows
to avoid orphaned TableBudget records: for each providerConfig being removed,
load its Budgets association (providerConfig.Budgets or
tx.Model(providerConfig).Association("Budgets").Find), collect their IDs, delete
those TableBudget rows via the transaction (e.g. tx.Where("id IN (?)",
ids).Delete(&configstoreTables.TableBudget{}) or call a delete method on
h.configStore if available), and ensure BudgetID is handled/cleared; do this in
the same transactional cleanup that currently only queues legacy BudgetID so
both association and underlying rows are removed atomically.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 444-459: The provider-level budgets are not being
numeric-normalized before submit, causing backend type errors; update
normalizeProviderConfigs (or the submit path that builds
updateData.provider_configs) to iterate each provider_config and replace budgets
with budgets.map(b => ({ max_limit: normalizeNumericField(b.max_limit)!,
reset_duration: b.reset_duration || "1M" })) while filtering out entries where
normalizeNumericField(...) is undefined, and also set the provider_config's
singular budget field (if you maintain legacy compatibility) from the first
normalized budget; use the same helpers/shape as used for virtualKey-level
conversion (normalizeNumericField, budgets, budget, provider_configs,
updateData) so provider_configs[].budgets[*].max_limit is a number not a string.
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 2558-2628: The reconciliation misses removing multi-budgets when a
ProviderConfig is deleted; after building existingProviderConfigs and processing
clone.ProviderConfigs, add a cleanup pass that finds provider configs present in
existingProviderConfigs but not in clone.ProviderConfigs and deletes their
associated budget IDs and rate limit IDs from the in-memory caches;
specifically, use the same identifiers used above (existingProviderConfigs,
clone.ProviderConfigs, gs.budgets.Delete and gs.rateLimits.Delete) to iterate
removed existingPC.Budgets (delete each oldBudget.ID), delete
existingPC.Budget.ID if non-nil, and delete existingPC.RateLimit.ID if non-nil
so orphaned budgets/rate-limits are removed from the global caches.
🪄 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: 5f5f4160-1c61-4d83-8e26-981f0c2b2091
📒 Files selected for processing (5)
framework/configstore/rdb.goplugins/governance/store.gotransports/bifrost-http/handlers/governance.goui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/lib/types/governance.ts
Comment on lines
+444
to
+462
| // Validate multi-budgets if provided | ||
| if len(req.Budgets) > 0 { | ||
| seenDurations := make(map[string]bool) | ||
| for _, b := range req.Budgets { | ||
| if b.MaxLimit < 0 { | ||
| SendError(ctx, 400, fmt.Sprintf("Budget max_limit cannot be negative: %.2f", b.MaxLimit)) | ||
| return | ||
| } | ||
| if _, err := configstoreTables.ParseDuration(b.ResetDuration); err != nil { | ||
| SendError(ctx, 400, fmt.Sprintf("Invalid reset duration format: %s", b.ResetDuration)) | ||
| return | ||
| } | ||
| if seenDurations[b.ResetDuration] { | ||
| SendError(ctx, 400, fmt.Sprintf("Duplicate reset_duration in budgets: %s", b.ResetDuration)) | ||
| return | ||
| } | ||
| seenDurations[b.ResetDuration] = true | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
Zero max_limit still turns into a 500 here.
validateBudget rejects max_limit == 0, but these new budgets[] branches only precheck < 0 or defer validation until after the transaction has started. A zero-limit payload therefore bubbles out as a server error instead of a 400. Reuse validateBudget as a badRequestError up front, or explicitly reject <= 0 in each branch.
Also applies to: 631-654, 863-875, 1090-1113, 1215-1227
🤖 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 444 - 462, The
new budgets[] pre-checks allow MaxLimit==0 to slip through and later cause a
500; update each budgets validation block (the one using req.Budgets,
seenDurations, configstoreTables.ParseDuration and SendError) to reject zero
limits up front by either calling validateBudget and treating its error as a
badRequestError before starting the transaction, or explicitly change the
MaxLimit check to <= 0 (SendError 400) instead of < 0; apply the same change to
the other similar blocks mentioned so all budget branches consistently return
400 for zero or invalid budgets.
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
a40a014 to
71634a1
Compare
akshaydeo
force-pushed
the
03-29-add_multiple_budgets_for_virtual_keys
branch
2 times, most recently
from
a731aa4 to
f9c60f3
Compare
akshaydeo
force-pushed
the
03-29-access_profiles
branch
from
71634a1 to
291045f
Compare
akshaydeo
force-pushed
the
03-29-add_multiple_budgets_for_virtual_keys
branch
from
f9c60f3 to
829e1bb
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 (4)
plugins/governance/store.go (1)
2505-2555:⚠️ Potential issue | 🟠 MajorDelete children for provider configs that disappear on update.
This branch only cleans budgets and rate limits for provider configs that still exist in
clone.ProviderConfigs. If an update removes a provider config entirely—or clears the slice—its oldpc.Budgets/pc.RateLimitstay in the maps.DumpBudgets()later persists every entry ings.budgets, so those orphan budgets keep syncing after the config is gone.🧹 Cleanup sketch
if clone.ProviderConfigs != nil { // Create a map of existing provider configs by ID for fast lookup existingProviderConfigs := make(map[uint]configstoreTables.TableVirtualKeyProviderConfig) + remainingProviderConfigIDs := make(map[uint]struct{}) if existingVK.ProviderConfigs != nil { for _, existingPC := range existingVK.ProviderConfigs { existingProviderConfigs[existingPC.ID] = existingPC } } // Process each new/updated provider config for i, pc := range clone.ProviderConfigs { + remainingProviderConfigIDs[pc.ID] = struct{}{} ... } + + for id, existingPC := range existingProviderConfigs { + if _, stillPresent := remainingProviderConfigIDs[id]; stillPresent { + continue + } + if existingPC.RateLimit != nil { + gs.rateLimits.Delete(existingPC.RateLimit.ID) + } + for _, oldBudget := range existingPC.Budgets { + gs.budgets.Delete(oldBudget.ID) + } + } }Based on learnings, budgets in the governance plugin are 1:1 with their parent entities, so removing a provider config should also remove only its owned budget records.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 2505 - 2555, The code fails to remove budgets and rate limits for provider configs that were deleted in the update; after building existingProviderConfigs and before/after looping clone.ProviderConfigs, compute which existing provider IDs are no longer present in clone.ProviderConfigs and for each missing existingPC delete its associated budget IDs from gs.budgets and its RateLimit.ID from gs.rateLimits (use existingPC.Budgets and existingPC.RateLimit for ownership) so orphaned budgets/rate-limits are removed; reference existingProviderConfigs, clone.ProviderConfigs, existingPC.Budgets, existingPC.RateLimit, gs.budgets and gs.rateLimits to locate where to add this cleanup.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (3)
878-910:⚠️ Potential issue | 🟡 MinorAdd test hooks for the new budget editors.
Both
MultiBudgetLinesusages introduce interactive controls, but neither exposes adata-testidhook the way the surrounding inputs and buttons do. Please plumb test IDs through the component so the VK and provider budget editors follow the existingvk-*-*selector scheme. As per coding guidelines,ui/**/*.{tsx,ts}must adddata-testidattributes to new interactive elements using the<entity>-<element>-<qualifier>pattern.Also applies to: 1195-1202
🤖 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 878 - 910, The MultiBudgetLines component usages need data-testid props so tests can target the new interactive budget editors; update the JSX where MultiBudgetLines is rendered (id={`providerBudget-${index}`} and the other usage at lines around 1195-1202) to pass a data-testid following the vk-<entity>-<element>-<qualifier> pattern (e.g., data-testid={`vk-provider-budget-${index}`}), and then thread that prop into the MultiBudgetLines component so its internal interactive elements receive the same data-testid attributes (adjust the MultiBudgetLines prop signature and its internal inputs/buttons to accept and apply a data-testid or testId prop). Ensure handleUpdateProviderConfig usage remains unchanged.
349-365:⚠️ Potential issue | 🔴 CriticalFix syntax and reference errors in helper block (lines 350–392).
The helper block is broken:
clearVirtualKeyRateLimitsnever closes beforenormalizeProviderConfigsis declared,normalizeProviderConfigsis missing itsconfigsparameter despite using it in the function body, andnormalizeIntegerFieldis called but not defined anywhere in the file. The file will not type-check and reset handlers will fail at runtime.Fix
const clearVirtualKeyRateLimits = () => { + form.setValue("tokenMaxLimit", "", { shouldDirty: true }); + form.setValue("tokenResetDuration", "1h", { shouldDirty: true }); + form.setValue("requestMaxLimit", "", { shouldDirty: true }); + form.setValue("requestResetDuration", "1h", { shouldDirty: true }); +}; const normalizeProviderConfigs = ( + configs: NonNullable<FormData["providerConfigs"]>, existingConfigs?: VirtualKey["provider_configs"], ): any[] => {+const normalizeIntegerField = (value: string | undefined): number | undefined => { + if (value === undefined || value === "") return undefined; + const num = Number.parseInt(value, 10); + return Number.isNaN(num) ? undefined : num; +};🤖 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 349 - 365, The helper block is malformed: close the clearVirtualKeyRateLimits function declaration (clearVirtualKeyRateLimits) before declaring normalizeProviderConfigs, add the missing parameter configs to normalizeProviderConfigs (normalizeProviderConfigs(existingConfigs?: VirtualKey["provider_configs"], configs: any[])) so it uses the array it maps over, and provide/inline or import a normalizeIntegerField implementation (or replace calls with the correct helper) used inside rate_limit normalization; ensure all references (clearVirtualKeyRateLimits, normalizeProviderConfigs, normalizeIntegerField) are declared and exported/visible in the module so the file type-checks and the reset handlers run.
1204-1253:⚠️ Potential issue | 🔴 CriticalRemove this calendar alignment code section or complete the multi-budget implementation.
Lines 1205–1253 reference undefined variables (
watchedBudgetMaxLimit,watchedBudgetResetDuration,watchedBudgetCalendarAligned,handleCalendarAlignedChange) that don't exist in the file. This causes a build break. Additionally, thebudgets[]schema andonSubmithandler lack anycalendar_alignedfield, so even after fixing the undefined variables, the toggle cannot persist data in multi-budget mode. The calendar alignment feature appears to be incomplete refactoring from single-budget logic and should either be removed or fully integrated into the multi-budget model with proper schema and serialization.🤖 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 1204 - 1253, This calendar-alignment UI is leftover single-budget logic and references undefined symbols (watchedBudgetMaxLimit, watchedBudgetResetDuration, watchedBudgetCalendarAligned, handleCalendarAlignedChange, showCalendarAlignWarning, setShowCalendarAlignWarning) and should be removed to restore build stability: delete the entire JSX block that renders the Align to calendar cycle Switch (the conditional starting with checked={watchedBudgetCalendarAligned} and its container) plus the subsequent AlertDialog block that shows the calendar-aligned warning (including AlertDialogContent/Header/Footer and the onClick that calls form.setValue("budgetCalendarAligned", ...)); also remove or update any related test IDs ("vk-budget-calendar-aligned-toggle", "vk-calendar-align-cancel-btn", "vk-calendar-align-enable-btn") and any leftover state/handlers named handleCalendarAlignedChange or showCalendarAlignWarning from this component, or alternatively if you prefer full integration, add a calendar_aligned boolean per-budget to the budgets[] schema and wire watchedBudget* state and form serialization to budgets[i].calendar_aligned and persist it in onSubmit instead of the removal above.
♻️ Duplicate comments (4)
framework/configstore/rdb.go (1)
2087-2146:⚠️ Potential issue | 🟠 MajorMulti-budget cleanup is still incomplete on delete paths.
Line 2089 still loads only
ProviderConfigs, so the new loops overpc.BudgetsandvirtualKey.Budgetsrun against empty slices and never delete the ownedTableBudgetrows.DeleteVirtualKeyProviderConfignow also skips budget cleanup entirely, so standalone provider-config deletes leak budgets too.Based on learnings, budgets and rate limits are created along with the parent and deleted together.🧹 Minimal fix sketch
- if err := tx.WithContext(ctx).Preload("ProviderConfigs").First(&virtualKey, "id = ?", id).Error; err != nil { + if err := tx.WithContext(ctx). + Preload("Budgets"). + Preload("ProviderConfigs"). + Preload("ProviderConfigs.Budgets"). + First(&virtualKey, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } @@ - if err := txDB.WithContext(ctx).First(&providerConfig, "id = ?", id).Error; err != nil { + if err := txDB.WithContext(ctx).Preload("Budgets").First(&providerConfig, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return ErrNotFound } return err } @@ if rateLimitID != nil { if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil { return err } } + for _, budget := range providerConfig.Budgets { + if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", budget.ID).Error; err != nil { + return err + } + } return nil }Also applies to: 2339-2350
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2087 - 2146, The delete path fails to remove related TableBudget and TableRateLimit rows because the transaction only Preloads "ProviderConfigs" and "Budgets" on the virtual key are never loaded, so loops over pc.Budgets and virtualKey.Budgets iterate empty slices; update the load and delete logic: in the transaction that reads the TableVirtualKey use chained Preload calls to load ProviderConfigs with their Budgets and RateLimit (e.g. Preload("ProviderConfigs").Preload("ProviderConfigs.Budgets").Preload("ProviderConfigs.RateLimit")) and also Preload("Budgets").Preload("RateLimit") on the virtual key so vkBudgetIDs and providerConfigBudgetIDs are populated, then delete those collected budget and rate limit IDs after removing the parent rows; also apply the same collection-and-delete fix in DeleteVirtualKeyProviderConfig so that deleting a standalone TableVirtualKeyProviderConfig collects and deletes its owned TableBudget and TableRateLimit entries (use pc.Budgets and pc.RateLimit references) before/after deleting the provider-config row.ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
353-384:⚠️ Potential issue | 🟠 MajorNormalize provider-config budgets before submit.
MultiBudgetLinesemits stringmax_limits, butnormalizeProviderConfigs()only coercesweightandrate_limit. Becausebudgetsride through...configuntouched,provider_configs[].budgets[*].max_limitis still sent as a string on both create and update. Please apply the same numeric normalization you already use for the top-leveldata.budgets, and preserve the[]case when clearing an existing provider budget set.Also applies to: 878-909
🤖 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 353 - 384, normalizeProviderConfigs currently leaves provider_configs[].budgets[*].max_limit as strings; update the map callback that transforms each config (the block that computes weight and rate_limit) to also normalize config.budgets: iterate config.budgets (if present) and replace each budget.max_limit with normalizeIntegerField(budget.max_limit) coercing to number or null just like top-level budgets normalization, and when budgets is intentionally cleared (empty array) ensure you return [] (not undefined) so existing provider budgets are removed; use normalizeIntegerField and refer to config.budgets and existingConfigs/existingConfig to decide whether to return [] vs undefined.transports/bifrost-http/handlers/governance.go (2)
208-210:⚠️ Potential issue | 🟠 MajorProvider-config budget cleanup depends on
Budgetsbeing loaded, but caller doesn’t preload it.
collectProviderConfigDeleteIDsnow readsconfig.Budgets, but existing provider configs are fetched withFind(&existingConfigs)only. Without preload, budget IDs can be missed during deletion cleanup.Suggested 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 }Also applies to: 862-865
🤖 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 208 - 210, collectProviderConfigDeleteIDs now reads config.Budgets but the caller fetches existing provider configs without preloading Budgets, so budget IDs can be missed; update the query that loads existing provider configs (the Find(&existingConfigs) call) to preload the Budgets association (e.g., use Preload("Budgets") or the ORM equivalent) so each ProviderConfig's Budgets slice is populated before calling collectProviderConfigDeleteIDs; ensure any other places that call collectProviderConfigDeleteIDs also load Budgets (see the other occurrence around lines 862-865) so cleanup sees all budget IDs.
440-457:⚠️ Potential issue | 🟡 Minor
max_limit == 0still escapes early validation and can surface as 500.Line 444 (and similar branches) only reject
< 0. A zero limit reachesvalidateBudget, and in several paths that error is returned as a generic transaction error, resulting in a 500 instead of a 400.Suggested fix pattern
- if b.MaxLimit < 0 { + if b.MaxLimit <= 0 { return &badRequestError{err: fmt.Errorf("budget max_limit cannot be negative or zero: %.2f", b.MaxLimit)} } budget := configstoreTables.TableBudget{ ... } - if err := validateBudget(&budget); err != nil { - return err - } + if err := validateBudget(&budget); err != nil { + return &badRequestError{err: err} + }Apply this consistently to VK-level and provider-config budget branches.
Also applies to: 577-593, 745-757, 930-947, 996-1003
🤖 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 440 - 457, The budget validation currently only rejects negative MaxLimit (b.MaxLimit < 0) so a zero MaxLimit slips through and later surfaces as a 500; update the validation in the request-budget handling (the loop over req.Budgets in governance.go) to treat MaxLimit <= 0 as invalid and return a 400 with a clear message, and apply the same check to the other budget validation branches (VK-level and provider-config budget paths) and any calls to validateBudget so zero limits are rejected early rather than propagating as transaction/internal errors.
🧹 Nitpick comments (4)
ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx (1)
306-322: Consider using budget ID as React key if available.The budget rendering logic correctly iterates over the budgets array. Using
idxas the key works here since the list is stable during render, but if theBudgettype has anidfield, usingb.idwould be more robust and follow React best practices for list keys.If Budget has an id field:
-{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 - 322, Replace the array index key used when rendering budgets with a stable identifier: in the vk.budgets.map callback inside virtualKeysTable.tsx, change the key from the loop index (idx) to the budget's unique id (b.id) if Budget includes an id field; ensure b.id is used for the key in the element rendered by the map to follow React best practices and avoid using idx as the key.plugins/governance/test_utils.go (1)
208-211: KeepbuildVirtualKeyWithMultiBudgetsconsistent with the other governance-limit builders.The neighboring helpers now add a default provider config specifically so resolver tests don’t stop at the provider gate, but this one still returns a virtual key with no providers. Matching the same preconditions here will make multi-budget tests less surprising.
🧪 Minimal tweak
func buildVirtualKeyWithMultiBudgets(id, value, name string, budgets []configstoreTables.TableBudget) *configstoreTables.TableVirtualKey { vk := buildVirtualKey(id, value, name, true) 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 208 - 211, buildVirtualKeyWithMultiBudgets returns a VK with budgets but no provider config unlike other builders; update it to mirror the other governance-limit helpers by adding the same default provider configuration (e.g., set vk.Providers or vk.ProviderConfig to the default provider used by buildVirtualKey or neighboring helpers) so resolver tests won't stop at the provider gate—locate buildVirtualKeyWithMultiBudgets and add the same default provider setup that the other helper functions apply to their returned *configstoreTables.TableVirtualKey.ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx (1)
151-187: Use the shared percentage helper for the new budget badges.Both new budget sections duplicate the same
current_usage / max_limitmath instead ofcalculateUsagePercentage, so budget badges can drift from the rate-limit UI on edge cases and future tweaks. A small shared budget renderer/helper would also remove the copy/paste between these two 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 - 187, The budget badge is computing the percent inline; replace the inline Math.round((b.current_usage / b.max_limit) * 100) expression with the shared helper calculateUsagePercentage (or calculateUsagePercentage(current, max) if that’s its signature) and ensure the Badge receives that value (preserving rounding/formatting done by the helper). Also extract the repeated budget display (the usage span + Badge) into a small reusable renderer/helper used in both budget sections to remove duplication; update references in the virtualKeyDetailsSheet render (where config.budgets is mapped) to call that helper so both places use calculateUsagePercentage consistently.transports/bifrost-http/handlers/governance.go (1)
496-518: Extract shared multi-budget create/replace flow to reduce drift.The same create/validate/associate logic is repeated in four places, and behavior already diverged (validation and calendar alignment). A shared helper would reduce future regressions.
Also applies to: 576-603, 929-956, 992-1045
🤖 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 496 - 518, Extract the repeated create/validate/associate multi-budget flow into a single helper (e.g., createAndAssociateBudgets) that accepts context, transaction (tx), the parent model reference (vk or other model), and the incoming req.Budgets slice; inside the helper perform the loop that builds TableBudget entries, calls validateBudget(&budget), calls h.configStore.CreateBudget(ctx, &budget, tx), accumulates budgetsToAssociate, and finally does tx.Model(parent).Association("Budgets").Replace(budgetsToAssociate) while returning detailed errors; then replace the four inline blocks (including the block using vk, and the blocks at the other locations) with calls to this helper to ensure consistent validation and calendar alignment 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 `@transports/bifrost-http/handlers/governance.go`:
- Around line 500-506: The multi-budget constructors are ignoring the request's
CalendarAligned and always setting LastReset to time.Now(), breaking
calendar-aligned semantics; update the TableBudget construction (e.g., where
TableBudget is created in the code blocks around the shown diff) to copy
CalendarAligned from the incoming budget struct (use CalendarAligned:
b.CalendarAligned) and set LastReset from the incoming value (LastReset:
b.LastReset) or compute the correct calendar-aligned reset time if the request
provides a calendar alignment flag, instead of always using time.Now(); apply
the same fix to the other TableBudget creations referenced (around the other
noted blocks).
---
Outside diff comments:
In `@plugins/governance/store.go`:
- Around line 2505-2555: The code fails to remove budgets and rate limits for
provider configs that were deleted in the update; after building
existingProviderConfigs and before/after looping clone.ProviderConfigs, compute
which existing provider IDs are no longer present in clone.ProviderConfigs and
for each missing existingPC delete its associated budget IDs from gs.budgets and
its RateLimit.ID from gs.rateLimits (use existingPC.Budgets and
existingPC.RateLimit for ownership) so orphaned budgets/rate-limits are removed;
reference existingProviderConfigs, clone.ProviderConfigs, existingPC.Budgets,
existingPC.RateLimit, gs.budgets and gs.rateLimits to locate where to add this
cleanup.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 878-910: The MultiBudgetLines component usages need data-testid
props so tests can target the new interactive budget editors; update the JSX
where MultiBudgetLines is rendered (id={`providerBudget-${index}`} and the other
usage at lines around 1195-1202) to pass a data-testid following the
vk-<entity>-<element>-<qualifier> pattern (e.g.,
data-testid={`vk-provider-budget-${index}`}), and then thread that prop into the
MultiBudgetLines component so its internal interactive elements receive the same
data-testid attributes (adjust the MultiBudgetLines prop signature and its
internal inputs/buttons to accept and apply a data-testid or testId prop).
Ensure handleUpdateProviderConfig usage remains unchanged.
- Around line 349-365: The helper block is malformed: close the
clearVirtualKeyRateLimits function declaration (clearVirtualKeyRateLimits)
before declaring normalizeProviderConfigs, add the missing parameter configs to
normalizeProviderConfigs (normalizeProviderConfigs(existingConfigs?:
VirtualKey["provider_configs"], configs: any[])) so it uses the array it maps
over, and provide/inline or import a normalizeIntegerField implementation (or
replace calls with the correct helper) used inside rate_limit normalization;
ensure all references (clearVirtualKeyRateLimits, normalizeProviderConfigs,
normalizeIntegerField) are declared and exported/visible in the module so the
file type-checks and the reset handlers run.
- Around line 1204-1253: This calendar-alignment UI is leftover single-budget
logic and references undefined symbols (watchedBudgetMaxLimit,
watchedBudgetResetDuration, watchedBudgetCalendarAligned,
handleCalendarAlignedChange, showCalendarAlignWarning,
setShowCalendarAlignWarning) and should be removed to restore build stability:
delete the entire JSX block that renders the Align to calendar cycle Switch (the
conditional starting with checked={watchedBudgetCalendarAligned} and its
container) plus the subsequent AlertDialog block that shows the calendar-aligned
warning (including AlertDialogContent/Header/Footer and the onClick that calls
form.setValue("budgetCalendarAligned", ...)); also remove or update any related
test IDs ("vk-budget-calendar-aligned-toggle", "vk-calendar-align-cancel-btn",
"vk-calendar-align-enable-btn") and any leftover state/handlers named
handleCalendarAlignedChange or showCalendarAlignWarning from this component, or
alternatively if you prefer full integration, add a calendar_aligned boolean
per-budget to the budgets[] schema and wire watchedBudget* state and form
serialization to budgets[i].calendar_aligned and persist it in onSubmit instead
of the removal above.
---
Duplicate comments:
In `@framework/configstore/rdb.go`:
- Around line 2087-2146: The delete path fails to remove related TableBudget and
TableRateLimit rows because the transaction only Preloads "ProviderConfigs" and
"Budgets" on the virtual key are never loaded, so loops over pc.Budgets and
virtualKey.Budgets iterate empty slices; update the load and delete logic: in
the transaction that reads the TableVirtualKey use chained Preload calls to load
ProviderConfigs with their Budgets and RateLimit (e.g.
Preload("ProviderConfigs").Preload("ProviderConfigs.Budgets").Preload("ProviderConfigs.RateLimit"))
and also Preload("Budgets").Preload("RateLimit") on the virtual key so
vkBudgetIDs and providerConfigBudgetIDs are populated, then delete those
collected budget and rate limit IDs after removing the parent rows; also apply
the same collection-and-delete fix in DeleteVirtualKeyProviderConfig so that
deleting a standalone TableVirtualKeyProviderConfig collects and deletes its
owned TableBudget and TableRateLimit entries (use pc.Budgets and pc.RateLimit
references) before/after deleting the provider-config row.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 208-210: collectProviderConfigDeleteIDs now reads config.Budgets
but the caller fetches existing provider configs without preloading Budgets, so
budget IDs can be missed; update the query that loads existing provider configs
(the Find(&existingConfigs) call) to preload the Budgets association (e.g., use
Preload("Budgets") or the ORM equivalent) so each ProviderConfig's Budgets slice
is populated before calling collectProviderConfigDeleteIDs; ensure any other
places that call collectProviderConfigDeleteIDs also load Budgets (see the other
occurrence around lines 862-865) so cleanup sees all budget IDs.
- Around line 440-457: The budget validation currently only rejects negative
MaxLimit (b.MaxLimit < 0) so a zero MaxLimit slips through and later surfaces as
a 500; update the validation in the request-budget handling (the loop over
req.Budgets in governance.go) to treat MaxLimit <= 0 as invalid and return a 400
with a clear message, and apply the same check to the other budget validation
branches (VK-level and provider-config budget paths) and any calls to
validateBudget so zero limits are rejected early rather than propagating as
transaction/internal errors.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 353-384: normalizeProviderConfigs currently leaves
provider_configs[].budgets[*].max_limit as strings; update the map callback that
transforms each config (the block that computes weight and rate_limit) to also
normalize config.budgets: iterate config.budgets (if present) and replace each
budget.max_limit with normalizeIntegerField(budget.max_limit) coercing to number
or null just like top-level budgets normalization, and when budgets is
intentionally cleared (empty array) ensure you return [] (not undefined) so
existing provider budgets are removed; use normalizeIntegerField and refer to
config.budgets and existingConfigs/existingConfig to decide whether to return []
vs undefined.
---
Nitpick comments:
In `@plugins/governance/test_utils.go`:
- Around line 208-211: buildVirtualKeyWithMultiBudgets returns a VK with budgets
but no provider config unlike other builders; update it to mirror the other
governance-limit helpers by adding the same default provider configuration
(e.g., set vk.Providers or vk.ProviderConfig to the default provider used by
buildVirtualKey or neighboring helpers) so resolver tests won't stop at the
provider gate—locate buildVirtualKeyWithMultiBudgets and add the same default
provider setup that the other helper functions apply to their returned
*configstoreTables.TableVirtualKey.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 496-518: Extract the repeated create/validate/associate
multi-budget flow into a single helper (e.g., createAndAssociateBudgets) that
accepts context, transaction (tx), the parent model reference (vk or other
model), and the incoming req.Budgets slice; inside the helper perform the loop
that builds TableBudget entries, calls validateBudget(&budget), calls
h.configStore.CreateBudget(ctx, &budget, tx), accumulates budgetsToAssociate,
and finally does
tx.Model(parent).Association("Budgets").Replace(budgetsToAssociate) while
returning detailed errors; then replace the four inline blocks (including the
block using vk, and the blocks at the other locations) with calls to this helper
to ensure consistent validation and calendar alignment behavior.
In `@ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx`:
- Around line 151-187: The budget badge is computing the percent inline; replace
the inline Math.round((b.current_usage / b.max_limit) * 100) expression with the
shared helper calculateUsagePercentage (or calculateUsagePercentage(current,
max) if that’s its signature) and ensure the Badge receives that value
(preserving rounding/formatting done by the helper). Also extract the repeated
budget display (the usage span + Badge) into a small reusable renderer/helper
used in both budget sections to remove duplication; update references in the
virtualKeyDetailsSheet render (where config.budgets is mapped) to call that
helper so both places use calculateUsagePercentage consistently.
In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx`:
- Around line 306-322: Replace the array index key used when rendering budgets
with a stable identifier: in the vk.budgets.map callback inside
virtualKeysTable.tsx, change the key from the loop index (idx) to the budget's
unique id (b.id) if Budget includes an id field; ensure b.id is used for the key
in the element rendered by the map to follow React best practices and avoid
using idx as the key.
🪄 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: 9021a479-f19f-47c5-b342-18b70e74d412
📒 Files selected for processing (20)
.agents/skills/expect/SKILL.md.claude/skills/expectframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.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/lib/config.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 (3)
- transports/bifrost-http/lib/config.go
- transports/config.schema.json
- framework/configstore/clientconfig.go
✅ Files skipped from review due to trivial changes (4)
- .claude/skills/expect
- plugins/governance/http_transport_prehook_test.go
- .agents/skills/expect/SKILL.md
- framework/configstore/migrations.go
🚧 Files skipped from review as they are similar to previous changes (1)
- ui/lib/types/governance.ts
akshaydeo
force-pushed
the
03-29-add_multiple_budgets_for_virtual_keys
branch
from
829e1bb to
d2cfff2
Compare
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 (3)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (1)
84-87:⚠️ Potential issue | 🟠 MajorPreserve
calendar_alignedwhen round-tripping budgets.The new form model only keeps
max_limitandreset_duration, so existing calendar-aligned budgets are hydrated without the flag and then saved back without it. Any edit to a virtual-key budget will silently disable calendar-based resets, and provider budgets cannot retain the flag at all with the currentMultiBudgetLinesshape. Please carrycalendar_alignedthrough the schema, default-value mapping, editor state, and create/update payloads.Also applies to: 121-124, 199-202, 222-224, 417-424, 461-467, 878-910, 1195-1201
🤖 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 - 87, The budgets schema and form flow drop the calendar_aligned flag causing calendar-aligned budgets to be lost; update the zod schema (the budgets: z.array(...) definition), the default-value mapping that hydrates existing budgets into the form model, the MultiBudgetLines editor state/shape, and the create/update payload builders used by the VirtualKeySheet component so that each budget object includes calendar_aligned (boolean, preserved when present, defaulting to false if absent). Specifically: add calendar_aligned to the z.object in the budgets array, ensure the hydration code that maps existing budget records into form defaults copies calendar_aligned, update MultiBudgetLines types/state to carry calendar_aligned through editing and UI (including any add/remove/modify flows), and include calendar_aligned in the payloads sent by your createVirtualKey/updateVirtualKey (or similar) functions so saved budgets retain the flag.transports/bifrost-http/lib/config_test.go (1)
15290-15300:⚠️ Potential issue | 🟠 MajorDon't exclude the new
budgetssurface from schema sync.Line 15294 and Line 15300 mark
budgetsas internal-only, but this stack addsBudgetsto virtual keys and provider configs. With these exclusions,TestConfigSchemaSyncwill pass even ifconfig.schema.jsonnever documents or validates the new field.Suggested 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 },If
budgetsis actually DB-only, the fix should be on the struct tags (json:"-") rather than by teaching the schema test to ignore a JSON-tagged field.🤖 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 15290 - 15300, The test currently marks "budgets" as internal-only in the exclusion map for tables.TableVirtualKey and tables.TableVirtualKeyProviderConfig, which hides the new Budgets surface from TestConfigSchemaSync; remove the "budgets": true entries from the exclusion map associated with tables.TableVirtualKey and tables.TableVirtualKeyProviderConfig so the schema sync test will require the JSON schema to document/validate Budgets, or if Budgets truly must be DB-only instead, add a json:"-" struct tag to the Budgets field on the corresponding struct rather than excluding it in the test.framework/configstore/migrations.go (1)
1004-1019:⚠️ Potential issue | 🔴 CriticalBlanket error suppression for
budget_idDDL creates a silent failure path on SQLite and downstream migrations.The
ALTER TABLEstatement at line 1004 usesIF NOT EXISTSsyntax, which SQLite does not support—the error will be silently ignored, leaving the column missing. The later migration at lines 5366–5375 checksHasColumnbefore backfilling, so it skips silently when the column doesn't exist. This is fragile and inconsistent with the safe approach used forrate_limit_id(lines 1009–1013).Replace the raw
tx.Exec()with explicit column existence check usingmigrator.HasColumn(), then add the column viamigrator.AddColumn()if needed—as already done forrate_limit_id. This ensures the migration either succeeds with the column present or fails visibly, preventing downstream data loss.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/migrations.go` around lines 1004 - 1019, Replace the raw tx.Exec("ALTER TABLE ... budget_id ...") call with the same existence-check pattern used for rate_limit_id: use migrator.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") and, if it returns false, call migrator.AddColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id"); ensure any AddColumn error is returned (e.g., wrap with fmt.Errorf) so the migration fails loudly rather than silently ignoring databases like SQLite that don't support IF NOT EXISTS.
♻️ Duplicate comments (1)
transports/bifrost-http/handlers/governance.go (1)
440-458:⚠️ Potential issue | 🟡 MinorZero
max_limitvalidation still inconsistent withvalidateBudget.The pre-check at line 444 uses
< 0butvalidateBudget(line 2145) rejectsMaxLimit <= 0. A zero value passes this pre-check but fails inside the transaction, returning a 500 instead of 400.Proposed fix: Change to <= 0
for _, b := range req.Budgets { - if b.MaxLimit < 0 { + if b.MaxLimit <= 0 { SendError(ctx, 400, fmt.Sprintf("Budget max_limit cannot be negative: %.2f", b.MaxLimit))🤖 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 440 - 458, The pre-check in the budgets validation loop currently allows MaxLimit == 0 but validateBudget enforces MaxLimit > 0, causing a later 500; update the MaxLimit check in the loop over req.Budgets to reject zero as well (use <= 0 instead of < 0) and keep using SendError with a 400 status and a clear message (e.g. "Budget max_limit must be > 0") so the pre-transaction validation matches validateBudget; leave the ResetDuration parsing and duplicate ResetDuration check (seenDurations and configstoreTables.ParseDuration) unchanged.
🧹 Nitpick comments (3)
framework/configstore/rdb_test.go (1)
524-544: These round-trips still don't exercise the new multi-budget surface.Both flows attach exactly one VK budget, and neither verifies
ProviderConfigs.Budgets. A regression that only preloads the first budget or skips provider-config budgets would still pass. Please add one VK case with two budgets and one provider-config case that asserts itsBudgetspreload.Also applies to: 1007-1041
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb_test.go` around lines 524 - 544, Add test cases that exercise multi-budget preload: for the virtual key path, create two Budget records and associate both to the same virtual key ID (use store.CreateVirtualKey, set both budgets' VirtualKeyID to vkID and call store.UpdateBudget or store.CreateBudget as appropriate), then call store.GetVirtualKey(ctx, vkID) and assert result.Budgets has length 2 and contains both budget IDs; for the provider-config path, create a ProviderConfig with two budgets attached (set ProviderConfig.Budgets or Budget.ProviderConfigID accordingly), persist them, call the provider-config retrieval helper (e.g., store.GetProviderConfig or the function used in the other tests) and assert ProviderConfig.Budgets is preloaded and contains both budgets; update the existing single-budget assertions near the shown vk-with-refs block and the other section (around the 1007-1041 area) to include these multi-budget assertions.transports/bifrost-http/handlers/governance_test.go (1)
274-290: Add a real multi-budget assertion here.Both updated cases still use a one-item
Budgetsslice, so a bug that only collects the first entry would still pass. Please add a case with two budgets and assert both IDs are returned/appended.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/handlers/governance_test.go` around lines 274 - 290, The test cases for TableVirtualKeyProviderConfig are insufficiently checking multi-budget behavior; add a new case (e.g., name: "handles multiple budgets") where config.Budgets contains two configstoreTables.TableBudget entries with distinct IDs (use a new budgetID2 alongside budgetID), set wantBudgetIDs to include both IDs (and for the append variant include initialBudgetIDs plus both IDs), and ensure the test asserts that the function under test returns/appends both budget IDs (not just the first). Target the table-driven cases adjacent to the existing entries that set initialBudgetIDs, wantBudgetIDs, and config.Budgets so the new case exercises both creation and append paths.transports/bifrost-http/handlers/governance.go (1)
508-510:validateBudgeterrors should be wrapped asbadRequestError.All
validateBudgetcalls return plain errors, which propagate through the transaction and result in 500 responses (via the fallback at line 634). Since validation failures are client errors, wrap them to ensure 400 responses:Proposed fix: Wrap validateBudget errors
if err := validateBudget(&budget); err != nil { - return err + return &badRequestError{err: err} }Apply this pattern to all
validateBudgetcalls within the transaction closures.Also applies to: 590-592, 771-773, 926-928, 1001-1003
🤖 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 508 - 510, Calls to validateBudget inside the transaction closures are returning plain errors (e.g., "return err") which become 500s; change those returns to wrap the validation error as a badRequestError so the HTTP response becomes a 400. Locate each validateBudget invocation (the one shown and the others at the ranges referenced) and replace the direct return of err with a wrapped badRequestError (use the existing badRequestError wrapper type/function in this file, e.g., badRequestError.Wrap(err) or the project’s standard badRequestError construction) so validation failures are propagated as client errors.
🤖 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 2392-2395: The backfill validation becomes the only effective
check after ignore_columns includes virtual_key_id and provider_config_id, but
the script currently only warns on mismatched legacy budget_id and on
provider-config presence; change the budget_id and provider-config mismatch
handlers to set failed=1 instead of just warning, and tighten the
provider-config check in the provider-config validation block to assert the
exact expected row value (compare to the known expected provider_config_id or
matching FK) rather than using IS NOT NULL; apply these same changes to the
second validation block that mirrors this logic (the other block around the
provider-config/backfill checks).
- Around line 2603-2675: The counts in verify_budget_migration_postgres() are
compared as scalars but use run_postgres_sql(), which returns formatted psql
output; create and use a scalar helper run_postgres_scalar() that calls psql
with -t -A (and preserves exit status) to return unadorned values, then replace
each run_postgres_sql(...) invocation in verify_budget_migration_postgres() (the
vk_budget_count, pc_budget_count, has_vk_col, has_pc_col, vk_has_budget_id,
pc_has_budget_id, junction_vk assignments) with run_postgres_scalar(...) and
remove the trailing redirection/whitespace-trimming so the if [ "$var" = "1" ]
checks work reliably.
In `@framework/configstore/migrations.go`:
- Around line 5360-5376: The error messages in the migration rollback use
uppercase acronyms ("VK" and "PC") which violates the lowercase error-string
convention; update the two fmt.Errorf messages inside the tx.Exec error returns
to use lowercase descriptions (e.g., "failed to backfill virtual key budget
virtual_key_id" and "failed to backfill provider config budget
provider_config_id") so they remain lowercase and descriptive; these occur in
the migration block that references tables.TableVirtualKeyProviderConfig and the
tx.Exec that updates governance_budgets/provider_config_id.
- Around line 5380-5390: The rollback currently removes only the new foreign
keys from governance_budgets (virtual_key_id/provider_config_id) but does not
restore the legacy budget_id columns on the parent tables, leaving a downgraded
binary unable to read budget relationships; update the Rollback in migrations.go
so it first re-creates the legacy budget_id columns on tables.TableVirtualKey
and tables.TableVirtualKeyProviderConfig (with the same type/nullability as
before) and backfill them from governance_budgets where possible before dropping
virtual_key_id/provider_config_id, or if you intend this to be irreversible,
explicitly make the migration one-way by returning a non-nil error in Rollback
with a clear message instead of performing partial changes.
In `@framework/configstore/tables/budget.go`:
- Around line 22-24: Add an ownership guard so a Budget cannot reference both
VirtualKeyID and ProviderConfigID: implement a BeforeSave method on the Budget
model that validates exactly one of VirtualKeyID or ProviderConfigID is non-nil
(error if both set or both nil), and update the DB migration for budgets to add
a durable CHECK constraint (e.g. ensuring (virtual_key_id IS NULL) !=
(provider_config_id IS NULL)) so the database enforces exclusive ownership as
well.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 909-933: The update path that iterates pc.Budgets is missing the
same pre-validation performed in the create path: validate incoming b.MaxLimit
and b.ResetDuration (and any other raw input rules used in the create branch)
before constructing the configstoreTables.TableBudget and calling
validateBudget/CreateBudget; add the same input checks you used in the create
flow (e.g. ensure b.MaxLimit is present/within allowed range and b.ResetDuration
is non-empty/parsable/calendar rules) for each item in pc.Budgets (within the
loop over pc.Budgets) and return a badRequestError on invalid inputs so
validateBudget and h.configStore.CreateBudget only receive already-validated
data.
- Around line 573-597: Add explicit MaxLimit validation for provider config
budgets before calling validateBudget and persisting: when iterating pc.Budgets
in the handler (the loop over "for _, b := range pc.Budgets"), check that
b.MaxLimit is > 0 and return a badRequestError (similar to the VK-level check)
if not; ensure this validation occurs prior to building the
configstoreTables.TableBudget, calling validateBudget, and invoking
h.configStore.CreateBudget so invalid max_limit produces a 400 instead of a 500.
---
Outside diff comments:
In `@framework/configstore/migrations.go`:
- Around line 1004-1019: Replace the raw tx.Exec("ALTER TABLE ... budget_id
...") call with the same existence-check pattern used for rate_limit_id: use
migrator.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") and, if
it returns false, call
migrator.AddColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id"); ensure
any AddColumn error is returned (e.g., wrap with fmt.Errorf) so the migration
fails loudly rather than silently ignoring databases like SQLite that don't
support IF NOT EXISTS.
In `@transports/bifrost-http/lib/config_test.go`:
- Around line 15290-15300: The test currently marks "budgets" as internal-only
in the exclusion map for tables.TableVirtualKey and
tables.TableVirtualKeyProviderConfig, which hides the new Budgets surface from
TestConfigSchemaSync; remove the "budgets": true entries from the exclusion map
associated with tables.TableVirtualKey and tables.TableVirtualKeyProviderConfig
so the schema sync test will require the JSON schema to document/validate
Budgets, or if Budgets truly must be DB-only instead, add a json:"-" struct tag
to the Budgets field on the corresponding struct rather than excluding it in the
test.
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 84-87: The budgets schema and form flow drop the calendar_aligned
flag causing calendar-aligned budgets to be lost; update the zod schema (the
budgets: z.array(...) definition), the default-value mapping that hydrates
existing budgets into the form model, the MultiBudgetLines editor state/shape,
and the create/update payload builders used by the VirtualKeySheet component so
that each budget object includes calendar_aligned (boolean, preserved when
present, defaulting to false if absent). Specifically: add calendar_aligned to
the z.object in the budgets array, ensure the hydration code that maps existing
budget records into form defaults copies calendar_aligned, update
MultiBudgetLines types/state to carry calendar_aligned through editing and UI
(including any add/remove/modify flows), and include calendar_aligned in the
payloads sent by your createVirtualKey/updateVirtualKey (or similar) functions
so saved budgets retain the flag.
---
Duplicate comments:
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 440-458: The pre-check in the budgets validation loop currently
allows MaxLimit == 0 but validateBudget enforces MaxLimit > 0, causing a later
500; update the MaxLimit check in the loop over req.Budgets to reject zero as
well (use <= 0 instead of < 0) and keep using SendError with a 400 status and a
clear message (e.g. "Budget max_limit must be > 0") so the pre-transaction
validation matches validateBudget; leave the ResetDuration parsing and duplicate
ResetDuration check (seenDurations and configstoreTables.ParseDuration)
unchanged.
---
Nitpick comments:
In `@framework/configstore/rdb_test.go`:
- Around line 524-544: Add test cases that exercise multi-budget preload: for
the virtual key path, create two Budget records and associate both to the same
virtual key ID (use store.CreateVirtualKey, set both budgets' VirtualKeyID to
vkID and call store.UpdateBudget or store.CreateBudget as appropriate), then
call store.GetVirtualKey(ctx, vkID) and assert result.Budgets has length 2 and
contains both budget IDs; for the provider-config path, create a ProviderConfig
with two budgets attached (set ProviderConfig.Budgets or Budget.ProviderConfigID
accordingly), persist them, call the provider-config retrieval helper (e.g.,
store.GetProviderConfig or the function used in the other tests) and assert
ProviderConfig.Budgets is preloaded and contains both budgets; update the
existing single-budget assertions near the shown vk-with-refs block and the
other section (around the 1007-1041 area) to include these multi-budget
assertions.
In `@transports/bifrost-http/handlers/governance_test.go`:
- Around line 274-290: The test cases for TableVirtualKeyProviderConfig are
insufficiently checking multi-budget behavior; add a new case (e.g., name:
"handles multiple budgets") where config.Budgets contains two
configstoreTables.TableBudget entries with distinct IDs (use a new budgetID2
alongside budgetID), set wantBudgetIDs to include both IDs (and for the append
variant include initialBudgetIDs plus both IDs), and ensure the test asserts
that the function under test returns/appends both budget IDs (not just the
first). Target the table-driven cases adjacent to the existing entries that set
initialBudgetIDs, wantBudgetIDs, and config.Budgets so the new case exercises
both creation and append paths.
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 508-510: Calls to validateBudget inside the transaction closures
are returning plain errors (e.g., "return err") which become 500s; change those
returns to wrap the validation error as a badRequestError so the HTTP response
becomes a 400. Locate each validateBudget invocation (the one shown and the
others at the ranges referenced) and replace the direct return of err with a
wrapped badRequestError (use the existing badRequestError wrapper type/function
in this file, e.g., badRequestError.Wrap(err) or the project’s standard
badRequestError construction) so validation failures are propagated as client
errors.
🪄 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: 6ca34073-df89-412e-b631-eca26814a9cd
📒 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
- plugins/governance/http_transport_prehook_test.go
- .agents/skills/expect/SKILL.md
🚧 Files skipped from review as they are similar to previous changes (7)
- plugins/governance/resolver.go
- plugins/governance/resolver_test.go
- plugins/governance/model_provider_governance_test.go
- plugins/governance/test_utils.go
- ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
- transports/config.schema.json
- framework/configstore/rdb.go
Comment on lines
+2392
to
+2395
| # - virtual_key_id, provider_config_id: new FK columns on governance_budgets (added by multi-budget migration) | ||
| # - status, description: key validation runs after migration, updating these fields | ||
| # for invalid/test keys (e.g., status becomes "list_models_failed") | ||
| local ignore_columns="updated_at config_hash created_at models_json weight allowed_models network_config_json concurrency_buffer_json proxy_config_json custom_provider_config_json budget_id rate_limit_id status description" | ||
| local ignore_columns="updated_at config_hash created_at models_json weight allowed_models network_config_json concurrency_buffer_json proxy_config_json custom_provider_config_json budget_id rate_limit_id virtual_key_id provider_config_id status description" |
Contributor
There was a problem hiding this comment.
Backfill misses should fail once these columns are ignored in the snapshot diff.
After Line 2395 starts ignoring virtual_key_id and provider_config_id, this helper becomes the only backfill validation. The faker seed in this script always populates the legacy budget_id links, but these branches only warn on mismatch, so a broken migration can still pass. Set failed=1 on both mismatches and tighten the provider-config check to the expected row instead of IS NOT NULL.
🛠️ Suggested fix
if [ "$vk_budget_count" = "1" ]; then
log_info " VK budget migration: budget-migration-test-1 → vk-migration-test-1 ✓"
else
- log_warn " VK budget migration: budget-migration-test-1 virtual_key_id not set (count=$vk_budget_count) — may be expected if old version didn't have budget_id on VK"
+ log_error " VK budget migration: budget-migration-test-1 virtual_key_id not set (count=$vk_budget_count)"
+ failed=1
fi
@@
- pc_budget_count=$(run_postgres_sql "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-2' AND provider_config_id IS NOT NULL" 2>/dev/null | tr -d '[:space:]')
+ pc_budget_count=$(run_postgres_scalar "SELECT COUNT(*) FROM governance_budgets b JOIN governance_virtual_key_provider_configs pc ON pc.id = b.provider_config_id WHERE b.id = 'budget-migration-test-2' AND pc.virtual_key_id = 'vk-migration-test-2' AND pc.provider = 'anthropic'" | tr -d '[:space:]')
if [ "$pc_budget_count" = "1" ]; then
log_info " PC budget migration: budget-migration-test-2 → provider_config ✓"
else
- log_warn " PC budget migration: budget-migration-test-2 provider_config_id not set (count=$pc_budget_count) — may be expected if old version didn't have budget_id on PC"
+ log_error " PC budget migration: budget-migration-test-2 provider_config_id not set (count=$pc_budget_count)"
+ failed=1
fiAlso applies to: 2607-2625
🤖 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 2392 - 2395,
The backfill validation becomes the only effective check after ignore_columns
includes virtual_key_id and provider_config_id, but the script currently only
warns on mismatched legacy budget_id and on provider-config presence; change the
budget_id and provider-config mismatch handlers to set failed=1 instead of just
warning, and tighten the provider-config check in the provider-config validation
block to assert the exact expected row value (compare to the known expected
provider_config_id or matching FK) rather than using IS NOT NULL; apply these
same changes to the second validation block that mirrors this logic (the other
block around the provider-config/backfill checks).
Comment on lines
+2603
to
+2675
| verify_budget_migration_postgres() { | ||
| log_info "Verifying budget migration (budget_id → virtual_key_id/provider_config_id)..." | ||
| local failed=0 | ||
|
|
||
| # Check: budget-migration-test-1 was linked to vk-migration-test-1 via budget_id | ||
| # After migration, governance_budgets.virtual_key_id should be set | ||
| local vk_budget_count | ||
| vk_budget_count=$(run_postgres_sql "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-1' AND virtual_key_id = 'vk-migration-test-1'" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$vk_budget_count" = "1" ]; then | ||
| log_info " VK budget migration: budget-migration-test-1 → vk-migration-test-1 ✓" | ||
| else | ||
| log_warn " VK budget migration: budget-migration-test-1 virtual_key_id not set (count=$vk_budget_count) — may be expected if old version didn't have budget_id on VK" | ||
| fi | ||
|
|
||
| # Check: budget-migration-test-2 was linked to provider config via budget_id | ||
| # After migration, governance_budgets.provider_config_id should be set | ||
| local pc_budget_count | ||
| pc_budget_count=$(run_postgres_sql "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-2' AND provider_config_id IS NOT NULL" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$pc_budget_count" = "1" ]; then | ||
| log_info " PC budget migration: budget-migration-test-2 → provider_config ✓" | ||
| else | ||
| log_warn " PC budget migration: budget-migration-test-2 provider_config_id not set (count=$pc_budget_count) — may be expected if old version didn't have budget_id on PC" | ||
| fi | ||
|
|
||
| # Check: virtual_key_id and provider_config_id columns exist on governance_budgets | ||
| local has_vk_col | ||
| has_vk_col=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_budgets' AND column_name = 'virtual_key_id'" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$has_vk_col" = "1" ]; then | ||
| log_info " Column governance_budgets.virtual_key_id exists ✓" | ||
| else | ||
| log_error " Column governance_budgets.virtual_key_id MISSING!" | ||
| failed=1 | ||
| fi | ||
|
|
||
| local has_pc_col | ||
| has_pc_col=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_budgets' AND column_name = 'provider_config_id'" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$has_pc_col" = "1" ]; then | ||
| log_info " Column governance_budgets.provider_config_id exists ✓" | ||
| else | ||
| log_error " Column governance_budgets.provider_config_id MISSING!" | ||
| failed=1 | ||
| fi | ||
|
|
||
| # Check: budget_id column should be dropped from governance_virtual_keys | ||
| local vk_has_budget_id | ||
| vk_has_budget_id=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_virtual_keys' AND column_name = 'budget_id'" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$vk_has_budget_id" = "0" ]; then | ||
| log_info " Column governance_virtual_keys.budget_id dropped ✓" | ||
| else | ||
| log_error " Column governance_virtual_keys.budget_id still exists!" | ||
| failed=1 | ||
| fi | ||
|
|
||
| # Check: budget_id column should be dropped from governance_virtual_key_provider_configs | ||
| local pc_has_budget_id | ||
| pc_has_budget_id=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'governance_virtual_key_provider_configs' AND column_name = 'budget_id'" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$pc_has_budget_id" = "0" ]; then | ||
| log_info " Column governance_virtual_key_provider_configs.budget_id dropped ✓" | ||
| else | ||
| log_error " Column governance_virtual_key_provider_configs.budget_id still exists!" | ||
| failed=1 | ||
| fi | ||
|
|
||
| # Check: junction tables should not exist | ||
| local junction_vk | ||
| junction_vk=$(run_postgres_sql "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'governance_virtual_key_budgets'" 2>/dev/null | tr -d '[:space:]') | ||
| if [ "$junction_vk" = "0" ]; then | ||
| log_info " Junction table governance_virtual_key_budgets dropped ✓" | ||
| else | ||
| log_warn " Junction table governance_virtual_key_budgets still exists (may not have existed in old version)" | ||
| fi | ||
|
|
||
| return $failed |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Does psql -c "SELECT COUNT(*) FROM ..."include headers/row-count metadata by default, and do-t -A flags return only the scalar value?
💡 Result:
Yes, psql -c "SELECT COUNT(*) FROM ..." includes headers and row-count metadata by default. No, the -t -A flags return only the scalar value.
Citations:
- 1: https://www.postgresql.org/docs/current/app-psql.html
- 2: https://stackoverflow.com/questions/9934264/how-to-hide-result-set-decoration-in-psql-output
- 3: https://stackoverflow.com/questions/45352374/psql-non-select-how-to-remove-formatting-and-show-only-certain-columns
- 4: https://stackoverflow.com/questions/67019292/postgresql-select-statement-without-header-footer-lines
🏁 Script executed:
# Find and examine the run_postgres_sql function
rg -A 10 "run_postgres_sql\(\)" .github/workflows/scripts/run-migration-tests.sh | head -30Repository: maximhq/bifrost
Length of output: 257
🏁 Script executed:
# Also check the entire function to see the exact psql command
rg -B 2 -A 15 "def run_postgres_sql|run_postgres_sql\s*\(\)" .github/workflows/scripts/run-migration-tests.shRepository: maximhq/bifrost
Length of output: 367
🏁 Script executed:
# Look for the function definition more carefully
grep -n "run_postgres_sql" .github/workflows/scripts/run-migration-tests.sh | head -20Repository: maximhq/bifrost
Length of output: 1546
run_postgres_sql() must use -t -A flags for scalar comparisons to work.
run_postgres_sql() invokes psql -c without -t -A flags, which returns formatted output including headers, column separators, and row-count metadata. Stripping whitespace with tr -d '[:space:]' removes only spaces and newlines—not the text content like count, ---, and (1 row). This causes all comparisons (if [ "$vk_budget_count" = "1" ], etc.) to fail silently, leaving the migration checks ineffective.
Add a dedicated scalar helper function:
Suggested fix
+run_postgres_scalar() {
+ local sql="$1"
+ local container
+ container=$(get_postgres_container)
+ if [ -z "$container" ]; then
+ log_error "PostgreSQL container not found"
+ return 1
+ fi
+ docker exec "$container" \
+ psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -A \
+ -c "$sql" 2>/dev/null
+}
+
verify_budget_migration_postgres() {
log_info "Verifying budget migration..."
- vk_budget_count=$(run_postgres_sql "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-1' AND virtual_key_id = 'vk-migration-test-1'" 2>/dev/null | tr -d '[:space:]')
+ vk_budget_count=$(run_postgres_scalar "SELECT COUNT(*) FROM governance_budgets WHERE id = 'budget-migration-test-1' AND virtual_key_id = 'vk-migration-test-1'")Then replace all scalar calls in verify_budget_migration_postgres() (lines 2610, 2620, 2629, 2638, 2648, 2658, 2668) from run_postgres_sql() to run_postgres_scalar() and remove the trailing 2>/dev/null | tr -d '[:space:]'.
🤖 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 2603 - 2675,
The counts in verify_budget_migration_postgres() are compared as scalars but use
run_postgres_sql(), which returns formatted psql output; create and use a scalar
helper run_postgres_scalar() that calls psql with -t -A (and preserves exit
status) to return unadorned values, then replace each run_postgres_sql(...)
invocation in verify_budget_migration_postgres() (the vk_budget_count,
pc_budget_count, has_vk_col, has_pc_col, vk_has_budget_id, pc_has_budget_id,
junction_vk assignments) with run_postgres_scalar(...) and remove the trailing
redirection/whitespace-trimming so the if [ "$var" = "1" ] checks work reliably.
Comment on lines
+5360
to
+5376
| `).Error; err != nil { | ||
| return fmt.Errorf("failed to backfill VK budget virtual_key_id: %w", err) | ||
| } | ||
| } | ||
|
|
||
| // Backfill: migrate existing VK single budgets to junction table | ||
| if err := tx.Exec(` | ||
| INSERT INTO governance_virtual_key_budgets (virtual_key_id, budget_id) | ||
| SELECT id, budget_id FROM governance_virtual_keys | ||
| WHERE budget_id IS NOT NULL AND budget_id != '' | ||
| AND NOT EXISTS ( | ||
| SELECT 1 FROM governance_virtual_key_budgets | ||
| WHERE governance_virtual_key_budgets.virtual_key_id = governance_virtual_keys.id | ||
| AND governance_virtual_key_budgets.budget_id = governance_virtual_keys.budget_id | ||
| ) | ||
| `).Error; err != nil { | ||
| return fmt.Errorf("failed to backfill VK budgets: %w", err) | ||
| // Backfill: set provider_config_id from legacy PC budget_id (if column still exists) | ||
| if mg.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") { | ||
| if err := tx.Exec(` | ||
| UPDATE governance_budgets SET provider_config_id = ( | ||
| SELECT id FROM governance_virtual_key_provider_configs | ||
| WHERE governance_virtual_key_provider_configs.budget_id = governance_budgets.id | ||
| ) WHERE provider_config_id IS NULL AND EXISTS ( | ||
| SELECT 1 FROM governance_virtual_key_provider_configs | ||
| WHERE governance_virtual_key_provider_configs.budget_id = governance_budgets.id | ||
| ) | ||
| `).Error; err != nil { | ||
| return fmt.Errorf("failed to backfill PC budget provider_config_id: %w", err) |
Contributor
There was a problem hiding this comment.
Keep the new wrapped errors lowercase.
VK and PC break the repo's error-string convention.
As per coding guidelines, "Error strings should be lowercase with no trailing punctuation (Go convention)."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@framework/configstore/migrations.go` around lines 5360 - 5376, The error
messages in the migration rollback use uppercase acronyms ("VK" and "PC") which
violates the lowercase error-string convention; update the two fmt.Errorf
messages inside the tx.Exec error returns to use lowercase descriptions (e.g.,
"failed to backfill virtual key budget virtual_key_id" and "failed to backfill
provider config budget provider_config_id") so they remain lowercase and
descriptive; these occur in the migration block that references
tables.TableVirtualKeyProviderConfig and the tx.Exec that updates
governance_budgets/provider_config_id.
Comment on lines
5380
to
5390
| // Drop legacy budget_id columns from VK and ProviderConfig | ||
| if mg.HasColumn(&tables.TableVirtualKey{}, "budget_id") { | ||
| if err := mg.DropColumn(&tables.TableVirtualKey{}, "budget_id"); err != nil { | ||
| return fmt.Errorf("failed to drop budget_id from governance_virtual_keys: %w", err) | ||
| } | ||
| } | ||
| if mg.HasColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id") { | ||
| if err := mg.DropColumn(&tables.TableVirtualKeyProviderConfig{}, "budget_id"); err != nil { | ||
| return fmt.Errorf("failed to drop budget_id from governance_virtual_key_provider_configs: %w", err) | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
Rollback no longer restores the pre-migration schema.
Migrate drops budget_id from both parent tables, but Rollback only removes virtual_key_id / provider_config_id from governance_budgets. After rollback, neither schema exists, so a downgraded binary has no budget relationship to read. Please re-add/backfill the legacy budget_id columns before dropping the new ones, or make this migration explicitly one-way instead of returning success.
Also applies to: 5394-5406
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@framework/configstore/migrations.go` around lines 5380 - 5390, The rollback
currently removes only the new foreign keys from governance_budgets
(virtual_key_id/provider_config_id) but does not restore the legacy budget_id
columns on the parent tables, leaving a downgraded binary unable to read budget
relationships; update the Rollback in migrations.go so it first re-creates the
legacy budget_id columns on tables.TableVirtualKey and
tables.TableVirtualKeyProviderConfig (with the same type/nullability as before)
and backfill them from governance_budgets where possible before dropping
virtual_key_id/provider_config_id, or if you intend this to be irreversible,
explicitly make the migration one-way by returning a non-nil error in Rollback
with a clear message instead of performing partial changes.
Comment on lines
+22
to
+24
| // Owner FKs: a budget belongs to at most one VK or one ProviderConfig | ||
| VirtualKeyID *string `gorm:"type:varchar(255);index" json:"virtual_key_id,omitempty"` | ||
| ProviderConfigID *uint `gorm:"index" json:"provider_config_id,omitempty"` |
Contributor
There was a problem hiding this comment.
Enforce exclusive ownership for the new budget FKs.
Nothing prevents the same budget row from carrying both VirtualKeyID and ProviderConfigID. That breaks the 1:1 ownership model and makes delete/cascade behavior ambiguous. A BeforeSave guard is a good start; a DB check in the migration would make it durable.
🛠️ Suggested guard
func (b *TableBudget) BeforeSave(tx *gorm.DB) error {
+ if b.VirtualKeyID != nil && b.ProviderConfigID != nil {
+ return fmt.Errorf("budget cannot belong to both virtual key and provider config")
+ }
+
// Validate that ResetDuration is in correct format (e.g., "30s", "5m", "1h", "1d", "1w", "1M", "1Y")
if d, err := ParseDuration(b.ResetDuration); err != nil {🤖 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 22 - 24, Add an
ownership guard so a Budget cannot reference both VirtualKeyID and
ProviderConfigID: implement a BeforeSave method on the Budget model that
validates exactly one of VirtualKeyID or ProviderConfigID is non-nil (error if
both set or both nil), and update the DB migration for budgets to add a durable
CHECK constraint (e.g. ensuring (virtual_key_id IS NULL) != (provider_config_id
IS NULL)) so the database enforces exclusive ownership as well.
akshaydeo
force-pushed
the
03-29-add_multiple_budgets_for_virtual_keys
branch
from
d2cfff2 to
4570b55
Compare
akshaydeo
changed the base branch from
03-29-access_profiles
to
graphite-base/2365
| budgets: virtualKey?.budgets && virtualKey.budgets.length > 0 | ||
| ? virtualKey.budgets.map((b) => ({ max_limit: String(b.max_limit ?? ""), reset_duration: b.reset_duration ?? "1M" })) | ||
| : [], | ||
| budgetCalendarAligned: virtualKey?.budgets?.some((b) => b.calendar_aligned) ?? false, |
Contributor
There was a problem hiding this comment.
budgetCalendarAligned always initializes to false
The form initializes budgetCalendarAligned by reading calendar_aligned from individual budget objects. But this PR moves calendar_aligned from a per-budget field to a VK-level field (TableVirtualKey.CalendarAligned). The Budget TS type still has calendar_aligned?: boolean, but the backend no longer populates it — so b.calendar_aligned is always undefined, and some(...) always returns false.
When editing an existing VK that had calendar alignment enabled, the toggle will always appear unchecked. The correct source is the VK's own calendar_aligned field:
Suggested change
| budgetCalendarAligned: virtualKey?.budgets?.some((b) => b.calendar_aligned) ?? false, | |
| budgetCalendarAligned: virtualKey?.calendar_aligned ?? false, |
Comment on lines
+5322
to
+5330
| // Add calendar_aligned to governance_virtual_keys (VK-level setting) | ||
| if !mg.HasColumn(&tables.TableVirtualKey{}, "calendar_aligned") { | ||
| if err := mg.AddColumn(&tables.TableVirtualKey{}, "CalendarAligned"); err != nil { | ||
| return fmt.Errorf("failed to add calendar_aligned column to governance_virtual_keys: %w", err) | ||
| } | ||
| } | ||
|
|
||
| // Create provider-config-level multi-budget junction table | ||
| if !mg.HasTable(&tables.TableVirtualKeyProviderConfigBudget{}) { | ||
| if err := mg.CreateTable(&tables.TableVirtualKeyProviderConfigBudget{}); err != nil { | ||
| return fmt.Errorf("failed to create governance_virtual_key_provider_config_budgets table: %w", err) | ||
| // Add FK columns on governance_budgets for multi-budget ownership | ||
| if !mg.HasColumn(&tables.TableBudget{}, "virtual_key_id") { |
Contributor
There was a problem hiding this comment.
calendar_aligned not backfilled from budgets to VKs during migration
The migration adds a calendar_aligned column to governance_virtual_keys (defaulting to false) and backfills virtual_key_id/provider_config_id on governance_budgets, but it never copies the existing per-budget calendar_aligned values up to the owning VK. Any deployment that already ran the previous add_budget_calendar_aligned_column migration and had budgets with calendar_aligned = true will silently lose that setting — all VKs will end up with calendar_aligned = false after the upgrade.
A backfill should be added before the legacy budget_id columns are dropped, for example:
-- Backfill VK calendar_aligned from any owned budget that had calendar_aligned = true
UPDATE governance_virtual_keys
SET calendar_aligned = true
WHERE id IN (
SELECT virtual_key_id
FROM governance_budgets
WHERE virtual_key_id IS NOT NULL AND calendar_aligned = true
);
This was referenced Mar 31, 2026
Closed
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 budgets with different reset durations (e.g., daily, weekly, monthly) to be configured simultaneously for more granular spending control.
Changes
Budgetsfield to virtual keys and provider configurations alongside existing singleBudgetfield for backward compatibilityGetVirtualKeys,GetVirtualKeysPaginated,GetVirtualKey, andGetVirtualKeyByValuebudgetsarray in create/update requests with validation for unique reset durationsMultiBudgetLinescomponentBudgetIDfieldType of change
Affected areas
How to test
Test multi-budget creation and management:
Verify that budget enforcement works across different time periods and that existing single-budget configurations continue to function.
Breaking changes
This change maintains full backward compatibility with existing single-budget configurations while adding new multi-budget capabilities.
Security considerations
Multi-budget validation ensures no duplicate reset durations and proper budget limit enforcement. Budget usage tracking maintains existing security boundaries and access controls.
Checklist
docs/contributing/README.mdand followed the guidelines