mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(coderd/database/dbpurge): allow disabling AI Bridge retention with 0 (#21062)
Previously setting AI Bridge retention to 0 would cause records to be deleted immediately since we didn't check for the zero value before calculating the deletion threshold. This adds a check for aibridgeRetention > 0 to skip deletion when retention is disabled, matching the pattern used for other retention settings (connection logs, audit logs, etc.). Also fixes the return type of DeleteOldAIBridgeRecords from int32 to int64 since COUNT(*) returns bigint in PostgreSQL. Refs #21055
This commit is contained in:
Mathias Fredriksson
@@ -118,11 +118,15 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
|
||||
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.
|
||||
purgedAIBridgeRecords, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("failed to delete old aibridge records: %w", err)
|
||||
var purgedAIBridgeRecords int64
|
||||
aibridgeRetention := vals.AI.BridgeConfig.Retention.Value()
|
||||
if aibridgeRetention > 0 {
|
||||
deleteAIBridgeRecordsBefore := start.Add(-aibridgeRetention)
|
||||
// nolint:gocritic // Needs to run as aibridge context.
|
||||
purgedAIBridgeRecords, err = tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("failed to delete old aibridge records: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var purgedConnectionLogs int64
|
||||
|
||||
@@ -973,170 +973,237 @@ func TestDeleteOldConnectionLogs(t *testing.T) {
|
||||
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
|
||||
type testFixtures struct {
|
||||
oldInterception database.AIBridgeInterception
|
||||
oldInterceptionWithRelated database.AIBridgeInterception
|
||||
recentInterception database.AIBridgeInterception
|
||||
nearThresholdInterception database.AIBridgeInterception
|
||||
}
|
||||
|
||||
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")
|
||||
testCases := []struct {
|
||||
name string
|
||||
retention time.Duration
|
||||
verify func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures)
|
||||
}{
|
||||
{
|
||||
name: "RetentionEnabled",
|
||||
retention: retentionPeriod,
|
||||
verify: func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) {
|
||||
t.Helper()
|
||||
|
||||
// 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")
|
||||
interceptions, err := db.GetAIBridgeInterceptions(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, interceptions, 2, "expected 2 interceptions remaining")
|
||||
|
||||
oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, oldUserPrompts, "old user prompts should be deleted")
|
||||
interceptionIDs := make([]uuid.UUID, len(interceptions))
|
||||
for i, interception := range interceptions {
|
||||
interceptionIDs[i] = interception.ID
|
||||
}
|
||||
|
||||
oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, oldToolUsages, "old tool usages should be deleted")
|
||||
require.NotContains(t, interceptionIDs, fixtures.oldInterception.ID, "old interception should be deleted")
|
||||
require.NotContains(t, interceptionIDs, fixtures.oldInterceptionWithRelated.ID, "old interception with related records should be deleted")
|
||||
require.Contains(t, interceptionIDs, fixtures.recentInterception.ID, "recent interception should be kept")
|
||||
require.Contains(t, interceptionIDs, fixtures.nearThresholdInterception.ID, "near threshold interception should be kept")
|
||||
|
||||
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 deleted for old interception.
|
||||
oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, oldTokenUsages, "old token usages should be deleted")
|
||||
|
||||
// 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")
|
||||
oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, oldUserPrompts, "old user prompts should 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")
|
||||
oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, oldToolUsages, "old tool usages should 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")
|
||||
// Verify related records were NOT deleted for near-threshold interception.
|
||||
newTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.nearThresholdInterception.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, newTokenUsages, 1, "near threshold token usages should not be deleted")
|
||||
|
||||
newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.nearThresholdInterception.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, newUserPrompts, 1, "near threshold user prompts should not be deleted")
|
||||
|
||||
newToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.nearThresholdInterception.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, newToolUsages, 1, "near threshold tool usages should not be deleted")
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "RetentionDisabled",
|
||||
retention: 0,
|
||||
verify: func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) {
|
||||
t.Helper()
|
||||
|
||||
interceptions, err := db.GetAIBridgeInterceptions(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, interceptions, 4, "expected all 4 interceptions to be retained")
|
||||
|
||||
interceptionIDs := make([]uuid.UUID, len(interceptions))
|
||||
for i, interception := range interceptions {
|
||||
interceptionIDs[i] = interception.ID
|
||||
}
|
||||
|
||||
require.Contains(t, interceptionIDs, fixtures.oldInterception.ID, "old interception should be kept")
|
||||
require.Contains(t, interceptionIDs, fixtures.oldInterceptionWithRelated.ID, "old interception with related records should be kept")
|
||||
require.Contains(t, interceptionIDs, fixtures.recentInterception.ID, "recent interception should be kept")
|
||||
require.Contains(t, interceptionIDs, fixtures.nearThresholdInterception.ID, "near threshold interception should be kept")
|
||||
|
||||
// Verify all related records were kept.
|
||||
oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, oldTokenUsages, 1, "old token usages should be kept")
|
||||
|
||||
oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, oldUserPrompts, 1, "old user prompts should be kept")
|
||||
|
||||
oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, oldToolUsages, 1, "old tool usages should be kept")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clk := quartz.NewMock(t)
|
||||
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 when retention enabled).
|
||||
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 when retention enabled).
|
||||
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,
|
||||
})
|
||||
|
||||
fixtures := testFixtures{
|
||||
oldInterception: oldInterception,
|
||||
oldInterceptionWithRelated: oldInterceptionWithRelated,
|
||||
recentInterception: recentInterception,
|
||||
nearThresholdInterception: nearThresholdInterception,
|
||||
}
|
||||
|
||||
// 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(tc.retention),
|
||||
},
|
||||
},
|
||||
}, clk)
|
||||
defer closer.Close()
|
||||
testutil.TryReceive(ctx, t, done)
|
||||
|
||||
tc.verify(t, ctx, db, fixtures)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteOldAuditLogs(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user