feat: add stateless AI Gateway proxy mode (#29349)

Adds proxy mode to the AI Gateway. No AI provider is supported in this
mode yet.

aibridge/proxy.go

@@ -0,0 +1,59 @@
+package aibridge
+
+import (
+	"net/http"
+	"slices"
+
+	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/aibridge/keypool"
+)
+
+// ProxyRouter is an [http.Handler] which serves the AI Gateway routes of an
+// immutable snapshot of providers. Unlike [RequestBridge] it holds no
+// per-user or per-request state, so a single instance serves every actor and
+// is replaced wholesale when the provider configuration changes.
+//
+// ProxyRouter has no concept of authentication or authorization; the caller is
+// responsible for authenticating requests before they reach the router.
+//
+// ProxyRouter is safe for concurrent use.
+type ProxyRouter struct {
+	mux *http.ServeMux
+
+	// providers is an owned copy of the provider list, fixed at construction.
+	providers []Provider
+}
+
+var _ http.Handler = (*ProxyRouter)(nil)
+
+// NewProxyRouter creates a [*ProxyRouter] for the given providers, whose
+// names must be valid and unique.
+//
+// Each configured-but-disabled provider serves a 503 sentinel on every path
+// under its name. Enabled providers have no routes registered yet, so their
+// requests reach the catch-all 404.
+func NewProxyRouter(providers []Provider, logger slog.Logger) (*ProxyRouter, error) {
+	if err := validateProviders(providers); err != nil {
+		return nil, err
+	}
+
+	snapshot := slices.Clone(providers)
+	mux := newProviderMux(snapshot, logger)
+
+	return &ProxyRouter{
+		mux:       mux,
+		providers: snapshot,
+	}, nil
+}
+
+// ServeHTTP serves the routes registered for the router's providers.
+func (p *ProxyRouter) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	// Cap the body as it is read; routes that do not read it retain their status.
+	r.Body = http.MaxBytesReader(rw, r.Body, maxRequestBodyBytes)
+	p.mux.ServeHTTP(rw, r)
+}
+
+// KeyPools returns the non-nil key pools from this router's provider snapshot.
+func (p *ProxyRouter) KeyPools() []*keypool.Pool {
+	return CollectKeyPools(p.providers)
+}

aibridge/proxy_internal_test.go

@@ -0,0 +1,33 @@
+package aibridge
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+)
+
+// Disabled providers return 503 without reading the body, even if it is oversized.
+func TestProxyRouterDisabledProviderOversizedBody(t *testing.T) {
+	t.Parallel()
+
+	router, err := NewProxyRouter(
+		[]Provider{NewDisabledProviderStub("disabled-openai", "openai")},
+		slogtest.Make(t, nil),
+	)
+	require.NoError(t, err)
+
+	req := httptest.NewRequest(http.MethodPost, "/disabled-openai/v1/chat/completions", strings.NewReader("x"))
+	req.ContentLength = maxRequestBodyBytes + 1
+
+	rec := httptest.NewRecorder()
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+	assert.Contains(t, rec.Body.String(), ErrorCodeProviderDisabled)
+}

aibridge/proxy_test.go

@@ -0,0 +1,143 @@
+package aibridge_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	promtestutil "github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/aibridge/config"
+	"github.com/coder/coder/v2/aibridge/internal/testutil"
+	"github.com/coder/coder/v2/aibridge/keypool"
+)
+
+type keyPoolProvider struct {
+	aibridge.Provider
+	pool *keypool.Pool
+}
+
+func (p keyPoolProvider) KeyPool() *keypool.Pool { return p.pool }
+
+func TestNewProxyRouterValidatesProviders(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name        string
+		providers   []aibridge.Provider
+		errContains string
+	}{
+		{
+			name:      "ValidNames",
+			providers: []aibridge.Provider{&testutil.MockProvider{NameStr: "openai"}, &testutil.MockProvider{NameStr: "anthropic-eu-1"}},
+		},
+		{
+			name:        "InvalidName",
+			providers:   []aibridge.Provider{&testutil.MockProvider{NameStr: "OpenAI_1"}},
+			errContains: `invalid provider name "OpenAI_1"`,
+		},
+		{
+			name:        "DuplicateName",
+			providers:   []aibridge.Provider{&testutil.MockProvider{NameStr: "openai"}, &testutil.MockProvider{NameStr: "openai"}},
+			errContains: `duplicate provider name: "openai"`,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			router, err := aibridge.NewProxyRouter(tc.providers, slogtest.Make(t, nil))
+			if tc.errContains != "" {
+				require.ErrorContains(t, err, tc.errContains)
+				require.Nil(t, router)
+				return
+			}
+			require.NoError(t, err)
+			require.NotNil(t, router)
+		})
+	}
+}
+
+// TestProxyRouterDisabledProvider asserts that every path under a disabled
+// provider's name serves the 503 sentinel, while a sibling enabled provider
+// has no routes yet and falls through to the catch-all.
+func TestProxyRouterDisabledProvider(t *testing.T) {
+	t.Parallel()
+
+	// Any upstream call would fail loudly: the base URL points nowhere and
+	// the mock provider has no interceptor.
+	enabled := &testutil.MockProvider{
+		NameStr:     "openai",
+		URL:         "http://127.0.0.1:1",
+		Bridged:     []string{"/v1/chat/completions"},
+		Passthrough: []string{"/v1/models"},
+	}
+	router, err := aibridge.NewProxyRouter(
+		[]aibridge.Provider{enabled, aibridge.NewDisabledProviderStub("disabled-openai", "openai")},
+		slogtest.Make(t, nil),
+	)
+	require.NoError(t, err)
+
+	for _, tc := range []struct {
+		name       string
+		path       string
+		wantStatus int
+		wantBody   string
+	}{
+		{name: "DisabledBridgedRoute", path: "/disabled-openai/v1/chat/completions", wantStatus: http.StatusServiceUnavailable, wantBody: aibridge.ErrorCodeProviderDisabled},
+		{name: "DisabledPassthroughRoute", path: "/disabled-openai/v1/models", wantStatus: http.StatusServiceUnavailable, wantBody: aibridge.ErrorCodeProviderDisabled},
+		{name: "DisabledUnknownRoute", path: "/disabled-openai/anything/else", wantStatus: http.StatusServiceUnavailable, wantBody: aibridge.ErrorCodeProviderDisabled},
+		{name: "EnabledBridgedRoute", path: "/openai/v1/chat/completions", wantStatus: http.StatusNotFound, wantBody: "route not supported"},
+		{name: "EnabledPassthroughRoute", path: "/openai/v1/models", wantStatus: http.StatusNotFound, wantBody: "route not supported"},
+		{name: "UnknownProvider", path: "/unknown/v1/models", wantStatus: http.StatusNotFound, wantBody: "route not supported"},
+		{name: "Root", path: "/", wantStatus: http.StatusNotFound, wantBody: "route not supported"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			resp := httptest.NewRecorder()
+			router.ServeHTTP(resp, httptest.NewRequest(http.MethodPost, tc.path, nil))
+
+			assert.Equal(t, tc.wantStatus, resp.Code)
+			assert.Contains(t, resp.Body.String(), tc.wantBody)
+		})
+	}
+}
+
+// TestProxyRouterSnapshotsProviders asserts the router routes and reports key
+// pools from its own snapshot. Mutating the caller's slice afterwards has
+// no effect.
+func TestProxyRouterSnapshotsProviders(t *testing.T) {
+	t.Parallel()
+
+	pool := testutil.SingleKeyPool(config.ProviderOpenAI, "test-key")
+	providers := []aibridge.Provider{
+		keyPoolProvider{Provider: aibridge.NewDisabledProviderStub("disabled-openai", "openai"), pool: pool},
+	}
+
+	router, err := aibridge.NewProxyRouter(providers, slogtest.Make(t, nil))
+	require.NoError(t, err)
+
+	// Replace the caller's entry with a provider the router never saw.
+	providers[0] = &testutil.MockProvider{NameStr: "swapped"}
+
+	require.NoError(t, promtestutil.CollectAndCompare(keypool.NewStateCollector(router.KeyPools), strings.NewReader(`
+# HELP key_pool_state The number of keys currently in each state (state: valid, temporary, permanent).
+# TYPE key_pool_state gauge
+key_pool_state{provider="openai",state="valid"} 1
+key_pool_state{provider="openai",state="temporary"} 0
+key_pool_state{provider="openai",state="permanent"} 0
+`), "key_pool_state"))
+
+	resp := httptest.NewRecorder()
+	router.ServeHTTP(resp, httptest.NewRequest(http.MethodPost, "/disabled-openai/v1/models", nil))
+	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
+
+	resp = httptest.NewRecorder()
+	router.ServeHTTP(resp, httptest.NewRequest(http.MethodPost, "/swapped/v1/models", nil))
+	assert.Equal(t, http.StatusNotFound, resp.Code)
+}

cli/aibridged.go

@@ -24,7 +24,7 @@ import (
 )
 
 // newAIBridgeDaemon constructs the in-memory aibridge daemon and wires
-// up a subscription that hot-reloads the provider pool over the in-memory
+// up a subscription that hot-reloads providers over the in-memory
 // RPC on every ai_providers change event. The returned unsubscribe
 // function tears down the subscription; callers must invoke it
 // alongside Server.Close on shutdown.
@@ -35,7 +35,7 @@ import (
 // hot-reload trigger.
 //
 // SubscribeProviderReload performs a best-effort initial reload synchronously,
-// so the pool is populated before this returns whenever the fetch succeeds.
+// so providers are populated before this returns whenever the fetch succeeds.
 // That reload blocks while acquiring a client, and it passes a background
 // context, so only the daemon lifecycle bounds the wait. That is acceptable
 // here: the embedded daemon's connection is an in-memory pipe that comes up
@@ -49,33 +49,24 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr
 	providerMetrics := aibridged.NewMetrics(reg)
 	tracer := coderAPI.TracerProvider.Tracer(tracing.TracerName)
 
-	// Create an empty pool for reusable stateful [aibridge.RequestBridge]
-	// instances (one per user). The reloader populates it via the initial
-	// reload below.
-	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger.Named("pool"), metrics, tracer) // TODO: configurable size.
-	if err != nil {
-		return nil, nil, xerrors.Errorf("create request pool: %w", err)
-	}
-
-	// Report current key pool state per provider at scrape time.
-	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, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
+	srv, err := aibridged.New(ctx, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
 		return coderAPI.CreateInMemoryAIBridgeServer(dialCtx)
-	}, logger, tracer)
+	}, logger, tracer, coderAPI.Experiments, metrics)
 	if err != nil {
 		return nil, nil, xerrors.Errorf("start in-memory aibridge daemon: %w", err)
 	}
 
-	// Subscribe to ai_providers change events so the pool tracks the database
+	reg.MustRegister(srv.KeyPoolStateCollector())
+
+	// Subscribe to ai_providers change events so the backend 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 := NewProviderRPCReloader(srv.ReplaceProviders, srv.Client, 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,
+		// Without the subscription the backend cannot track provider changes,
 		// so fail startup rather than serve a permanently stale snapshot.
 		_ = srv.Close()
 		return nil, nil, xerrors.Errorf("subscribe to ai providers change channel: %w", err)
@@ -84,44 +75,42 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr
 	return srv, unsubscribe, nil
 }
 
-// poolRPCReloader implements [aibridged.ProviderReloader] by fetching the
-// live provider set from coderd over a DRPC client and forwarding it to the
-// pool. It is shared by the embedded daemon (in-memory RPC, pubsub-triggered)
-// and the standalone gateway (WebSocket RPC, retried at startup) so the fetch,
+// providerRPCReloader implements [aibridged.ProviderReloader] by fetching the
+// live provider set from coderd over a DRPC client and publishing it to the
+// selected backend. Embedded (in-memory RPC, pubsub-triggered) and standalone
+// (WebSocket RPC, retried at startup) gateways share the fetch,
 // build, replace, and reload-metric accounting live in one place.
-type poolRPCReloader struct {
-	pool            aibridged.Pooler
-	client          aibridged.ClientFunc
-	cfg             codersdk.AIBridgeConfig
-	logger          slog.Logger
-	aibridgeMetrics *aibridge.Metrics
-	providerMetrics *aibridged.Metrics
+type providerRPCReloader struct {
+	replaceProviders func(context.Context, []aibridge.Provider) error
+	client           aibridged.ClientFunc
+	cfg              codersdk.AIBridgeConfig
+	logger           slog.Logger
+	aibridgeMetrics  *aibridge.Metrics
+	providerMetrics  *aibridged.Metrics
 }
 
-// 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. client receives
-// Reload's context, so a blocking acquisition unblocks when that context is
-// canceled.
-func NewPoolRPCReloader(
-	pool aibridged.Pooler,
+// NewProviderRPCReloader fetches and builds provider snapshots over DRPC and
+// publishes them through replaceProviders. Failed publication is reported as a
+// failed reload, preserving the selected backend's previous snapshot.
+func NewProviderRPCReloader(
+	replaceProviders func(context.Context, []aibridge.Provider) error,
 	client aibridged.ClientFunc,
 	cfg codersdk.AIBridgeConfig,
 	logger slog.Logger,
 	aibridgeMetrics *aibridge.Metrics,
 	providerMetrics *aibridged.Metrics,
 ) aibridged.ProviderReloader {
-	return &poolRPCReloader{
-		pool:            pool,
-		client:          client,
-		cfg:             cfg,
-		logger:          logger,
-		aibridgeMetrics: aibridgeMetrics,
-		providerMetrics: providerMetrics,
+	return &providerRPCReloader{
+		replaceProviders: replaceProviders,
+		client:           client,
+		cfg:              cfg,
+		logger:           logger,
+		aibridgeMetrics:  aibridgeMetrics,
+		providerMetrics:  providerMetrics,
 	}
 }
 
-func (r *poolRPCReloader) Reload(ctx context.Context) error {
+func (r *providerRPCReloader) Reload(ctx context.Context) error {
 	r.providerMetrics.RecordReloadAttempt()
 	// r.client blocks until the daemon connects to coderd or ctx is canceled.
 	client, err := r.client(ctx)
@@ -136,7 +125,9 @@ func (r *poolRPCReloader) Reload(ctx context.Context) error {
 		return xerrors.Errorf("fetch ai providers: %w", err)
 	}
 	providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), r.cfg, r.logger, r.aibridgeMetrics)
-	r.pool.ReplaceProviders(providers)
+	if err := r.replaceProviders(ctx, providers); err != nil {
+		return xerrors.Errorf("replace ai providers: %w", err)
+	}
 	r.providerMetrics.RecordReloadSuccess(outcomes)
 	return nil
 }

cli/aibridged_backend_internal_test.go

@@ -0,0 +1,58 @@
+//go:build !slim
+
+package cli
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/coderd/aibridged"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestNewAIBridgeDaemonBackend asserts the embedded daemon reaches its mode
+// selection from coderAPI.Experiments. The remaining mode matrix is shared with
+// the standalone E2E test and covered by coderd/aibridged.
+func TestNewAIBridgeDaemonBackend(t *testing.T) {
+	t.Parallel()
+	for _, tc := range []struct {
+		name       string
+		experiment bool
+		proxy      bool
+	}{
+		{name: "Interception"},
+		{name: "Proxy", experiment: true, proxy: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			dv := coderdtest.DeploymentValues(t)
+			if tc.experiment {
+				dv.Experiments = append(dv.Experiments, string(codersdk.ExperimentAIGatewayReverseProxy))
+			}
+			client, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{DeploymentValues: dv})
+			firstUser := coderdtest.CreateFirstUser(t, client)
+			srv, unsubscribe, err := newAIBridgeDaemon(api, dv.AI.BridgeConfig, prometheus.NewRegistry(), nil)
+			require.NoError(t, err)
+			t.Cleanup(func() { require.NoError(t, srv.Close()) })
+			t.Cleanup(unsubscribe)
+			handler, err := srv.GetRequestHandler(testutil.Context(t, testutil.WaitLong), aibridged.Request{InitiatorID: firstUser.UserID, SessionKey: client.SessionToken()})
+			require.NoError(t, err)
+			_, isBridge := handler.(*aibridge.RequestBridge)
+			require.Equal(t, !tc.proxy, isBridge, "interception serves from a RequestBridge")
+			if tc.proxy {
+				// A published router answers unregistered routes with 404;
+				// the not-ready handler would answer 503.
+				rec := httptest.NewRecorder()
+				handler.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil))
+				require.Equal(t, http.StatusNotFound, rec.Code)
+			}
+		})
+	}
+}

cli/aibridged_reload_internal_test.go

@@ -0,0 +1,61 @@
+//go:build !slim
+
+package cli
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	promtest "github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/coderd/aibridged"
+	"github.com/coder/coder/v2/coderd/aibridged/aibridgedmock"
+	"github.com/coder/coder/v2/coderd/aibridged/proto"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+func TestProviderRPCReloaderPublication(t *testing.T) {
+	t.Parallel()
+	for _, fail := range []bool{false, true} {
+		name := "Success"
+		if fail {
+			name = "Failure"
+		}
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			client := aibridgedmock.NewMockDRPCClient(gomock.NewController(t))
+			client.EXPECT().GetAIProviders(gomock.Any(), gomock.Any()).Return(&proto.GetAIProvidersResponse{}, nil)
+			metrics := aibridged.NewMetrics(prometheus.NewRegistry())
+			publishErr := xerrors.New("publish failed")
+			called := false
+			reloader := NewProviderRPCReloader(func(ctx context.Context, providers []aibridge.Provider) error {
+				called = true
+				require.Equal(t, t.Context(), ctx)
+				require.Empty(t, providers)
+				if fail {
+					return publishErr
+				}
+				return nil
+			}, func(context.Context) (aibridged.DRPCClient, error) {
+				return client, nil
+			}, codersdk.AIBridgeConfig{}, slogtest.Make(t, nil), nil, metrics)
+
+			err := reloader.Reload(t.Context())
+			require.True(t, called)
+			require.Positive(t, promtest.ToFloat64(metrics.ProvidersLastReloadTimestampSeconds))
+			if fail {
+				require.ErrorIs(t, err, publishErr)
+				require.Zero(t, promtest.ToFloat64(metrics.ProvidersLastReloadSuccessTimestampSeconds))
+			} else {
+				require.NoError(t, err)
+				require.Positive(t, promtest.ToFloat64(metrics.ProvidersLastReloadSuccessTimestampSeconds))
+			}
+		})
+	}
+}

coderd/aibridged/aibridged.go

@@ -9,34 +9,48 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/aibridge/keypool"
+	"github.com/coder/coder/v2/coderd/aibridged/proto"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/retry"
 )
 
+// mcpConfigTimeout bounds the startup MCP discovery call. The connect loop
+// retries on failure, so exceeding it only delays mode selection.
+const mcpConfigTimeout = 10 * time.Second
+
 var (
 	_ io.Closer = &Server{}
 
 	ErrShutdown = xerrors.New("aibridged server shutdown")
 )
 
-// Server provides the AI Bridge functionality.
-// It is responsible for:
-//   - receiving requests on /api/v2/aibridged/*
-//   - manipulating the requests
-//   - relaying requests to upstream AI services and relaying responses to caller
-//
-// It requires a [Dialer] to provide a [DRPCClient] implementation to
-// communicate with a [DRPCServer] implementation, to persist state and perform other functions.
+// Server manages the AI Gateway's connection to coderd and request handlers
+// for embedded and standalone deployments. It selects interception or proxy
+// mode at startup, applies provider snapshots, and coordinates shutdown.
 type Server struct {
 	clientDialer Dialer
 	clientCh     chan DRPCClient
 
-	// A pool of [aibridge.RequestBridge] instances, which service incoming requests.
-	requestBridgePool Pooler
+	// backend holds the current request handler. Mode is fixed at startup.
+	// Provider reloads update the pool in interception mode or atomically replace
+	// the handler in proxy mode. In-flight requests retain their original router.
+	backend   atomic.Pointer[backend]
+	backendMu sync.Mutex
+
+	// inflight tracks proxy requests across router snapshots. Shutdown waits
+	// for completion or cancels their contexts when its deadline expires.
+	inflight *aibridge.InflightGate
+
+	// reverseProxyExp is the experiment flag. Proxy mode also requires no MCP configs.
+	reverseProxyExp bool
+	metrics         *aibridge.Metrics
 
 	logger slog.Logger
 	tracer trace.Tracer
@@ -54,7 +68,18 @@ type Server struct {
 	shutdownOnce sync.Once
 }
 
-func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger, tracer trace.Tracer) (*Server, error) {
+// backend holds either an interception pool or a proxy router.
+// Its fields are fixed after publication. Both pool and proxyRouter are nil when
+// proxy mode is selected but providers are not yet loaded.
+type backend struct {
+	pool        Pooler                 // RequestBridge pool.
+	proxyRouter http.Handler           // Proxy router for this snapshot.
+	keyPools    func() []*keypool.Pool // Current key pools for metric scrapes.
+}
+
+// New starts a gateway server. Creates a request pool only when interception
+// mode is selected. Mode is fixed at startup and requires a restart to change.
+func New(ctx context.Context, rpcDialer Dialer, logger slog.Logger, tracer trace.Tracer, experiments codersdk.Experiments, metrics *aibridge.Metrics) (*Server, error) {
 	if rpcDialer == nil {
 		return nil, xerrors.Errorf("nil rpcDialer given")
 	}
@@ -68,7 +93,16 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger,
 		lifecycleCtx: ctx,
 		cancelFn:     cancel,
 
-		requestBridgePool: pool,
+		reverseProxyExp: experiments.Enabled(codersdk.ExperimentAIGatewayReverseProxy),
+		metrics:         metrics,
+		inflight:        aibridge.NewInflightGate(logger),
+	}
+
+	if !daemon.reverseProxyExp {
+		if err := daemon.initializeInterception(); err != nil {
+			cancel(err)
+			return nil, err
+		}
 	}
 
 	daemon.wg.Add(1)
@@ -137,6 +171,12 @@ connectLoop:
 			continue
 		}
 
+		if err := s.initializeBackend(s.lifecycleCtx, client); err != nil {
+			_ = client.DRPCConn().Close()
+			s.logger.Warn(s.lifecycleCtx, "failed to initialize gateway request handler", slog.Error(err))
+			continue
+		}
+
 		// Logged at info so operators of standalone (external) gateways
 		// can see initial connection and reconnection after a dial
 		// failure (paired with the warning logged above).
@@ -162,6 +202,48 @@ connectLoop:
 	}
 }
 
+// initializeBackend selects the backend on the server's first connection to coderd.
+// A failed or malformed MCP discovery is retried on the next connection, never read as an
+// empty configuration.
+func (s *Server) initializeBackend(ctx context.Context, client DRPCClient) error {
+	s.backendMu.Lock()
+	defer s.backendMu.Unlock()
+	if s.backend.Load() != nil {
+		return nil
+	}
+
+	configCtx, cancel := context.WithTimeout(ctx, mcpConfigTimeout)
+	defer cancel()
+	configs, err := client.GetMCPServerConfigs(configCtx, &proto.GetMCPServerConfigsRequest{})
+	if err != nil {
+		return xerrors.Errorf("get MCP server configs: %w", err)
+	}
+	if configs == nil {
+		return xerrors.New("nil MCP server configs response")
+	}
+
+	// MCP configuration requires interception mode.
+	if len(configs.GetExternalAuthMcpConfigs()) != 0 || configs.GetCoderMcpConfig() != nil {
+		s.logger.Info(ctx, "mcp configuration requires interception; restart to change gateway mode")
+		return s.initializeInterception()
+	}
+
+	// Otherwise, use proxy mode.
+	s.backend.Store(&backend{})
+	s.logger.Warn(ctx, "reverse proxy routing is not yet functional")
+	return nil
+}
+
+// initializeInterception creates and publishes the server-owned request pool.
+func (s *Server) initializeInterception() error {
+	pool, err := NewCachedBridgePool(DefaultPoolOptions, nil, s.logger.Named("pool"), s.metrics, s.tracer)
+	if err != nil {
+		return xerrors.Errorf("create request pool: %w", err)
+	}
+	s.backend.Store(&backend{pool: pool, keyPools: pool.KeyPools})
+	return nil
+}
+
 // Done returns a channel that is closed when the server lifecycle ends.
 // It closes on explicit shutdown and on fatal connection-loop exit.
 func (s *Server) Done() <-chan struct{} {
@@ -192,18 +274,25 @@ func (s *Server) Client(ctx context.Context) (DRPCClient, error) {
 	}
 }
 
-// GetRequestHandler retrieves a (possibly reused) [*aibridge.RequestBridge] from the pool, for the given user.
+// GetRequestHandler retrieves the selected gateway handler for the request.
 func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handler, error) {
-	if s.requestBridgePool == nil {
-		return nil, xerrors.New("nil requestBridgePool")
+	current := s.backend.Load()
+	if current == nil {
+		return http.HandlerFunc(notReadyHandler), nil
 	}
-
-	reqBridge, err := s.requestBridgePool.Acquire(ctx, req, s.Client, NewMCPProxyFactory(s.logger, s.tracer, s.Client))
-	if err != nil {
-		return nil, xerrors.Errorf("acquire request bridge: %w", err)
+	if current.pool != nil {
+		reqBridge, err := current.pool.Acquire(ctx, req, s.Client, NewMCPProxyFactory(s.logger, s.tracer, s.Client))
+		if err != nil {
+			return nil, xerrors.Errorf("acquire request bridge: %w", err)
+		}
+		return reqBridge, nil
 	}
-
-	return reqBridge, nil
+	// Proxy mode serves every request from one handler, which does not exist
+	// until the first provider snapshot arrives.
+	if current.proxyRouter == nil {
+		return http.HandlerFunc(notReadyHandler), nil
+	}
+	return current.proxyRouter, nil
 }
 
 // Ready reports whether the server currently has an active DRPC connection to coderd.
@@ -211,6 +300,52 @@ func (s *Server) Ready() bool {
 	return s.connected.Load()
 }
 
+// ReplaceProviders publishes a provider snapshot to the selected backend.
+// Proxy requests retain their router across reloads, and a failed construction
+// keeps the previous router serving. Returns an error if mode selection has
+// not completed.
+//
+// A replaced router is not drained: in-flight requests hold it by reference
+// and finish on the snapshot they started with. Shutdown drains them through
+// the server's in-flight tracking, whichever snapshot they run on.
+func (s *Server) ReplaceProviders(ctx context.Context, providers []aibridge.Provider) error {
+	s.backendMu.Lock()
+	defer s.backendMu.Unlock()
+	if s.isShutdown() {
+		return s.Err()
+	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	current := s.backend.Load()
+	if current == nil {
+		return xerrors.New("gateway mode not selected")
+	}
+	if current.pool != nil {
+		current.pool.ReplaceProviders(providers)
+		return nil
+	}
+	router, err := aibridge.NewProxyRouter(providers, s.logger)
+	if err != nil {
+		return xerrors.Errorf("create proxy router: %w", err)
+	}
+	s.backend.Store(&backend{proxyRouter: s.inflight.Middleware(router), keyPools: router.KeyPools})
+	return nil
+}
+
+// KeyPoolStateCollector reports key states from the current provider snapshot.
+func (s *Server) KeyPoolStateCollector() prometheus.Collector {
+	return keypool.NewStateCollector(s.keyPools)
+}
+
+func (s *Server) keyPools() []*keypool.Pool {
+	current := s.backend.Load()
+	if current == nil || current.keyPools == nil {
+		return nil
+	}
+	return current.keyPools()
+}
+
 // isShutdown returns whether the Server is shutdown or not.
 func (s *Server) isShutdown() bool {
 	select {
@@ -230,6 +365,21 @@ func (s *Server) Shutdown(ctx context.Context) error {
 		// Wait for any outstanding connections to terminate.
 		s.wg.Wait()
 
+		var pool Pooler
+		if current := s.backend.Load(); current != nil {
+			if current.pool != nil {
+				pool = current.pool
+			} else {
+				// Proxy snapshots share this gate. On deadline, cancel requests
+				// without waiting for their handlers to exit.
+				drainErr := s.inflight.Shutdown(ctx)
+				s.inflight.Close()
+				if drainErr != nil {
+					s.logger.Debug(ctx, "shutdown deadline passed, canceled in-flight proxy requests", slog.Error(drainErr))
+				}
+			}
+		}
+
 		select {
 		case <-ctx.Done():
 			s.logger.Warn(ctx, "graceful shutdown failed", slog.Error(ctx.Err()))
@@ -238,9 +388,12 @@ func (s *Server) Shutdown(ctx context.Context) error {
 		default:
 		}
 
-		s.logger.Info(ctx, "shutting down request pool")
-		if err = s.requestBridgePool.Shutdown(ctx); err != nil {
-			s.logger.Error(ctx, "request pool shutdown failed with error", slog.Error(err))
+		if pool != nil {
+			s.logger.Info(ctx, "shutting down request pool")
+			if err = pool.Shutdown(ctx); err != nil {
+				s.logger.Error(ctx, "request pool shutdown failed with error", slog.Error(err))
+				return
+			}
 		}
 
 		s.logger.Info(ctx, "gracefully shutdown")
@@ -254,3 +407,7 @@ func (s *Server) Close() error {
 	defer cancel()
 	return s.Shutdown(ctx)
 }
+
+func notReadyHandler(w http.ResponseWriter, _ *http.Request) {
+	http.Error(w, "AI Gateway is starting up, retry shortly", http.StatusServiceUnavailable)
+}

coderd/aibridged/aibridged_internal_test.go

@@ -0,0 +1,167 @@
+package aibridged
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// blockingHandler returns a handler that signals started, then blocks until
+// release is closed or the request context is canceled.
+func blockingHandler() (h http.Handler, started <-chan struct{}, release func()) {
+	startedCh := make(chan struct{}, 1)
+	releaseCh := make(chan struct{})
+	h = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		startedCh <- struct{}{}
+		select {
+		case <-releaseCh:
+			w.WriteHeader(http.StatusNoContent)
+		case <-r.Context().Done():
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}
+	})
+	return h, startedCh, func() { close(releaseCh) }
+}
+
+// serveAsync serves one request through the server and reports its status code.
+func serveAsync(tracker *Server, next http.Handler) <-chan int {
+	done := make(chan int, 1)
+	go func() {
+		rec := httptest.NewRecorder()
+		handler := tracker.inflight.Middleware(next)
+		handler.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/", nil))
+		done <- rec.Code
+	}()
+	return done
+}
+
+// TestServerProxy_DrainsAdmitted asserts a graceful shutdown waits for the
+// requests already admitted and does not cancel them.
+func TestServerProxy_DrainsAdmitted(t *testing.T) {
+	t.Parallel()
+
+	tracker := newProxyTestServer(t)
+	handler, started, release := blockingHandler()
+	done := serveAsync(tracker, handler)
+	<-started
+
+	shutdownDone := make(chan error, 1)
+	go func() { shutdownDone <- tracker.Shutdown(t.Context()) }()
+
+	select {
+	case err := <-shutdownDone:
+		require.Fail(t, "shutdown returned before the in-flight request finished", "err: %v", err)
+	default:
+	}
+
+	release()
+	require.Equal(t, http.StatusNoContent, <-done)
+	require.NoError(t, <-shutdownDone)
+}
+
+// TestServerProxy_CancelsOnDeadline asserts an expired shutdown context
+// cancels admitted requests instead of waiting on them, and is reported.
+func TestServerProxy_CancelsOnDeadline(t *testing.T) {
+	t.Parallel()
+
+	tracker := newProxyTestServer(t)
+	handler, started, _ := blockingHandler()
+	done := serveAsync(tracker, handler)
+	<-started
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+	require.ErrorIs(t, tracker.Shutdown(ctx), context.Canceled)
+	require.Equal(t, http.StatusServiceUnavailable, <-done, "the request must observe the forced cancellation")
+}
+
+// TestServerProxy_DeadlineDoesNotWaitForHandler ensures shutdown returns
+// even if an admitted handler cannot observe cancellation, such as a blocked
+// response write.
+func TestServerProxy_DeadlineDoesNotWaitForHandler(t *testing.T) {
+	t.Parallel()
+
+	tracker := newProxyTestServer(t)
+	started := make(chan context.Context, 1)
+	release := make(chan struct{})
+	done := serveAsync(tracker, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		started <- r.Context()
+		<-release
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(func() {
+		close(release)
+		select {
+		case status := <-done:
+			require.Equal(t, http.StatusNoContent, status)
+		case <-time.After(testutil.WaitLong):
+			t.Error("handler did not finish after release")
+		}
+	})
+	testCtx := testutil.Context(t, testutil.WaitLong)
+	requestCtx := testutil.RequireReceive(testCtx, t, started)
+
+	ctx, cancel := context.WithDeadline(t.Context(), time.Now().Add(-time.Second))
+	defer cancel()
+	shutdownDone := make(chan error, 1)
+	go func() { shutdownDone <- tracker.Shutdown(ctx) }()
+	require.ErrorIs(t, testutil.RequireReceive(testCtx, t, shutdownDone), context.DeadlineExceeded)
+	select {
+	case <-requestCtx.Done():
+	case <-testCtx.Done():
+		t.Fatal("request context was not canceled")
+	}
+	require.ErrorIs(t, requestCtx.Err(), context.Canceled)
+
+	select {
+	case <-done:
+		t.Fatal("handler returned before release")
+	default:
+	}
+
+	rec := httptest.NewRecorder()
+	handler := tracker.inflight.Middleware(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
+		t.Error("request admitted after shutdown")
+	}))
+	handler.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/", nil))
+	require.Equal(t, http.StatusServiceUnavailable, rec.Code)
+}
+
+// TestServerProxy_RefusesAfterShutdown asserts no request is admitted once
+// shutdown has begun.
+func TestServerProxy_RefusesAfterShutdown(t *testing.T) {
+	t.Parallel()
+
+	tracker := newProxyTestServer(t)
+	require.NoError(t, tracker.Shutdown(t.Context()))
+
+	handler, _, _ := blockingHandler()
+	rec := httptest.NewRecorder()
+	current := tracker.inflight.Middleware(handler)
+	current.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/", nil))
+	require.Equal(t, http.StatusServiceUnavailable, rec.Code)
+	require.Contains(t, rec.Body.String(), "shutting down")
+}
+
+// newProxyTestServer isolates proxy request handling from the DRPC connection.
+func newProxyTestServer(t *testing.T) *Server {
+	t.Helper()
+	ctx, cancel := context.WithCancelCause(t.Context())
+	s := &Server{
+		lifecycleCtx: ctx,
+		cancelFn:     cancel,
+		logger:       slogtest.Make(t, nil),
+		inflight:     aibridge.NewInflightGate(slogtest.Make(t, nil)),
+	}
+	s.backend.Store(&backend{})
+	t.Cleanup(func() { _ = s.Close() })
+	return s
+}

coderd/aibridged/aibridged_test.go

@@ -61,8 +61,9 @@ func newTestServerWithDialer(t *testing.T, dialer aibridged.Dialer, loggerOption
 			return client, nil
 		}
 	}
-	srv, err := aibridged.New(t.Context(), pool, dialer, logger, testTracer)
+	srv, err := aibridged.New(t.Context(), dialer, logger, testTracer, nil, nil)
 	require.NoError(t, err, "create new aibridged")
+	require.NoError(t, srv.SetPoolForTest(testutil.Context(t, testutil.WaitShort), t, pool))
 	t.Cleanup(func() {
 		srv.Shutdown(context.Background())
 	})
@@ -107,6 +108,19 @@ func sdkError(status int, message string) error {
 	})
 }
 
+// shutdownAndRequirePoolClosed verifies cleanup reaches the real interception pool.
+func shutdownAndRequirePoolClosed(t *testing.T, srv *aibridged.Server) {
+	t.Helper()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	require.NotNil(t, srv.InterceptionPoolForTest())
+	require.NoError(t, srv.Shutdown(ctx))
+
+	// Use a live context so cancellation cannot mask a pool left open.
+	handler, err := srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.ErrorContains(t, err, "pool shutting down")
+	require.Nil(t, handler)
+}
+
 func TestClient_TransientDialErrorRetries(t *testing.T) {
 	t.Parallel()
 
@@ -114,8 +128,6 @@ func TestClient_TransientDialErrorRetries(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	client := mock.NewMockDRPCClient(ctrl)
 	client.EXPECT().DRPCConn().AnyTimes().Return(&mockDRPCConn{})
-	pool := mock.NewMockPooler(ctrl)
-	pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
 	dialFc := func(context.Context) (aibridged.DRPCClient, error) {
 		if calls.Add(1) == 1 {
 			return nil, sdkError(http.StatusInternalServerError, "internal error")
@@ -123,9 +135,9 @@ func TestClient_TransientDialErrorRetries(t *testing.T) {
 		return client, nil
 	}
 
-	srv, err := aibridged.New(t.Context(), pool, dialFc, slogtest.Make(t, nil), testTracer)
+	srv, err := aibridged.New(t.Context(), dialFc, slogtest.Make(t, nil), testTracer, nil, nil)
 	require.NoError(t, err)
-	t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
+	t.Cleanup(func() { shutdownAndRequirePoolClosed(t, srv) })
 
 	_, err = srv.Client(testutil.Context(t, testutil.WaitShort))
 	require.NoError(t, err)
@@ -156,17 +168,14 @@ func TestClient_FatalDialErrors(t *testing.T) {
 			t.Parallel()
 
 			var calls atomic.Int32
-			ctrl := gomock.NewController(t)
-			pool := mock.NewMockPooler(ctrl)
-			pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
 			dialFc := func(context.Context) (aibridged.DRPCClient, error) {
 				calls.Add(1)
 				return nil, sdkError(tc.status, "dial rejected")
 			}
 
-			srv, err := aibridged.New(t.Context(), pool, dialFc, slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), testTracer)
+			srv, err := aibridged.New(t.Context(), dialFc, slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), testTracer, nil, nil)
 			require.NoError(t, err)
-			t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
+			t.Cleanup(func() { shutdownAndRequirePoolClosed(t, srv) })
 
 			ctx := testutil.Context(t, testutil.WaitShort)
 			if !tc.fatal {
@@ -853,8 +862,6 @@ func TestServeHTTP_ActorHeaders(t *testing.T) {
 				}, nil),
 			}
 
-			pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
-			require.NoError(t, err)
 			conn := &mockDRPCConn{}
 			client.EXPECT().DRPCConn().AnyTimes().Return(conn)
 
@@ -869,14 +876,16 @@ func TestServeHTTP_ActorHeaders(t *testing.T) {
 			client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).AnyTimes()
 
 			// Given: aibridged is started.
-			srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
+			srv, err := aibridged.New(t.Context(), func(ctx context.Context) (aibridged.DRPCClient, error) {
 				return client, nil
-			}, logger, testTracer)
+			}, logger, testTracer, nil, nil)
 			require.NoError(t, err, "create new aibridged")
 			t.Cleanup(func() {
 				_ = srv.Shutdown(testutil.Context(t, testutil.WaitShort))
 			})
 
+			require.NoError(t, srv.ReplaceProviders(t.Context(), providers))
+
 			// When: a request is made to aibridged.
 			ctx := testutil.Context(t, testutil.WaitShort)
 			req, err := http.NewRequestWithContext(ctx, http.MethodPost, tc.path, bytes.NewBufferString(`{}`))
@@ -956,8 +965,6 @@ func TestRouting(t *testing.T) {
 				aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: openaiSrv.URL, KeyPool: singleKeyPool(t, "openai", "test-key")}),
 				aibridgetest.NewAnthropicProvider(t, aibridge.AnthropicConfig{BaseURL: antSrv.URL, KeyPool: singleKeyPool(t, "anthropic", "test-key")}, nil),
 			}
-			pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
-			require.NoError(t, err)
 			conn := &mockDRPCConn{}
 			client.EXPECT().DRPCConn().AnyTimes().Return(conn)
 
@@ -974,14 +981,16 @@ func TestRouting(t *testing.T) {
 			client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).Times(tc.expectedHits)
 
 			// Given: aibridged is started.
-			srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
+			srv, err := aibridged.New(t.Context(), func(ctx context.Context) (aibridged.DRPCClient, error) {
 				return client, nil
-			}, logger, testTracer)
+			}, logger, testTracer, nil, nil)
 			require.NoError(t, err, "create new aibridged")
 			t.Cleanup(func() {
 				_ = srv.Shutdown(testutil.Context(t, testutil.WaitShort))
 			})
 
+			require.NoError(t, srv.ReplaceProviders(t.Context(), providers))
+
 			// When: a request is made to aibridged.
 			ctx := testutil.Context(t, testutil.WaitShort)
 			req, err := http.NewRequestWithContext(ctx, http.MethodPost, tc.path, bytes.NewBufferString(`{}`))
@@ -1076,9 +1085,6 @@ func TestReady(t *testing.T) {
 		t.Parallel()
 
 		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
-		ctrl := gomock.NewController(t)
-		pool := mock.NewMockPooler(ctrl)
-		pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
 
 		dialerCalled := make(chan struct{}, 1)
 		blockDialer := func(ctx context.Context) (aibridged.DRPCClient, error) {
@@ -1090,9 +1096,9 @@ func TestReady(t *testing.T) {
 			return nil, ctx.Err()
 		}
 
-		srv, err := aibridged.New(t.Context(), pool, blockDialer, logger, testTracer)
+		srv, err := aibridged.New(t.Context(), blockDialer, logger, testTracer, nil, nil)
 		require.NoError(t, err)
-		t.Cleanup(func() { srv.Close() })
+		t.Cleanup(func() { shutdownAndRequirePoolClosed(t, srv) })
 
 		testutil.RequireReceive(t.Context(), t, dialerCalled)
 		require.False(t, srv.Ready(), "expected not ready before first connection")
@@ -1113,8 +1119,6 @@ func TestReady(t *testing.T) {
 
 		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
 		ctrl := gomock.NewController(t)
-		pool := mock.NewMockPooler(ctrl)
-		pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
 
 		// allowDial gates the dialer. When open, dials succeed
 		// immediately. Replace with a fresh channel to block dials.
@@ -1143,9 +1147,9 @@ func TestReady(t *testing.T) {
 		// Start with dialer unblocked.
 		close(allowDial)
 
-		srv, err := aibridged.New(t.Context(), pool, dialer, logger, testTracer)
+		srv, err := aibridged.New(t.Context(), dialer, logger, testTracer, nil, nil)
 		require.NoError(t, err)
-		t.Cleanup(func() { srv.Close() })
+		t.Cleanup(func() { shutdownAndRequirePoolClosed(t, srv) })
 		srvNotReady := func() bool { return !srv.Ready() }
 
 		// Wait for the initial connection.

coderd/aibridged/aibridgedmock/poolmock.go

@@ -15,6 +15,7 @@ import (
 	reflect "reflect"
 
 	aibridge "github.com/coder/coder/v2/aibridge"
+	keypool "github.com/coder/coder/v2/aibridge/keypool"
 	aibridged "github.com/coder/coder/v2/coderd/aibridged"
 	gomock "go.uber.org/mock/gomock"
 )
@@ -58,6 +59,20 @@ func (mr *MockPoolerMockRecorder) Acquire(ctx, req, clientFn, mcpBootstrapper an
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Acquire", reflect.TypeOf((*MockPooler)(nil).Acquire), ctx, req, clientFn, mcpBootstrapper)
 }
 
+// KeyPools mocks base method.
+func (m *MockPooler) KeyPools() []*keypool.Pool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "KeyPools")
+	ret0, _ := ret[0].([]*keypool.Pool)
+	return ret0
+}
+
+// KeyPools indicates an expected call of KeyPools.
+func (mr *MockPoolerMockRecorder) KeyPools() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "KeyPools", reflect.TypeOf((*MockPooler)(nil).KeyPools))
+}
+
 // ReplaceProviders mocks base method.
 func (m *MockPooler) ReplaceProviders(providers []aibridge.Provider) {
 	m.ctrl.T.Helper()

coderd/aibridged/backend_internal_test.go

@@ -0,0 +1,162 @@
+package aibridged
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/aibridge/config"
+	"github.com/coder/coder/v2/aibridge/keypool"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/quartz"
+)
+
+func TestServerKeyPoolStateCollector(t *testing.T) {
+	t.Parallel()
+
+	newProvider := func(t *testing.T, name string) aibridge.Provider {
+		t.Helper()
+
+		pool, err := keypool.New(name, []string{"key"}, quartz.NewReal(), nil)
+		require.NoError(t, err)
+		return aibridge.NewOpenAIProvider(config.OpenAI{Name: name, BaseURL: "http://upstream.test", KeyPool: pool})
+	}
+	gather := func(t *testing.T, registry *prometheus.Registry) []*dto.MetricFamily {
+		t.Helper()
+
+		metrics, err := registry.Gather()
+		require.NoError(t, err)
+		return metrics
+	}
+
+	var server Server
+	registry := prometheus.NewRegistry()
+	require.NoError(t, registry.Register(server.KeyPoolStateCollector()))
+
+	// One registered key-state collector must safely handle startup before mode
+	// selection and proxy mode before the first provider snapshot arrives.
+	require.Empty(t, gather(t, registry))
+	server.backend.Store(&backend{})
+	require.Empty(t, gather(t, registry))
+
+	// Interception reloads update metrics from the same registered collector.
+	firstProvider := newProvider(t, "first")
+	pool, err := NewCachedBridgePool(DefaultPoolOptions, []aibridge.Provider{firstProvider}, slogtest.Make(t, nil), nil, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = pool.Shutdown(context.Background()) })
+	server.backend.Store(&backend{pool: pool, keyPools: pool.KeyPools})
+
+	metrics := gather(t, registry)
+	require.True(t, testutil.PromGaugeHasValue(t, metrics, 1, "key_pool_state", "first", "valid"))
+
+	secondProvider := newProvider(t, "second")
+	pool.ReplaceProviders([]aibridge.Provider{secondProvider})
+
+	metrics = gather(t, registry)
+	require.False(t, testutil.PromGaugeGathered(t, metrics, "key_pool_state", "first", "valid"))
+	require.True(t, testutil.PromGaugeHasValue(t, metrics, 1, "key_pool_state", "second", "valid"))
+
+	// Replacing proxy snapshots similarly changes the output without
+	// registering another collector.
+	thirdProvider := newProvider(t, "third")
+	firstRouter, err := aibridge.NewProxyRouter([]aibridge.Provider{thirdProvider}, slogtest.Make(t, nil))
+	require.NoError(t, err)
+	server.backend.Store(&backend{proxyRouter: firstRouter, keyPools: firstRouter.KeyPools})
+
+	metrics = gather(t, registry)
+	require.False(t, testutil.PromGaugeGathered(t, metrics, "key_pool_state", "second", "valid"))
+	require.True(t, testutil.PromGaugeHasValue(t, metrics, 1, "key_pool_state", "third", "valid"))
+
+	fourthProvider := newProvider(t, "fourth")
+	secondRouter, err := aibridge.NewProxyRouter([]aibridge.Provider{fourthProvider}, slogtest.Make(t, nil))
+	require.NoError(t, err)
+	server.backend.Store(&backend{proxyRouter: secondRouter, keyPools: secondRouter.KeyPools})
+
+	metrics = gather(t, registry)
+	require.False(t, testutil.PromGaugeGathered(t, metrics, "key_pool_state", "third", "valid"))
+	require.True(t, testutil.PromGaugeHasValue(t, metrics, 1, "key_pool_state", "fourth", "valid"))
+}
+
+// InterceptionPoolForTest returns the active interception pool, if selected.
+func (s *Server) InterceptionPoolForTest() Pooler {
+	current := s.backend.Load()
+	if current == nil {
+		return nil
+	}
+	return current.pool
+}
+
+// shutdownPool isolates pool shutdown from request handling.
+type shutdownPool struct {
+	Pooler
+	shutdown func(context.Context) error
+}
+
+func (p *shutdownPool) Shutdown(ctx context.Context) error { return p.shutdown(ctx) }
+
+func TestServerShutdownMode(t *testing.T) {
+	t.Parallel()
+	poolErr := xerrors.New("pool shutdown failed")
+	for _, tc := range []struct {
+		name         string
+		selected     bool
+		interception bool
+		poolErr      error
+	}{
+		{name: "Proxy", selected: true},
+		{name: "Interception", selected: true, interception: true},
+		{name: "InterceptionFailure", selected: true, interception: true, poolErr: poolErr},
+		{name: "BeforeSelection"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			lifecycleCtx, cancel := context.WithCancelCause(t.Context())
+			defer cancel(nil)
+			calls := 0
+			pool := &shutdownPool{shutdown: func(got context.Context) error {
+				require.Equal(t, ctx, got)
+				calls++
+				return tc.poolErr
+			}}
+			server := &Server{
+				lifecycleCtx: lifecycleCtx, cancelFn: cancel,
+				logger:   slogtest.Make(t, &slogtest.Options{IgnoreErrors: tc.poolErr != nil}),
+				inflight: aibridge.NewInflightGate(slogtest.Make(t, nil)),
+			}
+			if tc.selected {
+				current := &backend{}
+				if tc.interception {
+					current.pool = pool
+				}
+				server.backend.Store(current)
+			}
+			proxy := tc.selected && !tc.interception
+			if !proxy {
+				// An unused proxy gate must not delay interception or startup cleanup.
+				release, ok := server.inflight.Admit()
+				require.True(t, ok)
+				defer release()
+			}
+			require.ErrorIs(t, server.Shutdown(ctx), tc.poolErr)
+			expectedCalls := 0
+			if tc.interception {
+				expectedCalls = 1
+			}
+			require.Equal(t, expectedCalls, calls)
+			release, ok := server.inflight.Admit()
+			require.Equal(t, !proxy, ok, "only proxy mode shuts down the proxy gate")
+			if ok {
+				release()
+			}
+			_ = server.Shutdown(ctx)
+			require.Equal(t, expectedCalls, calls, "owned pool is shut down only once")
+		})
+	}
+}

coderd/aibridged/backend_test.go

@@ -0,0 +1,543 @@
+package aibridged_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+	"golang.org/x/xerrors"
+	"storj.io/drpc"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/aibridge/config"
+	"github.com/coder/coder/v2/aibridge/keypool"
+	"github.com/coder/coder/v2/coderd/aibridged"
+	mock "github.com/coder/coder/v2/coderd/aibridged/aibridgedmock"
+	"github.com/coder/coder/v2/coderd/aibridged/proto"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// countingDRPCConn tracks connections discarded after failed mode selection.
+type countingDRPCConn struct {
+	*mockDRPCConn
+	closes atomic.Int32
+}
+
+func newCountingDRPCConn(closedCh chan struct{}) *countingDRPCConn {
+	return &countingDRPCConn{mockDRPCConn: &mockDRPCConn{closedCh: closedCh}}
+}
+
+func (c *countingDRPCConn) Close() error {
+	c.closes.Add(1)
+	return nil
+}
+
+var _ drpc.Conn = &countingDRPCConn{}
+
+// mcpConfigsFn answers a GetMCPServerConfigs call. call is 1-based.
+type mcpConfigsFn func(call int32) (*proto.GetMCPServerConfigsResponse, error)
+
+// expectMCPConfigs counts discovery calls and supplies their responses.
+func expectMCPConfigs(client *mock.MockDRPCClient, fn mcpConfigsFn) *atomic.Int32 {
+	var calls atomic.Int32
+	client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
+		func(context.Context, *proto.GetMCPServerConfigsRequest) (*proto.GetMCPServerConfigsResponse, error) {
+			return fn(calls.Add(1))
+		})
+	return &calls
+}
+
+// noMCPConfigs is the response which selects the proxy backend.
+func noMCPConfigs() *proto.GetMCPServerConfigsResponse {
+	return &proto.GetMCPServerConfigsResponse{}
+}
+
+// staticConfigs answers every MCP discovery call with the same response.
+func staticConfigs(response *proto.GetMCPServerConfigsResponse) mcpConfigsFn {
+	return func(int32) (*proto.GetMCPServerConfigsResponse, error) {
+		return response, nil
+	}
+}
+
+func coderMCPConfig() *proto.GetMCPServerConfigsResponse {
+	return &proto.GetMCPServerConfigsResponse{
+		CoderMcpConfig: &proto.MCPServerConfig{Id: "coder", Url: "http://coder.test/api/experimental/mcp/http"},
+	}
+}
+
+func externalMCPConfig() *proto.GetMCPServerConfigsResponse {
+	return &proto.GetMCPServerConfigsResponse{
+		ExternalAuthMcpConfigs: []*proto.MCPServerConfig{{Id: "github", Url: "http://github.test/mcp"}},
+	}
+}
+
+// backendFixture holds the collaborators of a server under test.
+type backendFixture struct {
+	srv      *aibridged.Server
+	conn     *countingDRPCConn
+	mcpCalls *atomic.Int32
+}
+
+// newBackendServer starts a server whose single mock client answers MCP
+// discovery with fn. Interception pools are created by the server under test.
+func newBackendServer(t *testing.T, experiments codersdk.Experiments, fn mcpConfigsFn, ignoreErrors bool) *backendFixture {
+	t.Helper()
+
+	ctrl := gomock.NewController(t)
+	client := mock.NewMockDRPCClient(ctrl)
+	conn := newCountingDRPCConn(nil)
+	client.EXPECT().DRPCConn().AnyTimes().Return(conn)
+	calls := expectMCPConfigs(client, fn)
+
+	srv, err := aibridged.New(t.Context(),
+		func(context.Context) (aibridged.DRPCClient, error) { return client, nil },
+		slogtest.Make(t, &slogtest.Options{IgnoreErrors: ignoreErrors}), testTracer,
+		experiments, nil)
+	require.NoError(t, err, "create aibridged server")
+	t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
+
+	return &backendFixture{srv: srv, conn: conn, mcpCalls: calls}
+}
+
+func proxyExperiments() codersdk.Experiments {
+	return codersdk.Experiments{codersdk.ExperimentAIGatewayReverseProxy}
+}
+
+func waitReady(t *testing.T, srv *aibridged.Server) {
+	t.Helper()
+	require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast,
+		"server never reported an active connection")
+}
+
+func serveHandler(t *testing.T, h http.Handler, path string) *httptest.ResponseRecorder {
+	t.Helper()
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, path, nil))
+	return rec
+}
+
+func collectServerMetrics(t *testing.T, server *aibridged.Server) []*dto.MetricFamily {
+	t.Helper()
+
+	registry := prometheus.NewRegistry()
+	require.NoError(t, registry.Register(server.KeyPoolStateCollector()))
+	metrics, err := registry.Gather()
+	require.NoError(t, err)
+	return metrics
+}
+
+func requireKeyPoolState(t *testing.T, server *aibridged.Server, value float64, provider, state string) {
+	t.Helper()
+	require.True(t, testutil.PromGaugeHasValue(t, collectServerMetrics(t, server), value, "key_pool_state", provider, state))
+}
+
+func requireNoKeyPoolState(t *testing.T, server *aibridged.Server, provider, state string) {
+	t.Helper()
+	require.False(t, testutil.PromGaugeGathered(t, collectServerMetrics(t, server), "key_pool_state", provider, state))
+}
+
+func openAIProvider(name string, pool *keypool.Pool) aibridge.Provider {
+	return aibridge.NewOpenAIProvider(config.OpenAI{Name: name, BaseURL: "http://upstream.test", KeyPool: pool})
+}
+
+func TestBackendMode_Interception(t *testing.T) {
+	t.Parallel()
+	for _, tc := range []struct {
+		name        string
+		experiments codersdk.Experiments
+		response    *proto.GetMCPServerConfigsResponse
+		wantCalls   int32
+	}{
+		{name: "ExperimentDisabled", response: noMCPConfigs()},
+		{name: "CoderMCP", experiments: proxyExperiments(), response: coderMCPConfig(), wantCalls: 1},
+		{name: "ExternalAuthMCP", experiments: proxyExperiments(), response: externalMCPConfig(), wantCalls: 1},
+		{name: "Both", experiments: proxyExperiments(), wantCalls: 1, response: &proto.GetMCPServerConfigsResponse{
+			CoderMcpConfig:         coderMCPConfig().GetCoderMcpConfig(),
+			ExternalAuthMcpConfigs: externalMCPConfig().GetExternalAuthMcpConfigs(),
+		}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			f := newBackendServer(t, tc.experiments, func(int32) (*proto.GetMCPServerConfigsResponse, error) {
+				if tc.wantCalls == 0 {
+					t.Error("MCP discovery must not run when the experiment is disabled")
+				}
+				return tc.response, nil
+			}, false)
+			ctx := testutil.Context(t, testutil.WaitShort)
+			waitReady(t, f.srv)
+			require.NotNil(t, f.srv.InterceptionPoolForTest())
+			require.Equal(t, tc.wantCalls, f.mcpCalls.Load())
+
+			require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+				openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+			}))
+			requireKeyPoolState(t, f.srv, 1, "openai", "valid")
+		})
+	}
+}
+
+// Proxy mode returns 503 until providers load, then 404 for unregistered routes.
+func TestBackendMode_ProxyWhenNoMCPConfigs(t *testing.T) {
+	t.Parallel()
+
+	f := newBackendServer(t, proxyExperiments(), staticConfigs(noMCPConfigs()), false)
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	waitReady(t, f.srv)
+	require.Equal(t, int32(1), f.mcpCalls.Load())
+	require.Nil(t, f.srv.InterceptionPoolForTest())
+
+	h, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.NoError(t, err, "the proxy backend must answer requests instead of erroring")
+	rec := serveHandler(t, h, "/openai/v1/chat/completions")
+	require.Equal(t, http.StatusServiceUnavailable, rec.Code, "requests before the first snapshot must be refused")
+	require.Equal(t, "AI Gateway is starting up, retry shortly\n", rec.Body.String())
+	requireNoKeyPoolState(t, f.srv, "openai", "valid")
+
+	pool := singleKeyPool(t, "openai", "key")
+	require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{openAIProvider("openai", pool)}))
+
+	h, err = f.srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.NoError(t, err)
+	rec = serveHandler(t, h, "/openai/v1/chat/completions")
+	require.Equal(t, http.StatusNotFound, rec.Code, "the router has no provider routes registered yet")
+	requireKeyPoolState(t, f.srv, 1, "openai", "valid")
+}
+
+// Reconnects and reloads must not repeat mode selection.
+func TestBackendMode_SelectedOnceAcrossReconnects(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name     string
+		first    *proto.GetMCPServerConfigsResponse
+		second   *proto.GetMCPServerConfigsResponse
+		wantPool bool
+	}{
+		{name: "InterceptionStaysInterception", first: coderMCPConfig(), second: noMCPConfigs(), wantPool: true},
+		{name: "ProxyStaysProxy", first: noMCPConfigs(), second: coderMCPConfig(), wantPool: false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctrl := gomock.NewController(t)
+			client := mock.NewMockDRPCClient(ctrl)
+			closedCh := make(chan struct{})
+			// Redial switches to a fresh connection after the first is closed.
+			firstConn := newCountingDRPCConn(closedCh)
+			secondConn := newCountingDRPCConn(nil)
+			var activeConn atomic.Pointer[countingDRPCConn]
+			activeConn.Store(firstConn)
+			client.EXPECT().DRPCConn().AnyTimes().DoAndReturn(func() drpc.Conn { return activeConn.Load() })
+
+			calls := expectMCPConfigs(client, func(call int32) (*proto.GetMCPServerConfigsResponse, error) {
+				if call == 1 {
+					return tc.first, nil
+				}
+				return tc.second, nil
+			})
+
+			var dials atomic.Int32
+			srv, err := aibridged.New(t.Context(), func(context.Context) (aibridged.DRPCClient, error) {
+				if dials.Add(1) >= 2 {
+					activeConn.Store(secondConn)
+				}
+				return client, nil
+			}, slogtest.Make(t, nil), testTracer, proxyExperiments(), nil)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+			waitReady(t, srv)
+			require.Equal(t, int32(1), calls.Load())
+
+			require.NoError(t, srv.ReplaceProviders(ctx, []aibridge.Provider{
+				openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+			}))
+
+			// Force a reconnect and wait for the loop to dial again.
+			close(closedCh)
+			require.Eventually(t, func() bool { return dials.Load() > 1 }, testutil.WaitShort, testutil.IntervalFast,
+				"connection loop did not redial")
+
+			// Another reload after the reconnect must not change the mode.
+			require.NoError(t, srv.ReplaceProviders(ctx, []aibridge.Provider{
+				openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+			}))
+
+			require.Equal(t, int32(1), calls.Load(), "MCP discovery must run once, at startup")
+
+			if tc.wantPool {
+				require.NotNil(t, srv.InterceptionPoolForTest())
+				requireKeyPoolState(t, srv, 1, "openai", "valid")
+				return
+			}
+
+			h, err := srv.GetRequestHandler(ctx, aibridged.Request{})
+			require.NoError(t, err)
+			require.Nil(t, srv.InterceptionPoolForTest())
+			requireKeyPoolState(t, srv, 1, "openai", "valid")
+			require.Equal(t, http.StatusNotFound, serveHandler(t, h, "/openai/v1/chat/completions").Code)
+		})
+	}
+}
+
+// Failed discovery must discard the connection, not select proxy mode.
+func TestBackendMode_MCPDiscoveryFailureRetries(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name string
+		err  error
+	}{
+		{name: "RPCError", err: xerrors.New("boom")},
+		{name: "NilResponse"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			f := newBackendServer(t, proxyExperiments(), func(call int32) (*proto.GetMCPServerConfigsResponse, error) {
+				if call <= 2 {
+					return nil, tc.err // A nil error exercises a malformed response.
+				}
+				return noMCPConfigs(), nil
+			}, true)
+			ctx := testutil.Context(t, testutil.WaitShort)
+
+			// While selection keeps failing the server is not ready and
+			// requests are refused rather than served by a guessed backend.
+			require.Eventually(t, func() bool { return f.mcpCalls.Load() >= 1 }, testutil.WaitShort, testutil.IntervalFast)
+			require.False(t, f.srv.Ready())
+			h, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+			require.NoError(t, err)
+			require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, h, "/openai/v1/chat/completions").Code)
+			requireNoKeyPoolState(t, f.srv, "openai", "valid")
+			require.Nil(t, f.srv.InterceptionPoolForTest())
+
+			waitReady(t, f.srv)
+			require.Equal(t, int32(3), f.mcpCalls.Load(), "selection must be retried on each new connection")
+			require.GreaterOrEqual(t, f.conn.closes.Load(), int32(2), "a connection whose selection failed must be discarded")
+
+			require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+				openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+			}))
+			h, err = f.srv.GetRequestHandler(ctx, aibridged.Request{})
+			require.NoError(t, err)
+			require.Equal(t, http.StatusNotFound, serveHandler(t, h, "/openai/v1/chat/completions").Code)
+		})
+	}
+}
+
+// Reloads publish a new router while acquired handlers retain their snapshot.
+func TestReplaceProviders_SwapsRouter(t *testing.T) {
+	t.Parallel()
+
+	f := newBackendServer(t, proxyExperiments(), staticConfigs(noMCPConfigs()), false)
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	waitReady(t, f.srv)
+
+	firstPool := singleKeyPool(t, "openai", "key")
+	require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{openAIProvider("openai", firstPool)}))
+	requireKeyPoolState(t, f.srv, 1, "openai", "valid")
+
+	oldHandler, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.NoError(t, err)
+
+	secondPool := singleKeyPool(t, "anthropic", "key")
+	require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+		openAIProvider("openai", firstPool),
+		openAIProvider("anthropic", secondPool),
+		aibridge.NewDisabledProviderStub("disabled", "openai"),
+	}))
+	requireKeyPoolState(t, f.srv, 1, "openai", "valid")
+	requireKeyPoolState(t, f.srv, 1, "anthropic", "valid")
+
+	newHandler, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.NoError(t, err)
+
+	// Handlers already acquired keep serving from the router they hold.
+	require.Equal(t, http.StatusNotFound, serveHandler(t, oldHandler, "/disabled/v1/models").Code)
+	require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, newHandler, "/disabled/v1/models").Code)
+}
+
+// Invalid snapshots leave the previous router serving.
+func TestReplaceProviders_FailureRetainsRouter(t *testing.T) {
+	t.Parallel()
+
+	f := newBackendServer(t, proxyExperiments(), staticConfigs(noMCPConfigs()), false)
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	waitReady(t, f.srv)
+
+	good := singleKeyPool(t, "openai", "key")
+	require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+		openAIProvider("openai", good),
+		aibridge.NewDisabledProviderStub("disabled", "openai"),
+	}))
+	handler, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.NoError(t, err)
+
+	err = f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+		openAIProvider("dupe", singleKeyPool(t, "a", "key")),
+		openAIProvider("dupe", singleKeyPool(t, "b", "key")),
+	})
+	require.ErrorContains(t, err, "duplicate provider name")
+
+	requireKeyPoolState(t, f.srv, 1, "openai", "valid")
+	retained, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+	require.NoError(t, err)
+	require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, handler, "/disabled/v1/models").Code)
+	require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, retained, "/disabled/v1/models").Code)
+	require.Equal(t, http.StatusNotFound, serveHandler(t, retained, "/openai/v1/chat/completions").Code)
+}
+
+// Cancellation and shutdown prevent publication; shutdown also stops serving.
+func TestReplaceProviders_ContextAndShutdown(t *testing.T) {
+	t.Parallel()
+
+	t.Run("CanceledContext", func(t *testing.T) {
+		t.Parallel()
+
+		f := newBackendServer(t, proxyExperiments(), staticConfigs(noMCPConfigs()), false)
+		waitReady(t, f.srv)
+
+		ctx, cancel := context.WithCancel(t.Context())
+		cancel()
+		require.ErrorIs(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+			openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+		}), context.Canceled)
+		requireNoKeyPoolState(t, f.srv, "openai", "valid")
+	})
+
+	t.Run("AfterShutdown", func(t *testing.T) {
+		t.Parallel()
+
+		f := newBackendServer(t, proxyExperiments(), staticConfigs(noMCPConfigs()), false)
+		ctx := testutil.Context(t, testutil.WaitShort)
+		waitReady(t, f.srv)
+		require.NoError(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+			openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+		}))
+		retired, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+		require.NoError(t, err)
+
+		require.NoError(t, f.srv.Shutdown(context.Background()))
+		require.ErrorIs(t, f.srv.ReplaceProviders(ctx, []aibridge.Provider{
+			openAIProvider("openai", singleKeyPool(t, "openai", "key")),
+		}), aibridged.ErrShutdown)
+
+		// Shutdown refuses proxy requests, whether the handler was acquired
+		// before or after it.
+		h, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+		require.NoError(t, err)
+		require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, h, "/openai/v1/chat/completions").Code)
+		require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, retired, "/openai/v1/chat/completions").Code)
+	})
+
+	t.Run("CanceledShutdown", func(t *testing.T) {
+		t.Parallel()
+
+		f := newBackendServer(t, proxyExperiments(), staticConfigs(noMCPConfigs()), false)
+		waitReady(t, f.srv)
+		require.NoError(t, f.srv.ReplaceProviders(t.Context(), nil))
+		acquired, err := f.srv.GetRequestHandler(t.Context(), aibridged.Request{})
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(t.Context())
+		cancel()
+		require.ErrorIs(t, f.srv.Shutdown(ctx), context.Canceled)
+
+		// A shutdown that cannot finish gracefully still refuses proxy
+		// requests instead of serving from a snapshot that is being released.
+		handler, err := f.srv.GetRequestHandler(t.Context(), aibridged.Request{})
+		require.NoError(t, err)
+		require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, handler, "/openai/v1/models").Code)
+		require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, acquired, "/openai/v1/models").Code)
+	})
+}
+
+// Exercise serving, reloads, metrics, and shutdown together under -race.
+func TestBackend_ConcurrentUse(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name     string
+		response *proto.GetMCPServerConfigsResponse
+		proxy    bool
+	}{
+		{name: "Proxy", response: noMCPConfigs(), proxy: true},
+		{name: "Interception", response: coderMCPConfig()},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			f := newBackendServer(t, proxyExperiments(), staticConfigs(tc.response), true)
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+			waitReady(t, f.srv)
+
+			const workers = 8
+			registry := prometheus.NewRegistry()
+			require.NoError(t, registry.Register(f.srv.KeyPoolStateCollector()))
+			gatherErrors := make(chan error, workers)
+			var wg sync.WaitGroup
+			start := make(chan struct{})
+
+			providers := []aibridge.Provider{openAIProvider("openai", singleKeyPool(t, "openai", "key"))}
+			for range workers {
+				wg.Go(func() {
+					<-start
+					for range 20 {
+						h, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+						if err != nil {
+							continue
+						}
+						rec := httptest.NewRecorder()
+						h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil))
+					}
+				})
+
+				wg.Go(func() {
+					<-start
+					for range 20 {
+						// Errors are expected once shutdown wins the race.
+						_ = f.srv.ReplaceProviders(ctx, providers)
+						if _, err := registry.Gather(); err != nil {
+							gatherErrors <- err
+							return
+						}
+					}
+				})
+			}
+
+			wg.Go(func() {
+				<-start
+				_ = f.srv.Shutdown(context.Background())
+			})
+
+			close(start)
+			wg.Wait()
+			close(gatherErrors)
+			for err := range gatherErrors {
+				require.NoError(t, err)
+			}
+
+			if tc.proxy {
+				h, err := f.srv.GetRequestHandler(ctx, aibridged.Request{})
+				require.NoError(t, err)
+				require.Equal(t, http.StatusServiceUnavailable, serveHandler(t, h, "/openai/v1/chat/completions").Code,
+					"shutdown must refuse proxy requests")
+			}
+		})
+	}
+}

coderd/aibridged/metrics.go

@@ -21,7 +21,7 @@ type Metrics struct {
 	ProvidersLastReloadTimestampSeconds prometheus.Gauge
 
 	// ProvidersLastReloadSuccessTimestampSeconds is the unix timestamp
-	// of the last reload that successfully refreshed the pool. A gap
+	// of the last reload that successfully refreshed the providers. A gap
 	// against ProvidersLastReloadTimestampSeconds means the loop is
 	// firing but the refresh function is failing.
 	ProvidersLastReloadSuccessTimestampSeconds prometheus.Gauge
@@ -46,7 +46,7 @@ func NewMetrics(reg prometheus.Registerer) *Metrics {
 
 		ProvidersLastReloadSuccessTimestampSeconds: factory.NewGauge(prometheus.GaugeOpts{
 			Name: "providers_last_reload_success_timestamp_seconds",
-			Help: "Unix timestamp of the last provider reload that successfully refreshed the pool. A gap against the providers_last_reload_timestamp_seconds gauge means the loop is firing but the refresh function is failing.",
+			Help: "Unix timestamp of the last provider reload that successfully refreshed the providers. A gap against the providers_last_reload_timestamp_seconds gauge means the loop is firing but the refresh function is failing.",
 		}),
 	}
 }

coderd/aibridged/pool.go

@@ -31,6 +31,7 @@ const (
 // Pooler describes a pool of [*aibridge.RequestBridge] instances from which instances can be retrieved.
 // One [*aibridge.RequestBridge] instance is created per given key.
 type Pooler interface {
+	KeyPools() []*keypool.Pool
 	Acquire(ctx context.Context, req Request, clientFn ClientFunc, mcpBootstrapper MCPProxyBuilder) (http.Handler, error)
 	// ReplaceProviders swaps the providers used to construct future
 	// RequestBridge instances and clears the cache. Disabled providers

coderd/aibridged/testhooks.go

@@ -0,0 +1,29 @@
+package aibridged
+
+import (
+	"context"
+	"testing"
+
+	"golang.org/x/xerrors"
+)
+
+// SetPoolForTest closes the unused interception pool and installs a replacement.
+// Call it after interception is selected, before requests, reloads, or shutdown.
+// The server owns the replacement pool and shuts it down with the server.
+func (s *Server) SetPoolForTest(ctx context.Context, t testing.TB, pool Pooler) error {
+	t.Helper()
+	s.backendMu.Lock()
+	defer s.backendMu.Unlock()
+	if pool == nil {
+		return xerrors.New("nil request pool")
+	}
+	current := s.backend.Load()
+	if current == nil || current.pool == nil {
+		return xerrors.New("interception mode not selected")
+	}
+	if err := current.pool.Shutdown(ctx); err != nil {
+		return xerrors.Errorf("shutdown unused request pool: %w", err)
+	}
+	s.backend.Store(&backend{pool: pool, keyPools: pool.KeyPools})
+	return nil
+}

coderd/aibridgedtest/aibridgedtest.go

@@ -24,9 +24,8 @@ import (
 // (e.g. chattest.NewOpenAI) will have their requests proxied through the real
 // aibridged stack as they would in production.
 //
-// The daemon starts with an empty pool and fetches providers from coderd over
-// the in-memory DRPC, then refreshes on ai_providers change events, exactly
-// like cli.newAIBridgeDaemon.
+// The daemon fetches providers from coderd over the in-memory DRPC, then
+// refreshes on ai_providers change events, exactly like cli.newAIBridgeDaemon.
 //
 // metrics is the registry the daemon reports provider reload events to.
 // The caller owns the metrics instance and can assert on it after the daemon
@@ -62,15 +61,9 @@ func StartTestAIBridgeDaemonWithPubsub(
 		metrics = aibridged.NewMetrics(prometheus.NewRegistry())
 	}
 
-	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger.Named("pool"), nil, tracer)
-	if err != nil {
-		t.Fatalf("create bridge pool: %v", err)
-	}
-	t.Cleanup(func() { _ = pool.Shutdown(context.Background()) })
-
-	srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
+	srv, err := aibridged.New(ctx, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
 		return api.CreateInMemoryAIBridgeServer(dialCtx)
-	}, logger, tracer)
+	}, logger, tracer, api.Experiments, nil)
 	if err != nil {
 		t.Fatalf("create aibridged server: %v", err)
 	}
@@ -78,7 +71,7 @@ func StartTestAIBridgeDaemonWithPubsub(
 
 	// The reloader fetches providers from coderd over srv's DRPC client; the
 	// subscription drives an initial load and refreshes on change events.
-	reloader := cli.NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("reloader"), nil, metrics)
+	reloader := cli.NewProviderRPCReloader(srv.ReplaceProviders, srv.Client, cfg, logger.Named("reloader"), nil, metrics)
 	unsubscribe, err := aibridged.SubscribeProviderReload(ctx, ps, reloader, logger.Named("subscriber"))
 	if err != nil {
 		t.Fatalf("subscribe provider reload: %v", err)

docs/reference/cli/ai-gateway_start.md

@@ -161,6 +161,16 @@ Output JSON logs to a given file.
 
 Output Stackdriver compatible logs to a given file.
 
+### --experiments
+
+|             |                                 |
+|-------------|---------------------------------|
+| Type        | <code>string-array</code>       |
+| Environment | <code>$CODER_EXPERIMENTS</code> |
+| YAML        | <code>experiments</code>        |
+
+Enable one or more experiments. These are not ready for production. Separate multiple experiments with commas, or enter '*' to opt-in to all available experiments.
+
 ### -c, --config
 
 |             |                                 |

enterprise/aibridged_integration_test.go

@@ -185,14 +185,13 @@ func TestIntegration(t *testing.T) {
 
 	logger := testutil.Logger(t)
 	providers := []aibridge.Provider{aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: mockOpenAI.URL, KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-centralized-key")})}
-	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, tracer)
-	require.NoError(t, err)
 
-	// Given: aibridged is started.
-	srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
+	// Given: aibridged is started with its server-owned interception pool.
+	srv, err := aibridged.New(t.Context(), func(ctx context.Context) (aibridged.DRPCClient, error) {
 		return aiBridgeClient, nil
-	}, logger, tracer)
+	}, logger, tracer, nil, nil)
 	require.NoError(t, err, "create new aibridged")
+	require.NoError(t, srv.ReplaceProviders(ctx, providers))
 	t.Cleanup(func() {
 		_ = srv.Shutdown(ctx)
 	})
@@ -388,15 +387,12 @@ func TestIntegrationWithMetrics(t *testing.T) {
 	logger := testutil.Logger(t)
 	providers := []aibridge.Provider{aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: mockOpenAI.URL, KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-centralized-key")})}
 
-	// Create pool with metrics.
-	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, metrics, testTracer)
-	require.NoError(t, err)
-
-	// Given: aibridged is started.
-	srv, err := aibridged.New(ctx, pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
+	// Given: aibridged is started with its server-owned interception pool.
+	srv, err := aibridged.New(ctx, func(ctx context.Context) (aibridged.DRPCClient, error) {
 		return aiBridgeClient, nil
-	}, logger, testTracer)
+	}, logger, testTracer, nil, metrics)
 	require.NoError(t, err, "create new aibridged")
+	require.NoError(t, srv.ReplaceProviders(ctx, providers))
 	t.Cleanup(func() {
 		_ = srv.Shutdown(ctx)
 	})
@@ -504,15 +500,12 @@ func TestIntegrationCircuitBreaker(t *testing.T) {
 		}, nil),
 	}
 
-	// Create pool with metrics.
-	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, metrics, testTracer)
-	require.NoError(t, err)
-
-	// Given: aibridged is started.
-	srv, err := aibridged.New(ctx, pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
+	// Given: aibridged is started with its server-owned interception pool.
+	srv, err := aibridged.New(ctx, func(ctx context.Context) (aibridged.DRPCClient, error) {
 		return aiBridgeClient, nil
-	}, logger, testTracer)
+	}, logger, testTracer, nil, metrics)
 	require.NoError(t, err, "create new aibridged")
+	require.NoError(t, srv.ReplaceProviders(ctx, providers))
 	t.Cleanup(func() {
 		_ = srv.Shutdown(ctx)
 	})
@@ -602,13 +595,11 @@ func TestIntegrationRecordsUpstreamError(t *testing.T) {
 		aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: mockOpenAI.URL, KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-key")}),
 		aibridgetest.NewAnthropicProvider(t, aibridge.AnthropicConfig{BaseURL: mockAnthropic.URL, KeyPool: singleKeyPool(t, config.ProviderAnthropic, "test-key")}, nil),
 	}
-	pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
-	require.NoError(t, err)
-
-	srv, err := aibridged.New(ctx, pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
+	srv, err := aibridged.New(ctx, func(ctx context.Context) (aibridged.DRPCClient, error) {
 		return aiBridgeClient, nil
-	}, logger, testTracer)
+	}, logger, testTracer, nil, nil)
 	require.NoError(t, err, "create new aibridged")
+	require.NoError(t, srv.ReplaceProviders(ctx, providers))
 	t.Cleanup(func() { _ = srv.Shutdown(ctx) })
 
 	cases := []struct {

enterprise/cli/aigatewaystart.go

@@ -24,11 +24,11 @@ import (
 
 	"cdr.dev/slog/v3"
 	"github.com/coder/coder/v2/aibridge"
-	"github.com/coder/coder/v2/aibridge/keypool"
 	aibridgemetrics "github.com/coder/coder/v2/aibridge/metrics"
 	agpl "github.com/coder/coder/v2/cli"
 	"github.com/coder/coder/v2/cli/clilog"
 	"github.com/coder/coder/v2/cli/cliui"
+	agplcoderd "github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/aibridged"
 	coderdtracing "github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/codersdk"
@@ -72,6 +72,7 @@ var aiGatewayInheritedEnvs = map[string]struct{}{
 
 	// Config
 	"CODER_CONFIG_PATH": {},
+	"CODER_EXPERIMENTS": {},
 
 	// AI Gateway
 	"CODER_AI_GATEWAY_ALLOW_BYOK":                        {},
@@ -176,13 +177,6 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command {
 			}
 
 			gatewayLogger := logger.Named("ai-gateway")
-			// Standalone Gateway starts with an empty pool. Providers are
-			// fetched later via GetAIProviders DRPC and pool is updated.
-			pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, gatewayLogger.Named("pool"), metrics, tracer)
-			if err != nil {
-				return xerrors.Errorf("create request pool: %w", err)
-			}
-			gatewayRegisterer.MustRegister(keypool.NewStateCollector(pool.KeyPools))
 
 			return runStandaloneGateway(signalCtx, standaloneGatewayParams{
 				bridgeConfig: vals.AI.BridgeConfig,
@@ -191,8 +185,9 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command {
 				tlsCertFile:  tlsCertFile,
 				tlsKeyFile:   tlsKeyFile,
 
-				dialer: aibridged.NewWebsocketDialer(serverURL, transport, resolvedKey),
-				pool:   pool,
+				dialer:      aibridged.NewWebsocketDialer(serverURL, transport, resolvedKey),
+				experiments: agplcoderd.ReadExperiments(gatewayLogger, vals.Experiments.Value()),
+				registerer:  gatewayRegisterer,
 
 				logger:          gatewayLogger,
 				metrics:         metrics,
@@ -248,6 +243,8 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command {
 type standaloneGatewayParams struct {
 	// Configuration.
 	bridgeConfig codersdk.AIBridgeConfig
+	experiments  codersdk.Experiments
+	registerer   prometheus.Registerer
 	coderURL     string
 	httpAddress  string
 	tlsCertFile  string
@@ -255,7 +252,6 @@ type standaloneGatewayParams struct {
 
 	// Runtime dependencies.
 	dialer aibridged.Dialer
-	pool   aibridged.Pooler
 
 	// Observability.
 	// logger is the gateway-scoped logger; derived loggers (daemon,
@@ -314,15 +310,19 @@ func runStandaloneGateway(ctx context.Context, params standaloneGatewayParams) e
 func newStandaloneGateway(params standaloneGatewayParams) (*standaloneGateway, error) {
 	// The aibridged daemon must outlive the serving context so in-flight HTTP
 	// requests retain their DRPC connection during graceful HTTP shutdown.
-	daemon, err := aibridged.New(context.Background(), params.pool, params.dialer, params.logger.Named("aibridged"), params.tracer)
+	daemon, err := aibridged.New(context.Background(), params.dialer, params.logger.Named("aibridged"), params.tracer, params.experiments, params.metrics)
 	if err != nil {
 		return nil, xerrors.Errorf("start AI Gateway daemon: %w", err)
 	}
 
+	if params.registerer != nil {
+		params.registerer.MustRegister(daemon.KeyPoolStateCollector())
+	}
+
 	providerLogger := params.logger.Named("providers")
 	gateway := &standaloneGateway{
 		daemon:   daemon,
-		reloader: agpl.NewPoolRPCReloader(params.pool, daemon.Client, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics),
+		reloader: agpl.NewProviderRPCReloader(daemon.ReplaceProviders, daemon.Client, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics),
 
 		coderURL:    params.coderURL,
 		httpAddress: params.httpAddress,

enterprise/cli/aigatewaystart_internal_test.go

@@ -107,7 +107,6 @@ func newTestStandaloneGateway(t *testing.T, opts ...testGatewayOption) (*standal
 		httpAddress: "127.0.0.1:0",
 
 		dialer: blockingStandaloneDaemonDialer,
-		pool:   pool,
 
 		logger: logger,
 		tracer: tracer,
@@ -118,6 +117,7 @@ func newTestStandaloneGateway(t *testing.T, opts ...testGatewayOption) (*standal
 
 	gateway, err := newStandaloneGateway(params)
 	require.NoError(t, err)
+	require.NoError(t, gateway.daemon.SetPoolForTest(testutil.Context(t, testutil.WaitShort), t, pool))
 
 	t.Cleanup(func() {
 		require.NoError(t, shutdownWithTimeout(gateway.daemon.Shutdown, testutil.WaitShort))

enterprise/cli/aigatewaystart_test.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/coder/coder/v2/cli/clitest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/externalauth"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
 	"github.com/coder/coder/v2/enterprise/coderd/license"
@@ -183,10 +184,20 @@ func setupAIGatewayDeployment(ctx context.Context, t *testing.T, opts ...aiGatew
 	}
 }
 
+// withAIGatewayEnv sets an environment variable on the gateway invocation. The
+// command resolves its options from the invocation's environment, so a test
+// configures the gateway process the way a deployment would, without mutating
+// the test process's environment.
+func withAIGatewayEnv(name, value string) func(*serpent.Invocation) {
+	return func(inv *serpent.Invocation) {
+		inv.Environ.Set(name, value)
+	}
+}
+
 // startAIGatewayCommand runs `ai-gateway start` and returns the base URL of
 // its HTTP listener, discovered from the startup log line, together with the
-// command's error waiter.
-func startAIGatewayCommand(ctx context.Context, t *testing.T, coderURL, key string) (string, *clitest.ErrorWaiter) {
+// command's error waiter. Each option mutates the invocation before it starts.
+func startAIGatewayCommand(ctx context.Context, t *testing.T, coderURL, key string, opts ...func(*serpent.Invocation)) (string, *clitest.ErrorWaiter) {
 	t.Helper()
 
 	inv, _ := newCLI(t,
@@ -196,6 +207,9 @@ func startAIGatewayCommand(ctx context.Context, t *testing.T, coderURL, key stri
 		"--http-address", "127.0.0.1:0",
 	)
 	inv = inv.WithContext(ctx)
+	for _, opt := range opts {
+		opt(inv)
+	}
 	pty := ptytest.New(t).Attach(inv)
 	waiter := clitest.StartWithWaiter(t, inv)
 
@@ -294,6 +308,115 @@ func TestAIGatewayStartE2E(t *testing.T) {
 	require.NoError(t, waiter.Wait())
 }
 
+// TestAIGatewayStartE2E_ReverseProxyExperiment covers the startup mode
+// selection of the standalone gateway. The gateway process reads the
+// experiment from its own environment, and the selected mode decides which
+// handler serves LLM traffic: the reverse proxy router, whose provider routes
+// are not registered yet, or the interception pool.
+func TestAIGatewayStartE2E_ReverseProxyExperiment(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name string
+		// gatewayExperiments is CODER_EXPERIMENTS for the gateway process.
+		gatewayExperiments string
+		coderdOptions      func(*coderdenttest.Options)
+		wantStatus         int
+		wantBody           string
+		wantUpstreamHits   int32
+		wantSessions       int
+	}{
+		{
+			// Proxy mode is selected, and an enabled provider has no routes
+			// yet, so its requests fall through to the router's catch-all.
+			name:               "ProxyModeWithoutMCP",
+			gatewayExperiments: string(codersdk.ExperimentAIGatewayReverseProxy),
+			wantStatus:         http.StatusNotFound,
+			wantBody:           "route not supported",
+			wantUpstreamHits:   0,
+			wantSessions:       0,
+		},
+		{
+			// The experiment is a property of the gateway process, so coderd
+			// enabling it does not change the gateway's mode.
+			name: "InterceptionWhenGatewayExperimentUnset",
+			coderdOptions: func(opts *coderdenttest.Options) {
+				opts.DeploymentValues.Experiments = serpent.StringArray{string(codersdk.ExperimentAIGatewayReverseProxy)}
+			},
+			wantStatus:       http.StatusOK,
+			wantBody:         "standalone gateway e2e response",
+			wantUpstreamHits: 1,
+			wantSessions:     1,
+		},
+		{
+			// MCP injection requires interception, so a configured MCP server
+			// makes the gateway fall back to it despite the experiment.
+			name:               "InterceptionFallbackWhenMCPConfigured",
+			gatewayExperiments: string(codersdk.ExperimentAIGatewayReverseProxy),
+			coderdOptions: func(opts *coderdenttest.Options) {
+				opts.ExternalAuthConfigs = []*externalauth.Config{{
+					ID: "mcp-provider",
+					// The user holds no external auth link, so no MCP proxy
+					// is built for the request and this URL is never dialed.
+					MCPURL: "http://127.0.0.1:1/mcp",
+				}}
+			},
+			wantStatus:       http.StatusOK,
+			wantBody:         "standalone gateway e2e response",
+			wantUpstreamHits: 1,
+			wantSessions:     1,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testutil.Context(t, testutil.WaitLong)
+			// Given: a coderd entitled for the AI Gateway, with a key and an
+			// enabled provider backed by a mock upstream.
+			var depOpts []aiGatewayDeploymentOption
+			if tc.coderdOptions != nil {
+				depOpts = append(depOpts, withAIGatewayCoderdOptions(tc.coderdOptions))
+			}
+			dep := setupAIGatewayDeployment(ctx, t, depOpts...)
+
+			// When: the gateway starts with that environment.
+			var invOpts []func(*serpent.Invocation)
+			if tc.gatewayExperiments != "" {
+				invOpts = append(invOpts, withAIGatewayEnv("CODER_EXPERIMENTS", tc.gatewayExperiments))
+			}
+			baseURL, waiter := startAIGatewayCommand(ctx, t, dep.client.URL.String(), dep.key, invOpts...)
+
+			// Then: it becomes ready in either mode, which requires the mode
+			// selection and the initial provider load to have completed.
+			requireAIGatewayStatus(ctx, t, baseURL+aiGatewayHealthzPath, http.StatusOK)
+			requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)
+
+			// When: a user sends an LLM request to the gateway.
+			result := postChatCompletionAndRead(ctx, baseURL+aiGatewayChatCompletionPath,
+				dep.userClient.SessionToken(), aiGatewayChatCompletionRequest)
+
+			// Then: the selected handler serves it.
+			require.NoError(t, result.err)
+			require.Equal(t, tc.wantStatus, result.status, "body: %s", result.body)
+			require.Contains(t, string(result.body), tc.wantBody)
+			require.NotContains(t, string(result.body), "sk-e2e", "provider credentials must not reach the caller")
+			require.Equal(t, tc.wantUpstreamHits, dep.upstreamHits.Load())
+
+			// Then: coderd records an interception only for a served request.
+			sessions := requireAIGatewaySessions(ctx, t, dep, tc.wantSessions)
+			if tc.wantSessions > 0 {
+				require.Equal(t, dep.user.Username, sessions[0].Initiator.Username)
+			}
+
+			// When: the command is canceled. Then: it exits cleanly.
+			waiter.Cancel()
+			require.NoError(t, waiter.Wait())
+		})
+	}
+}
+
 // TestAIGatewayStartE2E_InvalidKey covers the fatal error plumbing: a gateway
 // started with a key coderd rejects must exit with the rejection instead of
 // retrying forever.

enterprise/cli/testdata/coder_ai-gateway_start_--help.golden

@@ -14,6 +14,11 @@ USAGE:
   or session token is not required.
 
 OPTIONS:
+      --experiments string-array, $CODER_EXPERIMENTS
+          Enable one or more experiments. These are not ready for production.
+          Separate multiple experiments with commas, or enter '*' to opt-in to
+          all available experiments.
+
       --http-address string, $CODER_AI_GATEWAY_HTTP_ADDRESS (default: 127.0.0.1:4001)
           The bind address to serve incoming AI Gateway client traffic.