feat: add /api/v2/aibridge/serve endpoint (#26506)

Adds a new enterprise-only `GET /api/v2/ai-gateway/serve` endpoint that standalone AI Gateway replicas use to connect to `coderd` over a DRPC-over-WebSocket transport, mirroring the existing in-memory path used by the embedded AI Bridge daemon.

- The endpoint upgrades the HTTP connection to a WebSocket, multiplexes it with yamux, and finally serves the three DRPC services (Recorder, MCPConfigurator, Authorizer).
- The `X-AI-Governance-Gateway-Key` header is used for authentication.
    - The key is looked up by its hashed secret
    - Missing or revoked keys return `401`.
- API version negotiation is enforced via a new `aibridged/proto` version (`v1.0`).
    - Incompatible versions return `400`.
- `FeatureAIBridge` entitlement is required.
- Key liveness (`last_used_at`) is recorded immediately on connection and refreshed every 60 seconds while the session remains open.
  - When key liveness detects the key was deleted (no rows where updated) session is closed.

#### Small refactors

* The three DRPC service registrations are extracted into `aibridgedserver.Register`, shared by both the in-memory and WebSocket paths.

* The literal `256 * 1024` used as the yamux-aligned WebSocket read limit is replaced with the named constant `drpcsdk.YamuxDefaultStreamWindowSize` in all call sites.
  * as noted in review comment https://github.com/coder/coder/pull/26506#discussion_r3461905223 order of `SetReadLimit` and `WebsocketNetConn` calls was fixed.

coderd/aibridged.go

@@ -6,7 +6,6 @@ import (
 	"io"
 	"net/http"
 
-	"golang.org/x/xerrors"
 	"storj.io/drpc/drpcmux"
 	"storj.io/drpc/drpcserver"
 
@@ -71,17 +70,8 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai
 	if err != nil {
 		return nil, err
 	}
-	err = aibridgedproto.DRPCRegisterRecorder(mux, srv)
-	if err != nil {
-		return nil, xerrors.Errorf("register recorder service: %w", err)
-	}
-	err = aibridgedproto.DRPCRegisterMCPConfigurator(mux, srv)
-	if err != nil {
-		return nil, xerrors.Errorf("register MCP configurator service: %w", err)
-	}
-	err = aibridgedproto.DRPCRegisterAuthorizer(mux, srv)
-	if err != nil {
-		return nil, xerrors.Errorf("register key validator service: %w", err)
+	if err := aibridgedserver.Register(mux, srv); err != nil {
+		return nil, err
 	}
 	server := drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux},
 		drpcserver.Options{

coderd/aibridged/proto/version.go

@@ -0,0 +1,19 @@
+package proto
+
+import "github.com/coder/coder/v2/apiversion"
+
+// Version history:
+//
+// API v1.0:
+//   - Initial version. Serves the Recorder, MCPConfigurator, and Authorizer
+//     services to embedded and standalone AI Gateway daemons.
+const (
+	CurrentMajor = 1
+	CurrentMinor = 0
+)
+
+// 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
+// above.
+var CurrentVersion = apiversion.New(CurrentMajor, CurrentMinor)

coderd/aibridgedserver/aibridgedserver.go

@@ -594,13 +594,13 @@ externalAuthLoop:
 // IsAuthorized validates a given Coder API key and returns the user ID to which it belongs (if valid).
 //
 // SECURITY: when in.KeyId is set (the "delegated" path), this method trusts the
-// caller's claim of identity and skips the key-secret check. This is safe only
-// because the DRPCServer is reachable solely via the in-process
-// [aibridged.MemTransportPipe]; the handler itself cannot tell whether it was
-// invoked over the in-memory pipe or a network socket. If this RPC is ever
-// exposed over a network boundary, any caller who knows a valid 10-char key ID
-// (which is not secret) could authenticate as the key's owner without the
-// secret. Do not bind this DRPCServer to a network listener.
+// caller's claim of identity and skips the key-secret check. This DRPCServer is
+// reachable both in-process via [aibridged.MemTransportPipe] and over the network
+// via the /api/v2/ai-gateway/serve endpoint. That endpoint admits only holders of
+// AI Gateway key, which are fully trusted. Standalone AI Gateway authenticates its
+// own users and acts on their behalf, much like a provisioner daemon. A Gateway key
+// holder can therefore act as any user without that user's secret. Per-user
+// authorization on this surface is a known gap.
 //
 // NOTE: this should really be using the code from [httpmw.ExtractAPIKey]. That function not only validates the key
 // but handles many other cases like updating last used, expiry, etc. This code does not currently use it for

coderd/aibridgedserver/register.go

@@ -0,0 +1,25 @@
+package aibridgedserver
+
+import (
+	"golang.org/x/xerrors"
+	"storj.io/drpc/drpcmux"
+
+	"github.com/coder/coder/v2/coderd/aibridged/proto"
+)
+
+// Register registers the Recorder, MCPConfigurator, and Authorizer DRPC
+// services backed by srv onto mux. It is shared by the embedded in-memory
+// server and the standalone /api/v2/ai-gateway/serve WebSocket handler so both
+// expose an identical service set.
+func Register(mux *drpcmux.Mux, srv *Server) error {
+	if err := proto.DRPCRegisterRecorder(mux, srv); err != nil {
+		return xerrors.Errorf("register recorder service: %w", err)
+	}
+	if err := proto.DRPCRegisterMCPConfigurator(mux, srv); err != nil {
+		return xerrors.Errorf("register MCP configurator service: %w", err)
+	}
+	if err := proto.DRPCRegisterAuthorizer(mux, srv); err != nil {
+		return xerrors.Errorf("register authorizer service: %w", err)
+	}
+	return nil
+}

coderd/apidoc/docs.go

@@ -1532,6 +1532,25 @@ const docTemplate = `{
                 ]
             }
         },
+        "/api/v2/ai-gateway/serve": {
+            "get": {
+                "tags": [
+                    "Enterprise"
+                ],
+                "summary": "AI Gateway serve",
+                "operationId": "ai-gateway-serve",
+                "responses": {
+                    "101": {
+                        "description": "Switching Protocols"
+                    }
+                },
+                "security": [
+                    {
+                        "AIGatewayKey": []
+                    }
+                ]
+            }
+        },
         "/api/v2/ai-gateway/sessions": {
             "get": {
                 "description": "Alias: also available at /api/v2/aibridge/sessions for backward compatibility.",
@@ -28912,6 +28931,11 @@ const docTemplate = `{
         }
     },
     "securityDefinitions": {
+        "AIGatewayKey": {
+            "type": "apiKey",
+            "name": "X-AI-Governance-Gateway-Key",
+            "in": "header"
+        },
         "Authorization": {
             "type": "apiKey",
             "name": "Authorizaiton",

coderd/apidoc/swagger.json

@@ -1355,6 +1355,23 @@
 				]
 			}
 		},
+		"/api/v2/ai-gateway/serve": {
+			"get": {
+				"tags": ["Enterprise"],
+				"summary": "AI Gateway serve",
+				"operationId": "ai-gateway-serve",
+				"responses": {
+					"101": {
+						"description": "Switching Protocols"
+					}
+				},
+				"security": [
+					{
+						"AIGatewayKey": []
+					}
+				]
+			}
+		},
 		"/api/v2/ai-gateway/sessions": {
 			"get": {
 				"description": "Alias: also available at /api/v2/aibridge/sessions for backward compatibility.",
@@ -26681,6 +26698,11 @@
 		}
 	},
 	"securityDefinitions": {
+		"AIGatewayKey": {
+			"type": "apiKey",
+			"name": "X-AI-Governance-Gateway-Key",
+			"in": "header"
+		},
 		"Authorization": {
 			"type": "apiKey",
 			"name": "Authorizaiton",

coderd/coderd.go

@@ -338,6 +338,10 @@ type Options struct {
 // @securitydefinitions.apiKey CoderSessionToken
 // @in header
 // @name Coder-Session-Token
+
+// @securitydefinitions.apiKey AIGatewayKey
+// @in header
+// @name X-AI-Governance-Gateway-Key
 // New constructs a Coder API handler.
 func New(options *Options) *API {
 	if options == nil {

coderd/coderdtest/swaggerparser.go

@@ -370,6 +370,11 @@ func assertSecurityDefined(t *testing.T, comment SwaggerComment) {
 		comment.router == "/api/v2/init-script/{os}/{arch}" {
 		return // endpoints do not require authorization
 	}
+	if comment.router == "/api/v2/ai-gateway/serve" {
+		assert.Equal(t, "AIGatewayKey", comment.security, "@Security must be AIGatewayKey")
+		return
+	}
+
 	assert.Containsf(t, authorizedSecurityTags, comment.security, "@Security must be either of these options: %v", authorizedSecurityTags)
 }
 

codersdk/client.go

@@ -96,6 +96,9 @@ const (
 	// ProvisionerDaemonKey contains the authentication key for an external provisioner daemon
 	ProvisionerDaemonKey = "Coder-Provisioner-Daemon-Key"
 
+	// AIGatewayKeyHeader contains the authentication key for a standalone AI Gateway replica.
+	AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key"
+
 	// BuildVersionHeader contains build information of Coder.
 	BuildVersionHeader = "X-Coder-Build-Version"
 

codersdk/drpcsdk/transport.go

@@ -18,6 +18,10 @@ const (
 	// MaxMessageSize is the maximum payload size that can be
 	// transported without error.
 	MaxMessageSize = 4 << 20
+
+	// YamuxDefaultStreamWindowSize matches hashicorp/yamux's unexported
+	// initialStreamWindow, which DefaultConfig uses as MaxStreamWindowSize.
+	YamuxDefaultStreamWindowSize = 256 * 1024
 )
 
 func DefaultDRPCOptions(options *drpcmanager.Options) drpcmanager.Options {

codersdk/provisionerdaemons.go

@@ -343,13 +343,11 @@ func (c *Client) ServeProvisionerDaemon(ctx context.Context, req ServeProvisione
 		}
 		return nil, ReadBodyAsError(res)
 	}
-	// Align with the frame size of yamux.
-	conn.SetReadLimit(256 * 1024)
-
 	config := yamux.DefaultConfig()
 	config.LogOutput = io.Discard
 	// Use background context because caller should close the client.
 	_, wsNetConn := WebsocketNetConn(context.Background(), conn, websocket.MessageBinary)
+	conn.SetReadLimit(drpcsdk.YamuxDefaultStreamWindowSize)
 	session, err := yamux.Client(wsNetConn, config)
 	if err != nil {
 		_ = conn.Close(websocket.StatusGoingAway, "")

docs/reference/api/enterprise.md

@@ -304,6 +304,26 @@ curl -X DELETE http://coder-server:8080/api/v2/ai-gateway/keys/{key} \
 
 To perform this operation, you must be authenticated. [Learn more](authentication.md).
 
+## AI Gateway serve
+
+### Code samples
+
+```shell
+# Example request using curl
+curl -X GET http://coder-server:8080/api/v2/ai-gateway/serve \
+  -H 'X-AI-Governance-Gateway-Key: API_KEY'
+```
+
+`GET /api/v2/ai-gateway/serve`
+
+### Responses
+
+| Status | Meaning                                                                  | Description         | Schema |
+|--------|--------------------------------------------------------------------------|---------------------|--------|
+| 101    | [Switching Protocols](https://tools.ietf.org/html/rfc7231#section-6.2.2) | Switching Protocols |        |
+
+To perform this operation, you must be authenticated. [Learn more](authentication.md).
+
 ## Get appearance
 
 ### Code samples

enterprise/coderd/aibridgeserve.go

@@ -0,0 +1,239 @@
+package coderd
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/hashicorp/yamux"
+	"golang.org/x/xerrors"
+	"storj.io/drpc/drpcmux"
+	"storj.io/drpc/drpcserver"
+
+	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/buildinfo"
+	aibridgedproto "github.com/coder/coder/v2/coderd/aibridged/proto"
+	"github.com/coder/coder/v2/coderd/aibridgedserver"
+	"github.com/coder/coder/v2/coderd/apikey"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/httpapi"
+	"github.com/coder/coder/v2/coderd/httpmw/loggermw"
+	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
+	"github.com/coder/websocket"
+)
+
+// aiGatewayKeyLastUsedInterval defines how often an active DRPC session refreshes
+// last_used_at for its authenticating key.
+const aiGatewayKeyLastUsedInterval = 60 * time.Second
+
+// aiGatewayServe upgrades the connection to a WebSocket and serves the DRPC
+// services (Recorder, MCPConfigurator, Authorizer) to a remote standalone AI
+// Gateway replica, mirroring the embedded case. AI Gateway key authentication is
+// enforced before the WebSocket upgrade. License entitlement is enforced by
+// middleware on the route.
+//
+// @Summary AI Gateway serve
+// @ID ai-gateway-serve
+// @Security AIGatewayKey
+// @Tags Enterprise
+// @Success 101
+// @Router /api/v2/ai-gateway/serve [get]
+func (api *API) aiGatewayServe(rw http.ResponseWriter, r *http.Request) {
+	key := r.Header.Get(codersdk.AIGatewayKeyHeader)
+	if key == "" {
+		httpapi.Write(r.Context(), rw, http.StatusUnauthorized, codersdk.Response{
+			Message: "AI Gateway key required.",
+		})
+		return
+	}
+
+	// nolint:gocritic // AI Gateway doesn't have Coder identity.System must look up the AI Gateway key to authenticate the request.
+	gatewayKey, err := api.Database.GetAIGatewayKeyByHashedSecret(dbauthz.AsSystemRestricted(r.Context()), apikey.HashSecret(key))
+	if err != nil {
+		if httpapi.Is404Error(err) {
+			httpapi.Write(r.Context(), rw, http.StatusUnauthorized, codersdk.Response{
+				Message: "AI Gateway key invalid.",
+			})
+			return
+		}
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to look up AI Gateway key.",
+		})
+		return
+	}
+
+	clientAPIVersion := r.URL.Query().Get("version")
+	clientCoderVersion := r.Header.Get(codersdk.BuildVersionHeader)
+	logger := api.Logger.Named("aigateway-serve").With(
+		slog.F("remote_addr", r.RemoteAddr),
+		slog.F("client_api_version", clientAPIVersion),
+		slog.F("client_build_version", clientCoderVersion),
+		slog.F("server_api_version", aibridgedproto.CurrentVersion.String()),
+		slog.F("server_build_version", buildinfo.Version),
+		slog.F("ai_gateway_key_id", gatewayKey.ID),
+		slog.F("ai_gateway_key_name", gatewayKey.Name),
+		slog.F("ai_gateway_key_prefix", gatewayKey.SecretPrefix),
+	)
+
+	// keyCtx bounds all work for this authenticated key. Canceling it terminates
+	// the websocket session and related background work.
+	keyCtx, keyCtxCancel := context.WithCancel(r.Context())
+	defer keyCtxCancel()
+
+	if err := aibridgedproto.CurrentVersion.Validate(clientAPIVersion); err != nil {
+		httpapi.Write(keyCtx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Incompatible or unparsable version",
+			Validations: []codersdk.ValidationError{
+				{Field: "version", Detail: err.Error()},
+				{Field: "client_api_version", Detail: clientAPIVersion},
+				{Field: "server_api_version", Detail: aibridgedproto.CurrentVersion.String()},
+			},
+		})
+		return
+	}
+
+	// Track the websocket so API shutdown waits for it to close.
+	api.AGPL.WebsocketWaitMutex.Lock()
+	api.AGPL.WebsocketWaitGroup.Add(1)
+	api.AGPL.WebsocketWaitMutex.Unlock()
+	defer api.AGPL.WebsocketWaitGroup.Done()
+
+	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		// Need to disable compression to avoid a data-race, yamux reads and writes concurrently.
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		if !xerrors.Is(err, context.Canceled) {
+			logger.Error(keyCtx, "websocket upgrade failed", slog.Error(err))
+		}
+		httpapi.Write(keyCtx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Failed to accept websocket connection.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	config := yamux.DefaultConfig()
+	config.LogOutput = io.Discard
+	connCtx, wsNetConn := codersdk.WebsocketNetConn(keyCtx, conn, websocket.MessageBinary)
+	conn.SetReadLimit(drpcsdk.YamuxDefaultStreamWindowSize)
+	defer wsNetConn.Close()
+	session, err := yamux.Server(wsNetConn, config)
+	if err != nil {
+		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("multiplex server: %s", err))
+		return
+	}
+
+	if _, err := aiGatewayUpdateKeyLastUsed(connCtx, api, gatewayKey.ID); err != nil {
+		logger.Warn(connCtx, "update ai gateway key last used", slog.Error(err))
+	}
+	go aiGatewayTrackKeyUsage(connCtx, keyCtxCancel, api, gatewayKey.ID, logger)
+
+	mux := drpcmux.New()
+	srv, err := aibridgedserver.NewServer(
+		connCtx,
+		api.Database,
+		logger,
+		api.AccessURL.String(),
+		api.DeploymentValues.AI.BridgeConfig,
+		api.ExternalAuthConfigs,
+		api.AGPL.Experiments,
+		api.AGPL.AISeatTracker,
+	)
+	if err != nil {
+		if !xerrors.Is(err, context.Canceled) {
+			logger.Error(connCtx, "server creation failed", slog.Error(err))
+		}
+		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("create ai gateway server: %s", err))
+		return
+	}
+	if err := aibridgedserver.Register(mux, srv); err != nil {
+		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("register ai gateway services: %s", err))
+		return
+	}
+
+	server := drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux},
+		drpcserver.Options{
+			Manager: drpcsdk.DefaultDRPCOptions(nil),
+			Log: func(err error) {
+				if xerrors.Is(err, io.EOF) {
+					return
+				}
+				logger.Debug(connCtx, "drpc server error", slog.Error(err))
+			},
+		},
+	)
+
+	// Log the request immediately instead of after it completes.
+	if rl := loggermw.RequestLoggerFromContext(connCtx); rl != nil {
+		rl.WriteLog(connCtx, http.StatusAccepted)
+	}
+
+	logger.Info(connCtx, "opened connection")
+	err = server.Serve(connCtx, session)
+	logger.Info(connCtx, "closed connection", slog.Error(err))
+	if err != nil && !xerrors.Is(err, io.EOF) {
+		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("serve: %s", err))
+		return
+	}
+	_ = conn.Close(websocket.StatusGoingAway, "")
+}
+
+// aiGatewayUpdateKeyLastUsed records liveness for keyID and returns whether
+// the key is still active. On error key is assumed to not be active.
+func aiGatewayUpdateKeyLastUsed(ctx context.Context, api *API, keyID uuid.UUID) (bool, error) {
+	// nolint:gocritic // Recording AI Gateway key liveness is an internal system write.
+	rows, err := api.Database.UpdateAIGatewayKeyLastUsedAt(dbauthz.AsSystemRestricted(ctx), keyID)
+	if err != nil {
+		return false, err
+	}
+	return rows > 0, nil
+}
+
+// aiGatewayTrackKeyUsage refreshes last_used_at for keyID on a fixed interval until ctx is canceled.
+func aiGatewayTrackKeyUsage(ctx context.Context, ctxCancel context.CancelFunc, api *API, keyID uuid.UUID, logger slog.Logger) {
+	ticker, done := api.NewTicker(aiGatewayKeyLastUsedInterval)
+	defer done()
+
+	consecutiveFailures := 0
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker:
+		}
+
+		active, err := aiGatewayUpdateKeyLastUsed(ctx, api, keyID)
+		if err == nil && !active {
+			logger.Info(ctx, "ai gateway key no longer exists, closing connection")
+			ctxCancel()
+			return
+		}
+
+		if err != nil {
+			if xerrors.Is(err, context.Canceled) {
+				return
+			}
+			consecutiveFailures++
+			// Log failures with exponential backoff (1, 2, 4, 8...).
+			// First failure logged at Debug, next failures escalate to Warn.
+			if consecutiveFailures&(consecutiveFailures-1) == 0 {
+				if consecutiveFailures == 1 {
+					logger.Debug(ctx, "update ai gateway key last used", slog.Error(err), slog.F("consecutive_failures", consecutiveFailures))
+				} else {
+					logger.Warn(ctx, "update ai gateway key last used", slog.Error(err), slog.F("consecutive_failures", consecutiveFailures))
+				}
+			}
+			continue
+		}
+		if consecutiveFailures > 1 {
+			logger.Info(ctx, "ai gateway key last used update recovered",
+				slog.F("consecutive_failures", consecutiveFailures))
+		}
+		consecutiveFailures = 0
+	}
+}

enterprise/coderd/aibridgeserve_test.go

@@ -0,0 +1,223 @@
+package coderd_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/yamux"
+	"github.com/stretchr/testify/require"
+
+	aibridgedproto "github.com/coder/coder/v2/coderd/aibridged/proto"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
+	"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
+	"github.com/coder/coder/v2/enterprise/coderd/license"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/serpent"
+	"github.com/coder/websocket"
+)
+
+// dialAIGatewayServe dials /api/v2/ai-gateway/serve, authenticating with the given
+// gateway key and API version. On a successful WebSocket upgrade it returns a
+// yamux session and http.StatusSwitchingProtocols. Otherwise it returns a nil
+// session and the HTTP status code coderd responded with.
+func dialAIGatewayServe(ctx context.Context, t *testing.T, client *codersdk.Client, key string, version string) (*yamux.Session, int) {
+	t.Helper()
+
+	serverURL, err := client.URL.Parse("/api/v2/ai-gateway/serve")
+	require.NoError(t, err)
+	query := serverURL.Query()
+	if version != "" {
+		query.Set("version", version)
+	}
+	serverURL.RawQuery = query.Encode()
+
+	headers := http.Header{}
+	if key != "" {
+		headers.Set(codersdk.AIGatewayKeyHeader, key)
+	}
+
+	conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{
+		HTTPClient:      &http.Client{Transport: client.HTTPClient.Transport},
+		CompressionMode: websocket.CompressionDisabled,
+		HTTPHeader:      headers,
+	})
+	if err != nil {
+		statusCode := 0
+		if res != nil {
+			statusCode = res.StatusCode
+			_ = res.Body.Close()
+		}
+		return nil, statusCode
+	}
+	cfg := yamux.DefaultConfig()
+	cfg.LogOutput = io.Discard
+	_, wsNetConn := codersdk.WebsocketNetConn(context.Background(), conn, websocket.MessageBinary)
+	conn.SetReadLimit(drpcsdk.YamuxDefaultStreamWindowSize)
+	session, err := yamux.Client(wsNetConn, cfg)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = session.Close()
+		_ = wsNetConn.Close()
+		_ = conn.Close(websocket.StatusNormalClosure, "")
+	})
+	return session, http.StatusSwitchingProtocols
+}
+
+func TestAIGatewayServeSuccess(t *testing.T) {
+	t.Parallel()
+
+	client, firstUser := coderdenttest.New(t, aibridgeOpts(t))
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	//nolint:gocritic // Owner role is needed for gateway key management.
+	created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-success"})
+	require.NoError(t, err)
+
+	session, status := dialAIGatewayServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String())
+	require.Equal(t, http.StatusSwitchingProtocols, status)
+	require.NotNil(t, session)
+
+	// The Authorizer service should be served and authorize the owner's
+	// session token, exercising a full DRPC round trip over the WebSocket.
+	authorizer := aibridgedproto.NewDRPCAuthorizerClient(drpcsdk.MultiplexedConn(session))
+	resp, err := authorizer.IsAuthorized(ctx, &aibridgedproto.IsAuthorizedRequest{
+		Key: client.SessionToken(),
+	})
+	require.NoError(t, err)
+	require.Equal(t, firstUser.UserID.String(), resp.GetOwnerId())
+
+	// The session records liveness for the authenticating key.
+	require.Eventually(t, func() bool {
+		//nolint:gocritic // Owner role is needed for gateway key management.
+		keys, err := client.ListAIGatewayKeys(ctx)
+		if err != nil {
+			return false
+		}
+		for _, k := range keys {
+			if k.ID == created.ID {
+				return k.LastUsedAt != nil
+			}
+		}
+		return false
+	}, testutil.WaitMedium, testutil.IntervalFast)
+}
+
+func TestAIGatewayServeKeyAndVersionValidationErr(t *testing.T) {
+	t.Parallel()
+
+	client, _ := coderdenttest.New(t, aibridgeOpts(t))
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	//nolint:gocritic // Owner role is needed for gateway key management.
+	created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-quick-failures"})
+	require.NoError(t, err)
+	validKey := created.Key
+
+	//nolint:gocritic // Owner role is needed for gateway key management.
+	revoked, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-revoked"})
+	require.NoError(t, err)
+	require.NoError(t, client.DeleteAIGatewayKey(ctx, revoked.ID))
+
+	tests := []struct {
+		name       string
+		key        string
+		version    string
+		wantStatus int
+	}{
+		{
+			name:       "MissingKey",
+			key:        "",
+			version:    aibridgedproto.CurrentVersion.String(),
+			wantStatus: http.StatusUnauthorized,
+		},
+		{
+			name:       "InvalidKey",
+			key:        "not-a-real-key",
+			version:    aibridgedproto.CurrentVersion.String(),
+			wantStatus: http.StatusUnauthorized,
+		},
+		{
+			name:       "RevokedKey",
+			key:        revoked.Key,
+			version:    aibridgedproto.CurrentVersion.String(),
+			wantStatus: http.StatusUnauthorized,
+		},
+		{
+			name:       "IncompatibleVersion",
+			key:        validKey,
+			version:    "999.0",
+			wantStatus: http.StatusBadRequest,
+		},
+		{
+			name:       "MissingVersion",
+			key:        validKey,
+			version:    "",
+			wantStatus: http.StatusBadRequest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			_, status := dialAIGatewayServe(t.Context(), t, client, tc.key, tc.version)
+			require.Equal(t, tc.wantStatus, status)
+		})
+	}
+}
+
+func TestAIGatewayServeMissingEntitlement(t *testing.T) {
+	t.Parallel()
+
+	// Enable the bridge config but do not grant the FeatureAIBridge license.
+	dv := coderdtest.DeploymentValues(t)
+	dv.AI.BridgeConfig.Enabled = serpent.Bool(true)
+	client, _ := coderdenttest.New(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{DeploymentValues: dv},
+		LicenseOptions: &coderdenttest.LicenseOptions{
+			Features: license.Features{},
+		},
+	})
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	_, status := dialAIGatewayServe(ctx, t, client, "any-key", aibridgedproto.CurrentVersion.String())
+	require.Equal(t, http.StatusForbidden, status)
+}
+
+func TestAIGatewayServeDeletedKeyClosesActiveSession(t *testing.T) {
+	t.Parallel()
+
+	tick := make(chan time.Time, 1)
+	opts := aibridgeOpts(t)
+	opts.Options.NewTicker = func(time.Duration) (<-chan time.Time, func()) {
+		return tick, func() {}
+	}
+
+	client, _ := coderdenttest.New(t, opts)
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	//nolint:gocritic // Owner role is needed for gateway key management.
+	created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-delete-active"})
+	require.NoError(t, err)
+
+	session, status := dialAIGatewayServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String())
+	require.Equal(t, http.StatusSwitchingProtocols, status)
+	require.NotNil(t, session)
+
+	//nolint:gocritic // Owner role is needed for gateway key management.
+	require.NoError(t, client.DeleteAIGatewayKey(ctx, created.ID))
+
+	tick <- time.Now() // trigger aiGatewayTrackKeyUsage.
+	require.Eventually(t, func() bool {
+		select {
+		case <-session.CloseChan():
+			return true
+		default:
+			return false
+		}
+	}, testutil.WaitShort, testutil.IntervalFast)
+}

enterprise/coderd/coderd.go

@@ -322,6 +322,17 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 		})
 	})
 
+	// /ai-gateway/serve provides the DRPC-over-WebSocket that standalone AI Gateway
+	// replicas connect to. It authenticates with a gateway key instead of a user session.
+	api.AGPL.APIHandler.Group(func(r chi.Router) {
+		r.Route("/ai-gateway/serve", func(r chi.Router) {
+			r.Use(
+				api.RequireFeatureMW(codersdk.FeatureAIBridge),
+			)
+			r.Get("/", api.aiGatewayServe)
+		})
+	})
+
 	api.AGPL.APIHandler.Group(func(r chi.Router) {
 		r.Get("/entitlements", api.serveEntitlements)
 		// /regions overrides the AGPL /regions endpoint

enterprise/coderd/provisionerdaemons.go

@@ -313,15 +313,13 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
 		})
 		return
 	}
-	// Align with the frame size of yamux.
-	conn.SetReadLimit(256 * 1024)
-
 	// Multiplexes the incoming connection using yamux.
 	// This allows multiple function calls to occur over
 	// the same connection.
 	config := yamux.DefaultConfig()
 	config.LogOutput = io.Discard
 	ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
+	conn.SetReadLimit(drpcsdk.YamuxDefaultStreamWindowSize)
 	defer wsNetConn.Close()
 	session, err := yamux.Server(wsNetConn, config)
 	if err != nil {

site/src/api/typesGenerated.ts

@@ -266,6 +266,12 @@ export interface AIGatewayKey {
 	readonly last_used_at?: string;
 }
 
+// From codersdk/client.go
+/**
+ * AIGatewayKeyHeader contains the authentication key for a standalone AI Gateway replica.
+ */
+export const AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key";
+
 // From codersdk/aiproviders.go
 /**
  * AIProvider represents an AI provider configuration row as returned