mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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
This commit is contained in:
@@ -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.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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{
|
||||
|
||||
Reference in New Issue
Block a user