fix(cli): close aibridge daemon before WebSocket shutdown wait (#25719)

> [!WARNING]
> The investigation and solution in this PR were done with
[Mux](https://mux.coder.com/). I've reviewed the investigation
methodology, evidence and solution, and it all appears sound.

## Summary

PR #25570 (`refactor: move aibridged out of enterprise to AGPL`, merged
2026-05-22) added an in-memory aibridge DRPC server in
`coderd/aibridged.go` that does `api.WebsocketWaitGroup.Add(1)` and only
releases `Done()` when its client session is closed. PR #25575 then
flipped `CODER_AI_GATEWAY_ENABLED` to default to `true`, so every
`cli.Server()` invocation now spins up that goroutine.

In `cli/server.go`, the only call to `aibridgeDaemon.Close()` was a
`defer` scheduled at function return. During graceful shutdown the code
first calls `coderAPICloser.Close()`, which waits on
`api.WebsocketWaitGroup`. That wait sits for the full 10s timeout in
`coderd/coderd.go` (`websocket shutdown timed out after 10 seconds`),
then returns, then the function unwinds, and only then does the deferred
`aibridgeDaemon.Close()` fire and let the goroutine call `Done()`.

The 10s tax was previously latent (aibridged was enterprise-only and
opt-in). After the two May 22 PRs it hit every `cli.Server()` test. On
Linux/macOS CI it just makes the suite slower; on the Depot Windows
runner, the ramdisk reservation leaves only ~17 GiB of headroom and the
~10s shutdown tails of multiple concurrent package binaries overlap into
an OOM, presenting as `test-go-pg (windows-2022)` jobs that die silently
at the ~600s watchdog with an empty `steps` array.

See Slack:
https://codercom.slack.com/archives/C05AE94121Z/p1779807717764189

## Fix

Close `aibridgeDaemon` explicitly during graceful shutdown, **before**
`coderAPICloser.Close()` waits on the WebSocket wait group. This matches
the existing ordered-shutdown pattern used for `tunnel` and
`notificationsManager`. The deferred `aibridgeDaemon.Close()` is
retained as a safety net for early-return paths, and is safe to
double-call because `aibridged.Server.Close()` is already idempotent via
`shutdownOnce` in `coderd/aibridged/aibridged.go`.

## Regression test

`TestServer_AIGatewayShutdownOrdering` boots a real `coder server` with
`--ai-gateway-enabled=true`, cancels its context, and asserts graceful
shutdown finishes in under 8s. With the fix the test runs in ~0.1s;
without the fix it fails deterministically at ~10.0s. The flag is passed
explicitly so the test continues to guard the ordering even if the
deployment default is ever flipped back.

## Evidence this fixes the OOM

On Linux the patched `cli` test package drops from 114 s back to its
pre-regression 30 s wall time at the same single-process peak RSS (~7.6
GiB), and the `websocket shutdown timed out after 10 seconds` log line
disappears from every server-test run. Since the Windows OOM is the sum
of multiple concurrent 10 s shutdown tails overlapping past the runner's
~17 GiB headroom, removing those tails returns the concurrent-RSS budget
to its pre-regression level. The Windows OOM was intermittent (a handful
of hits across many runs since May 22), so a single green `test-go-pg
(windows-2022)` job on this PR is not by itself proof; confirmation will
come from watching Windows runs on `main` over the next several days and
seeing the ~600 s silent-kill fingerprint stop recurring.

Relates to ENG-2771

cli/server.go

@@ -63,6 +63,7 @@ import (
 	"github.com/coder/coder/v2/cli/cliutil"
 	"github.com/coder/coder/v2/cli/config"
 	"github.com/coder/coder/v2/coderd"
+	"github.com/coder/coder/v2/coderd/aibridged"
 	"github.com/coder/coder/v2/coderd/autobuild"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/awsiamrds"
@@ -1014,6 +1015,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			if err != nil {
 				return xerrors.Errorf("create coder API: %w", err)
 			}
+			var aibridgeDaemon *aibridged.Server
 
 			// Both seed (writes) and build (reads) of AI providers need
 			// options.Database to be dbcrypt-wrapped, which only happens
@@ -1044,7 +1046,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				if err != nil {
 					return xerrors.Errorf("build AI providers: %w", err)
 				}
-				aibridgeDaemon, err := newAIBridgeDaemon(coderAPI, aibridgeProviders)
+				aibridgeDaemon, err = newAIBridgeDaemon(coderAPI, aibridgeProviders)
 				if err != nil {
 					return xerrors.Errorf("create aibridged: %w", err)
 				}
@@ -1310,6 +1312,11 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			}
 			wg.Wait()
 
+			// The in-memory aibridge server participates in the websocket
+			// wait group, so close its client before waiting for that group.
+			if aibridgeDaemon != nil {
+				_ = aibridgeDaemon.Close()
+			}
 			cliui.Info(inv.Stdout, "Waiting for WebSocket connections to close..."+"\n")
 			_ = coderAPICloser.Close()
 			cliui.Info(inv.Stdout, "Done waiting for WebSocket connections"+"\n")

cli/server_test.go

@@ -2184,6 +2184,53 @@ func TestServer_InterruptShutdown(t *testing.T) {
 	require.NoError(t, err)
 }
 
+// TestServer_AIGatewayShutdownOrdering is a regression test for a shutdown
+// ordering bug. The in-memory AI Gateway daemon registers itself with the
+// API WebsocketWaitGroup, so it must be closed before coderAPICloser.Close()
+// waits on that group. If it isn't, API.Close() blocks for the full 10s
+// WebsocketWaitGroup timeout, logs "websocket shutdown timed out after 10
+// seconds", and keeps heavy server-test state live for an extra 10s. On
+// Windows test-go-pg this extra shutdown tail overlapped across concurrent
+// package binaries and OOMed the runner.
+func TestServer_AIGatewayShutdownOrdering(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitLong))
+	defer cancel()
+
+	inv, cfg := clitest.New(t,
+		"server",
+		dbArg(t),
+		"--http-address", ":0",
+		"--access-url", "http://example.com",
+		"--cache-dir", t.TempDir(),
+		// Explicit so the test catches the regression even if the
+		// default for ai-gateway-enabled is ever flipped back to false.
+		"--ai-gateway-enabled=true",
+	)
+
+	serverErr := make(chan error, 1)
+	go func() {
+		serverErr <- inv.WithContext(ctx).Run()
+	}()
+
+	// Wait for the server to come up so the in-memory AI Gateway daemon
+	// is registered with the API and the WebsocketWaitGroup is nonzero.
+	_ = waitAccessURL(t, cfg)
+
+	// The WebsocketWaitGroup timeout in coderd.API.Close() is hard coded
+	// to 10s, so any value comfortably below 10s catches the regression
+	// while leaving headroom for slow CI runners.
+	shutdownStart := time.Now()
+	cancel()
+	if err := <-serverErr; err != nil {
+		require.ErrorIs(t, err, context.Canceled)
+	}
+	require.Less(t, time.Since(shutdownStart), 8*time.Second,
+		"graceful shutdown took too long; the in-memory AI Gateway daemon is "+
+			"likely not being closed before coderAPICloser.Close()")
+}
+
 func TestServer_GracefulShutdown(t *testing.T) {
 	t.Parallel()
 	if runtime.GOOS == "windows" {