From 6032a6e626e2a93177d8bc79330cd88ce28e04d4 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Wed, 19 Nov 2025 14:11:31 +0200 Subject: [PATCH 1/5] feat: add configurable retention for aibridge Signed-off-by: Danny Kopping --- cli/server.go | 2 +- cli/testdata/coder_server_--help.golden | 21 +- cli/testdata/server-config.yaml.golden | 15 +- coderd/apidoc/docs.go | 6 + coderd/apidoc/swagger.json | 6 + coderd/database/dbauthz/dbauthz.go | 7 + coderd/database/dbmetrics/querymetrics.go | 7 + coderd/database/dbmock/dbmock.go | 14 ++ coderd/database/dbpurge/dbpurge.go | 11 +- coderd/database/dbpurge/dbpurge_test.go | 198 +++++++++++++++++- coderd/database/querier.go | 1 + coderd/database/queries.sql.go | 25 +++ coderd/database/queries/aibridge.sql | 14 ++ codersdk/deployment.go | 28 ++- docs/reference/api/general.md | 4 +- docs/reference/api/schemas.md | 18 +- docs/reference/cli/server.md | 24 ++- .../cli/testdata/coder_server_--help.golden | 21 +- site/src/api/typesGenerated.ts | 2 + 19 files changed, 395 insertions(+), 29 deletions(-) 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..bcfa02a444044 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,25 @@ 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-retention-limit int, $CODER_AIBRIDGE_RETENTION_LIMIT (default: 12000) + Maximum number of records to purge per dbpurge cycle (10m). MUST be + set higher than expected base rate of interceptions * 10m. For + example, if you expect 50 interceptions per second you'll need a limit + of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. + Setting a value that's too high may slow down purging of other tables. + --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..f27979ca57fcd 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -747,7 +747,18 @@ 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 + # Maximum number of records to purge per dbpurge cycle (10m). MUST be set higher + # than expected base rate of interceptions * 10m. For example, if you expect 50 + # interceptions per second you'll need a limit of 50*10*60=30000 for dbpurge to + # keep up with the rate of insertions. Setting a value that's too high may slow + # down purging of other tables. + # (default: 12000, type: int) + retention_limit: 12000 diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index de7cd416f287d..dbe295d01d203 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -11705,6 +11705,12 @@ const docTemplate = `{ }, "openai": { "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig" + }, + "retention": { + "type": "integer" + }, + "retention_limit": { + "type": "integer" } } }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 80d705f335f13..13615624f3a90 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -10401,6 +10401,12 @@ }, "openai": { "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig" + }, + "retention": { + "type": "integer" + }, + "retention_limit": { + "type": "integer" } } }, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 87b5de36009bf..215a5b7e7caa0 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1723,6 +1723,13 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex return q.db.DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx, arg) } +func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, args database.DeleteOldAIBridgeRecordsParams) error { + if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceSystem); err != nil { + return err + } + return q.db.DeleteOldAIBridgeRecords(ctx, args) +} + 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/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index d841315924a15..cbb630c8fd566 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 database.DeleteOldAIBridgeRecordsParams) error { + start := time.Now() + r0 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime) + m.queryLatencies.WithLabelValues("DeleteOldAIBridgeRecords").Observe(time.Since(start).Seconds()) + return r0 +} + 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 313bb988979a1..44e4b84204308 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -709,6 +709,20 @@ 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, arg database.DeleteOldAIBridgeRecordsParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteOldAIBridgeRecords", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteOldAIBridgeRecords indicates an expected call of DeleteOldAIBridgeRecords. +func (mr *MockStoreMockRecorder) DeleteOldAIBridgeRecords(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAIBridgeRecords", reflect.TypeOf((*MockStore)(nil).DeleteOldAIBridgeRecords), ctx, arg) +} + // 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..834ba8fe572ba 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()) + if err := tx.DeleteOldAIBridgeRecords(ctx, database.DeleteOldAIBridgeRecordsParams{ + BeforeTime: deleteAIBridgeRecordsBefore, + LimitCount: int32(vals.AI.BridgeConfig.RetentionLimit.Value()), + }); err != nil { + return xerrors.Errorf("failed to delete old aibridge records: %w", err) + } + 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..b4b99bd6369ed 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,186 @@ func TestDeleteOldTelemetryHeartbeats(t *testing.T) { return totalCount == 2 && oldCount == 0 }, testutil.WaitShort, testutil.IntervalFast, "it should delete old telemetry heartbeats") } + +//nolint:paralleltest // It uses LockIDDBPurge. +func TestDeleteOldAIBridgeRecords(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitShort) + + clk := quartz.NewMock(t) + now := dbtime.Now() + 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{}) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + + // Create old AI Bridge interception (should be deleted) + oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + 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{String: apiKey.ID, Valid: true}, + InitiatorID: user.ID, + Provider: "openai", + Model: "gpt-4", + StartedAt: afterThreshold, + }, &afterThreshold) + + oldTokenUsage := dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + ID: uuid.New(), + InterceptionID: oldInterceptionWithRelated.ID, + ProviderResponseID: "resp-1", + InputTokens: 100, + OutputTokens: 50, + CreatedAt: afterThreshold, + }) + + oldUserPrompt := dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + ID: uuid.New(), + InterceptionID: oldInterceptionWithRelated.ID, + ProviderResponseID: "resp-1", + Prompt: "test prompt", + CreatedAt: afterThreshold, + }) + + oldToolUsage := 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{String: apiKey.ID, Valid: true}, + 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{String: apiKey.ID, Valid: true}, + InitiatorID: user.ID, + Provider: "anthropic", + Model: "claude-3-5-sonnet", + StartedAt: closeBeforeThreshold, + }, &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), + RetentionLimit: serpent.Int64(1000), + }, + }, + }, 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") + 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 also deleted + tokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Empty(t, tokenUsages, "old token usages should be deleted") + + userPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Empty(t, userPrompts, "old user prompts should be deleted") + + toolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Empty(t, toolUsages, "old tool usages should be deleted") + + // Silence unused variable warnings + _ = oldTokenUsage + _ = oldUserPrompt + _ = oldToolUsage +} + +func TestDeleteOldAIBridgeRecordsLimit(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + + db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure()) + user := dbgen.User(t, db, database.User{}) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + + now := dbtime.Now() + retentionPeriod := 30 * 24 * time.Hour + threshold := now.Add(-retentionPeriod) + + // Create 5 old AI Bridge interceptions + for i := range 5 { + startedAt := threshold.Add(-time.Duration(i+1) * time.Hour) + endedAt := startedAt.Add(1 * time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + InitiatorID: user.ID, + Provider: "anthropic", + Model: "claude-3-5-sonnet", + StartedAt: startedAt, + }, &endedAt) + } + + // Delete with limit of 1 + err := db.DeleteOldAIBridgeRecords(ctx, database.DeleteOldAIBridgeRecordsParams{ + BeforeTime: threshold, + LimitCount: 1, + }) + require.NoError(t, err) + + interceptions, err := db.GetAIBridgeInterceptions(ctx) + require.NoError(t, err) + require.Len(t, interceptions, 4, "should have 4 interceptions remaining after deleting 1") + + // Delete with limit of 100 (should delete all remaining) + err = db.DeleteOldAIBridgeRecords(ctx, database.DeleteOldAIBridgeRecordsParams{ + BeforeTime: threshold, + LimitCount: 100, + }) + require.NoError(t, err) + + interceptions, err = db.GetAIBridgeInterceptions(ctx) + require.NoError(t, err) + require.Len(t, interceptions, 0, "should have no interceptions remaining after deleting with high limit") +} diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 3e5771f96de04..9a256f7166063 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -102,6 +102,7 @@ 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 + DeleteOldAIBridgeRecords(ctx context.Context, arg DeleteOldAIBridgeRecordsParams) 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 21cb7b1874b5e..88411b763bb1b 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -324,6 +324,31 @@ func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAI return count, err } +const deleteOldAIBridgeRecords = `-- name: DeleteOldAIBridgeRecords :exec +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 + LIMIT $2::int + ), + -- CTEs are executed before the main statement, so all dependent records will delete before interceptions. + tool_usages AS (DELETE FROM aibridge_tool_usages WHERE interception_id IN (SELECT id FROM to_delete)), + token_usages AS (DELETE FROM aibridge_token_usages WHERE interception_id IN (SELECT id FROM to_delete)), + user_prompts AS (DELETE FROM aibridge_user_prompts WHERE interception_id IN (SELECT id FROM to_delete)) +DELETE FROM aibridge_interceptions WHERE id IN (SELECT id FROM to_delete) +` + +type DeleteOldAIBridgeRecordsParams struct { + BeforeTime time.Time `db:"before_time" json:"before_time"` + LimitCount int32 `db:"limit_count" json:"limit_count"` +} + +func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, arg DeleteOldAIBridgeRecordsParams) error { + _, err := q.db.ExecContext(ctx, deleteOldAIBridgeRecords, arg.BeforeTime, arg.LimitCount) + return 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..15810646cb238 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -326,3 +326,17 @@ FROM prompt_aggregates pa, tool_aggregates tool_agg ; + +-- name: DeleteOldAIBridgeRecords :exec +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 + LIMIT @limit_count::int + ), + -- CTEs are executed before the main statement, so all dependent records will delete before interceptions. + tool_usages AS (DELETE FROM aibridge_tool_usages WHERE interception_id IN (SELECT id FROM to_delete)), + token_usages AS (DELETE FROM aibridge_token_usages WHERE interception_id IN (SELECT id FROM to_delete)), + user_prompts AS (DELETE FROM aibridge_user_prompts WHERE interception_id IN (SELECT id FROM to_delete)) +DELETE FROM aibridge_interceptions WHERE id IN (SELECT id FROM to_delete); diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 073ab7faede3e..e7dd0e2fe0d29 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,28 @@ 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", + }, + { + Name: "AIBridge Data Retention Purge Limit", + Description: "Maximum number of records to purge per dbpurge cycle (10m). MUST be set higher than expected base rate of interceptions * 10m. " + + "For example, if you expect 50 interceptions per second you'll need a limit of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. " + + "Setting a value that's too high may slow down purging of other tables.", + Flag: "aibridge-retention-limit", + Env: "CODER_AIBRIDGE_RETENTION_LIMIT", + Value: &c.AI.BridgeConfig.RetentionLimit, + Default: "12000", // Assuming 20 requests per second. + Group: &deploymentGroupAIBridge, + YAML: "retention_limit", + }, { Name: "Enable Authorization Recordings", Description: "All api requests will have a header including all authorization calls made during the request. " + @@ -3373,6 +3395,8 @@ 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"` + RetentionLimit serpent.Int64 `json:"retention_limit" typescript:",notnull"` } type AIBridgeOpenAIConfig struct { diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index a0a24e69772e7..af8946cf5019b 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -179,7 +179,9 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "openai": { "base_url": "string", "key": "string" - } + }, + "retention": 0, + "retention_limit": 0 } }, "allow_workspace_renames": true, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 0f43255ad60c7..79b4d36eb28fa 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -393,7 +393,9 @@ "openai": { "base_url": "string", "key": "string" - } + }, + "retention": 0, + "retention_limit": 0 } ``` @@ -406,6 +408,8 @@ | `enabled` | boolean | false | | | | `inject_coder_mcp_tools` | boolean | false | | | | `openai` | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig) | false | | | +| `retention` | integer | false | | | +| `retention_limit` | integer | false | | | ## codersdk.AIBridgeInterception @@ -701,7 +705,9 @@ "openai": { "base_url": "string", "key": "string" - } + }, + "retention": 0, + "retention_limit": 0 } } ``` @@ -2858,7 +2864,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "openai": { "base_url": "string", "key": "string" - } + }, + "retention": 0, + "retention_limit": 0 } }, "allow_workspace_renames": true, @@ -3373,7 +3381,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "openai": { "base_url": "string", "key": "string" - } + }, + "retention": 0, + "retention_limit": 0 } }, "allow_workspace_renames": true, diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index bcebe05e7e070..5f195591d53b1 100644 --- a/docs/reference/cli/server.md +++ b/docs/reference/cli/server.md @@ -1762,4 +1762,26 @@ 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). + +### --aibridge-retention-limit + +| | | +|-------------|----------------------------------------------| +| Type | int | +| Environment | $CODER_AIBRIDGE_RETENTION_LIMIT | +| YAML | aibridge.retention_limit | +| Default | 12000 | + +Maximum number of records to purge per dbpurge cycle (10m). MUST be set higher than expected base rate of interceptions * 10m. For example, if you expect 50 interceptions per second you'll need a limit of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. Setting a value that's too high may slow down purging of other tables. diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index d272200609254..9622fdec13de5 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,25 @@ 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-retention-limit int, $CODER_AIBRIDGE_RETENTION_LIMIT (default: 12000) + Maximum number of records to purge per dbpurge cycle (10m). MUST be + set higher than expected base rate of interceptions * 10m. For + example, if you expect 50 interceptions per second you'll need a limit + of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. + Setting a value that's too high may slow down purging of other tables. + --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 c2c94aa314b3d..8412b5b25d3ab 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -32,6 +32,8 @@ export interface AIBridgeConfig { readonly anthropic: AIBridgeAnthropicConfig; readonly bedrock: AIBridgeBedrockConfig; readonly inject_coder_mcp_tools: boolean; + readonly retention: number; + readonly retention_limit: number; } // From codersdk/aibridge.go From 0832029eb49ec311480bfc495eecd40a2ca603f5 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Wed, 19 Nov 2025 16:23:44 +0200 Subject: [PATCH 2/5] chore: getting rid of limit, returning purge count Signed-off-by: Danny Kopping --- cli/testdata/coder_server_--help.golden | 7 --- cli/testdata/server-config.yaml.golden | 7 --- coderd/apidoc/docs.go | 3 -- coderd/apidoc/swagger.json | 3 -- coderd/database/dbauthz/dbauthz.go | 6 +-- coderd/database/dbauthz/dbauthz_test.go | 6 +++ coderd/database/dbmetrics/querymetrics.go | 6 +-- coderd/database/dbmock/dbmock.go | 13 ++--- coderd/database/dbpurge/dbpurge.go | 7 ++- coderd/database/dbpurge/dbpurge_test.go | 53 +------------------ coderd/database/querier.go | 3 +- coderd/database/queries.sql.go | 52 ++++++++++++------ coderd/database/queries/aibridge.sql | 35 +++++++++--- codersdk/deployment.go | 13 ----- docs/reference/api/general.md | 3 +- docs/reference/api/schemas.md | 13 ++--- docs/reference/cli/server.md | 11 ---- .../cli/testdata/coder_server_--help.golden | 7 --- site/src/api/typesGenerated.ts | 1 - 19 files changed, 93 insertions(+), 156 deletions(-) diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden index bcfa02a444044..1e418cf734d10 100644 --- a/cli/testdata/coder_server_--help.golden +++ b/cli/testdata/coder_server_--help.golden @@ -110,13 +110,6 @@ AIBRIDGE OPTIONS: Length of time to retain data such as interceptions and all related records (token, prompt, tool use). - --aibridge-retention-limit int, $CODER_AIBRIDGE_RETENTION_LIMIT (default: 12000) - Maximum number of records to purge per dbpurge cycle (10m). MUST be - set higher than expected base rate of interceptions * 10m. For - example, if you expect 50 interceptions per second you'll need a limit - of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. - Setting a value that's too high may slow down purging of other tables. - --aibridge-enabled bool, $CODER_AIBRIDGE_ENABLED (default: false) Whether to start an in-memory aibridged instance. diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index f27979ca57fcd..b5d0164c7dbb1 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -755,10 +755,3 @@ aibridge: # (token, prompt, tool use). # (default: 1440h, type: duration) retention: 1440h0m0s - # Maximum number of records to purge per dbpurge cycle (10m). MUST be set higher - # than expected base rate of interceptions * 10m. For example, if you expect 50 - # interceptions per second you'll need a limit of 50*10*60=30000 for dbpurge to - # keep up with the rate of insertions. Setting a value that's too high may slow - # down purging of other tables. - # (default: 12000, type: int) - retention_limit: 12000 diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index dbe295d01d203..c819db7d24c3e 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -11708,9 +11708,6 @@ const docTemplate = `{ }, "retention": { "type": "integer" - }, - "retention_limit": { - "type": "integer" } } }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 13615624f3a90..6b861da050d91 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -10404,9 +10404,6 @@ }, "retention": { "type": "integer" - }, - "retention_limit": { - "type": "integer" } } }, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 215a5b7e7caa0..822abdfc562e4 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1723,11 +1723,11 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex return q.db.DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx, arg) } -func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, args database.DeleteOldAIBridgeRecordsParams) error { +func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceSystem); err != nil { - return err + return -1, err } - return q.db.DeleteOldAIBridgeRecords(ctx, args) + return q.db.DeleteOldAIBridgeRecords(ctx, beforeTime) } func (q *querier) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index fc98700c548f6..f8dfa3d001bb7 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4647,6 +4647,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(nil).AnyTimes() + check.Args(t).Asserts(rbac.ResourceSystem, policy.ActionDelete) + })) } func (s *MethodTestSuite) TestTelemetry() { diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index cbb630c8fd566..c1bdf30c43d2e 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -389,11 +389,11 @@ func (m queryMetricsStore) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx conte return r0 } -func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime database.DeleteOldAIBridgeRecordsParams) error { +func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { start := time.Now() - r0 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime) + r0, r1 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime) m.queryLatencies.WithLabelValues("DeleteOldAIBridgeRecords").Observe(time.Since(start).Seconds()) - return r0 + return r0, r1 } func (m queryMetricsStore) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error { diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 44e4b84204308..b441ab1e55c50 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -710,17 +710,18 @@ func (mr *MockStoreMockRecorder) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx } // DeleteOldAIBridgeRecords mocks base method. -func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, arg database.DeleteOldAIBridgeRecordsParams) error { +func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "DeleteOldAIBridgeRecords", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 + 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, arg any) *gomock.Call { +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, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAIBridgeRecords", reflect.TypeOf((*MockStore)(nil).DeleteOldAIBridgeRecords), ctx, beforeTime) } // DeleteOldAuditLogConnectionEvents mocks base method. diff --git a/coderd/database/dbpurge/dbpurge.go b/coderd/database/dbpurge/dbpurge.go index 834ba8fe572ba..e7feeb5550e47 100644 --- a/coderd/database/dbpurge/dbpurge.go +++ b/coderd/database/dbpurge/dbpurge.go @@ -92,12 +92,11 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder } deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value()) - if err := tx.DeleteOldAIBridgeRecords(ctx, database.DeleteOldAIBridgeRecordsParams{ - BeforeTime: deleteAIBridgeRecordsBefore, - LimitCount: int32(vals.AI.BridgeConfig.RetentionLimit.Value()), - }); err != nil { + count, err := tx.DeleteOldAIBridgeRecords(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))) diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index b4b99bd6369ed..b3ce5fd87a4dc 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -849,8 +849,7 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{ AI: codersdk.AIConfig{ BridgeConfig: codersdk.AIBridgeConfig{ - Retention: serpent.Duration(retentionPeriod), - RetentionLimit: serpent.Int64(1000), + Retention: serpent.Duration(retentionPeriod), }, }, }, clk) @@ -891,53 +890,3 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { _ = oldUserPrompt _ = oldToolUsage } - -func TestDeleteOldAIBridgeRecordsLimit(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - - db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure()) - user := dbgen.User(t, db, database.User{}) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) - - now := dbtime.Now() - retentionPeriod := 30 * 24 * time.Hour - threshold := now.Add(-retentionPeriod) - - // Create 5 old AI Bridge interceptions - for i := range 5 { - startedAt := threshold.Add(-time.Duration(i+1) * time.Hour) - endedAt := startedAt.Add(1 * time.Minute) - dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ - ID: uuid.New(), - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, - InitiatorID: user.ID, - Provider: "anthropic", - Model: "claude-3-5-sonnet", - StartedAt: startedAt, - }, &endedAt) - } - - // Delete with limit of 1 - err := db.DeleteOldAIBridgeRecords(ctx, database.DeleteOldAIBridgeRecordsParams{ - BeforeTime: threshold, - LimitCount: 1, - }) - require.NoError(t, err) - - interceptions, err := db.GetAIBridgeInterceptions(ctx) - require.NoError(t, err) - require.Len(t, interceptions, 4, "should have 4 interceptions remaining after deleting 1") - - // Delete with limit of 100 (should delete all remaining) - err = db.DeleteOldAIBridgeRecords(ctx, database.DeleteOldAIBridgeRecordsParams{ - BeforeTime: threshold, - LimitCount: 100, - }) - require.NoError(t, err) - - interceptions, err = db.GetAIBridgeInterceptions(ctx) - require.NoError(t, err) - require.Len(t, interceptions, 0, "should have no interceptions remaining after deleting with high limit") -} diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 9a256f7166063..19f74258cdeda 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -102,7 +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 - DeleteOldAIBridgeRecords(ctx context.Context, arg DeleteOldAIBridgeRecordsParams) 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 88411b763bb1b..5e541db5b66ad 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -324,29 +324,47 @@ func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAI return count, err } -const deleteOldAIBridgeRecords = `-- name: DeleteOldAIBridgeRecords :exec +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 - LIMIT $2::int ), - -- CTEs are executed before the main statement, so all dependent records will delete before interceptions. - tool_usages AS (DELETE FROM aibridge_tool_usages WHERE interception_id IN (SELECT id FROM to_delete)), - token_usages AS (DELETE FROM aibridge_token_usages WHERE interception_id IN (SELECT id FROM to_delete)), - user_prompts AS (DELETE FROM aibridge_user_prompts WHERE interception_id IN (SELECT id FROM to_delete)) -DELETE FROM aibridge_interceptions WHERE id IN (SELECT id FROM to_delete) -` - -type DeleteOldAIBridgeRecordsParams struct { - BeforeTime time.Time `db:"before_time" json:"before_time"` - LimitCount int32 `db:"limit_count" json:"limit_count"` -} - -func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, arg DeleteOldAIBridgeRecordsParams) error { - _, err := q.db.ExecContext(ctx, deleteOldAIBridgeRecords, arg.BeforeTime, arg.LimitCount) - return err + -- 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 diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 15810646cb238..bfd08860cb39c 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -327,16 +327,37 @@ FROM tool_aggregates tool_agg ; --- name: DeleteOldAIBridgeRecords :exec +-- 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 - LIMIT @limit_count::int ), - -- CTEs are executed before the main statement, so all dependent records will delete before interceptions. - tool_usages AS (DELETE FROM aibridge_tool_usages WHERE interception_id IN (SELECT id FROM to_delete)), - token_usages AS (DELETE FROM aibridge_token_usages WHERE interception_id IN (SELECT id FROM to_delete)), - user_prompts AS (DELETE FROM aibridge_user_prompts WHERE interception_id IN (SELECT id FROM to_delete)) -DELETE FROM aibridge_interceptions WHERE id IN (SELECT id FROM to_delete); + -- 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 e7dd0e2fe0d29..2d9265f3f984b 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -3359,18 +3359,6 @@ Write out the current server config as YAML to stdout.`, Group: &deploymentGroupAIBridge, YAML: "retention", }, - { - Name: "AIBridge Data Retention Purge Limit", - Description: "Maximum number of records to purge per dbpurge cycle (10m). MUST be set higher than expected base rate of interceptions * 10m. " + - "For example, if you expect 50 interceptions per second you'll need a limit of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. " + - "Setting a value that's too high may slow down purging of other tables.", - Flag: "aibridge-retention-limit", - Env: "CODER_AIBRIDGE_RETENTION_LIMIT", - Value: &c.AI.BridgeConfig.RetentionLimit, - Default: "12000", // Assuming 20 requests per second. - Group: &deploymentGroupAIBridge, - YAML: "retention_limit", - }, { Name: "Enable Authorization Recordings", Description: "All api requests will have a header including all authorization calls made during the request. " + @@ -3396,7 +3384,6 @@ type AIBridgeConfig struct { Bedrock AIBridgeBedrockConfig `json:"bedrock" typescript:",notnull"` InjectCoderMCPTools serpent.Bool `json:"inject_coder_mcp_tools" typescript:",notnull"` Retention serpent.Duration `json:"retention" typescript:",notnull"` - RetentionLimit serpent.Int64 `json:"retention_limit" typescript:",notnull"` } type AIBridgeOpenAIConfig struct { diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index af8946cf5019b..2f30b4c1d67e2 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -180,8 +180,7 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "base_url": "string", "key": "string" }, - "retention": 0, - "retention_limit": 0 + "retention": 0 } }, "allow_workspace_renames": true, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 79b4d36eb28fa..ddbeadd5fa075 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -394,8 +394,7 @@ "base_url": "string", "key": "string" }, - "retention": 0, - "retention_limit": 0 + "retention": 0 } ``` @@ -409,7 +408,6 @@ | `inject_coder_mcp_tools` | boolean | false | | | | `openai` | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig) | false | | | | `retention` | integer | false | | | -| `retention_limit` | integer | false | | | ## codersdk.AIBridgeInterception @@ -706,8 +704,7 @@ "base_url": "string", "key": "string" }, - "retention": 0, - "retention_limit": 0 + "retention": 0 } } ``` @@ -2865,8 +2862,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "base_url": "string", "key": "string" }, - "retention": 0, - "retention_limit": 0 + "retention": 0 } }, "allow_workspace_renames": true, @@ -3382,8 +3378,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "base_url": "string", "key": "string" }, - "retention": 0, - "retention_limit": 0 + "retention": 0 } }, "allow_workspace_renames": true, diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index 5f195591d53b1..9f202932d1e88 100644 --- a/docs/reference/cli/server.md +++ b/docs/reference/cli/server.md @@ -1774,14 +1774,3 @@ Whether to inject Coder's MCP tools into intercepted AIBridge requests (requires | Default | 1440h | Length of time to retain data such as interceptions and all related records (token, prompt, tool use). - -### --aibridge-retention-limit - -| | | -|-------------|----------------------------------------------| -| Type | int | -| Environment | $CODER_AIBRIDGE_RETENTION_LIMIT | -| YAML | aibridge.retention_limit | -| Default | 12000 | - -Maximum number of records to purge per dbpurge cycle (10m). MUST be set higher than expected base rate of interceptions * 10m. For example, if you expect 50 interceptions per second you'll need a limit of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. Setting a value that's too high may slow down purging of other tables. diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index 9622fdec13de5..944a78b65680e 100644 --- a/enterprise/cli/testdata/coder_server_--help.golden +++ b/enterprise/cli/testdata/coder_server_--help.golden @@ -111,13 +111,6 @@ AIBRIDGE OPTIONS: Length of time to retain data such as interceptions and all related records (token, prompt, tool use). - --aibridge-retention-limit int, $CODER_AIBRIDGE_RETENTION_LIMIT (default: 12000) - Maximum number of records to purge per dbpurge cycle (10m). MUST be - set higher than expected base rate of interceptions * 10m. For - example, if you expect 50 interceptions per second you'll need a limit - of 50*10*60=30000 for dbpurge to keep up with the rate of insertions. - Setting a value that's too high may slow down purging of other tables. - --aibridge-enabled bool, $CODER_AIBRIDGE_ENABLED (default: false) Whether to start an in-memory aibridged instance. diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 8412b5b25d3ab..1be5e144b7178 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -33,7 +33,6 @@ export interface AIBridgeConfig { readonly bedrock: AIBridgeBedrockConfig; readonly inject_coder_mcp_tools: boolean; readonly retention: number; - readonly retention_limit: number; } // From codersdk/aibridge.go From f0a6d847b07d86bee6f37bee1cae3e88c5358b3e Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Wed, 19 Nov 2025 16:44:20 +0200 Subject: [PATCH 3/5] chore: appease CI Signed-off-by: Danny Kopping --- coderd/database/dbauthz/dbauthz_test.go | 2 +- codersdk/deployment.go | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index f8dfa3d001bb7..9664df65ea96b 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4650,7 +4650,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(nil).AnyTimes() + db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int32(0), nil).AnyTimes() check.Args(t).Asserts(rbac.ResourceSystem, policy.ActionDelete) })) } diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 2d9265f3f984b..6b394122858af 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -3358,6 +3358,7 @@ Write out the current server config as YAML to stdout.`, Default: "1440h", // 60 days. Group: &deploymentGroupAIBridge, YAML: "retention", + Annotations: serpent.Annotations{}.Mark(annotationFormatDuration, "true"), }, { Name: "Enable Authorization Recordings", From 3e4d7b22193b9faf6672643a5dffd1b877182e34 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Thu, 20 Nov 2025 14:21:34 +0200 Subject: [PATCH 4/5] chore: review feedback Signed-off-by: Danny Kopping --- coderd/database/dbpurge/dbpurge_test.go | 81 ++++++++++++++++++------- 1 file changed, 59 insertions(+), 22 deletions(-) diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index b3ce5fd87a4dc..0a4de8c922be9 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -759,12 +759,13 @@ func TestDeleteOldTelemetryHeartbeats(t *testing.T) { }, testutil.WaitShort, testutil.IntervalFast, "it should delete old telemetry heartbeats") } -//nolint:paralleltest // It uses LockIDDBPurge. func TestDeleteOldAIBridgeRecords(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) clk := quartz.NewMock(t) - now := dbtime.Now() + 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) @@ -774,12 +775,11 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure()) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) user := dbgen.User(t, db, database.User{}) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) // Create old AI Bridge interception (should be deleted) oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ ID: uuid.New(), - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + APIKeyID: sql.NullString{}, InitiatorID: user.ID, Provider: "anthropic", Model: "claude-3-5-sonnet", @@ -789,14 +789,14 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { // Create old interception with related records (should all be deleted) oldInterceptionWithRelated := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ ID: uuid.New(), - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + APIKeyID: sql.NullString{}, InitiatorID: user.ID, Provider: "openai", Model: "gpt-4", StartedAt: afterThreshold, }, &afterThreshold) - oldTokenUsage := dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + _ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ ID: uuid.New(), InterceptionID: oldInterceptionWithRelated.ID, ProviderResponseID: "resp-1", @@ -805,7 +805,7 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { CreatedAt: afterThreshold, }) - oldUserPrompt := dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + _ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ ID: uuid.New(), InterceptionID: oldInterceptionWithRelated.ID, ProviderResponseID: "resp-1", @@ -813,7 +813,7 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { CreatedAt: afterThreshold, }) - oldToolUsage := dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ + _ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ ID: uuid.New(), InterceptionID: oldInterceptionWithRelated.ID, ProviderResponseID: "resp-1", @@ -827,7 +827,7 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { // Create recent AI Bridge interception (should be kept) recentInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ ID: uuid.New(), - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + APIKeyID: sql.NullString{}, InitiatorID: user.ID, Provider: "anthropic", Model: "claude-3-5-sonnet", @@ -837,13 +837,41 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { // Create interception close to threshold (should be kept) nearThresholdInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ ID: uuid.New(), - APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + 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{ @@ -869,24 +897,33 @@ func TestDeleteOldAIBridgeRecords(t *testing.T) { 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") - 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 also deleted - tokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) + oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) require.NoError(t, err) - require.Empty(t, tokenUsages, "old token usages should be deleted") + require.Empty(t, oldTokenUsages, "old token usages should be deleted") - userPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID) + oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID) require.NoError(t, err) - require.Empty(t, userPrompts, "old user prompts should be deleted") + require.Empty(t, oldUserPrompts, "old user prompts should be deleted") - toolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) + oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) require.NoError(t, err) - require.Empty(t, toolUsages, "old tool usages should be deleted") + require.Empty(t, oldToolUsages, "old tool usages should be deleted") - // Silence unused variable warnings - _ = oldTokenUsage - _ = oldUserPrompt - _ = oldToolUsage + 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") } From 445648db9949c6d5952db86e66a3c6a8d32476c4 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Fri, 21 Nov 2025 11:14:25 +0200 Subject: [PATCH 5/5] chore: use correct RBAC context Signed-off-by: Danny Kopping --- coderd/database/dbauthz/dbauthz.go | 4 ++-- coderd/database/dbauthz/dbauthz_test.go | 2 +- coderd/database/dbpurge/dbpurge.go | 3 ++- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 68047700344e1..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{}, @@ -1724,7 +1724,7 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex } func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { - if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceSystem); err != nil { + if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAibridgeInterception); err != nil { return -1, err } return q.db.DeleteOldAIBridgeRecords(ctx, beforeTime) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index bd573669853c9..f367a4d13f34d 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4658,7 +4658,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() - check.Args(t).Asserts(rbac.ResourceSystem, policy.ActionDelete) + check.Args(t).Asserts(rbac.ResourceAibridgeInterception, policy.ActionDelete) })) } diff --git a/coderd/database/dbpurge/dbpurge.go b/coderd/database/dbpurge/dbpurge.go index e7feeb5550e47..f145281c8f65b 100644 --- a/coderd/database/dbpurge/dbpurge.go +++ b/coderd/database/dbpurge/dbpurge.go @@ -92,7 +92,8 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder } deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value()) - count, err := tx.DeleteOldAIBridgeRecords(ctx, deleteAIBridgeRecordsBefore) + // 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) }