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: make user AI budget and spend endpoints consistent (#27611)

Browse Source 
## Description

Makes the user AI cost control endpoints consistent.

## Changes

- Replaces the flat `spend_limit_micros` and `limit_source` fields on
`GET /users/{user}/ai/spend` with a nested `effective_budget`, reusing
the type behind `group_budget`. The flat pair made it possible to encode
a limit without a source.
- Renames `AIGroupBudget` to `AIBudgetLimit`, since it also carries
`user_override` limits and is no longer group-specific. The type name is
not part of the wire format.
- Moves `/users/{user}/ai/budget` to `/users/{user}/ai/budget/override`.
The endpoint only ever managed the per-user override, which the type,
the handlers, and the operation IDs all already said; the path was the
only place that didn't.

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by
@ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-29 14:01:32 +01:00
committed by GitHub
parent e71249a821
commit d6a5c8e9f8
14 changed files with 194 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/aibridge.go
+7 -5
View File
@@ -732,7 +732,7 @@ func (api *API) deleteGroupAIBudget(rw http.ResponseWriter, r *http.Request) {
// @Tags Enterprise
// @Param user path string true "User ID, username, or me"
// @Success 200 {object} codersdk.UserAIBudgetOverride
// @Router /api/v2/users/{user}/ai/budget [get]
// @Router /api/v2/users/{user}/ai/budget/override [get]
func (api *API) userAIBudgetOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
user := httpmw.UserParam(r)
@@ -760,7 +760,7 @@ func (api *API) userAIBudgetOverride(rw http.ResponseWriter, r *http.Request) {
// @Param user path string true "User ID, username, or me"
// @Param request body codersdk.UpsertUserAIBudgetOverrideRequest true "Upsert user AI budget override request"
// @Success 200 {object} codersdk.UserAIBudgetOverride
// @Router /api/v2/users/{user}/ai/budget [put]
// @Router /api/v2/users/{user}/ai/budget/override [put]
func (api *API) upsertUserAIBudgetOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
user := httpmw.UserParam(r)
@@ -857,7 +857,7 @@ func (api *API) upsertUserAIBudgetOverride(rw http.ResponseWriter, r *http.Reque
// @Tags Enterprise
// @Param user path string true "User ID, username, or me"
// @Success 204
// @Router /api/v2/users/{user}/ai/budget [delete]
// @Router /api/v2/users/{user}/ai/budget/override [delete]
func (api *API) deleteUserAIBudgetOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
user := httpmw.UserParam(r)
@@ -961,8 +961,10 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) {
if ok {
resp.EffectiveGroupID = &effectiveGroup.GroupID
if effectiveGroup.Limit != nil {
resp.SpendLimitMicros = &effectiveGroup.Limit.SpendLimitMicros
resp.LimitSource = &effectiveGroup.Limit.Source
resp.EffectiveBudget = &codersdk.AIBudgetLimit{
SpendLimitMicros: effectiveGroup.Limit.SpendLimitMicros,
LimitSource: effectiveGroup.Limit.Source,
}
}
logger = logger.With(slog.F("effective_group_id", effectiveGroup.GroupID))
enterprise/coderd/aibridge_test.go
+50 -42
View File
@@ -3230,33 +3230,38 @@ func TestUserAISpendStatus(t *testing.T) {
overrideLimit *int64 // nil = no user override configured
spent int64 // 0 = no spend seeded
wantHasEffectiveGroup bool
wantSpendLimitMicros *int64
wantLimitSource *codersdk.AIBudgetLimitSource
wantEffectiveBudget *codersdk.AIBudgetLimit
wantCurrentSpendMicros int64
}{
{
name: "GroupBudget/ZeroSpend",
groupBudget: ptr.Ref(int64(1_000_000_000)),
wantHasEffectiveGroup: true,
wantSpendLimitMicros: ptr.Ref(int64(1_000_000_000)),
wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup),
wantEffectiveBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 1_000_000_000,
LimitSource: codersdk.AIBudgetLimitSourceGroup,
},
},
{
name: "GroupBudget/PartialSpend",
groupBudget: ptr.Ref(int64(1_000_000_000)),
spent: 250_000_000,
wantHasEffectiveGroup: true,
wantSpendLimitMicros: ptr.Ref(int64(1_000_000_000)),
wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup),
name: "GroupBudget/PartialSpend",
groupBudget: ptr.Ref(int64(1_000_000_000)),
spent: 250_000_000,
wantHasEffectiveGroup: true,
wantEffectiveBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 1_000_000_000,
LimitSource: codersdk.AIBudgetLimitSourceGroup,
},
wantCurrentSpendMicros: 250_000_000,
},
{
name: "GroupBudget/SpendExceedsLimit",
groupBudget: ptr.Ref(int64(1_000_000_000)),
spent: 1_500_000_000,
wantHasEffectiveGroup: true,
wantSpendLimitMicros: ptr.Ref(int64(1_000_000_000)),
wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceGroup),
name: "GroupBudget/SpendExceedsLimit",
groupBudget: ptr.Ref(int64(1_000_000_000)),
spent: 1_500_000_000,
wantHasEffectiveGroup: true,
wantEffectiveBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 1_000_000_000,
LimitSource: codersdk.AIBudgetLimitSourceGroup,
},
wantCurrentSpendMicros: 1_500_000_000,
},
{
@@ -3264,27 +3269,33 @@ func TestUserAISpendStatus(t *testing.T) {
groupBudget: ptr.Ref(int64(5_000_000_000)),
overrideLimit: ptr.Ref(int64(200_000_000)),
wantHasEffectiveGroup: true,
wantSpendLimitMicros: ptr.Ref(int64(200_000_000)),
wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceUserOverride),
wantEffectiveBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 200_000_000,
LimitSource: codersdk.AIBudgetLimitSourceUserOverride,
},
},
{
name: "UserOverride/PartialSpend",
groupBudget: ptr.Ref(int64(5_000_000_000)),
overrideLimit: ptr.Ref(int64(200_000_000)),
spent: 50_000_000,
wantHasEffectiveGroup: true,
wantSpendLimitMicros: ptr.Ref(int64(200_000_000)),
wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceUserOverride),
name: "UserOverride/PartialSpend",
groupBudget: ptr.Ref(int64(5_000_000_000)),
overrideLimit: ptr.Ref(int64(200_000_000)),
spent: 50_000_000,
wantHasEffectiveGroup: true,
wantEffectiveBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 200_000_000,
LimitSource: codersdk.AIBudgetLimitSourceUserOverride,
},
wantCurrentSpendMicros: 50_000_000,
},
{
name: "UserOverride/SpendExceedsLimit",
groupBudget: ptr.Ref(int64(5_000_000_000)),
overrideLimit: ptr.Ref(int64(200_000_000)),
spent: 350_000_000,
wantHasEffectiveGroup: true,
wantSpendLimitMicros: ptr.Ref(int64(200_000_000)),
wantLimitSource: ptr.Ref(codersdk.AIBudgetLimitSourceUserOverride),
name: "UserOverride/SpendExceedsLimit",
groupBudget: ptr.Ref(int64(5_000_000_000)),
overrideLimit: ptr.Ref(int64(200_000_000)),
spent: 350_000_000,
wantHasEffectiveGroup: true,
wantEffectiveBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 200_000_000,
LimitSource: codersdk.AIBudgetLimitSourceUserOverride,
},
wantCurrentSpendMicros: 350_000_000,
},
}
@@ -3343,8 +3354,7 @@ func TestUserAISpendStatus(t *testing.T) {
wantEffectiveGroupID = &group.ID
}
require.Equal(t, wantEffectiveGroupID, got.EffectiveGroupID)
require.Equal(t, tt.wantSpendLimitMicros, got.SpendLimitMicros)
require.Equal(t, tt.wantLimitSource, got.LimitSource)
require.Equal(t, tt.wantEffectiveBudget, got.EffectiveBudget)
})
}
@@ -3377,8 +3387,7 @@ func TestUserAISpendStatus(t *testing.T) {
got, err := adminClient.UserAISpendStatus(ctx, targetUser.ID)
require.NoError(t, err)
require.Equal(t, &everyoneGroupID, got.EffectiveGroupID)
require.Nil(t, got.SpendLimitMicros)
require.Nil(t, got.LimitSource)
require.Nil(t, got.EffectiveBudget)
require.Equal(t, int64(100_000_000), got.CurrentSpendMicros)
})
@@ -3402,8 +3411,7 @@ func TestUserAISpendStatus(t *testing.T) {
got, err := adminClient.UserAISpendStatus(ctx, orglessUser.ID)
require.NoError(t, err)
require.Nil(t, got.EffectiveGroupID)
require.Nil(t, got.SpendLimitMicros)
require.Nil(t, got.LimitSource)
require.Nil(t, got.EffectiveBudget)
require.Equal(t, int64(0), got.CurrentSpendMicros)
})
}
@@ -5221,14 +5229,14 @@ func TestGroupMembersAISpend(t *testing.T) {
spent int64
wantEffectiveGroup bool
wantEffectiveEveryone bool
wantGroupBudget *codersdk.AIGroupBudget
wantGroupBudget *codersdk.AIBudgetLimit
wantSpendMicros int64
}{
{
name: "BudgetZeroSpend",
groupLimit: 1_000_000_000,
wantEffectiveGroup: true,
wantGroupBudget: &codersdk.AIGroupBudget{
wantGroupBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 1_000_000_000,
LimitSource: codersdk.AIBudgetLimitSourceGroup,
},
@@ -5238,7 +5246,7 @@ func TestGroupMembersAISpend(t *testing.T) {
groupLimit: 1_000_000_000,
spent: 250_000_000,
wantEffectiveGroup: true,
wantGroupBudget: &codersdk.AIGroupBudget{
wantGroupBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 1_000_000_000,
LimitSource: codersdk.AIBudgetLimitSourceGroup,
},
@@ -5248,7 +5256,7 @@ func TestGroupMembersAISpend(t *testing.T) {
name: "OverrideBudget",
overrideLimit: 500_000_000,
wantEffectiveGroup: true,
wantGroupBudget: &codersdk.AIGroupBudget{
wantGroupBudget: &codersdk.AIBudgetLimit{
SpendLimitMicros: 500_000_000,
LimitSource: codersdk.AIBudgetLimitSourceUserOverride,
},
enterprise/coderd/coderd.go
+1 -1
View File
@@ -698,7 +698,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
apiKeyMiddleware,
httpmw.ExtractUserParam(options.Database),
)
r.Route("/budget", func(r chi.Router) {
r.Route("/budget/override", func(r chi.Router) {
r.Get("/", api.userAIBudgetOverride)
r.Put("/", api.upsertUserAIBudgetOverride)
r.Delete("/", api.deleteUserAIBudgetOverride)