feat: add configurable retention for aibridge (#20828)

Closes https://github.com/coder/internal/issues/1134

---------

Signed-off-by: Danny Kopping <danny@coder.com>

cli/server.go

@@ -1029,7 +1029,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			defer shutdownConns()
 
 			// Ensures that old database entries are cleaned up over time!
-			purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, quartz.NewReal())
+			purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, options.DeploymentValues, quartz.NewReal())
 			defer purger.Close()
 
 			// Updates workspace usage

cli/testdata/coder_server_--help.golden

@@ -81,11 +81,6 @@ OPTIONS:
           check is performed once per day.
 
 AIBRIDGE OPTIONS: 
-      --aibridge-inject-coder-mcp-tools bool, $CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS (default: false)
-          Whether to inject Coder's MCP tools into intercepted AI Bridge
-          requests (requires the "oauth2" and "mcp-server-http" experiments to
-          be enabled).
-
       --aibridge-anthropic-base-url string, $CODER_AIBRIDGE_ANTHROPIC_BASE_URL (default: https://api.anthropic.com/)
           The base URL of the Anthropic API.
 
@@ -111,9 +106,18 @@ AIBRIDGE OPTIONS:
           See
           https://docs.claude.com/en/docs/claude-code/settings#environment-variables.
 
+      --aibridge-retention duration, $CODER_AIBRIDGE_RETENTION (default: 1440h)
+          Length of time to retain data such as interceptions and all related
+          records (token, prompt, tool use).
+
       --aibridge-enabled bool, $CODER_AIBRIDGE_ENABLED (default: false)
           Whether to start an in-memory aibridged instance.
 
+      --aibridge-inject-coder-mcp-tools bool, $CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS (default: false)
+          Whether to inject Coder's MCP tools into intercepted AIBridge requests
+          (requires the "oauth2" and "mcp-server-http" experiments to be
+          enabled).
+
       --aibridge-openai-base-url string, $CODER_AIBRIDGE_OPENAI_BASE_URL (default: https://api.openai.com/v1/)
           The base URL of the OpenAI API.
 

cli/testdata/server-config.yaml.golden

@@ -747,7 +747,11 @@ aibridge:
   # https://docs.claude.com/en/docs/claude-code/settings#environment-variables.
   # (default: global.anthropic.claude-haiku-4-5-20251001-v1:0, type: string)
   bedrock_small_fast_model: global.anthropic.claude-haiku-4-5-20251001-v1:0
-  # Whether to inject Coder's MCP tools into intercepted AI Bridge requests
-  # (requires the "oauth2" and "mcp-server-http" experiments to be enabled).
+  # Whether to inject Coder's MCP tools into intercepted AIBridge requests (requires
+  # the "oauth2" and "mcp-server-http" experiments to be enabled).
   # (default: false, type: bool)
   inject_coder_mcp_tools: false
+  # Length of time to retain data such as interceptions and all related records
+  # (token, prompt, tool use).
+  # (default: 1440h, type: duration)
+  retention: 1440h0m0s

coderd/apidoc/docs.go

@@ -11740,6 +11740,9 @@ const docTemplate = `{
                 },
                 "openai": {
                     "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig"
+                },
+                "retention": {
+                    "type": "integer"
                 }
             }
         },

coderd/apidoc/swagger.json

@@ -10432,6 +10432,9 @@
 				},
 				"openai": {
 					"$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig"
+				},
+				"retention": {
+					"type": "integer"
 				}
 			}
 		},

coderd/database/dbauthz/dbauthz.go

@@ -608,7 +608,7 @@ var (
 						policy.ActionReadPersonal, // Required to read users' external auth links. // TODO: this is too broad; reduce scope to just external_auth_links by creating separate resource.
 					},
 					rbac.ResourceApiKey.Type:               {policy.ActionRead}, // Validate API keys.
-					rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
+					rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
 				}),
 				User:    []rbac.Permission{},
 				ByOrgID: map[string]rbac.OrgPermissions{},
@@ -1723,6 +1723,13 @@ 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) {
+	if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAibridgeInterception); err != nil {
+		return -1, err
+	}
+	return q.db.DeleteOldAIBridgeRecords(ctx, beforeTime)
+}
+
 func (q *querier) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error {
 	// `ResourceSystem` is deprecated, but it doesn't make sense to add
 	// `policy.ActionDelete` to `ResourceAuditLog`, since this is the one and

coderd/database/dbauthz/dbauthz_test.go

@@ -4654,6 +4654,12 @@ func (s *MethodTestSuite) TestAIBridge() {
 		db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), params).Return(intc, nil).AnyTimes()
 		check.Args(params).Asserts(intc, policy.ActionUpdate).Returns(intc)
 	}))
+
+	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()
+		check.Args(t).Asserts(rbac.ResourceAibridgeInterception, policy.ActionDelete)
+	}))
 }
 
 func (s *MethodTestSuite) TestTelemetry() {

coderd/database/dbmetrics/querymetrics.go

@@ -389,6 +389,13 @@ func (m queryMetricsStore) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx conte
 	return r0
 }
 
+func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+	start := time.Now()
+	r0, r1 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime)
+	m.queryLatencies.WithLabelValues("DeleteOldAIBridgeRecords").Observe(time.Since(start).Seconds())
+	return r0, r1
+}
+
 func (m queryMetricsStore) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error {
 	start := time.Now()
 	r0 := m.s.DeleteOldAuditLogConnectionEvents(ctx, threshold)

coderd/database/dbmock/dbmock.go

@@ -709,6 +709,21 @@ func (mr *MockStoreMockRecorder) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOAuth2ProviderAppTokensByAppAndUserID", reflect.TypeOf((*MockStore)(nil).DeleteOAuth2ProviderAppTokensByAppAndUserID), ctx, arg)
 }
 
+// DeleteOldAIBridgeRecords mocks base method.
+func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DeleteOldAIBridgeRecords", ctx, beforeTime)
+	ret0, _ := ret[0].(int32)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteOldAIBridgeRecords indicates an expected call of DeleteOldAIBridgeRecords.
+func (mr *MockStoreMockRecorder) DeleteOldAIBridgeRecords(ctx, beforeTime any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAIBridgeRecords", reflect.TypeOf((*MockStore)(nil).DeleteOldAIBridgeRecords), ctx, beforeTime)
+}
+
 // DeleteOldAuditLogConnectionEvents mocks base method.
 func (m *MockStore) DeleteOldAuditLogConnectionEvents(ctx context.Context, arg database.DeleteOldAuditLogConnectionEventsParams) error {
 	m.ctrl.T.Helper()

coderd/database/dbpurge/dbpurge.go

@@ -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

coderd/database/dbpurge/dbpurge_test.go

@@ -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")
+}

coderd/database/querier.go

@@ -102,6 +102,8 @@ type sqlcQuerier interface {
 	DeleteOAuth2ProviderAppCodesByAppAndUserID(ctx context.Context, arg DeleteOAuth2ProviderAppCodesByAppAndUserIDParams) error
 	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)
 	DeleteOldAuditLogConnectionEvents(ctx context.Context, arg DeleteOldAuditLogConnectionEventsParams) error
 	// Delete all notification messages which have not been updated for over a week.
 	DeleteOldNotificationMessages(ctx context.Context) error

coderd/database/queries.sql.go

@@ -324,6 +324,49 @@ func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAI
 	return count, err
 }
 
+const deleteOldAIBridgeRecords = `-- name: DeleteOldAIBridgeRecords :one
+WITH
+  -- We don't have FK relationships between the dependent tables and aibridge_interceptions, so we can't rely on DELETE CASCADE.
+  to_delete AS (
+    SELECT id FROM aibridge_interceptions
+    WHERE started_at < $1::timestamp with time zone
+  ),
+  -- CTEs are executed in order.
+  tool_usages AS (
+    DELETE FROM aibridge_tool_usages
+    WHERE interception_id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  ),
+  token_usages AS (
+    DELETE FROM aibridge_token_usages
+    WHERE interception_id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  ),
+  user_prompts AS (
+    DELETE FROM aibridge_user_prompts
+    WHERE interception_id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  ),
+  interceptions AS (
+    DELETE FROM aibridge_interceptions
+    WHERE id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  )
+SELECT
+  (SELECT COUNT(*) FROM tool_usages) +
+  (SELECT COUNT(*) FROM token_usages) +
+  (SELECT COUNT(*) FROM user_prompts) +
+  (SELECT COUNT(*) FROM interceptions) as total_deleted
+`
+
+// Cumulative count.
+func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
+	row := q.db.QueryRowContext(ctx, deleteOldAIBridgeRecords, beforeTime)
+	var total_deleted int32
+	err := row.Scan(&total_deleted)
+	return total_deleted, err
+}
+
 const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
 SELECT
 	id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id

coderd/database/queries/aibridge.sql

@@ -326,3 +326,38 @@ FROM
     prompt_aggregates pa,
     tool_aggregates tool_agg
 ;
+
+-- name: DeleteOldAIBridgeRecords :one
+WITH
+  -- We don't have FK relationships between the dependent tables and aibridge_interceptions, so we can't rely on DELETE CASCADE.
+  to_delete AS (
+    SELECT id FROM aibridge_interceptions
+    WHERE started_at < @before_time::timestamp with time zone
+  ),
+  -- CTEs are executed in order.
+  tool_usages AS (
+    DELETE FROM aibridge_tool_usages
+    WHERE interception_id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  ),
+  token_usages AS (
+    DELETE FROM aibridge_token_usages
+    WHERE interception_id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  ),
+  user_prompts AS (
+    DELETE FROM aibridge_user_prompts
+    WHERE interception_id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  ),
+  interceptions AS (
+    DELETE FROM aibridge_interceptions
+    WHERE id IN (SELECT id FROM to_delete)
+    RETURNING 1
+  )
+-- Cumulative count.
+SELECT
+  (SELECT COUNT(*) FROM tool_usages) +
+  (SELECT COUNT(*) FROM token_usages) +
+  (SELECT COUNT(*) FROM user_prompts) +
+  (SELECT COUNT(*) FROM interceptions) as total_deleted;

codersdk/deployment.go

@@ -3340,8 +3340,8 @@ Write out the current server config as YAML to stdout.`,
 			YAML:        "bedrock_small_fast_model",
 		},
 		{
-			Name:        "AI Bridge Inject Coder MCP tools",
-			Description: "Whether to inject Coder's MCP tools into intercepted AI Bridge requests (requires the \"oauth2\" and \"mcp-server-http\" experiments to be enabled).",
+			Name:        "AIBridge Inject Coder MCP tools",
+			Description: "Whether to inject Coder's MCP tools into intercepted AIBridge requests (requires the \"oauth2\" and \"mcp-server-http\" experiments to be enabled).",
 			Flag:        "aibridge-inject-coder-mcp-tools",
 			Env:         "CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS",
 			Value:       &c.AI.BridgeConfig.InjectCoderMCPTools,
@@ -3349,6 +3349,17 @@ Write out the current server config as YAML to stdout.`,
 			Group:       &deploymentGroupAIBridge,
 			YAML:        "inject_coder_mcp_tools",
 		},
+		{
+			Name:        "AIBridge Data Retention Duration",
+			Description: "Length of time to retain data such as interceptions and all related records (token, prompt, tool use).",
+			Flag:        "aibridge-retention",
+			Env:         "CODER_AIBRIDGE_RETENTION",
+			Value:       &c.AI.BridgeConfig.Retention,
+			Default:     "1440h", // 60 days.
+			Group:       &deploymentGroupAIBridge,
+			YAML:        "retention",
+			Annotations: serpent.Annotations{}.Mark(annotationFormatDuration, "true"),
+		},
 		{
 			Name: "Enable Authorization Recordings",
 			Description: "All api requests will have a header including all authorization calls made during the request. " +
@@ -3373,6 +3384,7 @@ type AIBridgeConfig struct {
 	Anthropic           AIBridgeAnthropicConfig `json:"anthropic" typescript:",notnull"`
 	Bedrock             AIBridgeBedrockConfig   `json:"bedrock" typescript:",notnull"`
 	InjectCoderMCPTools serpent.Bool            `json:"inject_coder_mcp_tools" typescript:",notnull"`
+	Retention           serpent.Duration        `json:"retention" typescript:",notnull"`
 }
 
 type AIBridgeOpenAIConfig struct {

docs/reference/api/general.md

@@ -179,7 +179,8 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \
         "openai": {
           "base_url": "string",
           "key": "string"
-        }
+        },
+        "retention": 0
       }
     },
     "allow_workspace_renames": true,

docs/reference/api/schemas.md

@@ -393,7 +393,8 @@
   "openai": {
     "base_url": "string",
     "key": "string"
-  }
+  },
+  "retention": 0
 }
 ```
 
@@ -406,6 +407,7 @@
 | `enabled`                | boolean                                                              | false    |              |             |
 | `inject_coder_mcp_tools` | boolean                                                              | false    |              |             |
 | `openai`                 | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig)       | false    |              |             |
+| `retention`              | integer                                                              | false    |              |             |
 
 ## codersdk.AIBridgeInterception
 
@@ -701,7 +703,8 @@
     "openai": {
       "base_url": "string",
       "key": "string"
-    }
+    },
+    "retention": 0
   }
 }
 ```
@@ -2858,7 +2861,8 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
         "openai": {
           "base_url": "string",
           "key": "string"
-        }
+        },
+        "retention": 0
       }
     },
     "allow_workspace_renames": true,
@@ -3373,7 +3377,8 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
       "openai": {
         "base_url": "string",
         "key": "string"
-      }
+      },
+      "retention": 0
     }
   },
   "allow_workspace_renames": true,

docs/reference/cli/server.md

@@ -1762,4 +1762,15 @@ The small fast model to use when making requests to the AWS Bedrock API. Claude
 | YAML        | <code>aibridge.inject_coder_mcp_tools</code>        |
 | Default     | <code>false</code>                                  |
 
-Whether to inject Coder's MCP tools into intercepted AI Bridge requests (requires the "oauth2" and "mcp-server-http" experiments to be enabled).
+Whether to inject Coder's MCP tools into intercepted AIBridge requests (requires the "oauth2" and "mcp-server-http" experiments to be enabled).
+
+### --aibridge-retention
+
+|             |                                        |
+|-------------|----------------------------------------|
+| Type        | <code>duration</code>                  |
+| Environment | <code>$CODER_AIBRIDGE_RETENTION</code> |
+| YAML        | <code>aibridge.retention</code>        |
+| Default     | <code>1440h</code>                     |
+
+Length of time to retain data such as interceptions and all related records (token, prompt, tool use).

enterprise/cli/testdata/coder_server_--help.golden

@@ -82,11 +82,6 @@ OPTIONS:
           check is performed once per day.
 
 AIBRIDGE OPTIONS: 
-      --aibridge-inject-coder-mcp-tools bool, $CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS (default: false)
-          Whether to inject Coder's MCP tools into intercepted AI Bridge
-          requests (requires the "oauth2" and "mcp-server-http" experiments to
-          be enabled).
-
       --aibridge-anthropic-base-url string, $CODER_AIBRIDGE_ANTHROPIC_BASE_URL (default: https://api.anthropic.com/)
           The base URL of the Anthropic API.
 
@@ -112,9 +107,18 @@ AIBRIDGE OPTIONS:
           See
           https://docs.claude.com/en/docs/claude-code/settings#environment-variables.
 
+      --aibridge-retention duration, $CODER_AIBRIDGE_RETENTION (default: 1440h)
+          Length of time to retain data such as interceptions and all related
+          records (token, prompt, tool use).
+
       --aibridge-enabled bool, $CODER_AIBRIDGE_ENABLED (default: false)
           Whether to start an in-memory aibridged instance.
 
+      --aibridge-inject-coder-mcp-tools bool, $CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS (default: false)
+          Whether to inject Coder's MCP tools into intercepted AIBridge requests
+          (requires the "oauth2" and "mcp-server-http" experiments to be
+          enabled).
+
       --aibridge-openai-base-url string, $CODER_AIBRIDGE_OPENAI_BASE_URL (default: https://api.openai.com/v1/)
           The base URL of the OpenAI API.
 

site/src/api/typesGenerated.ts

@@ -32,6 +32,7 @@ export interface AIBridgeConfig {
 	readonly anthropic: AIBridgeAnthropicConfig;
 	readonly bedrock: AIBridgeBedrockConfig;
 	readonly inject_coder_mcp_tools: boolean;
+	readonly retention: number;
 }
 
 // From codersdk/aibridge.go