shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add allow-byok option for ai-gateway (#24274)

Browse Source 
## Summary                  
Adds `--ai-gateway-allow-byok` deployment option to control whether
users can use Bring Your Own Key (BYOK) mode with AI Gateway.
When disabled (`--ai-gateway-allow-byok=false`), BYOK requests are
rejected with a 403 and a message directing the admin to enable the
flag. Centralized key authentication works regardless of this setting.
Defaults to `true` (BYOK allowed).

---------

Co-authored-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Yevhenii Shcherbina
2026-04-15 14:16:49 -04:00
committed by GitHub
parent dd7397b42e
commit dd73ea54bd
13 changed files with 164 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/testdata/coder_server_--help.golden
Vendored
+4
View File
@@ -100,6 +100,10 @@ OPTIONS:
check is performed once per day.
AI BRIDGE OPTIONS:
--aibridge-allow-byok bool, $CODER_AIBRIDGE_ALLOW_BYOK (default: true)
Allow users to provide their own LLM API keys or subscriptions. When
disabled, only centralized key authentication is permitted.
--aibridge-anthropic-base-url string, $CODER_AIBRIDGE_ANTHROPIC_BASE_URL (default: https://api.anthropic.com/)
The base URL of the Anthropic API.
cli/testdata/server-config.yaml.golden
Vendored
+4
View File
@@ -816,6 +816,10 @@ aibridge:
# X-Ai-Bridge-Actor-Metadata-Username (their username).
# (default: false, type: bool)
send_actor_headers: false
# Allow users to provide their own LLM API keys or subscriptions. When disabled,
# only centralized key authentication is permitted.
# (default: true, type: bool)
allow_byok: true
# Enable the circuit breaker to protect against cascading failures from upstream
# AI provider rate limits (429, 503, 529 overloaded).
# (default: false, type: bool)
coderd/apidoc/docs.go
Generated
+3
View File
@@ -13102,6 +13102,9 @@ const docTemplate = `{
"codersdk.AIBridgeConfig": {
"type": "object",
"properties": {
"allow_byok": {
"type": "boolean"
},
"anthropic": {
"description": "Deprecated: Use Providers with indexed CODER_AIBRIDGE_PROVIDER_\u003cN\u003e_* env vars instead.",
"allOf": [
coderd/apidoc/swagger.json
Generated
+3
View File
@@ -11650,6 +11650,9 @@
"codersdk.AIBridgeConfig": {
"type": "object",
"properties": {
"allow_byok": {
"type": "boolean"
},
"anthropic": {
"description": "Deprecated: Use Providers with indexed CODER_AIBRIDGE_PROVIDER_\u003cN\u003e_* env vars instead.",
"allOf": [
codersdk/deployment.go
+11
View File
@@ -3811,6 +3811,16 @@ Write out the current server config as YAML to stdout.`,
Group: &deploymentGroupAIBridge,
YAML: "send_actor_headers",
},
{
Name: "AI Bridge Allow BYOK",
Description: "Allow users to provide their own LLM API keys or subscriptions. When disabled, only centralized key authentication is permitted.",
Flag: "aibridge-allow-byok",
Env: "CODER_AIBRIDGE_ALLOW_BYOK",
Value: &c.AI.BridgeConfig.AllowBYOK,
Default: "true",
Group: &deploymentGroupAIBridge,
YAML: "allow_byok",
},
{
Name: "AI Bridge Circuit Breaker Enabled",
Description: "Enable the circuit breaker to protect against cascading failures from upstream AI provider rate limits (429, 503, 529 overloaded).",
@@ -4062,6 +4072,7 @@ type AIBridgeConfig struct {
RateLimit serpent.Int64 `json:"rate_limit" typescript:",notnull"`
StructuredLogging serpent.Bool `json:"structured_logging" typescript:",notnull"`
SendActorHeaders serpent.Bool `json:"send_actor_headers" typescript:",notnull"`
AllowBYOK serpent.Bool `json:"allow_byok" typescript:",notnull"`
// Circuit breaker protects against cascading failures from upstream AI
// provider rate limits (429, 503, 529 overloaded).
CircuitBreakerEnabled serpent.Bool `json:"circuit_breaker_enabled" typescript:",notnull"`
docs/ai-coder/ai-gateway/clients/index.md
+10
View File
@@ -66,6 +66,16 @@ while allowing individual users to bring their own.
See individual client pages for configuration details.
### Enabling or disabling BYOK
BYOK is enabled by default. Administrators can disable it using `--aibridge-allow-byok=false` or `CODER_AIBRIDGE_ALLOW_BYOK=false`:
```sh
coder server --aibridge-allow-byok=false
```
When disabled, BYOK requests are rejected with a `403 Forbidden` response and only centralized key authentication is permitted.
## Compatibility
The table below shows tested AI clients and their compatibility with AI Gateway.
docs/reference/api/general.md
Generated
+1
View File
@@ -179,6 +179,7 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \
"upstream_proxy_ca": "string"
},
"bridge": {
"allow_byok": true,
"anthropic": {
"base_url": "string",
"key": "string"
docs/reference/api/schemas.md
Generated
+5
View File
@@ -431,6 +431,7 @@
```json
{
"allow_byok": true,
"anthropic": {
"base_url": "string",
"key": "string"
@@ -476,6 +477,7 @@
| Name | Type | Required | Restrictions | Description |
|-------------------------------------|-----------------------------------------------------------------------------|----------|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `allow_byok` | boolean | false | | |
| `anthropic` | [codersdk.AIBridgeAnthropicConfig](#codersdkaibridgeanthropicconfig) | false | | Deprecated: Use Providers with indexed CODER_AIBRIDGE_PROVIDER_<N>_* env vars instead. |
| `bedrock` | [codersdk.AIBridgeBedrockConfig](#codersdkaibridgebedrockconfig) | false | | Deprecated: Use Providers with indexed CODER_AIBRIDGE_PROVIDER_<N>_* env vars instead. |
| `circuit_breaker_enabled` | boolean | false | | Circuit breaker protects against cascading failures from upstream AI provider rate limits (429, 503, 529 overloaded). |
@@ -1245,6 +1247,7 @@
"upstream_proxy_ca": "string"
},
"bridge": {
"allow_byok": true,
"anthropic": {
"base_url": "string",
"key": "string"
@@ -3279,6 +3282,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
"upstream_proxy_ca": "string"
},
"bridge": {
"allow_byok": true,
"anthropic": {
"base_url": "string",
"key": "string"
@@ -3868,6 +3872,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
"upstream_proxy_ca": "string"
},
"bridge": {
"allow_byok": true,
"anthropic": {
"base_url": "string",
"key": "string"
docs/reference/cli/server.md
Generated
+11
View File
@@ -1879,6 +1879,17 @@ Emit structured logs for AI Bridge interception records. Use this for exporting
Once enabled, extra headers will be added to upstream requests to identify the user (actor) making requests to AI Bridge. This is only needed if you are using a proxy between AI Bridge and an upstream AI provider. This will send X-Ai-Bridge-Actor-Id (the ID of the user making the request) and X-Ai-Bridge-Actor-Metadata-Username (their username).
### --aibridge-allow-byok
| | |
|-------------|-----------------------------------------|
| Type | <code>bool</code> |
| Environment | <code>$CODER_AIBRIDGE_ALLOW_BYOK</code> |
| YAML | <code>aibridge.allow_byok</code> |
| Default | <code>true</code> |
Allow users to provide their own LLM API keys or subscriptions. When disabled, only centralized key authentication is permitted.
### --aibridge-circuit-breaker-enabled
| | |
enterprise/cli/testdata/coder_server_--help.golden
Vendored
+4
View File
@@ -101,6 +101,10 @@ OPTIONS:
check is performed once per day.
AI BRIDGE OPTIONS:
--aibridge-allow-byok bool, $CODER_AIBRIDGE_ALLOW_BYOK (default: true)
Allow users to provide their own LLM API keys or subscriptions. When
disabled, only centralized key authentication is permitted.
--aibridge-anthropic-base-url string, $CODER_AIBRIDGE_ANTHROPIC_BASE_URL (default: https://api.anthropic.com/)
The base URL of the Anthropic API.
enterprise/coderd/aibridge.go
+11
View File
@@ -15,6 +15,7 @@ import (
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/db2sdk"
"github.com/coder/coder/v2/coderd/httpapi"
@@ -74,6 +75,16 @@ func aibridgeHandler(api *API, middlewares ...func(http.Handler) http.Handler) f
return
}
// Reject BYOK requests when the deployment has not
// enabled bring-your-own-key mode.
if agplaibridge.IsBYOK(r.Header) && !bridgeCfg.AllowBYOK.Value() {
httpapi.Write(r.Context(), rw, http.StatusForbidden, codersdk.Response{
Message: "Bring Your Own Key (BYOK) mode is not enabled.",
Detail: "Contact your administrator to enable it with --aibridge-allow-byok.",
})
return
}
http.StripPrefix("/api/v2/aibridge", api.aibridgedHandler).ServeHTTP(rw, r)
})
})
enterprise/coderd/aibridge_test.go
+96
View File
@@ -12,6 +12,7 @@ import (
"github.com/stretchr/testify/require"
aiblib "github.com/coder/aibridge"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/db2sdk"
@@ -2287,3 +2288,98 @@ func TestAIBridgeGetSessionThreads(t *testing.T) {
require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
})
}
func TestAIBridgeAllowBYOK(t *testing.T) {
t.Parallel()
cases := []struct {
name string
allowBYOK bool
reqHeaders map[string]string
expectedStatus int
}{
{
name: "byok_enabled/centralized_request",
allowBYOK: true,
reqHeaders: map[string]string{
"Authorization": "Bearer coder-token",
},
expectedStatus: http.StatusOK,
},
{
name: "byok_enabled/byok_request",
allowBYOK: true,
reqHeaders: map[string]string{
agplaibridge.HeaderCoderToken: "coder-token",
"Authorization": "Bearer user-llm-key",
},
expectedStatus: http.StatusOK,
},
{
name: "byok_disabled/centralized_request",
allowBYOK: false,
reqHeaders: map[string]string{
"Authorization": "Bearer coder-token",
},
expectedStatus: http.StatusOK,
},
{
name: "byok_disabled/byok_request",
allowBYOK: false,
reqHeaders: map[string]string{
agplaibridge.HeaderCoderToken: "coder-token",
"Authorization": "Bearer user-llm-key",
},
expectedStatus: http.StatusForbidden,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
dv := coderdtest.DeploymentValues(t)
dv.AI.BridgeConfig.Enabled = serpent.Bool(true)
dv.AI.BridgeConfig.AllowBYOK = serpent.Bool(tc.allowBYOK)
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()
})
testHandler := http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
rw.WriteHeader(http.StatusOK)
})
api.RegisterInMemoryAIBridgedHTTPHandler(testHandler)
ctx := testutil.Context(t, testutil.WaitLong)
reqURL := client.URL.String() + "/api/v2/aibridge/test"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, reqURL, nil)
require.NoError(t, err)
req.Header.Set(codersdk.SessionTokenHeader, client.SessionToken())
for k, v := range tc.reqHeaders {
req.Header.Set(k, v)
}
resp, err := http.DefaultClient.Do(req)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, tc.expectedStatus, resp.StatusCode)
if tc.expectedStatus == http.StatusForbidden {
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
require.Contains(t, string(body), "Bring Your Own Key (BYOK) mode is not enabled.")
}
})
}
}
site/src/api/typesGenerated.ts
Generated
+1
View File
@@ -67,6 +67,7 @@ export interface AIBridgeConfig {
readonly rate_limit: number;
readonly structured_logging: boolean;
readonly send_actor_headers: boolean;
readonly allow_byok: boolean;
/**
* Circuit breaker protects against cascading failures from upstream AI
* provider rate limits (429, 503, 529 overloaded).