aibridgeHandler handles all aibridged-related endpoints.
(api *API, middlewares ...func(http.Handler) http.Handler)
| 45 | |
| 46 | // aibridgeHandler handles all aibridged-related endpoints. |
| 47 | func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) func(r chi.Router) { |
| 48 | // Build the overload protection middleware chain for the aibridged handler. |
| 49 | // These limits are applied per-replica. |
| 50 | bridgeCfg := api.DeploymentValues.AI.BridgeConfig |
| 51 | concurrencyLimiter := httpmw.ConcurrencyLimit(bridgeCfg.MaxConcurrency.Value(), "AI Bridge") |
| 52 | rateLimiter := httpmw.RateLimitByAuthToken(int(bridgeCfg.RateLimit.Value()), aiBridgeRateLimitWindow) |
| 53 | |
| 54 | return func(r chi.Router) { |
| 55 | r.Use(api.RequireFeatureMW(codersdk.FeatureAIBridge)) |
| 56 | r.Group(func(r chi.Router) { |
| 57 | r.Use(middlewares...) |
| 58 | r.Get("/interceptions", api.aiBridgeListInterceptions) |
| 59 | r.Get("/sessions", api.aiBridgeListSessions) |
| 60 | r.Get("/sessions/{session_id}", api.aiBridgeGetSessionThreads) |
| 61 | r.Get("/models", api.aiBridgeListModels) |
| 62 | r.Get("/clients", api.aiBridgeListClients) |
| 63 | }) |
| 64 | |
| 65 | // Apply overload protection middleware to the aibridged handler. |
| 66 | // Concurrency limit is checked first for faster rejection under load. |
| 67 | r.Group(func(r chi.Router) { |
| 68 | r.Use(concurrencyLimiter, rateLimiter) |
| 69 | // This is a bit funky but since aibridge only exposes a HTTP |
| 70 | // handler, this is how it has to be. |
| 71 | r.HandleFunc("/*", func(rw http.ResponseWriter, r *http.Request) { |
| 72 | if api.AGPL.GetAIBridgedHandler() == nil { |
| 73 | httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{ |
| 74 | Message: "aibridged handler not mounted", |
| 75 | }) |
| 76 | return |
| 77 | } |
| 78 | |
| 79 | // Reject BYOK requests when the deployment has not |
| 80 | // enabled bring-your-own-key mode. |
| 81 | if agplaibridge.IsBYOK(r.Header) && !bridgeCfg.AllowBYOK.Value() { |
| 82 | httpapi.Write(r.Context(), rw, http.StatusForbidden, codersdk.Response{ |
| 83 | Message: "Bring Your Own Key (BYOK) mode is not enabled.", |
| 84 | Detail: "Contact your administrator to enable it with --aibridge-allow-byok.", |
| 85 | }) |
| 86 | return |
| 87 | } |
| 88 | |
| 89 | api.AGPL.GetAIBridgedHandler().ServeHTTP(rw, r) |
| 90 | }) |
| 91 | }) |
| 92 | } |
| 93 | } |
| 94 | |
| 95 | // aiBridgeListInterceptions returns all AI Bridge interceptions a user can read. |
| 96 | // Optional filters with query params. |
no test coverage detected