From 61ab6a8f9d151c84377ba9f90152f257110cc5a6 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 11 Aug 2026 22:12:53 -0700 Subject: [PATCH] feat(gateway): make suspend/resume operator-usable end to end (#122100) * feat(gateway): make suspend/resume operator-usable end to end A prepared Gateway now accepts authenticated WebSocket connects while keeping every method except gateway.suspend.* fenced, so a fresh CLI or controller process can resume a suspension instead of dead-ending on a rejected upgrade until the two-minute lease expires. Restart drain, worker ingress, and desktop-observe streams stay fully closed. The gateway client surfaces non-101 upgrade responses (bounded body read) as typed retryable errors instead of an opaque 1006 close, and new openclaw gateway suspend / resume commands drive the whole handshake, including bounded --wait polling with blocker output. Live-verified on an isolated dev gateway: prepare, SIGSTOP/SIGCONT freeze, resume, over-TTL expiry self-heal, conflict and mismatch paths. * refactor(gateway-client): move wire-client contract types to protocol-client-contract The connectError addition pushed protocol-client.ts over the 700-line max-lines gate; split the adapter-facing contract types into their own module instead of suppressing. * refactor(gateway-client): keep contract-internal option types unexported Knip deadcode gates reject exported types with no importer; the connect and close decision shapes are only referenced inside the contract module. * chore(plugin-sdk): refresh gateway-runtime API baseline after rebase * fix(gateway-client): preserve hello type after rebase * test(gateway): support websocket upgrade rejection events * test(gateway): expect connection errors in close info * fix(gateway): keep prepared-suspension connects control-only Address ClawSweeper review: node and worker connects stay refused while suspension is prepared (only operator control connects pass), and the CLI never issues another suspend prepare after its --wait deadline. --- .../gateway-runtime.json | 2 +- docs/cli/gateway.md | 30 ++++ docs/gateway/external-apps.md | 16 +- docs/gateway/protocol.md | 2 +- .../src/client.handshake.test.ts | 109 ++++++++++- packages/gateway-client/src/client.ts | 69 +++++++ .../src/protocol-client-contract.ts | 114 ++++++++++++ .../gateway-client/src/protocol-client.ts | 127 ++----------- .../register.option-collisions.test.ts | 44 ++++- src/cli/gateway-cli/register.ts | 49 +++++ src/cli/gateway-cli/suspend-cli.test.ts | 169 ++++++++++++++++++ src/cli/gateway-cli/suspend-cli.ts | 157 ++++++++++++++++ src/gateway/call.test.ts | 28 +++ src/gateway/call.ts | 5 + src/gateway/client.test.ts | 6 + src/gateway/probe.device-auth-scope.test.ts | 3 +- src/gateway/server-http.ts | 16 +- src/gateway/server.preauth-hardening.test.ts | 75 +++++++- ...ssage-handler.suspension-admission.test.ts | 162 ++++++++++++----- .../server/ws-connection/message-handler.ts | 35 +++- 20 files changed, 1034 insertions(+), 184 deletions(-) create mode 100644 packages/gateway-client/src/protocol-client-contract.ts create mode 100644 src/cli/gateway-cli/suspend-cli.test.ts create mode 100644 src/cli/gateway-cli/suspend-cli.ts diff --git a/docs/.generated/plugin-sdk-api-baseline/gateway-runtime.json b/docs/.generated/plugin-sdk-api-baseline/gateway-runtime.json index 296271a135ae..1751f4ed359a 100644 --- a/docs/.generated/plugin-sdk-api-baseline/gateway-runtime.json +++ b/docs/.generated/plugin-sdk-api-baseline/gateway-runtime.json @@ -1 +1 @@ -{"contentHash":"85a649a9621f78be5d7e26b01ec4b8d465f7c353109571be30530b5c8c4bd5ac","entrypoint":"gateway-runtime","importSpecifier":"openclaw/plugin-sdk/gateway-runtime"} +{"contentHash":"e31574637f3e276338e3e05e4774f7ad79981042c8021e49fb53df7ffe901917","entrypoint":"gateway-runtime","importSpecifier":"openclaw/plugin-sdk/gateway-runtime"} diff --git a/docs/cli/gateway.md b/docs/cli/gateway.md index 046cb63cc8d0..6f0555978526 100644 --- a/docs/cli/gateway.md +++ b/docs/cli/gateway.md @@ -516,6 +516,36 @@ openclaw gateway call logs.tail --params '{"limit": 200}' `--params` must be valid JSON, and each method validates its own param shape (extra/misnamed fields are rejected). Use `--port` for a custom-port local Gateway; explicit `--url` targets still require explicit credentials. +### `gateway suspend` + +Prepare an idle Gateway for a cooperative host freeze or snapshot. Without +`--wait`, active work returns a nonzero exit with blocker details. With +`--wait`, the CLI retries until the bounded deadline using one stable request +ID. + +```bash +openclaw gateway suspend +openclaw gateway suspend --request-id snapshot-2026-08-11 --wait 30 +openclaw gateway suspend --port 18999 --json +``` + +The ready output includes the suspension ID, lease expiry, and the matching +resume command. Common RPC options such as `--url`, `--token`, `--password`, +`--timeout`, `--json`, and `--port` are supported. + +### `gateway resume ` + +Release a prepared suspension after thaw or when the host operation is +abandoned. + +```bash +openclaw gateway resume +openclaw gateway resume --port 18999 --json +``` + +An already expired or resumed lease is a successful no-op. A different active +suspension ID is rejected. + ## Manage the Gateway service ```bash diff --git a/docs/gateway/external-apps.md b/docs/gateway/external-apps.md index 9edae0d38813..4d7b4c2ca873 100644 --- a/docs/gateway/external-apps.md +++ b/docs/gateway/external-apps.md @@ -64,15 +64,15 @@ host-neutral suspension handshake: 4. If it is `ready`, save the returned `suspensionId`, then freeze or snapshot the process before `expiresAtMs`. 5. After thaw, or if suspension is abandoned, call `gateway.suspend.resume` - with that `suspensionId` over the existing WebSocket or Admin HTTP control - path. + with that `suspensionId` over the existing or a newly authenticated + WebSocket. The CLI equivalents are `openclaw gateway suspend` and + `openclaw gateway resume `. -A prepared Gateway rejects new WebSocket handshakes. A WebSocket controller -must keep its authenticated connection open across the host operation. If that -cannot be guaranteed, enable and use the -[Admin HTTP RPC plugin](/plugins/admin-http-rpc) before preparing. If the -control path is lost, wait for the two-minute lease to expire before -reconnecting; expiry reopens admission automatically. +A prepared Gateway accepts authenticated WebSocket connects, but fences every +method except `gateway.suspend.*`. Controllers may reconnect after thaw and +call resume. The [Admin HTTP RPC plugin](/plugins/admin-http-rpc) remains +available for hosts that cannot speak WebSocket at all. If every control path +is lost, the two-minute lease expiry reopens admission automatically. The RPC contract is: diff --git a/docs/gateway/protocol.md b/docs/gateway/protocol.md index 0323ed32660d..f52522db80d9 100644 --- a/docs/gateway/protocol.md +++ b/docs/gateway/protocol.md @@ -520,7 +520,7 @@ methods. Treat this as feature discovery, not a full enumeration of - `last-heartbeat` returns the latest persisted heartbeat event. - `set-heartbeats` toggles heartbeat processing on the gateway. - `gateway.restart.preflight` is a deprecated, read-only compatibility preview of restart-specific active work. It does not close admission, create a suspension lease, or provide the atomic full-work fence of `gateway.suspend.prepare`; new restart flows should call `gateway.restart.request`. - - `gateway.suspend.prepare` creates a short cooperative-suspension lease only when tracked Gateway work is idle. `gateway.suspend.status` checks that lease, and `gateway.suspend.resume` releases it after thaw or an aborted host operation. + - `gateway.suspend.prepare` creates a short cooperative-suspension lease only when tracked Gateway work is idle. While prepared, authenticated WebSocket connects remain available, but every method except `gateway.suspend.*` is fenced. `gateway.suspend.status` checks the lease, and `gateway.suspend.resume` releases it after thaw or an aborted host operation. diff --git a/packages/gateway-client/src/client.handshake.test.ts b/packages/gateway-client/src/client.handshake.test.ts index e5eb3f60a2cb..23f93891f7e1 100644 --- a/packages/gateway-client/src/client.handshake.test.ts +++ b/packages/gateway-client/src/client.handshake.test.ts @@ -1,4 +1,5 @@ // Gateway Client tests cover websocket opening-handshake timeout behavior. +import http from "node:http"; import net from "node:net"; import type { AddressInfo } from "node:net"; import { afterEach, describe, expect, it } from "vitest"; @@ -16,6 +17,9 @@ describe("GatewayClient websocket opening handshakeTimeout", () => { for (const socket of sockets.splice(0)) { socket.destroy(); } + for (const server of servers) { + (server as net.Server & { closeAllConnections?: () => void }).closeAllConnections?.(); + } await Promise.all( servers.splice(0).map( (server) => @@ -26,18 +30,22 @@ describe("GatewayClient websocket opening handshakeTimeout", () => { ); }); + async function listen(server: net.Server): Promise { + servers.push(server); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => resolve()); + }); + return (server.address() as AddressInfo).port; + } + it("fails when a peer accepts TCP but never completes the websocket upgrade", async () => { // Accept TCP but never complete the websocket upgrade so missing // handshakeTimeout would leave start() waiting forever for open. const server = net.createServer((socket) => { sockets.push(socket); }); - servers.push(server); - await new Promise((resolve, reject) => { - server.once("error", reject); - server.listen(0, "127.0.0.1", () => resolve()); - }); - const { port } = server.address() as AddressInfo; + const port = await listen(server); const handshakeTimeoutMs = 250; const startedAt = Date.now(); const outcome = await new Promise<{ @@ -89,4 +97,93 @@ describe("GatewayClient websocket opening handshakeTimeout", () => { }`, ); }); + + it("surfaces a rejected websocket upgrade body through the connection error", async () => { + let requestCount = 0; + const server = http.createServer((_req, res) => { + requestCount += 1; + res.writeHead(503, { "Content-Type": "text/plain" }); + res.end("Gateway websocket admission closed"); + }); + const port = await listen(server); + const errors: Error[] = []; + let resolveRetry = () => {}; + const retried = new Promise((resolve) => { + resolveRetry = resolve; + }); + const closed = new Promise<{ code: number; connectError?: Error }>((resolve) => { + const client = new GatewayClient({ + url: `ws://127.0.0.1:${port}`, + onConnectError: (error) => { + errors.push(error); + if (errors.length === 2) { + resolveRetry(); + } + }, + onClose: (code, _reason, info) => resolve({ code, connectError: info?.connectError }), + }); + clients.push(client); + client.start(); + }); + + await expect(closed).resolves.toMatchObject({ + code: 1006, + connectError: { + name: "GatewayClientRequestError", + message: + "gateway rejected websocket upgrade (HTTP 503): Gateway websocket admission closed", + gatewayCode: "UNAVAILABLE", + retryable: true, + }, + }); + await retried; + expect(requestCount).toBe(2); + expect(errors).toHaveLength(2); + expect(errors.map((error) => error.message)).toEqual([ + "gateway rejected websocket upgrade (HTTP 503): Gateway websocket admission closed", + "gateway rejected websocket upgrade (HTTP 503): Gateway websocket admission closed", + ]); + }); + + it("caps a rejected websocket upgrade body before the peer ends it", async () => { + const omittedTail = "omitted-tail-marker"; + const server = http.createServer((_req, res) => { + res.writeHead(503, { "Content-Type": "text/plain" }); + res.write(`${"x".repeat(3_000)}${omittedTail}`); + }); + const port = await listen(server); + const error = await new Promise((resolve) => { + const client = new GatewayClient({ + url: `ws://127.0.0.1:${port}`, + onConnectError: resolve, + }); + clients.push(client); + client.start(); + }); + + expect(error.message).toHaveLength( + "gateway rejected websocket upgrade (HTTP 503): ".length + 2 * 1024, + ); + expect(error.message).not.toContain(omittedTail); + }); + + it("times out while reading a stalled websocket upgrade response body", async () => { + const server = http.createServer((_req, res) => { + res.writeHead(503, { "Content-Type": "text/plain" }); + res.write("still suspending"); + }); + const port = await listen(server); + const startedAt = Date.now(); + const error = await new Promise((resolve) => { + const client = new GatewayClient({ + url: `ws://127.0.0.1:${port}`, + onConnectError: resolve, + }); + clients.push(client); + client.start(); + }); + + expect(error.message).toBe("gateway rejected websocket upgrade (HTTP 503): still suspending"); + expect(Date.now() - startedAt).toBeLessThan(1_500); + }); }); diff --git a/packages/gateway-client/src/client.ts b/packages/gateway-client/src/client.ts index d1799a4d763c..7b376c995730 100644 --- a/packages/gateway-client/src/client.ts +++ b/packages/gateway-client/src/client.ts @@ -1,4 +1,5 @@ import { randomUUID } from "node:crypto"; +import type { ClientRequest, IncomingMessage } from "node:http"; import { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES, @@ -229,6 +230,50 @@ type FingerprintCheckingClientOptions = Omit { + return await new Promise((resolve) => { + const chunks: Buffer[] = []; + let totalBytes = 0; + let settled = false; + const finish = () => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + response.off("data", onData); + response.off("end", finish); + response.off("error", finish); + response.off("aborted", finish); + resolve(Buffer.concat(chunks, totalBytes).toString("utf8").replace(/\s+/gu, " ").trim()); + }; + const stop = () => { + finish(); + response.destroy(); + }; + const onData = (chunk: Buffer | string) => { + const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); + const remaining = MAX_UPGRADE_ERROR_BODY_BYTES - totalBytes; + if (remaining > 0) { + const prefix = buffer.subarray(0, remaining); + chunks.push(prefix); + totalBytes += prefix.byteLength; + } + if (buffer.byteLength >= remaining) { + stop(); + } + }; + const timer = setTimeout(stop, UPGRADE_ERROR_BODY_TIMEOUT_MS); + timer.unref?.(); + response.on("data", onData); + response.once("end", finish); + response.once("error", finish); + response.once("aborted", finish); + }); +} export type GatewayReconnectPausedInfo = { code: number; @@ -241,6 +286,7 @@ export type GatewayClientCloseInfo = { socketOpened: boolean; transportValidated: boolean; transientPreHelloCleanClose: boolean; + connectError?: Error; }; export { GatewayClientRequestError } from "./request-error.js"; @@ -608,6 +654,7 @@ export class GatewayClient { } this.ws = ws; this.transportValidated = false; + let upgradeError: GatewayClientRequestError | undefined; ws.on("open", () => { handlers.open(); if (usesTls && this.opts.tlsFingerprint) { @@ -629,7 +676,28 @@ export class GatewayClient { this.resolvePendingStop(ws); handlers.close(code, reasonText); }); + ws.on("unexpected-response", (request: ClientRequest, response: IncomingMessage) => { + void readUpgradeErrorBody(response).then((body) => { + const statusCode = response.statusCode; + const message = `gateway rejected websocket upgrade (HTTP ${statusCode ?? "unknown"})${body ? `: ${body}` : ""}`; + upgradeError = new GatewayClientRequestError({ + code: "UNAVAILABLE", + message, + retryable: true, + details: { + reason: "websocket-upgrade-rejected", + ...(statusCode === undefined ? {} : { httpStatus: statusCode }), + }, + }); + handlers.error(upgradeError); + request.destroy(); + ws.close(); + }); + }); ws.on("error", (err) => { + if (upgradeError) { + return; + } this.logDebug(`gateway client error: ${formatGatewayClientErrorForLog(err)}`); handlers.error(err instanceof Error ? err : new Error(String(err))); }); @@ -1153,6 +1221,7 @@ export class GatewayClient { transportValidated: this.transportValidated, transientPreHelloCleanClose: !context.helloReceived && context.code === 1000 && context.reason === "", + ...(context.connectFailure?.error ? { connectError: context.connectFailure.error } : {}), }; } diff --git a/packages/gateway-client/src/protocol-client-contract.ts b/packages/gateway-client/src/protocol-client-contract.ts new file mode 100644 index 000000000000..52535b300137 --- /dev/null +++ b/packages/gateway-client/src/protocol-client-contract.ts @@ -0,0 +1,114 @@ +// Wire-client contract types shared by GatewayProtocolClient and its adapters. +import type { ErrorShape, EventFrame, HelloOk } from "@openclaw/gateway-protocol"; +import type { GatewayProtocolRequestTiming } from "./pending-request.js"; +import type { GatewayProtocolRequestError } from "./protocol-request.js"; + +export type GatewayProtocolSocket = { + isOpen: () => boolean; + send: (data: string) => void; + close: (code?: number, reason?: string) => void; +}; +export type GatewayProtocolSocketHandlers = { + open: () => void; + message: (data: string) => void; + close: (code: number, reason: string) => void; + error: (error: Error) => void; +}; +type GatewayProtocolConnectContext = { + generation: number; + nonce: string | null; + challengeTs: number | null | undefined; + plan: TPlan; +}; +export type GatewayProtocolCloseContext = { + code: number; + reason: string; + generation: number; + socketOpened: boolean; + helloReceived: boolean; + connectRequestSent: boolean; + connectFailure?: { error: Error; reconnectDelayMs?: number }; +}; +type GatewayProtocolConnectDecision = { + closeCode: number; + closeReason: string; + reconnectDelayMs?: number; + stop?: boolean; + error?: Error; +}; +type GatewayProtocolCloseDecision = { + retry: boolean; + notify: boolean; + reconnectDelayMs?: number; + pendingError?: Error; +}; +export type GatewayProtocolTiming = { + phase: + | "socket-open" + | "challenge" + | "fallback" + | "device-identity-ready" + | "connect-plan-ready" + | "request-sent" + | "hello" + | "failed"; + generation: number; + durationMs: number; + phaseDurationMs: number; + hasChallenge: boolean; + usedFallback: boolean; + plan?: TPlan; + detail?: unknown; +}; +export type GatewayProtocolClientOptions = { + createSocket: (handlers: GatewayProtocolSocketHandlers) => GatewayProtocolSocket; + createRequestId: () => string; + createRequestError?: (error: Partial) => GatewayProtocolRequestError; + createRequestTimeoutError?: (method: string, timeoutMs: number, requestSent: boolean) => Error; + createRequestAbortError?: (method: string) => Error; + buildConnectPlan: (params: { + nonce: string | null; + challengeTs: number | null | undefined; + generation: number; + }) => TPlan | Promise; + buildConnectParams: (plan: TPlan) => unknown; + onConnectPlanError?: (error: Error) => GatewayProtocolConnectDecision; + onConnectHello?: (hello: HelloOk, context: GatewayProtocolConnectContext) => void; + onHello?: (hello: HelloOk) => void; + onConnectFailure?: ( + error: GatewayProtocolRequestError, + context: GatewayProtocolConnectContext, + ) => GatewayProtocolConnectDecision; + resolveClose: (context: GatewayProtocolCloseContext) => GatewayProtocolCloseDecision; + onClose?: (context: GatewayProtocolCloseContext, decision: GatewayProtocolCloseDecision) => void; + notifyStoppedClose?: boolean; + onConnectError?: (error: Error) => void; + onSocketFactoryError?: (error: Error) => void; + onParseError?: (error: unknown) => void; + onEvent?: (event: EventFrame) => void; + onGap?: (info: { expected: number; received: number }) => void; + onActivity?: () => void; + onTiming?: (timing: GatewayProtocolTiming) => void; + onRequestTiming?: (timing: GatewayProtocolRequestTiming) => void; + onCallbackError?: (label: string, error: unknown) => void; + handshake: + | { mode: "fallback"; timeoutMs: number } + | { + mode: "require-challenge"; + timeoutMs: number; + timeoutMessage?: (elapsedMs: number) => string; + }; + reconnect: { initialMs: number; multiplier: number; maxMs: number }; + requestTimeoutMs?: number; + nowMs?: () => number; + shouldRetrySocketFactoryError?: (error: Error) => boolean; + rethrowSocketFactoryError?: (error: Error) => boolean; +}; +export type ConnectTimingState = { + generation: number; + startedAtMs: number; + lastAtMs: number; + hasChallenge: boolean; + usedFallback: boolean; +}; +export type CloseSnapshot = Omit; diff --git a/packages/gateway-client/src/protocol-client.ts b/packages/gateway-client/src/protocol-client.ts index 8ce4ca530fe8..9e32a8516d49 100644 --- a/packages/gateway-client/src/protocol-client.ts +++ b/packages/gateway-client/src/protocol-client.ts @@ -1,4 +1,4 @@ -import type { ErrorShape, EventFrame, HelloOk } from "@openclaw/gateway-protocol"; +import type { EventFrame, HelloOk } from "@openclaw/gateway-protocol"; import { isGatewayEventFrame, isGatewayResponseFrame, @@ -20,115 +20,21 @@ export { type GatewayProtocolRequestTiming, }; -export type GatewayProtocolSocket = { - isOpen: () => boolean; - send: (data: string) => void; - close: (code?: number, reason?: string) => void; -}; -export type GatewayProtocolSocketHandlers = { - open: () => void; - message: (data: string) => void; - close: (code: number, reason: string) => void; - error: (error: Error) => void; -}; -type GatewayProtocolConnectContext = { - generation: number; - nonce: string | null; - challengeTs: number | null | undefined; - plan: TPlan; -}; -export type GatewayProtocolCloseContext = { - code: number; - reason: string; - generation: number; - socketOpened: boolean; - helloReceived: boolean; - connectRequestSent: boolean; - connectFailure?: { error: Error; reconnectDelayMs?: number }; -}; -type GatewayProtocolConnectDecision = { - closeCode: number; - closeReason: string; - reconnectDelayMs?: number; - stop?: boolean; - error?: Error; -}; -type GatewayProtocolCloseDecision = { - retry: boolean; - notify: boolean; - reconnectDelayMs?: number; - pendingError?: Error; -}; -export type GatewayProtocolTiming = { - phase: - | "socket-open" - | "challenge" - | "fallback" - | "device-identity-ready" - | "connect-plan-ready" - | "request-sent" - | "hello" - | "failed"; - generation: number; - durationMs: number; - phaseDurationMs: number; - hasChallenge: boolean; - usedFallback: boolean; - plan?: TPlan; - detail?: unknown; -}; -type GatewayProtocolClientOptions = { - createSocket: (handlers: GatewayProtocolSocketHandlers) => GatewayProtocolSocket; - createRequestId: () => string; - createRequestError?: (error: Partial) => GatewayProtocolRequestError; - createRequestTimeoutError?: (method: string, timeoutMs: number, requestSent: boolean) => Error; - createRequestAbortError?: (method: string) => Error; - buildConnectPlan: (params: { - nonce: string | null; - challengeTs: number | null | undefined; - generation: number; - }) => TPlan | Promise; - buildConnectParams: (plan: TPlan) => unknown; - onConnectPlanError?: (error: Error) => GatewayProtocolConnectDecision; - onConnectHello?: (hello: HelloOk, context: GatewayProtocolConnectContext) => void; - onHello?: (hello: HelloOk) => void; - onConnectFailure?: ( - error: GatewayProtocolRequestError, - context: GatewayProtocolConnectContext, - ) => GatewayProtocolConnectDecision; - resolveClose: (context: GatewayProtocolCloseContext) => GatewayProtocolCloseDecision; - onClose?: (context: GatewayProtocolCloseContext, decision: GatewayProtocolCloseDecision) => void; - notifyStoppedClose?: boolean; - onConnectError?: (error: Error) => void; - onSocketFactoryError?: (error: Error) => void; - onParseError?: (error: unknown) => void; - onEvent?: (event: EventFrame) => void; - onGap?: (info: { expected: number; received: number }) => void; - onActivity?: () => void; - onTiming?: (timing: GatewayProtocolTiming) => void; - onRequestTiming?: (timing: GatewayProtocolRequestTiming) => void; - onCallbackError?: (label: string, error: unknown) => void; - handshake: - | { mode: "fallback"; timeoutMs: number } - | { - mode: "require-challenge"; - timeoutMs: number; - timeoutMessage?: (elapsedMs: number) => string; - }; - reconnect: { initialMs: number; multiplier: number; maxMs: number }; - requestTimeoutMs?: number; - nowMs?: () => number; - shouldRetrySocketFactoryError?: (error: Error) => boolean; - rethrowSocketFactoryError?: (error: Error) => boolean; -}; -type ConnectTimingState = { - generation: number; - startedAtMs: number; - lastAtMs: number; - hasChallenge: boolean; - usedFallback: boolean; -}; -type CloseSnapshot = Omit; +import type { + CloseSnapshot, + ConnectTimingState, + GatewayProtocolClientOptions, + GatewayProtocolCloseContext, + GatewayProtocolSocket, + GatewayProtocolTiming, +} from "./protocol-client-contract.js"; + +export type { + GatewayProtocolCloseContext, + GatewayProtocolSocket, + GatewayProtocolSocketHandlers, + GatewayProtocolTiming, +} from "./protocol-client-contract.js"; /** * Browser-safe gateway wire client. Environment adapters own transport and auth @@ -571,6 +477,7 @@ export class GatewayProtocolClient { if (!this.isActive(socket, generation) || this.connectSent) { return; } + this.connectFailure = { error }; this.opts.onConnectError?.(error); } diff --git a/src/cli/gateway-cli/register.option-collisions.test.ts b/src/cli/gateway-cli/register.option-collisions.test.ts index 81f0fdf8d817..c26e127a29a9 100644 --- a/src/cli/gateway-cli/register.option-collisions.test.ts +++ b/src/cli/gateway-cli/register.option-collisions.test.ts @@ -4,9 +4,21 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { registerGatewayCli } from "./register.js"; const mocks = vi.hoisted(() => ({ - callGatewayCli: vi.fn(async (_method: string, _opts: unknown, _params?: unknown) => ({ - ok: true, - })), + callGatewayCli: vi.fn(async (method: string, _opts: unknown, _params?: unknown) => { + if (method === "gateway.suspend.prepare") { + return { + status: "ready", + suspensionId: "suspension-1", + expiresAtMs: 1_800_000_000_000, + activeCount: 0, + blockers: [], + }; + } + if (method === "gateway.suspend.resume") { + return { ok: true, status: "running", resumed: true }; + } + return { ok: true }; + }), emitReachableGatewayAuthDiagnostic: vi.fn(async (_params: unknown) => false), formatHealthChannelLines: vi.fn(() => []), gatewayStatusCommand: vi.fn(async (_opts: unknown, _runtime: unknown) => {}), @@ -216,6 +228,32 @@ describe("gateway register option collisions", () => { expectLocalGatewayCall("health", 19085); }, }, + { + name: "projects gateway suspend --port and request id", + argv: ["gateway", "suspend", "--request-id", "host-operation", "--port", "19086", "--json"], + assert: () => { + expectLocalGatewayCall("gateway.suspend.prepare", 19086, { + requestId: "host-operation", + }); + expect(defaultRuntime.writeJson).toHaveBeenCalledWith( + expect.objectContaining({ status: "ready", requestId: "host-operation" }), + ); + }, + }, + { + name: "inherits parent --port for gateway resume", + argv: ["gateway", "--port", "19087", "resume", "suspension-1", "--json"], + assert: () => { + expectLocalGatewayCall("gateway.suspend.resume", 19087, { + suspensionId: "suspension-1", + }); + expect(defaultRuntime.writeJson).toHaveBeenCalledWith({ + ok: true, + status: "running", + resumed: true, + }); + }, + }, { name: "forwards --token to gateway probe when parent and child option names collide", argv: ["gateway", "probe", "--token", "tok_probe", "--json"], diff --git a/src/cli/gateway-cli/register.ts b/src/cli/gateway-cli/register.ts index b10ea56eb717..4643b0a9e85c 100644 --- a/src/cli/gateway-cli/register.ts +++ b/src/cli/gateway-cli/register.ts @@ -27,6 +27,7 @@ import type { GatewayDiscoverOpts } from "./discover.js"; import { isGatewayMachineOutput } from "./output-mode.js"; import { addGatewayRestartHandoffCommands } from "./register-restart-handoff.js"; import { addGatewayRunCommand } from "./run-command.js"; +import { runGatewayResume, runGatewaySuspend } from "./suspend-cli.js"; type GatewayRpcOpts = Parameters[1]; @@ -595,6 +596,54 @@ export function registerGatewayCli(program: Command, deps: GatewayCliDependencie }), ); + gatewayCallOpts( + gateway + .command("suspend") + .description("Prepare the Gateway for cooperative host suspension") + .option("--request-id ", "Stable suspension request id") + .option("--wait ", "Wait up to this many seconds for active work to drain") + .option("--port ", "Local Gateway port") + .action(async (opts, command) => { + await runGatewayCommand( + async () => { + const rpcOpts = await resolveGatewayRpcOptionsWithLocalPort(opts, command); + await runGatewaySuspend( + { + rpcOpts, + requestId: opts.requestId, + waitSeconds: opts.wait, + json: Boolean(rpcOpts.json), + }, + { callGateway: callGatewayCli, runtime: defaultRuntime }, + ); + }, + "Gateway suspend failed", + { json: Boolean(opts.json) }, + ); + }), + ); + + gatewayCallOpts( + gateway + .command("resume") + .description("Release a cooperative Gateway suspension") + .argument("", "Suspension id returned by gateway suspend") + .option("--port ", "Local Gateway port") + .action(async (suspensionId, opts, command) => { + await runGatewayCommand( + async () => { + const rpcOpts = await resolveGatewayRpcOptionsWithLocalPort(opts, command); + await runGatewayResume( + { rpcOpts, suspensionId: String(suspensionId), json: Boolean(rpcOpts.json) }, + { callGateway: callGatewayCli, runtime: defaultRuntime }, + ); + }, + "Gateway resume failed", + { json: Boolean(opts.json) }, + ); + }), + ); + gatewayCallOpts( gateway .command("usage-cost") diff --git a/src/cli/gateway-cli/suspend-cli.test.ts b/src/cli/gateway-cli/suspend-cli.test.ts new file mode 100644 index 000000000000..a5193f7fa706 --- /dev/null +++ b/src/cli/gateway-cli/suspend-cli.test.ts @@ -0,0 +1,169 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { OutputRuntimeEnv } from "../../runtime.js"; +import { runGatewayResume, runGatewaySuspend } from "./suspend-cli.js"; + +function createRuntime(): OutputRuntimeEnv { + return { + log: vi.fn(), + error: vi.fn(), + writeStdout: vi.fn(), + writeJson: vi.fn(), + exit: vi.fn(), + }; +} + +const readyResult = { + status: "ready" as const, + suspensionId: "suspension-1", + expiresAtMs: Date.parse("2026-08-11T12:00:00.000Z"), + activeCount: 0, + blockers: [], +}; + +const busyResult = { + status: "busy" as const, + reason: "active-work" as const, + retryAfterMs: 200, + activeCount: 1, + blockers: [{ kind: "root-request" as const, count: 1, message: "1 active request" }], +}; + +describe("gateway suspend CLI", () => { + beforeEach(() => vi.clearAllMocks()); + + it("prints a ready lease with the default CLI request id", async () => { + const callGateway = vi.fn(async () => readyResult); + const runtime = createRuntime(); + + await runGatewaySuspend({ rpcOpts: {} }, { callGateway, runtime }); + + expect(callGateway).toHaveBeenCalledWith( + "gateway.suspend.prepare", + {}, + { requestId: expect.stringMatching(/^cli-[0-9a-f]{8}$/u) }, + ); + expect(callGateway).toHaveBeenCalledOnce(); + expect(runtime.log).toHaveBeenCalledWith("Gateway suspension prepared."); + expect(runtime.log).toHaveBeenCalledWith("Suspension ID: suspension-1"); + expect(runtime.log).toHaveBeenCalledWith( + `Expires: 2026-08-11T12:00:00.000Z (${readyResult.expiresAtMs} ms)`, + ); + expect(runtime.log).toHaveBeenCalledWith("Resume with: openclaw gateway resume suspension-1"); + }); + + it("reports blockers without polling when --wait is omitted", async () => { + const callGateway = vi.fn(async () => busyResult); + + await expect( + runGatewaySuspend( + { rpcOpts: {}, requestId: "host-operation" }, + { callGateway, runtime: createRuntime() }, + ), + ).rejects.toThrow( + "Gateway suspension is busy (active-work; 1 active).\nBlockers:\n- 1 active request\nRetry later or use --wait .", + ); + expect(callGateway).toHaveBeenCalledOnce(); + }); + + it("polls with one stable request id until the Gateway is ready", async () => { + const callGateway = vi + .fn() + .mockResolvedValueOnce(busyResult) + .mockResolvedValueOnce(readyResult); + let now = 1_000; + const sleep = vi.fn(async (delayMs: number) => { + now += delayMs; + }); + + await runGatewaySuspend( + { rpcOpts: {}, requestId: "host-operation", waitSeconds: "2" }, + { callGateway, runtime: createRuntime(), nowMs: () => now, sleep }, + ); + + expect(sleep).toHaveBeenCalledExactlyOnceWith(200); + expect(callGateway).toHaveBeenCalledTimes(2); + expect(callGateway.mock.calls.map((call) => call[2])).toEqual([ + { requestId: "host-operation" }, + { requestId: "host-operation" }, + ]); + }); + + it("emits the latest busy result and exits nonzero in JSON mode", async () => { + const runtime = createRuntime(); + + await runGatewaySuspend( + { rpcOpts: { json: true }, requestId: "host-operation", json: true }, + { callGateway: vi.fn(async () => busyResult), runtime }, + ); + + expect(runtime.writeJson).toHaveBeenCalledWith({ + ...busyResult, + requestId: "host-operation", + }); + expect(runtime.exit).toHaveBeenCalledWith(1); + }); + + it("never issues another prepare after a sleep overshoots the deadline", async () => { + let now = 1_000; + const callGateway = vi.fn(async () => busyResult); + + await expect( + runGatewaySuspend( + { rpcOpts: {}, requestId: "host-operation", waitSeconds: "0.2" }, + { + callGateway, + runtime: createRuntime(), + nowMs: () => now, + sleep: async () => { + // A lagging clock can wake far past the advertised --wait window. + now += 10_000; + }, + }, + ), + ).rejects.toThrow("Timed out waiting for the Gateway to become idle."); + expect(callGateway).toHaveBeenCalledOnce(); + }); + + it("reports the latest blockers when the wait deadline expires", async () => { + let now = 1_000; + + await expect( + runGatewaySuspend( + { rpcOpts: {}, requestId: "host-operation", waitSeconds: "0.1" }, + { + callGateway: vi.fn(async () => busyResult), + runtime: createRuntime(), + nowMs: () => now, + sleep: async (delayMs) => { + now += delayMs; + }, + }, + ), + ).rejects.toThrow( + "Gateway suspension is busy (active-work; 1 active).\nBlockers:\n- 1 active request\nTimed out waiting for the Gateway to become idle.", + ); + }); +}); + +describe("gateway resume CLI", () => { + it.each([ + { resumed: true, message: "Gateway resumed." }, + { + resumed: false, + message: + "No matching suspension was held (lease already expired or resumed); gateway is running.", + }, + ])("prints the resumed=$resumed outcome", async ({ resumed, message }) => { + const runtime = createRuntime(); + const callGateway = vi.fn(async () => ({ ok: true, status: "running", resumed })); + + await runGatewayResume({ rpcOpts: {}, suspensionId: "suspension-1" }, { callGateway, runtime }); + + expect(callGateway).toHaveBeenCalledExactlyOnceWith( + "gateway.suspend.resume", + {}, + { suspensionId: "suspension-1" }, + ); + expect(runtime.log).toHaveBeenCalledExactlyOnceWith(message); + }); +}); diff --git a/src/cli/gateway-cli/suspend-cli.ts b/src/cli/gateway-cli/suspend-cli.ts new file mode 100644 index 000000000000..4700b2b0b362 --- /dev/null +++ b/src/cli/gateway-cli/suspend-cli.ts @@ -0,0 +1,157 @@ +import { randomBytes } from "node:crypto"; +import type { + GatewaySuspendPrepareResult, + GatewaySuspendResumeResult, +} from "../../../packages/gateway-protocol/src/index.js"; +import { colorize, isRich, theme } from "../../../packages/terminal-core/src/theme.js"; +import type { OutputRuntimeEnv } from "../../runtime.js"; +import type { callGatewayFromCliWithTransport } from "../gateway-rpc.js"; + +type SuspendRpcOpts = Parameters[1]; + +type SuspendRpcCall = (method: string, opts: SuspendRpcOpts, params?: unknown) => Promise; + +type SuspendCliDeps = { + callGateway: SuspendRpcCall; + runtime: OutputRuntimeEnv; + nowMs?: () => number; + sleep?: (delayMs: number) => Promise; +}; + +const MIN_SUSPEND_POLL_DELAY_MS = 50; + +function parseWaitMs(value: string | number | undefined): number | undefined { + if (value === undefined) { + return undefined; + } + const seconds = typeof value === "number" ? value : Number(value.trim()); + if (!Number.isFinite(seconds) || seconds < 0) { + throw new Error("--wait must be a non-negative number of seconds"); + } + const milliseconds = Math.floor(seconds * 1_000); + if (!Number.isSafeInteger(milliseconds)) { + throw new Error("--wait is too large"); + } + return milliseconds; +} + +function resolveRequestId(value: string | undefined): string { + if (value === undefined) { + return `cli-${randomBytes(4).toString("hex")}`; + } + const requestId = value.trim(); + if (!requestId || requestId.length > 128) { + throw new Error("--request-id must contain 1 to 128 characters"); + } + return requestId; +} + +function formatBusyResult( + result: Extract, +): string { + const blockers = result.blockers.map((blocker) => `- ${blocker.message}`); + return [ + `Gateway suspension is busy (${result.reason}; ${result.activeCount} active).`, + ...(blockers.length > 0 ? ["Blockers:", ...blockers] : []), + ].join("\n"); +} + +function writeSuspendJson( + runtime: OutputRuntimeEnv, + result: GatewaySuspendPrepareResult, + requestId: string, +): void { + runtime.writeJson({ ...result, requestId }); +} + +export async function runGatewaySuspend( + options: { + rpcOpts: SuspendRpcOpts; + requestId?: string; + waitSeconds?: string | number; + json?: boolean; + }, + deps: SuspendCliDeps, +): Promise { + const nowMs = deps.nowMs ?? Date.now; + const sleep = + deps.sleep ?? + (async (delayMs: number) => + await new Promise((resolve) => { + setTimeout(resolve, delayMs); + })); + const requestId = resolveRequestId(options.requestId); + const waitMs = parseWaitMs(options.waitSeconds); + const deadlineMs = waitMs === undefined ? undefined : nowMs() + waitMs; + const maxAttempts = waitMs === undefined ? 1 : Math.ceil(waitMs / MIN_SUSPEND_POLL_DELAY_MS) + 1; + let latest: GatewaySuspendPrepareResult | undefined; + + for (let attempt = 0; attempt < maxAttempts; attempt += 1) { + // A sleep can overshoot the deadline; never issue a prepare that could + // suspend the Gateway after the operator's advertised --wait window. + if (attempt > 0 && deadlineMs !== undefined && nowMs() >= deadlineMs) { + break; + } + latest = (await deps.callGateway("gateway.suspend.prepare", options.rpcOpts, { + requestId, + })) as GatewaySuspendPrepareResult; + if (latest.status === "ready") { + if (options.json) { + writeSuspendJson(deps.runtime, latest, requestId); + return; + } + const rich = isRich(); + deps.runtime.log(colorize(rich, theme.success, "Gateway suspension prepared.")); + deps.runtime.log(`${colorize(rich, theme.muted, "Suspension ID:")} ${latest.suspensionId}`); + deps.runtime.log( + `${colorize(rich, theme.muted, "Expires:")} ${new Date(latest.expiresAtMs).toISOString()} (${latest.expiresAtMs} ms)`, + ); + deps.runtime.log(`Resume with: openclaw gateway resume ${latest.suspensionId}`); + return; + } + + if (deadlineMs === undefined) { + if (options.json) { + writeSuspendJson(deps.runtime, latest, requestId); + deps.runtime.exit(1); + return; + } + throw new Error(`${formatBusyResult(latest)}\nRetry later or use --wait .`); + } + + const remainingMs = deadlineMs - nowMs(); + if (remainingMs <= 0) { + break; + } + const delayMs = Math.min(remainingMs, Math.max(MIN_SUSPEND_POLL_DELAY_MS, latest.retryAfterMs)); + await sleep(delayMs); + } + + if (!latest || latest.status !== "busy") { + throw new Error("Gateway suspension polling ended without a result"); + } + if (options.json) { + writeSuspendJson(deps.runtime, latest, requestId); + deps.runtime.exit(1); + return; + } + throw new Error(`${formatBusyResult(latest)}\nTimed out waiting for the Gateway to become idle.`); +} + +export async function runGatewayResume( + options: { rpcOpts: SuspendRpcOpts; suspensionId: string; json?: boolean }, + deps: Pick, +): Promise { + const result = (await deps.callGateway("gateway.suspend.resume", options.rpcOpts, { + suspensionId: options.suspensionId, + })) as GatewaySuspendResumeResult; + if (options.json) { + deps.runtime.writeJson(result); + return; + } + deps.runtime.log( + result.resumed + ? "Gateway resumed." + : "No matching suspension was held (lease already expired or resumed); gateway is running.", + ); +} diff --git a/src/gateway/call.test.ts b/src/gateway/call.test.ts index e8a87306eaf3..d27a71674e46 100644 --- a/src/gateway/call.test.ts +++ b/src/gateway/call.test.ts @@ -1703,6 +1703,34 @@ describe("callGateway error details", () => { }); }); + it("surfaces a websocket upgrade rejection carried by close info", async () => { + startMode = "silent"; + setLocalLoopbackGatewayConfig(); + const upgradeError = Object.assign( + new Error( + "gateway rejected websocket upgrade (HTTP 503): Gateway websocket admission closed", + ), + { + name: "GatewayClientRequestError", + gatewayCode: "UNAVAILABLE", + details: { reason: "websocket-upgrade-rejected", httpStatus: 503 }, + retryable: true, + }, + ); + + const request = callGateway({ method: "health" }); + await waitForFast(() => expect(lastClientOptions).not.toBeNull()); + lastClientOptions?.onClose?.(1006, "", { + phase: "pre-hello", + socketOpened: false, + transportValidated: false, + transientPreHelloCleanClose: false, + connectError: upgradeError, + }); + + await expect(request).rejects.toBe(upgradeError); + }); + it.each([ { name: "another structured auth rejection", diff --git a/src/gateway/call.ts b/src/gateway/call.ts index dd18b956761d..df7a9bade273 100644 --- a/src/gateway/call.ts +++ b/src/gateway/call.ts @@ -984,6 +984,11 @@ async function executeGatewayRequestWithScopes(params: { if (settled || ignoreClose) { return; } + if (info?.connectError) { + ignoreClose = true; + stop(info.connectError); + return; + } if ( !primaryRequestStarted && info?.transientPreHelloCleanClose === true && diff --git a/src/gateway/client.test.ts b/src/gateway/client.test.ts index b07ae3a0e43a..a2fee29aa28e 100644 --- a/src/gateway/client.test.ts +++ b/src/gateway/client.test.ts @@ -809,6 +809,7 @@ describe("GatewayClient close handling", () => { expect.objectContaining({ message: "gateway tls fingerprint mismatch" }), ); expect(onClose).toHaveBeenCalledWith(1008, "gateway tls fingerprint mismatch", { + connectError: expect.objectContaining({ message: "gateway tls fingerprint mismatch" }), phase: "pre-hello", socketOpened: true, transportValidated: false, @@ -2591,6 +2592,11 @@ describe("GatewayClient connect auth payload", () => { "gateway client reconnect paused handler error: Error: paused callback failed", ); expect(onClose).toHaveBeenCalledWith(1008, "connect failed", { + connectError: expect.objectContaining({ + details: { code: "AUTH_TOKEN_MISSING" }, + gatewayCode: "INVALID_REQUEST", + message: "unauthorized", + }), phase: "pre-hello", socketOpened: true, transportValidated: true, diff --git a/src/gateway/probe.device-auth-scope.test.ts b/src/gateway/probe.device-auth-scope.test.ts index 2d64190c7569..deda900418d0 100644 --- a/src/gateway/probe.device-auth-scope.test.ts +++ b/src/gateway/probe.device-auth-scope.test.ts @@ -7,7 +7,7 @@ import { loadOrCreateDeviceIdentity } from "../infra/device-identity.js"; import { closeOpenClawStateDatabaseForTest } from "../state/openclaw-state-db.js"; import { withTempDir } from "../test-utils/temp-dir.js"; -type WebSocketEvent = "open" | "message" | "close" | "error"; +type WebSocketEvent = "open" | "message" | "close" | "error" | "unexpected-response"; const webSockets = vi.hoisted((): ProbeWebSocket[] => []); @@ -25,6 +25,7 @@ class ProbeWebSocket { message: [], close: [], error: [], + "unexpected-response": [], }; constructor(_url: string, _options?: unknown) { diff --git a/src/gateway/server-http.ts b/src/gateway/server-http.ts index 3daacb523957..8d3307ecb8a2 100644 --- a/src/gateway/server-http.ts +++ b/src/gateway/server-http.ts @@ -18,7 +18,11 @@ import { createDiagnosticTraceContext, runWithDiagnosticTraceContext, } from "../infra/diagnostic-trace-context.js"; -import { isGatewayWorkAdmissionClosed } from "../process/gateway-work-admission.js"; +import { + getGatewaySuspendAdmissionPhase, + isGatewayRestartDraining, + isGatewayWorkAdmissionClosed, +} from "../process/gateway-work-admission.js"; import { createLazyRuntimeModule } from "../shared/lazy-runtime.js"; import { resolveAssistantIdentity } from "./assistant-identity.js"; import type { AuthRateLimiter } from "./auth-rate-limit.js"; @@ -782,7 +786,12 @@ function handleBudgetedGatewayWebSocketUpgrade(params: { prepareSocket?: (socket: GatewayIngressWebSocket) => void; }): void { const { req, socket, head, wss, preauthConnectionBudget, preauthBudgetKey, ingressName } = params; - if (isGatewayWorkAdmissionClosed()) { + if ( + isGatewayWorkAdmissionClosed() && + (ingressName === "Worker" || + isGatewayRestartDraining() || + getGatewaySuspendAdmissionPhase() !== "prepared") + ) { writeGatewayUpgradeServiceUnavailable(socket, `${ingressName} websocket admission closed`); socket.destroy(); return; @@ -962,8 +971,7 @@ export function attachGatewayUpgradeHandler(opts: { return; } // Plugin-owned upgrade routes have already had the opportunity to claim the socket. - // Core Gateway upgrades must stop at the HTTP boundary so a client cannot hold an - // untracked pre-connect socket after suspension or restart admission closes. + // Core Gateway control connections remain reachable while suspension is prepared. try { handleBudgetedGatewayWebSocketUpgrade({ req, diff --git a/src/gateway/server.preauth-hardening.test.ts b/src/gateway/server.preauth-hardening.test.ts index 3cce6ab91376..d18ae25f3291 100644 --- a/src/gateway/server.preauth-hardening.test.ts +++ b/src/gateway/server.preauth-hardening.test.ts @@ -9,7 +9,11 @@ import { resetDiagnosticEventsForTest, type DiagnosticEventPayload, } from "../infra/diagnostic-events.js"; -import { tryBeginGatewaySuspendAdmission } from "../process/gateway-work-admission.js"; +import { + markGatewayRestartDraining, + resetGatewayWorkAdmission, + tryBeginGatewaySuspendAdmission, +} from "../process/gateway-work-admission.js"; import { captureEnv, setTestEnvValue } from "../test-utils/env.js"; import type { ResolvedGatewayAuth } from "./auth.js"; import { MAX_PREAUTH_PAYLOAD_BYTES } from "./server-constants.js"; @@ -43,6 +47,7 @@ const PREAUTH_HANDSHAKE_TEST_CLOSE_LIMIT_MS = 5_000; const cleanupEnv: Array<() => void> = []; afterEach(async () => { + resetGatewayWorkAdmission(); while (cleanupEnv.length > 0) { cleanupEnv.pop()?.(); } @@ -152,6 +157,39 @@ describe("gateway pre-auth hardening", () => { } }); + it("rejects worker websocket upgrades after suspension is prepared", async () => { + const httpServer = http.createServer(); + const wss = new WebSocketServer({ maxPayload: 1024, noServer: true }); + wss.on("connection", (socket) => socket.close()); + attachWorkerGatewayUpgradeHandler({ + httpServer, + wss, + preauthConnectionBudget: createPreauthConnectionBudget(1), + }); + await new Promise((resolve) => { + httpServer.listen(0, "127.0.0.1", resolve); + }); + const address = httpServer.address(); + const port = typeof address === "object" && address ? address.port : 0; + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + + try { + await expect(requestUpgradeRejection(port)).resolves.toEqual({ + status: 503, + body: "Worker websocket admission closed", + }); + } finally { + suspension?.release(); + await new Promise((resolve) => { + wss.close(() => resolve()); + }); + await new Promise((resolve, reject) => { + httpServer.close((error) => (error ? reject(error) : resolve())); + }); + } + }); + it("rejects upgrades before websocket handlers attach (pre-auth budget enforced, then released)", async () => { const clients = new Set(); const resolvedAuth: ResolvedGatewayAuth = { mode: "none", allowTailscale: false }; @@ -196,18 +234,49 @@ describe("gateway pre-auth hardening", () => { } }); - it("rejects core websocket upgrades while suspension admission is closed", async () => { + it("accepts core websocket upgrades after suspension is prepared", async () => { const harness = await createGatewaySuiteHarness(); const suspension = tryBeginGatewaySuspendAdmission(() => {}); expect(suspension?.commit()).toBe(true); + try { + const ws = await harness.openWs(); + await expect(readConnectChallengeNonce(ws)).resolves.toEqual(expect.any(String)); + ws.close(); + await new Promise((resolve) => { + ws.once("close", () => resolve()); + }); + } finally { + suspension?.release(); + await harness.close(); + } + }); + + it("rejects core websocket upgrades while suspension is preparing", async () => { + const harness = await createGatewaySuiteHarness(); + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + + try { + await expect(requestUpgradeRejection(harness.port)).resolves.toEqual({ + status: 503, + body: "Gateway websocket admission closed", + }); + } finally { + suspension?.rollback(); + await harness.close(); + } + }); + + it("rejects core websocket upgrades during restart drain", async () => { + const harness = await createGatewaySuiteHarness(); + markGatewayRestartDraining(); + try { await expect(requestUpgradeRejection(harness.port)).resolves.toEqual({ status: 503, body: "Gateway websocket admission closed", }); } finally { - suspension?.release(); await harness.close(); } }); 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 8868483ad56c..89a211edf947 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 @@ -5,6 +5,7 @@ import type { WebSocket } from "ws"; import { PROTOCOL_VERSION } from "../../../../packages/gateway-protocol/src/index.js"; import { getActiveGatewayRootWorkCount, + markGatewayRestartDraining, resetGatewayWorkAdmission, tryBeginGatewaySuspendAdmission, } from "../../../process/gateway-work-admission.js"; @@ -149,6 +150,27 @@ function attachHarness(params: { deferSocketSend?: boolean } = {}) { }, }), ), + sendNodeConnect: () => + onMessage?.( + JSON.stringify({ + type: "req", + id: "node-connect-1", + method: "connect", + params: { + minProtocol: PROTOCOL_VERSION, + maxProtocol: PROTOCOL_VERSION, + client: { + id: "gateway-client", + version: "dev", + platform: "test", + mode: "backend", + }, + role: "node", + scopes: [], + caps: [], + }, + }), + ), sendWorkerConnect: () => onMessage?.( JSON.stringify({ @@ -173,54 +195,106 @@ beforeEach(() => { afterEach(resetGatewayWorkAdmission); describe("WebSocket connect suspension admission", () => { - it.each(["preparing", "prepared"] as const)( - "rejects a validated connect while suspension is %s before session mutations", - async (phase) => { - const suspension = tryBeginGatewaySuspendAdmission(() => {}); - expect(suspension).not.toBeNull(); - if (phase === "prepared") { - expect(suspension?.commit()).toBe(true); - } - const harness = attachHarness(); + it("rejects a validated connect while suspension is preparing before session mutations", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension).not.toBeNull(); + const harness = attachHarness(); - harness.sendConnect(); + harness.sendConnect(); - await vi.waitFor(() => { - expect(harness.socketSend).toHaveBeenCalledOnce(); - }); - const response = JSON.parse(harness.socketSend.mock.calls[0]?.[0] ?? "{}") as { - error?: { - code?: string; - retryable?: boolean; - retryAfterMs?: number; - details?: Record; - }; + await vi.waitFor(() => { + expect(harness.socketSend).toHaveBeenCalledOnce(); + }); + const response = JSON.parse(harness.socketSend.mock.calls[0]?.[0] ?? "{}") as { + error?: { + code?: string; + retryable?: boolean; + retryAfterMs?: number; + details?: Record; }; - expect(response.error).toMatchObject({ - code: "UNAVAILABLE", - retryable: true, - retryAfterMs: 1_000, - details: { - method: "connect", - reason: "gateway-suspending", - phase, - }, - }); - expect(harness.client).toBeNull(); - expect(harness.setClient).not.toHaveBeenCalled(); - expect(upsertPresenceMock).not.toHaveBeenCalled(); - expect(incrementPresenceVersionMock).not.toHaveBeenCalled(); - await vi.waitFor(() => { - expect(harness.close).toHaveBeenCalledWith(1013, "gateway suspension in progress"); - }); + }; + expect(response.error).toMatchObject({ + code: "UNAVAILABLE", + retryable: true, + retryAfterMs: 1_000, + details: { + method: "connect", + reason: "gateway-suspending", + phase: "preparing", + }, + }); + expect(harness.client).toBeNull(); + expect(harness.setClient).not.toHaveBeenCalled(); + expect(upsertPresenceMock).not.toHaveBeenCalled(); + expect(incrementPresenceVersionMock).not.toHaveBeenCalled(); + await vi.waitFor(() => { + expect(harness.close).toHaveBeenCalledWith(1013, "gateway suspension in progress"); + }); + suspension?.rollback(); + }); - if (phase === "prepared") { - suspension?.release(); - } else { - suspension?.rollback(); - } - }, - ); + it("accepts a validated connect while suspension is prepared", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + const harness = attachHarness(); + + harness.sendConnect(); + + await vi.waitFor(() => { + expect(harness.setClient).toHaveBeenCalledOnce(); + }); + expect(harness.client).not.toBeNull(); + expect(harness.close).not.toHaveBeenCalled(); + suspension?.release(); + }); + + it("rejects a node connect while suspension is prepared", async () => { + const suspension = tryBeginGatewaySuspendAdmission(() => {}); + expect(suspension?.commit()).toBe(true); + const harness = attachHarness(); + + harness.sendNodeConnect(); + + await vi.waitFor(() => { + expect(harness.socketSend).toHaveBeenCalledOnce(); + }); + const response = JSON.parse(harness.socketSend.mock.calls[0]?.[0] ?? "{}") as { + error?: { details?: Record }; + }; + expect(response.error?.details).toMatchObject({ + method: "connect", + reason: "gateway-suspending", + phase: "prepared", + }); + expect(harness.setClient).not.toHaveBeenCalled(); + expect(upsertPresenceMock).not.toHaveBeenCalled(); + await vi.waitFor(() => { + expect(harness.close).toHaveBeenCalledWith(1013, "gateway suspension in progress"); + }); + suspension?.release(); + }); + + it("rejects a validated connect during restart drain", async () => { + markGatewayRestartDraining(); + const harness = attachHarness(); + + harness.sendConnect(); + + await vi.waitFor(() => { + expect(harness.socketSend).toHaveBeenCalledOnce(); + }); + const response = JSON.parse(harness.socketSend.mock.calls[0]?.[0] ?? "{}") as { + error?: { details?: Record }; + }; + expect(response.error?.details).toMatchObject({ + method: "connect", + reason: "gateway-restarting", + }); + expect(harness.setClient).not.toHaveBeenCalled(); + await vi.waitFor(() => { + expect(harness.close).toHaveBeenCalledWith(1013, "gateway restart in progress"); + }); + }); it("keeps an accepted handshake visible as root work until hello is sent", async () => { const harness = attachHarness({ deferSocketSend: true }); diff --git a/src/gateway/server/ws-connection/message-handler.ts b/src/gateway/server/ws-connection/message-handler.ts index 560c034b559d..6aa2d785801b 100644 --- a/src/gateway/server/ws-connection/message-handler.ts +++ b/src/gateway/server/ws-connection/message-handler.ts @@ -402,21 +402,38 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar } }; - const rejectConnectForClosedAdmission = async (data: RawData): Promise => { + const parsePreauthConnectFrame = (data: RawData) => { if (isClosed() || rawDataByteLength(data) > MAX_PREAUTH_PAYLOAD_BYTES) { - return false; + return null; } let parsed: unknown; try { parsed = JSON.parse(rawDataToString(data)); } catch { - return false; + return null; } if ( !validateRequestFrame(parsed) || parsed.method !== "connect" || !validateConnectParams(parsed.params) ) { + return null; + } + return parsed; + }; + + const isPreparedControlConnect = (data: RawData): boolean => { + const parsed = parsePreauthConnectFrame(data); + if (!parsed) { + return false; + } + const connectParams = parsed.params as { role?: unknown }; + return connectParams.role !== "node" && !claimsWorkerConnectionIdentity(parsed.params); + }; + + const rejectConnectForClosedAdmission = async (data: RawData): Promise => { + const parsed = parsePreauthConnectFrame(data); + if (!parsed) { return false; } @@ -457,6 +474,18 @@ export function attachGatewayWsMessageHandler(params: GatewayWsMessageHandlerPar } const admission = tryBeginGatewayRootWorkAdmission(); if (!admission) { + if ( + !isGatewayRestartDraining() && + getGatewaySuspendAdmissionPhase() === "prepared" && + isPreparedControlConnect(data) + ) { + // Refuse-only suspension fences work, not control-plane visibility. Only + // operator connects are admitted while prepared, and they can only reach + // suspend-control methods after handshake; node and worker connects would + // attach presence/registry state, so they stay refused. + await handleMessage(data); + return; + } if (await rejectConnectForClosedAdmission(data)) { return; }