Skip to content

Commit 5a7d4f6

Browse files
dannykoppingdannykopping
authored
feat: add configurable retention for aibridge (#20828)
Closes coder/internal#1134 --------- Signed-off-by: Danny Kopping <danny@coder.com>
1 parent 83966e3 commit 5a7d4f6

File tree

20 files changed

+372
-30
lines changed

20 files changed

+372
-30
lines changed

‎cli/server.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1029,7 +1029,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
10291029
defer shutdownConns()
10301030

10311031
// Ensures that old database entries are cleaned up over time!
1032-
purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, quartz.NewReal())
1032+
purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, options.DeploymentValues, quartz.NewReal())
10331033
defer purger.Close()
10341034

10351035
// Updates workspace usage

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -81,11 +81,6 @@ OPTIONS:
8181
check is performed once per day.
8282

8383
AIBRIDGE OPTIONS:
84-
--aibridge-inject-coder-mcp-tools bool, $CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS (default: false)
85-
Whether to inject Coder's MCP tools into intercepted AI Bridge
86-
requests (requires the "oauth2" and "mcp-server-http" experiments to
87-
be enabled).
88-
8984
--aibridge-anthropic-base-url string, $CODER_AIBRIDGE_ANTHROPIC_BASE_URL (default: https://api.anthropic.com/)
9085
The base URL of the Anthropic API.
9186

@@ -111,9 +106,18 @@ AIBRIDGE OPTIONS:
111106
See
112107
https://docs.claude.com/en/docs/claude-code/settings#environment-variables.
113108

109+
--aibridge-retention duration, $CODER_AIBRIDGE_RETENTION (default: 1440h)
110+
Length of time to retain data such as interceptions and all related
111+
records (token, prompt, tool use).
112+
114113
--aibridge-enabled bool, $CODER_AIBRIDGE_ENABLED (default: false)
115114
Whether to start an in-memory aibridged instance.
116115

116+
--aibridge-inject-coder-mcp-tools bool, $CODER_AIBRIDGE_INJECT_CODER_MCP_TOOLS (default: false)
117+
Whether to inject Coder's MCP tools into intercepted AIBridge requests
118+
(requires the "oauth2" and "mcp-server-http" experiments to be
119+
enabled).
120+
117121
--aibridge-openai-base-url string, $CODER_AIBRIDGE_OPENAI_BASE_URL (default: https://api.openai.com/v1/)
118122
The base URL of the OpenAI API.
119123

‎cli/testdata/server-config.yaml.golden‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -747,7 +747,11 @@ aibridge:
747747
# https://docs.claude.com/en/docs/claude-code/settings#environment-variables.
748748
# (default: global.anthropic.claude-haiku-4-5-20251001-v1:0, type: string)
749749
bedrock_small_fast_model: global.anthropic.claude-haiku-4-5-20251001-v1:0
750-
# Whether to inject Coder's MCP tools into intercepted AI Bridge requests
751-
# (requires the "oauth2" and "mcp-server-http" experiments to be enabled).
750+
# Whether to inject Coder's MCP tools into intercepted AIBridge requests (requires
751+
# the "oauth2" and "mcp-server-http" experiments to be enabled).
752752
# (default: false, type: bool)
753753
inject_coder_mcp_tools: false
754+
# Length of time to retain data such as interceptions and all related records
755+
# (token, prompt, tool use).
756+
# (default: 1440h, type: duration)
757+
retention: 1440h0m0s

‎coderd/apidoc/docs.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -608,7 +608,7 @@ var (
608608
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.
609609
},
610610
rbac.ResourceApiKey.Type: {policy.ActionRead}, // Validate API keys.
611-
rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
611+
rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
612612
}),
613613
User: []rbac.Permission{},
614614
ByOrgID: map[string]rbac.OrgPermissions{},
@@ -1723,6 +1723,13 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex
17231723
return q.db.DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx, arg)
17241724
}
17251725

1726+
func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) {
1727+
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAibridgeInterception); err != nil {
1728+
return -1, err
1729+
}
1730+
return q.db.DeleteOldAIBridgeRecords(ctx, beforeTime)
1731+
}
1732+
17261733
func (q *querier) DeleteOldAuditLogConnectionEvents(ctx context.Context, threshold database.DeleteOldAuditLogConnectionEventsParams) error {
17271734
// `ResourceSystem` is deprecated, but it doesn't make sense to add
17281735
// `policy.ActionDelete` to `ResourceAuditLog`, since this is the one and

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4654,6 +4654,12 @@ func (s *MethodTestSuite) TestAIBridge() {
46544654
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), params).Return(intc, nil).AnyTimes()
46554655
check.Args(params).Asserts(intc, policy.ActionUpdate).Returns(intc)
46564656
}))
4657+
4658+
s.Run("DeleteOldAIBridgeRecords", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4659+
t := dbtime.Now()
4660+
db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int32(0), nil).AnyTimes()
4661+
check.Args(t).Asserts(rbac.ResourceAibridgeInterception, policy.ActionDelete)
4662+
}))
46574663
}
46584664

46594665
func (s *MethodTestSuite) TestTelemetry() {

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbpurge/dbpurge.go‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"github.com/coder/coder/v2/coderd/database/dbauthz"
1414
"github.com/coder/coder/v2/coderd/database/dbtime"
1515
"github.com/coder/coder/v2/coderd/pproflabel"
16+
"github.com/coder/coder/v2/codersdk"
1617
"github.com/coder/quartz"
1718
)
1819

@@ -36,7 +37,7 @@ const (
3637
// It is the caller's responsibility to call Close on the returned instance.
3738
//
3839
// This is for cleaning up old, unused resources from the database that take up space.
39-
func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.Clock) io.Closer {
40+
func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, clk quartz.Clock) io.Closer {
4041
closed := make(chan struct{})
4142

4243
ctx, cancelFunc := context.WithCancel(ctx)
@@ -90,6 +91,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.
9091
return xerrors.Errorf("failed to delete old audit log connection events: %w", err)
9192
}
9293

94+
deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value())
95+
// nolint:gocritic // Needs to run as aibridge context.
96+
count, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
97+
if err != nil {
98+
return xerrors.Errorf("failed to delete old aibridge records: %w", err)
99+
}
100+
logger.Debug(ctx, "purged aibridge entries", slog.F("count", count), slog.F("since", deleteAIBridgeRecordsBefore.Format(time.RFC3339)))
101+
93102
logger.Debug(ctx, "purged old database entries", slog.F("duration", clk.Since(start)))
94103

95104
return nil

0 commit comments

Comments
 (0)