fix(gateway): clarify websocket handshake phases

849261680 · vincentkoc · commit 022581d9bc68 · 2026-07-07T20:15:14.000-07:00
diff --git a/src/gateway/server/ws-connection.test.ts b/src/gateway/server/ws-connection.test.ts
@@ -262,7 +262,7 @@ describe("attachGatewayWsConnectionHandler", () => {
     ).toBe(true);
     handlerParams.setHandshakeState("connected");
     handlerParams.advanceHandshakePhase("session_attached");
-    handlerParams.advanceHandshakePhase("subscriptions_registered");
+    handlerParams.advanceHandshakePhase("hello_payload_prepared");
     handlerParams.advanceHandshakePhase("ready");
 
     socket.emit("close", 1000, Buffer.from("done"));
diff --git a/src/gateway/server/ws-connection/message-handler.post-connect-health.test.ts b/src/gateway/server/ws-connection/message-handler.post-connect-health.test.ts
@@ -212,6 +212,7 @@ function attachGatewayHarness(options: {
     mode: "none",
     allowTailscale: false,
   };
+  const advanceHandshakePhase = vi.fn();
   attachGatewayWsMessageHandler({
     socket,
     upgradeReq: {
@@ -244,7 +245,7 @@ function attachGatewayHarness(options: {
       return true;
     },
     setHandshakeState: vi.fn(),
-    advanceHandshakePhase: vi.fn(),
+    advanceHandshakePhase,
     setCloseCause: options.setCloseCause ?? createSetCloseCauseMock(),
     setLastFrameMeta: vi.fn(),
     originCheckMetrics: { hostHeaderFallbackAccepted: 0 },
@@ -257,6 +258,7 @@ function attachGatewayHarness(options: {
   }
   const sendMessage = onMessage;
   return {
+    advanceHandshakePhase,
     socketSend,
     sendRequest: (id: string, method: string, params: Record<string, unknown> = {}) => {
       sendMessage(
@@ -585,6 +587,46 @@ describe("attachGatewayWsMessageHandler post-connect health refresh", () => {
     expect(JSON.stringify(captured.events)).not.toContain("gateway-token");
   });
 
+  it("records credential and hello preparation phases during connect", async () => {
+    const harness = attachGatewayHarness({
+      connId: "conn-phases",
+      connectNonce: "nonce-phases",
+      resolvedAuth: {
+        mode: "token",
+        token: "gateway-token",
+        allowTailscale: false,
+      },
+    });
+
+    harness.sendConnect("connect-phases", {
+      minProtocol: PROTOCOL_VERSION,
+      maxProtocol: PROTOCOL_VERSION,
+      client: {
+        id: "gateway-client",
+        version: "dev",
+        platform: "test",
+        mode: "backend",
+      },
+      role: "operator",
+      scopes: [],
+      caps: [],
+      auth: {
+        token: "gateway-token",
+      },
+    });
+
+    await vi.waitFor(() => {
+      expect(harness.socketSend).toHaveBeenCalled();
+    });
+    expect(harness.advanceHandshakePhase.mock.calls.map(([phase]) => phase)).toEqual([
+      "auth_credentials_received",
+      "auth_validated",
+      "session_attached",
+      "hello_payload_prepared",
+      "ready",
+    ]);
+  });
+
   it("does not mark local backend self-pairing clients as approval runtimes", async () => {
     const refreshHealthSnapshot = vi.fn<GatewayRequestContext["refreshHealthSnapshot"]>(async () =>
       createHealthSummary(),
diff --git a/src/gateway/server/ws-connection/message-handler.ts b/src/gateway/server/ws-connection/message-handler.ts
@@ -930,6 +930,14 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar
           deviceRaw,
         });
         const device = controlUiAuthPolicy.device;
+        const hasRawHandshakeCredentials =
+          hasSharedAuth ||
+          Boolean(connectParams.auth?.bootstrapToken) ||
+          Boolean(connectParams.auth?.deviceToken) ||
+          Boolean(device);
+        if (hasRawHandshakeCredentials) {
+          advanceHandshakePhase("auth_credentials_received");
+        }
         const connectAuthState = await resolveConnectAuthState({
           resolvedAuth,
           connectAuth: connectParams.auth,
@@ -946,9 +954,6 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar
           deviceTokenCandidate,
           deviceTokenCandidateSource,
         } = connectAuthState;
-        if (hasSharedAuth || bootstrapTokenCandidate || deviceTokenCandidate || device) {
-          advanceHandshakePhase("auth_credentials_received");
-        }
         let { authResult, authOk, authMethod } = connectAuthState;
         const rejectUnauthorized = (failedAuth: GatewayAuthResult) => {
           const { authProvided, canRetryWithDeviceToken, recommendedNextStep } =
@@ -2153,7 +2158,6 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar
             );
         }
 
-        advanceHandshakePhase("subscriptions_registered");
         const snapshot = buildGatewaySnapshot({
           includeSensitive: scopes.includes(ADMIN_SCOPE),
         });
@@ -2198,6 +2202,7 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar
             tickIntervalMs: TICK_INTERVAL_MS,
           },
         };
+        advanceHandshakePhase("hello_payload_prepared");
 
         let revokedBootstrapTokenRecord:
           | Awaited<ReturnType<typeof revokeDeviceBootstrapToken>>["record"]
diff --git a/src/gateway/server/ws-types.ts b/src/gateway/server/ws-types.ts
@@ -33,7 +33,7 @@ export const WS_HANDSHAKE_PHASES = [
   "auth_credentials_received",
   "auth_validated",
   "session_attached",
-  "subscriptions_registered",
+  "hello_payload_prepared",
   "ready",
 ] as const;
 
