diff --git a/coderd/aibridged.go b/coderd/aibridged.go index cd97ef54fcf..a088c9a041c 100644 --- a/coderd/aibridged.go +++ b/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{ diff --git a/coderd/aibridged/proto/version.go b/coderd/aibridged/proto/version.go new file mode 100644 index 00000000000..914189515d0 --- /dev/null +++ b/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) diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 1690f8bac7e..469fac22194 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/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 diff --git a/coderd/aibridgedserver/register.go b/coderd/aibridgedserver/register.go new file mode 100644 index 00000000000..09f5a712f01 --- /dev/null +++ b/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 +} diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index bc3c7d32359..634dad9d0bf 100644 --- a/coderd/apidoc/docs.go +++ b/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.", @@ -28908,6 +28927,11 @@ const docTemplate = `{ } }, "securityDefinitions": { + "AIGatewayKey": { + "type": "apiKey", + "name": "X-AI-Governance-Gateway-Key", + "in": "header" + }, "Authorization": { "type": "apiKey", "name": "Authorizaiton", diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 38b3e2d94c6..62b0c4fe183 100644 --- a/coderd/apidoc/swagger.json +++ b/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.", @@ -26677,6 +26694,11 @@ } }, "securityDefinitions": { + "AIGatewayKey": { + "type": "apiKey", + "name": "X-AI-Governance-Gateway-Key", + "in": "header" + }, "Authorization": { "type": "apiKey", "name": "Authorizaiton", diff --git a/coderd/coderd.go b/coderd/coderd.go index 2688b55c033..ae0df0eb88b 100644 --- a/coderd/coderd.go +++ b/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 { diff --git a/coderd/coderdtest/swaggerparser.go b/coderd/coderdtest/swaggerparser.go index dcb65fac8b9..70d973c184d 100644 --- a/coderd/coderdtest/swaggerparser.go +++ b/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) } diff --git a/codersdk/client.go b/codersdk/client.go index b01b5e4fb3a..834dfa465e2 100644 --- a/codersdk/client.go +++ b/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" diff --git a/codersdk/drpcsdk/transport.go b/codersdk/drpcsdk/transport.go index 82a0921b410..8aef96db41c 100644 --- a/codersdk/drpcsdk/transport.go +++ b/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 { diff --git a/codersdk/provisionerdaemons.go b/codersdk/provisionerdaemons.go index 46238d7d484..1dced9ee738 100644 --- a/codersdk/provisionerdaemons.go +++ b/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, "") diff --git a/docs/reference/api/enterprise.md b/docs/reference/api/enterprise.md index a12c3247b24..1bf3cda5c02 100644 --- a/docs/reference/api/enterprise.md +++ b/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 diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go new file mode 100644 index 00000000000..283d4fc5f2d --- /dev/null +++ b/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 + } +} diff --git a/enterprise/coderd/aibridgeserve_test.go b/enterprise/coderd/aibridgeserve_test.go new file mode 100644 index 00000000000..96a37ddb9df --- /dev/null +++ b/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) +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index ac5b357c299..9f3860f3ca4 100644 --- a/enterprise/coderd/coderd.go +++ b/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 diff --git a/enterprise/coderd/provisionerdaemons.go b/enterprise/coderd/provisionerdaemons.go index 17a00d22421..c2ba568b969 100644 --- a/enterprise/coderd/provisionerdaemons.go +++ b/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 { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index aaf8c207d67..e8e3a7e16ab 100644 --- a/site/src/api/typesGenerated.ts +++ b/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