diff --git a/cli/server.go b/cli/server.go index 010e96d2fc693..6ae0d1aa70ecc 100644 --- a/cli/server.go +++ b/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 diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden index 49ce14b2f572f..1e418cf734d10 100644 --- a/cli/testdata/coder_server_--help.golden +++ b/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. diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index 33f5c56c43840..b5d0164c7dbb1 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/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 diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 398ce80970a3c..67be7d985dfe4 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -11740,6 +11740,9 @@ const docTemplate = `{ }, "openai": { "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig" + }, + "retention": { + "type": "integer" } } }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index dfe1c793811c8..888c88057e2a9 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -10432,6 +10432,9 @@ }, "openai": { "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig" + }, + "retention": { + "type": "integer" } } }, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 7cc074b5ce752..c3cc39116f656 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/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 diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 4e2f02d0d3c81..f367a4d13f34d 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/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() { diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 527db3953c2b6..fac1dee0536d2 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 4906b695b2724..588eb0b335a10 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/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() diff --git a/coderd/database/dbpurge/dbpurge.go b/coderd/database/dbpurge/dbpurge.go index 067fe1f0499e3..f145281c8f65b 100644 --- a/coderd/database/dbpurge/dbpurge.go +++ b/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 diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index 74bf36639fbb5..0a4de8c922be9 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/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") +} diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 39fa7dab120bf..bf6b0bc5c06cb 100644 --- a/coderd/database/querier.go +++ b/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 diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 557840db6794f..b4aa80cd6e4ab 100644 --- a/coderd/database/queries.sql.go +++ b/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 diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index bd85e2be1e465..bfd08860cb39c 100644 --- a/coderd/database/queries/aibridge.sql +++ b/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; diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 073ab7faede3e..6b394122858af 100644 --- a/codersdk/deployment.go +++ b/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 { diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index a0a24e69772e7..2f30b4c1d67e2 100644 --- a/docs/reference/api/general.md +++ b/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, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 3c7bcab5da580..55ac9b2bd58a7 100644 --- a/docs/reference/api/schemas.md +++ b/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, diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index bcebe05e7e070..9f202932d1e88 100644 --- a/docs/reference/cli/server.md +++ b/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 | 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). +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 | duration | +| Environment | $CODER_AIBRIDGE_RETENTION | +| YAML | aibridge.retention | +| Default | 1440h | + +Length of time to retain data such as interceptions and all related records (token, prompt, tool use). diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index d272200609254..944a78b65680e 100644 --- a/enterprise/cli/testdata/coder_server_--help.golden +++ b/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. diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index db4f3b2865d05..47f5b1254af27 100644 --- a/site/src/api/typesGenerated.ts +++ b/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