From 2ac829a059272ce4ebe6bfe555240040cdb504ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 17 Jun 2026 16:37:41 +0000 Subject: [PATCH 1/7] feat(coderd/aibridged): add API version --- coderd/aibridged/proto/version.go | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 coderd/aibridged/proto/version.go diff --git a/coderd/aibridged/proto/version.go b/coderd/aibridged/proto/version.go new file mode 100644 index 0000000000000..914189515d06b --- /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) From a05d464ecb7ca3bcb925a4e22a0bbc9da03f375a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 17 Jun 2026 16:58:15 +0000 Subject: [PATCH 2/7] feat(enterprise/coderd): add /api/v2/ai-gateway/serve endpoint --- coderd/aibridged.go | 14 +- coderd/aibridgedserver/aibridgedserver.go | 17 +- coderd/aibridgedserver/register.go | 25 +++ coderd/apidoc/docs.go | 24 +++ coderd/apidoc/swagger.json | 22 ++ coderd/coderd.go | 4 + codersdk/client.go | 4 + codersdk/drpcsdk/transport.go | 4 + codersdk/provisionerdaemons.go | 4 +- docs/reference/api/enterprise.md | 20 ++ enterprise/coderd/aibridgeserve.go | 247 ++++++++++++++++++++++ enterprise/coderd/aibridgeserve_test.go | 220 +++++++++++++++++++ enterprise/coderd/coderd.go | 13 ++ enterprise/coderd/provisionerdaemons.go | 4 +- site/src/api/typesGenerated.ts | 7 + 15 files changed, 604 insertions(+), 25 deletions(-) create mode 100644 coderd/aibridgedserver/register.go create mode 100644 enterprise/coderd/aibridgeserve.go create mode 100644 enterprise/coderd/aibridgeserve_test.go diff --git a/coderd/aibridged.go b/coderd/aibridged.go index cd97ef54fcfaa..a088c9a041c82 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/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 1690f8bac7eb5..2cbd55a5a6fda 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -594,13 +594,16 @@ 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, because the caller +// never holds the secret. 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 an +// operator-grade AI Gateway key, which are fully trusted: a 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, so Gateway keys must be guarded as high-value +// operator secrets. Per-user authorization on this surface is a known gap to +// tighten later. // // 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 0000000000000..09f5a712f01c5 --- /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 bc3c7d3235938..634dad9d0bfbc 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 38b3e2d94c6a4..62b0c4fe1836a 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 2688b55c033b4..ae0df0eb88bbc 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/codersdk/client.go b/codersdk/client.go index b01b5e4fb3a8f..0747edfaa3e24 100644 --- a/codersdk/client.go +++ b/codersdk/client.go @@ -96,6 +96,10 @@ 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 connecting to coderd. + AIGatewayKeyHeader = "X-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 82a0921b41057..8aef96db41c0c 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 46238d7d48478..1dced9ee73848 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 a12c3247b244f..1bf3cda5c0251 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 0000000000000..1c8e23bc978e8 --- /dev/null +++ b/enterprise/coderd/aibridgeserve.go @@ -0,0 +1,247 @@ +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 + +// aiBridgeServe upgrades the connection to a WebSocket and serves the aibridged +// 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) aiBridgeServe(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 // System must look up the AI Gateway key to authenticate the request. + keyID, err := api.Database.GetAIGatewayKeyIDByHashedSecret(dbauthz.AsSystemRestricted(r.Context()), apikey.HashSecret(key)) + if err != nil { + // The lookup is an exact match, missing row means key is invalid. + 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.", + Detail: err.Error(), + }) + return + } + + clientAPIVersion := r.URL.Query().Get("version") + clientCoderVersion := r.Header.Get(codersdk.BuildVersionHeader) + logger := api.Logger.Named("aibridge-serve").With( + slog.F("remote_addr", r.RemoteAddr), + slog.F("ai_gateway_client_api_version", clientAPIVersion), + slog.F("ai_gateway_client_build_version", clientCoderVersion), + slog.F("ai_gateway_server_api_version", aibridgedproto.CurrentVersion.String()), + slog.F("ai_gateway_server_build_version", buildinfo.Version), + slog.F("ai_gateway_key_id", keyID), + ) + + // keyCtx has the lifetime of the authenticated session. + // It is canceled when the request ends or when the + // authenticating key is deleted (from aiGatewayTrackKeyUsage). + // The connection and DRPC server contexts derive from it, + // canceling keyCtx tears down the session. + keyCtx, keyCtxCancel := context.WithCancel(r.Context()) + defer keyCtxCancel() + + // Mark key as used as soon as the request is authenticated. + if _, err := aiGatewayUpdateKeyLastUsed(keyCtx, api, keyID); err != nil { + logger.Warn(keyCtx, "update ai gateway key last used", slog.Error(err)) + } + go aiGatewayTrackKeyUsage(keyCtx, keyCtxCancel, api, keyID, logger) + + 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()}, + }, + }) + 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, "accept aibridge websocket conn", slog.Error(err)) + } + httpapi.Write(keyCtx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Failed to accept websocket connection.", + Detail: err.Error(), + }) + return + } + + // Multiplex the incoming connection using yamux. + 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 + } + + // The served DRPC services run with the aibridged authz subject. Layering it + // onto connCtx keeps srvCtx canceling together with keyCtx. + // nolint:gocritic // A standalone AI Gateway acts as the AI Bridge daemon. + srvCtx := dbauthz.AsAIBridged(connCtx) + + mux := drpcmux.New() + srv, err := aibridgedserver.NewServer( + srvCtx, + 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(srvCtx, "create aibridge server", slog.Error(err)) + } + _ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("create aibridge server: %s", err)) + return + } + if err := aibridgedserver.Register(mux, srv); err != nil { + _ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("register aibridge 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(srvCtx, "drpc server error", slog.Error(err)) + }, + }, + ) + + // Log the request immediately instead of after it completes. + if rl := loggermw.RequestLoggerFromContext(srvCtx); rl != nil { + rl.WriteLog(srvCtx, http.StatusAccepted) + } + + logger.Info(srvCtx, "opened AI Gateway connection") + err = server.Serve(srvCtx, session) + logger.Info(srvCtx, "closed AI Gateway 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, "") +} + +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 true, err + } + return rows > 0, nil +} + +// aiGatewayTrackKeyUsage refreshes last_used_at for keyID on a fixed interval +// until ctx is canceled. The caller records the initial usage; this loop +// handles subsequent refreshes and cancels the session via cancel if the key +// no longer exists. +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 { + 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 + } + if !active { + logger.Info(ctx, "ai gateway key no longer exists, closing connection") + ctxCancel() + return + } + } +} diff --git a/enterprise/coderd/aibridgeserve_test.go b/enterprise/coderd/aibridgeserve_test.go new file mode 100644 index 0000000000000..c6814dfc0f653 --- /dev/null +++ b/enterprise/coderd/aibridgeserve_test.go @@ -0,0 +1,220 @@ +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" +) + +// dialAIBridgeServe 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 dialAIBridgeServe(ctx context.Context, t *testing.T, client *codersdk.Client, key, 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 TestAIBridgeServe(t *testing.T) { + t.Parallel() + + t.Run("Success", func(t *testing.T) { + t.Parallel() + client, firstUser := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant for gateway key management here. + created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-success"}) + require.NoError(t, err) + + session, status := dialAIBridgeServe(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 irrelevant for gateway key management here. + 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) + }) + + t.Run("MissingKey", func(t *testing.T) { + t.Parallel() + client, _ := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + _, status := dialAIBridgeServe(ctx, t, client, "", aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusUnauthorized, status) + }) + + t.Run("InvalidKey", func(t *testing.T) { + t.Parallel() + client, _ := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + _, status := dialAIBridgeServe(ctx, t, client, "not-a-real-key", aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusUnauthorized, status) + }) + + t.Run("RevokedKey", func(t *testing.T) { + t.Parallel() + client, _ := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant for gateway key management here. + created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-revoked"}) + require.NoError(t, err) + //nolint:gocritic // Owner role is irrelevant for gateway key management here. + require.NoError(t, client.DeleteAIGatewayKey(ctx, created.ID)) + + _, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusUnauthorized, status) + }) + + t.Run("IncompatibleVersion", func(t *testing.T) { + t.Parallel() + client, _ := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant for gateway key management here. + created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-badversion"}) + require.NoError(t, err) + + _, status := dialAIBridgeServe(ctx, t, client, created.Key, "999.0") + require.Equal(t, http.StatusBadRequest, status) + }) + + t.Run("MissingVersion", func(t *testing.T) { + t.Parallel() + client, _ := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is irrelevant for gateway key management here. + created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-missing-version"}) + require.NoError(t, err) + + _, status := dialAIBridgeServe(ctx, t, client, created.Key, "") + require.Equal(t, http.StatusBadRequest, status) + }) + + t.Run("DeletedKeyClosesActiveSession", func(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 irrelevant for gateway key management here. + created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-delete-active"}) + require.NoError(t, err) + + session, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusSwitchingProtocols, status) + require.NotNil(t, session) + + //nolint:gocritic // Owner role is irrelevant for gateway key management here. + require.NoError(t, client.DeleteAIGatewayKey(ctx, created.ID)) + + tick <- time.Now() + require.Eventually(t, func() bool { + select { + case <-session.CloseChan(): + return true + default: + return false + } + }, testutil.WaitShort, testutil.IntervalFast) + }) + + t.Run("MissingEntitlement", func(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 := dialAIBridgeServe(ctx, t, client, "any-key", aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusForbidden, status) + }) +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index ac5b357c2992b..b96fa3fe34223 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -322,6 +322,19 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { }) }) + // /ai-gateway/serve is the DRPC-over-WebSocket endpoint that standalone AI + // Gateway replicas connect to. It authenticates with a gateway key instead + // of a user session, and deliberately sits outside the /ai-gateway catch-all + // so the catch-all's overload middleware does not apply to it. + 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.aiBridgeServe) + }) + }) + 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 17a00d22421b1..c2ba568b9691d 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 aaf8c207d676f..9d3f171f25e67 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -266,6 +266,13 @@ export interface AIGatewayKey { readonly last_used_at?: string; } +// From codersdk/client.go +/** + * AIGatewayKeyHeader contains the authentication key for a standalone AI + * Gateway replica connecting to coderd. + */ +export const AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key"; + // From codersdk/aiproviders.go /** * AIProvider represents an AI provider configuration row as returned From b93ad1836ed833ddefea50bdbd7c051dffe6c559 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Thu, 25 Jun 2026 09:41:50 +0000 Subject: [PATCH 3/7] agentic review 2 --- coderd/aibridgedserver/aibridgedserver.go | 17 +++++++---------- enterprise/coderd/aibridgeserve.go | 3 +-- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 2cbd55a5a6fda..3283f609c9a10 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -594,16 +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, because the caller -// never holds the secret. 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 an -// operator-grade AI Gateway key, which are fully trusted: a 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, so Gateway keys must be guarded as high-value -// operator secrets. Per-user authorization on this surface is a known gap to -// tighten later. +// 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 to tighten later. // // 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/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 1c8e23bc978e8..1d689426643a9 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -63,7 +63,6 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { } httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to look up AI Gateway key.", - Detail: err.Error(), }) return } @@ -236,8 +235,8 @@ func aiGatewayTrackKeyUsage(ctx context.Context, ctxCancel context.CancelFunc, a if consecutiveFailures > 1 { logger.Info(ctx, "ai gateway key last used update recovered", slog.F("consecutive_failures", consecutiveFailures)) - consecutiveFailures = 0 } + consecutiveFailures = 0 if !active { logger.Info(ctx, "ai gateway key no longer exists, closing connection") ctxCancel() From b12c2297f8b579ee6f6413178d9b95c863348534 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Thu, 25 Jun 2026 12:10:28 +0000 Subject: [PATCH 4/7] rebase changes --- enterprise/coderd/aibridgeserve.go | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 1d689426643a9..179808cadc870 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -52,7 +52,7 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { } // nolint:gocritic // System must look up the AI Gateway key to authenticate the request. - keyID, err := api.Database.GetAIGatewayKeyIDByHashedSecret(dbauthz.AsSystemRestricted(r.Context()), apikey.HashSecret(key)) + gatewayKey, err := api.Database.GetAIGatewayKeyByHashedSecret(dbauthz.AsSystemRestricted(r.Context()), apikey.HashSecret(key)) if err != nil { // The lookup is an exact match, missing row means key is invalid. if httpapi.Is404Error(err) { @@ -75,7 +75,9 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { slog.F("ai_gateway_client_build_version", clientCoderVersion), slog.F("ai_gateway_server_api_version", aibridgedproto.CurrentVersion.String()), slog.F("ai_gateway_server_build_version", buildinfo.Version), - slog.F("ai_gateway_key_id", keyID), + 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 has the lifetime of the authenticated session. @@ -87,10 +89,10 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { defer keyCtxCancel() // Mark key as used as soon as the request is authenticated. - if _, err := aiGatewayUpdateKeyLastUsed(keyCtx, api, keyID); err != nil { + if _, err := aiGatewayUpdateKeyLastUsed(keyCtx, api, gatewayKey.ID); err != nil { logger.Warn(keyCtx, "update ai gateway key last used", slog.Error(err)) } - go aiGatewayTrackKeyUsage(keyCtx, keyCtxCancel, api, keyID, logger) + go aiGatewayTrackKeyUsage(keyCtx, keyCtxCancel, api, gatewayKey.ID, logger) if err := aibridgedproto.CurrentVersion.Validate(clientAPIVersion); err != nil { httpapi.Write(keyCtx, rw, http.StatusBadRequest, codersdk.Response{ From 5f0d96d1d02e17fb942e39e5d7ed5b85e811085f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Thu, 25 Jun 2026 15:36:03 +0000 Subject: [PATCH 5/7] cleanup --- coderd/aibridgedserver/aibridgedserver.go | 2 +- codersdk/client.go | 3 +- enterprise/coderd/aibridgeserve.go | 32 +-- enterprise/coderd/aibridgeserve_test.go | 257 +++++++++++----------- enterprise/coderd/coderd.go | 6 +- site/src/api/typesGenerated.ts | 3 +- 6 files changed, 146 insertions(+), 157 deletions(-) diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 3283f609c9a10..469fac22194f0 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -600,7 +600,7 @@ externalAuthLoop: // 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 to tighten later. +// 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/codersdk/client.go b/codersdk/client.go index 0747edfaa3e24..66ddcf33c5d46 100644 --- a/codersdk/client.go +++ b/codersdk/client.go @@ -96,8 +96,7 @@ 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 connecting to coderd. + // AIGatewayKeyHeader contains the authentication key for a standalone AI Gateway replica. AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key" // BuildVersionHeader contains build information of Coder. diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 179808cadc870..b56bc8928f33f 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -51,10 +51,9 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { return } - // nolint:gocritic // System must look up the AI Gateway key to authenticate the request. + // 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 { - // The lookup is an exact match, missing row means key is invalid. if httpapi.Is404Error(err) { httpapi.Write(r.Context(), rw, http.StatusUnauthorized, codersdk.Response{ Message: "AI Gateway key invalid.", @@ -83,7 +82,7 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { // keyCtx has the lifetime of the authenticated session. // It is canceled when the request ends or when the // authenticating key is deleted (from aiGatewayTrackKeyUsage). - // The connection and DRPC server contexts derive from it, + // The connection / DRPC server context derive from it, // canceling keyCtx tears down the session. keyCtx, keyCtxCancel := context.WithCancel(r.Context()) defer keyCtxCancel() @@ -125,7 +124,6 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { return } - // Multiplex the incoming connection using yamux. config := yamux.DefaultConfig() config.LogOutput = io.Discard connCtx, wsNetConn := codersdk.WebsocketNetConn(keyCtx, conn, websocket.MessageBinary) @@ -137,14 +135,9 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { return } - // The served DRPC services run with the aibridged authz subject. Layering it - // onto connCtx keeps srvCtx canceling together with keyCtx. - // nolint:gocritic // A standalone AI Gateway acts as the AI Bridge daemon. - srvCtx := dbauthz.AsAIBridged(connCtx) - mux := drpcmux.New() srv, err := aibridgedserver.NewServer( - srvCtx, + connCtx, api.Database, logger, api.AccessURL.String(), @@ -155,7 +148,7 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { ) if err != nil { if !xerrors.Is(err, context.Canceled) { - logger.Error(srvCtx, "create aibridge server", slog.Error(err)) + logger.Error(connCtx, "create aibridge server", slog.Error(err)) } _ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("create aibridge server: %s", err)) return @@ -172,19 +165,19 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { if xerrors.Is(err, io.EOF) { return } - logger.Debug(srvCtx, "drpc server error", slog.Error(err)) + logger.Debug(connCtx, "drpc server error", slog.Error(err)) }, }, ) // Log the request immediately instead of after it completes. - if rl := loggermw.RequestLoggerFromContext(srvCtx); rl != nil { - rl.WriteLog(srvCtx, http.StatusAccepted) + if rl := loggermw.RequestLoggerFromContext(connCtx); rl != nil { + rl.WriteLog(connCtx, http.StatusAccepted) } - logger.Info(srvCtx, "opened AI Gateway connection") - err = server.Serve(srvCtx, session) - logger.Info(srvCtx, "closed AI Gateway connection", slog.Error(err)) + logger.Info(connCtx, "opened AI Gateway connection") + err = server.Serve(connCtx, session) + logger.Info(connCtx, "closed AI Gateway connection", slog.Error(err)) if err != nil && !xerrors.Is(err, io.EOF) { _ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("serve: %s", err)) return @@ -201,10 +194,7 @@ func aiGatewayUpdateKeyLastUsed(ctx context.Context, api *API, keyID uuid.UUID) return rows > 0, nil } -// aiGatewayTrackKeyUsage refreshes last_used_at for keyID on a fixed interval -// until ctx is canceled. The caller records the initial usage; this loop -// handles subsequent refreshes and cancels the session via cancel if the key -// no longer exists. +// 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() diff --git a/enterprise/coderd/aibridgeserve_test.go b/enterprise/coderd/aibridgeserve_test.go index c6814dfc0f653..e851a2a96535c 100644 --- a/enterprise/coderd/aibridgeserve_test.go +++ b/enterprise/coderd/aibridgeserve_test.go @@ -25,7 +25,7 @@ import ( // 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 dialAIBridgeServe(ctx context.Context, t *testing.T, client *codersdk.Client, key, version string) (*yamux.Session, int) { +func dialAIBridgeServe(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") @@ -68,153 +68,156 @@ func dialAIBridgeServe(ctx context.Context, t *testing.T, client *codersdk.Clien return session, http.StatusSwitchingProtocols } -func TestAIBridgeServe(t *testing.T) { +func TestAIBridgeServeSuccess(t *testing.T) { t.Parallel() - t.Run("Success", func(t *testing.T) { - t.Parallel() - client, firstUser := coderdenttest.New(t, aibridgeOpts(t)) - ctx := testutil.Context(t, testutil.WaitLong) + client, firstUser := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-success"}) - require.NoError(t, err) - - session, status := dialAIBridgeServe(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 irrelevant for gateway key management here. - 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) - }) - - t.Run("MissingKey", func(t *testing.T) { - t.Parallel() - client, _ := coderdenttest.New(t, aibridgeOpts(t)) - ctx := testutil.Context(t, testutil.WaitLong) - - _, status := dialAIBridgeServe(ctx, t, client, "", aibridgedproto.CurrentVersion.String()) - require.Equal(t, http.StatusUnauthorized, status) - }) + //nolint:gocritic // Owner role is needed for gateway key management. + created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-success"}) + require.NoError(t, err) - t.Run("InvalidKey", func(t *testing.T) { - t.Parallel() - client, _ := coderdenttest.New(t, aibridgeOpts(t)) - ctx := testutil.Context(t, testutil.WaitLong) + session, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusSwitchingProtocols, status) + require.NotNil(t, session) - _, status := dialAIBridgeServe(ctx, t, client, "not-a-real-key", aibridgedproto.CurrentVersion.String()) - require.Equal(t, http.StatusUnauthorized, status) + // 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()) - t.Run("RevokedKey", func(t *testing.T) { - t.Parallel() - client, _ := coderdenttest.New(t, aibridgeOpts(t)) - ctx := testutil.Context(t, testutil.WaitLong) - - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-revoked"}) - require.NoError(t, err) - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - require.NoError(t, client.DeleteAIGatewayKey(ctx, created.ID)) + // 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) +} - _, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) - require.Equal(t, http.StatusUnauthorized, status) - }) +func TestAIBridgeServeKeyAndVersionValidationErr(t *testing.T) { + t.Parallel() - t.Run("IncompatibleVersion", func(t *testing.T) { - t.Parallel() - client, _ := coderdenttest.New(t, aibridgeOpts(t)) - ctx := testutil.Context(t, testutil.WaitLong) + client, _ := coderdenttest.New(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-badversion"}) - require.NoError(t, err) + //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 - _, status := dialAIBridgeServe(ctx, t, client, created.Key, "999.0") - require.Equal(t, http.StatusBadRequest, status) - }) + //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, + }, + } - t.Run("MissingVersion", func(t *testing.T) { - t.Parallel() - client, _ := coderdenttest.New(t, aibridgeOpts(t)) - ctx := testutil.Context(t, testutil.WaitLong) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + _, status := dialAIBridgeServe(t.Context(), t, client, tc.key, tc.version) + require.Equal(t, tc.wantStatus, status) + }) + } +} - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-missing-version"}) - require.NoError(t, err) +func TestAIBridgeServeMissingEntitlement(t *testing.T) { + t.Parallel() - _, status := dialAIBridgeServe(ctx, t, client, created.Key, "") - require.Equal(t, http.StatusBadRequest, status) + // 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) - t.Run("DeletedKeyClosesActiveSession", func(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() {} - } + _, status := dialAIBridgeServe(ctx, t, client, "any-key", aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusForbidden, status) +} - client, _ := coderdenttest.New(t, opts) - ctx := testutil.Context(t, testutil.WaitLong) +func TestAIBridgeServeDeletedKeyClosesActiveSession(t *testing.T) { + t.Parallel() - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-delete-active"}) - require.NoError(t, err) + tick := make(chan time.Time, 1) + opts := aibridgeOpts(t) + opts.Options.NewTicker = func(time.Duration) (<-chan time.Time, func()) { + return tick, func() {} + } - session, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) - require.Equal(t, http.StatusSwitchingProtocols, status) - require.NotNil(t, session) + client, _ := coderdenttest.New(t, opts) + ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // Owner role is irrelevant for gateway key management here. - require.NoError(t, client.DeleteAIGatewayKey(ctx, created.ID)) + //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) - tick <- time.Now() - require.Eventually(t, func() bool { - select { - case <-session.CloseChan(): - return true - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast) - }) + session, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) + require.Equal(t, http.StatusSwitchingProtocols, status) + require.NotNil(t, session) - t.Run("MissingEntitlement", func(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) + //nolint:gocritic // Owner role is needed for gateway key management. + require.NoError(t, client.DeleteAIGatewayKey(ctx, created.ID)) - _, status := dialAIBridgeServe(ctx, t, client, "any-key", aibridgedproto.CurrentVersion.String()) - require.Equal(t, http.StatusForbidden, status) - }) + 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 b96fa3fe34223..3d5f372b6b84c 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -322,10 +322,8 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { }) }) - // /ai-gateway/serve is the DRPC-over-WebSocket endpoint that standalone AI - // Gateway replicas connect to. It authenticates with a gateway key instead - // of a user session, and deliberately sits outside the /ai-gateway catch-all - // so the catch-all's overload middleware does not apply to it. + // /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( diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 9d3f171f25e67..570b59ab40bad 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -268,8 +268,7 @@ export interface AIGatewayKey { // From codersdk/client.go /** - * AIGatewayKeyHeader contains the authentication key for a standalone AI - * Gateway replica connecting to coderd. + * AIGatewayKeyHeader contains the authentication key for a standalone AI Gateway replica. */ export const AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key"; From 26bf32f988986a597e7a6092d7aa0e14c77b4329 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Fri, 26 Jun 2026 12:45:59 +0000 Subject: [PATCH 6/7] review 1: fix header, logger, naming, update last_used_at after session established. --- codersdk/client.go | 2 +- enterprise/coderd/aibridgeserve.go | 69 +++++++++++++------------ enterprise/coderd/aibridgeserve_test.go | 20 +++---- enterprise/coderd/coderd.go | 2 +- site/src/api/typesGenerated.ts | 2 +- 5 files changed, 48 insertions(+), 47 deletions(-) diff --git a/codersdk/client.go b/codersdk/client.go index 66ddcf33c5d46..834dfa465e253 100644 --- a/codersdk/client.go +++ b/codersdk/client.go @@ -97,7 +97,7 @@ const ( ProvisionerDaemonKey = "Coder-Provisioner-Daemon-Key" // AIGatewayKeyHeader contains the authentication key for a standalone AI Gateway replica. - AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key" + AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key" // BuildVersionHeader contains build information of Coder. BuildVersionHeader = "X-Coder-Build-Version" diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index b56bc8928f33f..283d4fc5f2dbe 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -30,11 +30,11 @@ import ( // last_used_at for its authenticating key. const aiGatewayKeyLastUsedInterval = 60 * time.Second -// aiBridgeServe upgrades the connection to a WebSocket and serves the aibridged -// 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. +// 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 @@ -42,7 +42,7 @@ const aiGatewayKeyLastUsedInterval = 60 * time.Second // @Tags Enterprise // @Success 101 // @Router /api/v2/ai-gateway/serve [get] -func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { +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{ @@ -68,36 +68,29 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { clientAPIVersion := r.URL.Query().Get("version") clientCoderVersion := r.Header.Get(codersdk.BuildVersionHeader) - logger := api.Logger.Named("aibridge-serve").With( + logger := api.Logger.Named("aigateway-serve").With( slog.F("remote_addr", r.RemoteAddr), - slog.F("ai_gateway_client_api_version", clientAPIVersion), - slog.F("ai_gateway_client_build_version", clientCoderVersion), - slog.F("ai_gateway_server_api_version", aibridgedproto.CurrentVersion.String()), - slog.F("ai_gateway_server_build_version", buildinfo.Version), + 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 has the lifetime of the authenticated session. - // It is canceled when the request ends or when the - // authenticating key is deleted (from aiGatewayTrackKeyUsage). - // The connection / DRPC server context derive from it, - // canceling keyCtx tears down the session. + // 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() - // Mark key as used as soon as the request is authenticated. - if _, err := aiGatewayUpdateKeyLastUsed(keyCtx, api, gatewayKey.ID); err != nil { - logger.Warn(keyCtx, "update ai gateway key last used", slog.Error(err)) - } - go aiGatewayTrackKeyUsage(keyCtx, keyCtxCancel, api, gatewayKey.ID, logger) - 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 @@ -115,7 +108,7 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { }) if err != nil { if !xerrors.Is(err, context.Canceled) { - logger.Error(keyCtx, "accept aibridge websocket conn", slog.Error(err)) + logger.Error(keyCtx, "websocket upgrade failed", slog.Error(err)) } httpapi.Write(keyCtx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Failed to accept websocket connection.", @@ -135,6 +128,11 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { 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, @@ -148,13 +146,13 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { ) if err != nil { if !xerrors.Is(err, context.Canceled) { - logger.Error(connCtx, "create aibridge server", slog.Error(err)) + logger.Error(connCtx, "server creation failed", slog.Error(err)) } - _ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("create aibridge server: %s", 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 aibridge services: %s", err)) + _ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("register ai gateway services: %s", err)) return } @@ -175,9 +173,9 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { rl.WriteLog(connCtx, http.StatusAccepted) } - logger.Info(connCtx, "opened AI Gateway connection") + logger.Info(connCtx, "opened connection") err = server.Serve(connCtx, session) - logger.Info(connCtx, "closed AI Gateway connection", slog.Error(err)) + 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 @@ -185,11 +183,13 @@ func (api *API) aiBridgeServe(rw http.ResponseWriter, r *http.Request) { _ = 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 true, err + return false, err } return rows > 0, nil } @@ -208,6 +208,12 @@ func aiGatewayTrackKeyUsage(ctx context.Context, ctxCancel context.CancelFunc, a } 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 @@ -229,10 +235,5 @@ func aiGatewayTrackKeyUsage(ctx context.Context, ctxCancel context.CancelFunc, a slog.F("consecutive_failures", consecutiveFailures)) } consecutiveFailures = 0 - if !active { - logger.Info(ctx, "ai gateway key no longer exists, closing connection") - ctxCancel() - return - } } } diff --git a/enterprise/coderd/aibridgeserve_test.go b/enterprise/coderd/aibridgeserve_test.go index e851a2a96535c..96a37ddb9df5e 100644 --- a/enterprise/coderd/aibridgeserve_test.go +++ b/enterprise/coderd/aibridgeserve_test.go @@ -21,11 +21,11 @@ import ( "github.com/coder/websocket" ) -// dialAIBridgeServe dials /api/v2/ai-gateway/serve, authenticating with the given +// 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 dialAIBridgeServe(ctx context.Context, t *testing.T, client *codersdk.Client, key string, version string) (*yamux.Session, int) { +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") @@ -68,7 +68,7 @@ func dialAIBridgeServe(ctx context.Context, t *testing.T, client *codersdk.Clien return session, http.StatusSwitchingProtocols } -func TestAIBridgeServeSuccess(t *testing.T) { +func TestAIGatewayServeSuccess(t *testing.T) { t.Parallel() client, firstUser := coderdenttest.New(t, aibridgeOpts(t)) @@ -78,7 +78,7 @@ func TestAIBridgeServeSuccess(t *testing.T) { created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-success"}) require.NoError(t, err) - session, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) + session, status := dialAIGatewayServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) require.Equal(t, http.StatusSwitchingProtocols, status) require.NotNil(t, session) @@ -107,7 +107,7 @@ func TestAIBridgeServeSuccess(t *testing.T) { }, testutil.WaitMedium, testutil.IntervalFast) } -func TestAIBridgeServeKeyAndVersionValidationErr(t *testing.T) { +func TestAIGatewayServeKeyAndVersionValidationErr(t *testing.T) { t.Parallel() client, _ := coderdenttest.New(t, aibridgeOpts(t)) @@ -164,13 +164,13 @@ func TestAIBridgeServeKeyAndVersionValidationErr(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - _, status := dialAIBridgeServe(t.Context(), t, client, tc.key, tc.version) + _, status := dialAIGatewayServe(t.Context(), t, client, tc.key, tc.version) require.Equal(t, tc.wantStatus, status) }) } } -func TestAIBridgeServeMissingEntitlement(t *testing.T) { +func TestAIGatewayServeMissingEntitlement(t *testing.T) { t.Parallel() // Enable the bridge config but do not grant the FeatureAIBridge license. @@ -184,11 +184,11 @@ func TestAIBridgeServeMissingEntitlement(t *testing.T) { }) ctx := testutil.Context(t, testutil.WaitLong) - _, status := dialAIBridgeServe(ctx, t, client, "any-key", aibridgedproto.CurrentVersion.String()) + _, status := dialAIGatewayServe(ctx, t, client, "any-key", aibridgedproto.CurrentVersion.String()) require.Equal(t, http.StatusForbidden, status) } -func TestAIBridgeServeDeletedKeyClosesActiveSession(t *testing.T) { +func TestAIGatewayServeDeletedKeyClosesActiveSession(t *testing.T) { t.Parallel() tick := make(chan time.Time, 1) @@ -204,7 +204,7 @@ func TestAIBridgeServeDeletedKeyClosesActiveSession(t *testing.T) { created, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "serve-delete-active"}) require.NoError(t, err) - session, status := dialAIBridgeServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) + session, status := dialAIGatewayServe(ctx, t, client, created.Key, aibridgedproto.CurrentVersion.String()) require.Equal(t, http.StatusSwitchingProtocols, status) require.NotNil(t, session) diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 3d5f372b6b84c..9f3860f3ca4b7 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -329,7 +329,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Use( api.RequireFeatureMW(codersdk.FeatureAIBridge), ) - r.Get("/", api.aiBridgeServe) + r.Get("/", api.aiGatewayServe) }) }) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 570b59ab40bad..e8e3a7e16ab26 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -270,7 +270,7 @@ export interface AIGatewayKey { /** * AIGatewayKeyHeader contains the authentication key for a standalone AI Gateway replica. */ -export const AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key"; +export const AIGatewayKeyHeader = "X-Coder-AI-Governance-Gateway-Key"; // From codersdk/aiproviders.go /** From 14c2818c55f1297a6b736dbe4f3fe41752e78e05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Fri, 26 Jun 2026 15:11:04 +0000 Subject: [PATCH 7/7] test fix --- coderd/coderdtest/swaggerparser.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/coderd/coderdtest/swaggerparser.go b/coderd/coderdtest/swaggerparser.go index dcb65fac8b9f4..70d973c184d27 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) }