feat: add ai_user_daily_spend table and queries (#26562)

## Description Adds the spend tracking table and queries needed by [AIGOV-427](https://linear.app/codercom/issue/AIGOV-427/add-post-response-spend-accumulation) (post-response accumulation) and [AIGOV-428](https://linear.app/codercom/issue/AIGOV-428/add-pre-request-budget-enforcement) (pre-request enforcement). ## Changes - Add `ai_user_daily_spend` table to aggregate per-user, per-effective-group AI spend by UTC day. - Add `UpsertUserAIDailySpend` and `GetUserAISpendSince` queries. Closes https://linear.app/codercom/issue/AIGOV-426/add-daily-spend-table-and-queries > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-09-21 20:51:01 +08:00 · 2026-07-02 16:29:25 +01:00
parent dee41c34e6
commit fcdd029d74
15 changed files with 613 additions and 0 deletions
@@ -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
@@ -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)
 }
@@ -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()
@@ -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)
@@ -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()
@@ -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);
@@ -0,0 +1 @@
+DROP TABLE IF EXISTS ai_user_daily_spend CASCADE;
@@ -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);
@@ -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);
@@ -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.
@@ -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)
@@ -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() {
@@ -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
@@ -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;
@@ -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);
				`@@ -0,0 +1 @@`
				`DROP TABLE IF EXISTS ai_user_daily_spend CASCADE;`