diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go index aa0e9508bc0..b1ef9742db9 100644 --- a/coderd/aibridge/budget/budget.go +++ b/coderd/aibridge/budget/budget.go @@ -15,17 +15,6 @@ import ( "github.com/coder/coder/v2/codersdk" ) -// LimitSource identifies which tier produced an EffectiveBudget. -type LimitSource string - -const ( - // SourceUserOverride indicates the budget came from a per-user override. - SourceUserOverride LimitSource = "user_override" - // SourceGroup indicates the budget came from a group budget selected by the - // deployment policy. - SourceGroup LimitSource = "group" -) - // Store is the subset of database.Store needed to resolve a user's effective // AI budget. type Store interface { @@ -41,7 +30,7 @@ type EffectiveBudget struct { // SpendLimitMicros is the effective spend limit in micro-units // (1 unit = 1,000,000). SpendLimitMicros int64 - Source LimitSource + Source codersdk.AIBudgetLimitSource } // ResolveUserAIBudget returns the effective AI budget for userID. The second @@ -55,7 +44,7 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy return EffectiveBudget{ GroupID: override.GroupID, SpendLimitMicros: override.SpendLimitMicros, - Source: SourceUserOverride, + Source: codersdk.AIBudgetLimitSourceUserOverride, }, true, nil } if !errors.Is(err, sql.ErrNoRows) { @@ -75,7 +64,7 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy return EffectiveBudget{ GroupID: row.GroupID, SpendLimitMicros: row.SpendLimitMicros, - Source: SourceGroup, + Source: codersdk.AIBudgetLimitSourceGroup, }, true, nil default: return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy) diff --git a/coderd/aibridge/budget/budget_test.go b/coderd/aibridge/budget/budget_test.go index 9171ac2e389..4caaae02780 100644 --- a/coderd/aibridge/budget/budget_test.go +++ b/coderd/aibridge/budget/budget_test.go @@ -72,7 +72,7 @@ func TestResolveUserAIBudget(t *testing.T) { SpendLimitMicros: 1_000_000, }) require.NoError(t, err) - return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: budget.SourceUserOverride}, true + return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true }, }, { @@ -82,7 +82,7 @@ func TestResolveUserAIBudget(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{}) gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000) - return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: budget.SourceGroup}, true + return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true }, }, { @@ -94,7 +94,7 @@ func TestResolveUserAIBudget(t *testing.T) { budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000) budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000) high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000) - return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: budget.SourceGroup}, true + return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true }, }, { @@ -106,7 +106,7 @@ func TestResolveUserAIBudget(t *testing.T) { // Equal limits; "alpha" must win over "beta" by name ascending. alpha := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000) budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000) - return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: budget.SourceGroup}, true + return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true }, }, { @@ -124,7 +124,7 @@ func TestResolveUserAIBudget(t *testing.T) { if bytes.Compare(g2[:], g1[:]) < 0 { winner = g2 } - return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: budget.SourceGroup}, true + return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true }, }, { @@ -147,7 +147,7 @@ func TestResolveUserAIBudget(t *testing.T) { // Membership is via organization_members only (no group_members row), // exercising the org-members half of group_members_expanded. everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000) - return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: budget.SourceGroup}, true + return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true }, }, { @@ -165,7 +165,7 @@ func TestResolveUserAIBudget(t *testing.T) { SpendLimitMicros: 2_000_000, }) require.NoError(t, err) - return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: budget.SourceUserOverride}, true + return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true }, }, { diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 576ff8aad1d..f018cf7bc6e 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -9652,6 +9652,40 @@ const docTemplate = `{ ] } }, + "/api/v2/users/{user}/ai/spend": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "Enterprise" + ], + "summary": "Get user AI spend", + "operationId": "get-user-ai-spend", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.UserAISpendStatus" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/users/{user}/appearance": { "get": { "produces": [ @@ -15251,6 +15285,17 @@ const docTemplate = `{ } } }, + "codersdk.AIBudgetLimitSource": { + "type": "string", + "enum": [ + "user_override", + "group" + ], + "x-enum-varnames": [ + "AIBudgetLimitSourceUserOverride", + "AIBudgetLimitSourceGroup" + ] + }, "codersdk.AIConfig": { "type": "object", "properties": { @@ -25657,6 +25702,46 @@ const docTemplate = `{ } } }, + "codersdk.UserAISpendStatus": { + "type": "object", + "properties": { + "current_spend_micros": { + "description": "CurrentSpendMicros is the user's spend on their effective group over\nthe current budget period.", + "type": "integer" + }, + "effective_group_id": { + "description": "EffectiveGroupID is the group the spend is attributed to. Null when\nno budget applies.", + "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", + "format": "date-time" + }, + "period_start": { + "description": "PeriodStart is the inclusive lower bound of the current budget\nperiod.", + "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" + } + } + }, "codersdk.UserActivity": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 2993ba71292..699d7805c9c 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -8555,6 +8555,36 @@ ] } }, + "/api/v2/users/{user}/ai/spend": { + "get": { + "produces": ["application/json"], + "tags": ["Enterprise"], + "summary": "Get user AI spend", + "operationId": "get-user-ai-spend", + "parameters": [ + { + "type": "string", + "description": "User ID, username, or me", + "name": "user", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.UserAISpendStatus" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/users/{user}/appearance": { "get": { "produces": ["application/json"], @@ -13597,6 +13627,14 @@ } } }, + "codersdk.AIBudgetLimitSource": { + "type": "string", + "enum": ["user_override", "group"], + "x-enum-varnames": [ + "AIBudgetLimitSourceUserOverride", + "AIBudgetLimitSourceGroup" + ] + }, "codersdk.AIConfig": { "type": "object", "properties": { @@ -23585,6 +23623,46 @@ } } }, + "codersdk.UserAISpendStatus": { + "type": "object", + "properties": { + "current_spend_micros": { + "description": "CurrentSpendMicros is the user's spend on their effective group over\nthe current budget period.", + "type": "integer" + }, + "effective_group_id": { + "description": "EffectiveGroupID is the group the spend is attributed to. Null when\nno budget applies.", + "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", + "format": "date-time" + }, + "period_start": { + "description": "PeriodStart is the inclusive lower bound of the current budget\nperiod.", + "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" + } + } + }, "codersdk.UserActivity": { "type": "object", "properties": { diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index fd6adb1859a..e7e54791bd0 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -12,6 +12,49 @@ import ( "golang.org/x/xerrors" ) +// AIBudgetLimitSource identifies which tier produced the user's +// effective budget limit. +type AIBudgetLimitSource string + +const ( + // AIBudgetLimitSourceUserOverride indicates the limit came from a + // per-user override. + AIBudgetLimitSourceUserOverride AIBudgetLimitSource = "user_override" + // AIBudgetLimitSourceGroup indicates the limit came from a group + // budget selected by the deployment budget policy. + AIBudgetLimitSourceGroup AIBudgetLimitSource = "group" +) + +// UserAIBudgetSummary is the effective AI budget for a user. When no +// budget applies, all fields except UserID are null. +type UserAIBudgetSummary struct { + UserID uuid.UUID `json:"user_id" format:"uuid"` + // EffectiveGroupID is the group the spend is attributed to. Null when + // no budget applies. + 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"` +} + +// UserAISpendStatus is the current AI spend snapshot for a user within +// the active budget period. +type UserAISpendStatus struct { + UserAIBudgetSummary + // CurrentSpendMicros is the user's spend on their effective group over + // the current budget period. + CurrentSpendMicros int64 `json:"current_spend_micros"` + // PeriodStart is the inclusive lower bound of the current budget + // period. + PeriodStart time.Time `json:"period_start" format:"date-time"` + // PeriodEnd is the exclusive upper bound of the current budget + // period. + PeriodEnd time.Time `json:"period_end" format:"date-time"` +} + type AIBridgeSession struct { ID string `json:"id"` Initiator MinimalUser `json:"initiator"` @@ -375,3 +418,22 @@ func (c *Client) DeleteUserAIBudgetOverride(ctx context.Context, user uuid.UUID) } return nil } + +// UserAISpendStatus returns the current AI spend snapshot for the given user +// within the active budget period. +func (c *Client) UserAISpendStatus(ctx context.Context, user uuid.UUID) (UserAISpendStatus, error) { + res, err := c.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/users/%s/ai/spend", user.String()), + nil, + ) + if err != nil { + return UserAISpendStatus{}, xerrors.Errorf("make request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return UserAISpendStatus{}, ReadBodyAsError(res) + } + var resp UserAISpendStatus + return resp, json.NewDecoder(res.Body).Decode(&resp) +} diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index f6312412293..34915c902b5 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -3777,6 +3777,49 @@ curl -X DELETE http://coder-server:8080/api/v2/users/{user}/ai/budget \ To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get user AI spend + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/users/{user}/ai/spend \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/users/{user}/ai/spend` + +### Parameters + +| Name | In | Type | Required | Description | +|--------|------|--------|----------|--------------------------| +| `user` | path | string | true | User ID, username, or me | + +### Example responses + +> 200 Response + +```json +{ + "current_spend_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" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|--------------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.UserAISpendStatus](schemas.md#codersdkuseraispendstatus) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get user quiet hours schedule ### Code samples diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 4000260fdba..d374d10e115 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -905,6 +905,20 @@ | `server_url` | string | false | | | | `tool` | string | false | | | +## codersdk.AIBudgetLimitSource + +```json +"user_override" +``` + +### Properties + +#### Enumerated Values + +| Value(s) | +|--------------------------| +| `group`, `user_override` | + ## codersdk.AIConfig ```json @@ -13925,6 +13939,32 @@ If the schedule is empty, the user will be updated to use the default schedule.| | `updated_at` | string | false | | | | `user_id` | string | false | | | +## codersdk.UserAISpendStatus + +```json +{ + "current_spend_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. Null when no budget applies. | +| `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 | | | + ## codersdk.UserActivity ```json diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index d4f89982528..8736088a803 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -16,6 +16,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" + "github.com/coder/coder/v2/coderd/aibridge/budget" "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" @@ -586,7 +587,7 @@ func (api *API) groupAIBudget(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() group := httpmw.GroupParam(r) - budget, err := api.Database.GetGroupAIBudget(ctx, group.ID) + groupBudget, err := api.Database.GetGroupAIBudget(ctx, group.ID) if httpapi.Is404Error(err) { httpapi.ResourceNotFound(rw) return @@ -597,7 +598,7 @@ func (api *API) groupAIBudget(rw http.ResponseWriter, r *http.Request) { return } - httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(budget)) + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(groupBudget)) } // @Summary Upsert group AI budget @@ -875,3 +876,65 @@ func (api *API) deleteUserAIBudgetOverride(rw http.ResponseWriter, r *http.Reque rw.WriteHeader(http.StatusNoContent) } + +// @Summary Get user AI spend +// @ID get-user-ai-spend +// @Security CoderSessionToken +// @Produce json +// @Tags Enterprise +// @Param user path string true "User ID, username, or me" +// @Success 200 {object} codersdk.UserAISpendStatus +// @Router /api/v2/users/{user}/ai/spend [get] +func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + user := httpmw.UserParam(r) + logger := api.Logger.With(slog.F("user_id", user.ID)) + + periodWindow, err := budget.CurrentPeriod(api.Clock.Now(), codersdk.AIBudgetPeriodMonth) + if err != nil { + logger.Error(ctx, "failed to compute AI budget period", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + logger = logger.With( + slog.F("period_start", periodWindow.Start), + slog.F("period_end", periodWindow.End), + ) + + policy := codersdk.NewAIBudgetPolicyFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPolicy) + effectiveBudget, ok, err := budget.ResolveUserAIBudget(ctx, api.Database, user.ID, policy) + if err != nil { + logger.Error(ctx, "failed to resolve user AI budget", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + resp := codersdk.UserAISpendStatus{ + UserAIBudgetSummary: codersdk.UserAIBudgetSummary{ + UserID: user.ID, + }, + PeriodStart: periodWindow.Start, + PeriodEnd: periodWindow.End, + } + + if ok { + resp.EffectiveGroupID = &effectiveBudget.GroupID + resp.SpendLimitMicros = &effectiveBudget.SpendLimitMicros + resp.LimitSource = &effectiveBudget.Source + logger = logger.With(slog.F("effective_group_id", effectiveBudget.GroupID)) + + spend, err := api.Database.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: effectiveBudget.GroupID, + PeriodStart: periodWindow.Start, + }) + if err != nil { + logger.Error(ctx, "failed to get user AI spend", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + resp.CurrentSpendMicros = spend.SpendMicros + } + + httpapi.Write(ctx, rw, http.StatusOK, resp) +} diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 52540bf0c71..ccfe716864b 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -19,13 +19,16 @@ import ( "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" entaudit "github.com/coder/coder/v2/enterprise/audit" "github.com/coder/coder/v2/enterprise/audit/backends" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" "github.com/coder/coder/v2/enterprise/coderd/license" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" "github.com/coder/serpent" ) @@ -2337,7 +2340,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/CreatesAndUpdates", func(t *testing.T) { t.Parallel() - adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) // First upsert creates the override. @@ -2367,7 +2370,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/ReassignsGroup", func(t *testing.T) { t.Parallel() - adminClient, targetUser, groupA := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, groupA := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) // First upsert: attribute spend to groupA. @@ -2404,7 +2407,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/EveryoneGroup", func(t *testing.T) { t.Parallel() - adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, _ := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) // The Everyone group has id == organization_id, and the target user @@ -2427,7 +2430,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/AcceptsZeroSpendLimit", func(t *testing.T) { t.Parallel() - adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) // 0 is a valid value: it blocks all spend for the user. @@ -2442,7 +2445,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/RejectsNegativeSpend", func(t *testing.T) { t.Parallel() - adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ @@ -2457,7 +2460,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/RejectsUnknownGroup", func(t *testing.T) { t.Parallel() - adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, _ := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) // A group_id that doesn't exist (or that the caller can't see) @@ -2474,7 +2477,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Upsert/RejectsNonMemberGroup", func(t *testing.T) { t.Parallel() - adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, _ := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) // Create a second group the target is NOT a member of. @@ -2495,7 +2498,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Get/AbsentReturns404", func(t *testing.T) { t.Parallel() - adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, _ := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) _, err := adminClient.UserAIBudgetOverride(ctx, targetUser.ID) @@ -2507,7 +2510,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Get/UnknownUserReturns404", func(t *testing.T) { t.Parallel() - adminClient, _, _ := setupUserAIBudgetOverrideTest(t) + adminClient, _, _ := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) _, err := adminClient.UserAIBudgetOverride(ctx, uuid.New()) @@ -2519,7 +2522,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Delete/RoundTrip", func(t *testing.T) { t.Parallel() - adminClient, targetUser, group := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ @@ -2539,7 +2542,7 @@ func TestUserAIBudgetOverride(t *testing.T) { t.Run("Delete/AbsentReturns404", func(t *testing.T) { t.Parallel() - adminClient, targetUser, _ := setupUserAIBudgetOverrideTest(t) + adminClient, targetUser, _ := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-test-group"}) ctx := testutil.Context(t, testutil.WaitLong) err := adminClient.DeleteUserAIBudgetOverride(ctx, targetUser.ID) @@ -2974,10 +2977,185 @@ func TestUserAIBudgetOverrideDeletedOnMembershipRemoval(t *testing.T) { }) } -// setupUserAIBudgetOverrideTest returns an Admin client, a target user, and a -// group the target user is a member of. -func setupUserAIBudgetOverrideTest(t *testing.T) (adminClient *codersdk.Client, targetUser codersdk.User, group codersdk.Group) { - t.Helper() +func TestUserAISpendStatus(t *testing.T) { + t.Parallel() + + t.Run("RequiresLicenseFeature", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + client, _ := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{}, + }, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant here; the request is blocked before RBAC. + _, err := client.UserAISpendStatus(ctx, uuid.New()) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) + + t.Run("RequiresExperiment", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + client, _ := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureAIBridge: 1, + }, + }, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant here; the request is blocked before RBAC. + _, err := client.UserAISpendStatus(ctx, uuid.New()) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) + + tests := []struct { + name string + groupBudget *int64 // nil = no group budget configured + overrideLimit *int64 // nil = no user override configured + spent int64 // 0 = no spend seeded + wantHasEffectiveGroup bool + wantSpendLimitMicros *int64 + wantLimitSource *codersdk.AIBudgetLimitSource + wantCurrentSpendMicros int64 + }{ + { + name: "NoEffectiveGroup", + wantHasEffectiveGroup: false, + wantSpendLimitMicros: nil, + wantLimitSource: nil, + wantCurrentSpendMicros: 0, + }, + { + 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), + }, + { + 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), + 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), + wantCurrentSpendMicros: 1_500_000_000, + }, + { + name: "UserOverride/ZeroSpend", + 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), + }, + { + 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), + 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), + wantCurrentSpendMicros: 350_000_000, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "spend-test-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + // Use fixed dates to keep the test deterministic. + clock.Set(time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC)) + wantPeriodStart := time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC) + wantPeriodEnd := time.Date(2026, time.April, 1, 0, 0, 0, 0, time.UTC) + + if tt.groupBudget != nil { + _, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{ + SpendLimitMicros: *tt.groupBudget, + }) + require.NoError(t, err) + } + if tt.overrideLimit != nil { + _, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{ + GroupID: group.ID, + SpendLimitMicros: *tt.overrideLimit, + }) + require.NoError(t, err) + } + if tt.spent > 0 { + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: targetUser.ID, + EffectiveGroupID: group.ID, + Day: clock.Now(), + CostMicros: tt.spent, + }) + require.NoError(t, err) + } + + got, err := adminClient.UserAISpendStatus(ctx, targetUser.ID) + require.NoError(t, err) + require.Equal(t, targetUser.ID, got.UserID) + require.Equal(t, wantPeriodStart, got.PeriodStart) + require.Equal(t, wantPeriodEnd, got.PeriodEnd) + require.Equal(t, tt.wantCurrentSpendMicros, got.CurrentSpendMicros) + + var wantEffectiveGroupID *uuid.UUID + if tt.wantHasEffectiveGroup { + wantEffectiveGroupID = &group.ID + } + require.Equal(t, wantEffectiveGroupID, got.EffectiveGroupID) + require.Equal(t, tt.wantSpendLimitMicros, got.SpendLimitMicros) + require.Equal(t, tt.wantLimitSource, got.LimitSource) + }) + } +} + +func TestUserAISpendStatusRoleAccess(t *testing.T) { + t.Parallel() dv := coderdtest.DeploymentValues(t) dv.AI.BridgeConfig.Enabled = serpent.Bool(true) @@ -2991,12 +3169,88 @@ func setupUserAIBudgetOverrideTest(t *testing.T) (adminClient *codersdk.Client, }, }, }) - adminClient, _ = coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) - _, targetUser = coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + orgAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgAdmin(owner.OrganizationID)) + orgUserAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.ScopedRoleOrgUserAdmin(owner.OrganizationID)) + memberClient, memberUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + cases := []struct { + Name string + Client *codersdk.Client + Target uuid.UUID + WantCode int + }{ + {Name: "Owner", Client: ownerClient, Target: targetUser.ID, WantCode: http.StatusOK}, + {Name: "UserAdmin", Client: userAdminClient, Target: targetUser.ID, WantCode: http.StatusOK}, + {Name: "OrgAdmin", Client: orgAdminClient, Target: targetUser.ID, WantCode: http.StatusOK}, + {Name: "OrgUserAdmin", Client: orgUserAdminClient, Target: targetUser.ID, WantCode: http.StatusOK}, + {Name: "MemberReadsSelf", Client: memberClient, Target: memberUser.ID, WantCode: http.StatusOK}, + {Name: "MemberReadsOther", Client: memberClient, Target: targetUser.ID, WantCode: http.StatusNotFound}, + } + + for _, tc := range cases { + t.Run(tc.Name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := tc.Client.UserAISpendStatus(ctx, tc.Target) + if tc.WantCode == http.StatusOK { + require.NoError(t, err) + return + } + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, tc.WantCode, sdkErr.StatusCode()) + }) + } +} + +// aiCostControlTestOptions configures the setup of an AI cost control test +// deployment. GroupName is required. Clock, Database, and Pubsub are +// optional overrides (leave nil for defaults). +type aiCostControlTestOptions struct { + GroupName string + Clock quartz.Clock + Database database.Store + Pubsub pubsub.Pubsub +} + +// setupAICostControlTest builds a deployment with FeatureAIBridge licensed +// and the AI Gateway cost control experiment enabled, creates an admin +// client and target user, adds the target user to a group, and returns +// the admin client, target user, and group. +func setupAICostControlTest(t *testing.T, opts aiCostControlTestOptions) (*codersdk.Client, codersdk.User, codersdk.Group) { + t.Helper() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + coderdOpts := &coderdtest.Options{DeploymentValues: dv} + if opts.Clock != nil { + coderdOpts.Clock = opts.Clock + } + if opts.Database != nil { + coderdOpts.Database = opts.Database + } + if opts.Pubsub != nil { + coderdOpts.Pubsub = opts.Pubsub + } + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: coderdOpts, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + }, + }, + }) + adminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + _, targetUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) ctx := testutil.Context(t, testutil.WaitLong) g, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ - Name: "override-test-group", + Name: opts.GroupName, }) require.NoError(t, err) g, err = adminClient.PatchGroup(ctx, g.ID, codersdk.PatchGroupRequest{ diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 9f3860f3ca4..4f9e9fbc970 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -641,7 +641,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Get("/", api.userQuietHoursSchedule) r.Put("/", api.putUserQuietHoursSchedule) }) - r.Route("/users/{user}/ai/budget", func(r chi.Router) { + r.Route("/users/{user}/ai", func(r chi.Router) { // AI cost controls are a paid feature (AI Governance add-on). r.Use( // TODO(AIGOV-443): remove once AI Gateway cost control functionality is stable. @@ -650,9 +650,14 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { apiKeyMiddleware, httpmw.ExtractUserParam(options.Database), ) - r.Get("/", api.userAIBudgetOverride) - r.Put("/", api.upsertUserAIBudgetOverride) - r.Delete("/", api.deleteUserAIBudgetOverride) + r.Route("/budget", func(r chi.Router) { + r.Get("/", api.userAIBudgetOverride) + r.Put("/", api.upsertUserAIBudgetOverride) + r.Delete("/", api.deleteUserAIBudgetOverride) + }) + r.Route("/spend", func(r chi.Router) { + r.Get("/", api.userAISpendStatus) + }) }) r.Route("/prebuilds", func(r chi.Router) { r.Use( diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 760353d5e88..ee8d0d44b46 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -236,6 +236,14 @@ export interface AIBridgeToolCall { readonly created_at: string; } +// From codersdk/aibridge.go +export type AIBudgetLimitSource = "group" | "user_override"; + +export const AIBudgetLimitSources: AIBudgetLimitSource[] = [ + "group", + "user_override", +]; + // From codersdk/deployment.go export type AIBudgetPeriod = "month"; @@ -9647,6 +9655,30 @@ export interface UserAIBudgetOverride { readonly updated_at: string; } +// From codersdk/aibridge.go +/** + * UserAIBudgetSummary is the effective AI budget for a user. When no + * budget applies, all fields except UserID are null. + */ +export interface UserAIBudgetSummary { + readonly user_id: string; + /** + * EffectiveGroupID is the group the spend is attributed to. Null when + * no budget applies. + */ + readonly effective_group_id: string | null; + /** + * SpendLimitMicros is the effective spend limit in micro-units. + * Null when no budget applies to the user (unlimited). + */ + readonly spend_limit_micros: number | null; + /** + * LimitSource identifies which tier produced the limit. Null when no + * budget applies. + */ + readonly limit_source: AIBudgetLimitSource | null; +} + // From codersdk/chats.go /** * UserAIProviderKeyConfig is a provider summary from the current user's @@ -9659,6 +9691,29 @@ export interface UserAIProviderKeyConfig { readonly byok_enabled: boolean; } +// From codersdk/aibridge.go +/** + * UserAISpendStatus is the current AI spend snapshot for a user within + * the active budget period. + */ +export interface UserAISpendStatus extends UserAIBudgetSummary { + /** + * CurrentSpendMicros is the user's spend on their effective group over + * the current budget period. + */ + readonly current_spend_micros: number; + /** + * PeriodStart is the inclusive lower bound of the current budget + * period. + */ + readonly period_start: string; + /** + * PeriodEnd is the exclusive upper bound of the current budget + * period. + */ + readonly period_end: string; +} + // From codersdk/insights.go /** * UserActivity shows the session time for a user.