mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
refactor: use AI budget period from deployment config (#27117)
## Description
Read the AI budget period from the deployment config on both the RPC server and the `/users/{user}/ai/spend` endpoint, instead of hardcoding `month`. Drops the `period_start` RPC parameter that was incorrectly introduced in #26915: the period should have been derived from the deployment config from the start.
## Changes
- Add `codersdk.NewAIBudgetPeriodFromString`, mirroring `NewAIBudgetPolicyFromString`.
- `aibridgedserver.Server` takes a `quartz.Clock`, reads `BudgetPeriod` from the deployment config at construction, and computes the period window inside `IsBudgetExceeded`.
- Remove `period_start` from `IsBudgetExceededRequest` and stop sending it from the daemon.
- The `userAISpendStatus` endpoint reads the period from `AIBridgeConfig.BudgetPeriod` instead of hardcoding month.
> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
@@ -7,16 +7,12 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"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{}
|
||||
@@ -149,15 +145,8 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
logger = logger.With(slog.F("user_id", id))
|
||||
|
||||
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(periodWindow.Start),
|
||||
UserId: id.String(),
|
||||
})
|
||||
if err != nil {
|
||||
logger.Warn(ctx, "user AI budget check failed", slog.Error(err))
|
||||
|
||||
Reference in New Issue
Block a user