feat(coderd): add overload protection with rate limiting and concurrency control (#21161)

## Summary This adds configurable overload protection to the AI Bridge daemon to prevent the server from being overwhelmed during periods of high load. Partially addresses coder/internal#1153 (rate limits and concurrency control; circuit breakers are deferred to a follow-up). ## New Configuration Options | Option | Environment Variable | Description | Default | |--------|---------------------|-------------|---------| | `--aibridge-max-concurrency` | `CODER_AIBRIDGE_MAX_CONCURRENCY` | Maximum number of concurrent AI Bridge requests. Set to 0 to disable (unlimited). | `0` | | `--aibridge-rate-limit` | `CODER_AIBRIDGE_RATE_LIMIT` | Maximum number of AI Bridge requests per second. Set to 0 to disable rate limiting. | `0` | ## Behavior When limits are exceeded: - **Concurrency limit**: Returns HTTP `503 Service Unavailable` with message "AI Bridge is currently at capacity. Please try again later." - **Rate limit**: Returns HTTP `429 Too Many Requests` with `Retry-After` header. Both protections are optional and disabled by default (0 values). ## Implementation The overload protection is implemented as reusable middleware in `coderd/httpmw/ratelimit.go`: 1. **`RateLimitByAuthToken`**: Per-user rate limiting that uses `APITokenFromRequest` to extract the authentication token, with fallback to `X-Api-Key` header for AI provider compatibility (e.g., Anthropic). Falls back to IP-based rate limiting if no token is present. Includes `Retry-After` header for backpressure signaling. 2. **`ConcurrencyLimit`**: Uses an atomic counter to track in-flight requests and reject when at capacity. The middleware is applied in `enterprise/coderd/aibridge.go` via `r.Group` in the following order: 1. Concurrency check (faster rejection for load shedding) 2. Rate limit check **Note**: Rate limiting currently applies to all AI Bridge requests, including pass-through requests. Ideally only actual interceptions should count, but this would require changes in the aibridge library. ## Testing Added comprehensive tests for: - Rate limiting by auth token (Bearer token, X-Api-Key, no token fallback to IP) - Different tokens not rate limited against each other - Disabled when limit is zero - Retry-After header is set on 429 responses - Concurrency limiting (allows within limit, rejects over limit, disabled when zero)
2026-09-24 15:04:27 +08:00 · 2025-12-11 16:38:54 +01:00
parent 8ead6f795d
commit 6f86f67754
17 changed files with 568 additions and 37 deletions
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"time"

 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
@@ -23,10 +24,19 @@ import (
 const (
 	maxListInterceptionsLimit     = 1000
 	defaultListInterceptionsLimit = 100
+	// aiBridgeRateLimitWindow is the fixed duration for rate limiting AI Bridge
+	// requests. This is hardcoded to keep configuration simple.
+	aiBridgeRateLimitWindow = time.Second
 )

 // aibridgeHandler handles all aibridged-related endpoints.
 func aibridgeHandler(api *API, 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")
+	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) {
@@ -34,25 +44,30 @@ func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) f
 			r.Get("/interceptions", api.aiBridgeListInterceptions)
 		})

-		// 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.aibridgedHandler == nil {
-				httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{
-					Message: "aibridged handler not mounted",
-				})
-				return
-			}
-
-			// Strip either the experimental or stable prefix.
-			// TODO: experimental route is deprecated and must be removed with Beta.
-			prefixes := []string{"/api/experimental/aibridge", "/api/v2/aibridge"}
-			for _, prefix := range prefixes {
-				if strings.Contains(r.URL.String(), prefix) {
-					http.StripPrefix(prefix, api.aibridgedHandler).ServeHTTP(rw, r)
-					break
+		// Apply overload protection middleware to the aibridged handler.
+		// Concurrency limit is checked first for faster rejection under load.
+		r.Group(func(r chi.Router) {
+			r.Use(concurrencyLimiter, rateLimiter)
+			// 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.aibridgedHandler == nil {
+					httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{
+						Message: "aibridged handler not mounted",
+					})
+					return
 				}
-			}
+
+				// Strip either the experimental or stable prefix.
+				// TODO: experimental route is deprecated and must be removed with Beta.
+				prefixes := []string{"/api/experimental/aibridge", "/api/v2/aibridge"}
+				for _, prefix := range prefixes {
+					if strings.Contains(r.URL.String(), prefix) {
+						http.StripPrefix(prefix, api.aibridgedHandler).ServeHTTP(rw, r)
+						break
+					}
+				}
+			})
 		})
 	}
 }
@@ -703,3 +703,135 @@ func TestAIBridgeRouting(t *testing.T) {
 		})
 	}
 }
+
+func TestAIBridgeRateLimiting(t *testing.T) {
+	t.Parallel()
+
+	dv := coderdtest.DeploymentValues(t)
+	// Set a low rate limit for testing.
+	dv.AI.BridgeConfig.RateLimit = 2
+
+	client, closer, api, _ := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{
+			DeploymentValues: dv,
+		},
+		LicenseOptions: &coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureAIBridge: 1,
+			},
+		},
+	})
+	t.Cleanup(func() {
+		_ = closer.Close()
+	})
+
+	// Register a simple test handler.
+	testHandler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		rw.WriteHeader(http.StatusOK)
+	})
+	api.RegisterInMemoryAIBridgedHTTPHandler(testHandler)
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	httpClient := &http.Client{}
+	url := client.URL.String() + "/api/v2/aibridge/test"
+
+	// Make requests up to the limit - should succeed.
+	for range 2 {
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
+		require.NoError(t, err)
+		req.Header.Set(codersdk.SessionTokenHeader, client.SessionToken())
+
+		resp, err := httpClient.Do(req)
+		require.NoError(t, err)
+		_ = resp.Body.Close()
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+	}
+
+	// Next request should be rate limited.
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
+	require.NoError(t, err)
+	req.Header.Set(codersdk.SessionTokenHeader, client.SessionToken())
+
+	resp, err := httpClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusTooManyRequests, resp.StatusCode)
+	require.NotEmpty(t, resp.Header.Get("Retry-After"))
+}
+
+func TestAIBridgeConcurrencyLimiting(t *testing.T) {
+	t.Parallel()
+
+	dv := coderdtest.DeploymentValues(t)
+	// Set a low concurrency limit for testing.
+	dv.AI.BridgeConfig.MaxConcurrency = 1
+
+	client, closer, api, _ := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{
+			DeploymentValues: dv,
+		},
+		LicenseOptions: &coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureAIBridge: 1,
+			},
+		},
+	})
+	t.Cleanup(func() {
+		_ = closer.Close()
+	})
+
+	// Register a handler that blocks until signaled.
+	started := make(chan struct{})
+	unblock := make(chan struct{})
+	testHandler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		started <- struct{}{}
+		<-unblock
+		rw.WriteHeader(http.StatusOK)
+	})
+	api.RegisterInMemoryAIBridgedHTTPHandler(testHandler)
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	httpClient := &http.Client{}
+	url := client.URL.String() + "/api/v2/aibridge/test"
+
+	// Start a request that will block.
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
+		if err != nil {
+			return
+		}
+		req.Header.Set(codersdk.SessionTokenHeader, client.SessionToken())
+
+		resp, err := httpClient.Do(req)
+		if err == nil {
+			_ = resp.Body.Close()
+		}
+	}()
+
+	// Wait for the first request to start processing.
+	select {
+	case <-started:
+	case <-ctx.Done():
+		t.Fatal("timed out waiting for first request to start")
+	}
+
+	// Second request should be rejected with 503.
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
+	require.NoError(t, err)
+	req.Header.Set(codersdk.SessionTokenHeader, client.SessionToken())
+
+	resp, err := httpClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
+
+	// Unblock the first request and wait for it to complete.
+	close(unblock)
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Fatal("timed out waiting for first request to complete")
+	}
+}