diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go index 76b350816d757..fda7bc467577d 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 20a688c768a1a..2fbb0cd65342b 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 { @@ -5726,6 +5733,13 @@ func (q *querier) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UU return q.db.IncrementChatGenerationAttempt(ctx, id) } +func (q *querier) IncrementUserAIDailySpend(ctx context.Context, arg database.IncrementUserAIDailySpendParams) (database.AIUserDailySpend, error) { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAibridgeInterception); err != nil { + return database.AIUserDailySpend{}, err + } + return q.db.IncrementUserAIDailySpend(ctx, arg) +} + func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) { return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 68d5d222bb6f0..550ce623aaf90 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("GetUserAISpendSince", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + user := testutil.Fake(s.T(), faker, database.User{}) + arg := database.GetUserAISpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: uuid.New(), + PeriodStart: time.Now().UTC().Truncate(24 * time.Hour), + } + row := testutil.Fake(s.T(), faker, database.GetUserAISpendSinceRow{UserID: user.ID, EffectiveGroupID: arg.EffectiveGroupID, PeriodStart: arg.PeriodStart}) + dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes() + dbm.EXPECT().GetUserAISpendSince(gomock.Any(), arg).Return(row, nil).AnyTimes() + check.Args(arg).Asserts(user, policy.ActionRead).Returns(row) + })) + + s.Run("IncrementUserAIDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.IncrementUserAIDailySpendParams{ + UserID: uuid.New(), + EffectiveGroupID: uuid.New(), + Day: time.Now().UTC().Truncate(24 * time.Hour), + CostMicros: 1000, + } + row := testutil.Fake(s.T(), faker, database.AIUserDailySpend{UserID: arg.UserID, EffectiveGroupID: arg.EffectiveGroupID, Day: arg.Day}) + dbm.EXPECT().IncrementUserAIDailySpend(gomock.Any(), arg).Return(row, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceAibridgeInterception, policy.ActionUpdate).Returns(row) + })) + s.Run("GetAIProviderByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { provider := testutil.Fake(s.T(), faker, database.AIProvider{}) dbm.EXPECT().GetAIProviderByID(gomock.Any(), provider.ID).Return(provider, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 2095b41f5d9c5..5afa5fcdc1289 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) @@ -3889,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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 4b34989ae5b0a..97eacb11fa3bf 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() @@ -7285,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() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 35d63873d3e8b..7b659c63dd794 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/000536_ai_user_daily_spend.down.sql b/coderd/database/migrations/000536_ai_user_daily_spend.down.sql new file mode 100644 index 0000000000000..a559706ecde8a --- /dev/null +++ b/coderd/database/migrations/000536_ai_user_daily_spend.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS ai_user_daily_spend CASCADE; diff --git a/coderd/database/migrations/000536_ai_user_daily_spend.up.sql b/coderd/database/migrations/000536_ai_user_daily_spend.up.sql new file mode 100644 index 0000000000000..acc65d12fd9f5 --- /dev/null +++ b/coderd/database/migrations/000536_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/000536_ai_user_daily_spend.up.sql b/coderd/database/migrations/testdata/fixtures/000536_ai_user_daily_spend.up.sql new file mode 100644 index 0000000000000..5a5dc711ec097 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000536_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 c0dea30800cd8..2abae159c1caa 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 92dd60ebb5ead..0576f50e65b73 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 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. // The result can be filtered on template_ids, meaning only user data // from workspaces based on those templates will be included. @@ -995,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) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 2e8924e502aee..629e633c0c740 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11751,6 +11751,367 @@ func TestUpsertAISeats(t *testing.T) { require.False(t, alreadyExists) } +func TestIncrementUserAIDailySpend(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 + 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 { + 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 + var err error + for _, cost := range tt.costs { + row, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: day, + CostMicros: cost, + }) + if err != nil { + break + } + } + if tt.wantErr { + require.Error(t, err) + require.True(t, database.IsCheckViolation(err, database.CheckAIUserDailySpendSpendMicrosCheck)) + 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) + 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.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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + 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.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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + 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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + 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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + 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.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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: dayNonTruncated, CostMicros: 50, + }) + require.NoError(t, err) + require.Equal(t, int64(150), nonTruncatedRow.SpendMicros, + "non-midnight UTC time should accumulate on the truncated day's row") + 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. + 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 (2024-06-16). + localLate := time.Date(2024, 6, 15, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600)) + nonUTCRow, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: localLate, CostMicros: 100, + }) + require.NoError(t, err) + 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). + 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.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.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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + 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 TestGetUserAISpendSince(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 + + 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 []seedRow + wantSpend int64 + }{ + {name: "NoRows", wantSpend: 0}, + {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 { + 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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + Day: r.day, + CostMicros: r.spend, + }) + require.NoError(t, err) + } + + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + 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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: userA.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 100, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: userB.ID, EffectiveGroupID: group.ID, Day: today, CostMicros: 25, + }) + require.NoError(t, err) + + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + 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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupA.ID, Day: today, CostMicros: 100, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: groupB.ID, Day: today, CostMicros: 25, + }) + require.NoError(t, err) + + got, err := db.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + 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") + }) + + // 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.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999, + }) + require.NoError(t, err) + _, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ + 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.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + 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) { t.Parallel() if testing.Short() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 6e41d273b4257..4a5c315c63ec7 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2546,6 +2546,79 @@ func (q *sqlQuerier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UU return i, err } +const getUserAISpendSince = `-- name: GetUserAISpendSince :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 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 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"` + SpendMicros int64 `db:"spend_micros" json:"spend_micros"` +} + +// 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) + var i GetUserAISpendSinceRow + err := row.Scan( + &i.UserID, + &i.EffectiveGroupID, + &i.PeriodStart, + &i.SpendMicros, + ) + 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 diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index cad1ed645226f..00c4917371e11 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -79,3 +79,25 @@ ORDER BY -- (groups are unique on (organization_id, name), not name alone). gaib.group_id ASC LIMIT 1; + +-- 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) +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 + spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros +RETURNING *; + +-- name: GetUserAISpendSince :one +-- 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, + @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 68058923b3bd0..17adc9486d626 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);