opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(coderd): add budget.CurrentPeriod for [start, end) windows (#26972)

Browse Source 
## Description

Extracts the AI budget period computation into a shared
`budget.CurrentPeriod` helper. Pure refactor with no wire-value changes.

## Changes

- Add `budget.CurrentPeriod(now, period)` returning a
`PeriodWindow{Start, End}` in UTC. Unknown periods return an error,
matching the pattern used by `ResolveUserAIBudget` for unknown policies.
- Update the callers and respective tests to use `CurrentPeriod`.

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by
@ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-06 10:37:27 +01:00
committed by GitHub
parent f444c6f585
commit 64ca5ac9f8
4 changed files with 130 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/aibridge/budget/period.go
+32
View File
@@ -0,0 +1,32 @@
package budget
import (
"time"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/codersdk"
)
// PeriodWindow is the [Start, End) time window covered by an AI budget
// period. Bounds are in UTC.
type PeriodWindow struct {
// Start is the inclusive first instant of the window.
Start time.Time
// End is the exclusive first instant of the next window.
End time.Time
}
// CurrentPeriod returns the PeriodWindow containing now (normalized to UTC)
// for the given AI budget period. An unknown budget period returns an error.
func CurrentPeriod(now time.Time, period codersdk.AIBudgetPeriod) (PeriodWindow, error) {
nowUTC := now.UTC()
switch period {
case codersdk.AIBudgetPeriodMonth:
start := dbtime.StartOfMonth(nowUTC)
return PeriodWindow{Start: start, End: start.AddDate(0, 1, 0)}, nil
default:
return PeriodWindow{}, xerrors.Errorf("unsupported AI budget period: %q", period)
}
}
coderd/aibridge/budget/period_test.go
+85
View File
@@ -0,0 +1,85 @@
package budget_test
import (
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/coder/coder/v2/coderd/aibridge/budget"
"github.com/coder/coder/v2/codersdk"
)
func TestCurrentPeriod(t *testing.T) {
t.Parallel()
nonUTC := time.FixedZone("UTC-4", -4*60*60)
tests := []struct {
name string
now time.Time
period codersdk.AIBudgetPeriod
wantStart time.Time
wantEnd time.Time
wantErr string
}{
{
name: "MidMonthUTC",
now: time.Date(2026, time.March, 15, 12, 30, 45, 0, time.UTC),
period: codersdk.AIBudgetPeriodMonth,
wantStart: time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2026, time.April, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "FirstInstantOfMonthUTC",
now: time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC),
period: codersdk.AIBudgetPeriodMonth,
wantStart: time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2026, time.April, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "LastInstantOfMonthUTC",
now: time.Date(2026, time.March, 31, 23, 59, 59, 999_999_999, time.UTC),
period: codersdk.AIBudgetPeriodMonth,
wantStart: time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2026, time.April, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "DecemberRollsToJanuary",
now: time.Date(2026, time.December, 15, 12, 0, 0, 0, time.UTC),
period: codersdk.AIBudgetPeriodMonth,
wantStart: time.Date(2026, time.December, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2027, time.January, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "NonUTCNormalizedAcrossMonth",
// Non-UTC input must be normalized before computing the window:
// 2026-03-31 23:00 at UTC-4 is 2026-04-01 03:00 UTC.
now: time.Date(2026, time.March, 31, 23, 0, 0, 0, nonUTC),
period: codersdk.AIBudgetPeriodMonth,
wantStart: time.Date(2026, time.April, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2026, time.May, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "UnsupportedPeriod",
now: time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC),
period: codersdk.AIBudgetPeriod("unknown"),
wantErr: "unsupported AI budget period",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
got, err := budget.CurrentPeriod(tt.now, tt.period)
if tt.wantErr != "" {
require.ErrorContains(t, err, tt.wantErr)
return
}
require.NoError(t, err)
require.Equal(t, tt.wantStart, got.Start, "start")
require.Equal(t, tt.wantEnd, got.End, "end")
})
}
}
coderd/aibridged/http.go
+9 -2
View File
@@ -13,8 +13,10 @@ import (
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/recorder"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/coderd/aibridge/budget"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/codersdk"
)
var _ http.Handler = &Server{}
@@ -147,10 +149,15 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
}
logger = logger.With(slog.F("user_id", id))
periodStart := dbtime.StartOfMonth(dbtime.Now().UTC())
periodWindow, err := budget.CurrentPeriod(dbtime.Now(), codersdk.AIBudgetPeriodMonth)
if err != nil {
logger.Warn(ctx, "compute AI budget period", slog.Error(err))
http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError)
return
}
budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
UserId: id.String(),
PeriodStart: timestamppb.New(periodStart),
PeriodStart: timestamppb.New(periodWindow.Start),
})
if err != nil {
logger.Warn(ctx, "user AI budget check failed", slog.Error(err))
coderd/aibridgedserver/aibridgedserver_test.go
+4 -1
View File
@@ -27,6 +27,7 @@ import (
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogjson"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/coderd/aibridge/budget"
"github.com/coder/coder/v2/coderd/aibridged"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/aibridgedserver"
@@ -567,7 +568,9 @@ func TestIsBudgetExceeded(t *testing.T) {
req := &proto.IsBudgetExceededRequest{UserId: userIDStr}
if !tc.omitPeriodStart {
req.PeriodStart = timestamppb.New(dbtime.StartOfMonth(dbtime.Now().UTC()))
window, err := budget.CurrentPeriod(dbtime.Now(), codersdk.AIBudgetPeriodMonth)
require.NoError(t, err)
req.PeriodStart = timestamppb.New(window.Start)
}
resp, err := srv.IsBudgetExceeded(t.Context(), req)
if tc.wantErrContains != "" {