diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden index a698f163c7..286dfc0b1b 100644 --- a/cli/testdata/coder_server_--help.golden +++ b/cli/testdata/coder_server_--help.golden @@ -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. diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index bd1fc4cd68..abc5faa877 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -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) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index c0f6ecb23d..3c37594835 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -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": [ diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 31f5dd45e2..175fdd892c 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -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": [ diff --git a/codersdk/deployment.go b/codersdk/deployment.go index ddf2a08c7f..69cfc0d84b 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -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"` diff --git a/docs/ai-coder/ai-gateway/clients/index.md b/docs/ai-coder/ai-gateway/clients/index.md index c9e6b9bf0c..437e8f3e6d 100644 --- a/docs/ai-coder/ai-gateway/clients/index.md +++ b/docs/ai-coder/ai-gateway/clients/index.md @@ -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. diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index 7d2660395e..48ae6f8cf0 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -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" diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index e97796d2ee..fc08cb695e 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -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__* env vars instead. | | `bedrock` | [codersdk.AIBridgeBedrockConfig](#codersdkaibridgebedrockconfig) | false | | Deprecated: Use Providers with indexed CODER_AIBRIDGE_PROVIDER__* 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" diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index 43e293e56a..03c2f2c753 100644 --- a/docs/reference/cli/server.md +++ b/docs/reference/cli/server.md @@ -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 | bool | +| Environment | $CODER_AIBRIDGE_ALLOW_BYOK | +| YAML | 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-circuit-breaker-enabled | | | diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index c74f26fd7f..05a3790b16 100644 --- a/enterprise/cli/testdata/coder_server_--help.golden +++ b/enterprise/cli/testdata/coder_server_--help.golden @@ -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. diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 1d11315790..3b44ce2bae 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -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) }) }) diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index dea7030899..b06bc361a4 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -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.") + } + }) + } +} diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 4e7582e0a7..b85ec07ef2 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -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).