opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
> AI Tools were used to produce this PR

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

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

Some wiring used by this new command is added.

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

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

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

---------

Co-authored-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Paweł Banaszewski
2026-07-08 11:12:53 +02:00
committed by GitHub
co-authored by Danny Kopping
parent 195dffc651
commit ccba3969ab
18 changed files with 1418 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/aibridge.go
+30 -19
View File
@@ -72,12 +72,6 @@ func aiGatewayHTTPHandler(api *API, middlewares ...func(http.Handler) http.Handl
// under /aibridge. The stripPrefix parameter selects which URL prefix
// to strip before forwarding to the in-memory aibridged handler.
func aiBridgeRoutes(api *API, stripPrefix string, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) {
// Build the overload protection middleware chain for the aibridged handler.
// These limits are applied per-replica.
bridgeCfg := api.DeploymentValues.AI.BridgeConfig
concurrencyLimiter := httpmw.ConcurrencyLimit(bridgeCfg.MaxConcurrency.Value(), "AI Gateway")
rateLimiter := httpmw.RateLimitByAuthToken(int(bridgeCfg.RateLimit.Value()), aiBridgeRateLimitWindow)
return func(r chi.Router) {
r.Use(api.RequireFeatureMW(codersdk.FeatureAIBridge))
r.Group(func(r chi.Router) {
@@ -88,10 +82,10 @@ func aiBridgeRoutes(api *API, stripPrefix string, middlewares ...func(http.Handl
r.Get("/clients", api.aiBridgeListClients)
})
// Apply overload protection middleware to the aibridged handler.
// Concurrency limit is checked first for faster rejection under load.
// Apply the shared per-request data-plane middleware (per-replica
// overload protection plus BYOK gating) to the aibridged handler.
r.Group(func(r chi.Router) {
r.Use(concurrencyLimiter, rateLimiter)
r.Use(AIGatewayDataPlaneMiddleware(api.DeploymentValues.AI.BridgeConfig))
// This is a bit funky but since aibridge only exposes a HTTP
// handler, this is how it has to be.
r.HandleFunc("/*", func(rw http.ResponseWriter, r *http.Request) {
@@ -103,16 +97,6 @@ func aiBridgeRoutes(api *API, stripPrefix string, middlewares ...func(http.Handl
return
}
// Reject BYOK requests when the deployment has not
// enabled bring-your-own-key mode.
if agplaibridge.IsBYOK(r.Header) && !bridgeCfg.AllowBYOK.Value() {
httpapi.Write(r.Context(), rw, http.StatusForbidden, codersdk.Response{
Message: "Bring Your Own Key (BYOK) mode is not enabled.",
Detail: "Contact your administrator to enable it with --aibridge-allow-byok.",
})
return
}
// Strip the prefix and relay to the aibridged handler.
http.StripPrefix(stripPrefix, handler).ServeHTTP(rw, r)
})
@@ -120,6 +104,33 @@ func aiBridgeRoutes(api *API, stripPrefix string, middlewares ...func(http.Handl
}
}
// AIGatewayDataPlaneMiddleware returns the per-request middleware chain that
// guards the AI Gateway data-plane handler. It is the single source of truth
// shared by the embedded route and the standalone gateway.
func AIGatewayDataPlaneMiddleware(cfg codersdk.AIBridgeConfig) func(http.Handler) http.Handler {
concurrencyLimiter := httpmw.ConcurrencyLimit(cfg.MaxConcurrency.Value(), "AI Gateway")
rateLimiter := httpmw.RateLimitByAuthToken(int(cfg.RateLimit.Value()), aiBridgeRateLimitWindow)
byokGuard := aiGatewayBYOKGuard(cfg)
return func(next http.Handler) http.Handler {
return concurrencyLimiter(rateLimiter(byokGuard(next)))
}
}
func aiGatewayBYOKGuard(cfg codersdk.AIBridgeConfig) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
if agplaibridge.IsBYOK(r.Header) && !cfg.AllowBYOK.Value() {
httpapi.Write(r.Context(), rw, http.StatusForbidden, codersdk.Response{
Message: "Bring Your Own Key (BYOK) mode is not enabled.",
Detail: "Contact your administrator to enable it with --ai-gateway-allow-byok.",
})
return
}
next.ServeHTTP(rw, r)
})
}
}
// aiBridgeListSessions returns AI Bridge sessions (aggregated interceptions).
//
// @Summary List AI Gateway sessions