diff --git a/coderd/aibridge/budget/period.go b/coderd/aibridge/budget/period.go new file mode 100644 index 0000000000..227e804767 --- /dev/null +++ b/coderd/aibridge/budget/period.go @@ -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) + } +} diff --git a/coderd/aibridge/budget/period_test.go b/coderd/aibridge/budget/period_test.go new file mode 100644 index 0000000000..61d747e347 --- /dev/null +++ b/coderd/aibridge/budget/period_test.go @@ -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") + }) + } +} diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index c17c26587d..9927cb1b1d 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -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)) diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 5a5b3d153e..d0653c87ae 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -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 != "" {