diff --git a/docs/features/governance/budget-and-limits.mdx b/docs/features/governance/budget-and-limits.mdx index be86791efbb..ca242699134 100644 --- a/docs/features/governance/budget-and-limits.mdx +++ b/docs/features/governance/budget-and-limits.mdx @@ -171,9 +171,30 @@ Budgets and rate limits support flexible reset durations: By default, a budget **rolls**: after `reset_duration` elapses since `last_reset`, usage resets. With **`calendar_aligned`: `true`**, the budget resets at the **start of each calendar period in UTC** instead (same instant for every customer of that configuration). -**Supported `reset_duration` suffixes:** only day (`d`), week (`w`), month (`M`), quarter (`Q`), and year (`Y`). Examples: `1d` → midnight UTC each day; `1w` → Monday 00:00 UTC each week; `1M` → first day of each month; `1Q` → first day of the fiscal quarter; `1Y` → January 1 each year. Sub-day durations (for example `1h`, `30m`) **cannot** use calendar alignment; the API rejects invalid combinations. +**Supported `reset_duration` suffixes:** only day (`d`), week (`w`), month (`M`), quarter (`Q`), and year (`Y`). Examples: `1d` → midnight UTC each day; `1w` → Monday 00:00 UTC each week; `1M` → first day of each month; `1Q` → first day of the fiscal quarter; `1Y` → January 1 each year. Sub-day durations (for example `1h`, `30m`) have no calendar boundary to snap to. Setting `calendar_aligned` alongside one is **accepted rather than rejected**, and simply has no effect: that window keeps resetting on its rolling schedule. -Calendar alignment applies to budgets on **customers**, **teams**, **virtual keys**, and **per–provider-config** budgets. You can set it when creating a budget (`calendar_aligned` on create) or toggle it on update (`calendar_aligned` on the budget in `PUT` requests). Turning calendar alignment **on** for an existing budget resets **current usage to zero** and snaps **`last_reset`** to the current period start. +Calendar alignment applies to budgets on **customers**, **teams**, **virtual keys**, and **per-provider-config** budgets. `calendar_aligned` is an **owner-level** field, not a per-budget one: send it at the top level of the customer, team, virtual key, or provider-governance request body, on both `POST` (create) and `PUT` (update). Omitting it on a `PUT` leaves the current setting unchanged. + +The flag applies to **everything that owner resets**, not to budgets alone. That owner's [rate limits](#rate-limiting) follow the same rule, so its token and request counters align too. + +It sets the alignment **mode**, not a shared reset instant. An owner holds one window per budget plus one each for its token and request counters, and every one of them keeps its own `reset_duration` and its own `last_reset`. Each aligns to **its own** boundary: on an aligned owner, a `1M` budget resets on the 1st while a `1d` token limit resets at midnight. They do not reset together. + + +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: + +- **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. + +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. + +Because each window is judged separately, one switch can leave a `1d` limit untouched while clearing a `1M` budget on the same owner. + + +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. + ### Quarterly budgets and fiscal quarters diff --git a/docs/openapi/openapi.json b/docs/openapi/openapi.json index 5e5badc5c10..d81ac482668 100644 --- a/docs/openapi/openapi.json +++ b/docs/openapi/openapi.json @@ -82990,7 +82990,7 @@ "boolean", "null" ], - "description": "Set to true or false to enable or disable calendar-aligned resets. Only valid with reset durations that use day, week, month, quarter, or year suffixes (`d`, `w`, `M`, `Q`, `Y`); sub-day durations (e.g. `1h`, `30m`) are invalid with calendar alignment and the API rejects that combination. When enabling on an existing budget, current usage is reset to zero and last_reset snaps to the current period start. Changing quarter_start_month on a calendar-aligned quarterly budget does not move last_reset directly; the new definition shifts where the current window starts, so the budget converges onto the new fiscal boundary on the next reset tick.\n" + "description": "Set to true or false to enable or disable calendar-aligned resets. Only valid with reset durations that use day, week, month, quarter, or year suffixes (`d`, `w`, `M`, `Q`, `Y`); sub-day durations (e.g. `1h`, `30m`) are invalid with calendar alignment and the API rejects that combination. When enabling on an existing budget, alignment takes effect from the next period: the current window keeps its start and its accumulated usage, and the first aligned reset happens at the next boundary. Changing quarter_start_month on a calendar-aligned quarterly budget does not move last_reset directly; the new definition shifts where the current window starts, so the budget converges onto the new fiscal boundary on the next reset tick.\n" } } }, diff --git a/docs/openapi/schemas/management/governance.yaml b/docs/openapi/schemas/management/governance.yaml index 1cf9418290c..004946f1439 100644 --- a/docs/openapi/schemas/management/governance.yaml +++ b/docs/openapi/schemas/management/governance.yaml @@ -203,8 +203,9 @@ UpdateBudgetRequest: Set to true or false to enable or disable calendar-aligned resets. Only valid with reset durations that use day, week, month, quarter, or year suffixes (`d`, `w`, `M`, `Q`, `Y`); sub-day durations (e.g. `1h`, `30m`) are invalid with calendar alignment and the API rejects that combination. When - enabling on an existing budget, current usage is reset to zero and last_reset snaps to the current - period start. Changing quarter_start_month on a calendar-aligned quarterly budget does not move + enabling on an existing budget, alignment takes effect from the next period: the current window + keeps its start and its accumulated usage, and the first aligned reset happens at the next + boundary. Changing quarter_start_month on a calendar-aligned quarterly budget does not move last_reset directly; the new definition shifts where the current window starts, so the budget converges onto the new fiscal boundary on the next reset tick. diff --git a/plugins/governance/budgetcycle_test.go b/plugins/governance/budgetcycle_test.go index 812800171a0..0a792ee4825 100644 --- a/plugins/governance/budgetcycle_test.go +++ b/plugins/governance/budgetcycle_test.go @@ -841,3 +841,166 @@ func TestDumpBudgetsCannotUndoOperatorReset(t *testing.T) { assert.Equal(t, 7.5, persisted.CurrentUsage, "the dump after a reset must resume persisting usage, or a reset would stop accounting for good") } + +// TestEnablingCalendarAlignmentCanResetAtTheBoundaryAlreadyPassed records what +// enabling alignment actually does today, which is not what this feature's docs +// describe. It is a characterization test, not an endorsement. +// +// budgetResetTarget returns WindowStart(now) whenever that is after LastReset, and +// for an aligned budget WindowStart is the most recent calendar boundary. So a +// budget whose window opened before that boundary is already due the moment +// alignment is switched on, and the next sweep clears its usage. +// +// The documented promise is that alignment applies from the next period and the +// current window keeps its usage. That holds only when LastReset is newer than the +// most recent boundary, which is the case the existing coverage exercises: it +// 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. +func TestEnablingCalendarAlignmentCanResetAtTheBoundaryAlreadyPassed(t *testing.T) { + 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) + + alignedBudget := func(lastReset time.Time) *configstoreTables.TableBudget { + return &configstoreTables.TableBudget{ + ID: "align-boundary-budget", + MaxLimit: 100, + CurrentUsage: 42, + ResetDuration: "1M", + IsCalendarAligned: true, + CreatedAt: lastReset, + LastReset: lastReset, + } + } + + t.Run("a window opened before the boundary is due at once", func(t *testing.T) { + target := store.budgetResetTarget(alignedBudget(time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC)), now) + require.NotNil(t, target, + "current behaviour: the budget is due immediately, so its $42 of usage is cleared on the next reset evaluation - the sweep or the next request, whichever lands first") + assert.True(t, target.Equal(monthStart), + "the reset lands on the boundary that already passed, not on the next one") + }) + + // Being due is evaluated on two independent paths, not one. The subtest above + // asks budgetResetTarget the question the 10s ticker asks; BumpBudgetUsage asks + // the identical question on the request path and resets before recording the + // new cost. So an already-due window is cleared by the very next request even + // on a node whose sweep has not fired yet, which is why the transition is + // documented as the next reset *evaluation* rather than the next sweep. + t.Run("the request path clears an already-due window without any sweep", func(t *testing.T) { + ctx := context.Background() + budget := alignedBudget(time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC)) + budget.ID = "align-boundary-request-path-budget" + openedAt := budget.LastReset + store.budgets.Store(budget.ID, budget) + + // No sweepBudgetAt call anywhere in this subtest: the bump is the only + // thing that touches the budget. + require.NoError(t, store.BumpBudgetUsage(ctx, budget.ID, 2.5)) + + bumped := store.LoadBudget(ctx, budget.ID) + require.NotNil(t, bumped, "expected the budget to remain loaded after the bump") + assert.Equal(t, 2.5, bumped.CurrentUsage, + "the 42 already accumulated was cleared by the request itself rather than carried forward, so the new window holds only this request's cost") + assert.True(t, bumped.LastReset.After(openedAt), + "the request path advanced the window to its boundary, exactly as a sweep would") + }) + + t.Run("a window opened after the boundary is left alone", func(t *testing.T) { + target := store.budgetResetTarget(alignedBudget(time.Date(2026, time.February, 3, 9, 0, 0, 0, time.UTC)), now) + assert.Nil(t, target, + "this is the case the docs describe, and the only one existing coverage reaches") + }) + + // calendar_aligned is an owner-level flag: the same switch drives the owner's + // rate limits, and rateLimitResetTarget applies the identical + // "boundary after lastReset" rule. Pinned here so the documented transition + // cannot claim to cover rate limits while only budgets are actually checked - + // and so the eventual behaviour fix is reminded it owes them the same rule. + t.Run("rate limits carry the same transition", func(t *testing.T) { + duration := "1M" + + before := store.rateLimitResetTarget(&duration, true, + time.Time{}, time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC), now) + require.NotNil(t, before, + "a rate-limit window opened before the boundary is due at once, exactly like a budget") + assert.True(t, before.Equal(monthStart), + "the reset lands on the boundary that already passed") + + after := store.rateLimitResetTarget(&duration, true, + time.Time{}, time.Date(2026, time.February, 3, 9, 0, 0, 0, time.UTC), now) + assert.Nil(t, after, + "a window opened after the boundary is left alone, exactly like a budget") + }) + + // One owner holds several windows on independent cadences: every budget has its + // own ResetDuration, and a rate limit has two more in TokenResetDuration and + // RequestResetDuration, each with its own LastReset. The owner-level flag picks + // the alignment *mode*; it does not give them a shared boundary. Pinned because + // the documented behaviour is easy to state as "everything snaps together", + // which is wrong in both directions below. + t.Run("each window aligns on its own duration", func(t *testing.T) { + lastReset := time.Date(2026, time.January, 10, 9, 0, 0, 0, time.UTC) + monthly, daily := "1M", "1d" + + monthlyTarget := store.rateLimitResetTarget(&monthly, true, time.Time{}, lastReset, now) + dailyTarget := store.rateLimitResetTarget(&daily, true, time.Time{}, lastReset, now) + require.NotNil(t, monthlyTarget) + require.NotNil(t, dailyTarget) + assert.True(t, monthlyTarget.Equal(monthStart), + "a monthly window aligns to the month boundary") + assert.True(t, dailyTarget.Equal(time.Date(2026, time.February, 5, 0, 0, 0, 0, time.UTC)), + "a daily window aligns to midnight, not to the month boundary the budget uses") + assert.False(t, monthlyTarget.Equal(*dailyTarget), + "two windows on one aligned owner do not share a boundary") + }) + + // A sub-day counter has no calendar boundary to snap to, so rateLimitResetTarget + // falls through to the rolling branch and the owner-level flag changes nothing + // for it. Documenting alignment as owner-wide without this exception promises a + // behaviour the code does not have. + t.Run("sub-day durations stay rolling even when aligned", func(t *testing.T) { + hourly := "1h" + anchor := time.Date(2026, time.February, 5, 9, 30, 0, 0, time.UTC) + + aligned := store.rateLimitResetTarget(&hourly, true, anchor, anchor, now) + rolling := store.rateLimitResetTarget(&hourly, false, anchor, anchor, now) + require.NotNil(t, aligned) + require.NotNil(t, rolling) + assert.True(t, aligned.Equal(*rolling), + "calendar_aligned is inert on a sub-day window: it resets on its rolling anchor either way") + assert.False(t, aligned.Equal(time.Date(2026, time.February, 5, 0, 0, 0, 0, time.UTC)), + "and specifically it does not snap to midnight") + }) + + // The combination is accepted, not refused. BeforeSave validates owner count, + // duration format, a positive duration, max_limit, the override fields and + // reset_config - and nothing ties alignment to the duration, at the table layer + // or in the handlers. So "sub-day plus aligned" persists happily and is simply + // ignored at reset time, which is what the docs have to say. + t.Run("a sub-day aligned budget is accepted, not rejected", func(t *testing.T) { + ctx := context.Background() + logger := NewMockLogger() + configStore, err := configstore.NewConfigStore(ctx, &configstore.Config{ + Enabled: true, + Type: configstore.ConfigStoreTypeSQLite, + Config: &configstore.SQLiteConfig{Path: t.TempDir() + "/subdayaligned.db"}, + }, logger) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, configStore.Close(ctx)) }) + + budget := buildBudgetWithUsage("sub-day-aligned-budget", 100, 0, "1h") + budget.IsCalendarAligned = true + require.NoError(t, configStore.CreateBudget(ctx, budget), + "nothing validates alignment against the duration, so this must save") + + stored, err := configStore.GetBudget(ctx, budget.ID) + require.NoError(t, err) + assert.Equal(t, "1h", stored.ResetDuration, + "the sub-day duration is kept as written rather than corrected or refused") + }) +} diff --git a/tests/governance/customerbudget_test.go b/tests/governance/customerbudget_test.go index d2d0a77983e..d20f0c4a05d 100644 --- a/tests/governance/customerbudget_test.go +++ b/tests/governance/customerbudget_test.go @@ -3,6 +3,7 @@ package governance import ( "strconv" "testing" + "time" ) // TestCustomerBudgetExceededWithMultipleVKs tests that customer level budgets are enforced across multiple VKs @@ -335,3 +336,79 @@ func TestCustomerBudgetExceededWithMultipleTeams(t *testing.T) { t.Fatalf("Made %d requests but never hit customer budget limit (consumed $%.6f / $%.2f) - budget not being enforced", requestNum-1, consumedBudget, customerBudget) } + +// TestCustomerBudgetCalendarAlignmentAppliesFromNextPeriod mirrors the team case +// for the customer snap site. +// +// The three sites - team, customer and provider governance - were identical in +// shape and all discarded, but they persist through different store methods, so +// each is pinned separately rather than by analogy. +func TestCustomerBudgetCalendarAlignmentAppliesFromNextPeriod(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + createResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/customers", + Body: CreateCustomerRequest{ + Name: "test-customer-calendar-align-" + generateRandomID(), + Budgets: []BudgetRequest{{ + MaxLimit: 100, + ResetDuration: "1M", + }}, + }, + }) + if createResp.StatusCode != 200 { + t.Fatalf("Failed to create customer: status %d, body %v", createResp.StatusCode, createResp.Body) + } + customerID := ExtractIDFromResponse(t, createResp) + testData.AddCustomer(customerID) + + before := customerBudgetLastReset(t, customerID) + + aligned := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/customers/" + customerID, + Body: UpdateCustomerRequest{CalendarAligned: &aligned}, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to enable calendar alignment: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + after := customerBudgetLastReset(t, customerID) + if !after.Equal(before) { + t.Errorf("enabling calendar alignment moved last_reset from %s to %s; the current window must be left alone", + before.Format(time.RFC3339), after.Format(time.RFC3339)) + } +} + +// customerBudgetLastReset reads the first budget's last_reset off a customer. +func customerBudgetLastReset(t *testing.T, customerID string) time.Time { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/customers/" + customerID}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read customer %s: status %d, body %v", customerID, resp.StatusCode, resp.Body) + } + customer, ok := resp.Body["customer"].(map[string]interface{}) + if !ok { + customer = resp.Body + } + budgets, ok := customer["budgets"].([]interface{}) + if !ok || len(budgets) == 0 { + t.Fatalf("customer %s has no budgets: %v", customerID, customer) + } + budget, ok := budgets[0].(map[string]interface{}) + if !ok { + t.Fatalf("customer %s first budget is not an object: %v", customerID, budgets[0]) + } + raw, ok := budget["last_reset"].(string) + if !ok { + t.Fatalf("customer %s budget has no last_reset string: %v", customerID, budget) + } + parsed, err := time.Parse(time.RFC3339, raw) + if err != nil { + t.Fatalf("could not parse last_reset %q: %v", raw, err) + } + return parsed.UTC() +} diff --git a/tests/governance/teambudget_test.go b/tests/governance/teambudget_test.go index 01a1a5e1df7..001dfaeab51 100644 --- a/tests/governance/teambudget_test.go +++ b/tests/governance/teambudget_test.go @@ -3,6 +3,7 @@ package governance import ( "strconv" "testing" + "time" ) // TestTeamBudgetExceededWithMultipleVKs tests that team level budgets are enforced across multiple VKs @@ -159,3 +160,173 @@ func TestTeamBudgetExceededWithMultipleVKs(t *testing.T) { t.Fatalf("Made %d requests but never hit team budget limit (consumed $%.6f / $%.2f) - budget not being enforced", requestNum-1, consumedBudget, teamBudget) } + +// TestTeamBudgetCalendarAlignmentAppliesFromNextPeriod pins what enabling +// calendar alignment on an existing budget actually does. +// +// It leaves the current window alone. The docs and the handler used to promise +// that last_reset snapped back to the period start and usage was cleared, but +// neither ever happened: both values were written through UpdateBudget, which +// copies them back from the stored row on every config write so a config.json +// sync cannot replay stale values over live accounting. +// +// Deleting that dead code rather than making it work is deliberate. The intended +// snap moves the boundary backwards, and every persistence path guards last_reset +// forward-only so that cluster nodes agree on which window is current. Alignment +// still takes hold - from the next period boundary. +// +// Checked on a monthly budget with no usage, so it costs no inference: a budget +// created now is anchored at the creation instant, while the period start is +// midnight on the 1st. +func TestTeamBudgetCalendarAlignmentAppliesFromNextPeriod(t *testing.T) { + // Not parallel: this issues a PUT that contends with the other governance + // tests on the default SQLite store. + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + now := time.Now().UTC() + monthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + + createResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/teams", + Body: CreateTeamRequest{ + Name: "test-team-calendar-align-" + generateRandomID(), + Budgets: []BudgetRequest{{ + MaxLimit: 100, + ResetDuration: "1M", + }}, + }, + }) + if createResp.StatusCode != 200 { + t.Fatalf("Failed to create team: status %d, body %v", createResp.StatusCode, createResp.Body) + } + teamID := ExtractIDFromResponse(t, createResp) + testData.AddTeam(teamID) + + // A non-aligned budget anchors on its creation instant, not the 1st, so the + // two values below are distinguishable. + before := teamBudgetLastReset(t, teamID) + if before.Equal(monthStart) { + t.Fatalf("precondition failed: a non-aligned budget already sits on the period start %s", monthStart.Format(time.RFC3339)) + } + + aligned := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/teams/" + teamID, + Body: UpdateTeamRequest{CalendarAligned: &aligned}, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to enable calendar alignment: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + after := teamBudgetLastReset(t, teamID) + if !after.Equal(before) { + t.Errorf("enabling calendar alignment moved last_reset from %s to %s; the current window must be left alone", + before.Format(time.RFC3339), after.Format(time.RFC3339)) + } +} + +// teamBudgetLastReset reads the first budget's last_reset off a team. +func teamBudgetLastReset(t *testing.T, teamID string) time.Time { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/teams/" + teamID}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read team %s: status %d, body %v", teamID, resp.StatusCode, resp.Body) + } + team, ok := resp.Body["team"].(map[string]interface{}) + if !ok { + team = resp.Body + } + budgets, ok := team["budgets"].([]interface{}) + if !ok || len(budgets) == 0 { + t.Fatalf("team %s has no budgets: %v", teamID, team) + } + budget, ok := budgets[0].(map[string]interface{}) + if !ok { + t.Fatalf("team %s first budget is not an object: %v", teamID, budgets[0]) + } + raw, ok := budget["last_reset"].(string) + if !ok { + t.Fatalf("team %s budget has no last_reset string: %v", teamID, budget) + } + parsed, err := time.Parse(time.RFC3339, raw) + if err != nil { + t.Fatalf("could not parse last_reset %q: %v", raw, err) + } + return parsed.UTC() +} + +// TestTeamRateLimitCalendarAlignmentAppliesFromNextPeriod is the rate-limit half +// of the same contract. +// +// Each calendar-align block snapped its owner's rate limit as well as its +// budgets, through UpdateRateLimit, which carries both usage counters and both +// reset stamps forward exactly as UpdateBudget does. That half was discarded too, +// and is worth pinning separately: budgets and rate limits are persisted by +// different store methods, so one could regress without the other. +func TestTeamRateLimitCalendarAlignmentAppliesFromNextPeriod(t *testing.T) { + testData := NewGlobalTestData() + defer testData.Cleanup(t) + + monthly := "1M" + tokenLimit := int64(100000) + createResp := MakeRequest(t, APIRequest{ + Method: "POST", + Path: "/api/governance/teams", + Body: CreateTeamRequest{ + Name: "test-team-rl-calendar-align-" + generateRandomID(), + RateLimit: &CreateRateLimitRequest{TokenMaxLimit: &tokenLimit, TokenResetDuration: &monthly}, + }, + }) + if createResp.StatusCode != 200 { + t.Fatalf("Failed to create team: status %d, body %v", createResp.StatusCode, createResp.Body) + } + teamID := ExtractIDFromResponse(t, createResp) + testData.AddTeam(teamID) + + before := teamRateLimitTokenLastReset(t, teamID) + + aligned := true + updateResp := MakeRequest(t, APIRequest{ + Method: "PUT", + Path: "/api/governance/teams/" + teamID, + Body: UpdateTeamRequest{CalendarAligned: &aligned}, + }) + if updateResp.StatusCode != 200 { + t.Fatalf("Failed to enable calendar alignment: status %d, body %v", updateResp.StatusCode, updateResp.Body) + } + + after := teamRateLimitTokenLastReset(t, teamID) + if !after.Equal(before) { + t.Errorf("enabling calendar alignment moved token_last_reset from %s to %s; the current window must be left alone", + before.Format(time.RFC3339), after.Format(time.RFC3339)) + } +} + +// teamRateLimitTokenLastReset reads a team's rate-limit token reset stamp. +func teamRateLimitTokenLastReset(t *testing.T, teamID string) time.Time { + t.Helper() + resp := MakeRequest(t, APIRequest{Method: "GET", Path: "/api/governance/teams/" + teamID}) + if resp.StatusCode != 200 { + t.Fatalf("Failed to read team %s: status %d, body %v", teamID, resp.StatusCode, resp.Body) + } + team, ok := resp.Body["team"].(map[string]interface{}) + if !ok { + team = resp.Body + } + rateLimit, ok := team["rate_limit"].(map[string]interface{}) + if !ok { + t.Fatalf("team %s has no rate limit: %v", teamID, team) + } + raw, ok := rateLimit["token_last_reset"].(string) + if !ok { + t.Fatalf("team %s rate limit has no token_last_reset: %v", teamID, rateLimit) + } + parsed, err := time.Parse(time.RFC3339, raw) + if err != nil { + t.Fatalf("could not parse token_last_reset %q: %v", raw, err) + } + return parsed.UTC() +} diff --git a/tests/governance/test_utils.go b/tests/governance/test_utils.go index fb64d3a78df..cb37197be73 100644 --- a/tests/governance/test_utils.go +++ b/tests/governance/test_utils.go @@ -250,16 +250,19 @@ type BudgetResetConfigRequest struct { // CreateTeamRequest represents a request to create a team type CreateTeamRequest struct { - Name string `json:"name"` - CustomerID *string `json:"customer_id,omitempty"` - Budgets []BudgetRequest `json:"budgets,omitempty"` - CalendarAligned bool `json:"calendar_aligned,omitempty"` + Name string `json:"name"` + CustomerID *string `json:"customer_id,omitempty"` + Budgets []BudgetRequest `json:"budgets,omitempty"` + RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` + CalendarAligned bool `json:"calendar_aligned,omitempty"` } // CreateCustomerRequest represents a request to create a customer type CreateCustomerRequest struct { - Name string `json:"name"` - Budgets []BudgetRequest `json:"budgets,omitempty"` + Name string `json:"name"` + Budgets []BudgetRequest `json:"budgets,omitempty"` + RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"` + CalendarAligned bool `json:"calendar_aligned,omitempty"` } // UpdateBudgetRequest represents a request to update a budget @@ -304,8 +307,9 @@ type UpdateTeamRequest struct { // UpdateCustomerRequest represents a request to update a customer type UpdateCustomerRequest struct { - Name *string `json:"name,omitempty"` - Budgets []BudgetRequest `json:"budgets,omitempty"` + Name *string `json:"name,omitempty"` + Budgets []BudgetRequest `json:"budgets,omitempty"` + CalendarAligned *bool `json:"calendar_aligned,omitempty"` } // ChatCompletionRequest represents an OpenAI-compatible chat completion request diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 96158c613d4..5c5c07ca8b1 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -2570,11 +2570,9 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { // Resolve team-level calendar alignment for this update: // - explicit team-level field wins (req.CalendarAligned != nil) // - else leave existing team.CalendarAligned untouched - wasCalendarAligned := team.CalendarAligned if req.CalendarAligned != nil { team.CalendarAligned = *req.CalendarAligned } - calendarAlignmentJustEnabled := !wasCalendarAligned && team.CalendarAligned // Snap-to-calendar-period happens after budget/rate-limit reconciliation // below, so combined `calendar_aligned + budgets/rate_limit` updates see // the final persisted state. @@ -2691,44 +2689,7 @@ func (h *GovernanceHandler) updateTeam(ctx *fasthttp.RequestCtx) { team.RateLimit = &rateLimit } } - // Snap budgets and rate limit to the current calendar period when calendar - // alignment transitions false -> true in this request. Runs after budget/ - // rate-limit reconciliation so both the standalone-toggle and the combined - // (toggle + budgets/rate_limit in the same request) cases are covered, and - // only fires once per transition. - if calendarAlignmentJustEnabled { - now := time.Now() - for i := range team.Budgets { - b := &team.Budgets[i] - if !configstoreTables.IsCalendarAlignableDuration(b.ResetDuration) { - continue - } - b.LastReset = configstoreTables.GetCalendarPeriodStart(b.ResetDuration, now, b.QuarterStartMonth()) - b.CurrentUsage = 0 - if err := h.configStore.UpdateBudget(ctx, b, tx); err != nil { - return fmt.Errorf("failed to snap team budget %s on calendar-align enable: %w", b.ID, err) - } - } - if team.RateLimit != nil { - rl := team.RateLimit - snapped := false - if rl.TokenResetDuration != nil && configstoreTables.IsCalendarAlignableDuration(*rl.TokenResetDuration) { - rl.TokenLastReset = configstoreTables.GetCalendarPeriodStart(*rl.TokenResetDuration, now, configstoreTables.QuarterStartNotApplicable) - rl.TokenCurrentUsage = 0 - snapped = true - } - if rl.RequestResetDuration != nil && configstoreTables.IsCalendarAlignableDuration(*rl.RequestResetDuration) { - rl.RequestLastReset = configstoreTables.GetCalendarPeriodStart(*rl.RequestResetDuration, now, configstoreTables.QuarterStartNotApplicable) - rl.RequestCurrentUsage = 0 - snapped = true - } - if snapped { - if err := h.configStore.UpdateRateLimit(ctx, rl, tx); err != nil { - return fmt.Errorf("failed to snap team rate limit on calendar-align enable: %w", err) - } - } - } - } + if err := h.configStore.UpdateTeam(ctx, team, tx); err != nil { return err } @@ -2991,11 +2952,9 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { if req.Name != nil { customer.Name = *req.Name } - wasCalendarAligned := customer.CalendarAligned if req.CalendarAligned != nil { customer.CalendarAligned = *req.CalendarAligned } - calendarAlignmentJustEnabled := !wasCalendarAligned && customer.CalendarAligned // Handle budget updates: prefer Budgets slice; coerce legacy Budget if needed. effectiveBudgets := req.Budgets if effectiveBudgets == nil && req.Budget != nil { @@ -3065,42 +3024,7 @@ func (h *GovernanceHandler) updateCustomer(ctx *fasthttp.RequestCtx) { customer.RateLimit = &rateLimit } } - // Snap budgets and rate limit to the current calendar period when calendar - // alignment transitions false → true. Runs after reconciliation so combined - // "toggle + budgets" requests see the final reconciled state. - if calendarAlignmentJustEnabled { - now := time.Now() - for i := range customer.Budgets { - b := &customer.Budgets[i] - if !configstoreTables.IsCalendarAlignableDuration(b.ResetDuration) { - continue - } - b.LastReset = configstoreTables.GetCalendarPeriodStart(b.ResetDuration, now, b.QuarterStartMonth()) - b.CurrentUsage = 0 - if err := h.configStore.UpdateBudget(ctx, b, tx); err != nil { - return fmt.Errorf("failed to snap customer budget %s on calendar-align enable: %w", b.ID, err) - } - } - if customer.RateLimit != nil { - rl := customer.RateLimit - snapped := false - if rl.TokenResetDuration != nil && configstoreTables.IsCalendarAlignableDuration(*rl.TokenResetDuration) { - rl.TokenLastReset = configstoreTables.GetCalendarPeriodStart(*rl.TokenResetDuration, now, configstoreTables.QuarterStartNotApplicable) - rl.TokenCurrentUsage = 0 - snapped = true - } - if rl.RequestResetDuration != nil && configstoreTables.IsCalendarAlignableDuration(*rl.RequestResetDuration) { - rl.RequestLastReset = configstoreTables.GetCalendarPeriodStart(*rl.RequestResetDuration, now, configstoreTables.QuarterStartNotApplicable) - rl.RequestCurrentUsage = 0 - snapped = true - } - if snapped { - if err := h.configStore.UpdateRateLimit(ctx, rl, tx); err != nil { - return fmt.Errorf("failed to snap customer rate limit on calendar-align enable: %w", err) - } - } - } - } + if err := h.configStore.UpdateCustomer(ctx, customer, tx); err != nil { return err } @@ -3911,11 +3835,9 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) { var rateLimitIDToDelete string // Apply CalendarAligned if provided. - wasCalendarAligned := mc.CalendarAligned if req.CalendarAligned != nil { mc.CalendarAligned = *req.CalendarAligned } - calendarAlignmentJustEnabled := !wasCalendarAligned && mc.CalendarAligned // Rate limit lifecycle (mc references it via RateLimitID, so resolve it before // persisting the model config below). @@ -4015,43 +3937,6 @@ func (h *GovernanceHandler) updateProviderGovernance(ctx *fasthttp.RequestCtx) { } } - // Snap budgets and rate limit to the current calendar period when calendar - // alignment transitions false → true. Runs after reconciliation so combined - // "toggle + budgets" requests see the final reconciled state. - if !deleted && calendarAlignmentJustEnabled { - now := time.Now() - for i := range mc.Budgets { - b := &mc.Budgets[i] - if !configstoreTables.IsCalendarAlignableDuration(b.ResetDuration) { - continue - } - b.LastReset = configstoreTables.GetCalendarPeriodStart(b.ResetDuration, now, b.QuarterStartMonth()) - b.CurrentUsage = 0 - if err := h.configStore.UpdateBudget(ctx, b, tx); err != nil { - return fmt.Errorf("failed to snap provider budget %s on calendar-align enable: %w", b.ID, err) - } - } - if mc.RateLimit != nil { - rl := mc.RateLimit - snapped := false - if rl.TokenResetDuration != nil && configstoreTables.IsCalendarAlignableDuration(*rl.TokenResetDuration) { - rl.TokenLastReset = configstoreTables.GetCalendarPeriodStart(*rl.TokenResetDuration, now, configstoreTables.QuarterStartNotApplicable) - rl.TokenCurrentUsage = 0 - snapped = true - } - if rl.RequestResetDuration != nil && configstoreTables.IsCalendarAlignableDuration(*rl.RequestResetDuration) { - rl.RequestLastReset = configstoreTables.GetCalendarPeriodStart(*rl.RequestResetDuration, now, configstoreTables.QuarterStartNotApplicable) - rl.RequestCurrentUsage = 0 - snapped = true - } - if snapped { - if err := h.configStore.UpdateRateLimit(ctx, rl, tx); err != nil { - return fmt.Errorf("failed to snap provider rate limit on calendar-align enable: %w", err) - } - } - } - } - // Delete orphaned rate-limit row if it was unlinked. if rateLimitIDToDelete != "" { if err := tx.Delete(&configstoreTables.TableRateLimit{}, "id = ?", rateLimitIDToDelete).Error; err != nil { diff --git a/transports/bifrost-http/handlers/governance_test.go b/transports/bifrost-http/handlers/governance_test.go index e2efe59a492..eef058ad9a1 100644 --- a/transports/bifrost-http/handlers/governance_test.go +++ b/transports/bifrost-http/handlers/governance_test.go @@ -3062,7 +3062,23 @@ func (m *mockCustomerStore) CreateBudget(_ context.Context, budget *configstoreT m.createdBudgets = append(m.createdBudgets, budget) return nil } + +// UpdateBudget mirrors RDBConfigStore.UpdateBudget's contract: usage accounting is +// runtime-owned and is carried forward from the stored row, never authored by a +// configuration write. Without this the mock promises something the real store +// does not honour, and a handler test can assert a config write changed usage +// while production silently discards it - which is exactly how the dead +// calendar-alignment snap survived for so long. func (m *mockCustomerStore) UpdateBudget(_ context.Context, budget *configstoreTables.TableBudget, _ ...*gorm.DB) error { + for _, customer := range m.customers { + for i := range customer.Budgets { + if customer.Budgets[i].ID != budget.ID { + continue + } + budget.CurrentUsage = customer.Budgets[i].CurrentUsage + budget.LastReset = customer.Budgets[i].LastReset + } + } m.updatedBudgets = append(m.updatedBudgets, budget) return nil } @@ -3163,36 +3179,36 @@ func TestCreateCustomer_CalendarAligned_False(t *testing.T) { } } -// TestUpdateCustomer_CalendarAligned_SnapsExistingBudget verifies that toggling -// calendar_aligned from false to true snaps the existing budget's LastReset to the -// start of the current calendar period and resets CurrentUsage. -func TestUpdateCustomer_CalendarAligned_SnapsExistingBudget(t *testing.T) { +// TestUpdateCustomer_CalendarAligned_DoesNotTouchBudgets verifies that enabling +// calendar alignment leaves existing budgets alone. +// +// This replaces a test that asserted the opposite - that the toggle snapped +// LastReset to the period start and zeroed CurrentUsage - and passed for years +// while the behaviour never happened. It asserted on what the handler passed to +// UpdateBudget, against a mock that simply recorded the argument. The real store +// copies CurrentUsage and LastReset back from the stored row on every config +// write (rdb.go:4758-4769), so both values were discarded one layer below the +// mock. The mock below now carries them forward the same way, which is what +// keeps this class of test honest. +// +// Alignment still takes effect: the budget aligns from its next period boundary. +func TestUpdateCustomer_CalendarAligned_DoesNotTouchBudgets(t *testing.T) { SetLogger(&mockLogger{}) store := newMockCustomerStore() budgetID := "bud-snap" - budgetID2 := "bud-snap-2" - oldLastReset := time.Now().AddDate(0, -1, 0) // 1 month ago + oldLastReset := time.Now().AddDate(0, -1, 0) store.customers["cust-snap"] = &configstoreTables.TableCustomer{ ID: "cust-snap", Name: "Initech", CalendarAligned: false, - Budgets: []configstoreTables.TableBudget{ - { - ID: budgetID, - MaxLimit: 200.0, - ResetDuration: "1M", - LastReset: oldLastReset, - CurrentUsage: 99.0, - }, - { - ID: budgetID2, - MaxLimit: 500.0, - ResetDuration: "1Y", - LastReset: oldLastReset, - CurrentUsage: 150.0, - }, - }, + Budgets: []configstoreTables.TableBudget{{ + ID: budgetID, + MaxLimit: 200.0, + ResetDuration: "1M", + LastReset: oldLastReset, + CurrentUsage: 99.0, + }}, } h := &GovernanceHandler{configStore: store, governanceManager: &mockCustomerGovernanceManager{}} @@ -3201,31 +3217,21 @@ func TestUpdateCustomer_CalendarAligned_SnapsExistingBudget(t *testing.T) { ctx.Request.SetBody(body) ctx.SetUserValue("customer_id", "cust-snap") - snapBefore := time.Now() h.updateCustomer(ctx) if ctx.Response.StatusCode() != 200 { t.Fatalf("expected 200, got %d: %s", ctx.Response.StatusCode(), ctx.Response.Body()) } - // UpdateBudget must have been called once per budget (both snap). - if len(store.updatedBudgets) != 2 { - t.Fatalf("expected 2 UpdateBudget calls for snap, got %d", len(store.updatedBudgets)) + if got := store.customers["cust-snap"].CalendarAligned; !got { + t.Errorf("calendar_aligned should be enabled on the customer, got %v", got) } - snappedIDs := make(map[string]bool, 2) - for _, snapped := range store.updatedBudgets { - snappedIDs[snapped.ID] = true - if snapped.LastReset.Equal(oldLastReset) { - t.Errorf("budget %s LastReset was not snapped: still equals old value", snapped.ID) + for _, b := range store.customers["cust-snap"].Budgets { + if !b.LastReset.Equal(oldLastReset) { + t.Errorf("budget %s LastReset moved to %v; enabling alignment must not re-anchor the window", b.ID, b.LastReset) } - if snapped.LastReset.After(snapBefore) { - t.Errorf("budget %s snapped LastReset %v should be at the period start, not time.Now()", snapped.ID, snapped.LastReset) + if b.CurrentUsage != 99.0 { + t.Errorf("budget %s CurrentUsage changed to %v; enabling alignment must not clear usage", b.ID, b.CurrentUsage) } - if snapped.CurrentUsage != 0 { - t.Errorf("budget %s expected CurrentUsage reset to 0, got %v", snapped.ID, snapped.CurrentUsage) - } - } - if !snappedIDs[budgetID] || !snappedIDs[budgetID2] { - t.Errorf("expected both %q and %q to be snapped, got IDs: %v", budgetID, budgetID2, snappedIDs) } }