Function newAIBridgeDaemon

cli/aibridged.go:33–78

newAIBridgeDaemon constructs the in-memory aibridge daemon and wires up a subscription that hot-reloads the provider pool from the database on every ai_providers change event. The returned unsubscribe function tears down the subscription; callers must invoke it alongside Server.Close on shutdown.

(coderAPI *coderd.API, providers []aibridge.Provider, cfg codersdk.AIBridgeConfig)

// function tears down the subscription; callers must invoke it
// alongside Server.Close on shutdown.
func newAIBridgeDaemon(coderAPI *coderd.API, providers []aibridge.Provider, cfg codersdk.AIBridgeConfig) (*aibridged.Server, func(), error) {
	ctx := context.Background()
	coderAPI.Logger.Debug(ctx, "starting in-memory aibridge daemon")

	logger := coderAPI.Logger.Named("aibridged")

	reg := prometheus.WrapRegistererWithPrefix("coder_aibridged_", coderAPI.PrometheusRegistry)
	metrics := aibridge.NewMetrics(reg)
	providerMetrics := aibridged.NewMetrics(reg)
	tracer := coderAPI.TracerProvider.Tracer(tracing.TracerName)

	// Create pool for reusable stateful [aibridge.RequestBridge] instances (one per user).
	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger.Named("pool"), metrics, tracer) // TODO: configurable size.
	if err != nil {
		return nil, nil, xerrors.Errorf("create request pool: %w", err)
	}

	// Subscribe to ai_providers change events so the pool tracks the
	// database without a restart. The boot-time `providers` snapshot
	// derives from env config and serves as a fallback if the database
	// load fails inside the reloader.
	reloader := &poolDBReloader{
		pool:    pool,
		db:      coderAPI.Database,
		cfg:     cfg,
		logger:  logger.Named("provider-loader"),
		metrics: providerMetrics,
	}
	unsubscribe, err := aibridged.SubscribeProviderReload(ctx, coderAPI.Pubsub, reloader, logger.Named("provider-reload"))
	if err != nil {
		// Pool is still usable with the boot-time snapshot; subscription
		// failure is logged but not fatal so the daemon still serves.
		logger.Warn(ctx, "subscribe to ai providers change channel", slog.Error(err))
		unsubscribe = func() {}
	}

	// Create daemon.
	srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
		return coderAPI.CreateInMemoryAIBridgeServer(dialCtx)
	}, logger, tracer)
	if err != nil {
		unsubscribe()
		return nil, nil, xerrors.Errorf("start in-memory aibridge daemon: %w", err)
	}
	return srv, unsubscribe, nil
}

// poolDBReloader implements [aibridged.ProviderReloader] by loading
// the live provider set from the database and forwarding it to the