feat: add configurable retention for aibridge (#20828)

Closes https://github.com/coder/internal/issues/1134 --------- Signed-off-by: Danny Kopping <danny@coder.com>
2026-09-24 15:04:27 +08:00 · 2025-11-21 11:35:36 +02:00
parent 83966e346a
commit 5a7d4f69f6
20 changed files with 372 additions and 30 deletions
@@ -13,6 +13,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/pproflabel"
+	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/quartz"
 )

@@ -36,7 +37,7 @@ const (
 // It is the caller's responsibility to call Close on the returned instance.
 //
 // This is for cleaning up old, unused resources from the database that take up space.
-func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.Clock) io.Closer {
+func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, clk quartz.Clock) io.Closer {
 	closed := make(chan struct{})

 	ctx, cancelFunc := context.WithCancel(ctx)
@@ -90,6 +91,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.
 				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.
+			count, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
+			if err != nil {
+				return xerrors.Errorf("failed to delete old aibridge records: %w", err)
+			}
+			logger.Debug(ctx, "purged aibridge entries", slog.F("count", count), slog.F("since", deleteAIBridgeRecordsBefore.Format(time.RFC3339)))
+
 			logger.Debug(ctx, "purged old database entries", slog.F("duration", clk.Since(start)))

 			return nil
@@ -33,6 +33,7 @@ import (
 	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/testutil"
 	"github.com/coder/quartz"
+	"github.com/coder/serpent"
 )

 func TestMain(m *testing.M) {
@@ -51,7 +52,7 @@ func TestPurge(t *testing.T) {
 	done := awaitDoTick(ctx, t, clk)
 	mDB := dbmock.NewMockStore(gomock.NewController(t))
 	mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).Return(nil).Times(2)
-	purger := dbpurge.New(context.Background(), testutil.Logger(t), mDB, clk)
+	purger := dbpurge.New(context.Background(), testutil.Logger(t), mDB, &codersdk.DeploymentValues{}, clk)
 	<-done // wait for doTick() to run.
 	require.NoError(t, purger.Close())
 }
@@ -129,7 +130,7 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
 	})

 	// when
-	closer := dbpurge.New(ctx, logger, db, clk)
+	closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
 	defer closer.Close()

 	// then
@@ -154,7 +155,7 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {

 	// Start a new purger to immediately trigger delete after rollup.
 	_ = closer.Close()
-	closer = dbpurge.New(ctx, logger, db, clk)
+	closer = dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
 	defer closer.Close()

 	// then
@@ -245,7 +246,7 @@ func TestDeleteOldWorkspaceAgentLogs(t *testing.T) {
 	// After dbpurge completes, the ticker is reset. Trap this call.

 	done := awaitDoTick(ctx, t, clk)
-	closer := dbpurge.New(ctx, logger, db, clk)
+	closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
 	defer closer.Close()
 	<-done // doTick() has now run.

@@ -466,7 +467,7 @@ func TestDeleteOldProvisionerDaemons(t *testing.T) {
 	require.NoError(t, err)

 	// when
-	closer := dbpurge.New(ctx, logger, db, clk)
+	closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
 	defer closer.Close()

 	// then
@@ -570,7 +571,7 @@ func TestDeleteOldAuditLogConnectionEvents(t *testing.T) {

 	// Run the purge
 	done := awaitDoTick(ctx, t, clk)
-	closer := dbpurge.New(ctx, logger, db, clk)
+	closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
 	defer closer.Close()
 	// Wait for tick
 	testutil.TryReceive(ctx, t, done)
@@ -733,7 +734,7 @@ func TestDeleteOldTelemetryHeartbeats(t *testing.T) {
 	require.NoError(t, err)

 	done := awaitDoTick(ctx, t, clk)
-	closer := dbpurge.New(ctx, logger, db, clk)
+	closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
 	defer closer.Close()
 	<-done // doTick() has now run.

@@ -757,3 +758,172 @@ func TestDeleteOldTelemetryHeartbeats(t *testing.T) {
 		return totalCount == 2 && oldCount == 0
 	}, testutil.WaitShort, testutil.IntervalFast, "it should delete old telemetry heartbeats")
 }
+
+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
+	}
+
+	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")
+
+	// 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")
+
+	oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID)
+	require.NoError(t, err)
+	require.Empty(t, oldUserPrompts, "old user prompts should be deleted")
+
+	oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID)
+	require.NoError(t, err)
+	require.Empty(t, oldToolUsages, "old tool usages should be deleted")
+
+	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 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")
+
+	newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, nearThresholdInterception.ID)
+	require.NoError(t, err)
+	require.Len(t, newUserPrompts, 1, "near threshold user prompts should not 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")
+}