fix: fixes create virtual key double usage counting

[roroghost17]

Summary

Fixes a double-counting bug on newly created virtual keys where the in-memory store retained a direct pointer to the caller's TableVirtualKey. After CreateVirtualKeyInMemory returned, hydrateVKGovernance would mutate that same pointer — injecting model-config-owned rate-limit and budget identities onto the VK's fields — causing the usage tracker to count each request twice (once via the VK-scoped model path and again via the VK-hierarchy path).

Changes

• CreateVirtualKeyInMemory now works on a shallow clone of the incoming *TableVirtualKey before storing it, ensuring post-create mutations by the caller do not affect the tracked VK's hierarchy fields.
• Added a regression test (TestGovernanceStore_CreateVirtualKeyInMemory_DecouplesFromCallerPointer) that simulates the hydrateVKGovernance mutation pattern and asserts the stored VK remains unaffected.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

go test ./plugins/governance/... -run TestGovernanceStore_CreateVirtualKeyInMemory_DecouplesFromCallerPointer -v
go test ./plugins/governance/...

The new regression test will fail on the previous code and pass with the clone fix in place. Verify no other virtual key store tests regress.

Breaking changes

• Yes
• No

Related issues

Regression introduced by the hydrateVKGovernance post-store mutation pattern on newly created virtual keys.

Security considerations

No auth, secrets, or PII implications. The fix prevents incorrect usage accounting that could allow rate limits and budgets to be bypassed or incorrectly enforced on new virtual keys.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

Summary by CodeRabbit

• Bug Fixes

  • Stored isolated copies of virtual keys, teams, and customers so external mutations no longer affect tracked governance records; calendar-aligned flags are correctly applied to associated budgets and rate limits.

• Tests

  • Added a regression test ensuring in-memory governance records remain decoupled from caller-provided data after creation.

[CLAassistant]

All committers have signed the CLA.

[coderabbitai]

📝 Walkthrough

Walkthrough

CreateVirtualKeyInMemory, CreateTeamInMemory, and CreateCustomerInMemory now clone inputs, stamp IsCalendarAligned on cloned budgets/rate-limits (including provider-config entries), persist cloned entries into flat caches, and store pointers to the clones; a regression test verifies virtual key cloning decouples from the caller's pointer. Two test literals were reformatted.

Changes

In-memory cloning for governance entities

Layer / File(s)	Summary
VirtualKey clone and store plugins/governance/store.go	CreateVirtualKeyInMemory clones the incoming TableVirtualKey, stamps IsCalendarAligned on the cloned budgets and rate-limits (including provider-config), writes related entries keyed by IDs into in-memory maps, and stores &clone in gs.virtualKeys.
Team clone and store plugins/governance/store.go	CreateTeamInMemory clones the incoming team, stamps IsCalendarAligned on cloned team budgets and attached rate-limit, writes those entries keyed by IDs into in-memory maps, and stores &clone in gs.teams.
Customer clone and store plugins/governance/store.go	CreateCustomerInMemory clones the incoming customer, stamps IsCalendarAligned on cloned customer budgets and attached rate-limit, writes those entries keyed by IDs into in-memory maps, and stores &clone in gs.customers.
Regression test + test formatting plugins/governance/store_test.go	Adds TestGovernanceStore_CreateVirtualKeyInMemory_DecouplesFromCallerPointer to assert the stored VK does not reflect post-call mutations to the caller's VK; two TableRateLimit literals reformatted for alignment only.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• maximhq/bifrost#3452: Related governance changes around calendar-aligned state handling for VK-owned/provider-config budgets and rate limits.

Suggested reviewers

• Pratham-Mishra04

Poem

I'm a rabbit who copies with care, 🐇
I clone your keys and budgets fair.
No pointer leaks, no counts askew,
The store keeps snapshots safe and true.
Hop, test, and ship — a tidy view.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'fix: fixes create virtual key double usage counting' accurately describes the main bug fix—preventing double-counting of virtual key usage by storing a clone instead of a caller pointer.
Description check	✅ Passed	The PR description comprehensively covers the template sections: Summary explains the bug, Changes detail the fix and regression test, Type/Areas/Testing are checked, no breaking changes, and security implications are documented.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-05-fix_fixes_create_virtual_key_double_usage_counting

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[roroghost17]

• fix: fixes create virtual key double usage counting #4070 👈 (View in Graphite)
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 5/5

The change is safe to merge: it decouples stored governance records from caller-owned pointers without altering any public interfaces or update/delete paths.

All three Create*InMemory functions now store clones rather than the caller's pointer. The critical ProviderConfigs slice receives a proper element-level copy (the exact vector for the double-counting bug). The top-level Budgets slice is not deep-copied, but the actual post-store mutation in applyVKGovernanceFromModelConfigs is a field reassignment (vk.Budgets = mc.Budgets), not an in-place element mutation, so the shared backing array is never written through. The regression test directly exercises the mutation pattern described in the PR and all existing governance store tests remain intact.

No files require special attention.

Important Files Changed

Filename	Overview
plugins/governance/store.go	CreateVirtualKeyInMemory, CreateTeamInMemory, and CreateCustomerInMemory now work on shallow clones of the incoming pointers; ProviderConfigs receives a proper element-level copy to prevent the applyVKGovernanceFromModelConfigs in-place mutation from escaping into the stored VK's hierarchy fields.
plugins/governance/store_test.go	Adds a regression test that simulates the hydrateVKGovernance mutation pattern (both top-level reassignments and in-place ProviderConfigs element mutations) and asserts the stored VK clone remains decoupled.

_{Reviews (3): Last reviewed commit: "fix: fixes create virtual key double usa..." | Re-trigger Greptile}

[roroghost17]

@coderabbitai full-review

[coderabbitai]

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Get the diff for PR `#4070`
gh pr diff 4070 --repo maximhq/bifrost

echo "---FILE_CONTENT---"
cat plugins/governance/store.go

echo "---TEST_FILE---"
cat plugins/governance/store_test.go

Length of output: 213295

---

---

@roroghost17: Here is the full review of PR #4070.

---

Review of fix: fixes create virtual key double usage counting

Summary

The root cause is correctly identified and the fix is sound for the described bug pattern. CreateVirtualKeyInMemory (and by extension CreateTeamInMemory / CreateCustomerInMemory) used to store the caller's pointer directly into gs.virtualKeys. After the call returned, hydrateVKGovernance reassigned vk.RateLimit, vk.RateLimitID, and vk.Budgets on the same pointer, causing the stored VK to inherit model-config-owned IDs and produce double-counting on every subsequent request.

The shallow-clone fix (clone := *vk) decouples the stored struct from the caller's pointer so post-create slice/pointer reassignments are invisible to the store.

---

Correctness of the shallow clone

The three mutation types described for hydrateVKGovernance are all pointer/slice-header reassignments:

vk.RateLimit = hydratedRL       // replaces pointer in caller's struct
vk.RateLimitID = &hydratedRL.ID // replaces pointer in caller's struct
vk.Budgets = []...{...}         // replaces slice header in caller's struct

A shallow copy (clone := *vk) gives clone its own copy of these fields. Reassigning them on vk after the call does not touch clone. The fix is correct for all three cases.

One nuance worth noting: the shallow copy does not protect against in-place mutations through a non-nil shared pointer. Specifically, if vk.RateLimit != nil at call time, clone.RateLimit aliases the same TableRateLimit struct. If the caller subsequently mutates fields of that struct (e.g. vk.RateLimit.ID = "new-id") rather than reassigning the pointer, the stored clone's rate limit is also affected. The test covers only the nil initial RateLimit path (matching the actual hydrateVKGovernance scenario), so this is not a current bug, but it is a latent vulnerability if the caller contract changes.

Recommendation (non-blocking): Consider adding a comment on the shallow clone line explaining its scope:

// Shallow clone breaks the caller's pointer identity so post-create
// pointer/slice reassignments (e.g. hydrateVKGovernance) do not mutate the
// stored entry. Note: pointer fields (RateLimit) still alias the caller's
// struct if non-nil at create time; only reassignments are decoupled.
clone := *vk

---

Test coverage

The regression test TestGovernanceStore_CreateVirtualKeyInMemory_DecouplesFromCallerPointer is well-structured and correctly exercises the exact mutation pattern from hydrateVKGovernance.

Minor gap (non-blocking): CreateTeamInMemory and CreateCustomerInMemory receive the same clone-before-store treatment but have no equivalent regression tests. If a similar post-create mutation pattern exists for those callers, the coverage gap could hide regressions. A brief parallel test for each would round out confidence.

---

Logic review of CreateTeamInMemory

The CreateTeamInMemory change is functionally equivalent to the pre-fix code for the budget loop — budgets are still value-copied via b := clone.Budgets[i] before storing &b, so each budget gets its own stack copy. No behavioural change here.

---

Cosmetic changes

The reformatting of TableRateLimit struct literals in two existing tests (TestGovernanceStore_Customer_CalendarAligned_CreateInMemory and TestGovernanceStore_Customer_CalendarAligned_UpdateInMemory) is cosmetic alignment only — no functional change.

---

Overall verdict

The fix is correct and addresses the reported bug. The regression test will fail on the old code and pass with the fix. The application of the same pattern to CreateTeamInMemory and CreateCustomerInMemory is good defensive consistency. The two non-blocking suggestions above are optional improvements.

---

🧠 Learnings used

Learnt from: roroghost17
Repo: maximhq/bifrost PR: 3998
File: framework/configstore/migrations.go:9574-9591
Timestamp: 2026-06-02T16:47:23.453Z
Learning: In maximhq/bifrost reviews, full-table scans in one-time boot migrations should generally be treated as non-blocking optimizations unless there is a clear scale, locking, memory, or correctness risk. Do not escalate such scans as major issues solely because they are broader than necessary; note them as optional improvements when appropriate.

---

Learnt from: CR
Repo: maximhq/bifrost PR: 0
File: .greptile/rules.md:0-0
Timestamp: 2026-06-03T21:47:30.819Z
Learning: Applies to plugins/governance/**/*.go : Review budget, rate-limit, virtual key, and RBAC paths for fail-closed behavior where security is involved.

---

Learnt from: CR
Repo: maximhq/bifrost PR: 0
File: .greptile/rules.md:0-0
Timestamp: 2026-06-03T21:47:30.819Z
Learning: Applies to plugins/governance/**/*.go : Check that fallback and retry attempts do not double-count or undercount usage unexpectedly.

---

Learnt from: roroghost17
Repo: maximhq/bifrost PR: 3998
File: framework/configstore/migrations.go:9538-9562
Timestamp: 2026-06-02T16:35:46.964Z
Learning: In maximhq/bifrost, for customer budget migration logic in framework/configstore/migrations.go (migrationAddCustomerBudgetsToBudgetsTable), duplicate legacy governance_customers.budget_id references are not expected through supported API flows because each customer creates its own budget. Such duplicates would only arise from manual database edits or invalid config.json entries where multiple customers reference the same pre-declared budget ID, so do not treat a missing duplicate-budget_id preflight as a realistic migration blocker unless the PR explicitly supports or validates those unsupported data shapes.

---

Learnt from: BearTS
Repo: maximhq/bifrost PR: 4043
File: transports/bifrost-http/server/server.go:0-0
Timestamp: 2026-06-04T16:19:42.764Z
Learning: In `transports/bifrost-http/server/server.go` (`ReloadVirtualKey`), the VK-scoped model config stale-eviction logic (snapshot in-memory IDs → DB fetch → delete absent IDs) is intentionally best-effort and not protected by a per-VK lock. Concurrent `ReloadVirtualKey` calls can race and temporarily leave or delete in-memory entries incorrectly, but correctness is guaranteed by convergence: the next successful reload re-fetches from DB and restores the correct state. Do not flag this as a race requiring a lock or stricter ordering; the project explicitly rejects adding a VK-scoped lock for this path.

---

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 3934
File: plugins/modelcatalogresolver/main.go:84-118
Timestamp: 2026-06-04T20:54:37.494Z
Learning: In maximhq/bifrost, the "do not set Bifrost reserved context keys from handlers or plugins" rule (learning 3d4b3152) is scoped to `transports/**/*.go`. Plugins in `plugins/**` (including `plugins/modelcatalogresolver/main.go`) are permitted to call `ctx.AppendRoutingEngineLog(...)` and `schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, ...)` as part of normal routing-metadata emission. Do not flag these calls in plugin code as violating the reserved-key rule.

✅ Action performed

Full review finished.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@plugins/governance/store.go`:
- Around line 2897-2912: CreateTeamInMemory currently shallow-copies the team
causing caller-owned nested state to be shared: ensure nested mutable fields are
deep-copied before storing them so later mutations by the caller don't affect
in-memory data; specifically, in CreateTeamInMemory make a deep copy of
clone.Budgets (new slice and new Budget structs) and of clone.RateLimit (new
RateLimit struct) before mutating IsCalendarAligned and calling
gs.budgets.Store(b.ID, &b) and gs.rateLimits.Store(clone.RateLimit.ID,
clone.RateLimit), then store the fully independent clone via
gs.teams.Store(clone.ID, &clone) so gs.CheckTeamBudget and CheckTeamRateLimit
read immutable in-memory copies.
- Around line 3031-3040: CreateCustomerInMemory currently does a shallow copy
(clone := *customer) which can alias the caller's nested Budgets slice and
RateLimit pointer; instead deep-copy the nested structures before storing so
stored pointers reference newly allocated copies. Specifically, allocate a new
slice for Budgets, copy each budget element into the new slice, set
IsCalendarAligned on each new element and store pointers to those new elements
into gs.budgets; if clone.RateLimit != nil, allocate a new RateLimit copy, set
its IsCalendarAligned and store that into gs.rateLimits; finally store a pointer
to the new top-level customer struct (which contains the new Budgets slice and
RateLimit pointer) into gs.customers. Ensure all references saved to gs.budgets,
gs.rateLimits and gs.customers point to newly allocated objects to avoid
aliasing the input.

🪄 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: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 20b9bab5-431c-4dbc-a52a-4e137511465b

📥 Commits

Reviewing files that changed from the base of the PR and between 73d7e7e and a460f35.

📒 Files selected for processing (2)

• plugins/governance/store.go
• plugins/governance/store_test.go

[akshaydeo]

Merge activity

• Jun 5, 12:15 PM UTC: A user started a stack merge that includes this pull request via Graphite.