From e91bec85748189657d1e7e006d44c8aa4973622d Mon Sep 17 00:00:00 2001 From: Ethan Date: Wed, 27 May 2026 17:33:14 +1000 Subject: [PATCH] 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 | 9 ++++++++- cli/server_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/cli/server.go b/cli/server.go index c8ff2357f5..3fed22aeb5 100644 --- a/cli/server.go +++ b/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") diff --git a/cli/server_test.go b/cli/server_test.go index 5215eeb08c..89e0ba7048 100644 --- a/cli/server_test.go +++ b/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" {