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

refactor: use Options struct in aibridgedserver.NewServer (#27200)

Browse Source 
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
This commit is contained in:
Susana Ferreira
2026-07-16 16:09:25 +01:00
committed by GitHub
parent 77582be805
commit 101aee8ee0
5 changed files with 219 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/aibridged.go
+11 -2
View File
@@ -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
}
coderd/aibridgedserver/aibridgedserver.go
+30 -18
View File
@@ -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
}
coderd/aibridgedserver/aibridgedserver_test.go
+159 -31
View File
@@ -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)