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)

coderd/aibridge/aibridge.go

@@ -0,0 +1,24 @@
+// Package aibridge provides utilities for the AI Bridge feature.
+package aibridge
+
+import (
+	"net/http"
+	"strings"
+)
+
+// ExtractAuthToken extracts an authorization token from HTTP headers.
+// It checks the Authorization header (Bearer token) and X-Api-Key header,
+// which represent the different ways clients authenticate against AI providers.
+// If neither are present, an empty string is returned.
+func ExtractAuthToken(header http.Header) string {
+	if auth := strings.TrimSpace(header.Get("Authorization")); auth != "" {
+		fields := strings.Fields(auth)
+		if len(fields) == 2 && strings.EqualFold(fields[0], "Bearer") {
+			return fields[1]
+		}
+	}
+	if apiKey := strings.TrimSpace(header.Get("X-Api-Key")); apiKey != "" {
+		return apiKey
+	}
+	return ""
+}

coderd/apidoc/docs.go

@@ -11883,9 +11883,15 @@ const docTemplate = `{
                 "inject_coder_mcp_tools": {
                     "type": "boolean"
                 },
+                "max_concurrency": {
+                    "type": "integer"
+                },
                 "openai": {
                     "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig"
                 },
+                "rate_limit": {
+                    "type": "integer"
+                },
                 "retention": {
                     "type": "integer"
                 }

coderd/apidoc/swagger.json

@@ -10549,9 +10549,15 @@
 				"inject_coder_mcp_tools": {
 					"type": "boolean"
 				},
+				"max_concurrency": {
+					"type": "integer"
+				},
 				"openai": {
 					"$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig"
 				},
+				"rate_limit": {
+					"type": "integer"
+				},
 				"retention": {
 					"type": "integer"
 				}

coderd/httpmw/ratelimit.go

@@ -4,11 +4,13 @@ import (
 	"fmt"
 	"net/http"
 	"strconv"
+	"sync/atomic"
 	"time"
 
 	"github.com/go-chi/httprate"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/v2/coderd/aibridge"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/rbac"
@@ -70,3 +72,72 @@ func RateLimit(count int, window time.Duration) func(http.Handler) http.Handler
 		}),
 	)
 }
+
+// RateLimitByAuthToken returns a handler that limits requests based on the
+// authentication token in the request.
+//
+// This differs from [RateLimit] in several ways:
+//   - It extracts the token directly from request headers (Authorization Bearer
+//     or X-Api-Key) rather than from the request context, making it suitable for
+//     endpoints that handle authentication internally (like AI Bridge) rather than
+//     via [ExtractAPIKeyMW] middleware.
+//   - It does not support the bypass header for Owners.
+//   - It does not key by endpoint, so the limit applies across all endpoints using
+//     this middleware.
+//   - It includes a Retry-After header in 429 responses for backpressure signaling.
+//
+// If no token is found in the headers, it falls back to rate limiting by IP address.
+func RateLimitByAuthToken(count int, window time.Duration) func(http.Handler) http.Handler {
+	if count <= 0 {
+		return func(handler http.Handler) http.Handler {
+			return handler
+		}
+	}
+
+	return httprate.Limit(
+		count,
+		window,
+		httprate.WithKeyFuncs(func(r *http.Request) (string, error) {
+			// Try to extract auth token for per-user rate limiting using
+			// AI provider authentication headers (Authorization Bearer or X-Api-Key).
+			if token := aibridge.ExtractAuthToken(r.Header); token != "" {
+				return token, nil
+			}
+			// Fall back to IP-based rate limiting if no token present.
+			return httprate.KeyByIP(r)
+		}),
+		httprate.WithLimitHandler(func(w http.ResponseWriter, r *http.Request) {
+			// Add Retry-After header for backpressure signaling.
+			w.Header().Set("Retry-After", fmt.Sprintf("%d", int(window.Seconds())))
+			httpapi.Write(r.Context(), w, http.StatusTooManyRequests, codersdk.Response{
+				Message: "You've been rate limited. Please try again later.",
+			})
+		}),
+	)
+}
+
+// ConcurrencyLimit returns a handler that limits the number of concurrent
+// requests. When the limit is exceeded, it returns HTTP 503 Service Unavailable.
+func ConcurrencyLimit(maxConcurrent int64, resourceName string) func(http.Handler) http.Handler {
+	if maxConcurrent <= 0 {
+		return func(handler http.Handler) http.Handler {
+			return handler
+		}
+	}
+
+	var current atomic.Int64
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			c := current.Add(1)
+			defer current.Add(-1)
+
+			if c > maxConcurrent {
+				httpapi.Write(r.Context(), w, http.StatusServiceUnavailable, codersdk.Response{
+					Message: fmt.Sprintf("%s is currently at capacity. Please try again later.", resourceName),
+				})
+				return
+			}
+			next.ServeHTTP(w, r)
+		})
+	}
+}

coderd/httpmw/ratelimit_test.go

@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
 	"time"
 
@@ -17,6 +18,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
 )
 
 func randRemoteAddr() string {
@@ -145,3 +147,211 @@ func TestRateLimit(t *testing.T) {
 		}
 	})
 }
+
+func TestRateLimitByAuthToken(t *testing.T) {
+	t.Parallel()
+
+	t.Run("LimitsByAuthHeader", func(t *testing.T) {
+		t.Parallel()
+
+		tests := []struct {
+			name       string
+			headerName string
+			headerVal  string
+		}{
+			{
+				name:       "BearerToken",
+				headerName: "Authorization",
+				headerVal:  "Bearer test-token-123",
+			},
+			{
+				name:       "XApiKey",
+				headerName: "X-Api-Key",
+				headerVal:  "test-api-key-456",
+			},
+			{
+				name:       "NoToken",
+				headerName: "",
+				headerVal:  "",
+			},
+		}
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				t.Parallel()
+				rtr := chi.NewRouter()
+				rtr.Use(httpmw.RateLimitByAuthToken(2, time.Hour))
+				rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+					rw.WriteHeader(http.StatusOK)
+				})
+
+				// Same token (or IP if no token) should be rate limited after 2 requests.
+				for i := 0; i < 5; i++ {
+					req := httptest.NewRequest("GET", "/", nil)
+					if tt.headerName != "" {
+						req.Header.Set(tt.headerName, tt.headerVal)
+					}
+					rec := httptest.NewRecorder()
+					rtr.ServeHTTP(rec, req)
+					resp := rec.Result()
+					_ = resp.Body.Close()
+					if i < 2 {
+						require.Equal(t, http.StatusOK, resp.StatusCode, "request %d should succeed", i)
+					} else {
+						require.Equal(t, http.StatusTooManyRequests, resp.StatusCode, "request %d should be rate limited", i)
+						// Verify Retry-After header is set.
+						require.NotEmpty(t, resp.Header.Get("Retry-After"), "Retry-After header should be set")
+					}
+				}
+			})
+		}
+	})
+
+	t.Run("DifferentTokensNotLimited", func(t *testing.T) {
+		t.Parallel()
+		rtr := chi.NewRouter()
+		rtr.Use(httpmw.RateLimitByAuthToken(1, time.Hour))
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		// Different tokens should not be rate limited against each other.
+		for i := 0; i < 5; i++ {
+			req := httptest.NewRequest("GET", "/", nil)
+			req.Header.Set("Authorization", fmt.Sprintf("Bearer token-%d", i))
+			rec := httptest.NewRecorder()
+			rtr.ServeHTTP(rec, req)
+			resp := rec.Result()
+			_ = resp.Body.Close()
+			require.Equal(t, http.StatusOK, resp.StatusCode, "request %d should succeed", i)
+		}
+	})
+
+	t.Run("DisabledWhenZero", func(t *testing.T) {
+		t.Parallel()
+		rtr := chi.NewRouter()
+		rtr.Use(httpmw.RateLimitByAuthToken(0, time.Hour))
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		// Should not be rate limited when limit is 0.
+		for i := 0; i < 10; i++ {
+			req := httptest.NewRequest("GET", "/", nil)
+			req.Header.Set("Authorization", "Bearer same-token")
+			rec := httptest.NewRecorder()
+			rtr.ServeHTTP(rec, req)
+			resp := rec.Result()
+			_ = resp.Body.Close()
+			require.Equal(t, http.StatusOK, resp.StatusCode)
+		}
+	})
+}
+
+func TestConcurrencyLimit(t *testing.T) {
+	t.Parallel()
+
+	t.Run("LimitsConcurrentRequests", func(t *testing.T) {
+		t.Parallel()
+
+		const maxConcurrency = 2
+		rtr := chi.NewRouter()
+		rtr.Use(httpmw.ConcurrencyLimit(maxConcurrency, "Test"))
+
+		// Use a WaitGroup as a barrier to ensure all requests are in the handler
+		// before any of them proceed.
+		var handlersReady sync.WaitGroup
+		handlersReady.Add(maxConcurrency)
+		releaseHandler := make(chan struct{})
+
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			handlersReady.Done()
+			// Wait until released.
+			<-releaseHandler
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		server := httptest.NewServer(rtr)
+		defer server.Close()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+
+		// Start maxConcurrency requests that will block.
+		// We use channels to collect errors instead of require in goroutines.
+		type result struct {
+			statusCode int
+			err        error
+		}
+		results := make(chan result, maxConcurrency)
+
+		var wg sync.WaitGroup
+		for i := 0; i < maxConcurrency; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/", nil)
+				if err != nil {
+					results <- result{err: err}
+					return
+				}
+				resp, err := http.DefaultClient.Do(req)
+				if err != nil {
+					results <- result{err: err}
+					return
+				}
+				defer resp.Body.Close()
+				results <- result{statusCode: resp.StatusCode}
+			}()
+		}
+
+		// Wait for all requests to enter the handler with a timeout.
+		handlersReadyCh := make(chan struct{})
+		go func() {
+			handlersReady.Wait()
+			close(handlersReadyCh)
+		}()
+		select {
+		case <-handlersReadyCh:
+		case <-ctx.Done():
+			t.Fatal("timed out waiting for handlers to be ready")
+		}
+
+		// Next request should be rejected since we're at capacity.
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/", nil)
+		require.NoError(t, err)
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
+
+		// Release all blocked requests.
+		close(releaseHandler)
+		wg.Wait()
+		close(results)
+
+		// Check all goroutine results.
+		for res := range results {
+			require.NoError(t, res.err)
+			require.Equal(t, http.StatusOK, res.statusCode)
+		}
+	})
+
+	t.Run("DisabledWhenZero", func(t *testing.T) {
+		t.Parallel()
+		rtr := chi.NewRouter()
+		rtr.Use(httpmw.ConcurrencyLimit(0, "Test"))
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		// Should not be limited when maxConcurrency is 0.
+		for i := 0; i < 10; i++ {
+			req := httptest.NewRequest("GET", "/", nil)
+			rec := httptest.NewRecorder()
+			rtr.ServeHTTP(rec, req)
+			resp := rec.Result()
+			_ = resp.Body.Close()
+			require.Equal(t, http.StatusOK, resp.StatusCode)
+		}
+	})
+}