fix(gateway): preserve node reconnect state (#78351)

samzong · web-flow · commit 1819e41d2683 · 2026-05-08T12:32:18.000+08:00
Preserve node registry ownership across same-node WebSocket reconnect races so stale old-socket closes cannot clear the replacement session or complete the wrong pending invoke. Thanks @samzong.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -600,6 +600,7 @@ Docs: https://docs.openclaw.ai
 - Managed proxy/security: classify raw socket callsites and proxy runtime mutations in boundary checks so new direct egress or unmanaged proxy-state changes cannot land without explicit review. (#77126) Thanks @jesse-merhi.
 - Channels/iMessage: surface the silent group-allowlist drop at default log level by emitting a one-time `warn` per account at monitor startup when `channels.imessage.groupPolicy: "allowlist"` is set without a `channels.imessage.groups` block, plus a one-time `warn` per `chat_id` when the runtime gate drops a specific group, naming the exact `channels.imessage.groups[...]` key to add to allow it. Fixes #78749. (#79190) Thanks @omarshahine.
 - WhatsApp: stop Gateway-originated outbound echoes from advancing inbound activity in `openclaw channels status`, so outbound self-sends no longer look like handled inbound messages. Fixes #79056. (#79057) Thanks @ai-hpc and @bittoby.
+- Gateway/nodes: preserve the live node registry session and invoke ownership when an older same-node WebSocket closes after reconnecting. (#78351) Thanks @samzong.
 
 ## 2026.5.3-1
 
diff --git a/src/gateway/node-registry.test.ts b/src/gateway/node-registry.test.ts
@@ -0,0 +1,57 @@
+import { describe, expect, it } from "vitest";
+import { NodeRegistry } from "./node-registry.js";
+import type { GatewayWsClient } from "./server/ws-types.js";
+
+function makeClient(connId: string, nodeId: string, sent: string[] = []): GatewayWsClient {
+  return {
+    connId,
+    usesSharedGatewayAuth: false,
+    socket: {
+      send(frame: unknown) {
+        if (typeof frame === "string") {
+          sent.push(frame);
+        }
+      },
+    } as unknown as GatewayWsClient["socket"],
+    connect: {
+      client: { id: "openclaw-macos", version: "1.0.0", platform: "darwin", mode: "node" },
+      device: {
+        id: nodeId,
+        publicKey: "public-key",
+        signature: "signature",
+        signedAt: 1,
+        nonce: "nonce",
+      },
+    } as GatewayWsClient["connect"],
+  };
+}
+
+describe("gateway/node-registry", () => {
+  it("keeps a reconnected node when the old connection unregisters", async () => {
+    const registry = new NodeRegistry();
+    const oldFrames: string[] = [];
+    const newClient = makeClient("conn-new", "node-1");
+
+    registry.register(makeClient("conn-old", "node-1", oldFrames), {});
+    const oldInvoke = registry.invoke({
+      nodeId: "node-1",
+      command: "system.run",
+      timeoutMs: 1_000,
+    });
+    const oldDisconnected = oldInvoke.catch((err: unknown) => err);
+    const oldRequest = JSON.parse(oldFrames[0] ?? "{}") as { payload?: { id?: string } };
+    const newSession = registry.register(newClient, {});
+
+    expect(
+      registry.handleInvokeResult({
+        id: oldRequest.payload?.id ?? "",
+        nodeId: "node-1",
+        connId: "conn-new",
+        ok: true,
+      }),
+    ).toBe(false);
+    expect(registry.unregister("conn-old")).toBeNull();
+    expect(registry.get("node-1")).toBe(newSession);
+    await expect(oldDisconnected).resolves.toBeInstanceOf(Error);
+  });
+});
diff --git a/src/gateway/node-registry.ts b/src/gateway/node-registry.ts
@@ -24,6 +24,7 @@ export type NodeSession = {
 
 type PendingInvoke = {
   nodeId: string;
+  connId: string;
   command: string;
   resolve: (value: NodeInvokeResult) => void;
   reject: (err: Error) => void;
@@ -88,16 +89,19 @@ export class NodeRegistry {
       return null;
     }
     this.nodesByConn.delete(connId);
-    this.nodesById.delete(nodeId);
+    const unregistersCurrentNode = this.nodesById.get(nodeId)?.connId === connId;
+    if (unregistersCurrentNode) {
+      this.nodesById.delete(nodeId);
+    }
     for (const [id, pending] of this.pendingInvokes.entries()) {
-      if (pending.nodeId !== nodeId) {
+      if (pending.connId !== connId) {
         continue;
       }
       clearTimeout(pending.timer);
       pending.reject(new Error(`node disconnected (${pending.command})`));
       this.pendingInvokes.delete(id);
     }
-    return nodeId;
+    return unregistersCurrentNode ? nodeId : null;
   }
 
   listConnected(): NodeSession[] {
@@ -150,6 +154,7 @@ export class NodeRegistry {
       }, timeoutMs);
       this.pendingInvokes.set(requestId, {
         nodeId: params.nodeId,
+        connId: node.connId,
         command: params.command,
         resolve,
         reject,
@@ -161,6 +166,7 @@ export class NodeRegistry {
   handleInvokeResult(params: {
     id: string;
     nodeId: string;
+    connId: string | undefined;
     ok: boolean;
     payload?: unknown;
     payloadJSON?: string | null;
@@ -170,7 +176,7 @@ export class NodeRegistry {
     if (!pending) {
       return false;
     }
-    if (pending.nodeId !== params.nodeId) {
+    if (pending.nodeId !== params.nodeId || pending.connId !== params.connId) {
       return false;
     }
     clearTimeout(pending.timer);
diff --git a/src/gateway/server-methods/nodes.handlers.invoke-result.ts b/src/gateway/server-methods/nodes.handlers.invoke-result.ts
@@ -54,6 +54,7 @@ export const handleNodeInvokeResult: GatewayRequestHandler = async ({
   const ok = context.nodeRegistry.handleInvokeResult({
     id: p.id,
     nodeId: p.nodeId,
+    connId: client?.connId,
     ok: p.ok,
     payload: p.payload,
     payloadJSON: p.payloadJSON ?? null,
diff --git a/src/gateway/server/ws-connection.test.ts b/src/gateway/server/ws-connection.test.ts
@@ -3,13 +3,22 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { WebSocketServer } from "ws";
 import type { ResolvedGatewayAuth } from "../auth.js";
 
-const { attachGatewayWsMessageHandlerMock } = vi.hoisted(() => ({
-  attachGatewayWsMessageHandlerMock: vi.fn(),
-}));
+const { attachGatewayWsMessageHandlerMock, broadcastPresenceSnapshotMock, upsertPresenceMock } =
+  vi.hoisted(() => ({
+    attachGatewayWsMessageHandlerMock: vi.fn(),
+    broadcastPresenceSnapshotMock: vi.fn(),
+    upsertPresenceMock: vi.fn(),
+  }));
 
 vi.mock("./ws-connection/message-handler.js", () => ({
   attachGatewayWsMessageHandler: attachGatewayWsMessageHandlerMock,
 }));
+vi.mock("../../infra/system-presence.js", () => ({
+  upsertPresence: upsertPresenceMock,
+}));
+vi.mock("./presence-events.js", () => ({
+  broadcastPresenceSnapshot: broadcastPresenceSnapshotMock,
+}));
 
 import { attachGatewayWsConnectionHandler } from "./ws-connection.js";
 import { resolveSharedGatewaySessionGeneration } from "./ws-shared-generation.js";
@@ -122,6 +131,8 @@ async function connectTestWs(
 describe("attachGatewayWsConnectionHandler", () => {
   beforeEach(() => {
     attachGatewayWsMessageHandlerMock.mockReset();
+    broadcastPresenceSnapshotMock.mockReset();
+    upsertPresenceMock.mockReset();
   });
 
   afterEach(() => {
@@ -234,4 +245,78 @@ describe("attachGatewayWsConnectionHandler", () => {
     vi.advanceTimersByTime(25_000);
     expect(socket.ping).toHaveBeenCalledTimes(1);
   });
+
+  it("skips node presence disconnects for stale reconnected sockets", async () => {
+    const listeners = new Map<string, (...args: unknown[]) => void>();
+    const unregister = vi.fn(() => null);
+    const wss = {
+      on: vi.fn((event: string, handler: (...args: unknown[]) => void) => {
+        listeners.set(event, handler);
+      }),
+    } as unknown as WebSocketServer;
+    const socket = Object.assign(new EventEmitter(), {
+      _socket: {
+        remoteAddress: "127.0.0.1",
+        remotePort: 1234,
+        localAddress: "127.0.0.1",
+        localPort: 5678,
+      },
+      send: vi.fn(),
+      close: vi.fn(),
+    });
+    const upgradeReq = {
+      headers: { host: "127.0.0.1:19001" },
+      socket: { localAddress: "127.0.0.1" },
+    };
+
+    attachGatewayWsConnectionHandler({
+      wss,
+      clients: new Set(),
+      preauthConnectionBudget: { release: vi.fn() } as never,
+      port: 19001,
+      resolvedAuth: createResolvedAuth("token"),
+      gatewayMethods: [],
+      events: [],
+      refreshHealthSnapshot: vi.fn(),
+      logGateway: createLogger() as never,
+      logHealth: createLogger() as never,
+      logWsControl: createLogger() as never,
+      extraHandlers: {},
+      broadcast: vi.fn(),
+      buildRequestContext: () =>
+        ({
+          unsubscribeAllSessionEvents: vi.fn(),
+          nodeRegistry: { unregister },
+          nodeUnsubscribeAll: vi.fn(),
+        }) as never,
+    });
+
+    const onConnection = listeners.get("connection");
+    expect(onConnection).toBeTypeOf("function");
+    onConnection?.(socket, upgradeReq);
+    await waitForLazyMessageHandler();
+
+    const passed = attachGatewayWsMessageHandlerMock.mock.calls[0]?.[0] as {
+      setClient: (client: unknown) => boolean;
+    };
+    expect(
+      passed.setClient({
+        socket,
+        connect: {
+          role: "node",
+          client: { id: "openclaw-macos", mode: "node" },
+          device: { id: "node-1" },
+        },
+        connId: "conn-old",
+        presenceKey: "node-1",
+        usesSharedGatewayAuth: false,
+      }),
+    ).toBe(true);
+
+    socket.emit("close", 1000, Buffer.from("stale"));
+
+    expect(unregister).toHaveBeenCalledTimes(1);
+    expect(upsertPresenceMock).not.toHaveBeenCalled();
+    expect(broadcastPresenceSnapshotMock).not.toHaveBeenCalled();
+  });
 });
diff --git a/src/gateway/server/ws-connection.ts b/src/gateway/server/ws-connection.ts
@@ -387,19 +387,23 @@ export function attachGatewayWsConnectionHandler(params: AttachGatewayWsConnecti
           `webchat disconnected code=${code} reason=${logReason || "n/a"} conn=${connId}`,
         );
       }
-      if (client?.presenceKey) {
-        upsertPresence(client.presenceKey, { reason: "disconnect" });
-        broadcastPresenceSnapshot({ broadcast, incrementPresenceVersion, getHealthVersion });
-      }
       const context = buildRequestContext();
       context.unsubscribeAllSessionEvents(connId);
+      let currentDisconnectedNodeId: string | null = null;
       if (client?.connect?.role === "node") {
-        const nodeId = context.nodeRegistry.unregister(connId);
-        if (nodeId) {
-          removeRemoteNodeInfo(nodeId);
-          context.nodeUnsubscribeAll(nodeId);
-          clearNodeWakeState(nodeId);
-        }
+        currentDisconnectedNodeId = context.nodeRegistry.unregister(connId);
+      }
+      if (
+        client?.presenceKey &&
+        (client.connect.role !== "node" || currentDisconnectedNodeId !== null)
+      ) {
+        upsertPresence(client.presenceKey, { reason: "disconnect" });
+        broadcastPresenceSnapshot({ broadcast, incrementPresenceVersion, getHealthVersion });
+      }
+      if (currentDisconnectedNodeId) {
+        removeRemoteNodeInfo(currentDisconnectedNodeId);
+        context.nodeUnsubscribeAll(currentDisconnectedNodeId);
+        clearNodeWakeState(currentDisconnectedNodeId);
       }
       logWs("out", "close", {
         connId,
