test: stabilize gateway server shard (#77131)

src/gateway/server-http.request-trace.test.ts

@@ -88,9 +88,13 @@ describe("gateway HTTP request trace scope", () => {
       expect(activeTraceInHandler?.spanId).toMatch(/^[0-9a-f]{16}$/);
       expect(events).toEqual([{ trace: activeTraceInHandler, type: "message.queued" }]);
 
-      const [line] = fs.readFileSync(logPath, "utf8").trim().split("\n");
-      const record = JSON.parse(line ?? "{}") as Record<string, unknown>;
-      expect(record).toMatchObject({
+      const traceRecord = fs
+        .readFileSync(logPath, "utf8")
+        .trim()
+        .split("\n")
+        .map((line) => JSON.parse(line) as Record<string, unknown>)
+        .find((record) => record.message === "handled request trace");
+      expect(traceRecord).toMatchObject({
         traceId: activeTraceInHandler?.traceId,
         spanId: activeTraceInHandler?.spanId,
       });

src/gateway/server-network-runtime.e2e.test.ts

@@ -38,6 +38,14 @@ function isEnvHttpProxyDispatcher(dispatcher: unknown): boolean {
   );
 }
 
+async function closeTestDispatcher(dispatcher: unknown): Promise<void> {
+  const close = (dispatcher as { close?: () => Promise<void> | void } | undefined)?.close;
+  if (typeof close !== "function") {
+    return;
+  }
+  await close.call(dispatcher);
+}
+
 describe("gateway network runtime", () => {
   beforeEach(() => {
     clearRuntimeConfigSnapshot();
@@ -64,7 +72,8 @@ describe("gateway network runtime", () => {
     let server: Awaited<ReturnType<typeof startGatewayServer>> | undefined;
 
     try {
-      setGlobalDispatcher(new Agent());
+      const testDispatcher = new Agent();
+      setGlobalDispatcher(testDispatcher);
       for (const key of NETWORK_GATEWAY_ENV_KEYS) {
         delete process.env[key];
       }
@@ -101,7 +110,11 @@ describe("gateway network runtime", () => {
       expect(isEnvHttpProxyDispatcher(getGlobalDispatcher())).toBe(true);
     } finally {
       await server?.close({ reason: "gateway proxy bootstrap test complete" });
+      const dispatcherToClose = getGlobalDispatcher();
       setGlobalDispatcher(originalDispatcher);
+      if (dispatcherToClose !== originalDispatcher) {
+        await closeTestDispatcher(dispatcherToClose);
+      }
       await fs.rm(tempHome, { recursive: true, force: true });
       envSnapshot.restore();
     }

src/gateway/shared-auth.test-helpers.ts

@@ -2,10 +2,42 @@ import { expect } from "vitest";
 import { WebSocket } from "ws";
 import { connectOk, rpcReq, trackConnectChallengeNonce } from "./test-helpers.js";
 
-export async function openAuthenticatedGatewayWs(port: number, token: string): Promise<WebSocket> {
+export async function openAuthenticatedGatewayWs(
+  port: number,
+  token: string,
+  timeoutMs = 10_000,
+): Promise<WebSocket> {
   const ws = new WebSocket(`ws://127.0.0.1:${port}`);
   trackConnectChallengeNonce(ws);
-  await new Promise<void>((resolve) => ws.once("open", resolve));
+  await new Promise<void>((resolve, reject) => {
+    const cleanup = () => {
+      clearTimeout(timer);
+      ws.off("open", onOpen);
+      ws.off("error", onError);
+      ws.off("close", onClose);
+    };
+    const onOpen = () => {
+      cleanup();
+      resolve();
+    };
+    const onError = (error: unknown) => {
+      cleanup();
+      reject(error instanceof Error ? error : new Error(String(error)));
+    };
+    const onClose = (code: number, reason: Buffer) => {
+      cleanup();
+      reject(new Error(`gateway websocket closed before open (${code}: ${reason.toString()})`));
+    };
+    const timer = setTimeout(() => {
+      cleanup();
+      ws.close();
+      reject(new Error(`gateway websocket did not open within ${timeoutMs}ms`));
+    }, timeoutMs);
+    timer.unref?.();
+    ws.once("open", onOpen);
+    ws.once("error", onError);
+    ws.once("close", onClose);
+  });
   await connectOk(ws, { token });
   return ws;
 }
@@ -17,8 +49,11 @@ export async function waitForGatewayWsClose(
   return await new Promise((resolve, reject) => {
     const timer = setTimeout(() => {
       ws.off("close", onClose);
-      reject(new Error(`gateway websocket did not close within ${timeoutMs}ms`));
+      reject(
+        new Error(`gateway websocket did not close within ${timeoutMs}ms (state=${ws.readyState})`),
+      );
     }, timeoutMs);
+    timer.unref?.();
     const onClose = (code: number, reason: Buffer) => {
       clearTimeout(timer);
       resolve({ code, reason: reason.toString() });