fix(entitlement): preserve grant usage across balance snapshots - #4846
Open
GAlexIHU wants to merge 10 commits into
Open
fix(entitlement): preserve grant usage across balance snapshots#4846GAlexIHU wants to merge 10 commits into
GAlexIHU wants to merge 10 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThe credit engine now uses complete usage snapshots across balance creation, history, resets, persistence, and entitlement balance reads. Rollover transitions are represented explicitly. Snapshot cloning and validation preserve usage totals and unit configuration across runs. ChangesCredit snapshot contract and persistence
History and rollover transitions
Engine usage propagation
Entitlement balance integration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant EngineRun
participant Reset
participant History
participant SnapshotAdapter
participant EntitlementBalance
EngineRun->>Reset: create reset snapshot and rollover segment
Reset-->>EngineRun: return snapshot and rollover history
EngineRun->>History: build history from starting snapshot
History-->>EngineRun: return accumulated usage and grant usage
EngineRun->>SnapshotAdapter: persist complete usage snapshot
SnapshotAdapter-->>EntitlementBalance: provide latest valid snapshot
EntitlementBalance->>EntitlementBalance: calculate usage and available grant amount
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
GAlexIHU
force-pushed
the
codex/entitlement-balance-snapshots
branch
from
August 3, 2026 16:46
e365ec2 to
59058ea
Compare
GAlexIHU
marked this pull request as ready for review
August 3, 2026 16:48
![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: 3
🧹 Nitpick comments (2)
openmeter/credit/engine/history.go (1)
104-165: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the rollover-ordering comparator into a named helper.
The
sort.SliceStablecomparator at lines 123-129 encodes a real domain rule: rollover segments must sort before usage segments at the same timestamp. Move this into a named helper, for examplesegmentsLessForSortorrolloverPrecedesUsageAtSameTime. This keeps the domain rule visible outside the closure and makes it independently testable.The rest of
NewGrantBurnDownHistory(copy, per-segment ordering check, anchor validation, contiguity validation) is correct and matches the test coverage inhistory_test.go.Based on learnings, AGENTS.md states: "Do not hide type switching, validation, persistence mapping, or meaningful domain translation inside local closures. Use a named helper; reserve inline callbacks for obvious, tiny logic."
♻️ Proposed extraction
- sort.SliceStable(s, func(i, j int) bool { - if s[i].ClosedPeriod.From.Equal(s[j].ClosedPeriod.From) { - return s[i].TerminationReasons.Rollover && !s[j].TerminationReasons.Rollover - } - - return s[i].ClosedPeriod.From.Before(s[j].ClosedPeriod.From) - }) + sort.SliceStable(s, func(i, j int) bool { + return segmentLess(s[i], s[j]) + }) +} + +// segmentLess orders segments chronologically. Rollover transitions precede +// regular segments at the same timestamp so the reset transition is applied +// before new-period usage. +func segmentLess(a, b GrantBurnDownHistorySegment) bool { + if a.ClosedPeriod.From.Equal(b.ClosedPeriod.From) { + return a.TerminationReasons.Rollover && !b.TerminationReasons.Rollover + } + + return a.ClosedPeriod.From.Before(b.ClosedPeriod.From)🤖 Prompt for 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. In `@openmeter/credit/engine/history.go` around lines 104 - 165, Extract the domain-specific comparator from sort.SliceStable in NewGrantBurnDownHistory into a named helper such as segmentsLessForSort, preserving chronological ordering and ensuring rollover segments precede non-rollover segments at equal timestamps. Pass the helper to sort.SliceStable so the ordering rule is independently testable.Source: Coding guidelines
openmeter/credit/balance/balance_test.go (1)
93-114: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend
TestSnapshotCloneCopiesUsageSnapshotto coverBalancesandUnitConfigindependence.This test only checks that mutating the clone's
UsageSnapshotdoes not affect the original.Snapshot.Clone()also clonesBalancesandUnitConfigindependently. Add assertions for those two fields so a future regression in either clone path is caught here.♻️ Suggested extension
func TestSnapshotCloneCopiesUsageSnapshot(t *testing.T) { snapshot := balance.Snapshot{ + Balances: balance.Map{"grant-1": 100}, UsageSnapshot: &balance.UsageSnapshot{ Usage: 5, TotalGrantUsage: 10, }, } cloned := snapshot.Clone() + cloned.Balances["grant-1"] = 999 + assert.Equal(t, 100.0, snapshot.Balances["grant-1"]) require.NotNil(t, cloned.UsageSnapshot) ...As per path instructions, "Make sure the tests are comprehensive and cover the changes," and
Clone()is new behavior introduced by this PR without full field coverage.🤖 Prompt for 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. In `@openmeter/credit/balance/balance_test.go` around lines 93 - 114, Extend TestSnapshotCloneCopiesUsageSnapshot to initialize representative Balances and UnitConfig values, then assert the cloned fields match the original values and are independent. Mutate the clone’s Balances and UnitConfig after cloning and verify the corresponding fields on the original snapshot remain unchanged, alongside the existing UsageSnapshot assertions.Source: Path instructions
🤖 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 `@openmeter/credit/engine/reset.go`:
- Around line 51-86: Clone rolledOver before passing it to burnDownGrants in the
reset flow so in-place mutations cannot alter the original rollover balances.
Keep balances initialized from rolledOver for the reset snapshot, but pass a
cloned map as burnDownGrants’ balance input; preserve BalanceAtStart as the
pre-burn rolled-over state.
In `@openmeter/credit/engine/run.go`:
- Around line 251-277: Preserve the conversion configuration when constructing
the snapshot returned by runBetweenResets: copy
params.StartingSnapshot.UnitConfig into the new balance.Snapshot alongside the
existing fields. Also update reset() to copy UnitConfig into its returned reset
snapshot so subsequent engine segments retain the same configuration.
In `@openmeter/entitlement/README.md`:
- Around line 27-28: Update the sentence in the entitlement documentation so the
singular subject “Entitlement” uses the possessive pronoun “its” instead of
“their,” leaving the rest of the wording unchanged.
---
Nitpick comments:
In `@openmeter/credit/balance/balance_test.go`:
- Around line 93-114: Extend TestSnapshotCloneCopiesUsageSnapshot to initialize
representative Balances and UnitConfig values, then assert the cloned fields
match the original values and are independent. Mutate the clone’s Balances and
UnitConfig after cloning and verify the corresponding fields on the original
snapshot remain unchanged, alongside the existing UsageSnapshot assertions.
In `@openmeter/credit/engine/history.go`:
- Around line 104-165: Extract the domain-specific comparator from
sort.SliceStable in NewGrantBurnDownHistory into a named helper such as
segmentsLessForSort, preserving chronological ordering and ensuring rollover
segments precede non-rollover segments at equal timestamps. Pass the helper to
sort.SliceStable so the ordering rule is independently testable.
🪄 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: CHILL
Plan: Pro Plus
Run ID: d267b431-eec1-4a79-a976-00385739403a
⛔ Files ignored due to path filters (9)
openmeter/ent/db/balancesnapshot.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot/balancesnapshot.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot/where.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot_create.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot_update.gois excluded by!**/ent/db/**openmeter/ent/db/migrate/schema.gois excluded by!**/ent/db/**openmeter/ent/db/mutation.gois excluded by!**/ent/db/**openmeter/ent/db/runtime.gois excluded by!**/ent/db/**tools/migrate/migrations/atlas.sumis excluded by!**/*.sum,!**/*.sum
📒 Files selected for processing (34)
AGENTS.mdopenmeter/credit/README.mdopenmeter/credit/adapter/balance_snapshot.goopenmeter/credit/balance/balance.goopenmeter/credit/balance/balance_test.goopenmeter/credit/balance/repository.goopenmeter/credit/balance/service.goopenmeter/credit/balance/service_test.goopenmeter/credit/engine/engine.goopenmeter/credit/engine/engine_test.goopenmeter/credit/engine/history.goopenmeter/credit/engine/history_test.goopenmeter/credit/engine/reset.goopenmeter/credit/engine/reset_test.goopenmeter/credit/engine/run.goopenmeter/credit/engine/run_test.goopenmeter/credit/engine/runresult_test.goopenmeter/credit/engine/snapshot_test.goopenmeter/credit/helper.goopenmeter/ent/schema/balance_snapshot.goopenmeter/entitlement/README.mdopenmeter/entitlement/metered/balance.goopenmeter/entitlement/metered/balance_test.goopenmeter/entitlement/metered/balance_total_available_test.goopenmeter/entitlement/metered/balance_unitconfig_test.goopenmeter/entitlement/metered/lateevents_test.goopenmeter/entitlement/metered/reset_test.goopenmeter/entitlement/metered/utils_test.goopenmeter/registry/builder/entitlement.goopenmeter/subscription/README.mdtest/billing/subscription_suite.gotest/entitlement/regression/framework_test.gotools/migrate/migrations/20260803135655_add_balance_snapshot_usage_snapshot.down.sqltools/migrate/migrations/20260803135655_add_balance_snapshot_usage_snapshot.up.sql
💤 Files with no reviewable changes (1)
- openmeter/credit/balance/service_test.go
Comment on lines
+27
to
+28
| that materializes and supersedes subscription-managed entitlements. | ||
| Entitlement owns their persisted lifecycle and value resolution. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the pronoun-agreement slip.
"Entitlement" is singular, so "their" should be "its" in this sentence.
✏️ Proposed fix
-Entitlement owns their persisted lifecycle and value resolution.
+Entitlement owns its persisted lifecycle and value resolution.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| that materializes and supersedes subscription-managed entitlements. | |
| Entitlement owns their persisted lifecycle and value resolution. | |
| that materializes and supersedes subscription-managed entitlements. | |
| Entitlement owns its persisted lifecycle and value resolution. |
🤖 Prompt for 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.
In `@openmeter/entitlement/README.md` around lines 27 - 28, Update the sentence in
the entitlement documentation so the singular subject “Entitlement” uses the
possessive pronoun “its” instead of “their,” leaving the rest of the wording
unchanged.
Source: Path instructions
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
openmeter/credit/engine/snapshot_test.go (1)
327-353: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd lifecycle intent comments to the subtest.
Add concise
given,when, andthencomments. This test covers both normal execution and reset rollover behavior.As per coding guidelines, “Begin non-trivial service or lifecycle subtests with concise
given,when, andthenintent comments.”🤖 Prompt for 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. In `@openmeter/credit/engine/snapshot_test.go` around lines 327 - 353, Add concise given, when, and then intent comments inside the subtest around the starting snapshot and engine setup, the eng.Run invocation, and the snapshot unit-config assertions. Ensure the comments describe both normal execution and reset rollover behavior without changing test logic.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@openmeter/credit/engine/snapshot_test.go`:
- Around line 327-353: Add concise given, when, and then intent comments inside
the subtest around the starting snapshot and engine setup, the eng.Run
invocation, and the snapshot unit-config assertions. Ensure the comments
describe both normal execution and reset rollover behavior without changing test
logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 03dad9eb-620b-4498-8022-f16ffc861c2b
📒 Files selected for processing (3)
openmeter/credit/engine/reset.goopenmeter/credit/engine/run.goopenmeter/credit/engine/snapshot_test.go
🚧 Files skipped from review as they are similar to previous changes (2)
- openmeter/credit/engine/reset.go
- openmeter/credit/engine/run.go
chrisgacsal
previously approved these changes
Aug 4, 2026
GAlexIHU
force-pushed
the
codex/entitlement-balance-snapshots
branch
from
August 4, 2026 09:23
83dc569 to
67823cb
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
openmeter/credit/engine/history.go (1)
60-71: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueConsider validating
UsageSnapshot.Usagetoo.
validateStartingSnapshotguardsTotalGrantUsageagainst NaN, Inf, and negative values.UsageSnapshot.Usagegets no such guard, andgetUsageSnapshotAtStartOfSegmentaccumulates it into every reconstructed segment snapshot at line 207. A NaN or Inf there silently poisons every downstream usage figure the same way. Same check, same place, cheap to add.♻️ Optional: extend the finiteness check
+ usage := snapshot.UsageSnapshot.Usage + if usage < 0 || math.IsNaN(usage) || math.IsInf(usage, 0) { + return fmt.Errorf("starting snapshot usage must be a finite non-negative number, got %v", usage) + } + totalGrantUsage := snapshot.UsageSnapshot.TotalGrantUsage if totalGrantUsage < 0 || math.IsNaN(totalGrantUsage) || math.IsInf(totalGrantUsage, 0) { return fmt.Errorf("starting snapshot total grant usage must be a finite non-negative number, got %v", totalGrantUsage) }🤖 Prompt for 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. In `@openmeter/credit/engine/history.go` around lines 60 - 71, The validateStartingSnapshot function validates snapshot.UsageSnapshot.TotalGrantUsage for NaN, Inf, and negative values, but does not validate snapshot.UsageSnapshot.Usage with the same checks. Since Usage is accumulated into reconstructed segment snapshots downstream, an invalid Usage value (NaN or Inf) can poison all downstream usage figures. Add the same finiteness and non-negative validation check for UsageSnapshot.Usage in the validateStartingSnapshot function that already exists for TotalGrantUsage, using the same error-reporting pattern to maintain consistency.
🤖 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.
Nitpick comments:
In `@openmeter/credit/engine/history.go`:
- Around line 60-71: The validateStartingSnapshot function validates
snapshot.UsageSnapshot.TotalGrantUsage for NaN, Inf, and negative values, but
does not validate snapshot.UsageSnapshot.Usage with the same checks. Since Usage
is accumulated into reconstructed segment snapshots downstream, an invalid Usage
value (NaN or Inf) can poison all downstream usage figures. Add the same
finiteness and non-negative validation check for UsageSnapshot.Usage in the
validateStartingSnapshot function that already exists for TotalGrantUsage, using
the same error-reporting pattern to maintain consistency.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: c1ed8c37-c563-482a-af8a-dadbc62f70dc
⛔ Files ignored due to path filters (9)
openmeter/ent/db/balancesnapshot.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot/balancesnapshot.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot/where.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot_create.gois excluded by!**/ent/db/**openmeter/ent/db/balancesnapshot_update.gois excluded by!**/ent/db/**openmeter/ent/db/migrate/schema.gois excluded by!**/ent/db/**openmeter/ent/db/mutation.gois excluded by!**/ent/db/**openmeter/ent/db/runtime.gois excluded by!**/ent/db/**tools/migrate/migrations/atlas.sumis excluded by!**/*.sum,!**/*.sum
📒 Files selected for processing (34)
AGENTS.mdopenmeter/credit/README.mdopenmeter/credit/adapter/balance_snapshot.goopenmeter/credit/balance/balance.goopenmeter/credit/balance/balance_test.goopenmeter/credit/balance/repository.goopenmeter/credit/balance/service.goopenmeter/credit/balance/service_test.goopenmeter/credit/engine/engine.goopenmeter/credit/engine/engine_test.goopenmeter/credit/engine/history.goopenmeter/credit/engine/history_test.goopenmeter/credit/engine/reset.goopenmeter/credit/engine/reset_test.goopenmeter/credit/engine/run.goopenmeter/credit/engine/run_test.goopenmeter/credit/engine/runresult_test.goopenmeter/credit/engine/snapshot_test.goopenmeter/credit/helper.goopenmeter/ent/schema/balance_snapshot.goopenmeter/entitlement/README.mdopenmeter/entitlement/metered/balance.goopenmeter/entitlement/metered/balance_test.goopenmeter/entitlement/metered/balance_total_available_test.goopenmeter/entitlement/metered/balance_unitconfig_test.goopenmeter/entitlement/metered/lateevents_test.goopenmeter/entitlement/metered/reset_test.goopenmeter/entitlement/metered/utils_test.goopenmeter/registry/builder/entitlement.goopenmeter/subscription/README.mdtest/billing/subscription_suite.gotest/entitlement/regression/framework_test.gotools/migrate/migrations/20260803135655_add_balance_snapshot_usage_snapshot.down.sqltools/migrate/migrations/20260803135655_add_balance_snapshot_usage_snapshot.up.sql
💤 Files with no reviewable changes (1)
- openmeter/credit/balance/service_test.go
🚧 Files skipped from review as they are similar to previous changes (31)
- AGENTS.md
- openmeter/subscription/README.md
- openmeter/credit/engine/engine_test.go
- openmeter/credit/helper.go
- openmeter/entitlement/metered/balance_unitconfig_test.go
- test/billing/subscription_suite.go
- openmeter/entitlement/metered/balance.go
- openmeter/credit/balance/repository.go
- openmeter/credit/README.md
- test/entitlement/regression/framework_test.go
- openmeter/credit/engine/engine.go
- openmeter/credit/adapter/balance_snapshot.go
- openmeter/entitlement/metered/utils_test.go
- openmeter/ent/schema/balance_snapshot.go
- openmeter/credit/balance/service.go
- openmeter/credit/balance/balance.go
- openmeter/credit/balance/balance_test.go
- openmeter/entitlement/README.md
- openmeter/entitlement/metered/reset_test.go
- openmeter/entitlement/metered/balance_total_available_test.go
- openmeter/entitlement/metered/balance_test.go
- openmeter/credit/engine/snapshot_test.go
- openmeter/credit/engine/run.go
- openmeter/credit/engine/history_test.go
- openmeter/credit/engine/reset_test.go
- tools/migrate/migrations/20260803135655_add_balance_snapshot_usage_snapshot.up.sql
- openmeter/entitlement/metered/lateevents_test.go
- openmeter/credit/engine/reset.go
- openmeter/credit/engine/run_test.go
- openmeter/credit/engine/runresult_test.go
- openmeter/registry/builder/entitlement.go
Signed-off-by: Alex Goth <64845621+GAlexIHU@users.noreply.github.com>
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.
What changed
Why
totalAvailableGrantAmount was reconstructed from grant usage after the latest persisted snapshot. Grant usage from earlier in the same usage period was therefore lost, allowing API responses where usage exceeded total available grants while balance remained positive and overage stayed zero.
A persisted snapshot must be a complete, independent checkpoint. Its result must not depend on earlier snapshots still being present.
Impact
Entitlement balance reads now preserve the full usage-period grant history when resuming from snapshots. Existing snapshots without the new usage state are ignored and recalculated from measurement start until new checkpoints are written. The legacy field remains written so old and new binaries can overlap during deployment.
Checks
Deployment notes
usage_snapshot) are ignored after deploy; each owner's first balance read replays from measurement start until a new complete checkpoint lands. Expect a temporary balance-worker load increase and verify ClickHouse retention still covers the oldest active entitlements before deploying.preserveOverageAtResetbecause per-call reset overrides are not persisted. If explicit overrides were used, recomputed balances may differ from the previously checkpointed values.usage_snapshotin every snapshot read and write, so dropping the column first breaks entitlement balance APIs.Summary by CodeRabbit
New Features
Bug Fixes
Greptile Summary
The PR makes balance snapshots independent, resumable checkpoints that retain cumulative usage-period grant consumption.
Confidence Score: 5/5
The PR appears safe to merge.
No blocking failure remains.
Important Files Changed
Flowchart
%%{init: {'theme': 'neutral'}}%% flowchart LR Read[Balance read] --> Select{Compatible snapshot exists?} Select -->|Yes| Snapshot[Load complete snapshot] Select -->|No| Start[Create measurement-start snapshot] Snapshot --> Engine[Resume credit engine] Start --> Engine Engine --> Usage[Burn metered usage] Usage --> Reset{Usage reset?} Reset -->|Yes| Rollover[Apply rollover and preserved overage burn] Rollover --> History[Record explicit rollover segment] Reset -->|No| History History --> Result[Return balance and complete usage state] History --> Persist[Persist eligible resumable checkpoint]Reviews (5): Last reviewed commit: "Merge branch 'main' into codex/entitleme..." | Re-trigger Greptile
Context used: