From 101aee8ee03c39cc25e24d2511e767b935769acd Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Thu, 16 Jul 2026 16:09:25 +0100 Subject: [PATCH] refactor: use Options struct in aibridgedserver.NewServer (#27200) Refactor `aibridgedserver.NewServer` to take an `Options` struct instead of a long list of positional arguments. Follow-up to review feedback in https://github.com/coder/coder/pull/27117#discussion_r3571535760 --- cli/aibridged_internal_test.go | 9 +- coderd/aibridged.go | 13 +- coderd/aibridgedserver/aibridgedserver.go | 48 +++-- .../aibridgedserver/aibridgedserver_test.go | 190 +++++++++++++++--- enterprise/coderd/aibridgeserve.go | 23 +-- 5 files changed, 219 insertions(+), 64 deletions(-) diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index 59b2153c6f..7cd4f64d74 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -54,7 +54,14 @@ func buildFromEnv(t *testing.T, cfg codersdk.AIBridgeConfig) ([]aibridge.Provide // (providers, outcomes) the embedded reloader would observe. func buildFromDB(ctx context.Context, t *testing.T, db database.Store, cfg codersdk.AIBridgeConfig, logger slog.Logger) ([]aibridge.Provider, []aibridged.ProviderOutcome, error) { t.Helper() - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: cfg, + Logger: logger, + Clock: quartz.NewReal(), + }) if err != nil { return nil, nil, err } diff --git a/coderd/aibridged.go b/coderd/aibridged.go index 3784545f7c..d163fbfe09 100644 --- a/coderd/aibridged.go +++ b/coderd/aibridged.go @@ -65,8 +65,17 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai }() mux := drpcmux.New() - srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Pubsub, api.Logger.Named("aibridgedserver"), - api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker, api.Clock) + srv, err := aibridgedserver.NewServer(api.ctx, aibridgedserver.Options{ + Store: api.Database, + Pubsub: api.Pubsub, + AISeatTracker: api.AISeatTracker, + AccessURL: api.AccessURL.String(), + GatewayCfg: api.DeploymentValues.AI.BridgeConfig, + ExternalAuthConfigs: api.ExternalAuthConfigs, + Experiments: api.Experiments, + Logger: api.Logger.Named("aibridgedserver"), + Clock: api.Clock, + }) if err != nil { return nil, err } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 2aaf76d6bb..c17a087920 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -120,13 +120,25 @@ type Server struct { clock quartz.Clock } -func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logger slog.Logger, accessURL string, - bridgeCfg codersdk.AIBridgeConfig, externalAuthConfigs []*externalauth.Config, experiments codersdk.Experiments, - aiSeatTracker aiseats.SeatTracker, clock quartz.Clock, -) (*Server, error) { - eac := make(map[string]*externalauth.Config, len(externalAuthConfigs)) +// Options carries the dependencies required to construct an aibridged Server. +type Options struct { + Store store + Pubsub pubsub.Pubsub + AISeatTracker aiseats.SeatTracker - for _, cfg := range externalAuthConfigs { + AccessURL string + GatewayCfg codersdk.AIBridgeConfig + ExternalAuthConfigs []*externalauth.Config + Experiments codersdk.Experiments + + Logger slog.Logger + Clock quartz.Clock +} + +func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) { + eac := make(map[string]*externalauth.Config, len(opts.ExternalAuthConfigs)) + + for _, cfg := range opts.ExternalAuthConfigs { // Only External Auth configs which are configured with an MCP URL are relevant to aibridged. if cfg.MCPURL == "" { continue @@ -136,22 +148,22 @@ func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logg srv := &Server{ lifecycleCtx: lifecycleCtx, - store: store, - pubsub: ps, - logger: logger, + store: opts.Store, + pubsub: opts.Pubsub, + logger: opts.Logger, externalAuthConfigs: eac, - structuredLogging: bridgeCfg.StructuredLogging.Value(), - aiSeatTracker: aiSeatTracker, - budgetPolicy: codersdk.NewAIBudgetPolicyFromString(bridgeCfg.BudgetPolicy), - budgetPeriod: codersdk.NewAIBudgetPeriodFromString(bridgeCfg.BudgetPeriod), - clock: clock, + structuredLogging: opts.GatewayCfg.StructuredLogging.Value(), + aiSeatTracker: opts.AISeatTracker, + budgetPolicy: codersdk.NewAIBudgetPolicyFromString(opts.GatewayCfg.BudgetPolicy), + budgetPeriod: codersdk.NewAIBudgetPeriodFromString(opts.GatewayCfg.BudgetPeriod), + clock: opts.Clock, } - if bridgeCfg.InjectCoderMCPTools { - logger.Warn(lifecycleCtx, "inject MCP tools option is deprecated and will be removed in a future release") - coderMCPConfig, err := getCoderMCPServerConfig(experiments, accessURL) + if opts.GatewayCfg.InjectCoderMCPTools { + opts.Logger.Warn(lifecycleCtx, "inject MCP tools option is deprecated and will be removed in a future release") + coderMCPConfig, err := getCoderMCPServerConfig(opts.Experiments, opts.AccessURL) if err != nil { - logger.Warn(lifecycleCtx, "failed to retrieve coder MCP server config, Coder MCP will not be available", slog.Error(err)) + opts.Logger.Warn(lifecycleCtx, "failed to retrieve coder MCP server config, Coder MCP will not be available", slog.Error(err)) } srv.coderMCPConfig = coderMCPConfig } diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index e5b71803eb..bf28dbdf9e 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -209,7 +209,15 @@ func TestAuthorization(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) require.NotNil(t, srv) @@ -371,7 +379,15 @@ func TestAuthorization_Delegated(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) require.NotNil(t, srv) @@ -560,7 +576,15 @@ func TestIsBudgetExceeded(t *testing.T) { wantResp = tc.setupMocks(db, userID) } - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) req := &proto.IsBudgetExceededRequest{UserId: userIDStr} @@ -608,7 +632,15 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { }) require.NoError(t, err) - srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, clock) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: authzDB, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: clock, + }) require.NoError(t, err) return ctx, rawDB, srv, user, group @@ -771,9 +803,18 @@ func TestGetMCPServerConfigs(t *testing.T) { logger := testutil.Logger(t) accessURL := "https://my-cool-deployment.com" - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, accessURL, codersdk.AIBridgeConfig{ - InjectCoderMCPTools: serpent.Bool(!tc.disableCoderMCPInjection), - }, tc.externalAuthConfigs, tc.experiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: accessURL, + GatewayCfg: codersdk.AIBridgeConfig{ + InjectCoderMCPTools: serpent.Bool(!tc.disableCoderMCPInjection), + }, + ExternalAuthConfigs: tc.externalAuthConfigs, + Experiments: tc.experiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) require.NotNil(t, srv) @@ -811,19 +852,28 @@ func TestGetMCPServerAccessTokensBatch(t *testing.T) { logger := testutil.Logger(t) // Given: 2 external auth configured with MCP and 1 without. - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, []*externalauth.Config{ - { - ID: "1", - MCPURL: "1.com/mcp", + srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + ExternalAuthConfigs: []*externalauth.Config{ + { + ID: "1", + MCPURL: "1.com/mcp", + }, + { + ID: "2", + MCPURL: "2.com/mcp", + }, + { + ID: "3", + }, }, - { - ID: "2", - MCPURL: "2.com/mcp", - }, - { - ID: "3", - }, - }, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) require.NotNil(t, srv) @@ -2100,7 +2150,15 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC) // The server runs every store call as subjectAibridged via the authzDB. - srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: authzDB, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) _, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ @@ -2475,7 +2533,15 @@ func testRecordMethod[Req any, Resp any]( } ctx := testutil.Context(t, testutil.WaitLong) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) resp, err := callMethod(srv, ctx, tc.request) @@ -2795,9 +2861,17 @@ func TestStructuredLogging(t *testing.T) { tc.setupMocks(db, interceptionID) ctx := testutil.Context(t, testutil.WaitLong) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{ - StructuredLogging: serpent.Bool(tc.structuredLogging), - }, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{ + StructuredLogging: serpent.Bool(tc.structuredLogging), + }, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) err = tc.recordFn(srv, ctx, interceptionID) @@ -2839,7 +2913,15 @@ func TestInferredThreadsByToolCalls(t *testing.T) { user := dbgen.User(t, db, database.User{}) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) aID := uuid.New() @@ -2935,7 +3017,15 @@ func TestRecordToolUsageProviderItemID(t *testing.T) { user := dbgen.User(t, db, database.User{}) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Experiments: requiredExperiments, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) intcID := uuid.New() @@ -3065,7 +3155,14 @@ func TestGetAIProviders(t *testing.T) { Settings: sql.NullString{String: "{not valid json", Valid: true}, }) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) resp, err := srv.GetAIProviders(ctx, &proto.GetAIProvidersRequest{}) @@ -3128,7 +3225,14 @@ func TestGetAIProvidersBlocksOnSeedLock(t *testing.T) { BaseUrl: "https://api.openai.com/", }, "sk-openai") - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) // Simulate an in-flight env seed holding the advisory lock until released. @@ -3209,7 +3313,15 @@ func TestWatchAIProviders(t *testing.T) { // In-memory pubsub delivers Publish synchronously for deterministic signals. ps := pubsub.NewInMemory() - srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + Pubsub: ps, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) streamCtx, streamCancel := context.WithCancel(ctx) @@ -3247,7 +3359,15 @@ func TestWatchAIProvidersSignalsOnDeliveryError(t *testing.T) { logger := slogtest.Make(t, nil) ps := &captureListenerPubsub{listenerC: make(chan pubsub.ListenerWithErr, 1)} - srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{ + Store: db, + Pubsub: ps, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) streamCtx, streamCancel := context.WithCancel(ctx) @@ -3289,7 +3409,15 @@ func TestWatchAIProvidersStopsOnLifecycleCancel(t *testing.T) { // canceled while the stream stays open. lifecycleCtx, lifecycleCancel := context.WithCancel(ctx) defer lifecycleCancel() - srv, err := aibridgedserver.NewServer(lifecycleCtx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) + srv, err := aibridgedserver.NewServer(lifecycleCtx, aibridgedserver.Options{ + Store: db, + Pubsub: ps, + AISeatTracker: agplaiseats.Noop{}, + AccessURL: "/", + GatewayCfg: codersdk.AIBridgeConfig{}, + Logger: logger, + Clock: quartz.NewReal(), + }) require.NoError(t, err) streamCtx, streamCancel := context.WithCancel(ctx) diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 8e90ae6e71..14663c4610 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -134,18 +134,17 @@ func (api *API) aiGatewayServe(rw http.ResponseWriter, r *http.Request) { go aiGatewayCheckEntitlementAndTrackKeyUsage(connCtx, keyCtxCancel, api, gatewayKey.ID, logger) mux := drpcmux.New() - srv, err := aibridgedserver.NewServer( - connCtx, - api.Database, - api.AGPL.Pubsub, - logger, - api.AccessURL.String(), - api.DeploymentValues.AI.BridgeConfig, - api.ExternalAuthConfigs, - api.AGPL.Experiments, - api.AGPL.AISeatTracker, - api.AGPL.Clock, - ) + srv, err := aibridgedserver.NewServer(connCtx, aibridgedserver.Options{ + Store: api.Database, + Pubsub: api.AGPL.Pubsub, + AISeatTracker: api.AGPL.AISeatTracker, + AccessURL: api.AccessURL.String(), + GatewayCfg: api.DeploymentValues.AI.BridgeConfig, + ExternalAuthConfigs: api.ExternalAuthConfigs, + Experiments: api.AGPL.Experiments, + Logger: logger, + Clock: api.AGPL.Clock, + }) if err != nil { if !xerrors.Is(err, context.Canceled) { logger.Error(connCtx, "server creation failed", slog.Error(err))