From 77bd10da2aa6fdce30ef570c200f0e9330e0fdfb Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Mon, 22 Jun 2026 08:18:05 +0000 Subject: [PATCH 1/7] feat: add ai_user_daily_spend table and queries --- coderd/database/check_constraint.go | 1 + coderd/database/dbauthz/dbauthz.go | 15 + coderd/database/dbauthz/dbauthz_test.go | 25 ++ coderd/database/dbmetrics/querymetrics.go | 16 + coderd/database/dbmock/dbmock.go | 30 ++ coderd/database/dump.sql | 23 ++ .../000530_ai_user_daily_spend.down.sql | 1 + .../000530_ai_user_daily_spend.up.sql | 21 ++ .../000530_ai_user_daily_spend.up.sql | 7 + coderd/database/models.go | 12 + coderd/database/querier.go | 7 + coderd/database/querier_test.go | 340 ++++++++++++++++++ coderd/database/queries.sql.go | 75 ++++ coderd/database/queries/aicostcontrol.sql | 24 ++ coderd/database/unique_constraint.go | 1 + 15 files changed, 598 insertions(+) create mode 100644 coderd/database/migrations/000530_ai_user_daily_spend.down.sql create mode 100644 coderd/database/migrations/000530_ai_user_daily_spend.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000530_ai_user_daily_spend.up.sql diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go index 76b350816d7..fda7bc46757 100644 --- a/coderd/database/check_constraint.go +++ b/coderd/database/check_constraint.go @@ -14,6 +14,7 @@ const ( CheckAIModelPricesInputPriceCheck CheckConstraint = "ai_model_prices_input_price_check" // ai_model_prices CheckAIModelPricesOutputPriceCheck CheckConstraint = "ai_model_prices_output_price_check" // ai_model_prices CheckAIProvidersNameCheck CheckConstraint = "ai_providers_name_check" // ai_providers + CheckAIUserDailySpendSpendMicrosCheck CheckConstraint = "ai_user_daily_spend_spend_micros_check" // ai_user_daily_spend CheckAibridgeTokenUsagesCacheReadPriceMicrosCheck CheckConstraint = "aibridge_token_usages_cache_read_price_micros_check" // aibridge_token_usages CheckAibridgeTokenUsagesCacheWritePriceMicrosCheck CheckConstraint = "aibridge_token_usages_cache_write_price_micros_check" // aibridge_token_usages CheckAibridgeTokenUsagesCostMicrosCheck CheckConstraint = "aibridge_token_usages_cost_micros_check" // aibridge_token_usages diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 20a688c768a..00aba9a8c88 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5093,6 +5093,13 @@ func (q *querier) GetUserSkillByUserIDAndName(ctx context.Context, arg database. return q.db.GetUserSkillByUserIDAndName(ctx, arg) } +func (q *querier) GetUserSpendSince(ctx context.Context, arg database.GetUserSpendSinceParams) (database.GetUserSpendSinceRow, error) { + if _, err := q.GetUserByID(ctx, arg.UserID); err != nil { // AuthZ check + return database.GetUserSpendSinceRow{}, err + } + return q.db.GetUserSpendSince(ctx, arg) +} + func (q *querier) GetUserStatusCounts(ctx context.Context, arg database.GetUserStatusCountsParams) ([]database.GetUserStatusCountsRow, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUser); err != nil { return nil, err @@ -8936,6 +8943,14 @@ func (q *querier) UpsertUserChatPersonalModelOverride(ctx context.Context, arg d return q.db.UpsertUserChatPersonalModelOverride(ctx, arg) } +func (q *querier) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { + // Daily spend writes are made by the aibridged process. + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { + return database.AIUserDailySpend{}, err + } + return q.db.UpsertUserDailySpend(ctx, arg) +} + func (q *querier) UpsertWebpushVAPIDKeys(ctx context.Context, arg database.UpsertWebpushVAPIDKeysParams) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 68d5d222bb6..f30c9074ce3 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6802,6 +6802,31 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(user.ID).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override) })) + s.Run("GetUserSpendSince", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + user := testutil.Fake(s.T(), faker, database.User{}) + arg := database.GetUserSpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: uuid.New(), + PeriodStart: time.Now().UTC().Truncate(24 * time.Hour), + } + row := testutil.Fake(s.T(), faker, database.GetUserSpendSinceRow{UserID: user.ID, EffectiveGroupID: arg.EffectiveGroupID, PeriodStart: arg.PeriodStart}) + dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes() + dbm.EXPECT().GetUserSpendSince(gomock.Any(), arg).Return(row, nil).AnyTimes() + check.Args(arg).Asserts(user, policy.ActionRead).Returns(row) + })) + + s.Run("UpsertUserDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.UpsertUserDailySpendParams{ + 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().UpsertUserDailySpend(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() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 2095b41f5d9..d526862062a 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3321,6 +3321,14 @@ func (m queryMetricsStore) GetUserSkillByUserIDAndName(ctx context.Context, arg return r0, r1 } +func (m queryMetricsStore) GetUserSpendSince(ctx context.Context, arg database.GetUserSpendSinceParams) (database.GetUserSpendSinceRow, error) { + start := time.Now() + r0, r1 := m.s.GetUserSpendSince(ctx, arg) + m.queryLatencies.WithLabelValues("GetUserSpendSince").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserSpendSince").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetUserStatusCounts(ctx context.Context, arg database.GetUserStatusCountsParams) ([]database.GetUserStatusCountsRow, error) { start := time.Now() r0, r1 := m.s.GetUserStatusCounts(ctx, arg) @@ -6433,6 +6441,14 @@ func (m queryMetricsStore) UpsertUserChatPersonalModelOverride(ctx context.Conte return r0 } +func (m queryMetricsStore) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { + start := time.Now() + r0, r1 := m.s.UpsertUserDailySpend(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertUserDailySpend").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertUserDailySpend").Inc() + return r0, r1 +} + func (m queryMetricsStore) UpsertWebpushVAPIDKeys(ctx context.Context, arg database.UpsertWebpushVAPIDKeysParams) error { start := time.Now() r0 := m.s.UpsertWebpushVAPIDKeys(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 4b34989ae5b..4ddc598661a 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -6207,6 +6207,21 @@ func (mr *MockStoreMockRecorder) GetUserSkillByUserIDAndName(ctx, arg any) *gomo return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserSkillByUserIDAndName", reflect.TypeOf((*MockStore)(nil).GetUserSkillByUserIDAndName), ctx, arg) } +// GetUserSpendSince mocks base method. +func (m *MockStore) GetUserSpendSince(ctx context.Context, arg database.GetUserSpendSinceParams) (database.GetUserSpendSinceRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetUserSpendSince", ctx, arg) + ret0, _ := ret[0].(database.GetUserSpendSinceRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetUserSpendSince indicates an expected call of GetUserSpendSince. +func (mr *MockStoreMockRecorder) GetUserSpendSince(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserSpendSince", reflect.TypeOf((*MockStore)(nil).GetUserSpendSince), ctx, arg) +} + // GetUserStatusCounts mocks base method. func (m *MockStore) GetUserStatusCounts(ctx context.Context, arg database.GetUserStatusCountsParams) ([]database.GetUserStatusCountsRow, error) { m.ctrl.T.Helper() @@ -12024,6 +12039,21 @@ func (mr *MockStoreMockRecorder) UpsertUserChatPersonalModelOverride(ctx, arg an return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserChatPersonalModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertUserChatPersonalModelOverride), ctx, arg) } +// UpsertUserDailySpend mocks base method. +func (m *MockStore) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertUserDailySpend", ctx, arg) + ret0, _ := ret[0].(database.AIUserDailySpend) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UpsertUserDailySpend indicates an expected call of UpsertUserDailySpend. +func (mr *MockStoreMockRecorder) UpsertUserDailySpend(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserDailySpend", reflect.TypeOf((*MockStore)(nil).UpsertUserDailySpend), ctx, arg) +} + // UpsertWebpushVAPIDKeys mocks base method. func (m *MockStore) UpsertWebpushVAPIDKeys(ctx context.Context, arg database.UpsertWebpushVAPIDKeysParams) error { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 35d63873d3e..7b659c63dd7 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1517,6 +1517,24 @@ CREATE TABLE ai_seat_state ( updated_at timestamp with time zone NOT NULL ); +CREATE TABLE ai_user_daily_spend ( + user_id uuid NOT NULL, + effective_group_id uuid NOT NULL, + day date NOT NULL, + spend_micros bigint NOT NULL, + CONSTRAINT ai_user_daily_spend_spend_micros_check CHECK ((spend_micros >= 0)) +); + +COMMENT ON TABLE ai_user_daily_spend IS 'Daily AI spend per user and effective group.'; + +COMMENT ON COLUMN ai_user_daily_spend.user_id IS 'The user who incurred the spend.'; + +COMMENT ON COLUMN ai_user_daily_spend.effective_group_id IS 'The group this spend is attributed to for budget purposes.'; + +COMMENT ON COLUMN ai_user_daily_spend.day IS 'UTC calendar day the spend was incurred.'; + +COMMENT ON COLUMN ai_user_daily_spend.spend_micros IS 'Accumulated spend in micro-units (1 unit = 1,000,000).'; + CREATE TABLE aibridge_interceptions ( id uuid NOT NULL, initiator_id uuid NOT NULL, @@ -4120,6 +4138,9 @@ ALTER TABLE ONLY ai_providers ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); +ALTER TABLE ONLY ai_user_daily_spend + ADD CONSTRAINT ai_user_daily_spend_pkey PRIMARY KEY (user_id, effective_group_id, day); + ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_pkey PRIMARY KEY (id); @@ -4526,6 +4547,8 @@ CREATE INDEX idx_ai_provider_keys_provider_id ON ai_provider_keys USING btree (p CREATE INDEX idx_ai_providers_enabled ON ai_providers USING btree (enabled) WHERE (deleted = false); +CREATE INDEX idx_ai_user_daily_spend_effective_group_id_day ON ai_user_daily_spend USING btree (effective_group_id, day); + CREATE INDEX idx_aibridge_interceptions_agent_firewall_session_id ON aibridge_interceptions USING btree (agent_firewall_session_id) WHERE (agent_firewall_session_id IS NOT NULL); CREATE INDEX idx_aibridge_interceptions_client ON aibridge_interceptions USING btree (client); diff --git a/coderd/database/migrations/000530_ai_user_daily_spend.down.sql b/coderd/database/migrations/000530_ai_user_daily_spend.down.sql new file mode 100644 index 00000000000..a559706ecde --- /dev/null +++ b/coderd/database/migrations/000530_ai_user_daily_spend.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS ai_user_daily_spend CASCADE; diff --git a/coderd/database/migrations/000530_ai_user_daily_spend.up.sql b/coderd/database/migrations/000530_ai_user_daily_spend.up.sql new file mode 100644 index 00000000000..acc65d12fd9 --- /dev/null +++ b/coderd/database/migrations/000530_ai_user_daily_spend.up.sql @@ -0,0 +1,21 @@ +-- Aggregates a user's AI spend within their effective group, one row per +-- UTC day. Drives budget enforcement and reporting. +CREATE TABLE ai_user_daily_spend ( + -- No FK to users. Spend records persist after user deletion. + user_id UUID NOT NULL, + -- No FK to groups. Spend records persist after group deletion. + effective_group_id UUID NOT NULL, + day DATE NOT NULL, + spend_micros BIGINT NOT NULL CHECK (spend_micros >= 0), + PRIMARY KEY (user_id, effective_group_id, day) +); + +COMMENT ON TABLE ai_user_daily_spend IS 'Daily AI spend per user and effective group.'; +COMMENT ON COLUMN ai_user_daily_spend.user_id IS 'The user who incurred the spend.'; +COMMENT ON COLUMN ai_user_daily_spend.effective_group_id IS 'The group this spend is attributed to for budget purposes.'; +COMMENT ON COLUMN ai_user_daily_spend.day IS 'UTC calendar day the spend was incurred.'; +COMMENT ON COLUMN ai_user_daily_spend.spend_micros IS 'Accumulated spend in micro-units (1 unit = 1,000,000).'; + +-- For queries filtering by effective_group_id alone. +CREATE INDEX idx_ai_user_daily_spend_effective_group_id_day + ON ai_user_daily_spend (effective_group_id, day); diff --git a/coderd/database/migrations/testdata/fixtures/000530_ai_user_daily_spend.up.sql b/coderd/database/migrations/testdata/fixtures/000530_ai_user_daily_spend.up.sql new file mode 100644 index 00000000000..5a5dc711ec0 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000530_ai_user_daily_spend.up.sql @@ -0,0 +1,7 @@ +INSERT INTO ai_user_daily_spend ( + user_id, + effective_group_id, + day, + spend_micros +) VALUES + ('30095c71-380b-457a-8995-97b8ee6e5307', 'bb640d07-ca8a-4869-b6bc-ae61ebb2fda1', '2024-06-15', 100000); diff --git a/coderd/database/models.go b/coderd/database/models.go index c0dea30800c..2abae159c1c 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4678,6 +4678,18 @@ type AISeatState struct { UpdatedAt time.Time `db:"updated_at" json:"updated_at"` } +// Daily AI spend per user and effective group. +type AIUserDailySpend struct { + // The user who incurred the spend. + UserID uuid.UUID `db:"user_id" json:"user_id"` + // The group this spend is attributed to for budget purposes. + EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` + // UTC calendar day the spend was incurred. + Day time.Time `db:"day" json:"day"` + // Accumulated spend in micro-units (1 unit = 1,000,000). + SpendMicros int64 `db:"spend_micros" json:"spend_micros"` +} + type APIKey struct { ID string `db:"id" json:"id"` // hashed_secret contains a SHA256 hash of the key secret. This is considered a secret and MUST NOT be returned from the API as it is used for API key encryption in app proxying code. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 92dd60ebb5e..eff4242a99c 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -882,6 +882,9 @@ type sqlcQuerier interface { GetUserSecretsTelemetrySummary(ctx context.Context) (GetUserSecretsTelemetrySummaryRow, error) GetUserShellToolDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) GetUserSkillByUserIDAndName(ctx context.Context, arg GetUserSkillByUserIDAndNameParams) (UserSkill, error) + // Total spend for (user_id, effective_group_id) on or after period_start. + // The period_start parameter is normalized to its UTC calendar day. + GetUserSpendSince(ctx context.Context, arg GetUserSpendSinceParams) (GetUserSpendSinceRow, error) // GetUserStatusCounts returns the count of users in each status over time. // The time range is inclusively defined by the start_time and end_time parameters. GetUserStatusCounts(ctx context.Context, arg GetUserStatusCountsParams) ([]GetUserStatusCountsRow, error) @@ -1573,6 +1576,10 @@ type sqlcQuerier interface { UpsertUserAIProviderKey(ctx context.Context, arg UpsertUserAIProviderKeyParams) (UserAIProviderKey, error) UpsertUserChatDebugLoggingEnabled(ctx context.Context, arg UpsertUserChatDebugLoggingEnabledParams) error UpsertUserChatPersonalModelOverride(ctx context.Context, arg UpsertUserChatPersonalModelOverrideParams) error + // Adds cost_micros to the spend for (user_id, effective_group_id, day). + // The day parameter is normalized to its UTC calendar day before storage. + // Returns the resulting row. + UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) (AIUserDailySpend, error) UpsertWebpushVAPIDKeys(ctx context.Context, arg UpsertWebpushVAPIDKeysParams) error UpsertWorkspaceAgentContextResource(ctx context.Context, arg UpsertWorkspaceAgentContextResourceParams) (WorkspaceAgentContextResource, error) UpsertWorkspaceAgentContextSnapshot(ctx context.Context, arg UpsertWorkspaceAgentContextSnapshotParams) (WorkspaceAgentContextSnapshot, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 2e8924e502a..80d1c2fbb3c 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11751,6 +11751,346 @@ func TestUpsertAISeats(t *testing.T) { require.False(t, alreadyExists) } +func TestUpsertUserDailySpend(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + day := time.Date(2024, 6, 15, 0, 0, 0, 0, time.UTC) + nextDay := day.AddDate(0, 0, 1) + + // Given a sequence of costs upserted to the same (user, group, day), + // when applied in order, then they accumulate into a single row. + tests := []struct { + name string + costs []int64 + wantTotal int64 + }{ + {name: "InsertsNewRow", costs: []int64{100}, wantTotal: 100}, + {name: "AccumulatesAcrossCalls", costs: []int64{100, 50, 30, 20}, wantTotal: 200}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + var row database.AIUserDailySpend + for _, cost := range tt.costs { + var err error + row, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: day, + CostMicros: cost, + }) + require.NoError(t, err) + } + require.Equal(t, user.ID, row.UserID) + require.Equal(t, group.ID, row.EffectiveGroupID) + require.Equal(t, tt.wantTotal, row.SpendMicros) + require.True(t, row.Day.Equal(day), + "row.Day = %s, want = %s", row.Day, day) + }) + } + + // Given two users in the same group on the same day, when each upserts, then each gets its own row. + t.Run("SeparateRowPerUser", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + userA := dbgen.User(t, db, database.User{}) + userB := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + userARow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: userA.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, + }) + require.NoError(t, err) + require.Equal(t, int64(100), userARow.SpendMicros) + + userBRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: userB.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 25, + }) + require.NoError(t, err) + require.Equal(t, int64(25), userBRow.SpendMicros, + "userB row must not include userA spend") + }) + + // Given one user across two groups on the same day, when each upserts, then each gets its own row. + t.Run("SeparateRowPerEffectiveGroup", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + groupARow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupA.ID, Day: day, CostMicros: 100, + }) + require.NoError(t, err) + require.Equal(t, int64(100), groupARow.SpendMicros) + + groupBRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupB.ID, Day: day, CostMicros: 25, + }) + require.NoError(t, err) + require.Equal(t, int64(25), groupBRow.SpendMicros, + "groupB row must not include groupA spend") + }) + + // Given existing spend on day, when the same user upserts on the next day, then a new row is created. + t.Run("SeparateRowPerDay", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + dayRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, + }) + require.NoError(t, err) + require.Equal(t, int64(100), dayRow.SpendMicros) + + // The ON CONFLICT target is the full PK including day, so this upsert + // cannot modify the previous day's row by construction. + nextDayRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: nextDay, CostMicros: 25, + }) + require.NoError(t, err) + require.Equal(t, int64(25), nextDayRow.SpendMicros, + "nextDay row must not include day spend") + require.True(t, nextDayRow.Day.Equal(nextDay)) + }) + + // Given a non-midnight UTC time, when upserted, then it lands on the same row as the truncated day. + t.Run("TruncatesDayToUTCMidnight", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, + }) + require.NoError(t, err) + dayNonTruncated := day.Add(14*time.Hour + 30*time.Minute) + row, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: dayNonTruncated, CostMicros: 50, + }) + require.NoError(t, err) + require.Equal(t, int64(150), row.SpendMicros, + "non-midnight UTC time should accumulate on the truncated day's row") + require.True(t, row.Day.Equal(day), + "row.Day = %s, want truncated = %s", row.Day, day) + }) + + // Given a non-UTC time that crosses the UTC date boundary, when upserted, then it lands on the UTC calendar day. + t.Run("NormalizesNonUTCTimezones", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + // 2024-06-15 23:00 in UTC-5 is 2024-06-16 04:00 UTC, so this should land on nextDay. + localLate := time.Date(2024, 6, 15, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) + row, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: localLate, CostMicros: 100, + }) + require.NoError(t, err) + require.True(t, row.Day.Equal(nextDay), + "non-UTC input should land on the UTC calendar day (%s), got %s", nextDay, row.Day) + }) + + // Given a zero-cost upsert, when applied, then it is idempotent (creates a zero-spend row or leaves an existing one unchanged). + t.Run("ZeroCostIsIdempotent", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + // Zero-cost upsert on a fresh key creates a row with spend = 0. + newRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 0, + }) + require.NoError(t, err) + require.Equal(t, int64(0), newRow.SpendMicros) + + // After a real upsert, the row has spend = 100. + _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, + }) + require.NoError(t, err) + + // Zero-cost upsert on the existing row leaves spend unchanged. + sameRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 0, + }) + require.NoError(t, err) + require.Equal(t, int64(100), sameRow.SpendMicros, + "zero-cost upsert must not change existing spend") + }) +} + +func TestGetUserSpendSince(t *testing.T) { + t.Parallel() + + // Use fixed dates to keep the test deterministic. + monthStart := time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC) + today := monthStart.AddDate(0, 0, 14) // 2024-06-15 + prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31 + + // Given seeded rows for a single (user, group), when querying since + // monthStart, then the period sum is returned. + tests := []struct { + name string + rows []struct { + day time.Time + spend int64 + } + wantSpend int64 + }{ + {name: "NoRows", wantSpend: 0}, + { + name: "SingleRowOnToday", + rows: []struct { + day time.Time + spend int64 + }{{today, 100}}, + wantSpend: 100, + }, + { + name: "FirstOfMonthIncluded", + rows: []struct { + day time.Time + spend int64 + }{{monthStart, 50}}, + wantSpend: 50, + }, + { + name: "SumsMultipleDaysInMonth", + rows: []struct { + day time.Time + spend int64 + }{{monthStart, 50}, {today, 100}}, + wantSpend: 150, + }, + { + name: "ExcludesRowsBeforePeriodStart", + rows: []struct { + day time.Time + spend int64 + }{{prevMonthLastDay, 999}, {monthStart, 25}}, + wantSpend: 25, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + for _, r := range tt.rows { + _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: r.day, + CostMicros: r.spend, + }) + require.NoError(t, err) + } + + got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + PeriodStart: monthStart, + }) + require.NoError(t, err) + require.Equal(t, user.ID, got.UserID) + require.Equal(t, group.ID, got.EffectiveGroupID) + require.True(t, got.PeriodStart.Equal(monthStart), + "PeriodStart = %s, want = %s", got.PeriodStart, monthStart) + require.Equal(t, tt.wantSpend, got.SpendMicros) + }) + } + + // Given two users with spend in the same group on the same day, when querying one user, then the other's spend is excluded. + t.Run("SumExcludesOtherUsers", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + userA := dbgen.User(t, db, database.User{}) + userB := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: userA.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 100, + }) + require.NoError(t, err) + _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: userB.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 25, + }) + require.NoError(t, err) + + got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + UserID: userB.ID, + EffectiveGroupID: group.ID, + PeriodStart: monthStart, + }) + require.NoError(t, err) + require.Equal(t, int64(25), got.SpendMicros, + "userB sum must not include userA spend") + }) + + // Given one user with spend in two groups on the same day, when querying one group, then the other's spend is excluded. + t.Run("SumExcludesOtherEffectiveGroups", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupA.ID, Day: today, CostMicros: 100, + }) + require.NoError(t, err) + _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupB.ID, Day: today, CostMicros: 25, + }) + require.NoError(t, err) + + got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: groupB.ID, + PeriodStart: monthStart, + }) + require.NoError(t, err) + require.Equal(t, int64(25), got.SpendMicros, + "groupB sum must not include groupA spend") + }) +} + func TestChatPinOrderQueries(t *testing.T) { t.Parallel() if testing.Short() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 6e41d273b42..538e75b8599 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2546,6 +2546,45 @@ func (q *sqlQuerier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UU return i, err } +const getUserSpendSince = `-- name: GetUserSpendSince :one +SELECT + $1::uuid AS user_id, + $2::uuid AS effective_group_id, + (($3::timestamptz) AT TIME ZONE 'UTC')::date AS period_start, + COALESCE(SUM(spend_micros), 0)::BIGINT AS spend_micros +FROM ai_user_daily_spend +WHERE user_id = $1 + AND effective_group_id = $2 + AND day >= (($3::timestamptz) AT TIME ZONE 'UTC')::date +` + +type GetUserSpendSinceParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` + PeriodStart time.Time `db:"period_start" json:"period_start"` +} + +type GetUserSpendSinceRow struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` + PeriodStart time.Time `db:"period_start" json:"period_start"` + SpendMicros int64 `db:"spend_micros" json:"spend_micros"` +} + +// Total spend for (user_id, effective_group_id) on or after period_start. +// The period_start parameter is normalized to its UTC calendar day. +func (q *sqlQuerier) GetUserSpendSince(ctx context.Context, arg GetUserSpendSinceParams) (GetUserSpendSinceRow, error) { + row := q.db.QueryRowContext(ctx, getUserSpendSince, arg.UserID, arg.EffectiveGroupID, arg.PeriodStart) + var i GetUserSpendSinceRow + err := row.Scan( + &i.UserID, + &i.EffectiveGroupID, + &i.PeriodStart, + &i.SpendMicros, + ) + return i, err +} + const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec INSERT INTO ai_model_prices ( provider, model, input_price, output_price, cache_read_price, cache_write_price @@ -2629,6 +2668,42 @@ func (q *sqlQuerier) UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertU return i, err } +const upsertUserDailySpend = `-- name: UpsertUserDailySpend :one +INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) +VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4) +ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET + -- Add this call's cost to the running total for that day. + spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros +RETURNING user_id, effective_group_id, day, spend_micros +` + +type UpsertUserDailySpendParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` + Day time.Time `db:"day" json:"day"` + CostMicros int64 `db:"cost_micros" json:"cost_micros"` +} + +// Adds cost_micros to the spend for (user_id, effective_group_id, day). +// The day parameter is normalized to its UTC calendar day before storage. +// Returns the resulting row. +func (q *sqlQuerier) UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) (AIUserDailySpend, error) { + row := q.db.QueryRowContext(ctx, upsertUserDailySpend, + arg.UserID, + arg.EffectiveGroupID, + arg.Day, + arg.CostMicros, + ) + var i AIUserDailySpend + err := row.Scan( + &i.UserID, + &i.EffectiveGroupID, + &i.Day, + &i.SpendMicros, + ) + return i, err +} + const getActiveAISeatCount = `-- name: GetActiveAISeatCount :one SELECT COUNT(*) diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index cad1ed64522..3e25bc46d9a 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -79,3 +79,27 @@ ORDER BY -- (groups are unique on (organization_id, name), not name alone). gaib.group_id ASC LIMIT 1; + +-- name: UpsertUserDailySpend :one +-- Adds cost_micros to the spend for (user_id, effective_group_id, day). +-- The day parameter is normalized to its UTC calendar day before storage. +-- Returns the resulting row. +INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) +VALUES (@user_id, @effective_group_id, ((@day::timestamptz) AT TIME ZONE 'UTC')::date, @cost_micros) +ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET + -- Add this call's cost to the running total for that day. + spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros +RETURNING *; + +-- name: GetUserSpendSince :one +-- Total spend for (user_id, effective_group_id) on or after period_start. +-- The period_start parameter is normalized to its UTC calendar day. +SELECT + @user_id::uuid AS user_id, + @effective_group_id::uuid AS effective_group_id, + ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date AS period_start, + COALESCE(SUM(spend_micros), 0)::BIGINT AS spend_micros +FROM ai_user_daily_spend +WHERE user_id = @user_id + AND effective_group_id = @effective_group_id + AND day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date; diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 68058923b3b..17adc9486d6 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -12,6 +12,7 @@ const ( UniqueAIProviderKeysPkey UniqueConstraint = "ai_provider_keys_pkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id); UniqueAIProvidersPkey UniqueConstraint = "ai_providers_pkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_pkey PRIMARY KEY (id); UniqueAISeatStatePkey UniqueConstraint = "ai_seat_state_pkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); + UniqueAIUserDailySpendPkey UniqueConstraint = "ai_user_daily_spend_pkey" // ALTER TABLE ONLY ai_user_daily_spend ADD CONSTRAINT ai_user_daily_spend_pkey PRIMARY KEY (user_id, effective_group_id, day); UniqueAibridgeInterceptionsPkey UniqueConstraint = "aibridge_interceptions_pkey" // ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_pkey PRIMARY KEY (id); UniqueAibridgeTokenUsagesPkey UniqueConstraint = "aibridge_token_usages_pkey" // ALTER TABLE ONLY aibridge_token_usages ADD CONSTRAINT aibridge_token_usages_pkey PRIMARY KEY (id); UniqueAibridgeToolUsagesPkey UniqueConstraint = "aibridge_tool_usages_pkey" // ALTER TABLE ONLY aibridge_tool_usages ADD CONSTRAINT aibridge_tool_usages_pkey PRIMARY KEY (id); From 48ca411b3f0440b2a7693f6a9d21bc825cf7156b Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 23 Jun 2026 12:54:13 +0000 Subject: [PATCH 2/7] chore: address review comments --- coderd/database/dbauthz/dbauthz.go | 6 ++ coderd/database/dbauthz/dbauthz_test.go | 10 ++ coderd/database/querier.go | 1 - coderd/database/querier_test.go | 116 +++++++++++++--------- coderd/database/queries.sql.go | 2 - coderd/database/queries/aicostcontrol.sql | 2 - 6 files changed, 84 insertions(+), 53 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 00aba9a8c88..0fa0e6ab621 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -8944,6 +8944,12 @@ func (q *querier) UpsertUserChatPersonalModelOverride(ctx context.Context, arg d } func (q *querier) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { + // Reject negative deltas. The schema CHECK only fires when the result + // goes negative, so a negative cost_micros on an existing row would + // silently reduce accumulated spend and corrupt budget enforcement. + if arg.CostMicros < 0 { + return database.AIUserDailySpend{}, xerrors.Errorf("cost_micros must be non-negative, got %d", arg.CostMicros) + } // Daily spend writes are made by the aibridged process. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { return database.AIUserDailySpend{}, err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index f30c9074ce3..5e0076d8a38 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6827,6 +6827,16 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(arg).Asserts(rbac.ResourceAibridgeInterception, policy.ActionUpdate).Returns(row) })) + s.Run("NegativeCost/UpsertUserDailySpend", s.Mocked(func(_ *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.UpsertUserDailySpendParams{ + UserID: uuid.New(), + EffectiveGroupID: uuid.New(), + Day: time.Now().UTC().Truncate(24 * time.Hour), + CostMicros: -100, + } + check.Args(arg).Asserts().Errors(xerrors.New("cost_micros must be non-negative, got -100")) + })) + 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() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index eff4242a99c..f44489097d9 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1578,7 +1578,6 @@ type sqlcQuerier interface { UpsertUserChatPersonalModelOverride(ctx context.Context, arg UpsertUserChatPersonalModelOverrideParams) error // Adds cost_micros to the spend for (user_id, effective_group_id, day). // The day parameter is normalized to its UTC calendar day before storage. - // Returns the resulting row. UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) (AIUserDailySpend, error) UpsertWebpushVAPIDKeys(ctx context.Context, arg UpsertWebpushVAPIDKeysParams) error UpsertWorkspaceAgentContextResource(ctx context.Context, arg UpsertWorkspaceAgentContextResourceParams) (WorkspaceAgentContextResource, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 80d1c2fbb3c..eed8119e9c3 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11764,9 +11764,11 @@ func TestUpsertUserDailySpend(t *testing.T) { name string costs []int64 wantTotal int64 + wantErr bool }{ {name: "InsertsNewRow", costs: []int64{100}, wantTotal: 100}, {name: "AccumulatesAcrossCalls", costs: []int64{100, 50, 30, 20}, wantTotal: 200}, + {name: "SchemaRejectsNegativeSpend", costs: []int64{-100}, wantErr: true}, } for _, tt := range tests { @@ -11779,16 +11781,23 @@ func TestUpsertUserDailySpend(t *testing.T) { group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) var row database.AIUserDailySpend + var err error for _, cost := range tt.costs { - var err error row, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: cost, }) - require.NoError(t, err) + if err != nil { + break + } + } + if tt.wantErr { + require.Error(t, err) + return } + require.NoError(t, err) require.Equal(t, user.ID, row.UserID) require.Equal(t, group.ID, row.EffectiveGroupID) require.Equal(t, tt.wantTotal, row.SpendMicros) @@ -11884,15 +11893,16 @@ func TestUpsertUserDailySpend(t *testing.T) { UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) + dayNonTruncated := day.Add(14*time.Hour + 30*time.Minute) - row, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + nonTruncatedRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: dayNonTruncated, CostMicros: 50, }) require.NoError(t, err) - require.Equal(t, int64(150), row.SpendMicros, + require.Equal(t, int64(150), nonTruncatedRow.SpendMicros, "non-midnight UTC time should accumulate on the truncated day's row") - require.True(t, row.Day.Equal(day), - "row.Day = %s, want truncated = %s", row.Day, day) + require.True(t, nonTruncatedRow.Day.Equal(day), + "row.Day = %s, want truncated = %s", nonTruncatedRow.Day, day) }) // Given a non-UTC time that crosses the UTC date boundary, when upserted, then it lands on the UTC calendar day. @@ -11904,14 +11914,14 @@ func TestUpsertUserDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - // 2024-06-15 23:00 in UTC-5 is 2024-06-16 04:00 UTC, so this should land on nextDay. + // 2024-06-15 23:00 in UTC-5 is 2024-06-16 04:00 UTC, so this should land on nextDay (2024-06-16). localLate := time.Date(2024, 6, 15, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) - row, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + nonUTCRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: localLate, CostMicros: 100, }) require.NoError(t, err) - require.True(t, row.Day.Equal(nextDay), - "non-UTC input should land on the UTC calendar day (%s), got %s", nextDay, row.Day) + require.True(t, nonUTCRow.Day.Equal(nextDay), + "non-UTC input should land on the UTC calendar day (%s), got %s", nextDay, nonUTCRow.Day) }) // Given a zero-cost upsert, when applied, then it is idempotent (creates a zero-spend row or leaves an existing one unchanged). @@ -11931,10 +11941,11 @@ func TestUpsertUserDailySpend(t *testing.T) { require.Equal(t, int64(0), newRow.SpendMicros) // After a real upsert, the row has spend = 100. - _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + updatedRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) + require.Equal(t, int64(100), updatedRow.SpendMicros) // Zero-cost upsert on the existing row leaves spend unchanged. sameRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ @@ -11954,49 +11965,23 @@ func TestGetUserSpendSince(t *testing.T) { today := monthStart.AddDate(0, 0, 14) // 2024-06-15 prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31 + type seedRow struct { + day time.Time + spend int64 + } + // Given seeded rows for a single (user, group), when querying since // monthStart, then the period sum is returned. tests := []struct { - name string - rows []struct { - day time.Time - spend int64 - } + name string + rows []seedRow wantSpend int64 }{ {name: "NoRows", wantSpend: 0}, - { - name: "SingleRowOnToday", - rows: []struct { - day time.Time - spend int64 - }{{today, 100}}, - wantSpend: 100, - }, - { - name: "FirstOfMonthIncluded", - rows: []struct { - day time.Time - spend int64 - }{{monthStart, 50}}, - wantSpend: 50, - }, - { - name: "SumsMultipleDaysInMonth", - rows: []struct { - day time.Time - spend int64 - }{{monthStart, 50}, {today, 100}}, - wantSpend: 150, - }, - { - name: "ExcludesRowsBeforePeriodStart", - rows: []struct { - day time.Time - spend int64 - }{{prevMonthLastDay, 999}, {monthStart, 25}}, - wantSpend: 25, - }, + {name: "SingleRowOnToday", rows: []seedRow{{today, 100}}, wantSpend: 100}, + {name: "FirstOfMonthIncluded", rows: []seedRow{{monthStart, 50}}, wantSpend: 50}, + {name: "SumsMultipleDaysInMonth", rows: []seedRow{{monthStart, 50}, {today, 100}}, wantSpend: 150}, + {name: "ExcludesRowsBeforePeriodStart", rows: []seedRow{{prevMonthLastDay, 999}, {monthStart, 25}}, wantSpend: 25}, } for _, tt := range tests { @@ -12089,6 +12074,41 @@ func TestGetUserSpendSince(t *testing.T) { require.Equal(t, int64(25), got.SpendMicros, "groupB sum must not include groupA spend") }) + + // Given a non-UTC period_start that lands on the previous UTC day, when queried, then it normalizes and excludes the prior day's row. + t.Run("NormalizesNonUTCPeriodStart", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) + + // Seed a row on prevMonthLastDay (which lies on May 31 UTC). A naive + // query that does not normalize the period_start would include it. + _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999, + }) + require.NoError(t, err) + _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25, + }) + require.NoError(t, err) + + // 2024-05-31 23:00 in UTC-5 is 2024-06-01 04:00 UTC, so the + // normalized period_start lands on June 1. + localLate := time.Date(2024, 5, 31, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) + got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + PeriodStart: localLate, + }) + require.NoError(t, err) + require.True(t, got.PeriodStart.Equal(monthStart), + "PeriodStart should be normalized to 2024-06-01 UTC, got %s", got.PeriodStart) + require.Equal(t, int64(25), got.SpendMicros, + "sum must exclude prevMonthLastDay row after normalization") + }) } func TestChatPinOrderQueries(t *testing.T) { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 538e75b8599..e9d50f59290 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2672,7 +2672,6 @@ const upsertUserDailySpend = `-- name: UpsertUserDailySpend :one INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4) ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET - -- Add this call's cost to the running total for that day. spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros RETURNING user_id, effective_group_id, day, spend_micros ` @@ -2686,7 +2685,6 @@ type UpsertUserDailySpendParams struct { // Adds cost_micros to the spend for (user_id, effective_group_id, day). // The day parameter is normalized to its UTC calendar day before storage. -// Returns the resulting row. func (q *sqlQuerier) UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) (AIUserDailySpend, error) { row := q.db.QueryRowContext(ctx, upsertUserDailySpend, arg.UserID, diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 3e25bc46d9a..46c421db85a 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -83,11 +83,9 @@ LIMIT 1; -- name: UpsertUserDailySpend :one -- Adds cost_micros to the spend for (user_id, effective_group_id, day). -- The day parameter is normalized to its UTC calendar day before storage. --- Returns the resulting row. INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) VALUES (@user_id, @effective_group_id, ((@day::timestamptz) AT TIME ZONE 'UTC')::date, @cost_micros) ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET - -- Add this call's cost to the running total for that day. spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros RETURNING *; From d232bd1391525864b87b37857cd65a58ba99e5b8 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 23 Jun 2026 16:10:17 +0000 Subject: [PATCH 3/7] chore: rename queries to include AI Spend --- coderd/database/dbauthz/dbauthz.go | 42 ++++++++-------- coderd/database/dbauthz/dbauthz_test.go | 18 +++---- coderd/database/dbmetrics/querymetrics.go | 32 ++++++------ coderd/database/dbmock/dbmock.go | 60 +++++++++++------------ coderd/database/querier.go | 12 ++--- coderd/database/querier_test.go | 52 ++++++++++---------- coderd/database/queries.sql.go | 20 ++++---- coderd/database/queries/aicostcontrol.sql | 4 +- 8 files changed, 120 insertions(+), 120 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 0fa0e6ab621..974ff34c1dc 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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 { @@ -5093,13 +5100,6 @@ func (q *querier) GetUserSkillByUserIDAndName(ctx context.Context, arg database. return q.db.GetUserSkillByUserIDAndName(ctx, arg) } -func (q *querier) GetUserSpendSince(ctx context.Context, arg database.GetUserSpendSinceParams) (database.GetUserSpendSinceRow, error) { - if _, err := q.GetUserByID(ctx, arg.UserID); err != nil { // AuthZ check - return database.GetUserSpendSinceRow{}, err - } - return q.db.GetUserSpendSince(ctx, arg) -} - func (q *querier) GetUserStatusCounts(ctx context.Context, arg database.GetUserStatusCountsParams) ([]database.GetUserStatusCountsRow, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUser); err != nil { return nil, err @@ -8910,6 +8910,20 @@ func (q *querier) UpsertUserAIBudgetOverride(ctx context.Context, arg database.U return q.db.UpsertUserAIBudgetOverride(ctx, arg) } +func (q *querier) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { + // Reject negative deltas. The schema CHECK only fires when the result + // goes negative, so a negative cost_micros on an existing row would + // silently reduce accumulated spend and corrupt budget enforcement. + if arg.CostMicros < 0 { + return database.AIUserDailySpend{}, xerrors.Errorf("cost_micros must be non-negative, got %d", arg.CostMicros) + } + // Daily spend writes are made by the aibridged process. + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { + return database.AIUserDailySpend{}, err + } + return q.db.UpsertUserAIDailySpend(ctx, arg) +} + func (q *querier) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAIProviderKey, error) { u, err := q.db.GetUserByID(ctx, arg.UserID) if err != nil { @@ -8943,20 +8957,6 @@ func (q *querier) UpsertUserChatPersonalModelOverride(ctx context.Context, arg d return q.db.UpsertUserChatPersonalModelOverride(ctx, arg) } -func (q *querier) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { - // Reject negative deltas. The schema CHECK only fires when the result - // goes negative, so a negative cost_micros on an existing row would - // silently reduce accumulated spend and corrupt budget enforcement. - if arg.CostMicros < 0 { - return database.AIUserDailySpend{}, xerrors.Errorf("cost_micros must be non-negative, got %d", arg.CostMicros) - } - // Daily spend writes are made by the aibridged process. - if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { - return database.AIUserDailySpend{}, err - } - return q.db.UpsertUserDailySpend(ctx, arg) -} - func (q *querier) UpsertWebpushVAPIDKeys(ctx context.Context, arg database.UpsertWebpushVAPIDKeysParams) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 5e0076d8a38..79ddf7a3bf2 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6802,33 +6802,33 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(user.ID).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override) })) - s.Run("GetUserSpendSince", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + s.Run("GetUserAISpendSince", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { user := testutil.Fake(s.T(), faker, database.User{}) - arg := database.GetUserSpendSinceParams{ + arg := database.GetUserAISpendSinceParams{ UserID: user.ID, EffectiveGroupID: uuid.New(), PeriodStart: time.Now().UTC().Truncate(24 * time.Hour), } - row := testutil.Fake(s.T(), faker, database.GetUserSpendSinceRow{UserID: user.ID, EffectiveGroupID: arg.EffectiveGroupID, PeriodStart: arg.PeriodStart}) + 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().GetUserSpendSince(gomock.Any(), arg).Return(row, nil).AnyTimes() + dbm.EXPECT().GetUserAISpendSince(gomock.Any(), arg).Return(row, nil).AnyTimes() check.Args(arg).Asserts(user, policy.ActionRead).Returns(row) })) - s.Run("UpsertUserDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - arg := database.UpsertUserDailySpendParams{ + s.Run("UpsertUserAIDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.UpsertUserAIDailySpendParams{ 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().UpsertUserDailySpend(gomock.Any(), arg).Return(row, nil).AnyTimes() + dbm.EXPECT().UpsertUserAIDailySpend(gomock.Any(), arg).Return(row, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceAibridgeInterception, policy.ActionUpdate).Returns(row) })) - s.Run("NegativeCost/UpsertUserDailySpend", s.Mocked(func(_ *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - arg := database.UpsertUserDailySpendParams{ + s.Run("NegativeCost/UpsertUserAIDailySpend", s.Mocked(func(_ *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.UpsertUserAIDailySpendParams{ UserID: uuid.New(), EffectiveGroupID: uuid.New(), Day: time.Now().UTC().Truncate(24 * time.Hour), diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index d526862062a..5935dbac0ef 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3137,6 +3137,14 @@ func (m queryMetricsStore) GetUserAISeatStates(ctx context.Context, userIds []uu return r0, r1 } +func (m queryMetricsStore) GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) { + start := time.Now() + r0, r1 := m.s.GetUserAISpendSince(ctx, arg) + m.queryLatencies.WithLabelValues("GetUserAISpendSince").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserAISpendSince").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetUserActivityInsights(ctx context.Context, arg database.GetUserActivityInsightsParams) ([]database.GetUserActivityInsightsRow, error) { start := time.Now() r0, r1 := m.s.GetUserActivityInsights(ctx, arg) @@ -3321,14 +3329,6 @@ func (m queryMetricsStore) GetUserSkillByUserIDAndName(ctx context.Context, arg return r0, r1 } -func (m queryMetricsStore) GetUserSpendSince(ctx context.Context, arg database.GetUserSpendSinceParams) (database.GetUserSpendSinceRow, error) { - start := time.Now() - r0, r1 := m.s.GetUserSpendSince(ctx, arg) - m.queryLatencies.WithLabelValues("GetUserSpendSince").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserSpendSince").Inc() - return r0, r1 -} - func (m queryMetricsStore) GetUserStatusCounts(ctx context.Context, arg database.GetUserStatusCountsParams) ([]database.GetUserStatusCountsRow, error) { start := time.Now() r0, r1 := m.s.GetUserStatusCounts(ctx, arg) @@ -6417,6 +6417,14 @@ func (m queryMetricsStore) UpsertUserAIBudgetOverride(ctx context.Context, arg d return r0, r1 } +func (m queryMetricsStore) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { + start := time.Now() + r0, r1 := m.s.UpsertUserAIDailySpend(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertUserAIDailySpend").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertUserAIDailySpend").Inc() + return r0, r1 +} + func (m queryMetricsStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAIProviderKey, error) { start := time.Now() r0, r1 := m.s.UpsertUserAIProviderKey(ctx, arg) @@ -6441,14 +6449,6 @@ func (m queryMetricsStore) UpsertUserChatPersonalModelOverride(ctx context.Conte return r0 } -func (m queryMetricsStore) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { - start := time.Now() - r0, r1 := m.s.UpsertUserDailySpend(ctx, arg) - m.queryLatencies.WithLabelValues("UpsertUserDailySpend").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertUserDailySpend").Inc() - return r0, r1 -} - func (m queryMetricsStore) UpsertWebpushVAPIDKeys(ctx context.Context, arg database.UpsertWebpushVAPIDKeysParams) error { start := time.Now() r0 := m.s.UpsertWebpushVAPIDKeys(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 4ddc598661a..e440bb99518 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5862,6 +5862,21 @@ func (mr *MockStoreMockRecorder) GetUserAISeatStates(ctx, userIds any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserAISeatStates", reflect.TypeOf((*MockStore)(nil).GetUserAISeatStates), ctx, userIds) } +// GetUserAISpendSince mocks base method. +func (m *MockStore) GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetUserAISpendSince", ctx, arg) + ret0, _ := ret[0].(database.GetUserAISpendSinceRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetUserAISpendSince indicates an expected call of GetUserAISpendSince. +func (mr *MockStoreMockRecorder) GetUserAISpendSince(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserAISpendSince", reflect.TypeOf((*MockStore)(nil).GetUserAISpendSince), ctx, arg) +} + // GetUserActivityInsights mocks base method. func (m *MockStore) GetUserActivityInsights(ctx context.Context, arg database.GetUserActivityInsightsParams) ([]database.GetUserActivityInsightsRow, error) { m.ctrl.T.Helper() @@ -6207,21 +6222,6 @@ func (mr *MockStoreMockRecorder) GetUserSkillByUserIDAndName(ctx, arg any) *gomo return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserSkillByUserIDAndName", reflect.TypeOf((*MockStore)(nil).GetUserSkillByUserIDAndName), ctx, arg) } -// GetUserSpendSince mocks base method. -func (m *MockStore) GetUserSpendSince(ctx context.Context, arg database.GetUserSpendSinceParams) (database.GetUserSpendSinceRow, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetUserSpendSince", ctx, arg) - ret0, _ := ret[0].(database.GetUserSpendSinceRow) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetUserSpendSince indicates an expected call of GetUserSpendSince. -func (mr *MockStoreMockRecorder) GetUserSpendSince(ctx, arg any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserSpendSince", reflect.TypeOf((*MockStore)(nil).GetUserSpendSince), ctx, arg) -} - // GetUserStatusCounts mocks base method. func (m *MockStore) GetUserStatusCounts(ctx context.Context, arg database.GetUserStatusCountsParams) ([]database.GetUserStatusCountsRow, error) { m.ctrl.T.Helper() @@ -11996,6 +11996,21 @@ func (mr *MockStoreMockRecorder) UpsertUserAIBudgetOverride(ctx, arg any) *gomoc return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).UpsertUserAIBudgetOverride), ctx, arg) } +// UpsertUserAIDailySpend mocks base method. +func (m *MockStore) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertUserAIDailySpend", ctx, arg) + ret0, _ := ret[0].(database.AIUserDailySpend) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UpsertUserAIDailySpend indicates an expected call of UpsertUserAIDailySpend. +func (mr *MockStoreMockRecorder) UpsertUserAIDailySpend(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserAIDailySpend", reflect.TypeOf((*MockStore)(nil).UpsertUserAIDailySpend), ctx, arg) +} + // UpsertUserAIProviderKey mocks base method. func (m *MockStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAIProviderKey, error) { m.ctrl.T.Helper() @@ -12039,21 +12054,6 @@ func (mr *MockStoreMockRecorder) UpsertUserChatPersonalModelOverride(ctx, arg an return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserChatPersonalModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertUserChatPersonalModelOverride), ctx, arg) } -// UpsertUserDailySpend mocks base method. -func (m *MockStore) UpsertUserDailySpend(ctx context.Context, arg database.UpsertUserDailySpendParams) (database.AIUserDailySpend, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpsertUserDailySpend", ctx, arg) - ret0, _ := ret[0].(database.AIUserDailySpend) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// UpsertUserDailySpend indicates an expected call of UpsertUserDailySpend. -func (mr *MockStoreMockRecorder) UpsertUserDailySpend(ctx, arg any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserDailySpend", reflect.TypeOf((*MockStore)(nil).UpsertUserDailySpend), ctx, arg) -} - // UpsertWebpushVAPIDKeys mocks base method. func (m *MockStore) UpsertWebpushVAPIDKeys(ctx context.Context, arg database.UpsertWebpushVAPIDKeysParams) error { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index f44489097d9..0f56276d148 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -813,6 +813,9 @@ type sqlcQuerier interface { // Filters to active, non-deleted, non-system users to match the canonical // seat count query (GetActiveAISeatCount). GetUserAISeatStates(ctx context.Context, userIds []uuid.UUID) ([]uuid.UUID, error) + // Total spend for (user_id, effective_group_id) on or after period_start. + // The period_start parameter is normalized to its UTC calendar day. + GetUserAISpendSince(ctx context.Context, arg GetUserAISpendSinceParams) (GetUserAISpendSinceRow, error) // GetUserActivityInsights returns the ranking with top active users. // The result can be filtered on template_ids, meaning only user data // from workspaces based on those templates will be included. @@ -882,9 +885,6 @@ type sqlcQuerier interface { GetUserSecretsTelemetrySummary(ctx context.Context) (GetUserSecretsTelemetrySummaryRow, error) GetUserShellToolDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) GetUserSkillByUserIDAndName(ctx context.Context, arg GetUserSkillByUserIDAndNameParams) (UserSkill, error) - // Total spend for (user_id, effective_group_id) on or after period_start. - // The period_start parameter is normalized to its UTC calendar day. - GetUserSpendSince(ctx context.Context, arg GetUserSpendSinceParams) (GetUserSpendSinceRow, error) // GetUserStatusCounts returns the count of users in each status over time. // The time range is inclusively defined by the start_time and end_time parameters. GetUserStatusCounts(ctx context.Context, arg GetUserStatusCountsParams) ([]GetUserStatusCountsRow, error) @@ -1570,15 +1570,15 @@ type sqlcQuerier interface { // combination. The result is stored in the template_usage_stats table. UpsertTemplateUsageStats(ctx context.Context) error UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAIBudgetOverride, error) + // Adds cost_micros to the spend for (user_id, effective_group_id, day). + // The day parameter is normalized to its UTC calendar day before storage. + UpsertUserAIDailySpend(ctx context.Context, arg UpsertUserAIDailySpendParams) (AIUserDailySpend, error) // UpsertUserAIProviderKey preserves the original id and created_at when the // user/provider pair already exists. On conflict, callers provide id and // created_at for the insert path only. UpsertUserAIProviderKey(ctx context.Context, arg UpsertUserAIProviderKeyParams) (UserAIProviderKey, error) UpsertUserChatDebugLoggingEnabled(ctx context.Context, arg UpsertUserChatDebugLoggingEnabledParams) error UpsertUserChatPersonalModelOverride(ctx context.Context, arg UpsertUserChatPersonalModelOverrideParams) error - // Adds cost_micros to the spend for (user_id, effective_group_id, day). - // The day parameter is normalized to its UTC calendar day before storage. - UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) (AIUserDailySpend, error) UpsertWebpushVAPIDKeys(ctx context.Context, arg UpsertWebpushVAPIDKeysParams) error UpsertWorkspaceAgentContextResource(ctx context.Context, arg UpsertWorkspaceAgentContextResourceParams) (WorkspaceAgentContextResource, error) UpsertWorkspaceAgentContextSnapshot(ctx context.Context, arg UpsertWorkspaceAgentContextSnapshotParams) (WorkspaceAgentContextSnapshot, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index eed8119e9c3..e02e20924bb 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11751,7 +11751,7 @@ func TestUpsertAISeats(t *testing.T) { require.False(t, alreadyExists) } -func TestUpsertUserDailySpend(t *testing.T) { +func TestUpsertUserAIDailySpend(t *testing.T) { t.Parallel() // Use fixed dates to keep the test deterministic. @@ -11783,7 +11783,7 @@ func TestUpsertUserDailySpend(t *testing.T) { var row database.AIUserDailySpend var err error for _, cost := range tt.costs { - row, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + row, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, @@ -11816,13 +11816,13 @@ func TestUpsertUserDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - userARow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + userARow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: userA.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) require.Equal(t, int64(100), userARow.SpendMicros) - userBRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + userBRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: userB.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 25, }) require.NoError(t, err) @@ -11840,13 +11840,13 @@ func TestUpsertUserDailySpend(t *testing.T) { groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - groupARow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + groupARow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupA.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) require.Equal(t, int64(100), groupARow.SpendMicros) - groupBRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + groupBRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupB.ID, Day: day, CostMicros: 25, }) require.NoError(t, err) @@ -11863,7 +11863,7 @@ func TestUpsertUserDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dayRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + dayRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) @@ -11871,7 +11871,7 @@ func TestUpsertUserDailySpend(t *testing.T) { // The ON CONFLICT target is the full PK including day, so this upsert // cannot modify the previous day's row by construction. - nextDayRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + nextDayRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: nextDay, CostMicros: 25, }) require.NoError(t, err) @@ -11889,13 +11889,13 @@ func TestUpsertUserDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) dayNonTruncated := day.Add(14*time.Hour + 30*time.Minute) - nonTruncatedRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + nonTruncatedRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: dayNonTruncated, CostMicros: 50, }) require.NoError(t, err) @@ -11916,7 +11916,7 @@ func TestUpsertUserDailySpend(t *testing.T) { // 2024-06-15 23:00 in UTC-5 is 2024-06-16 04:00 UTC, so this should land on nextDay (2024-06-16). localLate := time.Date(2024, 6, 15, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) - nonUTCRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + nonUTCRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: localLate, CostMicros: 100, }) require.NoError(t, err) @@ -11934,21 +11934,21 @@ func TestUpsertUserDailySpend(t *testing.T) { group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) // Zero-cost upsert on a fresh key creates a row with spend = 0. - newRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + newRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 0, }) require.NoError(t, err) require.Equal(t, int64(0), newRow.SpendMicros) // After a real upsert, the row has spend = 100. - updatedRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + updatedRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) require.Equal(t, int64(100), updatedRow.SpendMicros) // Zero-cost upsert on the existing row leaves spend unchanged. - sameRow, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + sameRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 0, }) require.NoError(t, err) @@ -11957,7 +11957,7 @@ func TestUpsertUserDailySpend(t *testing.T) { }) } -func TestGetUserSpendSince(t *testing.T) { +func TestGetUserAISpendSince(t *testing.T) { t.Parallel() // Use fixed dates to keep the test deterministic. @@ -11994,7 +11994,7 @@ func TestGetUserSpendSince(t *testing.T) { group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) for _, r := range tt.rows { - _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: r.day, @@ -12003,7 +12003,7 @@ func TestGetUserSpendSince(t *testing.T) { require.NoError(t, err) } - got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ UserID: user.ID, EffectiveGroupID: group.ID, PeriodStart: monthStart, @@ -12027,16 +12027,16 @@ func TestGetUserSpendSince(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: userA.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 100, }) require.NoError(t, err) - _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: userB.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 25, }) require.NoError(t, err) - got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ UserID: userB.ID, EffectiveGroupID: group.ID, PeriodStart: monthStart, @@ -12056,16 +12056,16 @@ func TestGetUserSpendSince(t *testing.T) { groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupA.ID, Day: today, CostMicros: 100, }) require.NoError(t, err) - _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupB.ID, Day: today, CostMicros: 25, }) require.NoError(t, err) - got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ UserID: user.ID, EffectiveGroupID: groupB.ID, PeriodStart: monthStart, @@ -12086,11 +12086,11 @@ func TestGetUserSpendSince(t *testing.T) { // Seed a row on prevMonthLastDay (which lies on May 31 UTC). A naive // query that does not normalize the period_start would include it. - _, err := db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999, }) require.NoError(t, err) - _, err = db.UpsertUserDailySpend(ctx, database.UpsertUserDailySpendParams{ + _, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25, }) require.NoError(t, err) @@ -12098,7 +12098,7 @@ func TestGetUserSpendSince(t *testing.T) { // 2024-05-31 23:00 in UTC-5 is 2024-06-01 04:00 UTC, so the // normalized period_start lands on June 1. localLate := time.Date(2024, 5, 31, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) - got, err := db.GetUserSpendSince(ctx, database.GetUserSpendSinceParams{ + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ UserID: user.ID, EffectiveGroupID: group.ID, PeriodStart: localLate, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e9d50f59290..9df67f7778f 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2546,7 +2546,7 @@ func (q *sqlQuerier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UU return i, err } -const getUserSpendSince = `-- name: GetUserSpendSince :one +const getUserAISpendSince = `-- name: GetUserAISpendSince :one SELECT $1::uuid AS user_id, $2::uuid AS effective_group_id, @@ -2558,13 +2558,13 @@ WHERE user_id = $1 AND day >= (($3::timestamptz) AT TIME ZONE 'UTC')::date ` -type GetUserSpendSinceParams struct { +type GetUserAISpendSinceParams struct { UserID uuid.UUID `db:"user_id" json:"user_id"` EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` PeriodStart time.Time `db:"period_start" json:"period_start"` } -type GetUserSpendSinceRow struct { +type GetUserAISpendSinceRow struct { UserID uuid.UUID `db:"user_id" json:"user_id"` EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` PeriodStart time.Time `db:"period_start" json:"period_start"` @@ -2573,9 +2573,9 @@ type GetUserSpendSinceRow struct { // Total spend for (user_id, effective_group_id) on or after period_start. // The period_start parameter is normalized to its UTC calendar day. -func (q *sqlQuerier) GetUserSpendSince(ctx context.Context, arg GetUserSpendSinceParams) (GetUserSpendSinceRow, error) { - row := q.db.QueryRowContext(ctx, getUserSpendSince, arg.UserID, arg.EffectiveGroupID, arg.PeriodStart) - var i GetUserSpendSinceRow +func (q *sqlQuerier) GetUserAISpendSince(ctx context.Context, arg GetUserAISpendSinceParams) (GetUserAISpendSinceRow, error) { + row := q.db.QueryRowContext(ctx, getUserAISpendSince, arg.UserID, arg.EffectiveGroupID, arg.PeriodStart) + var i GetUserAISpendSinceRow err := row.Scan( &i.UserID, &i.EffectiveGroupID, @@ -2668,7 +2668,7 @@ func (q *sqlQuerier) UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertU return i, err } -const upsertUserDailySpend = `-- name: UpsertUserDailySpend :one +const upsertUserAIDailySpend = `-- name: UpsertUserAIDailySpend :one INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4) ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET @@ -2676,7 +2676,7 @@ ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET RETURNING user_id, effective_group_id, day, spend_micros ` -type UpsertUserDailySpendParams struct { +type UpsertUserAIDailySpendParams struct { UserID uuid.UUID `db:"user_id" json:"user_id"` EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` Day time.Time `db:"day" json:"day"` @@ -2685,8 +2685,8 @@ type UpsertUserDailySpendParams struct { // Adds cost_micros to the spend for (user_id, effective_group_id, day). // The day parameter is normalized to its UTC calendar day before storage. -func (q *sqlQuerier) UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) (AIUserDailySpend, error) { - row := q.db.QueryRowContext(ctx, upsertUserDailySpend, +func (q *sqlQuerier) UpsertUserAIDailySpend(ctx context.Context, arg UpsertUserAIDailySpendParams) (AIUserDailySpend, error) { + row := q.db.QueryRowContext(ctx, upsertUserAIDailySpend, arg.UserID, arg.EffectiveGroupID, arg.Day, diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 46c421db85a..f64ce1b9a5b 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -80,7 +80,7 @@ ORDER BY gaib.group_id ASC LIMIT 1; --- name: UpsertUserDailySpend :one +-- name: UpsertUserAIDailySpend :one -- Adds cost_micros to the spend for (user_id, effective_group_id, day). -- The day parameter is normalized to its UTC calendar day before storage. INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) @@ -89,7 +89,7 @@ ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros RETURNING *; --- name: GetUserSpendSince :one +-- name: GetUserAISpendSince :one -- Total spend for (user_id, effective_group_id) on or after period_start. -- The period_start parameter is normalized to its UTC calendar day. SELECT From 27efdd82aff8317ebee407420010d4279e29b7fb Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 23 Jun 2026 16:25:30 +0000 Subject: [PATCH 4/7] chore: remove negative cost check in dbauthz --- coderd/database/dbauthz/dbauthz.go | 6 ------ coderd/database/dbauthz/dbauthz_test.go | 10 ---------- 2 files changed, 16 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 974ff34c1dc..eab72fbe869 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -8911,12 +8911,6 @@ func (q *querier) UpsertUserAIBudgetOverride(ctx context.Context, arg database.U } func (q *querier) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { - // Reject negative deltas. The schema CHECK only fires when the result - // goes negative, so a negative cost_micros on an existing row would - // silently reduce accumulated spend and corrupt budget enforcement. - if arg.CostMicros < 0 { - return database.AIUserDailySpend{}, xerrors.Errorf("cost_micros must be non-negative, got %d", arg.CostMicros) - } // Daily spend writes are made by the aibridged process. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { return database.AIUserDailySpend{}, err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 79ddf7a3bf2..234c0954a02 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6827,16 +6827,6 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(arg).Asserts(rbac.ResourceAibridgeInterception, policy.ActionUpdate).Returns(row) })) - s.Run("NegativeCost/UpsertUserAIDailySpend", s.Mocked(func(_ *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - arg := database.UpsertUserAIDailySpendParams{ - UserID: uuid.New(), - EffectiveGroupID: uuid.New(), - Day: time.Now().UTC().Truncate(24 * time.Hour), - CostMicros: -100, - } - check.Args(arg).Asserts().Errors(xerrors.New("cost_micros must be non-negative, got -100")) - })) - 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() From dbdf64f1a08c162b3380eb32265535f2cecad409 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 24 Jun 2026 15:16:24 +0000 Subject: [PATCH 5/7] chore: address comments --- coderd/database/dbauthz/dbauthz.go | 15 +++-- coderd/database/dbauthz/dbauthz_test.go | 6 +- coderd/database/dbmetrics/querymetrics.go | 16 +++--- coderd/database/dbmock/dbmock.go | 30 +++++----- coderd/database/querier.go | 6 +- coderd/database/querier_test.go | 43 +++++++------- coderd/database/queries.sql.go | 68 +++++++++++------------ coderd/database/queries/aicostcontrol.sql | 2 +- 8 files changed, 93 insertions(+), 93 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index eab72fbe869..2fbb0cd6534 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5733,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) } @@ -8910,14 +8917,6 @@ func (q *querier) UpsertUserAIBudgetOverride(ctx context.Context, arg database.U return q.db.UpsertUserAIBudgetOverride(ctx, arg) } -func (q *querier) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { - // Daily spend writes are made by the aibridged process. - if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { - return database.AIUserDailySpend{}, err - } - return q.db.UpsertUserAIDailySpend(ctx, arg) -} - func (q *querier) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAIProviderKey, error) { u, err := q.db.GetUserByID(ctx, arg.UserID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 234c0954a02..550ce623aaf 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6815,15 +6815,15 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(arg).Asserts(user, policy.ActionRead).Returns(row) })) - s.Run("UpsertUserAIDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - arg := database.UpsertUserAIDailySpendParams{ + 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().UpsertUserAIDailySpend(gomock.Any(), arg).Return(row, nil).AnyTimes() + dbm.EXPECT().IncrementUserAIDailySpend(gomock.Any(), arg).Return(row, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceAibridgeInterception, policy.ActionUpdate).Returns(row) })) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 5935dbac0ef..5afa5fcdc12 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3897,6 +3897,14 @@ func (m queryMetricsStore) IncrementChatGenerationAttempt(ctx context.Context, i return r0, r1 } +func (m queryMetricsStore) IncrementUserAIDailySpend(ctx context.Context, arg database.IncrementUserAIDailySpendParams) (database.AIUserDailySpend, error) { + start := time.Now() + r0, r1 := m.s.IncrementUserAIDailySpend(ctx, arg) + m.queryLatencies.WithLabelValues("IncrementUserAIDailySpend").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "IncrementUserAIDailySpend").Inc() + return r0, r1 +} + func (m queryMetricsStore) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) { start := time.Now() r0, r1 := m.s.InsertAIBridgeInterception(ctx, arg) @@ -6417,14 +6425,6 @@ func (m queryMetricsStore) UpsertUserAIBudgetOverride(ctx context.Context, arg d return r0, r1 } -func (m queryMetricsStore) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { - start := time.Now() - r0, r1 := m.s.UpsertUserAIDailySpend(ctx, arg) - m.queryLatencies.WithLabelValues("UpsertUserAIDailySpend").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertUserAIDailySpend").Inc() - return r0, r1 -} - func (m queryMetricsStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAIProviderKey, error) { start := time.Now() r0, r1 := m.s.UpsertUserAIProviderKey(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index e440bb99518..97eacb11fa3 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -7300,6 +7300,21 @@ func (mr *MockStoreMockRecorder) IncrementChatGenerationAttempt(ctx, id any) *go return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementChatGenerationAttempt", reflect.TypeOf((*MockStore)(nil).IncrementChatGenerationAttempt), ctx, id) } +// IncrementUserAIDailySpend mocks base method. +func (m *MockStore) IncrementUserAIDailySpend(ctx context.Context, arg database.IncrementUserAIDailySpendParams) (database.AIUserDailySpend, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IncrementUserAIDailySpend", ctx, arg) + ret0, _ := ret[0].(database.AIUserDailySpend) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// IncrementUserAIDailySpend indicates an expected call of IncrementUserAIDailySpend. +func (mr *MockStoreMockRecorder) IncrementUserAIDailySpend(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementUserAIDailySpend", reflect.TypeOf((*MockStore)(nil).IncrementUserAIDailySpend), ctx, arg) +} + // InsertAIBridgeInterception mocks base method. func (m *MockStore) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) { m.ctrl.T.Helper() @@ -11996,21 +12011,6 @@ func (mr *MockStoreMockRecorder) UpsertUserAIBudgetOverride(ctx, arg any) *gomoc return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).UpsertUserAIBudgetOverride), ctx, arg) } -// UpsertUserAIDailySpend mocks base method. -func (m *MockStore) UpsertUserAIDailySpend(ctx context.Context, arg database.UpsertUserAIDailySpendParams) (database.AIUserDailySpend, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpsertUserAIDailySpend", ctx, arg) - ret0, _ := ret[0].(database.AIUserDailySpend) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// UpsertUserAIDailySpend indicates an expected call of UpsertUserAIDailySpend. -func (mr *MockStoreMockRecorder) UpsertUserAIDailySpend(ctx, arg any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserAIDailySpend", reflect.TypeOf((*MockStore)(nil).UpsertUserAIDailySpend), ctx, arg) -} - // UpsertUserAIProviderKey mocks base method. func (m *MockStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAIProviderKey, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 0f56276d148..0490ae91809 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -998,6 +998,9 @@ type sqlcQuerier interface { HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) error // Increments generation_attempt and returns the resulting value. IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) + // Adds cost_micros to the spend for (user_id, effective_group_id, day). + // The day parameter is normalized to its UTC calendar day before storage. + IncrementUserAIDailySpend(ctx context.Context, arg IncrementUserAIDailySpendParams) (AIUserDailySpend, error) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) InsertAIBridgeModelThought(ctx context.Context, arg InsertAIBridgeModelThoughtParams) (AIBridgeModelThought, error) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) @@ -1570,9 +1573,6 @@ type sqlcQuerier interface { // combination. The result is stored in the template_usage_stats table. UpsertTemplateUsageStats(ctx context.Context) error UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAIBudgetOverride, error) - // Adds cost_micros to the spend for (user_id, effective_group_id, day). - // The day parameter is normalized to its UTC calendar day before storage. - UpsertUserAIDailySpend(ctx context.Context, arg UpsertUserAIDailySpendParams) (AIUserDailySpend, error) // UpsertUserAIProviderKey preserves the original id and created_at when the // user/provider pair already exists. On conflict, callers provide id and // created_at for the insert path only. diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index e02e20924bb..629e633c0c7 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11751,7 +11751,7 @@ func TestUpsertAISeats(t *testing.T) { require.False(t, alreadyExists) } -func TestUpsertUserAIDailySpend(t *testing.T) { +func TestIncrementUserAIDailySpend(t *testing.T) { t.Parallel() // Use fixed dates to keep the test deterministic. @@ -11783,7 +11783,7 @@ func TestUpsertUserAIDailySpend(t *testing.T) { var row database.AIUserDailySpend var err error for _, cost := range tt.costs { - row, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + row, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, @@ -11795,6 +11795,7 @@ func TestUpsertUserAIDailySpend(t *testing.T) { } if tt.wantErr { require.Error(t, err) + require.True(t, database.IsCheckViolation(err, database.CheckAIUserDailySpendSpendMicrosCheck)) return } require.NoError(t, err) @@ -11816,13 +11817,13 @@ func TestUpsertUserAIDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - userARow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + userARow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: userA.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) require.Equal(t, int64(100), userARow.SpendMicros) - userBRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + userBRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: userB.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 25, }) require.NoError(t, err) @@ -11840,13 +11841,13 @@ func TestUpsertUserAIDailySpend(t *testing.T) { groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - groupARow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + groupARow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupA.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) require.Equal(t, int64(100), groupARow.SpendMicros) - groupBRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + groupBRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupB.ID, Day: day, CostMicros: 25, }) require.NoError(t, err) @@ -11863,7 +11864,7 @@ func TestUpsertUserAIDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - dayRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + dayRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) @@ -11871,7 +11872,7 @@ func TestUpsertUserAIDailySpend(t *testing.T) { // The ON CONFLICT target is the full PK including day, so this upsert // cannot modify the previous day's row by construction. - nextDayRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + nextDayRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: nextDay, CostMicros: 25, }) require.NoError(t, err) @@ -11889,13 +11890,13 @@ func TestUpsertUserAIDailySpend(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) dayNonTruncated := day.Add(14*time.Hour + 30*time.Minute) - nonTruncatedRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + nonTruncatedRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: dayNonTruncated, CostMicros: 50, }) require.NoError(t, err) @@ -11916,7 +11917,7 @@ func TestUpsertUserAIDailySpend(t *testing.T) { // 2024-06-15 23:00 in UTC-5 is 2024-06-16 04:00 UTC, so this should land on nextDay (2024-06-16). localLate := time.Date(2024, 6, 15, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) - nonUTCRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + nonUTCRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: localLate, CostMicros: 100, }) require.NoError(t, err) @@ -11934,21 +11935,21 @@ func TestUpsertUserAIDailySpend(t *testing.T) { group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) // Zero-cost upsert on a fresh key creates a row with spend = 0. - newRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + newRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 0, }) require.NoError(t, err) require.Equal(t, int64(0), newRow.SpendMicros) // After a real upsert, the row has spend = 100. - updatedRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + updatedRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 100, }) require.NoError(t, err) require.Equal(t, int64(100), updatedRow.SpendMicros) // Zero-cost upsert on the existing row leaves spend unchanged. - sameRow, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + sameRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: 0, }) require.NoError(t, err) @@ -11994,7 +11995,7 @@ func TestGetUserAISpendSince(t *testing.T) { group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) for _, r := range tt.rows { - _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: r.day, @@ -12027,11 +12028,11 @@ func TestGetUserAISpendSince(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: userA.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 100, }) require.NoError(t, err) - _, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: userB.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 25, }) require.NoError(t, err) @@ -12056,11 +12057,11 @@ func TestGetUserAISpendSince(t *testing.T) { groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID}) - _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupA.ID, Day: today, CostMicros: 100, }) require.NoError(t, err) - _, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: groupB.ID, Day: today, CostMicros: 25, }) require.NoError(t, err) @@ -12086,11 +12087,11 @@ func TestGetUserAISpendSince(t *testing.T) { // Seed a row on prevMonthLastDay (which lies on May 31 UTC). A naive // query that does not normalize the period_start would include it. - _, err := db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999, }) require.NoError(t, err) - _, err = db.UpsertUserAIDailySpend(ctx, database.UpsertUserAIDailySpendParams{ + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25, }) require.NoError(t, err) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 9df67f7778f..4e47103a0d7 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2585,6 +2585,40 @@ func (q *sqlQuerier) GetUserAISpendSince(ctx context.Context, arg GetUserAISpend return i, err } +const incrementUserAIDailySpend = `-- name: IncrementUserAIDailySpend :one +INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) +VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4) +ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET + spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros +RETURNING user_id, effective_group_id, day, spend_micros +` + +type IncrementUserAIDailySpendParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` + Day time.Time `db:"day" json:"day"` + CostMicros int64 `db:"cost_micros" json:"cost_micros"` +} + +// Adds cost_micros to the spend for (user_id, effective_group_id, day). +// The day parameter is normalized to its UTC calendar day before storage. +func (q *sqlQuerier) IncrementUserAIDailySpend(ctx context.Context, arg IncrementUserAIDailySpendParams) (AIUserDailySpend, error) { + row := q.db.QueryRowContext(ctx, incrementUserAIDailySpend, + arg.UserID, + arg.EffectiveGroupID, + arg.Day, + arg.CostMicros, + ) + var i AIUserDailySpend + err := row.Scan( + &i.UserID, + &i.EffectiveGroupID, + &i.Day, + &i.SpendMicros, + ) + return i, err +} + const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec INSERT INTO ai_model_prices ( provider, model, input_price, output_price, cache_read_price, cache_write_price @@ -2668,40 +2702,6 @@ func (q *sqlQuerier) UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertU return i, err } -const upsertUserAIDailySpend = `-- name: UpsertUserAIDailySpend :one -INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) -VALUES ($1, $2, (($3::timestamptz) AT TIME ZONE 'UTC')::date, $4) -ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET - spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros -RETURNING user_id, effective_group_id, day, spend_micros -` - -type UpsertUserAIDailySpendParams struct { - UserID uuid.UUID `db:"user_id" json:"user_id"` - EffectiveGroupID uuid.UUID `db:"effective_group_id" json:"effective_group_id"` - Day time.Time `db:"day" json:"day"` - CostMicros int64 `db:"cost_micros" json:"cost_micros"` -} - -// Adds cost_micros to the spend for (user_id, effective_group_id, day). -// The day parameter is normalized to its UTC calendar day before storage. -func (q *sqlQuerier) UpsertUserAIDailySpend(ctx context.Context, arg UpsertUserAIDailySpendParams) (AIUserDailySpend, error) { - row := q.db.QueryRowContext(ctx, upsertUserAIDailySpend, - arg.UserID, - arg.EffectiveGroupID, - arg.Day, - arg.CostMicros, - ) - var i AIUserDailySpend - err := row.Scan( - &i.UserID, - &i.EffectiveGroupID, - &i.Day, - &i.SpendMicros, - ) - return i, err -} - const getActiveAISeatCount = `-- name: GetActiveAISeatCount :one SELECT COUNT(*) diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index f64ce1b9a5b..9dd6e89ce1c 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -80,7 +80,7 @@ ORDER BY gaib.group_id ASC LIMIT 1; --- name: UpsertUserAIDailySpend :one +-- name: IncrementUserAIDailySpend :one -- Adds cost_micros to the spend for (user_id, effective_group_id, day). -- The day parameter is normalized to its UTC calendar day before storage. INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros) From 2a0bf88bc8b5a4a66899d00241177ec070cd4b91 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 25 Jun 2026 09:44:48 +0000 Subject: [PATCH 6/7] chore: improve GetUserAISpendSince comment --- ...daily_spend.down.sql => 000531_ai_user_daily_spend.down.sql} | 0 ...ser_daily_spend.up.sql => 000531_ai_user_daily_spend.up.sql} | 0 ...ser_daily_spend.up.sql => 000531_ai_user_daily_spend.up.sql} | 0 coderd/database/querier.go | 2 +- coderd/database/queries.sql.go | 2 +- coderd/database/queries/aicostcontrol.sql | 2 +- 6 files changed, 3 insertions(+), 3 deletions(-) rename coderd/database/migrations/{000530_ai_user_daily_spend.down.sql => 000531_ai_user_daily_spend.down.sql} (100%) rename coderd/database/migrations/{000530_ai_user_daily_spend.up.sql => 000531_ai_user_daily_spend.up.sql} (100%) rename coderd/database/migrations/testdata/fixtures/{000530_ai_user_daily_spend.up.sql => 000531_ai_user_daily_spend.up.sql} (100%) diff --git a/coderd/database/migrations/000530_ai_user_daily_spend.down.sql b/coderd/database/migrations/000531_ai_user_daily_spend.down.sql similarity index 100% rename from coderd/database/migrations/000530_ai_user_daily_spend.down.sql rename to coderd/database/migrations/000531_ai_user_daily_spend.down.sql diff --git a/coderd/database/migrations/000530_ai_user_daily_spend.up.sql b/coderd/database/migrations/000531_ai_user_daily_spend.up.sql similarity index 100% rename from coderd/database/migrations/000530_ai_user_daily_spend.up.sql rename to coderd/database/migrations/000531_ai_user_daily_spend.up.sql diff --git a/coderd/database/migrations/testdata/fixtures/000530_ai_user_daily_spend.up.sql b/coderd/database/migrations/testdata/fixtures/000531_ai_user_daily_spend.up.sql similarity index 100% rename from coderd/database/migrations/testdata/fixtures/000530_ai_user_daily_spend.up.sql rename to coderd/database/migrations/testdata/fixtures/000531_ai_user_daily_spend.up.sql diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 0490ae91809..0576f50e65b 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -813,7 +813,7 @@ type sqlcQuerier interface { // Filters to active, non-deleted, non-system users to match the canonical // seat count query (GetActiveAISeatCount). GetUserAISeatStates(ctx context.Context, userIds []uuid.UUID) ([]uuid.UUID, error) - // Total spend for (user_id, effective_group_id) on or after period_start. + // Total spend for (user_id, effective_group_id) on or after period_start until NOW. // The period_start parameter is normalized to its UTC calendar day. GetUserAISpendSince(ctx context.Context, arg GetUserAISpendSinceParams) (GetUserAISpendSinceRow, error) // GetUserActivityInsights returns the ranking with top active users. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 4e47103a0d7..4a5c315c63e 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2571,7 +2571,7 @@ type GetUserAISpendSinceRow struct { SpendMicros int64 `db:"spend_micros" json:"spend_micros"` } -// Total spend for (user_id, effective_group_id) on or after period_start. +// Total spend for (user_id, effective_group_id) on or after period_start until NOW. // The period_start parameter is normalized to its UTC calendar day. func (q *sqlQuerier) GetUserAISpendSince(ctx context.Context, arg GetUserAISpendSinceParams) (GetUserAISpendSinceRow, error) { row := q.db.QueryRowContext(ctx, getUserAISpendSince, arg.UserID, arg.EffectiveGroupID, arg.PeriodStart) diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 9dd6e89ce1c..00c4917371e 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -90,7 +90,7 @@ ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET RETURNING *; -- name: GetUserAISpendSince :one --- Total spend for (user_id, effective_group_id) on or after period_start. +-- Total spend for (user_id, effective_group_id) on or after period_start until NOW. -- The period_start parameter is normalized to its UTC calendar day. SELECT @user_id::uuid AS user_id, From ae35c716804ee6bc96c4883e4078eab69e9d5592 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Mon, 29 Jun 2026 13:31:15 +0000 Subject: [PATCH 7/7] chore: fix migration number --- ...r_daily_spend.down.sql => 000536_ai_user_daily_spend.down.sql} | 0 ..._user_daily_spend.up.sql => 000536_ai_user_daily_spend.up.sql} | 0 ..._user_daily_spend.up.sql => 000536_ai_user_daily_spend.up.sql} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename coderd/database/migrations/{000531_ai_user_daily_spend.down.sql => 000536_ai_user_daily_spend.down.sql} (100%) rename coderd/database/migrations/{000531_ai_user_daily_spend.up.sql => 000536_ai_user_daily_spend.up.sql} (100%) rename coderd/database/migrations/testdata/fixtures/{000531_ai_user_daily_spend.up.sql => 000536_ai_user_daily_spend.up.sql} (100%) diff --git a/coderd/database/migrations/000531_ai_user_daily_spend.down.sql b/coderd/database/migrations/000536_ai_user_daily_spend.down.sql similarity index 100% rename from coderd/database/migrations/000531_ai_user_daily_spend.down.sql rename to coderd/database/migrations/000536_ai_user_daily_spend.down.sql diff --git a/coderd/database/migrations/000531_ai_user_daily_spend.up.sql b/coderd/database/migrations/000536_ai_user_daily_spend.up.sql similarity index 100% rename from coderd/database/migrations/000531_ai_user_daily_spend.up.sql rename to coderd/database/migrations/000536_ai_user_daily_spend.up.sql diff --git a/coderd/database/migrations/testdata/fixtures/000531_ai_user_daily_spend.up.sql b/coderd/database/migrations/testdata/fixtures/000536_ai_user_daily_spend.up.sql similarity index 100% rename from coderd/database/migrations/testdata/fixtures/000531_ai_user_daily_spend.up.sql rename to coderd/database/migrations/testdata/fixtures/000536_ai_user_daily_spend.up.sql