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
coderd/apidoc/docs.go
Generated
+21 -25
View File
@@ -9950,7 +9950,7 @@ const docTemplate = `{
]
}
},
"/api/v2/users/{user}/ai/budget": {
"/api/v2/users/{user}/ai/budget/override": {
"get": {
"produces": [
"application/json"
@@ -15784,6 +15784,17 @@ const docTemplate = `{
}
}
},
"codersdk.AIBudgetLimit": {
"type": "object",
"properties": {
"limit_source": {
"$ref": "#/definitions/codersdk.AIBudgetLimitSource"
},
"spend_limit_micros": {
"type": "integer"
}
}
},
"codersdk.AIBudgetLimitSource": {
"type": "string",
"enum": [
@@ -15832,17 +15843,6 @@ const docTemplate = `{
}
}
},
"codersdk.AIGroupBudget": {
"type": "object",
"properties": {
"limit_source": {
"$ref": "#/definitions/codersdk.AIBudgetLimitSource"
},
"spend_limit_micros": {
"type": "integer"
}
}
},
"codersdk.AIProvider": {
"type": "object",
"properties": {
@@ -20838,7 +20838,7 @@ const docTemplate = `{
"description": "GroupBudget is the budget when the queried group is this user's\neffective budget source. Null when the user's budget resolves to another\ngroup or no budget applies to the user.",
"allOf": [
{
"$ref": "#/definitions/codersdk.AIGroupBudget"
"$ref": "#/definitions/codersdk.AIBudgetLimit"
}
]
},
@@ -26529,19 +26529,19 @@ const docTemplate = `{
"description": "CurrentSpendMicros is the user's spend on their effective group over\nthe current budget period.",
"type": "integer"
},
"effective_budget": {
"description": "EffectiveBudget is the spend limit that applies to the user, whether it\ncame from a group budget or a user override. Null when no budget\napplies, leaving the user's spend unlimited.",
"allOf": [
{
"$ref": "#/definitions/codersdk.AIBudgetLimit"
}
]
},
"effective_group_id": {
"description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization membership.",
"type": "string",
"format": "uuid"
},
"limit_source": {
"description": "LimitSource identifies which tier produced the limit. Null when no\nbudget applies.",
"allOf": [
{
"$ref": "#/definitions/codersdk.AIBudgetLimitSource"
}
]
},
"period_end": {
"description": "PeriodEnd is the exclusive upper bound of the current budget\nperiod.",
"type": "string",
@@ -26552,10 +26552,6 @@ const docTemplate = `{
"type": "string",
"format": "date-time"
},
"spend_limit_micros": {
"description": "SpendLimitMicros is the effective spend limit in micro-units.\nNull when no budget applies to the user (unlimited).",
"type": "integer"
},
"user_id": {
"type": "string",
"format": "uuid"
coderd/apidoc/swagger.json
Generated
+21 -25
View File
@@ -8823,7 +8823,7 @@
]
}
},
"/api/v2/users/{user}/ai/budget": {
"/api/v2/users/{user}/ai/budget/override": {
"get": {
"produces": ["application/json"],
"tags": ["Enterprise"],
@@ -14078,6 +14078,17 @@
}
}
},
"codersdk.AIBudgetLimit": {
"type": "object",
"properties": {
"limit_source": {
"$ref": "#/definitions/codersdk.AIBudgetLimitSource"
},
"spend_limit_micros": {
"type": "integer"
}
}
},
"codersdk.AIBudgetLimitSource": {
"type": "string",
"enum": ["user_override", "group"],
@@ -14123,17 +14134,6 @@
}
}
},
"codersdk.AIGroupBudget": {
"type": "object",
"properties": {
"limit_source": {
"$ref": "#/definitions/codersdk.AIBudgetLimitSource"
},
"spend_limit_micros": {
"type": "integer"
}
}
},
"codersdk.AIProvider": {
"type": "object",
"properties": {
@@ -18972,7 +18972,7 @@
"description": "GroupBudget is the budget when the queried group is this user's\neffective budget source. Null when the user's budget resolves to another\ngroup or no budget applies to the user.",
"allOf": [
{
"$ref": "#/definitions/codersdk.AIGroupBudget"
"$ref": "#/definitions/codersdk.AIBudgetLimit"
}
]
},
@@ -24386,19 +24386,19 @@
"description": "CurrentSpendMicros is the user's spend on their effective group over\nthe current budget period.",
"type": "integer"
},
"effective_budget": {
"description": "EffectiveBudget is the spend limit that applies to the user, whether it\ncame from a group budget or a user override. Null when no budget\napplies, leaving the user's spend unlimited.",
"allOf": [
{
"$ref": "#/definitions/codersdk.AIBudgetLimit"
}
]
},
"effective_group_id": {
"description": "EffectiveGroupID is the group the spend is attributed to, falling back to\nthe Everyone group when no budget applies. Null only when the user has no\norganization membership.",
"type": "string",
"format": "uuid"
},
"limit_source": {
"description": "LimitSource identifies which tier produced the limit. Null when no\nbudget applies.",
"allOf": [
{
"$ref": "#/definitions/codersdk.AIBudgetLimitSource"
}
]
},
"period_end": {
"description": "PeriodEnd is the exclusive upper bound of the current budget\nperiod.",
"type": "string",
@@ -24409,10 +24409,6 @@
"type": "string",
"format": "date-time"
},
"spend_limit_micros": {
"description": "SpendLimitMicros is the effective spend limit in micro-units.\nNull when no budget applies to the user (unlimited).",
"type": "integer"
},
"user_id": {
"type": "string",
"format": "uuid"
coderd/database/db2sdk/db2sdk.go
+1 -1
View File
@@ -1490,7 +1490,7 @@ func GroupMemberAISpend(row database.GetGroupMembersAISpendRow) codersdk.GroupMe
member.EffectiveGroupID = &row.EffectiveGroupID.UUID
}
if row.SpendLimitMicros.Valid {
member.GroupBudget = &codersdk.AIGroupBudget{
member.GroupBudget = &codersdk.AIBudgetLimit{
SpendLimitMicros: row.SpendLimitMicros.Int64,
LimitSource: codersdk.AIBudgetLimitSource(row.LimitSource.String),
}
codersdk/aibridge.go
+11 -13
View File
@@ -32,28 +32,26 @@ const (
AIBudgetLimitSourceGroup AIBudgetLimitSource = "group"
)
// AIGroupBudget is an AI spend limit and the tier that produced it. Both
// AIBudgetLimit is an AI spend limit and the tier that produced it. Both
// fields are always populated together.
type AIGroupBudget struct {
type AIBudgetLimit struct {
SpendLimitMicros int64 `json:"spend_limit_micros"`
LimitSource AIBudgetLimitSource `json:"limit_source"`
}
// UserAIBudgetSummary is the effective AI budget for a user. When no budget
// applies, the effective group falls back to the Everyone group with a null
// limit and source.
// budget.
type UserAIBudgetSummary struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
// EffectiveGroupID is the group the spend is attributed to, falling back to
// the Everyone group when no budget applies. Null only when the user has no
// organization membership.
EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"`
// SpendLimitMicros is the effective spend limit in micro-units.
// Null when no budget applies to the user (unlimited).
SpendLimitMicros *int64 `json:"spend_limit_micros"`
// LimitSource identifies which tier produced the limit. Null when no
// budget applies.
LimitSource *AIBudgetLimitSource `json:"limit_source"`
// EffectiveBudget is the spend limit that applies to the user, whether it
// came from a group budget or a user override. Null when no budget
// applies, leaving the user's spend unlimited.
EffectiveBudget *AIBudgetLimit `json:"effective_budget"`
}
// AISpendPeriodWindow is the [Start, End) window over which AI spend is
@@ -126,7 +124,7 @@ type GroupMemberAISpend struct {
// GroupBudget is the budget when the queried group is this user's
// effective budget source. Null when the user's budget resolves to another
// group or no budget applies to the user.
GroupBudget *AIGroupBudget `json:"group_budget"`
GroupBudget *AIBudgetLimit `json:"group_budget"`
// GroupSpendMicros is the user's spend attributed to the queried group
// over the current budget period.
GroupSpendMicros int64 `json:"group_spend_micros"`
@@ -498,7 +496,7 @@ type UpsertUserAIBudgetOverrideRequest struct {
// UserAIBudgetOverride returns the AI spend budget override configured for the given user.
func (c *Client) UserAIBudgetOverride(ctx context.Context, user uuid.UUID) (UserAIBudgetOverride, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()),
fmt.Sprintf("/api/v2/users/%s/ai/budget/override", user.String()),
nil,
)
if err != nil {
@@ -516,7 +514,7 @@ func (c *Client) UserAIBudgetOverride(ctx context.Context, user uuid.UUID) (User
// UpsertUserAIBudgetOverride creates or updates the AI spend budget override for the given user.
func (c *Client) UpsertUserAIBudgetOverride(ctx context.Context, user uuid.UUID, req UpsertUserAIBudgetOverrideRequest) (UserAIBudgetOverride, error) {
res, err := c.Request(ctx, http.MethodPut,
fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()),
fmt.Sprintf("/api/v2/users/%s/ai/budget/override", user.String()),
req,
)
if err != nil {
@@ -534,7 +532,7 @@ func (c *Client) UpsertUserAIBudgetOverride(ctx context.Context, user uuid.UUID,
// DeleteUserAIBudgetOverride removes the AI spend budget override for the given user.
func (c *Client) DeleteUserAIBudgetOverride(ctx context.Context, user uuid.UUID) error {
res, err := c.Request(ctx, http.MethodDelete,
fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()),
fmt.Sprintf("/api/v2/users/%s/ai/budget/override", user.String()),
nil,
)
if err != nil {
docs/reference/api/enterprise.md
Generated
+10 -8
View File
@@ -3982,12 +3982,12 @@ To perform this operation, you must be authenticated. [Learn more](authenticatio
```sh
# Example request using curl
curl -X GET http://coder-server:8080/api/v2/users/{user}/ai/budget \
curl -X GET http://coder-server:8080/api/v2/users/{user}/ai/budget/override \
-H 'Accept: application/json' \
-H 'Coder-Session-Token: API_KEY'
```
`GET /api/v2/users/{user}/ai/budget`
`GET /api/v2/users/{user}/ai/budget/override`
### Parameters
@@ -4023,13 +4023,13 @@ To perform this operation, you must be authenticated. [Learn more](authenticatio
```sh
# Example request using curl
curl -X PUT http://coder-server:8080/api/v2/users/{user}/ai/budget \
curl -X PUT http://coder-server:8080/api/v2/users/{user}/ai/budget/override \
-H 'Content-Type: application/json' \
-H 'Accept: application/json' \
-H 'Coder-Session-Token: API_KEY'
```
`PUT /api/v2/users/{user}/ai/budget`
`PUT /api/v2/users/{user}/ai/budget/override`
> Body parameter
@@ -4075,11 +4075,11 @@ To perform this operation, you must be authenticated. [Learn more](authenticatio
```sh
# Example request using curl
curl -X DELETE http://coder-server:8080/api/v2/users/{user}/ai/budget \
curl -X DELETE http://coder-server:8080/api/v2/users/{user}/ai/budget/override \
-H 'Coder-Session-Token: API_KEY'
```
`DELETE /api/v2/users/{user}/ai/budget`
`DELETE /api/v2/users/{user}/ai/budget/override`
### Parameters
@@ -4121,11 +4121,13 @@ curl -X GET http://coder-server:8080/api/v2/users/{user}/ai/spend \
```json
{
"current_spend_micros": 0,
"effective_budget": {
"limit_source": "user_override",
"spend_limit_micros": 0
},
"effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0",
"limit_source": "user_override",
"period_end": "2019-08-24T14:15:22Z",
"period_start": "2019-08-24T14:15:22Z",
"spend_limit_micros": 0,
"user_id": "a169451c-8525-4352-b8ca-070dd449a1a5"
}
```
docs/reference/api/schemas.md
Generated
+29 -28
View File
@@ -938,6 +938,22 @@
| `server_url` | string | false | | |
| `tool` | string | false | | |
## codersdk.AIBudgetLimit
```json
{
"limit_source": "user_override",
"spend_limit_micros": 0
}
```
### Properties
| Name | Type | Required | Restrictions | Description |
|----------------------|--------------------------------------------------------------|----------|--------------|-------------|
| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | |
| `spend_limit_micros` | integer | false | | |
## codersdk.AIBudgetLimitSource
```json
@@ -1071,22 +1087,6 @@
| `last_heartbeat_at` | string | false | | |
| `name` | string | false | | |
## codersdk.AIGroupBudget
```json
{
"limit_source": "user_override",
"spend_limit_micros": 0
}
```
### Properties
| Name | Type | Required | Restrictions | Description |
|----------------------|--------------------------------------------------------------|----------|--------------|-------------|
| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | |
| `spend_limit_micros` | integer | false | | |
## codersdk.AIProvider
```json
@@ -8029,7 +8029,7 @@ Only certain features set these fields: - FeatureManagedAgentLimit - FeatureAgen
| Name | Type | Required | Restrictions | Description |
|----------------------|--------------------------------------------------|----------|--------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `effective_group_id` | string | false | | Effective group ID is the user's effective budget group within the queried group's organization, falling back to the Everyone group when no budget applies. Null when the effective group belongs to a different organization than the queried group. |
| `group_budget` | [codersdk.AIGroupBudget](#codersdkaigroupbudget) | false | | Group budget is the budget when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user. |
| `group_budget` | [codersdk.AIBudgetLimit](#codersdkaibudgetlimit) | false | | Group budget is the budget when the queried group is this user's effective budget source. Null when the user's budget resolves to another group or no budget applies to the user. |
| `group_spend_micros` | integer | false | | Group spend micros is the user's spend attributed to the queried group over the current budget period. |
| `user_id` | string | false | | |
@@ -14478,26 +14478,27 @@ If the schedule is empty, the user will be updated to use the default schedule.|
```json
{
"current_spend_micros": 0,
"effective_budget": {
"limit_source": "user_override",
"spend_limit_micros": 0
},
"effective_group_id": "85e2b926-ddfb-4c66-b68e-b66e5acec6c0",
"limit_source": "user_override",
"period_end": "2019-08-24T14:15:22Z",
"period_start": "2019-08-24T14:15:22Z",
"spend_limit_micros": 0,
"user_id": "a169451c-8525-4352-b8ca-070dd449a1a5"
}
```
### Properties
| Name | Type | Required | Restrictions | Description |
|------------------------|--------------------------------------------------------------|----------|--------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `current_spend_micros` | integer | false | | Current spend micros is the user's spend on their effective group over the current budget period. |
| `effective_group_id` | string | false | | Effective group ID is the group the spend is attributed to, falling back to the Everyone group when no budget applies. Null only when the user has no organization membership. |
| `limit_source` | [codersdk.AIBudgetLimitSource](#codersdkaibudgetlimitsource) | false | | Limit source identifies which tier produced the limit. Null when no budget applies. |
| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. |
| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. |
| `spend_limit_micros` | integer | false | | Spend limit micros is the effective spend limit in micro-units. Null when no budget applies to the user (unlimited). |
| `user_id` | string | false | | |
| Name | Type | Required | Restrictions | Description |
|------------------------|--------------------------------------------------|----------|--------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `current_spend_micros` | integer | false | | Current spend micros is the user's spend on their effective group over the current budget period. |
| `effective_budget` | [codersdk.AIBudgetLimit](#codersdkaibudgetlimit) | false | | Effective budget is the spend limit that applies to the user, whether it came from a group budget or a user override. Null when no budget applies, leaving the user's spend unlimited. |
| `effective_group_id` | string | false | | Effective group ID is the group the spend is attributed to, falling back to the Everyone group when no budget applies. Null only when the user has no organization membership. |
| `period_end` | string | false | | Period end is the exclusive upper bound of the current budget period. |
| `period_start` | string | false | | Period start is the inclusive lower bound of the current budget period. |
| `user_id` | string | false | | |
## codersdk.UserActivity
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)
site/src/api/api.ts
+3 -3
View File
@@ -1615,7 +1615,7 @@ class ApiMethods {
userId: TypesGen.User["id"],
): Promise<TypesGen.UserAIBudgetOverride> => {
const response = await this.axios.get<TypesGen.UserAIBudgetOverride>(
`/api/v2/users/${encodeURIComponent(userId)}/ai/budget`,
`/api/v2/users/${encodeURIComponent(userId)}/ai/budget/override`,
);
return response.data;
@@ -1626,7 +1626,7 @@ class ApiMethods {
data: TypesGen.UpsertUserAIBudgetOverrideRequest,
): Promise<TypesGen.UserAIBudgetOverride> => {
const response = await this.axios.put<TypesGen.UserAIBudgetOverride>(
`/api/v2/users/${encodeURIComponent(userId)}/ai/budget`,
`/api/v2/users/${encodeURIComponent(userId)}/ai/budget/override`,
data,
);
@@ -1637,7 +1637,7 @@ class ApiMethods {
userId: TypesGen.User["id"],
): Promise<void> => {
await this.axios.delete(
`/api/v2/users/${encodeURIComponent(userId)}/ai/budget`,
`/api/v2/users/${encodeURIComponent(userId)}/ai/budget/override`,
);
};
site/src/api/typesGenerated.ts
Generated
+16 -20
View File
@@ -266,6 +266,16 @@ export interface AIBridgeToolCall {
readonly created_at: string;
}
// From codersdk/aibridge.go
/**
* AIBudgetLimit is an AI spend limit and the tier that produced it. Both
* fields are always populated together.
*/
export interface AIBudgetLimit {
readonly spend_limit_micros: number;
readonly limit_source: AIBudgetLimitSource;
}
// From codersdk/aibridge.go
export type AIBudgetLimitSource = "group" | "user_override";
@@ -310,16 +320,6 @@ export interface AIGatewayKey {
*/
export const AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key";
// From codersdk/aibridge.go
/**
* AIGroupBudget is an AI spend limit and the tier that produced it. Both
* fields are always populated together.
*/
export interface AIGroupBudget {
readonly spend_limit_micros: number;
readonly limit_source: AIBudgetLimitSource;
}
// From codersdk/aiproviders.go
/**
* AIProvider represents an AI provider configuration row as returned
@@ -5543,7 +5543,7 @@ export interface GroupMemberAISpend {
* effective budget source. Null when the user's budget resolves to another
* group or no budget applies to the user.
*/
readonly group_budget: AIGroupBudget | null;
readonly group_budget: AIBudgetLimit | null;
/**
* GroupSpendMicros is the user's spend attributed to the queried group
* over the current budget period.
@@ -10301,7 +10301,7 @@ export interface UserAIBudgetOverride {
/**
* UserAIBudgetSummary is the effective AI budget for a user. When no budget
* applies, the effective group falls back to the Everyone group with a null
* limit and source.
* budget.
*/
export interface UserAIBudgetSummary {
readonly user_id: string;
@@ -10312,15 +10312,11 @@ export interface UserAIBudgetSummary {
*/
readonly effective_group_id: string | null;
/**
* SpendLimitMicros is the effective spend limit in micro-units.
* Null when no budget applies to the user (unlimited).
* EffectiveBudget is the spend limit that applies to the user, whether it
* came from a group budget or a user override. Null when no budget
* applies, leaving the user's spend unlimited.
*/
readonly spend_limit_micros: number | null;
/**
* LimitSource identifies which tier produced the limit. Null when no
* budget applies.
*/
readonly limit_source: AIBudgetLimitSource | null;
readonly effective_budget: AIBudgetLimit | null;
}
// From codersdk/chats.go
site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.stories.tsx
+16 -5
View File
@@ -12,9 +12,11 @@ import { UserDropdown } from "./UserDropdown";
const mockAISpend: UserAISpendStatus = {
user_id: MockUserOwner.id,
spend_limit_micros: 1_200_000_000,
effective_group_id: "grp-789",
limit_source: "group",
effective_budget: {
spend_limit_micros: 1_200_000_000,
limit_source: "group",
},
current_spend_micros: 819_000_000,
period_start: "2026-06-01T00:00:00Z",
period_end: "2026-07-01T00:00:00Z",
@@ -159,7 +161,7 @@ export const AISpendUnlimited: Story = {
parameters: {
...aiCostControl,
queries: [
{ key: meAISpendKey, data: { ...mockAISpend, spend_limit_micros: null } },
{ key: meAISpendKey, data: { ...mockAISpend, effective_budget: null } },
],
},
play: async ({ canvasElement, step }) => {
@@ -205,7 +207,10 @@ export const AISpendZeroLimit: Story = {
data: {
...mockAISpend,
current_spend_micros: 0,
spend_limit_micros: 0,
effective_budget: {
spend_limit_micros: 0,
limit_source: "group",
},
},
},
],
@@ -281,7 +286,13 @@ export const AISpendHiddenOnNegativeLimit: Story = {
parameters: {
...aiCostControl,
queries: [
{ key: meAISpendKey, data: { ...mockAISpend, spend_limit_micros: -1 } },
{
key: meAISpendKey,
data: {
...mockAISpend,
effective_budget: { spend_limit_micros: -1, limit_source: "group" },
},
},
],
},
play: async ({ canvasElement, step }) => {
site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx
+4 -3
View File
@@ -44,16 +44,17 @@ export const UserDropdown: FC<UserDropdownProps> = ({
enabled: aibridgeVisible,
});
// A null limit is unlimited and still shown.
// A null budget is unlimited and still shown.
const hasValidSpend =
data !== undefined &&
data.current_spend_micros >= 0 &&
(data.spend_limit_micros === null || data.spend_limit_micros >= 0);
(data.effective_budget === null ||
data.effective_budget.spend_limit_micros >= 0);
const spend =
aibridgeVisible && !isError && hasValidSpend
? {
currentSpend: data.current_spend_micros,
spendLimit: data.spend_limit_micros,
spendLimit: data.effective_budget?.spend_limit_micros ?? null,
}
: null;
const severity =
site/src/pages/GroupsPage/GroupPage.stories.tsx
+4 -2
View File
@@ -100,9 +100,11 @@ const membersQuery = (data: unknown) => ({
/** period_end is exclusive. */
const mockUserAISpend: UserAISpendStatus = {
user_id: MockUserOwner.id,
spend_limit_micros: 9_000_000_000,
effective_group_id: MockGroupWithoutMembers.id,
limit_source: "group",
effective_budget: {
spend_limit_micros: 9_000_000_000,
limit_source: "group",
},
current_spend_micros: 1_345_000_000,
period_start: "2026-06-01T00:00:00Z",
period_end: "2026-07-01T00:00:00Z",