test(enterprise/cli): add standalone AI Gateway tests against a live coderd (#27863)

Stacked on #27860. Adds four connection tests that run the real `ai-gateway start` against `coderdenttest` and assert only what an operator or LLM client can observe. A `chaosProxy` between the gateway and coderd simulates outages by answering 503 and closing the connections it accepted, the latter because the DRPC websocket is hijacked and so out of reach of `httptest.Server`. - `RevokedKey`: revoking an in-use key closes the session, and the 401 on redial terminates the command. - `ReconnectAfterDisconnect`: LLM traffic and interception recording resume after a coderd outage, with no intervention. - `RequestWhileDisconnected`: a request arriving while disconnected is parked until the connection returns, not failed. The pre-flight DRPC calls block with the caller's context as the only bound, which is intentional: a caller willing to wait is served on reconnect, and `/readyz` has already withdrawn the replica. The RFC's "fails with 503 if pre-flight DRPC calls cannot complete" does not describe this and needs correcting. - `InFlightRequestSurvivesDisconnect`: a stream whose first chunk already reached the caller completes after the DRPC connection drops. No production code is changed. Refs https://linear.app/codercom/issue/AIGOV-320/write-connection-tests --- Generated with Coder Agents.
2026-09-24 15:04:27 +08:00 · 2026-08-18 14:03:41 +00:00
parent 995d7fe31b
commit 0a34a37314
1 changed files with 496 additions and 56 deletions
@@ -3,16 +3,26 @@
 package cli_test

 import (
+	"bufio"
 	"context"
 	"io"
+	"log"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httputil"
+	"net/url"
 	"regexp"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
+	"time"

+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"

 	"github.com/coder/coder/v2/cli/clitest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -25,12 +35,11 @@ import (
 )

 // 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.
+// against a real coderd and only observe the public surface.
+//
+// Readiness state transitions, shutdown ordering, and the classification of
+// dial errors are asserted directly, using the internals of gateway struct in
+// aigatewaystart_internal_test.go.

 // aiGatewayChatCompletionRequest is the LLM request sent through the gateway. The
 // model is asserted against the recorded interception.
@@ -53,6 +62,20 @@ const aiGatewayUpstreamResponse = `{
  "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
 }`

+const (
+	aiGatewayStreamFirstChunk = `data: {"id":"chatcmpl-e2e","object":"chat.completion.chunk","created":1753343279,"model":"gpt-4.1","choices":[{"index":0,"delta":{"role":"assistant","content":"first half"},"finish_reason":null}]}` + "\n\n"
+	aiGatewayStreamFinalChunk = `data: {"id":"chatcmpl-e2e","object":"chat.completion.chunk","created":1753343279,"model":"gpt-4.1","choices":[{"index":0,"delta":{"content":" second half"},"finish_reason":"stop"}]}` + "\n\n" + "data: [DONE]\n\n"
+)
+
+const (
+	aiGatewayHealthzPath = "/healthz"
+	aiGatewayReadyzPath  = "/readyz"
+	// aiGatewayChatCompletionPath is the LLM route on the gateway's own
+	// listener. The aibridge mux routes by provider name, so the first segment
+	// is the name the provider is registered under.
+	aiGatewayChatCompletionPath = "/openai/v1/chat/completions"
+)
+
 // 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.
@@ -61,35 +84,78 @@ type aiGatewayDeployment struct {
 	userClient   *codersdk.Client
 	user         codersdk.User
 	key          string
+	keyID        uuid.UUID
 	upstreamHits *atomic.Int32
 }

-func setupAIGatewayCoderdenttestDeployment(t *testing.T) (*codersdk.Client, codersdk.CreateFirstUserResponse) {
+func setupAIGatewayCoderdenttestDeployment(t *testing.T, mutate ...func(*coderdenttest.Options)) (*codersdk.Client, codersdk.CreateFirstUserResponse) {
 	dv := coderdtest.DeploymentValues(t)
 	dv.AI.BridgeConfig.Enabled = serpent.Bool(true)
-	return coderdenttest.New(t, &coderdenttest.Options{
+	opts := &coderdenttest.Options{
 		Options: &coderdtest.Options{DeploymentValues: dv},
 		LicenseOptions: &coderdenttest.LicenseOptions{
 			Features: license.Features{codersdk.FeatureAIBridge: 1},
 		},
-	})
+	}
+	for _, m := range mutate {
+		m(opts)
+	}
+	return coderdenttest.New(t, opts)
 }

-func setupAIGatewayDeployment(ctx context.Context, t *testing.T) *aiGatewayDeployment {
+// aiGatewayDeploymentConfig is the deployment shape a test needs.
+type aiGatewayDeploymentConfig struct {
+	// upstream serves the mock LLM API the provider points at.
+	upstream http.HandlerFunc
+	// coderdOptions mutates coderd's options before it starts.
+	coderdOptions func(*coderdenttest.Options)
+}
+
+type aiGatewayDeploymentOption func(*aiGatewayDeploymentConfig)
+
+// withAIGatewayUpstream replaces the mock LLM API. Requests are still counted
+// in aiGatewayDeployment.upstreamHits.
+func withAIGatewayUpstream(handler http.HandlerFunc) aiGatewayDeploymentOption {
+	return func(cfg *aiGatewayDeploymentConfig) {
+		cfg.upstream = handler
+	}
+}
+
+// withAIGatewayCoderdOptions mutates coderd's options.
+func withAIGatewayCoderdOptions(mutate func(*coderdenttest.Options)) aiGatewayDeploymentOption {
+	return func(cfg *aiGatewayDeploymentConfig) {
+		cfg.coderdOptions = mutate
+	}
+}
+
+func setupAIGatewayDeployment(ctx context.Context, t *testing.T, opts ...aiGatewayDeploymentOption) *aiGatewayDeployment {
 	t.Helper()

-	client, firstUser := setupAIGatewayCoderdenttestDeployment(t)
+	cfg := aiGatewayDeploymentConfig{
+		upstream: func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(aiGatewayUpstreamResponse))
+		},
+	}
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
+	var coderdMutations []func(*coderdenttest.Options)
+	if cfg.coderdOptions != nil {
+		coderdMutations = append(coderdMutations, cfg.coderdOptions)
+	}
+	client, firstUser := setupAIGatewayCoderdenttestDeployment(t, coderdMutations...)

 	//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) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		hits.Add(1)
-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(http.StatusOK)
-		_, _ = w.Write([]byte(aiGatewayUpstreamResponse))
+		cfg.upstream(w, r)
 	}))
 	t.Cleanup(upstream.Close)

@@ -110,6 +176,7 @@ func setupAIGatewayDeployment(ctx context.Context, t *testing.T) *aiGatewayDeplo
 		userClient:   userClient,
 		user:         user,
 		key:          key.Key,
+		keyID:        key.ID,
 		upstreamHits: &hits,
 	}
 }
@@ -138,15 +205,38 @@ func startAIGatewayCommand(ctx context.Context, t *testing.T, coderURL, key stri
 	return "http://" + matches[1], waiter
 }

-func getAIGatewayStatus(ctx context.Context, t *testing.T, url string) int {
+// aiGatewayStatus probes probeURL and reports transport errors.
+func aiGatewayStatus(ctx context.Context, probeURL string) (int, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, probeURL, nil)
+	if err != nil {
+		return 0, err
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return 0, err
+	}
+	defer resp.Body.Close()
+	return resp.StatusCode, nil
+}
+
+// requireAIGatewayStatus asserts the status of a single probe.
+func requireAIGatewayStatus(ctx context.Context, t *testing.T, probeURL string, want int) {
 	t.Helper()

-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	got, err := aiGatewayStatus(ctx, probeURL)
 	require.NoError(t, err)
-	resp, err := http.DefaultClient.Do(req)
-	require.NoError(t, err)
-	defer resp.Body.Close()
-	return resp.StatusCode
+	require.Equal(t, want, got, "unexpected status for %s", probeURL)
+}
+
+// requireEventualAIGatewayStatus waits for probeURL to return want, tolerating
+// transport errors while the gateway converges.
+func requireEventualAIGatewayStatus(ctx context.Context, t *testing.T, probeURL string, want int) {
+	t.Helper()
+
+	require.Eventuallyf(t, func() bool {
+		got, err := aiGatewayStatus(ctx, probeURL)
+		return err == nil && got == want
+	}, testutil.WaitLong, testutil.IntervalFast, "%s never returned %d", probeURL, want)
 }

 // TestAIGatewayStartE2E drives every part of the standalone gateway plumbing
@@ -158,44 +248,34 @@ func TestAIGatewayStartE2E(t *testing.T) {
 	t.Parallel()

 	ctx := testutil.Context(t, testutil.WaitLong)
+	// Given: a coderd entitled for the AI Gateway, with a key and a provider.
 	dep := setupAIGatewayDeployment(ctx, t)

+	// When: the gateway starts against that coderd.
 	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)
+	// Then: liveness holds as soon as the listener is bound, and readiness
+	// follows the DRPC connection and the provider load.
+	requireAIGatewayStatus(ctx, t, baseURL+aiGatewayHealthzPath, http.StatusOK)
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)

-	// 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")
+	// When: a user sends an LLM request to the gateway.
+	result := postChatCompletionAndRead(ctx, baseURL+aiGatewayChatCompletionPath,
+		dep.userClient.SessionToken(), aiGatewayChatCompletionRequest)
+
+	// Then: the upstream's response reaches the caller.
+	require.NoError(t, result.err)
+	require.Equal(t, http.StatusOK, result.status, "body: %s", result.body)
+	require.Contains(t, string(result.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)
+	// Then: coderd records the interception, attributed to that user.
+	sessions := requireAIGatewaySessions(ctx, t, dep, 1)
+	require.Equal(t, dep.user.Username, sessions[0].Initiator.Username)

-	// Graceful shutdown: canceling the command must produce a clean exit.
+	// When: the command is canceled.
 	waiter.Cancel()
+	// Then: it exits cleanly.
 	require.NoError(t, waiter.Wait())
 }

@@ -208,6 +288,7 @@ func TestAIGatewayStartE2E_InvalidKey(t *testing.T) {
 	ctx := testutil.Context(t, testutil.WaitLong)
 	client, _ := setupAIGatewayCoderdenttestDeployment(t)

+	// When: the gateway starts with invalid key.
 	inv, _ := newCLI(t,
 		"ai-gateway", "start",
 		"--url", client.URL.String(),
@@ -216,7 +297,9 @@ func TestAIGatewayStartE2E_InvalidKey(t *testing.T) {
 	)
 	inv = inv.WithContext(ctx)
 	waiter := clitest.StartWithWaiter(t, inv)
-	waiter.RequireContains("AI Gateway key invalid")
+
+	// Then: it exits with error without retrying.
+	require.ErrorContains(t, waiter.Wait(), "AI Gateway key invalid")
 }

 // TestAIGatewayStart_HealthBeforeReady covers the split between liveness and
@@ -225,17 +308,374 @@ func TestAIGatewayStartE2E_InvalidKey(t *testing.T) {
 func TestAIGatewayStart_HealthBeforeReady(t *testing.T) {
 	t.Parallel()

-	// Fake coderd that answers 503 so the daemon keeps retrying to connect.
+	// Given: a coderd that always answers 503
 	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)
+	// When: the gateway starts and binds its listener.
 	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"))
+	// Then: healthz is already 200 while readyz stays 503.
+	// The startup log line is emitted after the bind, so no retry is needed.
+	requireAIGatewayStatus(ctx, t, baseURL+aiGatewayHealthzPath, http.StatusOK)
+	requireAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusServiceUnavailable)
+}
+
+// TestAIGatewayStartE2E_RevokedKey covers revocation of a key that is already
+// in use. coderd closes the active session, the gateway redials, and the 401
+// on that redial must terminate the command.
+func TestAIGatewayStartE2E_RevokedKey(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	// Given: a ready gateway, and a coderd whose key check ticker the test drives.
+	keyCheck := make(chan time.Time, 1)
+	dep := setupAIGatewayDeployment(ctx, t, withAIGatewayCoderdOptions(func(opts *coderdenttest.Options) {
+		opts.Options.NewTicker = func(time.Duration) (<-chan time.Time, func()) {
+			return keyCheck, func() {}
+		}
+	}))
+
+	baseURL, waiter := startAIGatewayCommand(ctx, t, dep.client.URL.String(), dep.key)
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)
+
+	// When: the key is deleted and coderd runs its check.
+	//nolint:gocritic // Owner role is needed for gateway key management.
+	require.NoError(t, dep.client.DeleteAIGatewayKey(ctx, dep.keyID))
+	keyCheck <- time.Now()
+
+	// Then: the 401 on the gateway's redial terminates the command.
+	require.ErrorContains(t, waiter.Wait(), "AI Gateway key invalid")
+}
+
+// connTrackingListener records the connections it accepts so a test can close
+// them while the listener keeps accepting new ones.
+//
+// It is needed because the gateway's connection to coderd is a websocket, and
+// httputil.ReverseProxy serves the resulting 101 by hijacking its inbound
+// connection and splicing bytes for the life of the upgrade. httptest.Server
+// deletes a connection from its tracking set on http.StateHijacked, so
+// CloseClientConnections can no longer reach it, and the handler that would
+// observe the proxy being unhealthy never runs again for an established
+// websocket. Accept happens before any of that, so the raw net.Conn kept here
+// stays closable.
+type connTrackingListener struct {
+	net.Listener
+
+	mu    sync.Mutex
+	conns []net.Conn
+}
+
+func (l *connTrackingListener) Accept() (net.Conn, error) {
+	conn, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.conns = append(l.conns, conn)
+	return conn, nil
+}
+
+// closeConns closes every connection accepted so far and stops tracking them,
+// so connections established afterwards survive. The listener itself stays
+// open.
+func (l *connTrackingListener) closeConns() {
+	l.mu.Lock()
+	conns := l.conns
+	l.conns = nil
+	l.mu.Unlock()
+
+	for _, conn := range conns {
+		_ = conn.Close()
+	}
+}
+
+// chaosProxy sits between the standalone gateway and coderd so tests can
+// simulate network failures and coderd unavailability. While unhealthy
+// it answers 503, which the connect loop classifies as transient and retries.
+type chaosProxy struct {
+	srv      *httptest.Server
+	listener *connTrackingListener
+	healthy  atomic.Bool
+}
+
+func newChaosProxy(t *testing.T, target *url.URL) *chaosProxy {
+	t.Helper()
+
+	p := &chaosProxy{}
+	p.healthy.Store(true)
+	reverse := httputil.NewSingleHostReverseProxy(target)
+	reverse.ErrorLog = log.New(testutil.NewTestLogWriter(t), "chaosProxy: ", 0)
+	p.srv = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !p.healthy.Load() {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+		reverse.ServeHTTP(w, r)
+	}))
+	p.listener = &connTrackingListener{Listener: p.srv.Listener}
+	p.srv.Listener = p.listener
+	p.srv.Start()
+	t.Cleanup(p.srv.Close)
+	return p
+}
+
+func (p *chaosProxy) setHealthy(healthy bool) {
+	p.healthy.Store(healthy)
+}
+
+// disconnect makes coderd both unreachable and unavailable. Closing the proxy's
+// connections takes the gateway's multiplexed websocket down with them.
+func (p *chaosProxy) disconnect() {
+	p.setHealthy(false)
+	p.listener.closeConns()
+}
+
+// aiGatewayResponse is the outcome of an LLM request sent through a gateway.
+type aiGatewayResponse struct {
+	status int
+	body   []byte
+	err    error
+}
+
+// postChatCompletion sends an LLM request and returns the response with its body
+// unread. The caller owns the body.
+//
+//nolint:bodyclose // The caller owns and closes the body.
+func postChatCompletion(ctx context.Context, endpoint, token, requestBody string) (*http.Response, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(requestBody))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", "Bearer "+token)
+	req.Header.Set("Content-Type", "application/json")
+
+	return http.DefaultClient.Do(req)
+}
+
+// postChatCompletionAndRead sends an LLM request, reads the whole response, and
+// reports transport errors.
+func postChatCompletionAndRead(ctx context.Context, endpoint, token, requestBody string) aiGatewayResponse {
+	resp, err := postChatCompletion(ctx, endpoint, token, requestBody)
+	if err != nil {
+		return aiGatewayResponse{err: err}
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	return aiGatewayResponse{status: resp.StatusCode, body: body, err: err}
+}
+
+// requireAIGatewaySessions waits for coderd to have recorded 'want' sessions for
+// the deployment's user and returns them.
+func requireAIGatewaySessions(ctx context.Context, t *testing.T, dep *aiGatewayDeployment, want int) []codersdk.AIBridgeSession {
+	t.Helper()
+
+	var (
+		sessions []codersdk.AIBridgeSession
+		lastErr  error
+	)
+	require.Eventuallyf(t, func() bool {
+		//nolint:gocritic // Owner (or Auditor) role is needed to read sessions.
+		resp, err := dep.client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{
+			Initiator: dep.user.Username,
+		})
+		lastErr = err
+		if err != nil {
+			return false
+		}
+		sessions = resp.Sessions
+		return len(sessions) == want
+	}, testutil.WaitLong, testutil.IntervalFast,
+		"expected %d recorded session(s), got %d, last error: %v", want, len(sessions), lastErr)
+	return sessions
+}
+
+// TestAIGatewayStartE2E_ReconnectAfterDisconnect covers a coderd outage: the
+// gateway serves LLM traffic again once coderd returns, without any operator
+// intervention, which requires the provider cache and the recorder to recover
+// alongside the DRPC connection.
+func TestAIGatewayStartE2E_ReconnectAfterDisconnect(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	dep := setupAIGatewayDeployment(ctx, t)
+
+	// Given: a ready gateway, reaching coderd through a chaos proxy.
+	proxy := newChaosProxy(t, dep.client.URL)
+	baseURL, _ := startAIGatewayCommand(ctx, t, proxy.srv.URL, dep.key)
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)
+
+	// When: a request arrives while coderd is reachable.
+	before := postChatCompletionAndRead(ctx, baseURL+aiGatewayChatCompletionPath,
+		dep.userClient.SessionToken(), aiGatewayChatCompletionRequest)
+
+	// Then: it is served and recorded.
+	require.NoError(t, before.err)
+	require.Equal(t, http.StatusOK, before.status, "body: %s", before.body)
+	requireAIGatewaySessions(ctx, t, dep, 1)
+
+	// When: coderd becomes unreachable.
+	proxy.disconnect()
+
+	// Then: readiness withdraws.
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusServiceUnavailable)
+
+	// When: coderd is reachable again.
+	proxy.setHealthy(true)
+
+	// Then: readiness recovers with no intervention.
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)
+
+	// Then: a new request is served and both interceptions are recorded, so the
+	// provider cache and the recorder recovered with the connection.
+	after := postChatCompletionAndRead(ctx, baseURL+aiGatewayChatCompletionPath,
+		dep.userClient.SessionToken(), aiGatewayChatCompletionRequest)
+	require.NoError(t, after.err)
+	require.Equal(t, http.StatusOK, after.status, "body: %s", after.body)
+	require.Contains(t, string(after.body), "standalone gateway e2e response")
+	require.Equal(t, int32(2), dep.upstreamHits.Load())
+	requireAIGatewaySessions(ctx, t, dep, 2)
+}
+
+// TestAIGatewayStartE2E_RequestWhileDisconnected pins the behavior of an LLM
+// request that arrives while the gateway has no DRPC connection to coderd. The
+// request is parked until the connection returns rather than failing fast,
+// because the pre-flight calls block in [aibridged.Server.Client] with the
+// caller's context as the only bound.
+func TestAIGatewayStartE2E_RequestWhileDisconnected(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	dep := setupAIGatewayDeployment(ctx, t)
+
+	// Given: a ready gateway that then loses its connection to coderd.
+	proxy := newChaosProxy(t, dep.client.URL)
+	baseURL, _ := startAIGatewayCommand(ctx, t, proxy.srv.URL, dep.key)
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)
+	proxy.disconnect()
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusServiceUnavailable)
+
+	// When: a request arrives while the gateway has no connection to coderd.
+	responses := make(chan aiGatewayResponse, 1)
+	go func() {
+		responses <- postChatCompletionAndRead(ctx, baseURL+aiGatewayChatCompletionPath,
+			dep.userClient.SessionToken(), aiGatewayChatCompletionRequest)
+	}()
+
+	// Then: it does not fail fast, and it does not reach the upstream before it
+	// is authorized which requires a connection to coderd.
+	select {
+	case result := <-responses:
+		t.Fatalf("request completed while disconnected: status=%d, err=%v", result.status, result.err)
+	case <-time.After(testutil.IntervalMedium):
+	}
+	require.Equal(t, int32(0), dep.upstreamHits.Load(), "a request must not reach the upstream before it is authorized")
+
+	// When: coderd returns. Then: the parked request is served.
+	proxy.setHealthy(true)
+	result := testutil.RequireReceive(ctx, t, responses)
+	require.NoError(t, result.err)
+	require.Equal(t, http.StatusOK, result.status, "body: %s", result.body)
+	require.Contains(t, string(result.body), "standalone gateway e2e response")
+	require.Equal(t, int32(1), dep.upstreamHits.Load(), "the parked request reaches the upstream once it is authorized")
+}
+
+// readAIGatewayStreamEvent returns the payload of the next server-sent event,
+// blocking until the gateway forwards it.
+func readAIGatewayStreamEvent(t *testing.T, stream *bufio.Reader) string {
+	t.Helper()
+
+	for {
+		line, err := stream.ReadString('\n')
+		require.NoError(t, err, "read stream event")
+		payload, ok := strings.CutPrefix(strings.TrimSuffix(line, "\n"), "data: ")
+		if !ok {
+			// Blank line separating events.
+			continue
+		}
+		return payload
+	}
+}
+
+// readAIGatewayStreamDelta returns the content delta of the next event.
+func readAIGatewayStreamDelta(t *testing.T, stream *bufio.Reader) string {
+	t.Helper()
+
+	payload := readAIGatewayStreamEvent(t, stream)
+	return gjson.Get(payload, "choices.0.delta.content").String()
+}
+
+// blockingStreamAIGatewayUpstream returns a mock upstream that starts a
+// server-sent event stream, blocks until release is closed, then finishes the
+// stream. It lets a test drop the DRPC connection while a request is in flight,
+// after its pre-flight authorization has already succeeded and after the caller
+// has received part of the response.
+func blockingStreamAIGatewayUpstream(release <-chan struct{}) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		flusher, ok := w.(http.Flusher)
+		if !ok {
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.Header().Set("Cache-Control", "no-cache")
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(aiGatewayStreamFirstChunk))
+		flusher.Flush()
+
+		select {
+		case <-release:
+		case <-r.Context().Done():
+			return
+		}
+		_, _ = w.Write([]byte(aiGatewayStreamFinalChunk))
+		flusher.Flush()
+	}
+}
+
+// TestAIGatewayStartE2E_InFlightRequestSurvivesDisconnect covers a request that
+// has already passed pre-flight authorization and had part of its response
+// delivered when the DRPC connection to coderd drops. The rest of the stream
+// must reach the caller instead of being torn down along with the connection.
+func TestAIGatewayStartE2E_InFlightRequestSurvivesDisconnect(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	release := make(chan struct{})
+	dep := setupAIGatewayDeployment(ctx, t, withAIGatewayUpstream(blockingStreamAIGatewayUpstream(release)))
+
+	// Given: a ready gateway, and a streaming request whose first chunk the
+	// caller has already received.
+	proxy := newChaosProxy(t, dep.client.URL)
+	baseURL, _ := startAIGatewayCommand(ctx, t, proxy.srv.URL, dep.key)
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusOK)
+
+	req, err := sjson.Set(aiGatewayChatCompletionRequest, "stream", true)
+	require.NoError(t, err)
+	resp, err := postChatCompletion(ctx, baseURL+aiGatewayChatCompletionPath,
+		dep.userClient.SessionToken(), req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	stream := bufio.NewReader(resp.Body)
+	require.Equal(t, "first half", readAIGatewayStreamDelta(t, stream))
+
+	// When: the DRPC connection drops and the upstream finishes the stream.
+	proxy.disconnect()
+	requireEventualAIGatewayStatus(ctx, t, baseURL+aiGatewayReadyzPath, http.StatusServiceUnavailable)
+	close(release)
+
+	// Then: the rest of the stream reaches the caller, in order.
+	require.Equal(t, " second half", readAIGatewayStreamDelta(t, stream))
+	require.Equal(t, "[DONE]", readAIGatewayStreamEvent(t, stream), "the stream must be terminated")
+	require.Equal(t, int32(1), dep.upstreamHits.Load(), "the request must reach the upstream exactly once")
+
+	// The recording RPCs that follow the response share the connection that was
+	// dropped, so this interception's usage rows are expected to be lost. Only
+	// the caller-visible outcome is asserted.
 }