opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add ai_user_daily_spend table and queries (#26562)

Browse Source 
## Description

Adds the spend tracking table and queries needed by [AIGOV-427](https://linear.app/codercom/issue/AIGOV-427/add-post-response-spend-accumulation) (post-response accumulation) and [AIGOV-428](https://linear.app/codercom/issue/AIGOV-428/add-pre-request-budget-enforcement) (pre-request enforcement).

## Changes

- Add `ai_user_daily_spend` table to aggregate per-user, per-effective-group AI spend by UTC day.
- Add `UpsertUserAIDailySpend` and `GetUserAISpendSince` queries.

Closes https://linear.app/codercom/issue/AIGOV-426/add-daily-spend-table-and-queries

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-02 16:29:25 +01:00
committed by GitHub
parent dee41c34e6
commit fcdd029d74
15 changed files with 613 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+14
View File
@@ -4871,6 +4871,13 @@ func (q *querier) GetUserAISeatStates(ctx context.Context, userIDs []uuid.UUID)
return q.db.GetUserAISeatStates(ctx, userIDs)
}
func (q *querier) GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) {
if _, err := q.GetUserByID(ctx, arg.UserID); err != nil { // AuthZ check
return database.GetUserAISpendSinceRow{}, err
}
return q.db.GetUserAISpendSince(ctx, arg)
}
func (q *querier) GetUserActivityInsights(ctx context.Context, arg database.GetUserActivityInsightsParams) ([]database.GetUserActivityInsightsRow, error) {
// Used by insights endpoints. Need to check both for auditors and for regular users with template acl perms.
if err := q.authorizeContext(ctx, policy.ActionViewInsights, rbac.ResourceTemplate); err != nil {
@@ -5726,6 +5733,13 @@ func (q *querier) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UU
return q.db.IncrementChatGenerationAttempt(ctx, id)
}
func (q *querier) IncrementUserAIDailySpend(ctx context.Context, arg database.IncrementUserAIDailySpendParams) (database.AIUserDailySpend, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil {
return database.AIUserDailySpend{}, err
}
return q.db.IncrementUserAIDailySpend(ctx, arg)
}
func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) {
return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg)
}
coderd/database/dbauthz/dbauthz_test.go
+25
View File
@@ -6802,6 +6802,31 @@ func (s *MethodTestSuite) TestAIBridge() {
check.Args(user.ID).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override)
}))
s.Run("GetUserAISpendSince", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
user := testutil.Fake(s.T(), faker, database.User{})
arg := database.GetUserAISpendSinceParams{
UserID: user.ID,
EffectiveGroupID: uuid.New(),
PeriodStart: time.Now().UTC().Truncate(24 * time.Hour),
}
row := testutil.Fake(s.T(), faker, database.GetUserAISpendSinceRow{UserID: user.ID, EffectiveGroupID: arg.EffectiveGroupID, PeriodStart: arg.PeriodStart})
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
dbm.EXPECT().GetUserAISpendSince(gomock.Any(), arg).Return(row, nil).AnyTimes()
check.Args(arg).Asserts(user, policy.ActionRead).Returns(row)
}))
s.Run("IncrementUserAIDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
arg := database.IncrementUserAIDailySpendParams{
UserID: uuid.New(),
EffectiveGroupID: uuid.New(),
Day: time.Now().UTC().Truncate(24 * time.Hour),
CostMicros: 1000,
}
row := testutil.Fake(s.T(), faker, database.AIUserDailySpend{UserID: arg.UserID, EffectiveGroupID: arg.EffectiveGroupID, Day: arg.Day})
dbm.EXPECT().IncrementUserAIDailySpend(gomock.Any(), arg).Return(row, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceAibridgeInterception, policy.ActionUpdate).Returns(row)
}))
s.Run("GetAIProviderByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
provider := testutil.Fake(s.T(), faker, database.AIProvider{})
dbm.EXPECT().GetAIProviderByID(gomock.Any(), provider.ID).Return(provider, nil).AnyTimes()