diff --git a/docs/features/governance/budget-and-limits.mdx b/docs/features/governance/budget-and-limits.mdx
index b037e41205..e3d078dba4 100644
--- a/docs/features/governance/budget-and-limits.mdx
+++ b/docs/features/governance/budget-and-limits.mdx
@@ -183,18 +183,18 @@ It sets the alignment **mode**, not a shared reset instant. An owner holds one w
Alignment only applies to durations that have a calendar boundary: `1d`, `1w`, `1M`, `1Q` and `1Y`. A **sub-day** window such as a `1h` request limit stays on its rolling schedule even when its owner is aligned, and the flag changes nothing for it.
-Turning calendar alignment **on** for an existing owner never moves `last_reset` backwards, so an already-open window is not re-dated. What happens next is decided **per window**, against that window's own boundary:
+Turning calendar alignment **on** for an existing owner **keeps its accumulated usage**. Each open window is adopted onto the calendar grid instead of being reset: its start moves forward to the boundary it now follows, its usage carries over untouched, and its first aligned reset happens at the **next** boundary.
-- **Window opened after its most recent boundary** - nothing resets now. It keeps its start and its accumulated usage, and its first aligned reset happens at its next boundary.
-- **Window opened before its most recent boundary** - it is **already due** under the new rule, so the next reset evaluation resets it and **clears its current usage**: spend for a budget, the token or request counter for a rate limit. For a monthly window switched on partway through a month, this is the common case.
+So a monthly budget with $42 of spend, aligned on 5 February, keeps the $42 and resets on 1 March. Nothing is cleared at the switch, and nothing is forgiven.
-A reset evaluation is not only the background sweep. The same due-or-not check runs again on the request path, just before a request's usage is recorded, so an already-due window is cleared by the next request even if no sweep has fired yet.
+Two details follow from each window being adopted on its own terms:
-Because each window is judged separately, one switch can leave a `1d` limit untouched while clearing a `1M` budget on the same owner.
+- **Windows do not move together.** A `1M` budget adopts the month boundary while a `1d` counter beside it adopts midnight.
+- **Sub-day windows are not adopted at all**, because they have no boundary to adopt. A `1h` counter keeps its rolling schedule, as described above.
-
-Enabling alignment on an owner with accumulated usage can therefore clear that usage almost immediately, and which of its windows are affected depends on each one's own duration and start. There is generally no moment when every window is freshly reset, so to preserve them all set `calendar_aligned` at **create** time and avoid the switch-over entirely. To protect one specific window, enable alignment just after **that** window resets.
-
+
+The boundary only ever moves **forward**. A window that already opened after its most recent boundary is current, so it is left exactly as it is. This forward-only rule is what keeps every node in a cluster agreeing on which window is open.
+
### Quarterly budgets and fiscal quarters
diff --git a/framework/configstore/tables/budget.go b/framework/configstore/tables/budget.go
index 81ada9c430..0d7b700810 100644
--- a/framework/configstore/tables/budget.go
+++ b/framework/configstore/tables/budget.go
@@ -238,6 +238,40 @@ func (b *TableBudget) WindowStart(now time.Time) time.Time {
return RollingWindowStart(anchor, duration, now)
}
+// AdoptCalendarAlignment re-anchors an already-open window onto the calendar grid
+// it has just been told to follow, and reports whether it moved.
+//
+// Called once, on the transition from rolling to calendar-aligned. Without it the
+// switch is destructive: the reset path resets whenever WindowStart(now) is after
+// LastReset, and for a freshly aligned budget WindowStart is the most recent
+// boundary, so any window that opened before that boundary is instantly due and
+// its usage is cleared. That is most of the month for a monthly budget.
+//
+// Moving LastReset up to the boundary makes the window current instead of overdue,
+// so nothing resets now and the first aligned reset lands on the next boundary.
+// The move is forward-only, which is the one direction every persistence path
+// allows: rewinding would re-open a window the cluster already treated as spent.
+// CurrentUsage is deliberately untouched - the operator asked to change a
+// schedule, not to forgive spend.
+//
+// Returns false, changing nothing, when the budget is not aligned, when its
+// duration has no calendar boundary (sub-day windows stay rolling), or when the
+// window already opened at or after the boundary and is therefore current.
+func (b *TableBudget) AdoptCalendarAlignment(now time.Time) bool {
+ if b == nil || !b.IsCalendarAligned || !IsCalendarAlignableDuration(b.ResetDuration) {
+ return false
+ }
+ // Read through WindowStart rather than GetCalendarPeriodStart directly so the
+ // boundary adopted here can never drift from the one the reset path compares
+ // against.
+ start := b.WindowStart(now)
+ if !start.After(b.LastReset) {
+ return false
+ }
+ b.LastReset = start
+ return true
+}
+
// StampCalendarAlignment copies an owner's calendar-alignment flag onto the
// budgets and rate limit it owns.
//
diff --git a/framework/configstore/tables/budgetwindow_test.go b/framework/configstore/tables/budgetwindow_test.go
index c09bf038af..9dab489f0e 100644
--- a/framework/configstore/tables/budgetwindow_test.go
+++ b/framework/configstore/tables/budgetwindow_test.go
@@ -607,3 +607,145 @@ func TestQuarterlyBudgetResetIsNotPerpetuallyDue(t *testing.T) {
fmt.Sprintf("quarter start %s: budget became due an hour into its quarter", quarterStart))
}
}
+
+// TestAdoptCalendarAlignmentPreservesTheOpenWindow pins the switch-over contract:
+// turning calendar alignment on must not cost the owner its accumulated usage.
+//
+// Without adoption, a window opened before the current boundary is already due
+// the moment the flag flips, because budgetResetTarget resets whenever
+// WindowStart(now) is after LastReset. Adoption moves LastReset forward onto the
+// boundary instead, which the forward-only guard permits, so the window becomes
+// current rather than overdue and its usage survives to the next real boundary.
+func TestAdoptCalendarAlignmentPreservesTheOpenWindow(t *testing.T) {
+ now := time.Date(2026, time.February, 5, 12, 0, 0, 0, time.UTC)
+ monthStart := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC)
+
+ newBudget := func(duration string, lastReset time.Time) *TableBudget {
+ return &TableBudget{
+ ID: "adopt-budget",
+ MaxLimit: 100,
+ CurrentUsage: 42,
+ ResetDuration: duration,
+ IsCalendarAligned: true,
+ CreatedAt: lastReset,
+ LastReset: lastReset,
+ }
+ }
+
+ t.Run("a window opened before the boundary is adopted onto it", func(t *testing.T) {
+ budget := newBudget("1M", time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC))
+
+ assert.True(t, budget.AdoptCalendarAlignment(now), "the window moved, so adoption reports a change")
+ assert.True(t, budget.LastReset.Equal(monthStart), "the window is re-anchored on the boundary it now follows")
+ assert.Equal(t, 42.0, budget.CurrentUsage, "adoption must never spend the operator's usage")
+ assert.False(t, budget.WindowStart(now).After(budget.LastReset),
+ "and the budget must no longer read as due, which is the whole point")
+ })
+
+ t.Run("a window opened after the boundary is left alone", func(t *testing.T) {
+ lastReset := time.Date(2026, time.February, 3, 9, 0, 0, 0, time.UTC)
+ budget := newBudget("1M", lastReset)
+
+ assert.False(t, budget.AdoptCalendarAlignment(now), "nothing to move")
+ assert.True(t, budget.LastReset.Equal(lastReset), "an already-current window keeps its own start")
+ })
+
+ t.Run("the boundary is never moved backwards", func(t *testing.T) {
+ lastReset := time.Date(2026, time.February, 20, 9, 0, 0, 0, time.UTC)
+ budget := newBudget("1M", lastReset)
+
+ assert.False(t, budget.AdoptCalendarAlignment(now))
+ assert.True(t, budget.LastReset.Equal(lastReset),
+ "rewinding would re-open a window the cluster already agreed was spent")
+ })
+
+ t.Run("a sub-day window is not adopted", func(t *testing.T) {
+ lastReset := time.Date(2026, time.February, 5, 9, 30, 0, 0, time.UTC)
+ budget := newBudget("1h", lastReset)
+
+ assert.False(t, budget.AdoptCalendarAlignment(now),
+ "an hourly window has no calendar boundary, so alignment leaves it rolling")
+ assert.True(t, budget.LastReset.Equal(lastReset))
+ })
+
+ t.Run("an unaligned budget is not adopted", func(t *testing.T) {
+ lastReset := time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC)
+ budget := newBudget("1M", lastReset)
+ budget.IsCalendarAligned = false
+
+ assert.False(t, budget.AdoptCalendarAlignment(now),
+ "adoption belongs to the switch-over, not to every config write")
+ assert.True(t, budget.LastReset.Equal(lastReset))
+ })
+
+ t.Run("a quarterly budget adopts its fiscal boundary", func(t *testing.T) {
+ budget := newBudget("1Q", time.Date(2025, time.December, 10, 9, 0, 0, 0, time.UTC))
+ budget.ResetConfig = &BudgetResetConfig{QuarterStartMonth: int(time.February)}
+
+ assert.True(t, budget.AdoptCalendarAlignment(now))
+ assert.True(t, budget.LastReset.Equal(time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC)),
+ "a February fiscal year opens Q1 on Feb 1, not on the calendar quarter")
+ })
+}
+
+// TestAdoptCalendarAlignmentOnRateLimits mirrors the budget contract for the two
+// counters a rate limit carries.
+//
+// Token and request counters keep independent durations and independent
+// LastReset values, so adoption is decided per counter. rateLimitResetTarget
+// applies the same "boundary after last reset" rule as budgets, which means the
+// same switch-over would otherwise clear a counter the operator was still using.
+func TestAdoptCalendarAlignmentOnRateLimits(t *testing.T) {
+ now := time.Date(2026, time.February, 5, 12, 0, 0, 0, time.UTC)
+ monthStart := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC)
+ dayStart := time.Date(2026, time.February, 5, 0, 0, 0, 0, time.UTC)
+ stale := time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC)
+
+ newRateLimit := func(tokenDuration, requestDuration string) *TableRateLimit {
+ return &TableRateLimit{
+ ID: "adopt-rate-limit",
+ TokenResetDuration: &tokenDuration,
+ TokenCurrentUsage: 900,
+ TokenLastReset: stale,
+ RequestResetDuration: &requestDuration,
+ RequestCurrentUsage: 17,
+ RequestLastReset: stale,
+ IsCalendarAligned: true,
+ }
+ }
+
+ t.Run("each counter adopts its own boundary", func(t *testing.T) {
+ rl := newRateLimit("1M", "1d")
+
+ assert.True(t, rl.AdoptCalendarAlignment(now), "both counters moved")
+ assert.True(t, rl.TokenLastReset.Equal(monthStart), "the monthly counter adopts the month boundary")
+ assert.True(t, rl.RequestLastReset.Equal(dayStart), "the daily counter adopts midnight, not the month boundary")
+ assert.Equal(t, int64(900), rl.TokenCurrentUsage, "adoption must not clear a counter")
+ assert.Equal(t, int64(17), rl.RequestCurrentUsage)
+ })
+
+ t.Run("a sub-day counter is left rolling while its sibling adopts", func(t *testing.T) {
+ rl := newRateLimit("1M", "1h")
+
+ assert.True(t, rl.AdoptCalendarAlignment(now), "the monthly counter still moves")
+ assert.True(t, rl.TokenLastReset.Equal(monthStart))
+ assert.True(t, rl.RequestLastReset.Equal(stale),
+ "an hourly counter has no calendar boundary, so it keeps its rolling anchor")
+ })
+
+ t.Run("an unaligned rate limit is not adopted", func(t *testing.T) {
+ rl := newRateLimit("1M", "1d")
+ rl.IsCalendarAligned = false
+
+ assert.False(t, rl.AdoptCalendarAlignment(now))
+ assert.True(t, rl.TokenLastReset.Equal(stale))
+ assert.True(t, rl.RequestLastReset.Equal(stale))
+ })
+
+ t.Run("counters with no duration are skipped", func(t *testing.T) {
+ rl := &TableRateLimit{ID: "no-durations", IsCalendarAligned: true, TokenLastReset: stale, RequestLastReset: stale}
+
+ assert.False(t, rl.AdoptCalendarAlignment(now))
+ assert.True(t, rl.TokenLastReset.Equal(stale))
+ })
+}
diff --git a/framework/configstore/tables/ratelimit.go b/framework/configstore/tables/ratelimit.go
index 086a9449a2..d714f3d404 100644
--- a/framework/configstore/tables/ratelimit.go
+++ b/framework/configstore/tables/ratelimit.go
@@ -85,3 +85,36 @@ func (rl *TableRateLimit) BeforeSave(tx *gorm.DB) error {
return nil
}
+
+// AdoptCalendarAlignment re-anchors this rate limit's open counters onto the
+// calendar grid they have just been told to follow, and reports whether either
+// moved. Counterpart to TableBudget.AdoptCalendarAlignment; the reasoning for
+// moving the boundary rather than resetting is documented there.
+//
+// The two counters are decided independently: they carry their own durations and
+// their own LastReset values, so a monthly token limit can adopt the month
+// boundary while an hourly request limit stays rolling in the same call.
+// Usage is never cleared - the operator changed a schedule, not a quota.
+func (rl *TableRateLimit) AdoptCalendarAlignment(now time.Time) bool {
+ if rl == nil || !rl.IsCalendarAligned {
+ return false
+ }
+ adopt := func(duration *string, lastReset *time.Time) bool {
+ if duration == nil || !IsCalendarAlignableDuration(*duration) {
+ return false
+ }
+ // Rate limits carry no quarter definition, so the fiscal start is not
+ // applicable here; see QuarterStartNotApplicable.
+ start := GetCalendarPeriodStart(*duration, now, QuarterStartNotApplicable)
+ if !start.After(*lastReset) {
+ return false
+ }
+ *lastReset = start
+ return true
+ }
+ // Both counters are evaluated; || would short-circuit and leave the second
+ // one on a stale boundary whenever the first adopted.
+ tokenMoved := adopt(rl.TokenResetDuration, &rl.TokenLastReset)
+ requestMoved := adopt(rl.RequestResetDuration, &rl.RequestLastReset)
+ return tokenMoved || requestMoved
+}
diff --git a/plugins/governance/budgetcycle_test.go b/plugins/governance/budgetcycle_test.go
index 0a792ee482..026f07a585 100644
--- a/plugins/governance/budgetcycle_test.go
+++ b/plugins/governance/budgetcycle_test.go
@@ -857,9 +857,13 @@ func TestDumpBudgetsCannotUndoOperatorReset(t *testing.T) {
// creates the budget at test time, so LastReset is always inside the current
// period. Both cases are pinned below so the difference is visible.
//
-// Closing the gap needs either an activation marker or a dedicated write path,
-// because UpdateBudget copies last_reset back from the stored row on every config
-// write. Tracked as follow-up work; until then the docs describe this behaviour.
+// The gap is closed by adoption: the owner handlers now call
+// AdoptCalendarAlignmentInMemory on the switch-over, which moves each open window
+// forward onto its boundary so the sweep below never finds it overdue. These
+// assertions stay as they are on purpose - budgetResetTarget itself is unchanged,
+// and it is precisely its "reset whatever is overdue" rule that makes adoption
+// necessary. Read them as the reason the switch-over needs a step, not as a
+// description of what an operator sees.
func TestEnablingCalendarAlignmentCanResetAtTheBoundaryAlreadyPassed(t *testing.T) {
store := newStandaloneStore(t)
now := time.Date(2026, time.February, 5, 12, 0, 0, 0, time.UTC)
diff --git a/plugins/governance/store.go b/plugins/governance/store.go
index 0e459a176c..a5d249ce05 100644
--- a/plugins/governance/store.go
+++ b/plugins/governance/store.go
@@ -164,6 +164,13 @@ type GovernanceStore interface {
// ResetBudgetUsageInMemory clears one budget's usage on operator request,
// leaving its reset boundary untouched.
ResetBudgetUsageInMemory(ctx context.Context, budgetID string) (*configstoreTables.TableBudget, bool)
+ // AdoptCalendarAlignmentInMemory re-anchors a budget onto the calendar boundary
+ // it has just been told to follow, preserving its usage. Reports whether the
+ // boundary moved.
+ AdoptCalendarAlignmentInMemory(ctx context.Context, budgetID string, now time.Time) bool
+ // AdoptRateLimitCalendarAlignmentInMemory is the rate-limit counterpart,
+ // applying the rule to the token and request counters independently.
+ AdoptRateLimitCalendarAlignmentInMemory(ctx context.Context, rateLimitID string, now time.Time) bool
// DB sync for expired items
ResetExpiredRateLimits(ctx context.Context, resetRateLimits []*configstoreTables.TableRateLimit) error
ResetExpiredBudgets(ctx context.Context, resetBudgets []*configstoreTables.TableBudget) error
@@ -2127,6 +2134,69 @@ func (gs *LocalGovernanceStore) ResetBudgetUsageInMemory(ctx context.Context, bu
return reset, true
}
+// AdoptCalendarAlignmentInMemory re-anchors a budget onto the calendar grid when
+// its owner has just switched alignment on, and reports whether it moved.
+//
+// Called on the false-to-true transition only. Without it the switch is
+// destructive: budgetResetTarget resets whenever WindowStart(now) is after
+// LastReset, so a window that opened before the current boundary is instantly due
+// and the next sweep clears its usage. See TableBudget.AdoptCalendarAlignment.
+//
+// The usage carried into the swap is the one observed inside the CAS loop, not a
+// value read beforehand: a request bumping this budget between a read and the
+// write would otherwise have its spend dropped. No reset generation is bumped
+// here - nothing was reset, so an in-flight dump's rows are still accurate.
+func (gs *LocalGovernanceStore) AdoptCalendarAlignmentInMemory(ctx context.Context, budgetID string, now time.Time) bool {
+ for {
+ raw, exists := gs.budgets.Load(budgetID)
+ if !exists || raw == nil {
+ return false
+ }
+ old, ok := raw.(*configstoreTables.TableBudget)
+ if !ok || old == nil {
+ return false
+ }
+ clone := *old
+ // The owner's flag has already been applied upstream, but the in-memory
+ // copy predates it, so stamp it here or adoption declines its own work.
+ clone.IsCalendarAligned = true
+ if !clone.AdoptCalendarAlignment(now) {
+ return false
+ }
+ clone.RefreshOverrideCyclesRemaining()
+ if gs.budgets.CompareAndSwap(budgetID, raw, &clone) {
+ gs.logger.Debug(fmt.Sprintf("Adopted budget %s onto its calendar boundary %s (usage %.2f preserved)",
+ budgetID, clone.LastReset, clone.CurrentUsage))
+ return true
+ }
+ }
+}
+
+// AdoptRateLimitCalendarAlignmentInMemory is the rate-limit counterpart, applying
+// the same switch-over rule to the token and request counters independently.
+func (gs *LocalGovernanceStore) AdoptRateLimitCalendarAlignmentInMemory(ctx context.Context, rateLimitID string, now time.Time) bool {
+ for {
+ raw, exists := gs.rateLimits.Load(rateLimitID)
+ if !exists || raw == nil {
+ return false
+ }
+ old, ok := raw.(*configstoreTables.TableRateLimit)
+ if !ok || old == nil {
+ return false
+ }
+ clone := *old
+ clone.IsCalendarAligned = true
+ if !clone.AdoptCalendarAlignment(now) {
+ return false
+ }
+ if gs.rateLimits.CompareAndSwap(rateLimitID, raw, &clone) {
+ gs.logger.Debug(fmt.Sprintf("Adopted rate limit %s onto its calendar boundaries (token %s, request %s)",
+ rateLimitID, clone.TokenLastReset, clone.RequestLastReset))
+ return true
+ }
+ }
+}
+
// ResetExpiredBudgetsInMemory checks and resets budgets that have exceeded their reset duration.
// With no budgetIDs it scans every budget; with IDs it only checks those budgets.
// refreshReferences controls whether embedded owner references (VK, team, customer) are updated
diff --git a/plugins/governance/storeconcurrency_test.go b/plugins/governance/storeconcurrency_test.go
index d02269c090..5f51b5f8e6 100644
--- a/plugins/governance/storeconcurrency_test.go
+++ b/plugins/governance/storeconcurrency_test.go
@@ -145,3 +145,64 @@ func TestResetBudgetAt_ConcurrentResettersCollapse(t *testing.T) {
assert.True(t, final.LastReset.Equal(newLastReset))
assert.Equal(t, 4, final.OverrideCyclesRemaining, "the single winning reset should consume exactly one override cycle")
}
+
+// TestAdoptCalendarAlignmentInMemoryPreservesConcurrentSpend pins that adopting a
+// budget onto the calendar grid keeps whatever usage landed while the switch was
+// in flight.
+//
+// Adoption cannot read usage, then write it back: a request bumping the same
+// budget between those two steps would have its spend silently dropped. The CAS
+// loop has to carry the usage it observed at swap time, which is what this test
+// forces by bumping usage from another goroutine during the adoption.
+func TestAdoptCalendarAlignmentInMemoryPreservesConcurrentSpend(t *testing.T) {
+ ctx := context.Background()
+ store := newStandaloneStore(t)
+ now := time.Date(2026, time.February, 5, 12, 0, 0, 0, time.UTC)
+ monthStart := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC)
+
+ budget := &configstoreTables.TableBudget{
+ ID: "adopt-live-budget",
+ MaxLimit: 1000,
+ CurrentUsage: 0,
+ ResetDuration: "1M",
+ IsCalendarAligned: true,
+ CreatedAt: time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC),
+ LastReset: time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC),
+ }
+ store.budgets.Store(budget.ID, budget)
+
+ // Usage is bumped by direct CAS rather than through BumpBudgetUsage, which
+ // consults the real clock: a January window is long overdue against it, so the
+ // request path would reset the budget onto the current real month and the fixed
+ // `now` below could no longer move it. The property under test is that the
+ // adoption CAS carries whatever usage it observed, and a plain increment races
+ // it just as well without dragging real time into the fixture.
+ const bumps = 50
+ var wg sync.WaitGroup
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ for i := 0; i < bumps; i++ {
+ for {
+ raw, ok := store.budgets.Load(budget.ID)
+ require.True(t, ok)
+ current := raw.(*configstoreTables.TableBudget)
+ clone := *current
+ clone.CurrentUsage++
+ if store.budgets.CompareAndSwap(budget.ID, raw, &clone) {
+ break
+ }
+ }
+ }
+ }()
+ adopted := store.AdoptCalendarAlignmentInMemory(ctx, budget.ID, now)
+ wg.Wait()
+
+ assert.True(t, adopted, "a window opened before the boundary must be adopted")
+
+ live := store.LoadBudget(ctx, budget.ID)
+ require.NotNil(t, live)
+ assert.True(t, live.LastReset.Equal(monthStart), "the window is re-anchored on the boundary")
+ assert.Equal(t, float64(bumps), live.CurrentUsage,
+ "every concurrent bump survived: adoption changed the boundary, not the accounting")
+}
diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go
index 62278f0c94..8e259e3084 100644
--- a/transports/bifrost-http/handlers/governance.go
+++ b/transports/bifrost-http/handlers/governance.go
@@ -51,6 +51,12 @@ type GovernanceManager interface {
// enforces spend, leaving each reset boundary untouched. Enterprise also
// propagates the reset to cluster peers, which is what owner addresses.
ResetBudgetUsageInMemory(ctx context.Context, owner BudgetUsageResetOwner, budgetIDs []string) error
+ // AdoptCalendarAlignmentInMemory re-anchors the given budgets and rate limit
+ // onto the calendar boundaries they have just been told to follow, preserving
+ // their usage. Called only on the transition from rolling to calendar-aligned;
+ // without it that switch clears the usage of every window that opened before
+ // the current boundary.
+ AdoptCalendarAlignmentInMemory(ctx context.Context, owner BudgetUsageResetOwner, budgetIDs []string, rateLimitIDs []string) error
RemoveVirtualKey(ctx context.Context, id string) error
ReloadTeam(ctx context.Context, id string) (*configstoreTables.TableTeam, error)
RemoveTeam(ctx context.Context, id string) error
@@ -588,6 +594,38 @@ func (r *budgetUsageReset) apply(ctx context.Context, store configstore.ConfigSt
return nil
}
+// adoptCalendarAlignment re-anchors an owner's open windows after calendar
+// alignment was switched on, and is a no-op on every other write.
+//
+// The switch is otherwise destructive. The reset sweep treats a window as due
+// whenever its boundary is later than its last reset, and a freshly aligned window
+// is measured against the most recent calendar boundary, so anything that opened
+// before that boundary is due immediately and its usage is cleared on the next
+// tick. For a monthly budget that is most of the month.
+//
+// Adoption moves each window's boundary forward onto the grid instead, which the
+// forward-only guard permits, so the window reads as current and its usage
+// survives until the next real boundary. Every window is judged on its own
+// duration, so a monthly budget can adopt while an hourly counter beside it stays
+// rolling untouched.
+//
+// Runs after the in-memory reload, for the same reason the usage reset does: the
+// reload carries cached values forward and would otherwise overwrite the change.
+func (h *GovernanceHandler) adoptCalendarAlignment(ctx context.Context, switchedOn bool, owner BudgetUsageResetOwner, budgets []configstoreTables.TableBudget, rateLimitID *string) error {
+ if !switchedOn {
+ return nil
+ }
+ budgetIDs := make([]string, 0, len(budgets))
+ for _, budget := range budgets {
+ budgetIDs = append(budgetIDs, budget.ID)
+ }
+ var rateLimitIDs []string
+ if rateLimitID != nil && *rateLimitID != "" {
+ rateLimitIDs = append(rateLimitIDs, *rateLimitID)
+ }
+ return h.governanceManager.AdoptCalendarAlignmentInMemory(ctx, owner, budgetIDs, rateLimitIDs)
+}
+
// reconcileModelConfigBudgets upserts the desired set of budgets owned by a model config
// (via TableBudget.ModelConfigID), preserving usage on matched rows and deleting removed
// ones. It mutates mc.Budgets to the reconciled set. The model config row must already
@@ -1820,6 +1858,9 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
// Preserve / Reset dialog. Carried into budget reconciliation and collected
// there, so the in-memory store can be cleared once the transaction commits.
usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage}
+ // See adoptCalendarAlignment: captured before the update so open windows are
+ // re-anchored rather than reset when alignment is switched on.
+ alignmentSwitchedOn := false
// Parse expires_at when provided: a timestamp must be in the future, "" clears the expiry.
var newExpiresAt *time.Time
if req.ExpiresAt != nil && *req.ExpiresAt != "" {
@@ -1894,6 +1935,7 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
vk.ExpiresAt = newExpiresAt
}
if req.CalendarAligned != nil {
+ alignmentSwitchedOn = !vk.CalendarAligned && *req.CalendarAligned
vk.CalendarAligned = *req.CalendarAligned
}
// VK top-level and per-provider budgets/rate-limits are stored in VK-scoped model
@@ -2231,6 +2273,36 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
return
}
}
+ if alignmentSwitchedOn {
+ // The VK's own windows, plus every model config that inherits the flag from
+ // it: provider governance carries no alignment field of its own, so a VK
+ // switch is the only event that aligns those budgets and their rate limits.
+ // Reload is allowed to return no entity; fall back to the row already held.
+ adoptBudgets, adoptRateLimitID := vk.Budgets, vk.RateLimitID
+ if preloadedVk != nil {
+ adoptBudgets, adoptRateLimitID = preloadedVk.Budgets, preloadedVk.RateLimitID
+ }
+ if err := h.adoptCalendarAlignment(ctx, true,
+ BudgetUsageResetOwner{Kind: BudgetOwnerVirtualKey, ID: vk.ID}, adoptBudgets, adoptRateLimitID); err != nil {
+ logger.Error("failed to adopt calendar alignment for virtual key: %v", err)
+ SendError(ctx, 500, "Virtual key updated but calendar alignment did not take effect")
+ return
+ }
+ modelConfigs, err := h.configStore.GetModelConfigsByScopeAndScopeIDs(ctx, configstoreTables.ModelConfigScopeVirtualKey, []string{vk.ID})
+ if err != nil {
+ logger.Error("failed to load model configs for calendar alignment adoption: %v", err)
+ SendError(ctx, 500, "Virtual key updated but calendar alignment did not take effect")
+ return
+ }
+ for i := range modelConfigs {
+ if err := h.adoptCalendarAlignment(ctx, true,
+ BudgetUsageResetOwner{Kind: BudgetOwnerModelConfig, ID: modelConfigs[i].ID}, modelConfigs[i].Budgets, modelConfigs[i].RateLimitID); err != nil {
+ logger.Error("failed to adopt calendar alignment for model config %s: %v", modelConfigs[i].ID, err)
+ SendError(ctx, 500, "Virtual key updated but calendar alignment did not take effect")
+ return
+ }
+ }
+ }
// Per-user credential reconciliation when the VK's MCP allowlist
// changed. Mirrors the AP-propagation path: enterprise orphans /
@@ -2581,6 +2653,10 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) {
// reconciliation and collected there, so the in-memory store can be cleared
// once the transaction commits.
usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage}
+ // Whether this request switches alignment on. Captured before the update so the
+ // open windows can be adopted onto the calendar grid afterwards instead of
+ // being reset out from under the operator; see adoptCalendarAlignment.
+ alignmentSwitchedOn := false
// Fetching team from database
team, err := h.configStore.GetTeam(ctx, teamID)
if err != nil {
@@ -2611,6 +2687,7 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) {
// - explicit team-level field wins (req.CalendarAligned != nil)
// - else leave existing team.CalendarAligned untouched
if req.CalendarAligned != nil {
+ alignmentSwitchedOn = !team.CalendarAligned && *req.CalendarAligned
team.CalendarAligned = *req.CalendarAligned
}
// Snap-to-calendar-period happens after budget/rate-limit reconciliation
@@ -2784,6 +2861,17 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) {
return
}
}
+ // Reload is allowed to return no entity; fall back to the row already held.
+ adoptBudgets, adoptRateLimitID := team.Budgets, team.RateLimitID
+ if preloadedTeam != nil {
+ adoptBudgets, adoptRateLimitID = preloadedTeam.Budgets, preloadedTeam.RateLimitID
+ }
+ if err := h.adoptCalendarAlignment(ctx, alignmentSwitchedOn,
+ BudgetUsageResetOwner{Kind: BudgetOwnerTeam, ID: team.ID}, adoptBudgets, adoptRateLimitID); err != nil {
+ logger.Error("failed to adopt calendar alignment for Team: %v", err)
+ SendError(ctx, 500, "Team updated but calendar alignment did not take effect")
+ return
+ }
SendJSON(ctx, map[string]interface{}{
"message": "Team updated successfully",
"team": preloadedTeam,
@@ -2998,6 +3086,9 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) {
// reconciliation and collected there, so the in-memory store can be cleared
// once the transaction commits.
usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage}
+ // See adoptCalendarAlignment: captured before the update so open windows are
+ // re-anchored rather than reset when alignment is switched on.
+ alignmentSwitchedOn := false
if req.Budgets != nil && req.Budget != nil {
SendError(ctx, 400, "only one of 'budget' or 'budgets' may be set")
return
@@ -3021,6 +3112,7 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) {
customer.Name = *req.Name
}
if req.CalendarAligned != nil {
+ alignmentSwitchedOn = !customer.CalendarAligned && *req.CalendarAligned
customer.CalendarAligned = *req.CalendarAligned
}
// Handle budget updates: prefer Budgets slice; coerce legacy Budget if needed.
@@ -3133,6 +3225,18 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) {
return
}
}
+ // Read the windows off the reload when it produced one, else off the row this
+ // handler already holds: ReloadCustomer is allowed to return no entity.
+ adoptBudgets, adoptRateLimitID := customer.Budgets, customer.RateLimitID
+ if preloadedCustomer != nil {
+ adoptBudgets, adoptRateLimitID = preloadedCustomer.Budgets, preloadedCustomer.RateLimitID
+ }
+ if err := h.adoptCalendarAlignment(ctx, alignmentSwitchedOn,
+ BudgetUsageResetOwner{Kind: BudgetOwnerCustomer, ID: customer.ID}, adoptBudgets, adoptRateLimitID); err != nil {
+ logger.Error("failed to adopt calendar alignment for Customer: %v", err)
+ SendError(ctx, 500, "Customer updated but calendar alignment did not take effect")
+ return
+ }
SendJSON(ctx, map[string]interface{}{
"message": "Customer updated successfully",
@@ -3888,6 +3992,9 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) {
// reconciliation and collected there, so the in-memory store can be cleared
// once the transaction commits.
usageReset := &budgetUsageReset{requested: req.ResetBudgetUsage != nil && *req.ResetBudgetUsage}
+ // See adoptCalendarAlignment: captured before the update so open windows are
+ // re-anchored rather than reset when alignment is switched on.
+ alignmentSwitchedOn := false
if req.Budget != nil && req.Budgets != nil {
SendError(ctx, 400, "only one of 'budget' or 'budgets' may be set")
return
@@ -3941,6 +4048,7 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) {
// Apply CalendarAligned if provided.
if req.CalendarAligned != nil {
+ alignmentSwitchedOn = !mc.CalendarAligned && *req.CalendarAligned
mc.CalendarAligned = *req.CalendarAligned
}
@@ -4090,6 +4198,14 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) {
return
}
}
+ if !deleted {
+ if err := h.adoptCalendarAlignment(ctx, alignmentSwitchedOn,
+ BudgetUsageResetOwner{Kind: BudgetOwnerModelConfig, ID: mc.ID}, mc.Budgets, mc.RateLimitID); err != nil {
+ logger.Error("failed to adopt calendar alignment for provider governance: %v", err)
+ SendError(ctx, 500, "Provider governance updated but calendar alignment did not take effect")
+ return
+ }
+ }
SendJSON(ctx, map[string]interface{}{
"message": "Provider governance updated successfully",
"provider": resp,
diff --git a/transports/bifrost-http/handlers/governance_test.go b/transports/bifrost-http/handlers/governance_test.go
index eef058ad9a..b9d1b49a03 100644
--- a/transports/bifrost-http/handlers/governance_test.go
+++ b/transports/bifrost-http/handlers/governance_test.go
@@ -3095,6 +3095,17 @@ func (m *mockCustomerStore) DeleteBudget(_ context.Context, _ string, _ ...*gorm
type mockCustomerGovernanceManager struct {
GovernanceManager
+ // adoptedBudgetIDs records what the handler asked to be re-anchored onto the
+ // calendar grid, so a test can tell "alignment was switched on" apart from
+ // "alignment was already on and nothing needed adopting".
+ adoptedBudgetIDs []string
+ adoptCalls int
+}
+
+func (m *mockCustomerGovernanceManager) AdoptCalendarAlignmentInMemory(_ context.Context, _ BudgetUsageResetOwner, budgetIDs []string, _ []string) error {
+ m.adoptCalls++
+ m.adoptedBudgetIDs = append(m.adoptedBudgetIDs, budgetIDs...)
+ return nil
}
func (m *mockCustomerGovernanceManager) ReloadCustomer(_ context.Context, _ string) (*configstoreTables.TableCustomer, error) {
@@ -3210,7 +3221,8 @@ func TestUpdateCustomer_CalendarAligned_DoesNotTouchBudgets(t *testing.T) {
CurrentUsage: 99.0,
}},
}
- h := &GovernanceHandler{configStore: store, governanceManager: &mockCustomerGovernanceManager{}}
+ governanceManager := &mockCustomerGovernanceManager{}
+ h := &GovernanceHandler{configStore: store, governanceManager: governanceManager}
body, _ := json.Marshal(map[string]any{"calendar_aligned": true})
ctx := &fasthttp.RequestCtx{}
@@ -3233,6 +3245,16 @@ func TestUpdateCustomer_CalendarAligned_DoesNotTouchBudgets(t *testing.T) {
t.Errorf("budget %s CurrentUsage changed to %v; enabling alignment must not clear usage", b.ID, b.CurrentUsage)
}
}
+ // The config write leaves the window alone, but on its own that is exactly the
+ // bug: the reset sweep would find the window overdue against the new boundary
+ // and clear the 99.0 above. The handler must hand the budget to the in-memory
+ // adoption, which re-anchors it forward instead.
+ if governanceManager.adoptCalls != 1 {
+ t.Errorf("expected exactly one calendar adoption call, got %d", governanceManager.adoptCalls)
+ }
+ if len(governanceManager.adoptedBudgetIDs) != 1 || governanceManager.adoptedBudgetIDs[0] != budgetID {
+ t.Errorf("expected budget %s to be adopted onto its calendar boundary, got %v", budgetID, governanceManager.adoptedBudgetIDs)
+ }
}
// TestUpdateCustomer_CalendarAligned_NoSnapWhenAlreadyEnabled verifies that if
@@ -3262,7 +3284,8 @@ func TestUpdateCustomer_CalendarAligned_NoSnapWhenAlreadyEnabled(t *testing.T) {
},
},
}
- h := &GovernanceHandler{configStore: store, governanceManager: &mockCustomerGovernanceManager{}}
+ governanceManager := &mockCustomerGovernanceManager{}
+ h := &GovernanceHandler{configStore: store, governanceManager: governanceManager}
body, _ := json.Marshal(map[string]any{"calendar_aligned": true})
ctx := &fasthttp.RequestCtx{}
@@ -3277,6 +3300,12 @@ func TestUpdateCustomer_CalendarAligned_NoSnapWhenAlreadyEnabled(t *testing.T) {
if len(store.updatedBudgets) != 0 {
t.Errorf("expected no UpdateBudget call when calendar_aligned was already true, got %d", len(store.updatedBudgets))
}
+ // Adoption belongs to the switch-over only. A budget already on the calendar
+ // grid must not be re-anchored on every unrelated config write, which would
+ // keep pushing its boundary forward and stop it ever resetting.
+ if governanceManager.adoptCalls != 0 {
+ t.Errorf("alignment was already enabled, so no adoption should be requested; got %d calls", governanceManager.adoptCalls)
+ }
}
// TestApplyVKGovernanceFromModelConfigs_PreservesDirectlyAttachedBudget is a
@@ -3438,6 +3467,80 @@ func TestProviderGovernance_UnknownProviderStill404(t *testing.T) {
}
}
+// providerGovernanceAdoptionManager records which budget IDs the provider
+// governance handler asked to re-anchor onto the calendar grid. It reuses the
+// pricing-override manager for everything else, including ReloadModelConfig,
+// which returns no entity - the handler must not depend on the reload to know
+// which windows to adopt.
+type providerGovernanceAdoptionManager struct {
+ pricingOverrideTestGovernanceManager
+ adoptedBudgetIDs []string
+ adoptCalls int
+}
+
+func (m *providerGovernanceAdoptionManager) AdoptCalendarAlignmentInMemory(_ context.Context, _ BudgetUsageResetOwner, budgetIDs []string, _ []string) error {
+ m.adoptCalls++
+ m.adoptedBudgetIDs = append(m.adoptedBudgetIDs, budgetIDs...)
+ return nil
+}
+
+// TestUpdateProviderGovernance_AdoptsReconciledBudgetsNotStaleOnes pins the
+// contract that makes it safe for the adoption call to read mc.Budgets rather
+// than a reloaded model config: reconcileModelConfigBudgets writes the
+// reconciled set back onto mc, so by the time alignment is adopted, mc.Budgets
+// holds the rows that actually exist.
+//
+// The switch-on request replaces a monthly budget with a quarterly one, so the
+// same call deletes one row and creates another. Adoption must be addressed to
+// the row that now exists. Addressed to the deleted one it is a silent no-op,
+// and the newly created window is never anchored, so the next evaluation clears
+// usage the operator was promised would carry over.
+func TestUpdateProviderGovernance_AdoptsReconciledBudgetsNotStaleOnes(t *testing.T) {
+ SetLogger(&mockLogger{})
+ ctx := context.Background()
+ store := setupPricingOverrideHandlerStore(t)
+ manager := &providerGovernanceAdoptionManager{}
+ handler := &GovernanceHandler{configStore: store, governanceManager: manager}
+
+ const providerName = "openai"
+ require.NoError(t, store.AddProvider(ctx, schemas.ModelProvider(providerName), configstore.ProviderConfig{}))
+
+ // A rolling monthly budget, alignment off.
+ createCtx := newGovernanceProviderNameCtx(providerName, `{"budgets":[{"max_limit":10,"reset_duration":"1M"}],"calendar_aligned":false}`)
+ handler.updateProviderGovernance(createCtx)
+ require.Equal(t, fasthttp.StatusOK, createCtx.Response.StatusCode(),
+ "seed PUT failed; body=%s", createCtx.Response.Body())
+ require.Zero(t, manager.adoptCalls,
+ "alignment was never switched on, so nothing should have been adopted yet")
+
+ pn := providerName
+ before, err := store.GetModelConfig(ctx, configstoreTables.ModelConfigScopeGlobal, nil, configstoreTables.ModelConfigAllModels, &pn)
+ require.NoError(t, err)
+ require.Len(t, before.Budgets, 1)
+ staleBudgetID := before.Budgets[0].ID
+
+ // Switch alignment on and swap the duration in the same request, so
+ // reconciliation cannot match the existing row and must delete it and create
+ // a replacement.
+ switchCtx := newGovernanceProviderNameCtx(providerName, `{"budgets":[{"max_limit":25,"reset_duration":"1Q"}],"calendar_aligned":true}`)
+ handler.updateProviderGovernance(switchCtx)
+ require.Equal(t, fasthttp.StatusOK, switchCtx.Response.StatusCode(),
+ "switch-on PUT failed; body=%s", switchCtx.Response.Body())
+
+ after, err := store.GetModelConfig(ctx, configstoreTables.ModelConfigScopeGlobal, nil, configstoreTables.ModelConfigAllModels, &pn)
+ require.NoError(t, err)
+ require.Len(t, after.Budgets, 1)
+ freshBudgetID := after.Budgets[0].ID
+ require.NotEqual(t, staleBudgetID, freshBudgetID,
+ "this test is only meaningful if reconciliation actually swapped the row")
+
+ require.Equal(t, 1, manager.adoptCalls, "switching alignment on must adopt exactly once")
+ assert.Equal(t, []string{freshBudgetID}, manager.adoptedBudgetIDs,
+ "adoption must name the budget that survived reconciliation")
+ assert.NotContains(t, manager.adoptedBudgetIDs, staleBudgetID,
+ "adoption named the deleted budget, so the surviving window was left un-anchored")
+}
+
// TestProviderGovernance_MalformedEncodingReturns400 locks in the fail-closed
// contract: when the provider name is not valid percent-encoding (e.g. a stray
// "%2"), url.PathUnescape fails and both handlers must respond 400 rather than
diff --git a/transports/bifrost-http/handlers/pricing_override_test.go b/transports/bifrost-http/handlers/pricing_override_test.go
index 4cb571e6e3..7cfa2794ce 100644
--- a/transports/bifrost-http/handlers/pricing_override_test.go
+++ b/transports/bifrost-http/handlers/pricing_override_test.go
@@ -26,6 +26,9 @@ func (pricingOverrideTestGovernanceManager) GetGovernanceData(ctx context.Contex
func (pricingOverrideTestGovernanceManager) ResetBudgetUsageInMemory(context.Context, BudgetUsageResetOwner, []string) error {
return nil
}
+func (pricingOverrideTestGovernanceManager) AdoptCalendarAlignmentInMemory(context.Context, BudgetUsageResetOwner, []string, []string) error {
+ return nil
+}
func (pricingOverrideTestGovernanceManager) ReloadVirtualKey(context.Context, string) (*configstoreTables.TableVirtualKey, error) {
return nil, nil
}
diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go
index 72750d39a0..cc4fa842b2 100644
--- a/transports/bifrost-http/server/server.go
+++ b/transports/bifrost-http/server/server.go
@@ -105,6 +105,11 @@ type ServerCallbacks interface {
// owns them so enterprise can address the cluster broadcast that propagates the
// reset to peers.
ResetBudgetUsageInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string) error
+ // AdoptCalendarAlignmentInMemory re-anchors an owner's budgets and rate limits
+ // onto their calendar boundaries after alignment was switched on, preserving
+ // usage. owner identifies the entity so enterprise can address the cluster
+ // broadcast, exactly as the usage reset above does.
+ AdoptCalendarAlignmentInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string, rateLimitIDs []string) error
RemoveVirtualKey(ctx context.Context, id string) error
// Provider related callbacks
GetModelsForProvider(provider schemas.ModelProvider) []string
@@ -499,6 +504,41 @@ func (s *BifrostHTTPServer) ResetBudgetUsageInMemory(ctx context.Context, owner
return nil
}
+// AdoptCalendarAlignmentInMemory re-anchors an owner's budgets and rate limits
+// onto their calendar boundaries after alignment was switched on.
+//
+// Only meaningful on the false-to-true transition. The in-memory store is what
+// the reset sweep reads, and a window that opened before the current boundary
+// reads as already due the moment the flag flips, so without this the operator
+// loses the usage they were still accumulating. Adoption moves the boundary
+// forward instead, which every persistence path allows, and the next dump tick
+// carries the new boundary to the database the same way an ordinary reset does.
+//
+// Missing entries are not an error: a budget or rate limit can legitimately have
+// been deleted in the same request that turned alignment on.
+func (s *BifrostHTTPServer) AdoptCalendarAlignmentInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string, rateLimitIDs []string) error {
+ if len(budgetIDs) == 0 && len(rateLimitIDs) == 0 {
+ return nil
+ }
+ governancePlugin, err := s.getGovernancePlugin()
+ if err != nil {
+ return err
+ }
+ store := governancePlugin.GetGovernanceStore()
+ now := time.Now()
+ for _, budgetID := range budgetIDs {
+ if !store.AdoptCalendarAlignmentInMemory(ctx, budgetID, now) {
+ logger.Debug("budget %s needed no calendar adoption (absent, unalignable, or already current)", budgetID)
+ }
+ }
+ for _, rateLimitID := range rateLimitIDs {
+ if !store.AdoptRateLimitCalendarAlignmentInMemory(ctx, rateLimitID, now) {
+ logger.Debug("rate limit %s needed no calendar adoption (absent, unalignable, or already current)", rateLimitID)
+ }
+ }
+ return nil
+}
+
// RemoveVirtualKey removes a virtual key from the in-memory store
func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) error {
governancePlugin, err := s.getGovernancePlugin()