fix(gateway): drain admitted node results before disconnect (#119246)

* fix(gateway): drain admitted node results on close

* fix(gateway): keep idle node cleanup synchronous

* fix(gateway): preserve node revocation during drain

* refactor(gateway): keep websocket handler within line budget

src/gateway/server.node-result-close.test.ts

@@ -0,0 +1,219 @@
+// Node result/close ordering tests keep admitted terminal frames authoritative.
+import { randomUUID } from "node:crypto";
+import { afterEach, expect, test, vi } from "vitest";
+import { writeConfigFile } from "../config/config.js";
+import { approveNodePairing, requestNodePairing } from "../infra/node-pairing.js";
+import { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES } from "../utils/message-channel.js";
+import { pairDeviceIdentity } from "./device-authz.test-helpers.js";
+import { GatewayNodeLifecycleDispatchTracker } from "./server/ws-connection/node-lifecycle-dispatch.js";
+import { connectGatewayClient } from "./test-helpers.e2e.js";
+import { installGatewayTestHooks, startServer } from "./test-helpers.js";
+
+const pairingRead = vi.hoisted(() => ({
+  blocked: null as Promise<void> | null,
+  onBlocked: null as (() => void) | null,
+  release: null as (() => void) | null,
+}));
+
+vi.mock("../infra/node-pairing-state.js", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("../infra/node-pairing-state.js")>();
+  return {
+    ...actual,
+    resolveCurrentNodePairingBinding: async (nodeId: string) => {
+      const current = await actual.resolveCurrentNodePairingBinding(nodeId);
+      if (pairingRead.blocked) {
+        pairingRead.onBlocked?.();
+        await pairingRead.blocked;
+      }
+      return current;
+    },
+  };
+});
+
+installGatewayTestHooks({ scope: "suite" });
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  pairingRead.blocked = null;
+  pairingRead.onBlocked = null;
+  pairingRead.release = null;
+});
+
+test.each([
+  ["a terminal node result admitted before close wins over disconnect cleanup", false],
+  ["pairing removal still fences a terminal node result while close drains", true],
+] as const)("%s", async (_name, removePairingDuringDrain) => {
+  const pairedNode = await pairDeviceIdentity({
+    name: "node-result-before-close",
+    role: "node",
+    scopes: [],
+    clientId: GATEWAY_CLIENT_NAMES.NODE_HOST,
+    clientMode: GATEWAY_CLIENT_MODES.NODE,
+  });
+  const pairing = await requestNodePairing({
+    nodeId: pairedNode.identity.deviceId,
+    platform: "linux",
+    deviceFamily: "Linux",
+    commands: ["camera.list"],
+  });
+  await approveNodePairing(pairing.request.requestId, {
+    callerScopes: ["operator.pairing", "operator.write"],
+  });
+  await writeConfigFile({
+    gateway: { nodes: { commands: { allow: ["camera.list"] } } },
+  });
+
+  const { port, server } = await startServer("secret");
+  const url = `ws://127.0.0.1:${port}`;
+  let operator: Awaited<ReturnType<typeof connectGatewayClient>> | undefined;
+  let node: Awaited<ReturnType<typeof connectGatewayClient>> | undefined;
+  let resolveInvokeFrame:
+    | ((frame: { id: string; nodeId: string; command: string }) => void)
+    | undefined;
+  const invokeFrame = new Promise<{ id: string; nodeId: string; command: string }>((resolve) => {
+    resolveInvokeFrame = resolve;
+  });
+
+  try {
+    operator = await connectGatewayClient({
+      url,
+      token: "secret",
+      role: "operator",
+      clientName: GATEWAY_CLIENT_NAMES.GATEWAY_CLIENT,
+      clientDisplayName: "node result close operator",
+      mode: GATEWAY_CLIENT_MODES.BACKEND,
+      scopes: ["operator.admin", "operator.read", "operator.write"],
+    });
+    node = await connectGatewayClient({
+      url,
+      token: "secret",
+      role: "node",
+      clientName: GATEWAY_CLIENT_NAMES.NODE_HOST,
+      clientDisplayName: "node result close host",
+      mode: GATEWAY_CLIENT_MODES.NODE,
+      platform: "linux",
+      deviceFamily: "Linux",
+      scopes: [],
+      commands: ["camera.list"],
+      deviceIdentity: pairedNode.identity,
+      onEvent: (event) => {
+        if (event.event !== "node.invoke.request" || !event.payload) {
+          return;
+        }
+        resolveInvokeFrame?.(event.payload as { id: string; nodeId: string; command: string });
+      },
+    });
+    await vi.waitFor(async () => {
+      const listed = await operator?.request<{
+        nodes?: Array<{ nodeId?: string; connected?: boolean; commands?: string[] }>;
+      }>("node.list", {}, { timeoutMs: 10_000 });
+      expect(listed?.nodes?.find((entry) => entry.nodeId === pairedNode.identity.deviceId)).toEqual(
+        expect.objectContaining({
+          connected: true,
+          commands: ["camera.list"],
+        }),
+      );
+    });
+
+    const invoked = operator.request<{
+      ok: boolean;
+      nodeId: string;
+      command: string;
+      payload: unknown;
+    }>(
+      "node.invoke",
+      {
+        nodeId: pairedNode.identity.deviceId,
+        command: "camera.list",
+        timeoutMs: 10_000,
+        idempotencyKey: randomUUID(),
+      },
+      { timeoutMs: 10_000 },
+    );
+    const frame = await Promise.race([
+      invokeFrame,
+      invoked.then(
+        () => {
+          throw new Error("node.invoke settled without sending a node request");
+        },
+        (error: unknown) => {
+          throw error instanceof Error ? error : new Error(String(error));
+        },
+      ),
+    ]);
+
+    pairingRead.blocked = new Promise<void>((resolve) => {
+      pairingRead.release = resolve;
+    });
+    const pairingReadStarted = new Promise<void>((resolve) => {
+      pairingRead.onBlocked = resolve;
+    });
+    const drainSpy = vi.spyOn(GatewayNodeLifecycleDispatchTracker.prototype, "drain");
+    const resultAck = node
+      .request(
+        "node.invoke.result",
+        {
+          id: frame.id,
+          nodeId: frame.nodeId,
+          ok: true,
+          payloadJSON: JSON.stringify({ completed: "before-close" }),
+        },
+        { timeoutMs: 10_000 },
+      )
+      .catch((error: unknown) => error);
+    await pairingReadStarted;
+    const rawNodeSocket = Reflect.get(node, "ws") as { terminate?: () => void } | null;
+    const stopped = node.stopAndWait({ timeoutMs: 1_000 });
+    rawNodeSocket?.terminate?.();
+    await stopped;
+    node = undefined;
+    await vi.waitFor(() => expect(drainSpy).toHaveBeenCalledOnce());
+    if (removePairingDuringDrain) {
+      await operator.request(
+        "node.pair.remove",
+        { nodeId: pairedNode.identity.deviceId },
+        { timeoutMs: 10_000 },
+      );
+    }
+    pairingRead.release?.();
+
+    if (removePairingDuringDrain) {
+      await expect(invoked).rejects.toThrow("node pairing changed while invocation was active");
+    } else {
+      await expect(invoked).resolves.toMatchObject({
+        ok: true,
+        nodeId: pairedNode.identity.deviceId,
+        command: "camera.list",
+        payload: { completed: "before-close" },
+      });
+    }
+    await resultAck;
+    await vi.waitFor(async () => {
+      const listed = await operator?.request<{
+        nodes?: Array<{ nodeId?: string; connected?: boolean }>;
+      }>("node.list", {}, { timeoutMs: 10_000 });
+      const listedNode = listed?.nodes?.find(
+        (entry) => entry.nodeId === pairedNode.identity.deviceId,
+      );
+      if (removePairingDuringDrain) {
+        expect(listedNode).toBeUndefined();
+      } else {
+        expect(listedNode?.connected).toBe(false);
+      }
+    });
+  } finally {
+    releaseBlockedPairingRead();
+    await Promise.allSettled([
+      ...(node ? [node.stopAndWait({ timeoutMs: 1_000 })] : []),
+      ...(operator ? [operator.stopAndWait({ timeoutMs: 1_000 })] : []),
+    ]);
+    await server.close();
+  }
+});
+
+function releaseBlockedPairingRead(): void {
+  pairingRead.release?.();
+  pairingRead.onBlocked = null;
+  pairingRead.blocked = null;
+  pairingRead.release = null;
+}

src/gateway/server/ws-connection.ts

@@ -40,6 +40,10 @@ import {
   shouldLimitMissingCredentialAuthLog,
 } from "./ws-connection/handshake-auth-log-limiter.js";
 import type { WsOriginCheckMetrics } from "./ws-connection/message-handler.js";
+import {
+  GatewayNodeLifecycleDispatchTracker,
+  NODE_LIFECYCLE_DISPATCH_DRAIN_TIMEOUT_MS,
+} from "./ws-connection/node-lifecycle-dispatch.js";
 import {
   attachWorkerWsMessageHandler,
   type WorkerConnectionService,
@@ -318,6 +322,7 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
     let lastFrameMethod: string | undefined;
     let lastFrameId: string | undefined;
     let hasReceivedPreauthFrame = false;
+    const nodeLifecycleDispatch = new GatewayNodeLifecycleDispatchTracker();
 
     socket.once("message", () => {
       hasReceivedPreauthFrame = true;
@@ -357,6 +362,7 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
     let pingTimer: ReturnType<typeof setInterval> | undefined;
     let cleanupWorkerConnection: (() => void) | undefined;
     let awaitingPong = false;
+    let retainClientUntilNodeDrain = false;
     const handshakeTimeoutMs = resolvePreauthHandshakeTimeoutMs({
       configuredTimeoutMs: params.preauthHandshakeTimeoutMs,
     });
@@ -379,20 +385,15 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
       }
     }, handshakeTimeoutMs);
 
-    const close = (code = 1000, reason?: string) => {
+    const retireTransport = (code = 1000, reason?: string) => {
       if (closed) {
         return;
       }
       closed = true;
       clearTimeout(handshakeTimer);
-      if (pingTimer !== undefined) {
-        clearInterval(pingTimer);
-      }
+      clearInterval(pingTimer);
       cleanupWorkerConnection?.();
       releasePreauthBudget();
-      if (client) {
-        clients.delete(client);
-      }
       try {
         socket.close(code, reason);
       } catch {
@@ -400,6 +401,13 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
       }
     };
 
+    const close = (code = 1000, reason?: string) => {
+      retireTransport(code, reason);
+      if (client && !retainClientUntilNodeDrain) {
+        clients.delete(client);
+      }
+    };
+
     const send = (obj: unknown) => {
       if (closed) {
         return;
@@ -476,7 +484,7 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
       normalizeLowercaseStringOrEmpty(requestUserAgent).startsWith("openclaw/") &&
       isLoopbackAddress(remoteAddr);
 
-    socket.once("close", (code, reason) => {
+    const handleSocketClose = async (code: number, reason: Buffer) => {
       const durationMs = Date.now() - openedAt;
       const logForwardedFor = sanitizeLogValue(forwardedFor);
       const logOrigin = sanitizeLogValue(requestOrigin);
@@ -565,7 +573,23 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
         context.terminalSessions?.handleDisconnect(connId);
         let currentDisconnectedNodeId: string | null = null;
         if (client?.connect?.role === "node") {
-          currentDisconnectedNodeId = context.nodeRegistry.unregister(connId);
+          // Retire I/O immediately, but keep the client revocable until admitted
+          // lifecycle work drains; pairing/token removal must still fence it.
+          retainClientUntilNodeDrain = true;
+          retireTransport();
+          try {
+            if (nodeLifecycleDispatch.hasActive()) {
+              const drained = await nodeLifecycleDispatch.drain();
+              if (!drained) {
+                logGateway.warn(
+                  `node lifecycle dispatch drain timed out after ${NODE_LIFECYCLE_DISPATCH_DRAIN_TIMEOUT_MS}ms conn=${connId}`,
+                );
+              }
+            }
+            currentDisconnectedNodeId = context.nodeRegistry.unregister(connId);
+          } finally {
+            retainClientUntilNodeDrain = false;
+          }
         }
         if (
           client?.presenceKey &&
@@ -597,6 +621,12 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
         endpoint,
       });
       close();
+    };
+    socket.once("close", (code, reason) => {
+      void handleSocketClose(code, reason).catch((error: unknown) => {
+        logGateway.error(`websocket close cleanup failed conn=${connId}: ${formatError(error)}`);
+        close();
+      });
     });
 
     const setClient = (next: GatewayWsClient) => {
@@ -697,6 +727,7 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
       extraHandlers,
       getMethodRegistry,
       buildRequestContext,
+      nodeLifecycleDispatch,
       refreshHealthSnapshot,
       send,
       close,

src/gateway/server/ws-connection/authenticated-request-dispatch.ts

@@ -187,12 +187,17 @@ export function createGatewayAuthenticatedRequestDispatcher(params: {
       }
     };
     const upstreamTrace = parseDiagnosticTraceparent(req.traceparent);
-    const requestDispatch = upstreamTrace
-      ? runWithDiagnosticTraceContext(
-          createChildDiagnosticTraceContext(upstreamTrace),
-          executeRequest,
-        )
-      : executeRequest();
+    const dispatchRequest = () =>
+      upstreamTrace
+        ? runWithDiagnosticTraceContext(
+            createChildDiagnosticTraceContext(upstreamTrace),
+            executeRequest,
+          )
+        : executeRequest();
+    const requestDispatch =
+      client.connect.role === "node"
+        ? params.handler.nodeLifecycleDispatch.dispatch(req.method, dispatchRequest)
+        : dispatchRequest();
     if (DEVICE_CREDENTIAL_INVALIDATING_METHODS.has(req.method)) {
       const barrier = requestDispatch.finally(() => {
         if (deviceCredentialMutationBarrier === barrier) {

src/gateway/server/ws-connection/message-handler-types.ts

@@ -20,6 +20,7 @@ import type { GatewayRequestContext, GatewayRequestHandlers } from "../../server
 import type { GatewayWsClient, WsHandshakePhase } from "../ws-types.js";
 import type { resolveControlUiAuthPolicy } from "./connect-policy.js";
 import type { resolvePairingLocality } from "./handshake-auth-helpers.js";
+import type { GatewayNodeLifecycleDispatchTracker } from "./node-lifecycle-dispatch.js";
 
 type SubsystemLogger = ReturnType<typeof createSubsystemLogger>;
 type ControlUiAuthPolicy = ReturnType<typeof resolveControlUiAuthPolicy>;
@@ -60,6 +61,7 @@ export type GatewayWsMessageHandlerParams = {
   extraHandlers: GatewayRequestHandlers;
   getMethodRegistry?: () => GatewayMethodRegistry;
   buildRequestContext: () => GatewayRequestContext;
+  nodeLifecycleDispatch: GatewayNodeLifecycleDispatchTracker;
   refreshHealthSnapshot: GatewayRequestContext["refreshHealthSnapshot"];
   send: (obj: unknown) => void;
   close: (code?: number, reason?: string) => void;

src/gateway/server/ws-connection/message-handler.post-connect-health.test.ts

@@ -22,6 +22,7 @@ import type { HealthSummary } from "../../health/types.js";
 import { getOperatorApprovalRuntimeToken } from "../../operator-approval-runtime-token.js";
 import { handleGatewayRequest } from "../../server-methods.js";
 import type { GatewayRequestContext } from "../../server-methods/types.js";
+import { GatewayNodeLifecycleDispatchTracker } from "./node-lifecycle-dispatch.js";
 
 const {
   buildGatewaySnapshotMock,
@@ -257,6 +258,7 @@ function attachGatewayHarness(options: {
     events: [],
     extraHandlers: {},
     buildRequestContext: () => ({}) as GatewayRequestContext,
+    nodeLifecycleDispatch: new GatewayNodeLifecycleDispatchTracker(),
     refreshHealthSnapshot:
       options.refreshHealthSnapshot ?? vi.fn(async () => createHealthSummary()),
     send,

src/gateway/server/ws-connection/message-handler.suspension-admission.test.ts

@@ -9,6 +9,7 @@ import {
   tryBeginGatewaySuspendAdmission,
 } from "../../../process/gateway-work-admission.js";
 import type { GatewayRequestContext } from "../../server-methods/types.js";
+import { GatewayNodeLifecycleDispatchTracker } from "./node-lifecycle-dispatch.js";
 
 const { incrementPresenceVersionMock, loadConfigMock, upsertPresenceMock } = vi.hoisted(() => ({
   incrementPresenceVersionMock: vi.fn(() => 2),
@@ -100,6 +101,7 @@ function attachHarness(params: { deferSocketSend?: boolean } = {}) {
     events: [],
     extraHandlers: {},
     buildRequestContext: () => ({}) as GatewayRequestContext,
+    nodeLifecycleDispatch: new GatewayNodeLifecycleDispatchTracker(),
     refreshHealthSnapshot: vi.fn(async () => ({}) as never),
     send,
     close,

src/gateway/server/ws-connection/node-lifecycle-dispatch.test.ts

@@ -0,0 +1,50 @@
+import { describe, expect, it, vi } from "vitest";
+import { GatewayNodeLifecycleDispatchTracker } from "./node-lifecycle-dispatch.js";
+
+describe("GatewayNodeLifecycleDispatchTracker", () => {
+  it("drains every admitted node progress and terminal result", async () => {
+    const tracker = new GatewayNodeLifecycleDispatchTracker();
+    const events: string[] = [];
+    let releaseProgress: (() => void) | undefined;
+    const progressGate = new Promise<void>((resolve) => {
+      releaseProgress = resolve;
+    });
+
+    const progress = tracker.dispatch("node.invoke.progress", async () => {
+      events.push("progress-start");
+      await progressGate;
+      events.push("progress-end");
+    });
+    const result = tracker.dispatch("node.invoke.result", async () => {
+      events.push("result");
+    });
+    const drained = tracker.drain(1_000);
+
+    await vi.waitFor(() => expect(events).toEqual(["progress-start", "result"]));
+    releaseProgress?.();
+
+    await expect(Promise.all([progress, result])).resolves.toEqual([undefined, undefined]);
+    await expect(drained).resolves.toBe(true);
+    expect(events).toEqual(["progress-start", "result", "progress-end"]);
+  });
+
+  it("does not queue unrelated methods and bounds a stuck lifecycle drain", async () => {
+    const tracker = new GatewayNodeLifecycleDispatchTracker();
+    let releaseResult: (() => void) | undefined;
+    const resultGate = new Promise<void>((resolve) => {
+      releaseResult = resolve;
+    });
+    const result = tracker.dispatch("node.invoke.result", async () => {
+      await resultGate;
+    });
+
+    await expect(tracker.dispatch("node.event", async () => undefined)).resolves.toBeUndefined();
+    expect(tracker.hasActive()).toBe(true);
+    await expect(tracker.drain(1)).resolves.toBe(false);
+
+    releaseResult?.();
+    await expect(result).resolves.toBeUndefined();
+    expect(tracker.hasActive()).toBe(false);
+    await expect(tracker.drain(1)).resolves.toBe(true);
+  });
+});

src/gateway/server/ws-connection/node-lifecycle-dispatch.ts

@@ -0,0 +1,56 @@
+const NODE_LIFECYCLE_METHODS = new Set(["node.invoke.progress", "node.invoke.result"]);
+
+export const NODE_LIFECYCLE_DISPATCH_DRAIN_TIMEOUT_MS = 1_000;
+
+/**
+ * Tracks admitted node progress/result requests so physical disconnect cleanup
+ * drains their existing work before retiring invokes. It does not add a queue.
+ */
+export class GatewayNodeLifecycleDispatchTracker {
+  private active = new Set<Promise<void>>();
+
+  hasActive(): boolean {
+    return this.active.size > 0;
+  }
+
+  dispatch(method: string, run: () => Promise<void>): Promise<void> {
+    const execution = run();
+    if (!NODE_LIFECYCLE_METHODS.has(method)) {
+      return execution;
+    }
+    const settled = execution.then(
+      () => undefined,
+      () => undefined,
+    );
+    this.active.add(settled);
+    void settled.finally(() => {
+      this.active.delete(settled);
+    });
+    return execution;
+  }
+
+  async drain(timeoutMs = NODE_LIFECYCLE_DISPATCH_DRAIN_TIMEOUT_MS): Promise<boolean> {
+    const deadlineAt = Date.now() + Math.max(0, timeoutMs);
+    while (this.active.size > 0) {
+      const remainingMs = deadlineAt - Date.now();
+      if (remainingMs <= 0) {
+        return false;
+      }
+      let timeout: ReturnType<typeof setTimeout> | undefined;
+      const timedOut = Symbol("node-lifecycle-dispatch-timeout");
+      const result = await Promise.race([
+        Promise.allSettled(this.active),
+        new Promise<typeof timedOut>((resolve) => {
+          timeout = setTimeout(() => resolve(timedOut), remainingMs);
+        }),
+      ]);
+      if (timeout) {
+        clearTimeout(timeout);
+      }
+      if (result === timedOut) {
+        return false;
+      }
+    }
+    return true;
+  }
+}