opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-22 05:05:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add configurable retention for aibridge (#20828)

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

---------

Signed-off-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Danny Kopping
2025-11-21 11:35:36 +02:00
committed by GitHub
parent 83966e346a
commit 5a7d4f69f6
20 changed files with 372 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/apidoc/docs.go
Generated
+3
View File
@@ -11740,6 +11740,9 @@ const docTemplate = `{
},
"openai": {
"$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig"
},
"retention": {
"type": "integer"
}
}
},
coderd/apidoc/swagger.json
Generated
+3
View File
@@ -10432,6 +10432,9 @@
},
"openai": {
"$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig"
},
"retention": {
"type": "integer"
}
}
},
coderd/database/dbauthz/dbauthz.go
+8 -1
View File
@@ -608,7 +608,7 @@ var (
policy.ActionReadPersonal, // Required to read users' external auth links. // TODO: this is too broad; reduce scope to just external_auth_links by creating separate resource.
},
rbac.ResourceApiKey.Type: {policy.ActionRead}, // Validate API keys.
rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
}),
User: []rbac.Permission{},
ByOrgID: map[string]rbac.OrgPermissions{},
@@ -1723,6 +1723,13 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex
return q.db.DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx, arg)
}
func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAibridgeInterception); err != nil {
return -1, err
}
return q.db.DeleteOldAIBridgeRecords(ctx, beforeTime)
}
func (q *querier) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error {
// `ResourceSystem` is deprecated, but it doesn't make sense to add
// `policy.ActionDelete` to `ResourceAuditLog`, since this is the one and
coderd/database/dbauthz/dbauthz_test.go
+6
View File
@@ -4654,6 +4654,12 @@ func (s *MethodTestSuite) TestAIBridge() {
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), params).Return(intc, nil).AnyTimes()
check.Args(params).Asserts(intc, policy.ActionUpdate).Returns(intc)
}))
s.Run("DeleteOldAIBridgeRecords", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
t := dbtime.Now()
db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int32(0), nil).AnyTimes()
check.Args(t).Asserts(rbac.ResourceAibridgeInterception, policy.ActionDelete)
}))
}
func (s *MethodTestSuite) TestTelemetry() {
coderd/database/dbmetrics/querymetrics.go
+7
View File
@@ -389,6 +389,13 @@ func (m queryMetricsStore) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx conte
return r0
}
func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
start := time.Now()
r0, r1 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime)
m.queryLatencies.WithLabelValues("DeleteOldAIBridgeRecords").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m queryMetricsStore) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error {
start := time.Now()
r0 := m.s.DeleteOldAuditLogConnectionEvents(ctx, threshold)
coderd/database/dbmock/dbmock.go
+15
View File
@@ -709,6 +709,21 @@ func (mr *MockStoreMockRecorder) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOAuth2ProviderAppTokensByAppAndUserID", reflect.TypeOf((*MockStore)(nil).DeleteOAuth2ProviderAppTokensByAppAndUserID), ctx, arg)
}
// DeleteOldAIBridgeRecords mocks base method.
func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteOldAIBridgeRecords", ctx, beforeTime)
ret0, _ := ret[0].(int32)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// DeleteOldAIBridgeRecords indicates an expected call of DeleteOldAIBridgeRecords.
func (mr *MockStoreMockRecorder) DeleteOldAIBridgeRecords(ctx, beforeTime any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAIBridgeRecords", reflect.TypeOf((*MockStore)(nil).DeleteOldAIBridgeRecords), ctx, beforeTime)
}
// DeleteOldAuditLogConnectionEvents mocks base method.
func (m *MockStore) DeleteOldAuditLogConnectionEvents(ctx context.Context, arg database.DeleteOldAuditLogConnectionEventsParams) error {
m.ctrl.T.Helper()
coderd/database/dbpurge/dbpurge.go
+10 -1
View File
@@ -13,6 +13,7 @@ import (
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/pproflabel"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/quartz"
)
@@ -36,7 +37,7 @@ const (
// It is the caller's responsibility to call Close on the returned instance.
//
// This is for cleaning up old, unused resources from the database that take up space.
func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.Clock) io.Closer {
func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, clk quartz.Clock) io.Closer {
closed := make(chan struct{})
ctx, cancelFunc := context.WithCancel(ctx)
@@ -90,6 +91,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.
return xerrors.Errorf("failed to delete old audit log connection events: %w", err)
}
deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value())
// nolint:gocritic // Needs to run as aibridge context.
count, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
if err != nil {
return xerrors.Errorf("failed to delete old aibridge records: %w", err)
}
logger.Debug(ctx, "purged aibridge entries", slog.F("count", count), slog.F("since", deleteAIBridgeRecordsBefore.Format(time.RFC3339)))
logger.Debug(ctx, "purged old database entries", slog.F("duration", clk.Since(start)))
return nil
coderd/database/dbpurge/dbpurge_test.go
+177 -7
View File
@@ -33,6 +33,7 @@ import (
"github.com/coder/coder/v2/provisionersdk"
"github.com/coder/coder/v2/testutil"
"github.com/coder/quartz"
"github.com/coder/serpent"
)
func TestMain(m *testing.M) {
@@ -51,7 +52,7 @@ func TestPurge(t *testing.T) {
done := awaitDoTick(ctx, t, clk)
mDB := dbmock.NewMockStore(gomock.NewController(t))
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).Return(nil).Times(2)
purger := dbpurge.New(context.Background(), testutil.Logger(t), mDB, clk)
purger := dbpurge.New(context.Background(), testutil.Logger(t), mDB, &codersdk.DeploymentValues{}, clk)
<-done // wait for doTick() to run.
require.NoError(t, purger.Close())
}
@@ -129,7 +130,7 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
})
// when
closer := dbpurge.New(ctx, logger, db, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
defer closer.Close()
// then
@@ -154,7 +155,7 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
// Start a new purger to immediately trigger delete after rollup.
_ = closer.Close()
closer = dbpurge.New(ctx, logger, db, clk)
closer = dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
defer closer.Close()
// then
@@ -245,7 +246,7 @@ func TestDeleteOldWorkspaceAgentLogs(t *testing.T) {
// After dbpurge completes, the ticker is reset. Trap this call.
done := awaitDoTick(ctx, t, clk)
closer := dbpurge.New(ctx, logger, db, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
defer closer.Close()
<-done // doTick() has now run.
@@ -466,7 +467,7 @@ func TestDeleteOldProvisionerDaemons(t *testing.T) {
require.NoError(t, err)
// when
closer := dbpurge.New(ctx, logger, db, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
defer closer.Close()
// then
@@ -570,7 +571,7 @@ func TestDeleteOldAuditLogConnectionEvents(t *testing.T) {
// Run the purge
done := awaitDoTick(ctx, t, clk)
closer := dbpurge.New(ctx, logger, db, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
defer closer.Close()
// Wait for tick
testutil.TryReceive(ctx, t, done)
@@ -733,7 +734,7 @@ func TestDeleteOldTelemetryHeartbeats(t *testing.T) {
require.NoError(t, err)
done := awaitDoTick(ctx, t, clk)
closer := dbpurge.New(ctx, logger, db, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk)
defer closer.Close()
<-done // doTick() has now run.
@@ -757,3 +758,172 @@ func TestDeleteOldTelemetryHeartbeats(t *testing.T) {
return totalCount == 2 && oldCount == 0
}, testutil.WaitShort, testutil.IntervalFast, "it should delete old telemetry heartbeats")
}
func TestDeleteOldAIBridgeRecords(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
clk := quartz.NewMock(t)
now := time.Date(2025, 1, 15, 7, 30, 0, 0, time.UTC)
retentionPeriod := 30 * 24 * time.Hour // 30 days
afterThreshold := now.Add(-retentionPeriod).Add(-24 * time.Hour) // 31 days ago (older than threshold)
beforeThreshold := now.Add(-15 * 24 * time.Hour) // 15 days ago (newer than threshold)
closeBeforeThreshold := now.Add(-retentionPeriod).Add(24 * time.Hour) // 29 days ago
clk.Set(now).MustWait(ctx)
db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
user := dbgen.User(t, db, database.User{})
// Create old AI Bridge interception (should be deleted)
oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
ID: uuid.New(),
APIKeyID: sql.NullString{},
InitiatorID: user.ID,
Provider: "anthropic",
Model: "claude-3-5-sonnet",
StartedAt: afterThreshold,
}, &afterThreshold)
// Create old interception with related records (should all be deleted)
oldInterceptionWithRelated := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
ID: uuid.New(),
APIKeyID: sql.NullString{},
InitiatorID: user.ID,
Provider: "openai",
Model: "gpt-4",
StartedAt: afterThreshold,
}, &afterThreshold)
_ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
ID: uuid.New(),
InterceptionID: oldInterceptionWithRelated.ID,
ProviderResponseID: "resp-1",
InputTokens: 100,
OutputTokens: 50,
CreatedAt: afterThreshold,
})
_ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{
ID: uuid.New(),
InterceptionID: oldInterceptionWithRelated.ID,
ProviderResponseID: "resp-1",
Prompt: "test prompt",
CreatedAt: afterThreshold,
})
_ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{
ID: uuid.New(),
InterceptionID: oldInterceptionWithRelated.ID,
ProviderResponseID: "resp-1",
Tool: "test-tool",
ServerUrl: sql.NullString{String: "http://test", Valid: true},
Input: "{}",
Injected: true,
CreatedAt: afterThreshold,
})
// Create recent AI Bridge interception (should be kept)
recentInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
ID: uuid.New(),
APIKeyID: sql.NullString{},
InitiatorID: user.ID,
Provider: "anthropic",
Model: "claude-3-5-sonnet",
StartedAt: beforeThreshold,
}, &beforeThreshold)
// Create interception close to threshold (should be kept)
nearThresholdInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
ID: uuid.New(),
APIKeyID: sql.NullString{},
InitiatorID: user.ID,
Provider: "anthropic",
Model: "claude-3-5-sonnet",
StartedAt: closeBeforeThreshold,
}, &closeBeforeThreshold)
_ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
ID: uuid.New(),
InterceptionID: nearThresholdInterception.ID,
ProviderResponseID: "resp-1",
InputTokens: 100,
OutputTokens: 50,
CreatedAt: closeBeforeThreshold,
})
_ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{
ID: uuid.New(),
InterceptionID: nearThresholdInterception.ID,
ProviderResponseID: "resp-1",
Prompt: "test prompt",
CreatedAt: closeBeforeThreshold,
})
_ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{
ID: uuid.New(),
InterceptionID: nearThresholdInterception.ID,
ProviderResponseID: "resp-1",
Tool: "test-tool",
ServerUrl: sql.NullString{String: "http://test", Valid: true},
Input: "{}",
Injected: true,
CreatedAt: closeBeforeThreshold,
})
// Run the purge with configured retention period
done := awaitDoTick(ctx, t, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{
AI: codersdk.AIConfig{
BridgeConfig: codersdk.AIBridgeConfig{
Retention: serpent.Duration(retentionPeriod),
},
},
}, clk)
defer closer.Close()
// Wait for tick
testutil.TryReceive(ctx, t, done)
// Verify results by querying all AI Bridge records
interceptions, err := db.GetAIBridgeInterceptions(ctx)
require.NoError(t, err)
// Extract interception IDs for comparison
interceptionIDs := make([]uuid.UUID, len(interceptions))
for i, interception := range interceptions {
interceptionIDs[i] = interception.ID
}
require.NotContains(t, interceptionIDs, oldInterception.ID, "old interception should be deleted")
require.NotContains(t, interceptionIDs, oldInterceptionWithRelated.ID, "old interception with related records should be deleted")
// Verify related records were also deleted
oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID)
require.NoError(t, err)
require.Empty(t, oldTokenUsages, "old token usages should be deleted")
oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID)
require.NoError(t, err)
require.Empty(t, oldUserPrompts, "old user prompts should be deleted")
oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID)
require.NoError(t, err)
require.Empty(t, oldToolUsages, "old tool usages should be deleted")
require.Contains(t, interceptionIDs, recentInterception.ID, "recent interception should be kept")
require.Contains(t, interceptionIDs, nearThresholdInterception.ID, "near threshold interception should be kept")
// Verify related records were NOT deleted
newTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, nearThresholdInterception.ID)
require.NoError(t, err)
require.Len(t, newTokenUsages, 1, "near threshold token usages should not be deleted")
newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, nearThresholdInterception.ID)
require.NoError(t, err)
require.Len(t, newUserPrompts, 1, "near threshold user prompts should not be deleted")
newToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, nearThresholdInterception.ID)
require.NoError(t, err)
require.Len(t, newToolUsages, 1, "near threshold tool usages should not be deleted")
}
coderd/database/querier.go
+2
View File
@@ -102,6 +102,8 @@ type sqlcQuerier interface {
DeleteOAuth2ProviderAppCodesByAppAndUserID(ctx context.Context, arg DeleteOAuth2ProviderAppCodesByAppAndUserIDParams) error
DeleteOAuth2ProviderAppSecretByID(ctx context.Context, id uuid.UUID) error
DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Context, arg DeleteOAuth2ProviderAppTokensByAppAndUserIDParams) error
// Cumulative count.
DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error)
DeleteOldAuditLogConnectionEvents(ctx context.Context, arg DeleteOldAuditLogConnectionEventsParams) error
// Delete all notification messages which have not been updated for over a week.
DeleteOldNotificationMessages(ctx context.Context) error
coderd/database/queries.sql.go
+43
View File
@@ -324,6 +324,49 @@ func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAI
return count, err
}
const deleteOldAIBridgeRecords = `-- name: DeleteOldAIBridgeRecords :one
WITH
-- We don't have FK relationships between the dependent tables and aibridge_interceptions, so we can't rely on DELETE CASCADE.
to_delete AS (
SELECT id FROM aibridge_interceptions
WHERE started_at < $1::timestamp with time zone
),
-- CTEs are executed in order.
tool_usages AS (
DELETE FROM aibridge_tool_usages
WHERE interception_id IN (SELECT id FROM to_delete)
RETURNING 1
),
token_usages AS (
DELETE FROM aibridge_token_usages
WHERE interception_id IN (SELECT id FROM to_delete)
RETURNING 1
),
user_prompts AS (
DELETE FROM aibridge_user_prompts
WHERE interception_id IN (SELECT id FROM to_delete)
RETURNING 1
),
interceptions AS (
DELETE FROM aibridge_interceptions
WHERE id IN (SELECT id FROM to_delete)
RETURNING 1
)
SELECT
(SELECT COUNT(*) FROM tool_usages) +
(SELECT COUNT(*) FROM token_usages) +
(SELECT COUNT(*) FROM user_prompts) +
(SELECT COUNT(*) FROM interceptions) as total_deleted
`
// Cumulative count.
func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
row := q.db.QueryRowContext(ctx, deleteOldAIBridgeRecords, beforeTime)
var total_deleted int32
err := row.Scan(&total_deleted)
return total_deleted, err
}
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
coderd/database/queries/aibridge.sql
+35
View File
@@ -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;