Skip to content

feat: add /api/v2/ai-gateway API route aliases #26475

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ssncferreira wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: add /api/v2/ai-gateway API route aliases #26475

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 8 additions & 9 deletions coderd/aibridged.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,12 @@ import (
"github.com/coder/coder/v2/codersdk/drpcsdk"
)

// GetAIBridgedHandler returns the in-memory aibridge HTTP handler set by
// [API.RegisterInMemoryAIBridgedHTTPHandler], or nil if the daemon has not
// been wired in. Used by the enterprise /api/v2/aibridge route (license-gated)
// to forward requests into the same in-memory handler that chatd dispatches
// to in-process.
func (api *API) GetAIBridgedHandler() http.Handler {
return api.aibridgedHandler
// GetAIGatewayHandler returns the raw in-memory AI Gateway HTTP handler
// set by [API.RegisterInMemoryAIBridgedHTTPHandler], or nil if the daemon
// has not been wired in. Callers must apply their own [http.StripPrefix]
// for the route prefix they are mounting under.
func (api *API) GetAIGatewayHandler() http.Handler {
return api.aiGatewayHandler
}

// RegisterInMemoryAIBridgedHTTPHandler mounts [aibridged.Server]'s HTTP router onto
Expand All @@ -42,9 +41,9 @@ func (api *API) RegisterInMemoryAIBridgedHTTPHandler(srv http.Handler) {
panic("aibridged cannot be nil")
}

api.aibridgedHandler = http.StripPrefix("/api/v2/aibridge", srv)
api.aiGatewayHandler = srv

factory := aibridged.NewTransportFactory(api.aibridgedHandler)
factory := aibridged.NewTransportFactory(http.StripPrefix("/api/v2/ai-gateway", srv))
var asInterface agplaibridge.TransportFactory = factory
api.AIBridgeTransportFactory.Store(&asInterface)
}
Expand Down
10 changes: 5 additions & 5 deletions coderd/aibridged/transport.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,12 @@ import (
"github.com/coder/coder/v2/coderd/aibridge"
)

// aibridgeRootPath is the URL prefix the in-memory aibridged handler
// aiGatewayRootPath is the URL prefix the in-memory AI Gateway handler
// registers all of its routes under. The in-process round-tripper
// prepends this plus the provider name to every request before
// dispatch so callers can hand it upstream-shaped requests without
// knowing the daemon's mount layout.
const aibridgeRootPath = "/api/v2/aibridge"
const aiGatewayRootPath = "/api/v2/ai-gateway"

// NewTransportFactory returns an [aibridge.TransportFactory] whose RoundTripper
// dispatches requests to handler in-process, streaming the response body
Expand All @@ -37,7 +37,7 @@ type transportFactory struct {
// TransportFor returns an in-process [http.RoundTripper] that dispatches
// requests through the aibridged handler. The provider name is the routing
// key the daemon mounts on; the round-tripper rewrites each request's URL
// path to "/api/v2/aibridge/<providerName>/..." before dispatching so
// path to "/api/v2/ai-gateway/<providerName>/..." before dispatching so
// callers can build upstream-shaped requests and stay agnostic of the
// daemon's mount layout. The source is attached to the request context for
// downstream logging; routing does not depend on it.
Expand Down Expand Up @@ -69,10 +69,10 @@ func (t *inMemoryRoundTripper) RoundTrip(req *http.Request) (*http.Response, err
}

// Adapt the caller's upstream-shaped URL to the daemon's mount layout:
// "/api/v2/aibridge/<providerName>/<original-path>". Done here so
// "/api/v2/ai-gateway/<providerName>/<original-path>". Done here so
// callers do not need to encode the mount prefix or the provider
// routing key into the requests they hand to the transport.
newPath, err := url.JoinPath(aibridgeRootPath, t.providerName, req.URL.Path)
newPath, err := url.JoinPath(aiGatewayRootPath, t.providerName, req.URL.Path)
if err != nil {
return nil, xerrors.Errorf("rewrite request URL for provider %q: %w", t.providerName, err)
}
Expand Down
2 changes: 1 addition & 1 deletion coderd/aibridged/transport_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ func TestTransportFactory_TransportFor(t *testing.T) {
require.NoError(t, err)
defer resp.Body.Close()

require.Equal(t, "/api/v2/aibridge/my-anthropic/v1/messages", <-got)
require.Equal(t, "/api/v2/ai-gateway/my-anthropic/v1/messages", <-got)
require.Equal(t, origPath, req.URL.Path,
"caller's request URL must not be mutated by RoundTrip")
})
Expand Down
11 changes: 6 additions & 5 deletions coderd/coderd.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2229,11 +2229,12 @@ type API struct {
// providers directly. Registered by coderd at startup once aibridged is
// wired in-memory.
AIBridgeTransportFactory atomic.Pointer[aibridge.TransportFactory]
// aibridgedHandler is the in-memory aibridge HTTP handler. Set by
// RegisterInMemoryAIBridgedHTTPHandler; read both by the enterprise
// /api/v2/aibridge route (license-gated) and by the in-memory transport
// (used by chatd, license-exempt).
aibridgedHandler http.Handler
// aiGatewayHandler is the raw in-memory AI Gateway HTTP handler
// (no prefix stripping). Set by RegisterInMemoryAIBridgedHTTPHandler;
// used by the enterprise /api/v2/aibridge and /api/v2/ai-gateway
// routes (license-gated) which apply their own StripPrefix, and by
// the in-memory transport (used by chatd, license-exempt).
aiGatewayHandler http.Handler

UpdatesProvider tailnet.WorkspaceUpdatesProvider

Expand Down
10 changes: 5 additions & 5 deletions enterprise/aibridgeproxyd/aibridgeproxyd.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ type RoundTripDumper interface {
const (
// ProxyAuthRealm is the realm used in Proxy-Authenticate challenges.
// The realm helps clients identify which credentials to use.
ProxyAuthRealm = `"Coder AI Bridge Proxy"`
ProxyAuthRealm = `"Coder AI Gateway Proxy"`
)

// proxyAuthRequiredMsg is the response body for 407 responses.
Expand Down Expand Up @@ -912,7 +912,7 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.
)

resp := goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusProxyAuthRequired, "Proxy authentication required")
resp.Header.Set("Proxy-Authenticate", `Basic realm="Coder AI Bridge Proxy"`)
resp.Header.Set("Proxy-Authenticate", `Basic realm="Coder AI Gateway Proxy"`)
return req, resp
}

Expand Down Expand Up @@ -968,16 +968,16 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Proxy misconfigured")
}

aiBridgeURL, err := url.JoinPath(s.coderAccessURL.String(), "api/v2/aibridge", reqCtx.Provider, originalPath)
aiBridgeURL, err := url.JoinPath(s.coderAccessURL.String(), "api/v2/ai-gateway", reqCtx.Provider, originalPath)
if err != nil {
logger.Error(s.ctx, "failed to build aibridged URL", slog.Error(err))
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to build AI Bridge URL")
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to build AI Gateway URL")
}

aiBridgeParsedURL, err := url.Parse(aiBridgeURL)
if err != nil {
logger.Error(s.ctx, "failed to parse aibridged URL", slog.Error(err))
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to parse AI Bridge URL")
return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusInternalServerError, "Failed to parse AI Gateway URL")
}

// Preserve query parameters from the original request.
Expand Down
32 changes: 23 additions & 9 deletions enterprise/coderd/aibridge.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,24 @@ var errInvalidCursor = xerrors.New("invalid pagination cursor")
// check_constraint.go.
const userAIBudgetOverridesMustBeGroupMemberConstraint database.CheckConstraint = "user_ai_budget_overrides_must_be_group_member"

// aibridgeHandler handles all aibridged-related endpoints.
// aibridgeHandler handles all aibridged-related endpoints under /api/v2/aibridge.
func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) {
return aiGatewayRoutes(api, "/api/v2/aibridge", middlewares...)
}

// aiGatewayHTTPHandler handles all aibridged-related endpoints under /api/v2/ai-gateway.
func aiGatewayHTTPHandler(api *API, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) {
return aiGatewayRoutes(api, "/api/v2/ai-gateway", middlewares...)
}

// aiGatewayRoutes builds the route tree for AI Gateway endpoints.
// The stripPrefix parameter is the URL prefix to strip before forwarding
// to the raw in-memory handler.
func aiGatewayRoutes(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 Bridge")
concurrencyLimiter := httpmw.ConcurrencyLimit(bridgeCfg.MaxConcurrency.Value(), "AI Gateway")
rateLimiter := httpmw.RateLimitByAuthToken(int(bridgeCfg.RateLimit.Value()), aiBridgeRateLimitWindow)

return func(r chi.Router) {
Expand All @@ -72,7 +84,8 @@ func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) f
// 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) {
if api.AGPL.GetAIBridgedHandler() == nil {
raw := api.AGPL.GetAIGatewayHandler()
if raw == nil {
httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{
Message: "aibridged handler not mounted",
})
Expand All @@ -89,7 +102,8 @@ func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) f
return
}

api.AGPL.GetAIBridgedHandler().ServeHTTP(rw, r)
// Strip the prefix and relay to the raw aibridged handler.
http.StripPrefix(stripPrefix, raw).ServeHTTP(rw, r)
})
})
}
Expand Down Expand Up @@ -203,7 +217,7 @@ func (api *API) aiBridgeListSessions(rw http.ResponseWriter, r *http.Request) {
})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error getting AI Bridge sessions.",
Message: "Internal error getting AI Gateway sessions.",
Detail: err.Error(),
})
return
Expand Down Expand Up @@ -446,7 +460,7 @@ func (api *API) aiBridgeListModels(rw http.ResponseWriter, r *http.Request) {

if len(errs) > 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid AI Bridge models search query.",
Message: "Invalid AI Gateway models search query.",
Validations: errs,
})
return
Expand All @@ -455,7 +469,7 @@ func (api *API) aiBridgeListModels(rw http.ResponseWriter, r *http.Request) {
models, err := api.Database.ListAIBridgeModels(ctx, filter)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error getting AI Bridge models.",
Message: "Internal error getting AI Gateway models.",
Detail: err.Error(),
})
return
Expand Down Expand Up @@ -498,7 +512,7 @@ func (api *API) aiBridgeListClients(rw http.ResponseWriter, r *http.Request) {

if len(errs) > 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid AI Bridge clients search query.",
Message: "Invalid AI Gateway clients search query.",
Validations: errs,
})
return
Expand All @@ -507,7 +521,7 @@ func (api *API) aiBridgeListClients(rw http.ResponseWriter, r *http.Request) {
clients, err := api.Database.ListAIBridgeClients(ctx, filter)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error getting AI Bridge clients.",
Message: "Internal error getting AI Gateway clients.",
Detail: err.Error(),
})
return
Expand Down
19 changes: 15 additions & 4 deletions enterprise/coderd/aibridgeproxy.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,19 @@ func (api *API) RegisterInMemoryAIBridgeProxydHTTPHandler(srv http.Handler) {
api.aibridgeproxydHandler = srv
}

// aibridgeproxyHandler handles AI Bridge Proxy endpoints.
// aibridgeproxyHandler handles legacy AI Bridge Proxy endpoints under /api/v2/aibridge/proxy.
// Kept for backward compatibility. New endpoints should be added to /ai-gateway only.
func aibridgeproxyHandler(api *API, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) {
return aiGatewayProxyRoutes(api, "/api/v2/aibridge/proxy", middlewares...)
}

// aiGatewayProxyHTTPHandler handles AI Gateway Proxy endpoints under /api/v2/ai-gateway/proxy.
func aiGatewayProxyHTTPHandler(api *API, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) {
return aiGatewayProxyRoutes(api, "/api/v2/ai-gateway/proxy", middlewares...)
}

// aiGatewayProxyRoutes builds the route tree for AI Gateway Proxy endpoints.
func aiGatewayProxyRoutes(api *API, stripPrefix string, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) {
return func(r chi.Router) {
r.Use(api.RequireFeatureMW(codersdk.FeatureAIBridge))
r.Use(middlewares...)
Expand All @@ -30,21 +41,21 @@ func aibridgeproxyHandler(api *API, middlewares ...func(http.Handler) http.Handl
// Check if the proxy is enabled.
if !api.DeploymentValues.AI.BridgeProxyConfig.Enabled.Value() {
httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{
Message: "AI Bridge Proxy is not enabled.",
Message: "AI Gateway Proxy is not enabled.",
})
return
}

// Check if the handler is registered.
if api.aibridgeproxydHandler == nil {
httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{
Message: "AI Bridge Proxy handler not mounted.",
Message: "AI Gateway Proxy handler not mounted.",
})
return
}

// Strip the prefix and relay to the aibridgeproxyd handler.
http.StripPrefix("/api/v2/aibridge/proxy", api.aibridgeproxydHandler).ServeHTTP(rw, r)
http.StripPrefix(stripPrefix, api.aibridgeproxydHandler).ServeHTTP(rw, r)
})
}
}
13 changes: 12 additions & 1 deletion enterprise/coderd/coderd.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,8 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
return api.refreshEntitlements(ctx)
}

// Legacy aibridge routes: kept for backward compatibility.
// New endpoints should be added to /ai-gateway only.
api.AGPL.APIHandler.Group(func(r chi.Router) {
r.Route("/aibridge", aibridgeHandler(api, apiKeyMiddleware))
})
Expand All @@ -299,8 +301,17 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
r.Route("/aibridge/proxy", aibridgeproxyHandler(api, apiKeyMiddleware))
})

// AI Gateway routes: canonical aliases for the aibridge endpoints.
api.AGPL.APIHandler.Group(func(r chi.Router) {
r.Route("/aibridge/keys", func(r chi.Router) {
r.Route("/ai-gateway", aiGatewayHTTPHandler(api, apiKeyMiddleware))
})

api.AGPL.APIHandler.Group(func(r chi.Router) {
r.Route("/ai-gateway/proxy", aiGatewayProxyHTTPHandler(api, apiKeyMiddleware))
})

api.AGPL.APIHandler.Group(func(r chi.Router) {
r.Route("/ai-gateway/keys", func(r chi.Router) {
r.Use(
apiKeyMiddleware,
api.RequireFeatureMW(codersdk.FeatureAIBridge),
Expand Down
Loading