Thanks to visit codestin.com
Credit goes to github.com

Skip to content

feat: synchronise provider changes with WatchAIProviders #26797

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dannykopping merged 6 commits into from
Jul 8, 2026
Merged

feat: synchronise provider changes with WatchAIProviders #26797

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
049761f
feat(coderd/aibridged): push provider changes to gateways via WatchAI…
dannykopping Jun 26, 2026
94a444b
review: address coder-agents-review feedback on WatchAIProviders
dannykopping Jun 26, 2026
b1f9597
review: fix shutdown deadlock in standalone gateway watch drain
dannykopping Jun 29, 2026
1aa5d2e
review: address coder-agents-review feedback on provider reload
dannykopping Jun 29, 2026
8bc4680
test(coderd/aibridgedserver): cover lifecycle-cancel and pubsub inter…
dannykopping Jun 29, 2026
3b84826
fix: make provider reload client acquisition context-aware
dannykopping Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 9 additions & 7 deletions cli/aibridged.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr
reg.MustRegister(keypool.NewStateCollector(pool.KeyPools))

// Create daemon. Construct it before subscribing so the reloader can use
// srv.Client() to fetch providers over the in-memory RPC.
// srv.ClientContext to fetch providers over the in-memory RPC.
srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
return coderAPI.CreateInMemoryAIBridgeServer(dialCtx)
}, logger, tracer)
Expand All @@ -72,7 +72,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr
// Subscribe to ai_providers change events so the pool tracks the database
// without a restart, and perform the initial reload. The reload data path
// is the in-memory RPC.
reloader := NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("provider-loader"), metrics, providerMetrics)
reloader := NewPoolRPCReloader(pool, srv.ClientContext, cfg, logger.Named("provider-loader"), metrics, providerMetrics)
unsubscribe, err := aibridged.SubscribeProviderReload(ctx, coderAPI.Pubsub, reloader, logger.Named("provider-reload"))
if err != nil {
// Without the subscription the pool can never track provider changes,
Expand All @@ -91,7 +91,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr
// build, replace, and reload-metric accounting live in one place.
type poolRPCReloader struct {
pool *aibridged.CachedBridgePool
client func() (aibridged.DRPCClient, error)
client aibridged.ClientFuncWithContext
cfg codersdk.AIBridgeConfig
logger slog.Logger
aibridgeMetrics *aibridge.Metrics
Expand All @@ -100,10 +100,12 @@ type poolRPCReloader struct {

// NewPoolRPCReloader builds an [aibridged.ProviderReloader] that fetches the
// provider set over the DRPC client returned by client and replaces pool's
// providers, recording reload metrics against providerMetrics.
// providers, recording reload metrics against providerMetrics. client receives
// Reload's context, so a blocking acquisition unblocks when that context is
// canceled.
func NewPoolRPCReloader(
pool *aibridged.CachedBridgePool,
client func() (aibridged.DRPCClient, error),
client aibridged.ClientFuncWithContext,
cfg codersdk.AIBridgeConfig,
logger slog.Logger,
aibridgeMetrics *aibridge.Metrics,
Expand All @@ -121,8 +123,8 @@ func NewPoolRPCReloader(

func (r *poolRPCReloader) Reload(ctx context.Context) error {
r.providerMetrics.RecordReloadAttempt()
// r.client() blocks until the daemon is connected to coderd.
client, err := r.client()
// r.client blocks until the daemon connects to coderd or ctx is canceled.
client, err := r.client(ctx)
if err != nil {
return xerrors.Errorf("get ai-gateway client: %w", err)
}
Expand Down
2 changes: 1 addition & 1 deletion cli/aibridged_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ 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, logger, "/", cfg, nil, nil, agplaiseats.Noop{})
srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{})
if err != nil {
return nil, nil, err
}
Expand Down
2 changes: 1 addition & 1 deletion coderd/aibridged.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai
}()

mux := drpcmux.New()
srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Logger.Named("aibridgedserver"),
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)
if err != nil {
return nil, err
Expand Down
15 changes: 15 additions & 0 deletions coderd/aibridged/aibridgedmock/clientmock.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions coderd/aibridged/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,11 @@ type Dialer func(ctx context.Context) (DRPCClient, error)

type ClientFunc func() (DRPCClient, error)

// ClientFuncWithContext acquires a DRPCClient, honoring the passed context so a
// blocking acquisition (e.g. waiting for the daemon to connect to coderd)
// unblocks when the context is canceled. Server.ClientContext satisfies it.
type ClientFuncWithContext func(context.Context) (DRPCClient, error)

// DRPCClient is the union of various service interfaces the client must support.
type DRPCClient interface {
proto.DRPCRecorderClient
Expand Down
Loading
Loading