fix(coderd/database/dbpurge): allow disabling AI Bridge retention with 0 (#21062)

Previously setting AI Bridge retention to 0 would cause records to be
deleted immediately since we didn't check for the zero value before
calculating the deletion threshold.

This adds a check for aibridgeRetention > 0 to skip deletion when
retention is disabled, matching the pattern used for other retention
settings (connection logs, audit logs, etc.).

Also fixes the return type of DeleteOldAIBridgeRecords from int32 to
int64 since COUNT(*) returns bigint in PostgreSQL.

Refs #21055

coderd/database/dbauthz/dbauthz.go

@@ -1732,7 +1732,7 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex
 	return q.db.DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx, arg)
 }
 
-func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) {
 	if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAibridgeInterception); err != nil {
 		return -1, err
 	}

coderd/database/dbauthz/dbauthz_test.go

@@ -4705,7 +4705,7 @@ func (s *MethodTestSuite) TestAIBridge() {
 
 	s.Run("DeleteOldAIBridgeRecords", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		t := dbtime.Now()
-		db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int32(0), nil).AnyTimes()
+		db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int64(0), nil).AnyTimes()
 		check.Args(t).Asserts(rbac.ResourceAibridgeInterception, policy.ActionDelete)
 	}))
 }

coderd/database/dbmetrics/querymetrics.go

@@ -396,7 +396,7 @@ func (m queryMetricsStore) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx conte
 	return r0
 }
 
-func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) {
 	start := time.Now()
 	r0, r1 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime)
 	m.queryLatencies.WithLabelValues("DeleteOldAIBridgeRecords").Observe(time.Since(start).Seconds())

coderd/database/dbmock/dbmock.go

@@ -725,10 +725,10 @@ func (mr *MockStoreMockRecorder) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx
 }
 
 // DeleteOldAIBridgeRecords mocks base method.
-func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "DeleteOldAIBridgeRecords", ctx, beforeTime)
-	ret0, _ := ret[0].(int32)
+	ret0, _ := ret[0].(int64)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }

coderd/database/dbpurge/dbpurge.go

@@ -118,11 +118,15 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
 				return xerrors.Errorf("failed to delete old audit log connection events: %w", err)
 			}
 
-			deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value())
-			// nolint:gocritic // Needs to run as aibridge context.
-			purgedAIBridgeRecords, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
-			if err != nil {
-				return xerrors.Errorf("failed to delete old aibridge records: %w", err)
+			var purgedAIBridgeRecords int64
+			aibridgeRetention := vals.AI.BridgeConfig.Retention.Value()
+			if aibridgeRetention > 0 {
+				deleteAIBridgeRecordsBefore := start.Add(-aibridgeRetention)
+				// nolint:gocritic // Needs to run as aibridge context.
+				purgedAIBridgeRecords, err = tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
+				if err != nil {
+					return xerrors.Errorf("failed to delete old aibridge records: %w", err)
+				}
 			}
 
 			var purgedConnectionLogs int64

coderd/database/dbpurge/dbpurge_test.go

@@ -973,170 +973,237 @@ func TestDeleteOldConnectionLogs(t *testing.T) {
 func TestDeleteOldAIBridgeRecords(t *testing.T) {
 	t.Parallel()
 
-	ctx := testutil.Context(t, testutil.WaitShort)
-
-	clk := quartz.NewMock(t)
 	now := time.Date(2025, 1, 15, 7, 30, 0, 0, time.UTC)
 	retentionPeriod := 30 * 24 * time.Hour                                // 30 days
 	afterThreshold := now.Add(-retentionPeriod).Add(-24 * time.Hour)      // 31 days ago (older than threshold)
 	beforeThreshold := now.Add(-15 * 24 * time.Hour)                      // 15 days ago (newer than threshold)
 	closeBeforeThreshold := now.Add(-retentionPeriod).Add(24 * time.Hour) // 29 days ago
-	clk.Set(now).MustWait(ctx)
 
-	db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
-	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
-	user := dbgen.User(t, db, database.User{})
-
-	// Create old AI Bridge interception (should be deleted)
-	oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
-		ID:          uuid.New(),
-		APIKeyID:    sql.NullString{},
-		InitiatorID: user.ID,
-		Provider:    "anthropic",
-		Model:       "claude-3-5-sonnet",
-		StartedAt:   afterThreshold,
-	}, &afterThreshold)
-
-	// Create old interception with related records (should all be deleted)
-	oldInterceptionWithRelated := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
-		ID:          uuid.New(),
-		APIKeyID:    sql.NullString{},
-		InitiatorID: user.ID,
-		Provider:    "openai",
-		Model:       "gpt-4",
-		StartedAt:   afterThreshold,
-	}, &afterThreshold)
-
-	_ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
-		ID:                 uuid.New(),
-		InterceptionID:     oldInterceptionWithRelated.ID,
-		ProviderResponseID: "resp-1",
-		InputTokens:        100,
-		OutputTokens:       50,
-		CreatedAt:          afterThreshold,
-	})
-
-	_ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{
-		ID:                 uuid.New(),
-		InterceptionID:     oldInterceptionWithRelated.ID,
-		ProviderResponseID: "resp-1",
-		Prompt:             "test prompt",
-		CreatedAt:          afterThreshold,
-	})
-
-	_ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{
-		ID:                 uuid.New(),
-		InterceptionID:     oldInterceptionWithRelated.ID,
-		ProviderResponseID: "resp-1",
-		Tool:               "test-tool",
-		ServerUrl:          sql.NullString{String: "http://test", Valid: true},
-		Input:              "{}",
-		Injected:           true,
-		CreatedAt:          afterThreshold,
-	})
-
-	// Create recent AI Bridge interception (should be kept)
-	recentInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
-		ID:          uuid.New(),
-		APIKeyID:    sql.NullString{},
-		InitiatorID: user.ID,
-		Provider:    "anthropic",
-		Model:       "claude-3-5-sonnet",
-		StartedAt:   beforeThreshold,
-	}, &beforeThreshold)
-
-	// Create interception close to threshold (should be kept)
-	nearThresholdInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
-		ID:          uuid.New(),
-		APIKeyID:    sql.NullString{},
-		InitiatorID: user.ID,
-		Provider:    "anthropic",
-		Model:       "claude-3-5-sonnet",
-		StartedAt:   closeBeforeThreshold,
-	}, &closeBeforeThreshold)
-
-	_ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
-		ID:                 uuid.New(),
-		InterceptionID:     nearThresholdInterception.ID,
-		ProviderResponseID: "resp-1",
-		InputTokens:        100,
-		OutputTokens:       50,
-		CreatedAt:          closeBeforeThreshold,
-	})
-
-	_ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{
-		ID:                 uuid.New(),
-		InterceptionID:     nearThresholdInterception.ID,
-		ProviderResponseID: "resp-1",
-		Prompt:             "test prompt",
-		CreatedAt:          closeBeforeThreshold,
-	})
-
-	_ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{
-		ID:                 uuid.New(),
-		InterceptionID:     nearThresholdInterception.ID,
-		ProviderResponseID: "resp-1",
-		Tool:               "test-tool",
-		ServerUrl:          sql.NullString{String: "http://test", Valid: true},
-		Input:              "{}",
-		Injected:           true,
-		CreatedAt:          closeBeforeThreshold,
-	})
-
-	// Run the purge with configured retention period
-	done := awaitDoTick(ctx, t, clk)
-	closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{
-		AI: codersdk.AIConfig{
-			BridgeConfig: codersdk.AIBridgeConfig{
-				Retention: serpent.Duration(retentionPeriod),
-			},
-		},
-	}, clk)
-	defer closer.Close()
-	// Wait for tick
-	testutil.TryReceive(ctx, t, done)
-
-	// Verify results by querying all AI Bridge records
-	interceptions, err := db.GetAIBridgeInterceptions(ctx)
-	require.NoError(t, err)
-
-	// Extract interception IDs for comparison
-	interceptionIDs := make([]uuid.UUID, len(interceptions))
-	for i, interception := range interceptions {
-		interceptionIDs[i] = interception.ID
+	type testFixtures struct {
+		oldInterception            database.AIBridgeInterception
+		oldInterceptionWithRelated database.AIBridgeInterception
+		recentInterception         database.AIBridgeInterception
+		nearThresholdInterception  database.AIBridgeInterception
 	}
 
-	require.NotContains(t, interceptionIDs, oldInterception.ID, "old interception should be deleted")
-	require.NotContains(t, interceptionIDs, oldInterceptionWithRelated.ID, "old interception with related records should be deleted")
+	testCases := []struct {
+		name      string
+		retention time.Duration
+		verify    func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures)
+	}{
+		{
+			name:      "RetentionEnabled",
+			retention: retentionPeriod,
+			verify: func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) {
+				t.Helper()
 
-	// Verify related records were also deleted
-	oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID)
-	require.NoError(t, err)
-	require.Empty(t, oldTokenUsages, "old token usages should be deleted")
+				interceptions, err := db.GetAIBridgeInterceptions(ctx)
+				require.NoError(t, err)
+				require.Len(t, interceptions, 2, "expected 2 interceptions remaining")
 
-	oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID)
-	require.NoError(t, err)
-	require.Empty(t, oldUserPrompts, "old user prompts should be deleted")
+				interceptionIDs := make([]uuid.UUID, len(interceptions))
+				for i, interception := range interceptions {
+					interceptionIDs[i] = interception.ID
+				}
 
-	oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID)
-	require.NoError(t, err)
-	require.Empty(t, oldToolUsages, "old tool usages should be deleted")
+				require.NotContains(t, interceptionIDs, fixtures.oldInterception.ID, "old interception should be deleted")
+				require.NotContains(t, interceptionIDs, fixtures.oldInterceptionWithRelated.ID, "old interception with related records should be deleted")
+				require.Contains(t, interceptionIDs, fixtures.recentInterception.ID, "recent interception should be kept")
+				require.Contains(t, interceptionIDs, fixtures.nearThresholdInterception.ID, "near threshold interception should be kept")
 
-	require.Contains(t, interceptionIDs, recentInterception.ID, "recent interception should be kept")
-	require.Contains(t, interceptionIDs, nearThresholdInterception.ID, "near threshold interception should be kept")
+				// Verify related records were deleted for old interception.
+				oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
+				require.NoError(t, err)
+				require.Empty(t, oldTokenUsages, "old token usages should be deleted")
 
-	// Verify related records were NOT deleted
-	newTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, nearThresholdInterception.ID)
-	require.NoError(t, err)
-	require.Len(t, newTokenUsages, 1, "near threshold token usages should not be deleted")
+				oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
+				require.NoError(t, err)
+				require.Empty(t, oldUserPrompts, "old user prompts should be deleted")
 
-	newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, nearThresholdInterception.ID)
-	require.NoError(t, err)
-	require.Len(t, newUserPrompts, 1, "near threshold user prompts should not be deleted")
+				oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
+				require.NoError(t, err)
+				require.Empty(t, oldToolUsages, "old tool usages should be deleted")
 
-	newToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, nearThresholdInterception.ID)
-	require.NoError(t, err)
-	require.Len(t, newToolUsages, 1, "near threshold tool usages should not be deleted")
+				// Verify related records were NOT deleted for near-threshold interception.
+				newTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.nearThresholdInterception.ID)
+				require.NoError(t, err)
+				require.Len(t, newTokenUsages, 1, "near threshold token usages should not be deleted")
+
+				newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.nearThresholdInterception.ID)
+				require.NoError(t, err)
+				require.Len(t, newUserPrompts, 1, "near threshold user prompts should not be deleted")
+
+				newToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.nearThresholdInterception.ID)
+				require.NoError(t, err)
+				require.Len(t, newToolUsages, 1, "near threshold tool usages should not be deleted")
+			},
+		},
+		{
+			name:      "RetentionDisabled",
+			retention: 0,
+			verify: func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) {
+				t.Helper()
+
+				interceptions, err := db.GetAIBridgeInterceptions(ctx)
+				require.NoError(t, err)
+				require.Len(t, interceptions, 4, "expected all 4 interceptions to be retained")
+
+				interceptionIDs := make([]uuid.UUID, len(interceptions))
+				for i, interception := range interceptions {
+					interceptionIDs[i] = interception.ID
+				}
+
+				require.Contains(t, interceptionIDs, fixtures.oldInterception.ID, "old interception should be kept")
+				require.Contains(t, interceptionIDs, fixtures.oldInterceptionWithRelated.ID, "old interception with related records should be kept")
+				require.Contains(t, interceptionIDs, fixtures.recentInterception.ID, "recent interception should be kept")
+				require.Contains(t, interceptionIDs, fixtures.nearThresholdInterception.ID, "near threshold interception should be kept")
+
+				// Verify all related records were kept.
+				oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
+				require.NoError(t, err)
+				require.Len(t, oldTokenUsages, 1, "old token usages should be kept")
+
+				oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
+				require.NoError(t, err)
+				require.Len(t, oldUserPrompts, 1, "old user prompts should be kept")
+
+				oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
+				require.NoError(t, err)
+				require.Len(t, oldToolUsages, 1, "old tool usages should be kept")
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+			clk := quartz.NewMock(t)
+			clk.Set(now).MustWait(ctx)
+
+			db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
+			logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+			user := dbgen.User(t, db, database.User{})
+
+			// Create old AI Bridge interception (should be deleted when retention enabled).
+			oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
+				ID:          uuid.New(),
+				APIKeyID:    sql.NullString{},
+				InitiatorID: user.ID,
+				Provider:    "anthropic",
+				Model:       "claude-3-5-sonnet",
+				StartedAt:   afterThreshold,
+			}, &afterThreshold)
+
+			// Create old interception with related records (should all be deleted when retention enabled).
+			oldInterceptionWithRelated := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
+				ID:          uuid.New(),
+				APIKeyID:    sql.NullString{},
+				InitiatorID: user.ID,
+				Provider:    "openai",
+				Model:       "gpt-4",
+				StartedAt:   afterThreshold,
+			}, &afterThreshold)
+
+			_ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
+				ID:                 uuid.New(),
+				InterceptionID:     oldInterceptionWithRelated.ID,
+				ProviderResponseID: "resp-1",
+				InputTokens:        100,
+				OutputTokens:       50,
+				CreatedAt:          afterThreshold,
+			})
+
+			_ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{
+				ID:                 uuid.New(),
+				InterceptionID:     oldInterceptionWithRelated.ID,
+				ProviderResponseID: "resp-1",
+				Prompt:             "test prompt",
+				CreatedAt:          afterThreshold,
+			})
+
+			_ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{
+				ID:                 uuid.New(),
+				InterceptionID:     oldInterceptionWithRelated.ID,
+				ProviderResponseID: "resp-1",
+				Tool:               "test-tool",
+				ServerUrl:          sql.NullString{String: "http://test", Valid: true},
+				Input:              "{}",
+				Injected:           true,
+				CreatedAt:          afterThreshold,
+			})
+
+			// Create recent AI Bridge interception (should be kept).
+			recentInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
+				ID:          uuid.New(),
+				APIKeyID:    sql.NullString{},
+				InitiatorID: user.ID,
+				Provider:    "anthropic",
+				Model:       "claude-3-5-sonnet",
+				StartedAt:   beforeThreshold,
+			}, &beforeThreshold)
+
+			// Create interception close to threshold (should be kept).
+			nearThresholdInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
+				ID:          uuid.New(),
+				APIKeyID:    sql.NullString{},
+				InitiatorID: user.ID,
+				Provider:    "anthropic",
+				Model:       "claude-3-5-sonnet",
+				StartedAt:   closeBeforeThreshold,
+			}, &closeBeforeThreshold)
+
+			_ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
+				ID:                 uuid.New(),
+				InterceptionID:     nearThresholdInterception.ID,
+				ProviderResponseID: "resp-1",
+				InputTokens:        100,
+				OutputTokens:       50,
+				CreatedAt:          closeBeforeThreshold,
+			})
+
+			_ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{
+				ID:                 uuid.New(),
+				InterceptionID:     nearThresholdInterception.ID,
+				ProviderResponseID: "resp-1",
+				Prompt:             "test prompt",
+				CreatedAt:          closeBeforeThreshold,
+			})
+
+			_ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{
+				ID:                 uuid.New(),
+				InterceptionID:     nearThresholdInterception.ID,
+				ProviderResponseID: "resp-1",
+				Tool:               "test-tool",
+				ServerUrl:          sql.NullString{String: "http://test", Valid: true},
+				Input:              "{}",
+				Injected:           true,
+				CreatedAt:          closeBeforeThreshold,
+			})
+
+			fixtures := testFixtures{
+				oldInterception:            oldInterception,
+				oldInterceptionWithRelated: oldInterceptionWithRelated,
+				recentInterception:         recentInterception,
+				nearThresholdInterception:  nearThresholdInterception,
+			}
+
+			// Run the purge with configured retention period.
+			done := awaitDoTick(ctx, t, clk)
+			closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{
+				AI: codersdk.AIConfig{
+					BridgeConfig: codersdk.AIBridgeConfig{
+						Retention: serpent.Duration(tc.retention),
+					},
+				},
+			}, clk)
+			defer closer.Close()
+			testutil.TryReceive(ctx, t, done)
+
+			tc.verify(t, ctx, db, fixtures)
+		})
+	}
 }
 
 func TestDeleteOldAuditLogs(t *testing.T) {

coderd/database/querier.go

@@ -104,7 +104,7 @@ type sqlcQuerier interface {
 	DeleteOAuth2ProviderAppSecretByID(ctx context.Context, id uuid.UUID) error
 	DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Context, arg DeleteOAuth2ProviderAppTokensByAppAndUserIDParams) error
 	// Cumulative count.
-	DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error)
+	DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error)
 	DeleteOldAuditLogConnectionEvents(ctx context.Context, arg DeleteOldAuditLogConnectionEventsParams) error
 	// Deletes old audit logs based on retention policy, excluding deprecated
 	// connection events (connect, disconnect, open, close) which are handled

coderd/database/queries.sql.go

@@ -354,17 +354,18 @@ WITH
     WHERE id IN (SELECT id FROM to_delete)
     RETURNING 1
   )
-SELECT
+SELECT (
   (SELECT COUNT(*) FROM tool_usages) +
   (SELECT COUNT(*) FROM token_usages) +
   (SELECT COUNT(*) FROM user_prompts) +
-  (SELECT COUNT(*) FROM interceptions) as total_deleted
+  (SELECT COUNT(*) FROM interceptions)
+)::bigint as total_deleted
 `
 
 // Cumulative count.
-func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) {
 	row := q.db.QueryRowContext(ctx, deleteOldAIBridgeRecords, beforeTime)
-	var total_deleted int32
+	var total_deleted int64
 	err := row.Scan(&total_deleted)
 	return total_deleted, err
 }

coderd/database/queries/aibridge.sql

@@ -360,8 +360,9 @@ WITH
     RETURNING 1
   )
 -- Cumulative count.
-SELECT
+SELECT (
   (SELECT COUNT(*) FROM tool_usages) +
   (SELECT COUNT(*) FROM token_usages) +
   (SELECT COUNT(*) FROM user_prompts) +
-  (SELECT COUNT(*) FROM interceptions) as total_deleted;
+  (SELECT COUNT(*) FROM interceptions)
+)::bigint as total_deleted;