mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
Generated
+2
@@ -176,10 +176,12 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \
|
||||
},
|
||||
"enabled": true,
|
||||
"inject_coder_mcp_tools": true,
|
||||
"max_concurrency": 0,
|
||||
"openai": {
|
||||
"base_url": "string",
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
}
|
||||
},
|
||||
|
||||
Generated
+10
@@ -390,10 +390,12 @@
|
||||
},
|
||||
"enabled": true,
|
||||
"inject_coder_mcp_tools": true,
|
||||
"max_concurrency": 0,
|
||||
"openai": {
|
||||
"base_url": "string",
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
}
|
||||
```
|
||||
@@ -406,7 +408,9 @@
|
||||
| `bedrock` | [codersdk.AIBridgeBedrockConfig](#codersdkaibridgebedrockconfig) | false | | |
|
||||
| `enabled` | boolean | false | | |
|
||||
| `inject_coder_mcp_tools` | boolean | false | | |
|
||||
| `max_concurrency` | integer | false | | |
|
||||
| `openai` | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig) | false | | |
|
||||
| `rate_limit` | integer | false | | |
|
||||
| `retention` | integer | false | | |
|
||||
|
||||
## codersdk.AIBridgeInterception
|
||||
@@ -700,10 +704,12 @@
|
||||
},
|
||||
"enabled": true,
|
||||
"inject_coder_mcp_tools": true,
|
||||
"max_concurrency": 0,
|
||||
"openai": {
|
||||
"base_url": "string",
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
}
|
||||
}
|
||||
@@ -2860,10 +2866,12 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
|
||||
},
|
||||
"enabled": true,
|
||||
"inject_coder_mcp_tools": true,
|
||||
"max_concurrency": 0,
|
||||
"openai": {
|
||||
"base_url": "string",
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
}
|
||||
},
|
||||
@@ -3383,10 +3391,12 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
|
||||
},
|
||||
"enabled": true,
|
||||
"inject_coder_mcp_tools": true,
|
||||
"max_concurrency": 0,
|
||||
"openai": {
|
||||
"base_url": "string",
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
}
|
||||
},
|
||||
|
||||
Generated
+22
@@ -1781,6 +1781,28 @@ Whether to inject Coder's MCP tools into intercepted AI Bridge requests (require
|
||||
|
||||
Length of time to retain data such as interceptions and all related records (token, prompt, tool use).
|
||||
|
||||
### --aibridge-max-concurrency
|
||||
|
||||
| | |
|
||||
|-------------|----------------------------------------------|
|
||||
| Type | <code>int</code> |
|
||||
| Environment | <code>$CODER_AIBRIDGE_MAX_CONCURRENCY</code> |
|
||||
| YAML | <code>aibridge.maxConcurrency</code> |
|
||||
| Default | <code>0</code> |
|
||||
|
||||
Maximum number of concurrent AI Bridge requests per replica. Set to 0 to disable (unlimited).
|
||||
|
||||
### --aibridge-rate-limit
|
||||
|
||||
| | |
|
||||
|-------------|-----------------------------------------|
|
||||
| Type | <code>int</code> |
|
||||
| Environment | <code>$CODER_AIBRIDGE_RATE_LIMIT</code> |
|
||||
| YAML | <code>aibridge.rateLimit</code> |
|
||||
| Default | <code>0</code> |
|
||||
|
||||
Maximum number of AI Bridge requests per second per replica. Set to 0 to disable (unlimited).
|
||||
|
||||
### --audit-logs-retention
|
||||
|
||||
| | |
|
||||
|
||||
Reference in New Issue
Block a user