Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 6580cdc

Browse filesBrowse files
ssncferreirassncferreira
authored
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
1 parent 199b593 commit 6580cdc
Copy full SHA for 6580cdc

10 files changed

+236-248Lines changed: 236 additions & 248 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎cli/aibridged_internal_test.go‎

Copy file name to clipboardExpand all lines: cli/aibridged_internal_test.go
+2-1Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2525
"github.com/coder/coder/v2/codersdk"
2626
"github.com/coder/coder/v2/testutil"
27+
"github.com/coder/quartz"
2728
"github.com/coder/serpent"
2829
)
2930

@@ -53,7 +54,7 @@ func buildFromEnv(t *testing.T, cfg codersdk.AIBridgeConfig) ([]aibridge.Provide
5354
// (providers, outcomes) the embedded reloader would observe.
5455
func buildFromDB(ctx context.Context, t *testing.T, db database.Store, cfg codersdk.AIBridgeConfig, logger slog.Logger) ([]aibridge.Provider, []aibridged.ProviderOutcome, error) {
5556
t.Helper()
56-
srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{})
57+
srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{}, quartz.NewReal())
5758
if err != nil {
5859
return nil, nil, err
5960
}
Collapse file

‎coderd/aibridged.go‎

Copy file name to clipboardExpand all lines: coderd/aibridged.go
+1-1Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai
6666

6767
mux := drpcmux.New()
6868
srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Pubsub, api.Logger.Named("aibridgedserver"),
69-
api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker)
69+
api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker, api.Clock)
7070
if err != nil {
7171
return nil, err
7272
}
Collapse file

‎coderd/aibridged/http.go‎

Copy file name to clipboardExpand all lines: coderd/aibridged/http.go
+1-12Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,12 @@ import (
77

88
"github.com/google/uuid"
99
"golang.org/x/xerrors"
10-
"google.golang.org/protobuf/types/known/timestamppb"
1110

1211
"cdr.dev/slog/v3"
1312
"github.com/coder/coder/v2/aibridge"
1413
"github.com/coder/coder/v2/aibridge/recorder"
1514
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
16-
"github.com/coder/coder/v2/coderd/aibridge/budget"
1715
"github.com/coder/coder/v2/coderd/aibridged/proto"
18-
"github.com/coder/coder/v2/coderd/database/dbtime"
19-
"github.com/coder/coder/v2/codersdk"
2016
)
2117

2218
var _ http.Handler = &Server{}
@@ -149,15 +145,8 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
149145
}
150146
logger = logger.With(slog.F("user_id", id))
151147

152-
periodWindow, err := budget.CurrentPeriod(dbtime.Now(), codersdk.AIBudgetPeriodMonth)
153-
if err != nil {
154-
logger.Warn(ctx, "compute AI budget period", slog.Error(err))
155-
http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError)
156-
return
157-
}
158148
budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
159-
UserId: id.String(),
160-
PeriodStart: timestamppb.New(periodWindow.Start),
149+
UserId: id.String(),
161150
})
162151
if err != nil {
163152
logger.Warn(ctx, "user AI budget check failed", slog.Error(err))
Collapse file

‎coderd/aibridged/proto/aibridged.pb.go‎

Copy file name to clipboardExpand all lines: coderd/aibridged/proto/aibridged.pb.go
+156-170Lines changed: 156 additions & 170 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Collapse file

‎coderd/aibridged/proto/aibridged.proto‎

Copy file name to clipboardExpand all lines: coderd/aibridged/proto/aibridged.proto
+1-3Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ service Authorizer {
3535
// TODO: add authorization; currently only key validation takes place.
3636
rpc IsAuthorized(IsAuthorizedRequest) returns (IsAuthorizedResponse);
3737
// IsBudgetExceeded reports whether the user's AI spend has reached their
38-
// effective limit over [period_start, now].
38+
// effective limit for the current deployment-configured budget period.
3939
rpc IsBudgetExceeded(IsBudgetExceededRequest) returns (IsBudgetExceededResponse);
4040
}
4141

@@ -189,8 +189,6 @@ message IsAuthorizedResponse {
189189

190190
message IsBudgetExceededRequest {
191191
string user_id = 1; // UUID
192-
// The spend aggregation window is [period_start, now].
193-
google.protobuf.Timestamp period_start = 2;
194192
}
195193

196194
message IsBudgetExceededResponse {
Collapse file

‎coderd/aibridgedserver/aibridgedserver.go‎

Copy file name to clipboardExpand all lines: coderd/aibridgedserver/aibridgedserver.go
+15-8Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
3434
"github.com/coder/coder/v2/coderd/util/ptr"
3535
"github.com/coder/coder/v2/codersdk"
36+
"github.com/coder/quartz"
3637
)
3738

3839
var (
@@ -113,11 +114,15 @@ type Server struct {
113114
// budgetPolicy selects the effective group when a user belongs to multiple
114115
// budgeted groups, used for cost attribution on token usage records.
115116
budgetPolicy codersdk.AIBudgetPolicy
117+
// budgetPeriod is the deployment-configured budgeting period used to
118+
// derive the window over which user AI spend is aggregated.
119+
budgetPeriod codersdk.AIBudgetPeriod
120+
clock quartz.Clock
116121
}
117122

118123
func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logger slog.Logger, accessURL string,
119124
bridgeCfg codersdk.AIBridgeConfig, externalAuthConfigs []*externalauth.Config, experiments codersdk.Experiments,
120-
aiSeatTracker aiseats.SeatTracker,
125+
aiSeatTracker aiseats.SeatTracker, clock quartz.Clock,
121126
) (*Server, error) {
122127
eac := make(map[string]*externalauth.Config, len(externalAuthConfigs))
123128

@@ -138,6 +143,8 @@ func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logg
138143
structuredLogging: bridgeCfg.StructuredLogging.Value(),
139144
aiSeatTracker: aiSeatTracker,
140145
budgetPolicy: codersdk.NewAIBudgetPolicyFromString(bridgeCfg.BudgetPolicy),
146+
budgetPeriod: codersdk.NewAIBudgetPeriodFromString(bridgeCfg.BudgetPeriod),
147+
clock: clock,
141148
}
142149

143150
if bridgeCfg.InjectCoderMCPTools {
@@ -747,7 +754,8 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest
747754
}
748755

749756
// IsBudgetExceeded reports whether the user's AI spend has reached their
750-
// effective limit over [PeriodStart, now].
757+
// effective limit over [periodStart, now], where periodStart is the start of
758+
// the current deployment-configured budget period.
751759
func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (*proto.IsBudgetExceededResponse, error) {
752760
//nolint:gocritic // AIBridged has specific authz rules.
753761
ctx = dbauthz.AsAIBridged(ctx)
@@ -756,14 +764,13 @@ func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceede
756764
if err != nil {
757765
return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err)
758766
}
759-
// An unset PeriodStart deserializes to time.Unix(0, 0), which would
760-
// incorrectly aggregate the user's lifetime spend against a period budget.
761-
if in.PeriodStart == nil {
762-
return nil, xerrors.New("period_start is required")
767+
768+
periodWindow, err := budget.CurrentPeriod(s.clock.Now(), s.budgetPeriod)
769+
if err != nil {
770+
return nil, xerrors.Errorf("compute AI budget period: %w", err)
763771
}
764-
periodStart := in.GetPeriodStart().AsTime()
765772

766-
userBudget, err := s.checkUserAIBudget(ctx, userID, periodStart)
773+
userBudget, err := s.checkUserAIBudget(ctx, userID, periodWindow.Start)
767774
if err != nil {
768775
return nil, err
769776
}

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.