feat(coderd): enforce ai budget on pre-request path (#26915)

## Description

Adds pre-request AI budget enforcement to `aibridged`. Requests are rejected with HTTP 403 when the user's aggregated spend for the current period has reached their effective limit.

## Changes

- Add `IsBudgetExceeded` RPC to `aibridgedserver`. Resolves the user's effective budget, aggregates spend over the caller-supplied `[period_start, now]` window, and returns whether the limit has been reached along with the effective limit.
- Wire the check into `aibridged`'s HTTP handler. The caller computes the period start (monthly for now) and passes it in the request.
- Reject exceeded requests with HTTP 403 Forbidden and a message directing the user to contact an administrator.
- Add `dbtime.StartOfMonth` alongside `StartOfDay` for period computation.
- Add real-DB tests covering the enforcement path: month-boundary excludes prior-period spend, and a new user override unblocks a previously-exceeded user.

Closes https://linear.app/codercom/issue/AIGOV-428/add-pre-request-budget-enforcement

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-02 16:53:36 +01:00
committed by GitHub
parent be9c95c8f5
commit 1989db0e2b
9 changed files with 885 additions and 181 deletions
+30 -7
View File
@@ -1,17 +1,20 @@
package aibridged
import (
"fmt"
"net/http"
"strings"
"github.com/google/uuid"
"golang.org/x/xerrors"
"google.golang.org/protobuf/types/known/timestamppb"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/recorder"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/database/dbtime"
)
var _ http.Handler = &Server{}
@@ -21,6 +24,7 @@ var (
ErrConnect = xerrors.New("could not connect to coderd")
ErrUnauthorized = xerrors.New("unauthorized")
ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler")
ErrBudgetCheck = xerrors.New("internal server error checking user AI budget")
)
// ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge.
@@ -135,6 +139,32 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
return
}
id, err := uuid.Parse(resp.GetOwnerId())
if err != nil {
logger.Warn(ctx, "failed to parse user ID", slog.Error(err), slog.F("id", resp.GetOwnerId()))
http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden)
return
}
logger = logger.With(slog.F("user_id", id))
periodStart := dbtime.StartOfMonth(dbtime.Now().UTC())
budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
UserId: id.String(),
PeriodStart: timestamppb.New(periodStart),
})
if err != nil {
logger.Warn(ctx, "user AI budget check failed", slog.Error(err))
http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError)
return
}
if budgetResp.GetExceeded() {
http.Error(rw, fmt.Sprintf(
"AI budget of US$%.2f exceeded. Please contact an administrator for more details.",
float64(budgetResp.GetSpendLimitMicros())/1_000_000,
), http.StatusForbidden)
return
}
// Rewire request context to include actor.
//
// [NOTE]
@@ -144,13 +174,6 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
"Username": resp.GetUsername(),
}))
id, err := uuid.Parse(resp.GetOwnerId())
if err != nil {
logger.Warn(ctx, "failed to parse user ID", slog.Error(err), slog.F("id", resp.GetOwnerId()))
http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden)
return
}
handler, err := s.GetRequestHandler(ctx, Request{
SessionKey: key,
APIKeyID: resp.ApiKeyId,