feat: add ai-gateway start command (#26605)

> AI Tools were used to produce this PR

This PR adds `coder ai-gateway start` command that runs the AI Gateway
as an independent process.

- Standalone process doesn't have access to DB. Uses DRPC services under
`/api/v2/ai-gateway/serve`for auth, recording and provider
initialization.
- It only handles LLM traffic, other endpoints (eg. `/sessions`) are
only available though `coderd`.
- The standalone gateway reuses applicable flags from AI Gateway
deployment options. Provider-seeding and coderd-only options are
excluded.
- Only added to fat build, the slim build stub rejects the command.

Some wiring used by this new command is added.

**`NewWebsocketDialer`** - implements the standalone gateway's
connection to coderd's `/api/v2/ai-gateway/serve` endpoint. It upgrades
to a WebSocket, multiplexes with yamux, and wires all DRPC services.

**`AIGatewayDataPlaneMiddleware`** - extracts the per-request middleware
chain (concurrency limiting, rate limiting, BYOK gating) into a shared
function used by both the embedded route and the standalone gateway.

**`RootCmd.ResolveClientConnection`** - resolve the deployment URL and
builds an HTTP transport without requiring a session token. Used in
`ai-gateway start`command as it authenticates using different credential
type.

---------

Co-authored-by: Danny Kopping <danny@coder.com>

coderd/aibridged/aibridged.go

@@ -16,7 +16,11 @@ import (
 	"github.com/coder/retry"
 )
 
-var _ io.Closer = &Server{}
+var (
+	_ io.Closer = &Server{}
+
+	ErrShutdown = xerrors.New("aibridged server shutdown")
+)
 
 // Server provides the AI Bridge functionality.
 // It is responsible for:
@@ -42,10 +46,11 @@ type Server struct {
 	initConnectionCh   chan struct{}
 	initConnectionOnce sync.Once
 
-	// lifecycleCtx is canceled when we start closing.
+	// lifecycleCtx is canceled when we start closing or when the
+	// connection loop exits permanently.
 	lifecycleCtx context.Context
-	// cancelFn closes the lifecycleCtx.
-	cancelFn func()
+	// cancelFn closes the lifecycleCtx with the reason it closed.
+	cancelFn context.CancelCauseFunc
 
 	shutdownOnce sync.Once
 }
@@ -55,7 +60,7 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger,
 		return nil, xerrors.Errorf("nil rpcDialer given")
 	}
 
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancelCause(ctx)
 	daemon := &Server{
 		logger:           logger,
 		tracer:           tracer,
@@ -78,6 +83,11 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger,
 func (s *Server) connect() {
 	defer s.logger.Debug(s.lifecycleCtx, "connect loop exited")
 	defer s.wg.Done()
+	defer func() {
+		if s.lifecycleCtx.Err() == nil {
+			s.cancelFn(xerrors.New("connect loop exited"))
+		}
+	}()
 
 	logConnect := s.logger.With(slog.F("context", "aibridged.server")).Debug
 	// An exponential back-off occurs when the connection is failing to dial.
@@ -93,13 +103,25 @@ connectLoop:
 		client, err := s.clientDialer(s.lifecycleCtx)
 		if err != nil {
 			if errors.Is(err, context.Canceled) {
+				if s.lifecycleCtx.Err() == nil {
+					s.cancelFn(err)
+				}
 				return
 			}
 			var sdkErr *codersdk.Error
-			// If something is wrong with our auth, stop trying to connect.
-			if errors.As(err, &sdkErr) && sdkErr.StatusCode() == http.StatusForbidden {
-				s.logger.Error(s.lifecycleCtx, "not authorized to dial coderd", slog.Error(err))
-				return
+			// If something is wrong with configuration, stop trying to connect.
+			if errors.As(err, &sdkErr) {
+				switch sdkErr.StatusCode() {
+				// These statuses are terminal failures from the /api/v2/ai-gateway/serve
+				// handshake: wrong gateway key, incompatible API version, or entitlement failure.
+				case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden:
+					err = xerrors.Errorf("dial coderd: %w", err)
+					s.logger.Error(s.lifecycleCtx, "fatal error dialing coderd", slog.Error(err))
+					s.cancelFn(err)
+					return
+				default:
+					err = xerrors.Errorf("unexpected HTTP response dialing coderd: %w", err)
+				}
 			}
 			if s.isShutdown() {
 				return
@@ -133,9 +155,32 @@ connectLoop:
 	}
 }
 
+// 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{} {
+	return s.lifecycleCtx.Done()
+}
+
+// Err returns the reason the server lifecycle ended.
+func (s *Server) Err() error {
+	if cause := context.Cause(s.lifecycleCtx); cause != nil {
+		return cause
+	}
+	return s.lifecycleCtx.Err()
+}
+
 func (s *Server) Client() (DRPCClient, error) {
+	return s.ClientContext(context.Background())
+}
+
+func (s *Server) ClientContext(ctx context.Context) (DRPCClient, error) {
 	select {
-	case <-s.lifecycleCtx.Done():
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-s.Done():
+		if err := s.Err(); err != nil {
+			return nil, err
+		}
 		return nil, xerrors.New("context closed")
 	case client := <-s.clientCh:
 		return client, nil
@@ -170,7 +215,7 @@ func (s *Server) isShutdown() bool {
 func (s *Server) Shutdown(ctx context.Context) error {
 	var err error
 	s.shutdownOnce.Do(func() {
-		s.cancelFn()
+		s.cancelFn(ErrShutdown)
 
 		// Wait for any outstanding connections to terminate.
 		s.wg.Wait()

coderd/aibridged/aibridged_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 
 	"github.com/google/uuid"
@@ -40,8 +41,13 @@ func singleKeyPool(t *testing.T, name, key string) *keypool.Pool {
 
 func newTestServer(t *testing.T) (*aibridged.Server, *mock.MockDRPCClient, *mock.MockPooler) {
 	t.Helper()
+	return newTestServerWithDialer(t, nil, nil)
+}
 
-	logger := slogtest.Make(t, nil)
+func newTestServerWithDialer(t *testing.T, dialer aibridged.Dialer, loggerOptions *slogtest.Options) (*aibridged.Server, *mock.MockDRPCClient, *mock.MockPooler) {
+	t.Helper()
+
+	logger := slogtest.Make(t, loggerOptions)
 	ctrl := gomock.NewController(t)
 	client := mock.NewMockDRPCClient(ctrl)
 	pool := mock.NewMockPooler(ctrl)
@@ -50,12 +56,12 @@ func newTestServer(t *testing.T) (*aibridged.Server, *mock.MockDRPCClient, *mock
 	client.EXPECT().DRPCConn().AnyTimes().Return(conn)
 	pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
 
-	srv, err := aibridged.New(
-		t.Context(),
-		pool,
-		func(ctx context.Context) (aibridged.DRPCClient, error) {
+	if dialer == nil {
+		dialer = func(ctx context.Context) (aibridged.DRPCClient, error) {
 			return client, nil
-		}, logger, testTracer)
+		}
+	}
+	srv, err := aibridged.New(t.Context(), pool, dialer, logger, testTracer)
 	require.NoError(t, err, "create new aibridged")
 	t.Cleanup(func() {
 		srv.Shutdown(context.Background())
@@ -79,6 +85,39 @@ func (*mockDRPCConn) NewStream(ctx context.Context, rpc string, enc drpc.Encodin
 	return nil, nil
 }
 
+func sdkError(status int, message string) error {
+	return codersdk.ReadBodyAsError(&http.Response{
+		StatusCode: status,
+		Header:     http.Header{"Content-Type": []string{"application/json"}},
+		Body:       io.NopCloser(bytes.NewBufferString(`{"message":"` + message + `"}`)),
+	})
+}
+
+func TestClient_TransientDialErrorRetries(t *testing.T) {
+	t.Parallel()
+
+	var calls atomic.Int32
+	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")
+		}
+		return client, nil
+	}
+
+	srv, err := aibridged.New(t.Context(), pool, dialFc, slogtest.Make(t, nil), testTracer)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
+
+	_, err = srv.ClientContext(testutil.Context(t, testutil.WaitShort))
+	require.NoError(t, err)
+	require.Equal(t, int32(2), calls.Load())
+}
+
 func TestServeHTTP_FailureModes(t *testing.T) {
 	t.Parallel()
 
@@ -91,6 +130,7 @@ func TestServeHTTP_FailureModes(t *testing.T) {
 		applyMocksFn   func(client *mock.MockDRPCClient, pool *mock.MockPooler)
 		dialerFn       aibridged.Dialer
 		contextFn      func() context.Context
+		ignoreLogs     bool
 		expectedErr    error
 		expectedStatus int
 	}{
@@ -127,7 +167,34 @@ func TestServeHTTP_FailureModes(t *testing.T) {
 			expectedStatus: http.StatusForbidden,
 		},
 
-		// TODO: coderd connection-related failures.
+		// Coderd connection-related failures.
+		{
+			name: "fatal bad request dial error",
+			dialerFn: func(context.Context) (aibridged.DRPCClient, error) {
+				return nil, sdkError(http.StatusBadRequest, "bad request")
+			},
+			ignoreLogs:     true,
+			expectedErr:    aibridged.ErrConnect,
+			expectedStatus: http.StatusServiceUnavailable,
+		},
+		{
+			name: "fatal unauthorized dial error",
+			dialerFn: func(context.Context) (aibridged.DRPCClient, error) {
+				return nil, sdkError(http.StatusUnauthorized, "unauthorized")
+			},
+			ignoreLogs:     true,
+			expectedErr:    aibridged.ErrConnect,
+			expectedStatus: http.StatusServiceUnavailable,
+		},
+		{
+			name: "fatal forbidden dial error",
+			dialerFn: func(context.Context) (aibridged.DRPCClient, error) {
+				return nil, sdkError(http.StatusForbidden, "forbidden")
+			},
+			ignoreLogs:     true,
+			expectedErr:    aibridged.ErrConnect,
+			expectedStatus: http.StatusServiceUnavailable,
+		},
 
 		// Budget-related failures.
 		{
@@ -173,7 +240,11 @@ func TestServeHTTP_FailureModes(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			srv, client, pool := newTestServer(t)
+			var loggerOptions *slogtest.Options
+			if tc.ignoreLogs {
+				loggerOptions = &slogtest.Options{IgnoreErrors: true}
+			}
+			srv, client, pool := newTestServerWithDialer(t, tc.dialerFn, loggerOptions)
 			conn := &mockDRPCConn{}
 			client.EXPECT().DRPCConn().AnyTimes().Return(conn)
 

coderd/aibridged/dialer.go

@@ -0,0 +1,101 @@
+package aibridged
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net/http"
+	"net/url"
+
+	"github.com/hashicorp/yamux"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/buildinfo"
+	aibridgedproto "github.com/coder/coder/v2/coderd/aibridged/proto"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
+	"github.com/coder/websocket"
+)
+
+// NewWebsocketDialer returns a [Dialer] that connects a standalone AI
+// Gateway to coderd's /api/v2/ai-gateway/serve endpoint over a WebSocket,
+// multiplexes it with yamux, and exposes the aibridged DRPC services
+// (Recorder, MCPConfigurator, Authorizer, ProviderConfigurator) over it.
+// This is the standalone counterpart to API.CreateInMemoryAIBridgeServer,
+// which wires the same services over an in-memory pipe for the embedded
+// daemon.
+//
+// The gateway authenticates with an AI Gateway key
+// (codersdk.AIGatewayKeyHeader), advertises its API version via the
+// "version" query parameter, and reports its build version via
+// codersdk.BuildVersionHeader (used by coderd for observability only).
+// TLS for this connection is governed by the scheme of serverURL and any
+// TLS configuration baked into transport.
+//
+// On a failed upgrade the coderd HTTP error is returned as a
+// *codersdk.Error so [Server.connect] can distinguish fatal
+// auth/entitlement failures from transient ones.
+func readAIGatewayServeError(res *http.Response) error {
+	err := codersdk.ReadBodyAsError(res)
+
+	var sdkErr *codersdk.Error
+	if errors.As(err, &sdkErr) && res.StatusCode == http.StatusUnauthorized {
+		// /ai-gateway/serve authenticates with an AI Gateway key, not a user
+		// session. Generic user-login helpers are misleading here.
+		sdkErr.Helper = ""
+	}
+	return err
+}
+
+func NewWebsocketDialer(serverURL *url.URL, transport http.RoundTripper, key string) Dialer {
+	return func(ctx context.Context) (DRPCClient, error) {
+		serveURL, err := serverURL.Parse("/api/v2/ai-gateway/serve")
+		if err != nil {
+			return nil, xerrors.Errorf("parse url: %w", err)
+		}
+		query := serveURL.Query()
+		query.Add(aibridgedproto.VersionQueryParam, aibridgedproto.CurrentVersion.String())
+		serveURL.RawQuery = query.Encode()
+
+		headers := http.Header{}
+		headers.Set(codersdk.BuildVersionHeader, buildinfo.Version())
+		headers.Set(codersdk.AIGatewayKeyHeader, key)
+
+		httpClient := &http.Client{
+			Transport: transport,
+		}
+		// nolint:bodyclose // ReadBodyAsError closes the body; success path hands off to the websocket conn.
+		conn, res, err := websocket.Dial(ctx, serveURL.String(), &websocket.DialOptions{
+			HTTPClient:      httpClient,
+			CompressionMode: websocket.CompressionDisabled,
+			HTTPHeader:      headers,
+		})
+		if err != nil {
+			if res == nil {
+				return nil, err
+			}
+			return nil, readAIGatewayServeError(res)
+		}
+		config := yamux.DefaultConfig()
+		config.LogOutput = io.Discard
+		// Use a background context because the caller closes the client
+		// (and thus the multiplexed session) explicitly.
+		_, wsNetConn := codersdk.WebsocketNetConn(context.Background(), conn, websocket.MessageBinary)
+		conn.SetReadLimit(drpcsdk.YamuxDefaultStreamWindowSize)
+		session, err := yamux.Client(wsNetConn, config)
+		if err != nil {
+			_ = conn.Close(websocket.StatusGoingAway, "")
+			_ = wsNetConn.Close()
+			return nil, xerrors.Errorf("multiplex client: %w", err)
+		}
+
+		dconn := drpcsdk.MultiplexedConn(session)
+		return &Client{
+			Conn:                           dconn,
+			DRPCRecorderClient:             aibridgedproto.NewDRPCRecorderClient(dconn),
+			DRPCMCPConfiguratorClient:      aibridgedproto.NewDRPCMCPConfiguratorClient(dconn),
+			DRPCAuthorizerClient:           aibridgedproto.NewDRPCAuthorizerClient(dconn),
+			DRPCProviderConfiguratorClient: aibridgedproto.NewDRPCProviderConfiguratorClient(dconn),
+		}, nil
+	}
+}

coderd/aibridged/http.go

@@ -117,7 +117,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 		r.Header.Del("X-Api-Key")
 	}
 
-	client, err := s.Client()
+	client, err := s.ClientContext(ctx)
 	if err != nil {
 		logger.Warn(ctx, "failed to connect to coderd", slog.Error(err))
 		http.Error(rw, ErrConnect.Error(), http.StatusServiceUnavailable)

coderd/aibridged/proto/version.go

@@ -17,6 +17,11 @@ const (
 	CurrentMinor = 1
 )
 
+// VersionQueryParam is the URL query parameter the standalone AI Gateway
+// uses to advertise its aibridged API version when dialing coderd's serve
+// endpoint, and that coderd reads to negotiate compatibility.
+const VersionQueryParam = "version"
+
 // CurrentVersion is the current aibridged API version.
 // Breaking changes to the aibridged API **MUST** increment CurrentMajor above.
 // Non-breaking changes to the aibridged API **MUST** increment CurrentMinor