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)
This commit is contained in:
Kacper Sawicki
2025-12-11 16:38:54 +01:00
committed by GitHub
parent 8ead6f795d
commit 6f86f67754
17 changed files with 568 additions and 37 deletions
+22
View File
@@ -3401,6 +3401,26 @@ Write out the current server config as YAML to stdout.`,
YAML: "retention",
Annotations: serpent.Annotations{}.Mark(annotationFormatDuration, "true"),
},
{
Name: "AI Bridge Max Concurrency",
Description: "Maximum number of concurrent AI Bridge requests per replica. Set to 0 to disable (unlimited).",
Flag: "aibridge-max-concurrency",
Env: "CODER_AIBRIDGE_MAX_CONCURRENCY",
Value: &c.AI.BridgeConfig.MaxConcurrency,
Default: "0",
Group: &deploymentGroupAIBridge,
YAML: "maxConcurrency",
},
{
Name: "AI Bridge Rate Limit",
Description: "Maximum number of AI Bridge requests per second per replica. Set to 0 to disable (unlimited).",
Flag: "aibridge-rate-limit",
Env: "CODER_AIBRIDGE_RATE_LIMIT",
Value: &c.AI.BridgeConfig.RateLimit,
Default: "0",
Group: &deploymentGroupAIBridge,
YAML: "rateLimit",
},
// Retention settings
{
Name: "Audit Logs Retention",
@@ -3471,6 +3491,8 @@ type AIBridgeConfig struct {
Bedrock AIBridgeBedrockConfig `json:"bedrock" typescript:",notnull"`
InjectCoderMCPTools serpent.Bool `json:"inject_coder_mcp_tools" typescript:",notnull"`
Retention serpent.Duration `json:"retention" typescript:",notnull"`
MaxConcurrency serpent.Int64 `json:"max_concurrency" typescript:",notnull"`
RateLimit serpent.Int64 `json:"rate_limit" typescript:",notnull"`
}
type AIBridgeOpenAIConfig struct {