From 89740ab41ecd4c6d9946c6c4dcdf505361512859 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Thu, 6 Aug 2026 16:04:38 +0000 Subject: [PATCH 1/3] test(enterprise/cli): add standalone AI Gateway connection tests --- .../cli/aigatewaystart_internal_test.go | 343 ++++++------------ enterprise/cli/aigatewaystart_test.go | 244 +++++++++++++ 2 files changed, 363 insertions(+), 224 deletions(-) create mode 100644 enterprise/cli/aigatewaystart_test.go diff --git a/enterprise/cli/aigatewaystart_internal_test.go b/enterprise/cli/aigatewaystart_internal_test.go index 9e4615f101b..3b5af611d5b 100644 --- a/enterprise/cli/aigatewaystart_internal_test.go +++ b/enterprise/cli/aigatewaystart_internal_test.go @@ -5,7 +5,6 @@ package cli import ( "context" "errors" - "fmt" "net" "net/http" "net/http/httptest" @@ -21,40 +20,30 @@ import ( "storj.io/drpc" "cdr.dev/slog/v3" - "github.com/coder/coder/v2/cli/clitest" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/aibridged" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) -// failThenSucceedReloader fails the first failUntil reloads, then succeeds, -// modeling a coderd connection or provider fetch that recovers after a few -// transient failures. -type failThenSucceedReloader struct { +// mockReloader fails the first failUntil reloads with err, then succeeds. +// after runs at the end of every reload, letting a test hang the retry loop. +type mockReloader struct { calls atomic.Int32 failUntil int32 + err error + after func() } -func (r *failThenSucceedReloader) Reload(_ context.Context) error { - if r.calls.Add(1) <= r.failUntil { - return xerrors.New("transient failure") - } - return nil -} - -type failingReloader struct { - after func() - calls atomic.Int32 - err error -} - -func (r *failingReloader) Reload(context.Context) error { - r.calls.Add(1) +func (r *mockReloader) Reload(context.Context) error { + failed := r.calls.Add(1) <= r.failUntil if r.after != nil { r.after() } - return r.err + if failed { + return r.err + } + return nil } type connectedDRPCConn struct { @@ -95,12 +84,14 @@ func (p *controlledShutdownPool) Shutdown(ctx context.Context) error { return errors.Join(p.CachedBridgePool.Shutdown(ctx), p.err) } -type standaloneGatewayTestParams struct { - params standaloneGatewayParams - pool *controlledShutdownPool -} +// testGatewayOption mutates the default standaloneGatewayParams before the +// gateway is constructed. +type testGatewayOption func(*standaloneGatewayParams) -func newStandaloneGatewayTestParams(t *testing.T) *standaloneGatewayTestParams { +// newTestStandaloneGateway constructs a standalone gateway for +// testing, with a controllable shutdown pool and optional customizations. +// Uses blockingStandaloneDaemonDialer to mock coderd. +func newTestStandaloneGateway(t *testing.T, opts ...testGatewayOption) (*standaloneGateway, *controlledShutdownPool) { t.Helper() logger := slog.Make() @@ -112,18 +103,26 @@ func newStandaloneGatewayTestParams(t *testing.T) *standaloneGatewayTestParams { }) pool := &controlledShutdownPool{CachedBridgePool: cachedPool} - return &standaloneGatewayTestParams{ - params: standaloneGatewayParams{ - httpAddress: "127.0.0.1:0", + params := standaloneGatewayParams{ + httpAddress: "127.0.0.1:0", - dialer: blockingStandaloneDaemonDialer, - pool: pool, + dialer: blockingStandaloneDaemonDialer, + pool: pool, - logger: logger, - tracer: tracer, - }, - pool: pool, + logger: logger, + tracer: tracer, + } + for _, m := range opts { + m(¶ms) } + + gateway, err := newStandaloneGateway(params) + require.NoError(t, err) + + t.Cleanup(func() { + require.NoError(t, shutdownWithTimeout(gateway.daemon.Shutdown, testutil.WaitShort)) + }) + return gateway, pool } func TestStandaloneGatewayLoadProviders(t *testing.T) { @@ -131,40 +130,37 @@ func TestStandaloneGatewayLoadProviders(t *testing.T) { reloadErr := xerrors.New("reload failed") tests := []struct { - name string - setup func(*testing.T, *aibridged.Server, context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) - wantErr error - wantCalls int32 - wantLoaded bool + name string + reloaderFailUntil int32 + reloaderErr error + reloaderAfter func(t *testing.T, daemon *aibridged.Server, cancel context.CancelFunc) + wantErr error + wantCalls int32 + wantLoaded bool }{ { - name: "Retry succeeds", - setup: func(_ *testing.T, _ *aibridged.Server, _ context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) { - reloader := &failThenSucceedReloader{failUntil: 2} - return reloader, &reloader.calls - }, - wantCalls: 3, - wantLoaded: true, + name: "Retry succeeds", + reloaderFailUntil: 2, + reloaderErr: xerrors.New("transient failure"), + wantCalls: 3, + wantLoaded: true, }, { - name: "Daemon stops retry", - setup: func(t *testing.T, daemon *aibridged.Server, _ context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) { - reloader := &failingReloader{ - after: func() { - require.NoError(t, daemon.Close()) - }, - err: reloadErr, - } - return reloader, &reloader.calls + name: "Daemon stops retry", + reloaderFailUntil: 1, + reloaderErr: reloadErr, + reloaderAfter: func(t *testing.T, daemon *aibridged.Server, _ context.CancelFunc) { + require.NoError(t, daemon.Close()) }, wantErr: reloadErr, wantCalls: 1, }, { - name: "Context cancellation stops retry", - setup: func(_ *testing.T, _ *aibridged.Server, cancel context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) { - reloader := &failingReloader{after: cancel, err: reloadErr} - return reloader, &reloader.calls + name: "Context cancellation stops retry", + reloaderFailUntil: 1, + reloaderErr: reloadErr, + reloaderAfter: func(_ *testing.T, _ *aibridged.Server, cancel context.CancelFunc) { + cancel() }, wantErr: context.Canceled, wantCalls: 1, @@ -177,13 +173,13 @@ func TestStandaloneGatewayLoadProviders(t *testing.T) { ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort)) defer cancel() - logger := slog.Make() - daemon := newTestStandaloneDaemon(t, logger) - reloader, calls := tc.setup(t, daemon, cancel) - gateway := &standaloneGateway{ - daemon: daemon, - providerLogger: logger, - reloader: reloader, + reloader := &mockReloader{failUntil: tc.reloaderFailUntil, err: tc.reloaderErr} + gateway, _ := newTestStandaloneGateway(t) + gateway.reloader = reloader + // reloaderAfter needs daemon which only exists once the gateway is constructed, + // nothing reloads until loadProviders below. + if tc.reloaderAfter != nil { + reloader.after = func() { tc.reloaderAfter(t, gateway.daemon, cancel) } } err := gateway.loadProviders(ctx) @@ -192,7 +188,7 @@ func TestStandaloneGatewayLoadProviders(t *testing.T) { } else { require.ErrorIs(t, err, tc.wantErr) } - require.Equal(t, tc.wantCalls, calls.Load()) + require.Equal(t, tc.wantCalls, reloader.calls.Load()) require.Equal(t, tc.wantLoaded, gateway.providersLoaded.Load()) }) } @@ -202,34 +198,19 @@ func TestStandaloneGatewayHealthAndReadiness(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) - logger := slog.Make() - tracer := sdktrace.NewTracerProvider().Tracer("test") - pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, tracer) - require.NoError(t, err) connections := make(chan drpc.Conn, 2) - dialer := func(ctx context.Context) (aibridged.DRPCClient, error) { - select { - case conn := <-connections: - return &aibridged.Client{Conn: conn}, nil - case <-ctx.Done(): - return nil, ctx.Err() + modifyDialer := func(p *standaloneGatewayParams) { + p.dialer = func(ctx context.Context) (aibridged.DRPCClient, error) { + select { + case conn := <-connections: + return &aibridged.Client{Conn: conn}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } } } - daemon, err := aibridged.New(ctx, pool, dialer, logger, tracer) - require.NoError(t, err) - t.Cleanup(func() { - require.NoError(t, shutdownWithTimeout(daemon.Shutdown, testutil.WaitShort)) - }) - - gateway := &standaloneGateway{ - daemon: daemon, - providerLogger: logger, - reloader: &failThenSucceedReloader{}, - } - gateway.httpServer = &http.Server{ - Handler: newGatewayMux(daemon, gateway.ready, func(next http.Handler) http.Handler { return next }), - ReadHeaderTimeout: testutil.WaitShort, - } + gateway, _ := newTestStandaloneGateway(t, modifyDialer) + gateway.reloader = &mockReloader{} // The HTTP server is healthy before the daemon connects or providers load. require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) @@ -238,7 +219,7 @@ func TestStandaloneGatewayHealthAndReadiness(t *testing.T) { // A daemon connection alone does not make the gateway ready. firstConn := &connectedDRPCConn{closed: make(chan struct{})} connections <- firstConn - require.Eventually(t, daemon.Ready, testutil.WaitShort, testutil.IntervalFast) + require.Eventually(t, gateway.daemon.Ready, testutil.WaitShort, testutil.IntervalFast) require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) require.Equal(t, http.StatusServiceUnavailable, readyzStatus(t, gateway)) @@ -249,13 +230,13 @@ func TestStandaloneGatewayHealthAndReadiness(t *testing.T) { // Losing the daemon connection affects readiness but not HTTP health. require.NoError(t, firstConn.Close()) - require.Eventually(t, func() bool { return !daemon.Ready() }, testutil.WaitShort, testutil.IntervalFast) + require.Eventually(t, func() bool { return !gateway.daemon.Ready() }, testutil.WaitShort, testutil.IntervalFast) require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) require.Equal(t, http.StatusServiceUnavailable, readyzStatus(t, gateway)) // Readiness recovers when the daemon reconnects; providers remain loaded. connections <- &connectedDRPCConn{closed: make(chan struct{})} - require.Eventually(t, daemon.Ready, testutil.WaitShort, testutil.IntervalFast) + require.Eventually(t, gateway.daemon.Ready, testutil.WaitShort, testutil.IntervalFast) require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) require.Equal(t, http.StatusOK, readyzStatus(t, gateway)) } @@ -277,73 +258,13 @@ func probeStatus(t *testing.T, gateway *standaloneGateway, path string) int { return rec.Code } -func TestAIGatewayStart_HealthBeforeReady(t *testing.T) { - t.Parallel() - - // Fake coderd that answers 503 so the daemon keeps retrying to connect. - coderSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusServiceUnavailable) - })) - t.Cleanup(coderSrv.Close) - gatewayAddress := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) - - var root RootCmd - cmd, err := root.Command(root.enterpriseOnly()) - require.NoError(t, err) - inv, _ := clitest.NewWithCommand(t, cmd, - "--url", coderSrv.URL, - "ai-gateway", "start", - "--key", "test-key", - "--http-address", gatewayAddress, - ) - ctx := testutil.Context(t, testutil.WaitShort) - // Watch the command for an early exit so a clash on gatewayAddress is - // reported as a bind failure instead of a probe timeout. - cmdDone := make(chan error, 1) - clitest.StartWithAssert(t, inv.WithContext(ctx), func(_ *testing.T, err error) { - cmdDone <- err - }) - - // healthz check - client := &http.Client{Timeout: testutil.WaitShort} - baseURL := "http://" + gatewayAddress - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case err := <-cmdDone: - t.Fatalf("ai-gateway start exited before serving: %v", err) - default: - } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+healthzPath, nil) - if err != nil { - return false - } - resp, err := client.Do(req) - if err != nil { - return false - } - defer resp.Body.Close() - return resp.StatusCode == http.StatusOK - }, testutil.IntervalFast) - - // readyz check (unavailable due to no connection to coderd) - req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+readyzPath, nil) - require.NoError(t, err) - resp, err := client.Do(req) - require.NoError(t, err) - defer resp.Body.Close() - require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) -} - func TestRunStandaloneGateway_ContextCanceled(t *testing.T) { t.Parallel() testCtx := testutil.Context(t, testutil.WaitShort) runCtx, cancelRun := context.WithCancel(testCtx) defer cancelRun() - test := newStandaloneGatewayTestParams(t) - - gateway, err := newStandaloneGateway(test.params) - require.NoError(t, err) + gateway, _ := newTestStandaloneGateway(t) runDone := make(chan error, 1) go func() { runDone <- gateway.run(runCtx) @@ -360,14 +281,13 @@ func TestRunStandaloneGateway_ContextCanceled(t *testing.T) { func TestRunStandaloneGateway_DaemonExited(t *testing.T) { t.Parallel() - test := newStandaloneGatewayTestParams(t) - test.params.dialer = func(context.Context) (aibridged.DRPCClient, error) { - return nil, codersdk.NewError(http.StatusUnauthorized, codersdk.Response{Message: "invalid gateway key"}) + modifyDialer := func(p *standaloneGatewayParams) { + p.dialer = func(context.Context) (aibridged.DRPCClient, error) { + return nil, codersdk.NewError(http.StatusUnauthorized, codersdk.Response{Message: "invalid gateway key"}) + } } - - gateway, err := newStandaloneGateway(test.params) - require.NoError(t, err) - err = gateway.run(testutil.Context(t, testutil.WaitShort)) + gateway, _ := newTestStandaloneGateway(t, modifyDialer) + err := gateway.run(testutil.Context(t, testutil.WaitShort)) require.ErrorContains(t, err, "AI Gateway daemon exited") require.True(t, gateway.drpcClosed.Load(), "DRPC connection must be closed before run returns") require.True(t, gateway.providerRefreshStopped.Load(), "provider refresh must stop before run returns") @@ -378,18 +298,18 @@ func TestRunStandaloneGateway_HTTPStopsBeforeDaemonShutdown(t *testing.T) { t.Parallel() testCtx := testutil.Context(t, testutil.WaitShort) - test := newStandaloneGatewayTestParams(t) + modifyTLS := func(p *standaloneGatewayParams) { + p.tlsCertFile = filepath.Join(t.TempDir(), "missing.crt") + p.tlsKeyFile = filepath.Join(t.TempDir(), "missing.key") + } + gateway, pool := newTestStandaloneGateway(t, modifyTLS) shutdownErr := xerrors.New("pool shutdown failed") shutdownStarted := make(chan struct{}, 1) shutdownRelease := make(chan struct{}) - test.pool.err = shutdownErr - test.pool.started = shutdownStarted - test.pool.release = shutdownRelease - test.params.tlsCertFile = filepath.Join(t.TempDir(), "missing.crt") - test.params.tlsKeyFile = filepath.Join(t.TempDir(), "missing.key") + pool.err = shutdownErr + pool.started = shutdownStarted + pool.release = shutdownRelease - gateway, err := newStandaloneGateway(test.params) - require.NoError(t, err) runDone := make(chan error, 1) go func() { runDone <- gateway.run(testCtx) @@ -400,7 +320,7 @@ func TestRunStandaloneGateway_HTTPStopsBeforeDaemonShutdown(t *testing.T) { require.False(t, gateway.drpcClosed.Load(), "DRPC connection must remain open until daemon shutdown completes") close(shutdownRelease) - err = testutil.RequireReceive(testCtx, t, runDone) + err := testutil.RequireReceive(testCtx, t, runDone) require.False(t, gateway.drpcClosed.Load(), "DRPC connection shutdown must not be marked successful after an error") require.ErrorContains(t, err, "serve:") require.ErrorContains(t, err, "shutdown AI Gateway daemon:") @@ -410,17 +330,20 @@ func TestRunStandaloneGateway_HTTPStopsBeforeDaemonShutdown(t *testing.T) { func TestRunStandaloneGateway_ListenAndShutdownErrors(t *testing.T) { t.Parallel() - test := newStandaloneGatewayTestParams(t) - shutdownErr := xerrors.New("pool shutdown failed") - test.pool.err = shutdownErr listener, err := net.Listen("tcp", "127.0.0.1:0") require.NoError(t, err) t.Cleanup(func() { _ = listener.Close() }) - test.params.httpAddress = listener.Addr().String() + // Occupy an address so binding the gateway listener fails. + modifyAddr := func(p *standaloneGatewayParams) { + p.httpAddress = listener.Addr().String() + } + gateway, pool := newTestStandaloneGateway(t, modifyAddr) + shutdownErr := xerrors.New("pool shutdown failed") + pool.err = shutdownErr - err = runStandaloneGateway(testutil.Context(t, testutil.WaitShort), test.params) + err = gateway.run(testutil.Context(t, testutil.WaitShort)) require.NoError(t, listener.Close()) require.ErrorContains(t, err, "listen on") require.ErrorContains(t, err, "shutdown AI Gateway daemon:") @@ -432,45 +355,31 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { // Set up a running daemon, provider reloader, and blocked HTTP request. testCtx := testutil.Context(t, testutil.WaitShort) - logger := slog.Make() - pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, sdktrace.NewTracerProvider().Tracer("test")) - require.NoError(t, err) - - daemon, err := aibridged.New(context.Background(), pool, blockingStandaloneDaemonDialer, logger, sdktrace.NewTracerProvider().Tracer("test")) - require.NoError(t, err) - httpAddress := "127.0.0.1:0" // inFlightPath scopes the blocking handler to the request this test keeps in // flight. Another test may probe a port the OS later assigns to this // listener, and such traffic must not stand in for that request. const inFlightPath = "/in-flight" - reloader := &failThenSucceedReloader{} + reloader := &mockReloader{} handlerStarted := make(chan struct{}, 1) httpShutdownStarted := make(chan struct{}, 1) releaseHandler := make(chan struct{}) - gateway := &standaloneGateway{ - daemon: daemon, - httpServer: &http.Server{ - Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != inFlightPath { - w.WriteHeader(http.StatusNotFound) - return - } - select { - case handlerStarted <- struct{}{}: - default: - } - <-releaseHandler - w.WriteHeader(http.StatusNoContent) - }), - ReadHeaderTimeout: testutil.WaitShort, - }, - httpAddress: httpAddress, - logger: logger, - providerLogger: logger, - reloader: reloader, - listenerReady: make(chan struct{}), - } + gateway, _ := newTestStandaloneGateway(t) + gateway.reloader = reloader + // The gateway mux is replaced with a handler this test can block, so an + // in-flight request is observable during shutdown. + gateway.httpServer.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != inFlightPath { + w.WriteHeader(http.StatusNotFound) + return + } + select { + case handlerStarted <- struct{}{}: + default: + } + <-releaseHandler + w.WriteHeader(http.StatusNoContent) + }) gateway.httpServer.RegisterOnShutdown(func() { httpShutdownStarted <- struct{}{} }) @@ -539,20 +448,6 @@ func requireListening(ctx context.Context, t *testing.T, gateway *standaloneGate return "" } -func newTestStandaloneDaemon(t *testing.T, logger slog.Logger) *aibridged.Server { - t.Helper() - - tracer := sdktrace.NewTracerProvider().Tracer("test") - pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, tracer) - require.NoError(t, err) - daemon, err := aibridged.New(context.Background(), pool, blockingStandaloneDaemonDialer, logger, tracer) - require.NoError(t, err) - t.Cleanup(func() { - require.NoError(t, daemon.Close()) - }) - return daemon -} - func blockingStandaloneDaemonDialer(ctx context.Context) (aibridged.DRPCClient, error) { <-ctx.Done() return nil, ctx.Err() diff --git a/enterprise/cli/aigatewaystart_test.go b/enterprise/cli/aigatewaystart_test.go new file mode 100644 index 00000000000..2daa8574cf7 --- /dev/null +++ b/enterprise/cli/aigatewaystart_test.go @@ -0,0 +1,244 @@ +//go:build !slim + +package cli_test + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "regexp" + "strings" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/enterprise/coderd/license" + "github.com/coder/coder/v2/pty/ptytest" + "github.com/coder/coder/v2/testutil" + "github.com/coder/serpent" +) + +// The end-to-end tests in this file run the real `ai-gateway start` command +// against a real coderd and only observe the public surface: CLI flags, log +// output, the gateway's HTTP endpoints, coderd's API, and the command's exit +// error. Each part of the standalone plumbing executes at least once. Detailed +// behavior (readiness transitions, reconnect semantics, shutdown ordering) is +// covered by the internal tests in aigatewaystart_internal_test.go and the +// reconnection tests, which construct the gateway directly. + +// aiGatewayChatCompletionRequest is the LLM request sent through the gateway. The +// model is asserted against the recorded interception. +const aiGatewayChatCompletionRequest = `{"messages":[{"role":"user","content":"standalone gateway e2e"}],"model":"gpt-4.1"}` + +// aiGatewayUpstreamResponse is the fixed completion the mock upstream LLM API +// returns, so the test can recognize it in the gateway's response. +const aiGatewayUpstreamResponse = `{ + "id": "chatcmpl-e2e", + "object": "chat.completion", + "created": 1753343279, + "model": "gpt-4.1", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "standalone gateway e2e response"}, + "finish_reason": "stop" + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} +}` + +// aiGatewayDeployment is a coderd entitled for the AI Gateway, with a gateway key, a +// configured provider backed by a mock upstream, and a member user whose +// session token authenticates LLM traffic. +type aiGatewayDeployment struct { + client *codersdk.Client + userClient *codersdk.Client + user codersdk.User + key string + upstreamHits *atomic.Int32 +} + +func setupAIGatewayDeployment(ctx context.Context, t *testing.T) *aiGatewayDeployment { + t.Helper() + + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + client, firstUser := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{codersdk.FeatureAIBridge: 1}, + }, + }) + + //nolint:gocritic // Owner role is needed for gateway key management. + key, err := client.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "e2e"}) + require.NoError(t, err) + + var hits atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(aiGatewayUpstreamResponse)) + })) + t.Cleanup(upstream.Close) + + //nolint:gocritic // Owner role is needed for provider management. + _, err = client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Type: codersdk.AIProviderTypeOpenAI, + Name: "openai", + Enabled: true, + BaseURL: upstream.URL, + APIKeys: []string{"sk-e2e"}, + }) + require.NoError(t, err) + + userClient, user := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID) + + return &aiGatewayDeployment{ + client: client, + userClient: userClient, + user: user, + key: key.Key, + upstreamHits: &hits, + } +} + +// startAIGatewayCommand runs `ai-gateway start` and returns the base URL of +// its HTTP listener, discovered from the startup log line, together with the +// command's error waiter. +func startAIGatewayCommand(ctx context.Context, t *testing.T, coderURL, key string) (string, *clitest.ErrorWaiter) { + t.Helper() + + inv, _ := newCLI(t, + "ai-gateway", "start", + "--url", coderURL, + "--key", key, + "--http-address", "127.0.0.1:0", + ) + inv = inv.WithContext(ctx) + pty := ptytest.New(t).Attach(inv) + waiter := clitest.StartWithWaiter(t, inv) + + // Extract bound address from the startup log. + pty.ExpectMatch(ctx, "standalone AI Gateway listening") + line := pty.ReadLine(ctx) + matches := regexp.MustCompile(`address=([0-9.]+:[0-9]+)`).FindStringSubmatch(line) + require.Len(t, matches, 2, "listener address not found in startup log: %q", line) + return "http://" + matches[1], waiter +} + +func getAIGatewayStatus(ctx context.Context, t *testing.T, url string) int { + t.Helper() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + require.NoError(t, err) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + return resp.StatusCode +} + +// TestAIGatewayStartE2E drives every part of the standalone gateway plumbing +// once through public surface only: the CLI starts with flags, connects to +// coderd with a gateway key, reports health and readiness, proxies an LLM +// request from a real client to a real upstream, records the interception in +// coderd, and shuts down cleanly. +func TestAIGatewayStartE2E(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + dep := setupAIGatewayDeployment(ctx, t) + + baseURL, waiter := startAIGatewayCommand(ctx, t, dep.client.URL.String(), dep.key) + + // Liveness holds as soon as the listener is up; readiness follows once the + // DRPC connection is established and providers are loaded. + require.Equal(t, http.StatusOK, getAIGatewayStatus(ctx, t, baseURL+"/healthz")) + require.Eventually(t, func() bool { + return getAIGatewayStatus(ctx, t, baseURL+"/readyz") == http.StatusOK + }, testutil.WaitLong, testutil.IntervalFast) + + // One LLM request through the gateway's own listener. + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + baseURL+"/openai/v1/chat/completions", strings.NewReader(aiGatewayChatCompletionRequest)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+dep.userClient.SessionToken()) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + body, err := io.ReadAll(resp.Body) + _ = resp.Body.Close() + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "body: %s", body) + require.Contains(t, string(body), "standalone gateway e2e response") + require.Equal(t, int32(1), dep.upstreamHits.Load()) + + // The interception is recorded in coderd. Recording is asynchronous, so + // the assertion has to be eventual. + require.Eventually(t, func() bool { + //nolint:gocritic // Owner role is needed to list every user's sessions. + sessions, err := dep.client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{ + Initiator: dep.user.Username, + }) + return err == nil && len(sessions.Sessions) == 1 + }, testutil.WaitLong, testutil.IntervalFast) + + // Graceful shutdown: canceling the command must produce a clean exit. + waiter.Cancel() + require.NoError(t, waiter.Wait()) +} + +// TestAIGatewayStartE2E_InvalidKey covers the fatal error plumbing: a gateway +// started with a key coderd rejects must exit with the rejection instead of +// retrying forever. +func TestAIGatewayStartE2E_InvalidKey(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + dv := coderdtest.DeploymentValues(t) + dv.AI.BridgeConfig.Enabled = serpent.Bool(true) + client, _ := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{DeploymentValues: dv}, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{codersdk.FeatureAIBridge: 1}, + }, + }) + + inv, _ := newCLI(t, + "ai-gateway", "start", + "--url", client.URL.String(), + "--key", "not-a-valid-key", + "--http-address", "127.0.0.1:0", + ) + inv = inv.WithContext(ctx) + waiter := clitest.StartWithWaiter(t, inv) + waiter.RequireContains("AI Gateway key invalid") +} + +// TestAIGatewayStart_HealthBeforeReady covers the split between liveness and +// readiness: the listener serves /healthz as soon as it is bound, while +// /readyz stays 503 until the daemon reaches coderd. +func TestAIGatewayStart_HealthBeforeReady(t *testing.T) { + t.Parallel() + + // Fake coderd that answers 503 so the daemon keeps retrying to connect. + coderSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + t.Cleanup(coderSrv.Close) + + ctx := testutil.Context(t, testutil.WaitShort) + baseURL, _ := startAIGatewayCommand(ctx, t, coderSrv.URL, "test-key") + + // The startup log line is emitted after the listener is bound, so no retry + // loop is needed. + require.Equal(t, http.StatusOK, getAIGatewayStatus(ctx, t, baseURL+"/healthz")) + require.Equal(t, http.StatusServiceUnavailable, getAIGatewayStatus(ctx, t, baseURL+"/readyz")) +} From 9bcc0a2b9f56241219b9c13c30a3502e3e139125 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 11 Aug 2026 08:34:28 +0000 Subject: [PATCH 2/3] docs(ai-coder/ai-gateway): use approximate spend and add Everyone group tip Align cost control docs with the UI, which now labels spend as approximate rather than estimated. Also note that an admin can use the organization's Everyone group Members tab to look up any user's effective group. --- docs/ai-coder/ai-gateway/cost-controls.md | 38 ++++++++++++----------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/docs/ai-coder/ai-gateway/cost-controls.md b/docs/ai-coder/ai-gateway/cost-controls.md index 52c4fa13c44..4ab0edca51c 100644 --- a/docs/ai-coder/ai-gateway/cost-controls.md +++ b/docs/ai-coder/ai-gateway/cost-controls.md @@ -17,20 +17,20 @@ AI Governance Cost Control requires: - AI Gateway [enabled and configured](./setup.md) with at least one provider. > [!NOTE] -> AI Governance Cost Control reports estimated spend rather than billed cost. -> Estimates will not match your provider invoices exactly. For details, see -> [How spend is estimated](#how-spend-is-estimated). +> AI Governance Cost Control reports approximate spend rather than billed cost. +> These figures will not match your provider invoices exactly. For details, see +> [How spend is approximated](#how-spend-is-approximated). These terms appear throughout this page and in the Coder dashboard: -| Term | What it means | Where it is set | -|---------------------|---------------------------------------------------------------------------------------|----------------------------------------| -| **Budget period** | The window spend accumulates in before it resets. Defaults to the UTC calendar month. | Deployment settings | -| **Budget policy** | The rule that selects a user's effective group. Defaults to the highest budget. | Deployment settings | -| **Group budget** | A spend limit granted to each member of a group. | **Groups** > {group} > **Settings** | -| **User override** | A spend limit for one user that supersedes their group budget. | **Groups** > {group} > **Members** tab | -| **Effective group** | The group that supplies a user's budget and has their spend associated with it. | Resolved automatically | -| **Estimated spend** | A user's approximate spend in the current budget period. | Estimated from usage | +| Term | What it means | Where it is set | +|-----------------------|---------------------------------------------------------------------------------------|----------------------------------------| +| **Budget period** | The window spend accumulates in before it resets. Defaults to the UTC calendar month. | Deployment settings | +| **Budget policy** | The rule that selects a user's effective group. Defaults to the highest budget. | Deployment settings | +| **Group budget** | A spend limit granted to each member of a group. | **Groups** > {group} > **Settings** | +| **User override** | A spend limit for one user that supersedes their group budget. | **Groups** > {group} > **Members** tab | +| **Effective group** | The group that supplies a user's budget and has their spend associated with it. | Resolved automatically | +| **Approximate spend** | A user's approximate spend in the current budget period. | Calculated from usage | ## Deployment settings @@ -181,7 +181,7 @@ affected user: For delivery methods, see [Notifications](../../admin/monitoring/notifications/index.md). -## How spend is estimated +## How spend is approximated Coder multiplies the token usage of each request by the published price of the model that served it. Prices come from a curated [models.dev](https://models.dev) @@ -200,10 +200,10 @@ https://github.com/coder/coder/blob/release//coderd/aibridge/prices/dat Replace `` with your Coder minor version, for example `2.36`. > [!IMPORTANT] -> Estimated spend can differ from provider-reported amounts, and some usage might -> not count toward spend: +> Approximate spend can differ from provider-reported amounts, and some usage +> might not count toward spend: > -> - Estimates exclude negotiated discounts, committed-use pricing, and +> - Approximations exclude negotiated discounts, committed-use pricing, and > provider-specific billing rules. > - Requests to models that are missing from the price table record token usage > but add nothing to a user's spend. A user who only calls unpriced models is @@ -236,7 +236,9 @@ Visibility follows the viewer's role: group budget. If the effective group is another group in the same organization, the row shows `Budget managed by another group`. If the effective group is in a different organization, the row shows a dash and explains that the group is not - visible there. + visible there. Because the organization's `Everyone` group includes every + member, its **Members** tab is a quick way to look up any user's effective + group. - The avatar menu reports the signed-in user's own spend for the budget period as `$ / $ USD`, or `$ / Unlimited USD` when no budget applies. @@ -247,7 +249,7 @@ endpoint to see a user's current effective group. ### CSV Export -Users who can read group-member data for the organization can export estimated +Users who can read group-member data for the organization can export approximate spend for reporting and internal cost allocation. The export is available through the API only. @@ -294,7 +296,7 @@ Expect the following differences: Control applied the lowest. - Budgets cover priced AI Gateway traffic. Chat, IDE extensions, and CLI agents draw on the same budget when their provider and model are priced. See - [How spend is estimated](#how-spend-is-estimated). + [How spend is approximated](#how-spend-is-approximated). - Recorded spend does not carry over. Every user starts the first period at $0 USD. - Coder Agents users who exceed their budget see a usage limit error in chat. From 58ccbef7557cad28ff949dde2dff5b0bbaa7154c Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Tue, 11 Aug 2026 08:36:53 +0000 Subject: [PATCH 3/3] docs(ai-coder/ai-gateway): rename spend section to How spend is calculated --- docs/ai-coder/ai-gateway/cost-controls.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/ai-coder/ai-gateway/cost-controls.md b/docs/ai-coder/ai-gateway/cost-controls.md index 4ab0edca51c..bd55dd4f4db 100644 --- a/docs/ai-coder/ai-gateway/cost-controls.md +++ b/docs/ai-coder/ai-gateway/cost-controls.md @@ -19,7 +19,7 @@ AI Governance Cost Control requires: > [!NOTE] > AI Governance Cost Control reports approximate spend rather than billed cost. > These figures will not match your provider invoices exactly. For details, see -> [How spend is approximated](#how-spend-is-approximated). +> [How spend is calculated](#how-spend-is-calculated). These terms appear throughout this page and in the Coder dashboard: @@ -181,7 +181,7 @@ affected user: For delivery methods, see [Notifications](../../admin/monitoring/notifications/index.md). -## How spend is approximated +## How spend is calculated Coder multiplies the token usage of each request by the published price of the model that served it. Prices come from a curated [models.dev](https://models.dev) @@ -203,7 +203,7 @@ Replace `` with your Coder minor version, for example `2.36`. > Approximate spend can differ from provider-reported amounts, and some usage > might not count toward spend: > -> - Approximations exclude negotiated discounts, committed-use pricing, and +> - Approximate spend excludes negotiated discounts, committed-use pricing, and > provider-specific billing rules. > - Requests to models that are missing from the price table record token usage > but add nothing to a user's spend. A user who only calls unpriced models is @@ -296,7 +296,7 @@ Expect the following differences: Control applied the lowest. - Budgets cover priced AI Gateway traffic. Chat, IDE extensions, and CLI agents draw on the same budget when their provider and model are priced. See - [How spend is approximated](#how-spend-is-approximated). + [How spend is calculated](#how-spend-is-calculated). - Recorded spend does not carry over. Every user starts the first period at $0 USD. - Coder Agents users who exceed their budget see a usage limit error in chat.