From 1cf1af4aed3dedb010a885fcc9dc652f6d6d93d2 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 4 Aug 2026 09:25:34 -0700 Subject: [PATCH] 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 | 219 ++++++++++++++++++ src/gateway/server/ws-connection.ts | 49 +++- .../authenticated-request-dispatch.ts | 17 +- .../ws-connection/message-handler-types.ts | 2 + ...essage-handler.post-connect-health.test.ts | 2 + ...ssage-handler.suspension-admission.test.ts | 2 + .../node-lifecycle-dispatch.test.ts | 50 ++++ .../ws-connection/node-lifecycle-dispatch.ts | 56 +++++ 8 files changed, 382 insertions(+), 15 deletions(-) create mode 100644 src/gateway/server.node-result-close.test.ts create mode 100644 src/gateway/server/ws-connection/node-lifecycle-dispatch.test.ts create mode 100644 src/gateway/server/ws-connection/node-lifecycle-dispatch.ts diff --git a/src/gateway/server.node-result-close.test.ts b/src/gateway/server.node-result-close.test.ts new file mode 100644 index 000000000000..ebe7e437b64a --- /dev/null +++ b/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 | null, + onBlocked: null as (() => void) | null, + release: null as (() => void) | null, +})); + +vi.mock("../infra/node-pairing-state.js", async (importOriginal) => { + const actual = await importOriginal(); + 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> | undefined; + let node: Awaited> | 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((resolve) => { + pairingRead.release = resolve; + }); + const pairingReadStarted = new Promise((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; +} diff --git a/src/gateway/server/ws-connection.ts b/src/gateway/server/ws-connection.ts index 201ac8fec0ee..6fefbef1a3a9 100644 --- a/src/gateway/server/ws-connection.ts +++ b/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 | 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, diff --git a/src/gateway/server/ws-connection/authenticated-request-dispatch.ts b/src/gateway/server/ws-connection/authenticated-request-dispatch.ts index 6d749900c1a8..cc48da86db2f 100644 --- a/src/gateway/server/ws-connection/authenticated-request-dispatch.ts +++ b/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) { diff --git a/src/gateway/server/ws-connection/message-handler-types.ts b/src/gateway/server/ws-connection/message-handler-types.ts index e7833db30b41..7d6daa8a2f85 100644 --- a/src/gateway/server/ws-connection/message-handler-types.ts +++ b/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; type ControlUiAuthPolicy = ReturnType; @@ -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; 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 index e9f904c88c91..d062a793badb 100644 --- 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 @@ -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, diff --git a/src/gateway/server/ws-connection/message-handler.suspension-admission.test.ts b/src/gateway/server/ws-connection/message-handler.suspension-admission.test.ts index 8ee3d836ca63..8868483ad56c 100644 --- a/src/gateway/server/ws-connection/message-handler.suspension-admission.test.ts +++ b/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, diff --git a/src/gateway/server/ws-connection/node-lifecycle-dispatch.test.ts b/src/gateway/server/ws-connection/node-lifecycle-dispatch.test.ts new file mode 100644 index 000000000000..008dd81d3884 --- /dev/null +++ b/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((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((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); + }); +}); diff --git a/src/gateway/server/ws-connection/node-lifecycle-dispatch.ts b/src/gateway/server/ws-connection/node-lifecycle-dispatch.ts new file mode 100644 index 000000000000..6b696790ce4c --- /dev/null +++ b/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>(); + + hasActive(): boolean { + return this.active.size > 0; + } + + dispatch(method: string, run: () => Promise): Promise { + 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 { + 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 | undefined; + const timedOut = Symbol("node-lifecycle-dispatch-timeout"); + const result = await Promise.race([ + Promise.allSettled(this.active), + new Promise((resolve) => { + timeout = setTimeout(() => resolve(timedOut), remainingMs); + }), + ]); + if (timeout) { + clearTimeout(timeout); + } + if (result === timedOut) { + return false; + } + } + return true; + } +}