From c3895ff9c01d7334ccd42d051775784bd63fb575 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Tue, 28 Jul 2026 10:58:38 +0100 Subject: [PATCH] feat: add CSV export for AI spend data (#27491) ## Description Adds `GET /api/v2/organizations/{organization}/ai/spend/export`, returning `text/csv` with per-user, per-group, per-model, per-provider aggregated AI spend. The data is built from the raw AI Gateway token usage tables rather than the `ai_user_daily_spend` rollup, but stays consistent with it: spend is attributed through the token usage's effective group and bucketed by the token usage `created_at`, the same values the daily rollup derives from. The period defaults to the current UTC month, narrowed to the configured AI Gateway retention window when the month begins before retained data does. Explicit `period_start`/`period_end` params must be provided together, are interpreted as UTC, and may span at most 31 days. Unlike the default period, an explicit period that begins before the retention window is rejected rather than narrowed. Every row echoes the applied bounds, so a narrowed window is visible in the export. The endpoint requires organization-level admin permissions. ## Changes - Add the `ExportOrganizationAISpend` query aggregating `aibridge_token_usages` joined to `aibridge_interceptions`, scoped to the organization via the effective group, resolving the username, group name, and organization name alongside their IDs. - Add the `exportOrganizationAISpend` handler and route, gated by the `aigateway-cost-control` experiment and the `AIBridge` feature, returning the CSV in a single response. - Add the `ExportOrganizationAISpend` codersdk client method. - Require organization-wide `ResourceGroupMember` read, since the export aggregates every user in the organization. The per-row filter stays in `dbauthz` as defence in depth. - Escape leading formula characters in the free-text columns, so a model or provider name recorded from an intercepted request cannot be evaluated when the CSV is opened in a spreadsheet. - Add an index on `aibridge_token_usages (effective_group_id, created_at)`, which the period and group predicates otherwise cannot use. Closes https://linear.app/codercom/issue/AIGOV-293/add-csv-export-for-ai-spend-data > [!NOTE] > Generated by Coder Agents on behalf of @ssncferreira --- coderd/apidoc/docs.go | 47 + coderd/apidoc/swagger.json | 43 + coderd/coderdtest/swaggerparser.go | 1 + coderd/database/dbauthz/dbauthz.go | 4 + coderd/database/dbauthz/dbauthz_test.go | 16 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 + coderd/database/dump.sql | 2 + ...ge_token_usage_spend_export_index.down.sql | 1 + ...idge_token_usage_spend_export_index.up.sql | 5 + coderd/database/modelmethods.go | 4 + coderd/database/querier.go | 5 + coderd/database/queries.sql.go | 102 ++ coderd/database/queries/aicostcontrol.sql | 40 + codersdk/aibridge.go | 31 + docs/reference/api/enterprise.md | 33 + enterprise/coderd/aibridge.go | 210 +++- enterprise/coderd/aibridge_test.go | 982 ++++++++++++++++++ enterprise/coderd/coderd.go | 11 + 19 files changed, 1558 insertions(+), 2 deletions(-) create mode 100644 coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.down.sql create mode 100644 coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.up.sql diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 7be7b900dd..13eea683db 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -4788,6 +4788,53 @@ const docTemplate = `{ ] } }, + "/api/v2/organizations/{organization}/ai/spend/export": { + "get": { + "description": "Returns per-user, per-group, per-model, per-provider aggregated AI spend for the organization as CSV, built from raw AI Gateway token usage.\nThe optional period_start and period_end query parameters bound the period and are interpreted as UTC. They must be provided together and span at most 31 days. When both are omitted, the current UTC monthly period is used.\nAn explicit period_start must fall within the configured AI Gateway data retention window, since older token usage is purged. The default period is narrowed to that window instead, and every row echoes the applied bounds.\nRequires organization-level administrator permissions.", + "produces": [ + "text/csv" + ], + "tags": [ + "Enterprise" + ], + "summary": "Export organization AI spend as CSV", + "operationId": "export-organization-ai-spend-as-csv", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Organization ID", + "name": "organization", + "in": "path", + "required": true + }, + { + "type": "string", + "format": "date-time", + "description": "Inclusive lower bound (RFC3339)", + "name": "period_start", + "in": "query" + }, + { + "type": "string", + "format": "date-time", + "description": "Exclusive upper bound (RFC3339)", + "name": "period_end", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/organizations/{organization}/groups": { "get": { "produces": [ diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index eeb9d48059..a50b9dabee 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -4227,6 +4227,49 @@ ] } }, + "/api/v2/organizations/{organization}/ai/spend/export": { + "get": { + "description": "Returns per-user, per-group, per-model, per-provider aggregated AI spend for the organization as CSV, built from raw AI Gateway token usage.\nThe optional period_start and period_end query parameters bound the period and are interpreted as UTC. They must be provided together and span at most 31 days. When both are omitted, the current UTC monthly period is used.\nAn explicit period_start must fall within the configured AI Gateway data retention window, since older token usage is purged. The default period is narrowed to that window instead, and every row echoes the applied bounds.\nRequires organization-level administrator permissions.", + "produces": ["text/csv"], + "tags": ["Enterprise"], + "summary": "Export organization AI spend as CSV", + "operationId": "export-organization-ai-spend-as-csv", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Organization ID", + "name": "organization", + "in": "path", + "required": true + }, + { + "type": "string", + "format": "date-time", + "description": "Inclusive lower bound (RFC3339)", + "name": "period_start", + "in": "query" + }, + { + "type": "string", + "format": "date-time", + "description": "Exclusive upper bound (RFC3339)", + "name": "period_end", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK" + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/v2/organizations/{organization}/groups": { "get": { "produces": ["application/json"], diff --git a/coderd/coderdtest/swaggerparser.go b/coderd/coderdtest/swaggerparser.go index 00dd9d9dc7..c5cd690aaf 100644 --- a/coderd/coderdtest/swaggerparser.go +++ b/coderd/coderdtest/swaggerparser.go @@ -426,6 +426,7 @@ func assertProduce(t *testing.T, comment SwaggerComment) { (comment.router == "/api/v2/debug/tailnet" && comment.method == "get") || (comment.router == "/api/v2/workspaces/{workspace}/acl" && comment.method == "patch") || (comment.router == "/api/v2/init-script/{os}/{arch}" && comment.method == "get") || + (comment.router == "/api/v2/organizations/{organization}/ai/spend/export" && comment.method == "get") || (comment.router == "/api/v2/templatebuilder/compose" && comment.method == "post") { return // Exception: HTTP 200 is returned without response entity } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index a5418d3307..5b8036958c 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2704,6 +2704,10 @@ func (q *querier) ExpirePrebuildsAPIKeys(ctx context.Context, now time.Time) err return q.db.ExpirePrebuildsAPIKeys(ctx, now) } +func (q *querier) ExportOrganizationAISpend(ctx context.Context, arg database.ExportOrganizationAISpendParams) ([]database.ExportOrganizationAISpendRow, error) { + return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ExportOrganizationAISpend)(ctx, arg) +} + func (q *querier) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error { fetch := func(ctx context.Context, id uuid.UUID) (database.Workspace, error) { return q.db.GetWorkspaceByID(ctx, id) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 671ea09e93..ef0e45eaeb 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6992,6 +6992,22 @@ func (s *MethodTestSuite) TestAIBridge() { Returns([]database.GetGroupMembersAISpendRow{row1, row2}) })) + s.Run("ExportOrganizationAISpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + org := testutil.Fake(s.T(), faker, database.Organization{}) + row1 := testutil.Fake(s.T(), faker, database.ExportOrganizationAISpendRow{OrganizationID: org.ID}) + row2 := testutil.Fake(s.T(), faker, database.ExportOrganizationAISpendRow{OrganizationID: org.ID}) + arg := database.ExportOrganizationAISpendParams{ + OrganizationID: org.ID, + PeriodStart: time.Now().UTC().Truncate(24 * time.Hour), + PeriodEnd: time.Now().UTC(), + } + dbm.EXPECT().ExportOrganizationAISpend(gomock.Any(), arg). + Return([]database.ExportOrganizationAISpendRow{row1, row2}, nil).AnyTimes() + check.Args(arg). + Asserts(row1, policy.ActionRead, row2, policy.ActionRead). + Returns([]database.ExportOrganizationAISpendRow{row1, row2}) + })) + s.Run("GetGroupAIBudget", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { g := testutil.Fake(s.T(), faker, database.Group{}) b := testutil.Fake(s.T(), faker, database.GroupAIBudget{GroupID: g.ID}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 3e50d07714..db91a9772b 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1009,6 +1009,14 @@ func (m queryMetricsStore) ExpirePrebuildsAPIKeys(ctx context.Context, now time. return r0 } +func (m queryMetricsStore) ExportOrganizationAISpend(ctx context.Context, arg database.ExportOrganizationAISpendParams) ([]database.ExportOrganizationAISpendRow, error) { + start := time.Now() + r0, r1 := m.s.ExportOrganizationAISpend(ctx, arg) + m.queryLatencies.WithLabelValues("ExportOrganizationAISpend").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ExportOrganizationAISpend").Inc() + return r0, r1 +} + func (m queryMetricsStore) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error { start := time.Now() r0 := m.s.FavoriteWorkspace(ctx, id) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index be7a1db6a2..a038a74957 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1724,6 +1724,21 @@ func (mr *MockStoreMockRecorder) ExpirePrebuildsAPIKeys(ctx, now any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExpirePrebuildsAPIKeys", reflect.TypeOf((*MockStore)(nil).ExpirePrebuildsAPIKeys), ctx, now) } +// ExportOrganizationAISpend mocks base method. +func (m *MockStore) ExportOrganizationAISpend(ctx context.Context, arg database.ExportOrganizationAISpendParams) ([]database.ExportOrganizationAISpendRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ExportOrganizationAISpend", ctx, arg) + ret0, _ := ret[0].([]database.ExportOrganizationAISpendRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ExportOrganizationAISpend indicates an expected call of ExportOrganizationAISpend. +func (mr *MockStoreMockRecorder) ExportOrganizationAISpend(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExportOrganizationAISpend", reflect.TypeOf((*MockStore)(nil).ExportOrganizationAISpend), ctx, arg) +} + // FavoriteWorkspace mocks base method. func (m *MockStore) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 90bf283c17..5127d0159d 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -4698,6 +4698,8 @@ CREATE INDEX idx_aibridge_interceptions_thread_root_id ON aibridge_interceptions CREATE INDEX idx_aibridge_model_thoughts_interception_id ON aibridge_model_thoughts USING btree (interception_id); +CREATE INDEX idx_aibridge_token_usages_effective_group_id_created_at ON aibridge_token_usages USING btree (effective_group_id, created_at) WHERE (effective_group_id IS NOT NULL); + CREATE INDEX idx_aibridge_token_usages_interception_id ON aibridge_token_usages USING btree (interception_id); CREATE INDEX idx_aibridge_token_usages_provider_response_id ON aibridge_token_usages USING btree (provider_response_id); diff --git a/coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.down.sql b/coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.down.sql new file mode 100644 index 0000000000..4bb830a52a --- /dev/null +++ b/coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS idx_aibridge_token_usages_effective_group_id_created_at; diff --git a/coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.up.sql b/coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.up.sql new file mode 100644 index 0000000000..ffcc0fab3c --- /dev/null +++ b/coderd/database/migrations/000554_aibridge_token_usage_spend_export_index.up.sql @@ -0,0 +1,5 @@ +-- Serves spend queries that filter token usage by effective group over a time +-- range. Rows with no effective group are excluded. +CREATE INDEX idx_aibridge_token_usages_effective_group_id_created_at + ON aibridge_token_usages (effective_group_id, created_at) + WHERE effective_group_id IS NOT NULL; diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go index 3576c98276..fae247adbf 100644 --- a/coderd/database/modelmethods.go +++ b/coderd/database/modelmethods.go @@ -474,6 +474,10 @@ func (r GetGroupMembersAISpendRow) RBACObject() rbac.Object { return rbac.ResourceGroupMember.WithID(r.UserID).InOrg(r.OrganizationID).WithOwner(r.UserID.String()) } +func (r ExportOrganizationAISpendRow) RBACObject() rbac.Object { + return rbac.ResourceGroupMember.WithID(r.UserID).InOrg(r.OrganizationID).WithOwner(r.UserID.String()) +} + // PrebuiltWorkspaceResource defines the interface for types that can be identified as prebuilt workspaces // and converted to their corresponding prebuilt workspace RBAC object. type PrebuiltWorkspaceResource interface { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index fd914b2af9..f82ef6ea14 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -265,6 +265,11 @@ type sqlcQuerier interface { // Next, collect api_keys that belong to the prebuilds user but have no token name. // These were most likely created via 'coder login' as the prebuilds user. ExpirePrebuildsAPIKeys(ctx context.Context, now time.Time) error + // Returns per-user, per-group, per-model, per-provider aggregated AI spend for + // @organization_id over the [period_start, period_end) window. Spend is + // attributed through the token usage's effective group, and rows are bucketed + // by the token usage created_at, matching how ai_user_daily_spend is derived. + ExportOrganizationAISpend(ctx context.Context, arg ExportOrganizationAISpendParams) ([]ExportOrganizationAISpendRow, error) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error FetchMemoryResourceMonitorsByAgentID(ctx context.Context, agentID uuid.UUID) (WorkspaceAgentMemoryResourceMonitor, error) FetchMemoryResourceMonitorsUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]WorkspaceAgentMemoryResourceMonitor, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 79ab8e0d94..1428f39983 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2525,6 +2525,108 @@ func (q *sqlQuerier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid return i, err } +const exportOrganizationAISpend = `-- name: ExportOrganizationAISpend :many +SELECT + ai.initiator_id AS user_id, + users.username AS username, + tu.effective_group_id AS group_id, + groups.name AS group_name, + groups.organization_id AS organization_id, + organizations.name AS organization_name, + ai.model AS model, + ai.provider AS provider, + ai.provider_name AS provider_name, + COALESCE(SUM(tu.input_tokens), 0)::BIGINT AS input_tokens, + COALESCE(SUM(tu.output_tokens), 0)::BIGINT AS output_tokens, + COALESCE(SUM(tu.cache_read_input_tokens), 0)::BIGINT AS cache_read_tokens, + COALESCE(SUM(tu.cache_write_input_tokens), 0)::BIGINT AS cache_write_tokens, + COALESCE(SUM(tu.cost_micros), 0)::BIGINT AS cost_micros +FROM aibridge_token_usages tu +JOIN aibridge_interceptions ai ON ai.id = tu.interception_id +JOIN users ON users.id = ai.initiator_id +JOIN groups ON groups.id = tu.effective_group_id +JOIN organizations ON organizations.id = groups.organization_id +WHERE groups.organization_id = $1 + AND tu.created_at >= $2::timestamptz + AND tu.created_at < $3::timestamptz +GROUP BY + ai.initiator_id, + users.username, + tu.effective_group_id, + groups.name, + groups.organization_id, + organizations.name, + ai.model, + ai.provider, + ai.provider_name +ORDER BY ai.initiator_id, tu.effective_group_id, ai.provider, ai.provider_name, ai.model +` + +type ExportOrganizationAISpendParams struct { + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + PeriodStart time.Time `db:"period_start" json:"period_start"` + PeriodEnd time.Time `db:"period_end" json:"period_end"` +} + +type ExportOrganizationAISpendRow struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + Username string `db:"username" json:"username"` + GroupID uuid.NullUUID `db:"group_id" json:"group_id"` + GroupName string `db:"group_name" json:"group_name"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + OrganizationName string `db:"organization_name" json:"organization_name"` + Model string `db:"model" json:"model"` + Provider string `db:"provider" json:"provider"` + ProviderName string `db:"provider_name" json:"provider_name"` + InputTokens int64 `db:"input_tokens" json:"input_tokens"` + OutputTokens int64 `db:"output_tokens" json:"output_tokens"` + CacheReadTokens int64 `db:"cache_read_tokens" json:"cache_read_tokens"` + CacheWriteTokens int64 `db:"cache_write_tokens" json:"cache_write_tokens"` + CostMicros int64 `db:"cost_micros" json:"cost_micros"` +} + +// Returns per-user, per-group, per-model, per-provider aggregated AI spend for +// @organization_id over the [period_start, period_end) window. Spend is +// attributed through the token usage's effective group, and rows are bucketed +// by the token usage created_at, matching how ai_user_daily_spend is derived. +func (q *sqlQuerier) ExportOrganizationAISpend(ctx context.Context, arg ExportOrganizationAISpendParams) ([]ExportOrganizationAISpendRow, error) { + rows, err := q.db.QueryContext(ctx, exportOrganizationAISpend, arg.OrganizationID, arg.PeriodStart, arg.PeriodEnd) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ExportOrganizationAISpendRow + for rows.Next() { + var i ExportOrganizationAISpendRow + if err := rows.Scan( + &i.UserID, + &i.Username, + &i.GroupID, + &i.GroupName, + &i.OrganizationID, + &i.OrganizationName, + &i.Model, + &i.Provider, + &i.ProviderName, + &i.InputTokens, + &i.OutputTokens, + &i.CacheReadTokens, + &i.CacheWriteTokens, + &i.CostMicros, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at FROM ai_model_prices diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 5f65dff8b2..6d9753ff6b 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -305,3 +305,43 @@ FROM user_spend WHERE current_spend_micros >= spend_limit_micros GROUP BY effective_group_id ORDER BY effective_group_id; + +-- name: ExportOrganizationAISpend :many +-- Returns per-user, per-group, per-model, per-provider aggregated AI spend for +-- @organization_id over the [period_start, period_end) window. Spend is +-- attributed through the token usage's effective group, and rows are bucketed +-- by the token usage created_at, matching how ai_user_daily_spend is derived. +SELECT + ai.initiator_id AS user_id, + users.username AS username, + tu.effective_group_id AS group_id, + groups.name AS group_name, + groups.organization_id AS organization_id, + organizations.name AS organization_name, + ai.model AS model, + ai.provider AS provider, + ai.provider_name AS provider_name, + COALESCE(SUM(tu.input_tokens), 0)::BIGINT AS input_tokens, + COALESCE(SUM(tu.output_tokens), 0)::BIGINT AS output_tokens, + COALESCE(SUM(tu.cache_read_input_tokens), 0)::BIGINT AS cache_read_tokens, + COALESCE(SUM(tu.cache_write_input_tokens), 0)::BIGINT AS cache_write_tokens, + COALESCE(SUM(tu.cost_micros), 0)::BIGINT AS cost_micros +FROM aibridge_token_usages tu +JOIN aibridge_interceptions ai ON ai.id = tu.interception_id +JOIN users ON users.id = ai.initiator_id +JOIN groups ON groups.id = tu.effective_group_id +JOIN organizations ON organizations.id = groups.organization_id +WHERE groups.organization_id = @organization_id + AND tu.created_at >= @period_start::timestamptz + AND tu.created_at < @period_end::timestamptz +GROUP BY + ai.initiator_id, + users.username, + tu.effective_group_id, + groups.name, + groups.organization_id, + organizations.name, + ai.model, + ai.provider, + ai.provider_name +ORDER BY ai.initiator_id, tu.effective_group_id, ai.provider, ai.provider_name, ai.model; diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 34572240ab..ff3fcff369 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "strings" "time" @@ -368,6 +369,36 @@ func (c *Client) AIBridgeListClients(ctx context.Context) ([]string, error) { return clients, json.NewDecoder(res.Body).Decode(&clients) } +// ExportOrganizationAISpend returns a CSV of per-user, per-group, per-model, +// per-provider AI spend for the organization over the requested period. Both +// bounds are optional and interpreted as UTC, and zero values fall back to the +// current budget period on the server. The caller is responsible for closing +// the returned ReadCloser. +func (c *Client) ExportOrganizationAISpend(ctx context.Context, organization uuid.UUID, opts AISpendPeriodWindow) (io.ReadCloser, error) { + res, err := c.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/organizations/%s/ai/spend/export", organization.String()), + nil, + func(r *http.Request) { + q := r.URL.Query() + if !opts.PeriodStart.IsZero() { + q.Set("period_start", opts.PeriodStart.UTC().Format(time.RFC3339Nano)) + } + if !opts.PeriodEnd.IsZero() { + q.Set("period_end", opts.PeriodEnd.UTC().Format(time.RFC3339Nano)) + } + r.URL.RawQuery = q.Encode() + }, + ) + if err != nil { + return nil, xerrors.Errorf("make request: %w", err) + } + if res.StatusCode != http.StatusOK { + defer res.Body.Close() + return nil, ReadBodyAsError(res) + } + return res.Body, nil +} + type GroupAIBudget struct { GroupID uuid.UUID `json:"group_id" format:"uuid"` SpendLimitMicros int64 `json:"spend_limit_micros"` diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index eeaf8032ec..cff040db9a 100644 --- a/docs/reference/api/enterprise.md +++ b/docs/reference/api/enterprise.md @@ -1715,6 +1715,39 @@ curl -X DELETE http://coder-server:8080/api/v2/oauth2-provider/apps/{app}/secret To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Export organization AI spend as CSV + +### Code samples + +```sh +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/organizations/{organization}/ai/spend/export \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/v2/organizations/{organization}/ai/spend/export` + +Returns per-user, per-group, per-model, per-provider aggregated AI spend for the organization as CSV, built from raw AI Gateway token usage. +The optional period_start and period_end query parameters bound the period and are interpreted as UTC. They must be provided together and span at most 31 days. When both are omitted, the current UTC monthly period is used. +An explicit period_start must fall within the configured AI Gateway data retention window, since older token usage is purged. The default period is narrowed to that window instead, and every row echoes the applied bounds. +Requires organization-level administrator permissions. + +### Parameters + +| Name | In | Type | Required | Description | +|----------------|-------|-------------------|----------|---------------------------------| +| `organization` | path | string(uuid) | true | Organization ID | +| `period_start` | query | string(date-time) | false | Inclusive lower bound (RFC3339) | +| `period_end` | query | string(date-time) | false | Exclusive upper bound (RFC3339) | + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|--------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get groups by organization ### Code samples diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index f07b6bb18d..5b6f256e68 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -1,12 +1,15 @@ package coderd import ( + "bytes" "context" "database/sql" + "encoding/csv" "errors" "fmt" "net/http" "strconv" + "strings" "time" "github.com/go-chi/chi/v5" @@ -23,6 +26,8 @@ import ( "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/coderd/httpmw" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/searchquery" "github.com/coder/coder/v2/codersdk" ) @@ -39,6 +44,9 @@ const ( aiBridgeRateLimitWindow = time.Second maxOrganizationGroupsAISpendGroupIDs = 100 maxGroupMembersAISpendUserIDs = 100 + // maxAISpendExportPeriod bounds an explicit AI spend export window to at + // most 31 days, matching the maximum length of the monthly default period. + maxAISpendExportPeriod = 31 * 24 * time.Hour ) // errInvalidCursor is returned when a pagination cursor does not @@ -910,8 +918,8 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { slog.F("period_end", periodWindow.End), ) - policy := codersdk.NewAIBudgetPolicyFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPolicy) - effectiveGroup, ok, err := budget.ResolveUserEffectiveGroup(ctx, api.Database, user.ID, policy) + budgetPolicy := codersdk.NewAIBudgetPolicyFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPolicy) + effectiveGroup, ok, err := budget.ResolveUserEffectiveGroup(ctx, api.Database, user.ID, budgetPolicy) if err != nil { logger.Error(ctx, "failed to resolve user AI budget", slog.Error(err)) httpapi.InternalServerError(rw, err) @@ -1026,6 +1034,204 @@ func (api *API) organizationGroupsAISpend(rw http.ResponseWriter, r *http.Reques httpapi.Write(ctx, rw, http.StatusOK, resp) } +// AISpendExportCSVHeader is the CSV column order for the AI spend export. +var AISpendExportCSVHeader = []string{ + "user_id", "username", "group_id", "group_name", "organization_id", "organization_name", + "model", "provider", "provider_name", + "input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens", + "cost_micros", "period_start", "period_end", +} + +// csvFormulaPrefixes are the leading characters a spreadsheet treats as the +// start of a formula rather than text. +const csvFormulaPrefixes = "=+-@\t\r" + +// escapeCSVCell prefixes a leading formula character with a single quote, which +// spreadsheets strip on display, so the value renders as its original text +// instead of being evaluated. +func escapeCSVCell(value string) string { + if value == "" || !strings.ContainsRune(csvFormulaPrefixes, rune(value[0])) { + return value + } + return "'" + value +} + +// aiSpendExportPeriod resolves the export window from the request. When neither +// start nor end is supplied it defaults to the current UTC monthly budget +// period, narrowed to the retention window. Both bounds must be supplied +// together and are interpreted as UTC, and an explicit window must be non-empty, +// span at most 31 days, and begin within the retention window. On invalid input +// it writes the error response and returns ok=false. +func (api *API) aiSpendExportPeriod(ctx context.Context, rw http.ResponseWriter, r *http.Request) (start, end time.Time, ok bool) { + query := r.URL.Query() + hasStart := query.Has("period_start") + hasEnd := query.Has("period_end") + + // retentionStart is the oldest token usage still available, since anything + // older has been purged. A retention of zero disables purging. + retention := api.DeploymentValues.AI.BridgeConfig.Retention.Value() + hasRetention := retention > 0 + var retentionStart time.Time + if hasRetention { + retentionStart = api.Clock.Now().Add(-retention) + } + + switch { + case !hasStart && !hasEnd: + // No period was requested, so start at the budget period or the + // retention window, whichever is later. + window, err := api.currentAIBudgetWindow() + if err != nil { + api.Logger.Error(ctx, "failed to compute AI budget period", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return time.Time{}, time.Time{}, false + } + start, end = window.Start, window.End + if hasRetention && start.Before(retentionStart) { + start = retentionStart + } + case hasStart != hasEnd: + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Query parameters \"period_start\" and \"period_end\" must be provided together.", + }) + return time.Time{}, time.Time{}, false + default: + // The caller asked for this period, so validate it. + parser := httpapi.NewQueryParamParser() + start = parser.Time3339Nano(query, time.Time{}, "period_start") + end = parser.Time3339Nano(query, time.Time{}, "period_end") + parser.ErrorExcessParams(query) + if len(parser.Errors) > 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Query parameters have invalid values.", + Validations: parser.Errors, + }) + return time.Time{}, time.Time{}, false + } + if !start.Before(end) { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Query parameter \"period_start\" must be before \"period_end\".", + }) + return time.Time{}, time.Time{}, false + } + if end.Sub(start) > maxAISpendExportPeriod { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Query period must not exceed 31 days.", + }) + return time.Time{}, time.Time{}, false + } + // Fail if the period starts before the oldest retained data + if hasRetention && start.Before(retentionStart) { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: fmt.Sprintf("Query parameter \"period_start\" is older than the configured AI Gateway data retention window (%s).", retention), + }) + return time.Time{}, time.Time{}, false + } + } + + return start, end, true +} + +// @Summary Export organization AI spend as CSV +// @Description Returns per-user, per-group, per-model, per-provider aggregated AI spend for the organization as CSV, built from raw AI Gateway token usage. +// @Description The optional period_start and period_end query parameters bound the period and are interpreted as UTC. They must be provided together and span at most 31 days. When both are omitted, the current UTC monthly period is used. +// @Description An explicit period_start must fall within the configured AI Gateway data retention window, since older token usage is purged. The default period is narrowed to that window instead, and every row echoes the applied bounds. +// @Description Requires organization-level administrator permissions. +// @ID export-organization-ai-spend-as-csv +// @Security CoderSessionToken +// @Produce text/csv +// @Tags Enterprise +// @Param organization path string true "Organization ID" format(uuid) +// @Param period_start query string false "Inclusive lower bound (RFC3339)" format(date-time) +// @Param period_end query string false "Exclusive upper bound (RFC3339)" format(date-time) +// @Success 200 +// @Router /api/v2/organizations/{organization}/ai/spend/export [get] +func (api *API) exportOrganizationAISpend(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + org := httpmw.OrganizationParam(r) + logger := api.Logger.With(slog.F("organization_id", org.ID)) + + // The export aggregates the whole organization, so require organization-wide + // read rather than letting the per-row filter narrow it to the caller. + if !api.Authorize(r, policy.ActionRead, rbac.ResourceGroupMember.InOrg(org.ID)) { + httpapi.Forbidden(rw) + return + } + + periodStart, periodEnd, ok := api.aiSpendExportPeriod(ctx, rw, r) + if !ok { + return + } + logger = logger.With( + slog.F("period_start", periodStart), + slog.F("period_end", periodEnd), + ) + + rows, err := api.Database.ExportOrganizationAISpend(ctx, database.ExportOrganizationAISpendParams{ + OrganizationID: org.ID, + PeriodStart: periodStart, + PeriodEnd: periodEnd, + }) + if err != nil { + logger.Error(ctx, "failed to export organization AI spend", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + start := periodStart.UTC().Format(time.RFC3339) + end := periodEnd.UTC().Format(time.RFC3339) + + var buf bytes.Buffer + cw := csv.NewWriter(&buf) + if err := cw.Write(AISpendExportCSVHeader); err != nil { + logger.Error(ctx, "failed to write AI spend export header", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + for _, row := range rows { + if err := cw.Write([]string{ + row.UserID.String(), + escapeCSVCell(row.Username), + row.GroupID.UUID.String(), + escapeCSVCell(row.GroupName), + row.OrganizationID.String(), + escapeCSVCell(row.OrganizationName), + escapeCSVCell(row.Model), + escapeCSVCell(row.Provider), + escapeCSVCell(row.ProviderName), + strconv.FormatInt(row.InputTokens, 10), + strconv.FormatInt(row.OutputTokens, 10), + strconv.FormatInt(row.CacheReadTokens, 10), + strconv.FormatInt(row.CacheWriteTokens, 10), + strconv.FormatInt(row.CostMicros, 10), + start, + end, + }); err != nil { + logger.Error(ctx, "failed to write AI spend export row", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + } + cw.Flush() + if err := cw.Error(); err != nil { + logger.Error(ctx, "failed to build AI spend export", slog.Error(err)) + httpapi.InternalServerError(rw, err) + return + } + + // Name the file after the organization and period so separate exports stay + // distinguishable once downloaded. + filename := fmt.Sprintf("ai-spend-export-%s-%s-to-%s.csv", + org.Name, periodStart.UTC().Format(time.DateOnly), periodEnd.UTC().Format(time.DateOnly)) + rw.Header().Set("Content-Type", "text/csv; charset=utf-8") + rw.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", filename)) + rw.Header().Set("Content-Length", strconv.Itoa(buf.Len())) + rw.WriteHeader(http.StatusOK) + if _, err := rw.Write(buf.Bytes()); err != nil { + logger.Error(ctx, "failed to write AI spend export", slog.Error(err)) + } +} + // @Summary Get group members AI spend by organization // @Description Returns aggregate AI spend attributed to the group per requested user. // @Description A maximum of 100 user IDs may be requested per call, and requests with more are rejected, so callers are expected to batch across multiple requests. diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index ca270ece8d..38e467f0c6 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -1,10 +1,15 @@ package coderd_test import ( + "bytes" + "context" "database/sql" + "encoding/csv" "encoding/json" + "fmt" "io" "net/http" + "strconv" "testing" "time" @@ -26,6 +31,7 @@ import ( "github.com/coder/coder/v2/codersdk" entaudit "github.com/coder/coder/v2/enterprise/audit" "github.com/coder/coder/v2/enterprise/audit/backends" + entcoderd "github.com/coder/coder/v2/enterprise/coderd" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" "github.com/coder/coder/v2/enterprise/coderd/license" "github.com/coder/coder/v2/testutil" @@ -3693,6 +3699,976 @@ func TestOrganizationGroupsAISpendRoleAccess(t *testing.T) { } } +// readAISpendExportCSV parses a CSV export body into its records. +func readAISpendExportCSV(t *testing.T, body io.Reader) [][]string { + t.Helper() + records, err := csv.NewReader(body).ReadAll() + require.NoError(t, err) + return records +} + +// readAISpendExportResponse asserts the response is sent once with an accurate +// Content-Length and returns the parsed CSV records. +func readAISpendExportResponse(t *testing.T, res *http.Response) [][]string { + t.Helper() + body, err := io.ReadAll(res.Body) + require.NoError(t, err) + require.Equal(t, strconv.Itoa(len(body)), res.Header.Get("Content-Length")) + return readAISpendExportCSV(t, bytes.NewReader(body)) +} + +// requestAISpendExport issues a raw export request so callers can inspect the +// status code, headers, and CSV body directly. +func requestAISpendExport(ctx context.Context, t *testing.T, client *codersdk.Client, orgID uuid.UUID, params map[string]string) *http.Response { + t.Helper() + res, err := client.Request(ctx, http.MethodGet, + fmt.Sprintf("/api/v2/organizations/%s/ai/spend/export", orgID), + nil, + func(r *http.Request) { + q := r.URL.Query() + for k, v := range params { + q.Set(k, v) + } + r.URL.RawQuery = q.Encode() + }, + ) + require.NoError(t, err) + return res +} + +func TestExportOrganizationAISpend(t *testing.T) { + t.Parallel() + + t.Run("Enablement", func(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + experiments []string + features license.Features + wantMsgContains string + }{ + { + name: "RequiresLicenseFeature", + experiments: []string{string(codersdk.ExperimentAIGatewayCostControl)}, + features: license.Features{}, + wantMsgContains: "AI Gateway is a Premium feature", + }, + { + name: "RequiresExperiment", + experiments: nil, + features: license.Features{codersdk.FeatureAIBridge: 1}, + wantMsgContains: "ai-gateway-cost-control", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + if len(tc.experiments) > 0 { + dv.Experiments = tc.experiments + } + client, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{Features: tc.features}, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant because the request is blocked before RBAC. + _, err := client.ExportOrganizationAISpend(ctx, owner.OrganizationID, codersdk.AISpendPeriodWindow{}) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Message, tc.wantMsgContains) + }) + } + }) + + t.Run("PeriodValidation", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-period-validation-group", + Clock: clock, + }) + start := time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC) + // Older than the default 60d retention window relative to the clock. + beforeRetention := time.Date(2025, time.December, 1, 0, 0, 0, 0, time.UTC) + + // wantMsgContains pins the branch each case exercises, since every + // rejection returns 400 and would otherwise be indistinguishable. + cases := []struct { + name string + params map[string]string + wantStatus int + wantMsgContains string + }{ + { + name: "OnlyStart", + params: map[string]string{"period_start": start.Format(time.RFC3339Nano)}, + wantStatus: http.StatusBadRequest, + wantMsgContains: "must be provided together", + }, + { + name: "OnlyEnd", + params: map[string]string{"period_end": start.Format(time.RFC3339Nano)}, + wantStatus: http.StatusBadRequest, + wantMsgContains: "must be provided together", + }, + { + name: "StartEqualsEnd", + params: map[string]string{"period_start": start.Format(time.RFC3339Nano), "period_end": start.Format(time.RFC3339Nano)}, + wantStatus: http.StatusBadRequest, + wantMsgContains: `"period_start" must be before "period_end"`, + }, + { + name: "InvalidFormat", + params: map[string]string{"period_start": "not-a-date", "period_end": start.AddDate(0, 0, 1).Format(time.RFC3339Nano)}, + wantStatus: http.StatusBadRequest, + wantMsgContains: "have invalid values", + }, + { + name: "PeriodTooLong", + params: map[string]string{"period_start": start.Format(time.RFC3339Nano), "period_end": start.AddDate(0, 0, 32).Format(time.RFC3339Nano)}, + wantStatus: http.StatusBadRequest, + wantMsgContains: "must not exceed 31 days", + }, + { + name: "MaxPeriodAllowed", + params: map[string]string{"period_start": start.Format(time.RFC3339Nano), "period_end": start.AddDate(0, 0, 31).Format(time.RFC3339Nano)}, + wantStatus: http.StatusOK, + }, + { + // period_start predates the retention window, so the raw + // token usage would be purged and results incomplete. + name: "BeforeRetentionWindow", + params: map[string]string{"period_start": beforeRetention.Format(time.RFC3339Nano), "period_end": beforeRetention.AddDate(0, 0, 1).Format(time.RFC3339Nano)}, + wantStatus: http.StatusBadRequest, + wantMsgContains: "retention window", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, tc.params) + defer res.Body.Close() + require.Equal(t, tc.wantStatus, res.StatusCode) + if tc.wantMsgContains == "" { + return + } + var sdkErr *codersdk.Error + require.ErrorAs(t, codersdk.ReadBodyAsError(res), &sdkErr) + require.Contains(t, sdkErr.Message, tc.wantMsgContains) + }) + } + }) + + t.Run("DefaultsToCurrentMonthAndAggregates", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-default-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + groupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + // Two claude-4 interceptions for the same user aggregate into one row. + for _, tu := range []database.InsertAIBridgeTokenUsageParams{ + {InputTokens: 100, OutputTokens: 50, CacheReadInputTokens: 10, CacheWriteInputTokens: 5, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}}, + {InputTokens: 200, OutputTokens: 100, CacheReadInputTokens: 20, CacheWriteInputTokens: 10, CostMicros: sql.NullInt64{Int64: 2000, Valid: true}}, + } { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: inMonth, + }, nil) + tu.InterceptionID = intc.ID + tu.CreatedAt = inMonth + tu.EffectiveGroupID = groupID + dbgen.AIBridgeTokenUsage(t, db, tu) + } + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Equal(t, entcoderd.AISpendExportCSVHeader, records[0]) + require.Len(t, records, 2) // header + single aggregated row + + // The default window echoes the current UTC month. + require.Equal(t, []string{ + targetUser.ID.String(), targetUser.Username, + group.ID.String(), group.Name, + group.OrganizationID.String(), group.OrganizationName, + "claude-4", "anthropic", "anthropic-prod", "300", "150", "30", "15", "3000", + "2026-03-01T00:00:00Z", "2026-04-01T00:00:00Z", + }, records[1]) + }) + + t.Run("DefaultPeriodClampedToRetention", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 20, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + retention := 14 * 24 * time.Hour + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-retention-clamp-group", + Clock: clock, + Database: db, + Pubsub: ps, + Retention: &retention, + }) + ctx := testutil.Context(t, testutil.WaitLong) + // Retention starts on 6 March 2026 12:00 UTC. + retentionStart := now.Add(-retention) + groupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + // Usage before the retention start falls outside the narrowed period. + for _, seed := range []struct { + at time.Time + inputTokens int64 + outputTokens int64 + costMicros int64 + }{ + {at: retentionStart.Add(-time.Hour), inputTokens: 999, outputTokens: 999, costMicros: 9999}, + {at: retentionStart.Add(time.Hour), inputTokens: 100, outputTokens: 50, costMicros: 1000}, + } { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: seed.at, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: seed.at, EffectiveGroupID: groupID, + InputTokens: seed.inputTokens, OutputTokens: seed.outputTokens, + CostMicros: sql.NullInt64{Int64: seed.costMicros, Valid: true}, + }) + } + + // Start: 6 March 2026 12:00 UTC (inclusive). + // End: 1 April 2026 00:00 UTC (exclusive). + // Now: 20 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 2) // header + only the retained row + require.Equal(t, []string{ + targetUser.ID.String(), targetUser.Username, + group.ID.String(), group.Name, + group.OrganizationID.String(), group.OrganizationName, + "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", + "2026-03-06T12:00:00Z", "2026-04-01T00:00:00Z", + }, records[1]) + }) + + t.Run("CustomPeriodHalfOpen", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-custom-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + effectiveGroupID := uuid.NullUUID{UUID: group.ID, Valid: true} + start := time.Date(2026, time.March, 10, 0, 0, 0, 0, time.UTC) + end := time.Date(2026, time.March, 11, 0, 0, 0, 0, time.UTC) + + // Usage at start is included, usage at end is excluded. + for _, at := range []time.Time{start, end} { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: at, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: at, EffectiveGroupID: effectiveGroupID, + InputTokens: 100, OutputTokens: 50, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + } + + // Start: 10 March 2026 00:00 UTC (inclusive). + // End: 11 March 2026 00:00 UTC (exclusive). + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, map[string]string{ + "period_start": start.Format(time.RFC3339Nano), + "period_end": end.Format(time.RFC3339Nano), + }) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 2) // header + only the start-boundary row + require.Equal(t, []string{ + targetUser.ID.String(), targetUser.Username, + group.ID.String(), group.Name, + group.OrganizationID.String(), group.OrganizationName, + "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", + start.Format(time.RFC3339), end.Format(time.RFC3339), + }, records[1]) + }) + + t.Run("ExplicitPeriodBeforeRetentionRejected", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 20, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + retention := 14 * 24 * time.Hour + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-retention-reject-group", + Clock: clock, + Retention: &retention, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + // The requested start is an hour before the retention window begins, so + // the request fails instead of being shortened like the default period. + start := now.Add(-retention).Add(-time.Hour) + + _, err := adminClient.ExportOrganizationAISpend(ctx, group.OrganizationID, codersdk.AISpendPeriodWindow{ + PeriodStart: start, + PeriodEnd: now, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Message, "retention window") + require.Contains(t, sdkErr.Message, retention.String()) + }) + + t.Run("RetentionDisabledAllowsOldPeriod", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 20, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + // A retention of zero disables purging, so no period is too old to + // export and neither the narrowing nor the rejection applies. + retention := time.Duration(0) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-retention-disabled-group", + Clock: clock, + Database: db, + Pubsub: ps, + Retention: &retention, + }) + ctx := testutil.Context(t, testutil.WaitLong) + // Two years before the request, far outside any retention window. + at := time.Date(2024, time.March, 10, 8, 0, 0, 0, time.UTC) + effectiveGroupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: at, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: at, EffectiveGroupID: effectiveGroupID, + InputTokens: 100, OutputTokens: 50, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + + // Start: 10 March 2024 07:00 UTC (inclusive). + // End: 10 March 2024 09:00 UTC (exclusive). + // Now: 20 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, map[string]string{ + "period_start": at.Add(-time.Hour).Format(time.RFC3339Nano), + "period_end": at.Add(time.Hour).Format(time.RFC3339Nano), + }) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 2) // header + the retained row + require.Equal(t, []string{ + targetUser.ID.String(), targetUser.Username, + group.ID.String(), group.Name, + group.OrganizationID.String(), group.OrganizationName, + "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", + "2024-03-10T07:00:00Z", "2024-03-10T09:00:00Z", + }, records[1]) + }) + + t.Run("SeparateRowPerModel", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-per-model-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + effectiveGroupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + // Usage for two different models produces one row each. + claudeIntc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: claudeIntc.ID, CreatedAt: inMonth, EffectiveGroupID: effectiveGroupID, + InputTokens: 100, OutputTokens: 50, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + gptIntc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "openai", ProviderName: "openai-prod", Model: "gpt-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: gptIntc.ID, CreatedAt: inMonth, EffectiveGroupID: effectiveGroupID, + InputTokens: 500, OutputTokens: 250, CostMicros: sql.NullInt64{Int64: 5000, Valid: true}, + }) + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 3) // header + one row per model + + userID := targetUser.ID.String() + username := targetUser.Username + groupID := group.ID.String() + groupName := group.Name + orgID := group.OrganizationID.String() + orgName := group.OrganizationName + periodStart := "2026-03-01T00:00:00Z" + periodEnd := "2026-04-01T00:00:00Z" + // Ordered by provider then model: anthropic/claude-4, then openai/gpt-4. + require.Equal(t, []string{userID, username, groupID, groupName, orgID, orgName, "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", periodStart, periodEnd}, records[1]) + require.Equal(t, []string{userID, username, groupID, groupName, orgID, orgName, "gpt-4", "openai", "openai-prod", "500", "250", "0", "0", "5000", periodStart, periodEnd}, records[2]) + }) + + t.Run("SeparateRowPerProviderName", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-per-provider-name-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + effectiveGroupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + // Two configurations of the same provider, same model. Spend is reported + // per configuration rather than merged into one provider row. + for _, seed := range []struct { + providerName string + inputTokens int64 + outputTokens int64 + costMicros int64 + }{ + {providerName: "anthropic-dev", inputTokens: 100, outputTokens: 50, costMicros: 1000}, + {providerName: "anthropic-prod", inputTokens: 500, outputTokens: 250, costMicros: 5000}, + } { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: seed.providerName, Model: "claude-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: inMonth, EffectiveGroupID: effectiveGroupID, + InputTokens: seed.inputTokens, OutputTokens: seed.outputTokens, + CostMicros: sql.NullInt64{Int64: seed.costMicros, Valid: true}, + }) + } + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 3) // header + one row per provider name + + userID := targetUser.ID.String() + username := targetUser.Username + groupID := group.ID.String() + groupName := group.Name + orgID := group.OrganizationID.String() + orgName := group.OrganizationName + periodStart := "2026-03-01T00:00:00Z" + periodEnd := "2026-04-01T00:00:00Z" + // Ordered by provider name: anthropic-dev, then anthropic-prod. + require.Equal(t, []string{userID, username, groupID, groupName, orgID, orgName, "claude-4", "anthropic", "anthropic-dev", "100", "50", "0", "0", "1000", periodStart, periodEnd}, records[1]) + require.Equal(t, []string{userID, username, groupID, groupName, orgID, orgName, "claude-4", "anthropic", "anthropic-prod", "500", "250", "0", "0", "5000", periodStart, periodEnd}, records[2]) + }) + + t.Run("SeparateRowPerGroup", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-per-group-first", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + + // A second group in the same organization. The user is not added to it: + // the effective group is snapshotted on each token usage, so spend from + // before a membership or budget change stays attributed to the old group. + secondGroup, err := adminClient.CreateGroup(ctx, group.OrganizationID, codersdk.CreateGroupRequest{ + Name: "export-per-group-second", + }) + require.NoError(t, err) + + for _, seed := range []struct { + groupID uuid.UUID + inputTokens int64 + outputTokens int64 + costMicros int64 + }{ + {groupID: group.ID, inputTokens: 100, outputTokens: 50, costMicros: 1000}, + {groupID: secondGroup.ID, inputTokens: 500, outputTokens: 250, costMicros: 5000}, + } { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: inMonth, + EffectiveGroupID: uuid.NullUUID{UUID: seed.groupID, Valid: true}, + InputTokens: seed.inputTokens, OutputTokens: seed.outputTokens, + CostMicros: sql.NullInt64{Int64: seed.costMicros, Valid: true}, + }) + } + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 3) // header + one row per group + + userID := targetUser.ID.String() + username := targetUser.Username + orgID := group.OrganizationID.String() + orgName := group.OrganizationName + periodStart := "2026-03-01T00:00:00Z" + periodEnd := "2026-04-01T00:00:00Z" + // Rows are ordered by group ID, which is a random UUID, so compare + // without depending on which group sorts first. + require.ElementsMatch(t, [][]string{ + {userID, username, group.ID.String(), group.Name, orgID, orgName, "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", periodStart, periodEnd}, + {userID, username, secondGroup.ID.String(), secondGroup.Name, orgID, orgName, "claude-4", "anthropic", "anthropic-prod", "500", "250", "0", "0", "5000", periodStart, periodEnd}, + }, records[1:]) + }) + + t.Run("PreviousMonthExcluded", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-prev-month-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + prevMonth := time.Date(2026, time.February, 20, 8, 0, 0, 0, time.UTC) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + groupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + // Previous-month usage falls outside the default window and is excluded. + prevIntc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: prevMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: prevIntc.ID, CreatedAt: prevMonth, EffectiveGroupID: groupID, + InputTokens: 999, OutputTokens: 999, CostMicros: sql.NullInt64{Int64: 9999, Valid: true}, + }) + + // Current-month usage is included. + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: inMonth, EffectiveGroupID: groupID, + InputTokens: 100, OutputTokens: 50, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + // Only the current-month usage is present, so the row carries none of the + // previous month's tokens or cost. + require.Len(t, records, 2) // header + current-month row + require.Equal(t, []string{ + targetUser.ID.String(), targetUser.Username, + group.ID.String(), group.Name, + group.OrganizationID.String(), group.OrganizationName, + "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", + "2026-03-01T00:00:00Z", "2026-04-01T00:00:00Z", + }, records[1]) + }) + + t.Run("ExcludesNullEffectiveGroup", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-null-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + at := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + + // Usage with no effective group cannot be attributed to an organization, + // so it is excluded entirely and the export returns no rows at all. + nullIntc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: at, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: nullIntc.ID, CreatedAt: at, + InputTokens: 500, CostMicros: sql.NullInt64{Int64: 5000, Valid: true}, + }) + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Equal(t, entcoderd.AISpendExportCSVHeader, records[0]) + require.Len(t, records, 1) // header only + }) + + t.Run("ExcludesOtherOrganizations", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + // Built inline rather than through setupAICostControlTest, which + // licenses a single organization and returns no owner client. + db, ps := dbtestutil.NewDB(t) + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv, Database: db, Pubsub: ps, Clock: clock}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + ctx := testutil.Context(t, testutil.WaitLong) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + + otherOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{}) + _, otherOrgMember := coderdtest.CreateAnotherUser(t, ownerClient, otherOrg.ID) + userAdminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin()) + group, err := userAdminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "export-org-scope-group", + }) + require.NoError(t, err) + otherOrgGroup, err := userAdminClient.CreateGroup(ctx, otherOrg.ID, codersdk.CreateGroupRequest{ + Name: "export-org-scope-other-org-group", + }) + require.NoError(t, err) + + // Usage in each organization, attributed through that organization's + // group. + for _, seed := range []struct { + initiator uuid.UUID + groupID uuid.UUID + }{ + {initiator: owner.UserID, groupID: group.ID}, + {initiator: otherOrgMember.ID, groupID: otherOrgGroup.ID}, + } { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: seed.initiator, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: inMonth, + EffectiveGroupID: uuid.NullUUID{UUID: seed.groupID, Valid: true}, + InputTokens: 100, OutputTokens: 50, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + } + + // The owner can read group members in both organizations, so only the + // query's organization filter keeps the other organization out. + // Now: 15 March 2026 12:00 UTC. + //nolint:gocritic // The owner is required to rule out RBAC filtering. + res := requestAISpendExport(ctx, t, ownerClient, owner.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 2) // header + the requested organization's row + require.Equal(t, []string{ + owner.UserID.String(), coderdtest.FirstUserParams.Username, + group.ID.String(), group.Name, + owner.OrganizationID.String(), group.OrganizationName, + "claude-4", "anthropic", "anthropic-prod", "100", "50", "0", "0", "1000", + "2026-03-01T00:00:00Z", "2026-04-01T00:00:00Z", + }, records[1]) + }) + + t.Run("EscapesFormulaCells", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-formula-escape-group", + Clock: clock, + Database: db, + Pubsub: ps, + }) + ctx := testutil.Context(t, testutil.WaitLong) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + groupID := uuid.NullUUID{UUID: group.ID, Valid: true} + + // Model, provider, and provider name are recorded verbatim from the + // intercepted request, so a leading formula character must be escaped. + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: targetUser.ID, + Provider: "+openai", + ProviderName: "@prod", + Model: `=HYPERLINK("http://insecure/","invoice")`, + StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, CreatedAt: inMonth, EffectiveGroupID: groupID, + InputTokens: 100, OutputTokens: 50, CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, nil) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + records := readAISpendExportResponse(t, res) + require.Len(t, records, 2) // header + the escaped row + require.Equal(t, []string{ + targetUser.ID.String(), targetUser.Username, + group.ID.String(), group.Name, + group.OrganizationID.String(), group.OrganizationName, + `'=HYPERLINK("http://insecure/","invoice")`, "'+openai", "'@prod", + "100", "50", "0", "0", "1000", + "2026-03-01T00:00:00Z", "2026-04-01T00:00:00Z", + }, records[1]) + }) + + t.Run("DownloadHeaders", func(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + adminClient, _, group := setupAICostControlTest(t, aiCostControlTestOptions{ + GroupName: "export-download-headers-group", + Clock: clock, + }) + ctx := testutil.Context(t, testutil.WaitLong) + + // Start: 10 March 2026 00:00 UTC (inclusive). + // End: 11 March 2026 00:00 UTC (exclusive). + // Now: 15 March 2026 12:00 UTC. + res := requestAISpendExport(ctx, t, adminClient, group.OrganizationID, map[string]string{ + "period_start": time.Date(2026, time.March, 10, 0, 0, 0, 0, time.UTC).Format(time.RFC3339Nano), + "period_end": time.Date(2026, time.March, 11, 0, 0, 0, 0, time.UTC).Format(time.RFC3339Nano), + }) + defer res.Body.Close() + require.Equal(t, http.StatusOK, res.StatusCode) + + require.Equal(t, "text/csv; charset=utf-8", res.Header.Get("Content-Type")) + // The filename carries the organization name and the exported period. + require.Equal(t, + fmt.Sprintf(`attachment; filename="ai-spend-export-%s-2026-03-10-to-2026-03-11.csv"`, group.OrganizationName), + res.Header.Get("Content-Disposition")) + + // No usage is seeded, so only the column header is written. Spelled out + // rather than compared against the handler's own variable, since the + // column names are a published contract that a rename would break. + records := readAISpendExportResponse(t, res) + require.Equal(t, []string{ + "user_id", "username", "group_id", "group_name", "organization_id", "organization_name", + "model", "provider", "provider_name", + "input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens", + "cost_micros", "period_start", "period_end", + }, records[0]) + require.Len(t, records, 1) + }) +} + +func TestExportOrganizationAISpendRoleAccess(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. Seeding at time.Now() + // against the default month period fails when a run crosses a UTC month + // boundary. + now := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + inMonth := time.Date(2026, time.March, 10, 8, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + clock.Set(now) + + db, ps := dbtestutil.NewDB(t) + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv, Database: db, Pubsub: ps, Clock: clock}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureAIBridge: 1, + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + 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, member := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID) + + otherOrg := coderdenttest.CreateOrganization(t, ownerClient, coderdenttest.CreateOrganizationOptions{}) + otherOrgMemberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, otherOrg.ID) + + ctx := testutil.Context(t, testutil.WaitLong) + group, err := userAdminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{ + Name: "export-role-access-group", + }) + require.NoError(t, err) + + // Seed spend for two users in the current month: the owner and a regular + // member. Both are attributed to the group. + for _, initiator := range []uuid.UUID{owner.UserID, member.ID} { + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: initiator, Provider: "anthropic", ProviderName: "anthropic-prod", Model: "claude-4", StartedAt: inMonth, + }, nil) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: intc.ID, + CreatedAt: inMonth, + EffectiveGroupID: uuid.NullUUID{UUID: group.ID, Valid: true}, + InputTokens: 100, + OutputTokens: 50, + CostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + } + + cases := []struct { + name string + client *codersdk.Client + wantUserIDs []string // expected user_id column values when wantStatus is unset + wantStatus int // non-zero means the request is rejected with this status + }{ + // Admins see every user's rows. + {name: "Owner", client: ownerClient, wantUserIDs: []string{owner.UserID.String(), member.ID.String()}}, + {name: "UserAdmin", client: userAdminClient, wantUserIDs: []string{owner.UserID.String(), member.ID.String()}}, + {name: "OrgAdmin", client: orgAdminClient, wantUserIDs: []string{owner.UserID.String(), member.ID.String()}}, + {name: "OrgUserAdmin", client: orgUserAdminClient, wantUserIDs: []string{owner.UserID.String(), member.ID.String()}}, + // The export covers the whole organization, so a regular member is + // rejected rather than served their own rows. + {name: "Member", client: memberClient, wantStatus: http.StatusForbidden}, + // A member of another org cannot read this org at all, so it fails + // earlier, when the organization is resolved. + {name: "OtherOrgMember", client: otherOrgMemberClient, wantStatus: http.StatusNotFound}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + body, err := tc.client.ExportOrganizationAISpend(ctx, owner.OrganizationID, codersdk.AISpendPeriodWindow{}) + if tc.wantStatus != 0 { + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, tc.wantStatus, sdkErr.StatusCode()) + return + } + require.NoError(t, err) + defer body.Close() + + records := readAISpendExportCSV(t, body) + var gotUserIDs []string + for _, row := range records[1:] { + gotUserIDs = append(gotUserIDs, row[0]) + } + require.ElementsMatch(t, tc.wantUserIDs, gotUserIDs) + }) + } +} + func TestGroupMembersAISpend(t *testing.T) { t.Parallel() @@ -4214,6 +5190,9 @@ type aiCostControlTestOptions struct { Clock quartz.Clock Database database.Store Pubsub pubsub.Pubsub + // Retention overrides the AI Gateway data retention duration. Nil leaves the + // deployment default in place, and zero disables purging. + Retention *time.Duration } // setupAICostControlTest builds a deployment with FeatureAIBridge licensed @@ -4226,6 +5205,9 @@ func setupAICostControlTest(t *testing.T, opts aiCostControlTestOptions) (*coder dv := coderdtest.DeploymentValues(t) dv.AI.BridgeConfig.Enabled = serpent.Bool(true) dv.Experiments = []string{string(codersdk.ExperimentAIGatewayCostControl)} + if opts.Retention != nil { + dv.AI.BridgeConfig.Retention = serpent.Duration(*opts.Retention) + } coderdOpts := &coderdtest.Options{DeploymentValues: dv} if opts.Clock != nil { coderdOpts.Clock = opts.Clock diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 857161b4ed..0a6118214d 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -530,6 +530,17 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { }) }) }) + r.Route("/organizations/{organization}/ai/spend", func(r chi.Router) { + // AI cost controls are a paid feature (AI Governance add-on). + r.Use( + apiKeyMiddleware, + httpmw.ExtractOrganizationParam(api.Database), + // TODO(AIGOV-443): remove once AI Gateway cost control functionality is stable. + httpmw.RequireExperiment(api.AGPL.Experiments, codersdk.ExperimentAIGatewayCostControl), + api.RequireFeatureMW(codersdk.FeatureAIBridge), + ) + r.Get("/export", api.exportOrganizationAISpend) + }) r.Route("/provisionerkeys", func(r chi.Router) { r.Use( httpmw.ExtractProvisionerDaemonAuthenticated(httpmw.ExtractProvisionerAuthConfig{