feat(coderd): enforce ai budget on pre-request path - #26915
#26915feat(coderd): enforce ai budget on pre-request path#26915ssncferreira merged 6 commits intomaincoder/coder:mainfrom ssncf/ai-budget-enforcementcoder/coder:ssncf/ai-budget-enforcementCopy head branch name to clipboard
Conversation
This stack of pull requests is managed by Graphite. Learn more about stacking. |
8fe4e34 to
e83bc03
Compare
f1c6469 to
13eb001
Compare
a90d3f2 to
af61bce
Compare
af61bce to
90b1aaa
Compare
|
/coder-agents-review |
|
Chat: Review posted | View chat Review history
deep-review v0.9.0 | Round 1 | Last posted: Round 1, 9 findings (4 P3, 3 Nit, 2 Note), COMMENT. Review Finding inventoryFindings
Round logRound 1Netero: no findings. Panel: Bisky, Chopper, Ging-Go, Gon, Hisoka, Killua, Kite, Knov, Komugi, Leorio, Mafu-san, Mafuuu, Meruem, Pariston, Ryosuke, Zoro. 4 P3, 3 Nit, 2 Note. Reviewed against 13eb001..90b1aaa. About deep-reviewCRF = Coder Review Finding (P0-P4, Nit, Note)
|
There was a problem hiding this comment.
Clean, focused PR that adds a real enforcement gate where none existed. Test density is 65%, coverage is thorough (9 unit scenarios + a real-DB boundary test), and the code reuses budget.ResolveUserAIBudget from the post-request path. Well-decomposed.
Severity count: 4x P3, 3x Nit, 2x Note. No P0-P2.
This is a first panel round. All findings are below P0-P1, so no changes requested.
"The PR title says 'enforce.' The implementation checks, but enforcement requires the check to exclude the gap." — Hisoka
🤖 This review was automatically generated with Coder Agents.
|
|
||
| // TODO: make the budget period configurable; monthly for now. | ||
| periodStart := dbtime.StartOfMonth(dbtime.Now().UTC()) | ||
| spendStatus, err := client.GetUserAISpendStatus(ctx, &proto.GetUserAISpendStatusRequest{ |
There was a problem hiding this comment.
In here, I opted for a different approach from the other RPC calls. On the other RPC calls, we don't distinguish between internal server errors and logic errors (like authentication), both return the same response, e.g., ErrUnauthorized.Error() . To distinguish between internal server errors and the actual budget error, err handles only the first case. When the budget is exceeded, we check spendStatus.GetExceeded() and send the appropriate error response. It is not really consistent with the other cases, but I thought it would result in clearer errors. Let me know what you think.
| // - Overage is bounded by request cost × concurrency; once the accumulated | ||
| // spend crosses the limit, subsequent requests are blocked. | ||
| // - Fail-open is the deliberate design choice for this feature. | ||
| func (s *Server) resolveUserAISpendStatus(ctx context.Context, userID uuid.UUID, periodStart time.Time) (userAISpendStatus, error) { |
There was a problem hiding this comment.
Enforcement is a pre-LLM-request check on the hot path. In the worst case, in addition to the authorization database queries, we do 3 more queries:
GetUserAIBudgetOverrideGetHighestGroupAIBudgetByUserGetUserAISpendSince
I think it's clearer to keep it this way for now as each query is focused and self-contained. If this becomes a bottleneck, we can consolidate into a single optimized query.
3243363 to
957b043
Compare
Documentation CheckUpdates Needed
Automated review via Coder Agents |
| PeriodStart: timestamppb.New(periodStart), | ||
| }) | ||
| if err != nil { | ||
| logger.Warn(ctx, "ai spend status check failed", slog.Error(err)) |
There was a problem hiding this comment.
This should be an error not a warning I reckon.
There was a problem hiding this comment.
Correct, it should. But all the other logs here use Warn as well for error message. And my assumption is that this is due to the tests, because then we would need to update them to ignore errors. I can address this in a follow-up and update all the logs as well as tests if that makes sense
9891e0b to
24ec993
Compare
24ec993 to
15f3e9a
Compare
13eb001 to
dbb6f44
Compare
15f3e9a to
6708bd9
Compare
dbb6f44 to
04e366f
Compare
Merge activity
|
a76e58c to
e753f18
Compare
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
IsBudgetExceededRPC toaibridgedserver. 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.aibridged's HTTP handler. The caller computes the period start (monthly for now) and passes it in the request.dbtime.StartOfMonthalongsideStartOfDayfor period computation.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