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

## 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>

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.
 

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)
 			})
 		})

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.")
+			}
+		})
+	}
+}